Update model.py

model.py

@@ -1,50 +1,39 @@
-import streamlit as st
+import torch
+from diffusers import DiffusionPipeline
+import trimesh
+import numpy as np
 from PIL import Image
-from model import load_pipeline, generate_3d_model, convert_to_gif
-import os
+from io import BytesIO
 
-# Load pipeline
-pipe = load_pipeline()
+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
 
-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 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
 
-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"
-            )
+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