Update app.py

app.py

@@ -1,57 +1,46 @@
 import torch
-from diffusers import ShapEPipeline
-import trimesh
 import numpy as np
+from diffusers import DiffusionPipeline
+import streamlit as st
 
-def generate_3d_model(prompt, output_path="assistant_3d.obj"):
-    """
-    Generate a 3D model using ShapE and export it in a Blender-compatible format
-    """
+# Load the ShapE pipeline on CPU
+pipeline = DiffusionPipeline.from_pretrained(
+    "openai/shap-e",
+    torch_dtype=torch.float32,
+    trust_remote_code=True,
+    custom_pipeline="openai/shap-e",  # Assuming it works with custom_pipeline param
+).to("cpu")
+
+# Define the function to generate and save a 3D model
+def generate_3d_model(prompt, output_path="/tmp/output.ply"):
+    # Run the pipeline with a text prompt
+    result = pipeline(prompt, None)
+    
+    # Try to save the result as a 3D model
     try:
-        # Load pipeline optimized for CPU usage
-        pipe = ShapEPipeline.from_pretrained(
-            "openai/shap-e",
-            torch_dtype=torch.float32,
-            low_cpu_mem_usage=True
-        ).to("cpu")
-        
-        # Generate model
-        outputs = pipe(
-            prompt=prompt,
-            num_inference_steps=8,  # Reduced for CPU efficiency
-            guidance_scale=5.0      # Lower guidance to conserve resources
-        )
-        
-        # Check output structure for debugging
-        print("Output structure:", outputs.keys())
-        
-        # Attempt to extract vertices and faces
-        try:
-            vertices = outputs["vertices"][0].detach().cpu().numpy()
-            faces = outputs["faces"][0].detach().cpu().numpy()
-        except KeyError as ke:
-            print(f"Key error: {ke}")
-            print(f"Available keys in output: {list(outputs.keys())}")
-            return None
-        
-        # Construct the 3D mesh object
-        mesh_obj = trimesh.Trimesh(vertices=vertices, faces=faces, process=True)
-        
-        # Export mesh object
-        if output_path.endswith('.obj'):
-            mesh_obj.export(output_path, include_normals=True)
-        else:
-            mesh_obj.export(output_path)
-        print(f"Successfully exported 3D model to: {output_path}")
+        pipeline.save_ply(result, output_path)
+        print(f"Model saved to {output_path}")
         return output_path
-        
     except Exception as e:
-        print(f"Error during generation: {e}")
-        raise
+        print(f"Failed to save model: {e}")
+        return None
 
-if __name__ == "__main__":
-    prompt = "a simple 3D ring, perfect circle, clean geometry"
-    try:
-        generate_3d_model(prompt, "assistant_3d.obj")
-    except Exception as e:
-        print(f"Generation failed: {e}")
+# Streamlit interface
+st.title("3D Model Generator")
+prompt = st.text_input("Enter a prompt to generate a 3D model:", "a cat statue")
+
+if st.button("Generate Model"):
+    with st.spinner("Generating model..."):
+        model_path = generate_3d_model(prompt)
+        if model_path:
+            st.success("Model generated successfully!")
+            # Display download link
+            with open(model_path, "rb") as file:
+                st.download_button(
+                    label="Download 3D Model",
+                    data=file,
+                    file_name="generated_model.ply",
+                    mime="application/octet-stream"
+                )
+        else:
+            st.error("Model generation failed.")