Update ActionMesh space

app.py

@@ -1,3 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
 """
 ActionMesh Gradio Demo
 
@@ -205,11 +211,11 @@ from actionmesh.io.mesh_io import save_deformation
 
 # --- Import ActionMesh modules after setup ---
 from actionmesh.io.video_input import load_frames
-from actionmesh.pipeline import ActionMeshPipeline
 from actionmesh.render.utils import save_rgba_video
+from gradio_pipeline import GradioPipeline
 
 # Global pipeline instance (loaded on CPU at startup)
-pipeline: ActionMeshPipeline | None = None
+pipeline: GradioPipeline | None = None
 
 
 def get_available_examples() -> list[tuple[str, str]]:
@@ -275,7 +281,7 @@ def load_example_images(evt: gr.SelectData) -> list[str]:
     return []
 
 
-def load_pipeline_cpu() -> ActionMeshPipeline:
+def load_pipeline_cpu() -> GradioPipeline:
     """Load the ActionMesh pipeline on CPU (called once at module load)."""
     global pipeline
     if pipeline is None:
@@ -283,7 +289,7 @@ def load_pipeline_cpu() -> ActionMeshPipeline:
         # Get config path from actionmesh cache directory
         cache_dir = Path.home() / ".cache" / "actionmesh"
         config_dir = str(cache_dir / "actionmesh" / "configs")
-        pipeline = ActionMeshPipeline(
+        pipeline = GradioPipeline(
             config_name="actionmesh.yaml",
             config_dir=config_dir,
         )
@@ -355,7 +361,7 @@ def _run_actionmesh_impl(
             torch.cuda.empty_cache()
 
         # Run inference
-        progress(0.3, desc="Running ActionMesh inference...")
+        progress(None, desc="Starting pipeline...")
 
         # Set steps based on quality mode
         if quality_mode == "⚡ Fast":
@@ -365,12 +371,17 @@ def _run_actionmesh_impl(
             stage_0_steps = 100
             stage_1_steps = 30
 
+        # Create progress callback for the pipeline
+        def pipeline_progress_callback(value: float, desc: str) -> None:
+            progress(value, desc=desc)
+
         meshes = pipe(
             input=input_data,
             anchor_idx=reference_frame - 1,  # Convert from 1-indexed UI to 0-indexed
             stage_0_steps=stage_0_steps,
             stage_1_steps=stage_1_steps,
             seed=seed,
+            progress_callback=pipeline_progress_callback,
         )
 
         # Save input video
@@ -381,7 +392,7 @@ def _run_actionmesh_impl(
             return None, None, None, "Error: No meshes generated."
 
         # Save deformations and create animated GLB
-        progress(0.9, desc="Creating animated GLB...")
+        progress(1.0, desc="Creating animated GLB...")
 
         vertices_path, faces_path = save_deformation(
             meshes, path=f"{output_dir}/deformations"

gradio_pipeline.py

@@ -0,0 +1,115 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+GradioPipeline: ActionMesh pipeline with Gradio progress tracking.
+
+This module provides a subclass of ActionMeshPipeline that adds progress
+callbacks for integration with Gradio's progress bar.
+"""
+
+from typing import Callable, Optional
+
+import torch
+import trimesh
+from actionmesh.io.video_input import ActionMeshInput
+from actionmesh.pipeline import ActionMeshPipeline
+
+ProgressCallback = Callable[[float, str], None]
+
+
+class GradioPipeline(ActionMeshPipeline):
+    """
+    ActionMesh pipeline with Gradio progress tracking support.
+
+    Progress breakdown:
+        - 0% -> 10%: Anchor 3D generation (image_to_3d)
+        - 10% -> 90%: Stage 1 - Flow matching denoising (step-by-step)
+        - 90% -> 100%: Stage 2 - Mesh decoding (step-by-step)
+    """
+
+    def __call__(
+        self,
+        input: ActionMeshInput,
+        seed: int = 44,
+        stage_0_steps: int | None = None,
+        face_decimation: int | None = None,
+        floaters_threshold: float | None = None,
+        stage_1_steps: int | None = None,
+        guidance_scales: list[float] | None = None,
+        anchor_idx: int | None = None,
+        progress_callback: Optional[ProgressCallback] = None,
+    ) -> list[trimesh.Trimesh]:
+        """Generate an animated mesh sequence with progress tracking."""
+        # Apply parameter overrides
+        if stage_0_steps is not None:
+            self.cfg.model.image_to_3D_denoiser.num_inference_steps = stage_0_steps
+        if stage_1_steps is not None:
+            self.scheduler.num_inference_steps = stage_1_steps
+        if guidance_scales is not None:
+            self.cf_guidance.guidance_scales = guidance_scales
+        if face_decimation is not None:
+            self.mesh_process.face_decimation = face_decimation
+        if floaters_threshold is not None:
+            self.mesh_process.floaters_threshold = floaters_threshold
+        if anchor_idx is not None:
+            self.cfg.anchor_idx = anchor_idx
+
+        # -- Preprocessing: remove background
+        input.frames = self.background_removal.process_images(input.frames)
+
+        # -- Preprocessing: grouped cropping & padding
+        input.frames = self.image_process.process_images(input.frames)
+
+        with torch.inference_mode():
+            # -- Stage 0: generate anchor 3D mesh & latent from single frame
+            latent_bank, mesh_bank = self.init_banks_from_anchor(input, seed)
+
+            if progress_callback is not None:
+                progress_callback(0.10, "Anchor 3D generated, starting Stage 1...")
+
+            # Stage 1 callback: 10% -> 90%
+            def stage1_callback(
+                step: int, total_steps: int, window_idx: int, total_windows: int
+            ) -> None:
+                if progress_callback is not None:
+                    window_progress = (window_idx + step / total_steps) / total_windows
+                    progress_callback(
+                        0.10 + 0.80 * window_progress,
+                        f"Stage 1: step {step}/{total_steps} ",
+                    )
+
+            # Stage 2 callback: 90% -> 100%
+            def stage2_callback(
+                step: int, total_steps: int, window_idx: int, total_windows: int
+            ) -> None:
+                if progress_callback is not None:
+                    window_progress = (window_idx + step / total_steps) / total_windows
+                    progress_callback(
+                        0.90 + 0.10 * window_progress,
+                        f"Stage 2: step {step}/{total_steps} ",
+                    )
+
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                # -- Stage I: denoise synchronized 3D latents
+                latent_bank = self.generate_3d_latents(
+                    input,
+                    latent_bank=latent_bank,
+                    seed=seed,
+                    step_callback=stage1_callback,
+                )
+
+                # -- Stage II: decode latents into mesh displacements
+                mesh_bank = self.generate_mesh_animation(
+                    latent_bank=latent_bank,
+                    mesh_bank=mesh_bank,
+                    step_callback=stage2_callback,
+                )
+
+            if progress_callback is not None:
+                progress_callback(1.0, "Pipeline complete!")
+
+        return mesh_bank.get_ordered(device="cpu")[0]