Update model.py

model.py

@@ -1,39 +1,50 @@
-import torch
-from diffusers import DiffusionPipeline
-import trimesh
-import numpy as np
+import streamlit as st
 from PIL import Image
-from io import BytesIO
+from model import load_pipeline, generate_3d_model, convert_to_gif
+import os
 
-def load_pipeline():
-    """
-    Load the stable-zero123 model pipeline from Hugging Face.
-    """
-    ckpt_id = "stabilityai/stable-zero123"
-    pipe = DiffusionPipeline.from_pretrained(ckpt_id, torch_dtype=torch.float32).to("cpu")
-    return pipe
+# Load pipeline
+pipe = load_pipeline()
 
-def generate_3d_model(pipe, prompt, output_path="output.obj", guidance_scale=7.5, num_inference_steps=32):
-    """
-    Generate a 3D model from the prompt and save it in a Blender-compatible format (.obj).
-    """
-    # Generate the model output
-    outputs = pipe(prompt=prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps)
-    
-    # Extract mesh data if the output structure allows
-    vertices = outputs["vertices"][0].detach().cpu().numpy()
-    faces = outputs["faces"][0].detach().cpu().numpy()
-    
-    # Create and save the mesh using trimesh
-    mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=True)
-    mesh.export(output_path)
-    return output_path
+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"
 
-def convert_to_gif(images, gif_path="output.gif"):
-    """
-    Convert a list of images into a GIF.
-    """
-    images[0].save(
-        gif_path, save_all=True, append_images=images[1:], loop=0, duration=100
-    )
-    return gif_path
+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"
+            )