Update app.py

app.py

@@ -12,7 +12,7 @@ from typing import List, Tuple
 from PIL import Image
 from easydict import EasyDict as edict
 
-# Add missing imports for RigNet API
+# Add missing imports for MagicArticulate API
 from gradio_client import Client, handle_file
 from gradio_client.exceptions import AppError
 
@@ -21,8 +21,7 @@ from trellis.pipelines import TrellisImageTo3DPipeline
 from trellis.representations import Gaussian, MeshExtractResult
 from trellis.utils import render_utils, postprocessing_utils
 
-RIGNET_SPACE = os.getenv("RIGNET_SPACE", "ckc99u/spaceB")
-RIGNET_AVAILABLE = True
+MAGIC_ARTICULATE_URL = "https://62157d4fdd84d3addb.gradio.live/"
 
 # Configuration
 MAX_SEED = np.iinfo(np.int32).max
@@ -97,7 +96,7 @@ def unpack_state(state: dict) -> Tuple[Gaussian, edict]:
     gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
     gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
     gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
-    
+
     mesh = edict(
         vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
         faces=torch.tensor(state['mesh']['faces'], device='cuda'),
@@ -108,43 +107,51 @@ def get_seed(randomize_seed: bool, seed: int) -> int:
     """Get random seed for generation"""
     return np.random.randint(0, MAX_SEED) if randomize_seed else seed
 
-def call_rignet_api(obj_path: str, bandwidth: float = 0.0429, rignet_space: str = "ckc99u/spaceB") -> Tuple[str, str]:
+def call_magic_articulate_api(obj_path: str, api_url: str = MAGIC_ARTICULATE_URL) -> Tuple[str, str, str]:
     """
-    Call RigNet space API to generate rigging from OBJ mesh
-    
+    Call MagicArticulate Colab API to generate rigging and skeleton from OBJ mesh
+
     Args:
-        obj_path: Path to OBJ file (must be 1K-5K vertices)
-        bandwidth: Rigging bandwidth parameter (0.01-0.1)
-        rignet_space: RigNet space identifier (e.g., "username/space_name")
-    
+        obj_path: Path to OBJ file
+        api_url: MagicArticulate Colab gradio URL
+
     Returns:
-        Tuple of (rig_file_path, rig_info_text)
-        - rig_file_path: Path to generated TXT rigging file
-        - rig_info_text: Info about rigging (joint count, etc.)
+        Tuple of (rig_pred_path, skeleton_obj_path, info_text)
+        - rig_pred_path: Path to generated rig prediction TXT file
+        - skeleton_obj_path: Path to generated skeleton OBJ file
+        - info_text: Info about rigging results
     """
     try:
-        print(f"🦴 Connecting to RigNet API ({rignet_space})...")
-        rignet_client = Client(rignet_space)
-        
-        print("📤 Uploading OBJ to RigNet...")
-        rig_result = rignet_client.predict(
-            obj_path=handle_file(obj_path),
-            bandwidth=float(bandwidth),
+        print(f"🦴 Connecting to MagicArticulate API ({api_url})...")
+        magic_client = Client(api_url)
+
+        print("📤 Uploading OBJ to MagicArticulate...")
+        result = magic_client.predict(
+            input_mesh=handle_file(obj_path),
             api_name="/predict"
         )
-        
-        # RigNet returns (rig_file_path, info_text)
-        rig_file = rig_result[0]
-        rig_info = rig_result[1]
-        print("✅ RigNet generation successful!")
-        return rig_file, rig_info
-        
+
+        # MagicArticulate returns (rig_pred.txt, skeleton.obj, normalized_mesh.obj)
+        rig_pred_file = result[0]
+        skeleton_file = result[1]
+
+        print("✅ MagicArticulate generation successful!")
+
+        # Read skeleton info
+        info_text = "Skeleton generated with hierarchical bone ordering"
+        if skeleton_file and os.path.exists(skeleton_file):
+            skeleton_mesh = trimesh.load(skeleton_file, force='mesh')
+            num_vertices = len(skeleton_mesh.vertices)
+            info_text = f"Joints: {num_vertices // 2}, Hierarchical structure"
+
+        return rig_pred_file, skeleton_file, info_text
+
     except AppError as e:
         error_msg = str(e)
-        print(f"⚠️ RigNet error: {error_msg}")
+        print(f"⚠️ MagicArticulate error: {error_msg}")
         raise
     except Exception as e:
-        print(f"⚠️ RigNet API error: {str(e)}")
+        print(f"⚠️ MagicArticulate API error: {str(e)}")
         raise
 
 @spaces.GPU(duration=180)
@@ -158,18 +165,17 @@ def generate_3d_with_rigging(
     mesh_simplify: float,
     texture_size: int,
     req: gr.Request,
-) -> Tuple[dict, str, str, str, str, str]:
+) -> Tuple[dict, str, str, str, str, str, str]:
     """
-    Complete pipeline: Image -> 3D Model (TRELLIS) -> OBJ -> Rigging (RigNet on CPU)
-    All processing happens in this single GPU-decorated function to avoid quota issues.
+    Complete pipeline: Image -> 3D Model (TRELLIS) -> OBJ -> Rigging (MagicArticulate)
     """
     try:
         user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-        
+
         # ============ STEP 1: TRELLIS 3D GENERATION ============
         print("🎨 Generating 3D model with TRELLIS...")
         init_pipeline()
-        
+
         outputs = pipeline.run(
             image,
             seed=seed,
@@ -184,25 +190,26 @@ def generate_3d_with_rigging(
                 "cfg_strength": slat_guidance_strength,
             },
         )
-        
+
         # Extract Gaussian and Mesh
         gs = outputs['gaussian'][0]
         mesh = outputs['mesh'][0]
-        
+
         # ============ STEP 2: RENDER VIDEO ============
         print("📹 Rendering 360° preview video...")
         video = render_utils.render_video(gs, num_frames=120)['color']
         video_geo = render_utils.render_video(mesh, num_frames=120)['normal']
         video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+
         video_path = os.path.join(user_dir, 'sample.mp4')
         imageio.mimsave(video_path, video, fps=15)
-        
+
         # ============ STEP 3: EXTRACT GLB ============
         print("🎁 Extracting GLB with textures...")
         glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
         glb_path = os.path.join(user_dir, 'sample.glb')
         glb.export(glb_path)
-        
+
         # ============ STEP 4: CONVERT GLB TO OBJ ============
         print("🔄 Converting GLB to OBJ format...")
         obj_path = os.path.join(user_dir, "model.obj")
@@ -210,7 +217,7 @@ def generate_3d_with_rigging(
         original_vertices = len(mesh_trimesh.vertices)
         original_faces = len(mesh_trimesh.faces)
         mesh_trimesh.export(obj_path)
-        
+
         mesh_info = f"""
 📊 Mesh Statistics:
 • Vertices: {original_vertices:,}
@@ -218,50 +225,56 @@ def generate_3d_with_rigging(
 • Texture Size: {texture_size}px
 • Status: ✓ Ready for rigging
 """
-        
-        # ============ STEP 5: RIGNET RIGGING (CPU) ============
-        print("🦴 Calling RigNet API for automatic rigging...")
+
+        # ============ STEP 5: MAGIC ARTICULATE RIGGING ============
+        print("🦴 Calling MagicArticulate API for automatic skeleton generation...")
         rig_info = ""
         rig_file = None
-        
+        skeleton_file = None
+
         try:
-            # Call RigNet space API
-            rig_result_path, rig_info_text = call_rignet_api(
+            # Call MagicArticulate Colab API
+            rig_result_path, skeleton_result_path, rig_info_text = call_magic_articulate_api(
                 obj_path=obj_path,
-                bandwidth=0.0429,
-                rignet_space="ckc99u/spaceB"  # ← Update with your RigNet space
+                api_url=MAGIC_ARTICULATE_URL
             )
-            
+
             if rig_result_path and os.path.exists(rig_result_path):
-                # Copy rigging file to user directory
-                rig_file = os.path.join(user_dir, 'rig_result.txt')
+                # Copy rig prediction file to user directory
+                rig_file = os.path.join(user_dir, 'rig_pred.txt')
                 shutil.copy(rig_result_path, rig_file)
-                rig_info = f"""✅ RigNet Rigging Generated:
+
+            if skeleton_result_path and os.path.exists(skeleton_result_path):
+                # Copy skeleton file to user directory
+                skeleton_file = os.path.join(user_dir, 'skeleton.obj')
+                shutil.copy(skeleton_result_path, skeleton_file)
+
+            rig_info = f"""✅ MagicArticulate Skeleton Generated:
 {rig_info_text}
 
-📥 Download rig_result.txt to import into Blender/Maya
+📥 Downloads:
+• rig_pred.txt - Joint positions & bone hierarchy
+• skeleton.obj - 3D skeleton visualization
+
+🔧 Import into Blender/Maya for animation
 """
-            else:
-                # Partial failure: Create fallback
-                rig_file = os.path.join(user_dir, 'rig_result.txt')
-                with open(rig_file, 'w') as f:
-                    f.write(rig_info_text)
-                rig_info = rig_info_text
-                
+
         except Exception as e:
-            print(f"⚠️ RigNet API error: {str(e)}")
+            print(f"⚠️ MagicArticulate API error: {str(e)}")
             # Create error file with instructions
-            rig_file = os.path.join(user_dir, 'rig_result.txt')
+            rig_file = os.path.join(user_dir, 'rig_pred.txt')
             with open(rig_file, 'w') as f:
-                f.write(f"RigNet Error: {str(e)}\n\n")
+                f.write(f"MagicArticulate Error: {str(e)}\n\n")
                 f.write("Workaround: Download OBJ and rig manually in Blender.")
-            rig_info = f"⚠️ RigNet API unavailable: {str(e)}\n\n**Solution:** Download OBJ and use Blender Rigify add-on"
-        
+
+            rig_info = f"⚠️ MagicArticulate API unavailable: {str(e)}\n\n**Solution:** Download OBJ and use Blender Rigify add-on"
+            skeleton_file = None
+
         # ============ STEP 6: PACK RESULTS ============
         print("📦 Packaging results...")
         state = pack_state(gs, mesh)
         torch.cuda.empty_cache()
-        
+
         combined_info = f"""
 🎨 TRELLIS Generation:
 • Seed: {seed}
@@ -278,23 +291,24 @@ def generate_3d_with_rigging(
 ✓ Video preview (360° rotation)
 ✓ GLB file (textured 3D model)
 ✓ OBJ file (standard 3D format)
-✓ Rigging data (TXT)
+✓ Rig prediction (TXT)
+✓ Skeleton (OBJ)
 
 🔧 Next Steps:
-1. Download OBJ file
+1. Download OBJ + Skeleton files
 2. Import into Blender/Maya/C4D
-3. Apply automatic rigging (Rigify, etc.)
-4. Rig and animate your model
+3. Apply rigging from rig_pred.txt
+4. Animate your model
 
 💡 Pro Tips:
-• Use GLB for quick preview
-• Use OBJ for professional workflows
-• The model is optimized for animation
+• Skeleton shows joint hierarchy visually
+• Rig prediction contains exact joint coordinates
+• Model is optimized for animation workflow
 """
-        
+
         print("✅ All processing complete!")
-        return state, video_path, glb_path, obj_path, rig_file, combined_info
-        
+        return state, video_path, glb_path, obj_path, rig_file, skeleton_file, combined_info
+
     except Exception as e:
         import traceback
         error_detail = traceback.format_exc()
@@ -313,32 +327,32 @@ def extract_gaussian(state: dict, req: gr.Request) -> Tuple[str, str]:
     return gaussian_path, gaussian_path
 
 # ============ GRADIO UI ============
-with gr.Blocks(title="Image to Rigged 3D Model (Integrated)", delete_cache=(600, 600)) as demo:
+with gr.Blocks(title="Image to Rigged 3D Model (MagicArticulate)", delete_cache=(600, 600)) as demo:
     gr.Markdown("""
-# 🎭 Image → 3D → Rigging (Integrated Pipeline)
+# 🎭 Image → 3D → Rigging (MagicArticulate Pipeline)
 
-**Skip API calls, run everything locally in one space!**
+**Automated 3D generation with hierarchical skeleton rigging!**
 
 This unified pipeline combines:
 - **TRELLIS** (Image-to-3D, Microsoft Research)
-- **RigNet** (Auto-rigging, SIGGRAPH 2020)
+- **MagicArticulate** (Auto-skeleton generation, CVPR 2025)
 
 ### 🚀 Workflow:
 1. 📤 Upload image of object/character
 2. 🎨 TRELLIS generates high-quality 3D mesh (GPU)
 3. 🔄 Convert to OBJ format
-4. 🦴 RigNet generates skeletal rigging (CPU)
-5. 💾 Download mesh + rigging for animation
+4. 🦴 MagicArticulate generates hierarchical skeleton
+5. 💾 Download mesh + rigging + skeleton for animation
 
 ### ✨ Benefits:
-- ✅ No API quota errors (runs directly in space)
-- ✅ Faster processing (no inter-space communication)
-- ✅ Complete control over parameters
+- ✅ Hierarchical bone ordering for better animation
+- ✅ Automatic joint placement and bone connections
 - ✅ Production-ready output for Blender/Maya
+- ✅ Visual skeleton + rig data included
 
 ⏱️ **Estimated time:** 2-5 minutes
 """)
-    
+
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("### 📥 Input")
@@ -349,11 +363,11 @@ This unified pipeline combines:
                 type="pil",
                 height=300
             )
-            
+
             with gr.Accordion("⚙️ TRELLIS Parameters", open=False):
                 seed = gr.Slider(0, MAX_SEED, label="Seed", value=0, step=1)
                 randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                
+
                 gr.Markdown("**Stage 1: Sparse Structure**")
                 ss_guidance = gr.Slider(
                     0.0, 10.0,
@@ -367,7 +381,7 @@ This unified pipeline combines:
                     value=12,
                     step=1
                 )
-                
+
                 gr.Markdown("**Stage 2: Structured Latent**")
                 slat_guidance = gr.Slider(
                     0.0, 10.0,
@@ -381,7 +395,7 @@ This unified pipeline combines:
                     value=12,
                     step=1
                 )
-            
+
             with gr.Accordion("⚙️ Output Settings", open=False):
                 mesh_simplify = gr.Slider(
                     0.9, 0.98,
@@ -395,21 +409,21 @@ This unified pipeline combines:
                     value=1024,
                     step=512
                 )
-            
+
             generate_btn = gr.Button(
                 "🚀 Generate Rigged Model",
                 variant="primary",
                 size="lg"
             )
-            
+
             extract_gs_btn = gr.Button(
                 "📥 Extract Gaussian (PLY)",
                 interactive=False
             )
-        
+
         with gr.Column(scale=1):
             gr.Markdown("### 📤 Outputs")
-            
+
             with gr.Tabs():
                 with gr.Tab("📹 Preview"):
                     video_output = gr.Video(
@@ -418,13 +432,13 @@ This unified pipeline combines:
                         loop=True,
                         height=300
                     )
-                
+
                 with gr.Tab("🎨 3D Viewer"):
                     model_output = gr.Model3D(
                         label="GLB Viewer",
                         height=400
                     )
-                
+
                 with gr.Tab("📦 Files"):
                     glb_download = gr.DownloadButton(
                         label="📥 Download GLB",
@@ -435,34 +449,38 @@ This unified pipeline combines:
                         interactive=False
                     )
                     rig_download = gr.DownloadButton(
-                        label="🦴 Download RIG (TXT)",
+                        label="🦴 Download Rig Prediction (TXT)",
+                        interactive=False
+                    )
+                    skeleton_download = gr.DownloadButton(
+                        label="🦴 Download Skeleton (OBJ)",
                         interactive=False
                     )
                     gs_download = gr.DownloadButton(
                         label="✨ Download Gaussian (PLY)",
                         interactive=False
                     )
-                
+
                 with gr.Tab("ℹ️ Info"):
                     info_output = gr.Textbox(
                         label="Pipeline Information",
                         lines=20,
                         max_lines=30
                     )
-    
+
     # State management
     output_buf = gr.State()
-    
+
     # Event handlers
     demo.load(start_session)
     demo.unload(end_session)
-    
+
     input_image.upload(
         preprocess_image,
         inputs=[input_image],
         outputs=[input_image],
     )
-    
+
     generate_btn.click(
         get_seed,
         inputs=[randomize_seed, seed],
@@ -475,27 +493,29 @@ This unified pipeline combines:
             slat_guidance, slat_steps,
             mesh_simplify, texture_size
         ],
-        outputs=[output_buf, video_output, model_output, obj_download, rig_download, info_output],
+        outputs=[output_buf, video_output, model_output, obj_download, rig_download, skeleton_download, info_output],
     ).then(
         lambda: (
             gr.Button(interactive=True),
             gr.DownloadButton(interactive=True),
             gr.DownloadButton(interactive=True),
             gr.DownloadButton(interactive=True),
+            gr.DownloadButton(interactive=True),
         ),
-        outputs=[extract_gs_btn, glb_download, obj_download, rig_download],
+        outputs=[extract_gs_btn, glb_download, obj_download, rig_download, skeleton_download],
     )
-    
+
     video_output.clear(
         lambda: (
             gr.Button(interactive=False),
             gr.DownloadButton(interactive=False),
             gr.DownloadButton(interactive=False),
             gr.DownloadButton(interactive=False),
+            gr.DownloadButton(interactive=False),
         ),
-        outputs=[extract_gs_btn, glb_download, obj_download, rig_download],
+        outputs=[extract_gs_btn, glb_download, obj_download, rig_download, skeleton_download],
     )
-    
+
     extract_gs_btn.click(
         extract_gaussian,
         inputs=[output_buf],
@@ -507,4 +527,4 @@ This unified pipeline combines:
 
 if __name__ == "__main__":
     init_pipeline()
-    demo.launch(mcp_server=True)
+    demo.launch()