Rename app.py to model.py

app.py

@@ -1,46 +0,0 @@
-import torch
-import numpy as np
-from diffusers import DiffusionPipeline
-import streamlit as st
-
-# 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:
-        pipeline.save_ply(result, output_path)
-        print(f"Model saved to {output_path}")
-        return output_path
-    except Exception as e:
-        print(f"Failed to save model: {e}")
-        return None
-
-# 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.")

model.py

@@ -0,0 +1,24 @@
+import torch
+from diffusers import ShapEPipeline
+from diffusers.utils import export_to_gif
+
+# Define checkpoint ID and load pipeline on CPU
+ckpt_id = "openai/shap-e"
+pipe = ShapEPipeline.from_pretrained(ckpt_id).to("cpu")
+
+# Define generation parameters
+guidance_scale = 10.0  # Lowered for efficiency on CPU
+num_inference_steps = 32  # Reduced steps for CPU performance
+prompt = "a shark"
+
+# Generate images from the prompt with reduced settings
+images = pipe(
+    prompt=prompt,
+    guidance_scale=guidance_scale,
+    num_inference_steps=num_inference_steps,
+    size=256,  # Image size for the model
+).images
+
+# Export images to GIF format
+gif_path = export_to_gif(images, "shark_3d.gif")
+print(f"GIF saved at {gif_path}")