Create generate.py

generate.py

@@ -0,0 +1,26 @@
+import torch
+from PIL import Image
+from shap_e.diffusion.sample import sample_latents
+from shap_e.diffusion.eval.image_embedder import CLIPImageEmbedder
+from shap_e.models.download import load_model
+from shap_e.util.notebooks import decode_latent_mesh, create_pan_cameras, render_views
+
+device = torch.device('cpu')  # Spaces use CPU
+
+# Load models
+model = load_model('text300M', device=device)
+image_embedder = CLIPImageEmbedder(device=device)
+
+def generate_model_from_text(prompt: str, filename="model.glb"):
+    batch = [prompt]
+    latents = sample_latents(model, batch, guidance_scale=15.0, device=device)
+    mesh = decode_latent_mesh(latents[0], device=device)
+    mesh.export(filename)
+    return filename
+
+def generate_model_from_image(image: Image.Image, filename="model.glb"):
+    image_embed = image_embedder.embed_image(image)
+    latents = sample_latents(model, [image_embed], guidance_scale=15.0, device=device, is_image=True)
+    mesh = decode_latent_mesh(latents[0], device=device)
+    mesh.export(filename)
+    return filename