Adapt to HF ZeroGPU

app.py

@@ -9,8 +9,9 @@ from pathlib import Path
 from huggingface_hub import hf_hub_download
 import zipfile
 from datetime import datetime
+import spaces
 
-from infer_asset import infer_single_asset
+from infer_asset import infer_single_asset, save_articulated_meshes
 from particulate.models import Articulate3D_B
 from particulate.data_utils import load_obj_raw_preserve
 from particulate.export_utils import export_urdf, export_mjcf
@@ -31,15 +32,11 @@ class ParticulateApp:
         self.model = Articulate3D_B(**self.model_config)
         self.model.eval()
         
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        if device == "cpu":
-            print("WARNING: CUDA is not available. This application requires CUDA for full functionality as infer_asset.py assumes CUDA.")
-            # We attempt to use CUDA anyway because infer_asset.py hardcodes it in prepare_inputs
-        
+        # Always load to CPU initially to support Zero GPU and avoid VRAM usage when idle
         print("Downloading/Loading model from Hugging Face...")
         self.model_checkpoint = hf_hub_download(repo_id="rayli/Particulate", filename="model.pt")
-        self.model.load_state_dict(torch.load(self.model_checkpoint, map_location="cuda" if torch.cuda.is_available() else "cpu"))
-        self.model.to('cuda')
+        self.model.load_state_dict(torch.load(self.model_checkpoint, map_location="cpu"))
+        # Model stays on CPU until inference
 
         model_dir = os.path.join("PartField", "model")
         os.makedirs(model_dir, exist_ok=True)
@@ -78,9 +75,16 @@ class ParticulateApp:
     ):
         if mesh is None:
             return None, "Please upload a 3D model."
-        
-        with tempfile.TemporaryDirectory() as temp_dir:
-            try:
+            
+        try:
+            outputs, face_indices = self._predict_impl(
+                mesh,
+                min_part_confidence,
+                num_points,
+                up_dir
+            )
+
+            with tempfile.TemporaryDirectory() as temp_dir:
                 (
                     mesh_parts_original,
                     unique_part_ids,
@@ -91,15 +95,11 @@ class ParticulateApp:
                     revolute_range,
                     prismatic_axis,
                     prismatic_range
-                ) = infer_single_asset(
-                    mesh=mesh,
-                    up_dir=up_dir,
-                    model=self.model,
-                    num_points=int(num_points),
+                ) = save_articulated_meshes(
+                    mesh, face_indices, outputs,
                     strict=strict,
-                    output_path=temp_dir,
                     animation_frames=int(animation_frames),
-                    min_part_confidence=min_part_confidence,
+                    output_path=temp_dir
                 )
 
                 animated_glb_file = os.path.join(temp_dir, "animated_textured.glb")
@@ -130,13 +130,29 @@ class ParticulateApp:
                         *[None] * 9
                     )
 
-            except Exception as e:
-                import traceback
-                traceback.print_exc()
-                return (
-                    None, None, f"Error: {str(e)}",
-                    *[None] * 9
-                )
+        except Exception as e:
+            import traceback
+            traceback.print_exc()
+            return (
+                None, None, f"Error: {str(e)}",
+                *[None] * 9
+            )
+
+    @spaces.GPU
+    def _predict_impl(
+        self,
+        mesh,
+        min_part_confidence,
+        num_points,
+        up_dir
+    ):
+        return infer_single_asset(
+            mesh=mesh,
+            up_dir=up_dir,
+            model=self.model.to('cuda'),
+            num_points=int(num_points),
+            min_part_confidence=min_part_confidence,
+        )
 
     def export_urdf(
         self,

infer_asset.py

@@ -445,9 +445,6 @@ def infer_single_asset(
     up_dir,
     model,
     num_points,
-    strict,
-    output_path,
-    animation_frames,
     min_part_confidence=0.0
 ):
     if up_dir is ["x", "X"]:
@@ -491,14 +488,18 @@ def infer_single_asset(
             min_part_confidence=min_part_confidence
         )
     
-    part_ids = outputs[0]['part_ids']
-    motion_hierarchy = outputs[0]['motion_hierarchy']
-    is_part_revolute = outputs[0]['is_part_revolute']
-    is_part_prismatic = outputs[0]['is_part_prismatic']
-    revolute_plucker = outputs[0]['revolute_plucker']
-    revolute_range = outputs[0]['revolute_range']
-    prismatic_axis = outputs[0]['prismatic_axis']
-    prismatic_range = outputs[0]['prismatic_range']
+    return outputs, face_indices
+
+
+def save_articulated_meshes(mesh, face_indices, outputs, output_path, strict, animation_frames: int = 50, hyp_idx: int = 0):
+    part_ids = outputs[hyp_idx]['part_ids']
+    motion_hierarchy = outputs[hyp_idx]['motion_hierarchy']
+    is_part_revolute = outputs[hyp_idx]['is_part_revolute']
+    is_part_prismatic = outputs[hyp_idx]['is_part_prismatic']
+    revolute_plucker = outputs[hyp_idx]['revolute_plucker']
+    revolute_range = outputs[hyp_idx]['revolute_range']
+    prismatic_axis = outputs[hyp_idx]['prismatic_axis']
+    prismatic_range = outputs[hyp_idx]['prismatic_range']
 
     _, face_part_ids_refined_strict, face_part_ids_refined = find_part_ids_for_faces(
         mesh,
@@ -510,13 +511,6 @@ def infer_single_asset(
     num_parts = len(unique_part_ids)
     print(f"Found {num_parts} unique parts")
     
-    # Check if original mesh has texture/UV coordinates
-    has_original_texture = (
-        hasattr(mesh.visual, 'uv') and 
-        mesh.visual.uv is not None and 
-        len(mesh.visual.uv) > 0
-    )
-    
     mesh_parts_original = [mesh.submesh([face_part_ids == part_id], append=True) for part_id in unique_part_ids]
     mesh_parts_segmented = create_textured_mesh_parts([mp.copy() for mp in mesh_parts_original])