import streamlit as st
from PIL import Image
from model import load_pipeline, generate_3d_model, convert_to_gif
import os

# Load pipeline
pipe = load_pipeline()

st.title("3D Model Generator with GIF Export")
prompt = st.text_input("Enter a prompt for the 3D model:", "a shark")
output_obj_path = "generated_3d_model.obj"
output_gif_path = "generated_3d_model.gif"

if st.button("Generate 3D Model and Convert to GIF"):
    with st.spinner("Generating 3D model..."):
        obj_path = generate_3d_model(pipe, prompt, output_path=output_obj_path)
        
        # Load the 3D model in trimesh for rendering as images
        mesh = trimesh.load(obj_path)
        images = []
        
        # Rotate and render views for GIF
        for angle in range(0, 360, 30):  # Capture images every 30 degrees
            scene = mesh.scene()
            scene.set_camera(angles=(angle, angle))
            data = scene.save_image(resolution=(256, 256))
            images.append(Image.open(BytesIO(data)))
        
        # Convert to GIF
        gif_path = convert_to_gif(images, gif_path=output_gif_path)
        
        # Display GIF
        st.image(gif_path, caption="Generated 3D Model GIF")
        
        # Provide download options
        with open(gif_path, "rb") as file:
            st.download_button(
                label="Download GIF",
                data=file,
                file_name="generated_3d_model.gif",
                mime="image/gif"
            )

        with open(obj_path, "rb") as file:
            st.download_button(
                label="Download 3D Model (.obj)",
                data=file,
                file_name="generated_3d_model.obj",
                mime="application/octet-stream"
            )