Update gradio_app.py

gradio_app.py

@@ -67,11 +67,11 @@ def create_batch(input_image: Image) -> dict[str, Any]:
     }
     return batch
 
-def generate_and_process_3d(prompt: str, seed: int = 42, width: int = 1024, height: int = 1024) -> str:
+def generate_and_process_3d(prompt: str, seed: int = 42, width: int = 1024, height: int = 1024) -> tuple[str, Image.Image]:
     """Generate image from prompt and convert to 3D model."""
     try:
         # Generate image using FLUX
-        generator = torch.Generator().manual_seed(seed)
+        generator = torch.Generator(device=device).manual_seed(seed)
         generated_image = flux_pipe(
             prompt=prompt,
             width=width,
@@ -84,12 +84,13 @@ def generate_and_process_3d(prompt: str, seed: int = 42, width: int = 1024, heig
         # Convert PIL image to RGBA
         input_image = generated_image.convert("RGBA")
         
-        # Remove background if needed
-        input_image = bg_remover.process(input_image.convert("RGB"))
+        # Remove background
+        rgba_image = bg_remover.process(input_image.convert("RGB"))
+        rgba_image.putalpha(255)  # Add alpha channel
         
         # Auto crop
         input_image = spar3d_utils.foreground_crop(
-            input_image,
+            rgba_image,
             crop_ratio=1.3,
             newsize=(COND_WIDTH, COND_HEIGHT),
             no_crop=False
@@ -101,7 +102,7 @@ def generate_and_process_3d(prompt: str, seed: int = 42, width: int = 1024, heig
 
         # Generate mesh
         with torch.no_grad():
-            with torch.autocast(device_type=device, dtype=torch.bfloat16):
+            with torch.autocast(device_type='cuda' if torch.cuda.is_available() else 'cpu', dtype=torch.bfloat16):
                 trimesh_mesh, _ = spar3d_model.generate_mesh(
                     batch,
                     1024,  # texture_resolution
@@ -112,16 +113,17 @@ def generate_and_process_3d(prompt: str, seed: int = 42, width: int = 1024, heig
                 trimesh_mesh = trimesh_mesh[0]
 
         # Export to GLB
-        temp_file = tempfile.NamedTemporaryFile(suffix='.glb', delete=False)
-        trimesh_mesh.export(temp_file.name, file_type="glb", include_normals=True)
+        temp_dir = tempfile.mkdtemp()
+        output_path = os.path.join(temp_dir, 'output.glb')
+        trimesh_mesh.export(output_path, file_type="glb", include_normals=True)
         
-        return temp_file.name, generated_image
+        return output_path, generated_image
         
     except Exception as e:
-        return str(e), None
-
-
+        print(f"Error: {str(e)}")
+        return None, None
 
+# Create Gradio interface
 demo = gr.Interface(
     fn=generate_and_process_3d,
     inputs=[
@@ -153,8 +155,8 @@ demo = gr.Interface(
     ],
     outputs=[
         gr.File(
-            label="Download GLB",
-            file_types=[".glb"],
+            label="Download 3D Model",
+            file_types=[".glb"]
         ),
         gr.Image(
             label="Generated Image",
@@ -166,4 +168,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue().launch()