Update app.py

app.py

@@ -1,290 +1,137 @@
 import os
-import sys
-import subprocess
-from pathlib import Path
-
-# Download models before importing anything else
-def download_models_if_needed():
-    """Download models if they don't exist"""
-    checkpoint_path = Path("checkpoints/spatial_order.pt")
-
-    if not checkpoint_path.exists():
-        print("=" * 70)
-        print("🔄 First run: Downloading model checkpoints...")
-        print("=" * 70)
-
-        try:
-            # Run the download script
-            result = subprocess.run(
-                [sys.executable, "download_models.py"],
-                check=True,
-                capture_output=True,
-                text=True
-            )
-            print(result.stdout)
-
-            if result.returncode != 0:
-                print(f"Error downloading models: {result.stderr}")
-                raise RuntimeError("Failed to download model checkpoints")
-
-        except Exception as e:
-            print(f"❌ Error during model download: {e}")
-            print("\nPlease ensure:")
-            print("1. You have internet connectivity")
-            print("2. Hugging Face Hub is accessible")
-            print("3. The download_models.py script is present")
-            raise
-    else:
-        print("✅ Model checkpoints already exist, skipping download")
-
-# Download models before importing heavy dependencies
-download_models_if_needed()
-
-# Now import the rest
+import gradio as gr
 import torch
 import trimesh
-import numpy as np
-import gradio as gr
 import tempfile
-
-from skeleton_models.skeletongen import SkeletonGPT
-from data_utils.save_npz import normalize_to_unit_cube
-from utils.mesh_to_pc import MeshProcessor
-from utils.save_utils import (
-    pred_joints_and_bones, 
-    save_skeleton_to_txt, 
-    merge_duplicate_joints_and_fix_bones,
-    save_skeleton_obj,
-    save_mesh
-)
-
-# Global model variable
-model = None
-args_config = None
-
-def initialize_model(use_hierarchical=False):
-    """Initialize the model once at startup"""
-    global model, args_config
-
-    # Check if we need to reload for different ordering
-    if model is not None and hasattr(args_config, 'hier_order'):
-        if args_config.hier_order == use_hierarchical:
-            return
-
-    print(f"Initializing MagicArticulate model (hierarchical={use_hierarchical})...")
-
-    # Create a simple args object with default parameters
-    class Args:
-        def __init__(self, hier_order=False):
-            self.input_pc_num = 8192
-            self.num_beams = 1
-            self.llm = "facebook/opt-350m"
-            self.pad_id = -1
-            self.n_discrete_size = 128
-            self.n_max_bones = 100
-            self.seed = 0
-            self.precision = "fp16"
-            self.hier_order = hier_order
-            # Select checkpoint based on ordering type
-            if hier_order:
-                self.pretrained_weights = "checkpoints/hier_order.pt"
-            else:
-                self.pretrained_weights = "checkpoints/spatial_order.pt"
-
-    args_config = Args(hier_order=use_hierarchical)
-
-    # Determine device
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print(f"Using device: {device}")
-
-    # Load model
-    model = SkeletonGPT(args_config).to(device)
-
-    # Load pretrained weights
-    if os.path.exists(args_config.pretrained_weights):
-        print(f"Loading checkpoint: {args_config.pretrained_weights}")
-        pkg = torch.load(args_config.pretrained_weights, map_location=device)
-        model.load_state_dict(pkg["model"])
-        model.eval()
-        print("✅ Model loaded successfully!")
-    else:
-        error_msg = f"Checkpoint not found: {args_config.pretrained_weights}"
-        print(f"❌ {error_msg}")
-        raise FileNotFoundError(error_msg)
+import numpy as np
+from pathlib import Path
+import gc
+
+# Import after installing dependencies
+from inference import SkeletonInferencer
+
+# Initialize model globally to avoid reloading
+MODEL = None
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+def initialize_model():
+    """Initialize model once at startup"""
+    global MODEL
+    if MODEL is None:
+        print(f"Initializing model on {DEVICE}...")
+        MODEL = SkeletonInferencer(
+            pretrained_weights="skeleton_ckpt/checkpoint_trainonv2_spatial.pth",
+            device=DEVICE,
+            precision="fp16" if DEVICE == "cuda" else "fp32"
+        )
+        print("Model initialized successfully!")
+    return MODEL
 
 def process_mesh(
     input_file,
-    apply_marching_cubes,
-    hier_order,
-    octree_depth
+    num_points=8192,
+    use_marching_cubes=False,
+    sequence_type="spatial",
+    progress=gr.Progress()
 ):
     """
-    Process the input mesh and generate rigging prediction
-
+    Process uploaded mesh and generate skeleton rigging
+    
     Args:
-        input_file: Uploaded mesh file (.obj, .ply, or .stl)
-        apply_marching_cubes: Whether to apply marching cubes
-        hier_order: Whether to use hierarchical ordering
-        octree_depth: Depth for octree (if using marching cubes)
-
-    Returns:
-        Tuple of (skeleton obj file, rig txt file, normalized mesh file, status message)
+        input_file: Uploaded mesh file (.obj, .ply, .stl)
+        num_points: Number of points for point cloud sampling
+        use_marching_cubes: Apply marching cubes preprocessing
+        sequence_type: "spatial" or "hierarchical" ordering
+        progress: Gradio progress tracker
     """
     try:
-        # Initialize/reload model if ordering changed
-        initialize_model(use_hierarchical=hier_order)
-
-        # Determine device
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-        # Create temporary output directory
-        output_dir = tempfile.mkdtemp()
-
-        # Get file information
-        file_name = Path(input_file).stem
-        file_ext = Path(input_file).suffix.lower()
-
-        # Check file type
-        if file_ext not in ['.obj', '.ply', '.stl']:
-            return None, None, None, f"Error: Unsupported file type {file_ext}. Please upload .obj, .ply, or .stl file."
-
-        # Load mesh
-        print(f"Loading mesh: {file_name}{file_ext}")
-        mesh = trimesh.load(input_file, force='mesh')
-        print(f"Mesh loaded: {len(mesh.vertices)} vertices, {len(mesh.faces)} faces")
-
-        # Convert mesh to point cloud
-        print(f"Converting to point cloud (marching_cubes={apply_marching_cubes})...")
-        pc_list = MeshProcessor.convert_meshes_to_point_clouds(
-            [mesh], 
-            args_config.input_pc_num, 
-            apply_marching_cubes=apply_marching_cubes, 
-            octree_depth=octree_depth
-        )
-        pc_normal = pc_list[0]
-
-        # Normalize point cloud
-        pc_coor = pc_normal[:, :3]
-        normals = pc_normal[:, 3:]
-        pc_coor, center, scale = normalize_to_unit_cube(pc_coor, scale_factor=0.9995)
-
-        pc_coor = pc_coor.astype(np.float32)
-        normals = normals.astype(np.float32)
-
-        # Calculate transform parameters
-        bounds = np.array([pc_coor.min(axis=0), pc_coor.max(axis=0)])
-        pc_center = (bounds[0] + bounds[1])[None, :] / 2
-        pc_scale = ((bounds[1] - bounds[0]).max() + 1e-5)
-
-        transform_params = torch.tensor([
-            center[0], center[1], center[2],
-            scale,
-            pc_center[0][0], pc_center[0][1], pc_center[0][2],
-            pc_scale
-        ], dtype=torch.float32)
-
-        # Prepare batch data
-        pc_normal_tensor = torch.from_numpy(
-            np.concatenate([pc_coor, normals], axis=-1).astype(np.float16)
-        ).unsqueeze(0).to(device)
-
-        batch_data = {
-            'pc_normal': pc_normal_tensor,
-            'file_name': [file_name],
-            'transform_params': transform_params.unsqueeze(0).to(device),
-            'vertices': torch.from_numpy(mesh.vertices).unsqueeze(0).to(device),
-            'faces': torch.from_numpy(mesh.faces).unsqueeze(0).to(device)
-        }
-
-        # Generate skeleton
-        print("Generating skeleton...")
-        with torch.no_grad():
-            pred_bone_coords = model.generate(batch_data)
-
-        # Process predictions
-        skeleton = pred_bone_coords[0].cpu().numpy()
-        pred_joints, pred_bones = pred_joints_and_bones(skeleton.squeeze())
-
-        print(f"Generated {len(pred_joints)} joints and {len(pred_bones)} bones")
-
-        # Post-process: merge duplicate joints
-        if hier_order:
-            pred_root_index = pred_bones[0][0]
-            pred_joints, pred_bones, pred_root_index = merge_duplicate_joints_and_fix_bones(
-                pred_joints, pred_bones, root_index=pred_root_index
+        if input_file is None:
+            return None, None, None, "❌ Please upload a mesh file"
+        
+        progress(0.1, desc="Loading model...")
+        model = initialize_model()
+        
+        # Create temporary directory for outputs
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_path = Path(temp_dir)
+            
+            progress(0.2, desc="Processing mesh...")
+            
+            # Get file extension
+            file_ext = Path(input_file.name).suffix.lower()
+            if file_ext not in ['.obj', '.ply', '.stl']:
+                return None, None, None, f"❌ Unsupported file format: {file_ext}"
+            
+            progress(0.3, desc="Generating skeleton...")
+            
+            # Run inference
+            results = model.infer(
+                input_path=input_file.name,
+                output_dir=str(temp_path),
+                input_pc_num=num_points,
+                apply_marching_cubes=use_marching_cubes,
+                sequence_type=sequence_type
             )
-        else:
-            pred_joints, pred_bones = merge_duplicate_joints_and_fix_bones(pred_joints, pred_bones)
-            pred_root_index = None
-
-        print(f"After post-processing: {len(pred_joints)} joints, {len(pred_bones)} bones")
-
-        # Denormalize joints for rig file
-        transform_params_np = transform_params.cpu().numpy()
-        trans = transform_params_np[:3]
-        scale_val = transform_params_np[3]
-        pc_trans = transform_params_np[4:7]
-        pc_scale_val = transform_params_np[7]
-
-        pred_joints_denorm = pred_joints * pc_scale_val + pc_trans
-        pred_joints_denorm = pred_joints_denorm / scale_val + trans
-
-        # Save outputs
-        skel_obj_path = os.path.join(output_dir, f'{file_name}_skel.obj')
-        rig_txt_path = os.path.join(output_dir, f'{file_name}_pred.txt')
-        mesh_obj_path = os.path.join(output_dir, f'{file_name}_mesh.obj')
-
-        # Save skeleton
-        save_skeleton_obj(
-            pred_joints, 
-            pred_bones, 
-            skel_obj_path, 
-            pred_root_index if hier_order else None,
-            use_cone=hier_order
-        )
-        print(f"Saved skeleton: {skel_obj_path}")
-
-        # Save rig
-        vertices_np = mesh.vertices
-        save_skeleton_to_txt(
-            pred_joints_denorm, 
-            pred_bones, 
-            pred_root_index, 
-            hier_order, 
-            vertices_np, 
-            rig_txt_path
-        )
-        print(f"Saved rig: {rig_txt_path}")
-
-        # Save normalized mesh
-        vertices_norm = (vertices_np - trans) * scale_val
-        vertices_norm = (vertices_norm - pc_trans) / pc_scale_val
-        save_mesh(vertices_norm, mesh.faces, mesh_obj_path)
-        print(f"Saved mesh: {mesh_obj_path}")
-
-        status_msg = f"✅ Success! Generated skeleton with {len(pred_joints)} joints and {len(pred_bones)} bones."
-
-        return skel_obj_path, rig_txt_path, mesh_obj_path, status_msg
-
+            
+            progress(0.7, desc="Saving results...")
+            
+            # Get output file paths
+            file_name = Path(input_file.name).stem
+            skeleton_file = temp_path / f"{file_name}_skel.obj"
+            rig_file = temp_path / f"{file_name}_pred.txt"
+            mesh_file = temp_path / f"{file_name}_mesh.obj"
+            
+            # Prepare outputs
+            output_files = []
+            info_text = "✅ **Processing Complete!**\n\n"
+            
+            if skeleton_file.exists():
+                output_files.append(str(skeleton_file))
+                info_text += f"- Skeleton: `{skeleton_file.name}`\n"
+            
+            if rig_file.exists():
+                output_files.append(str(rig_file))
+                with open(rig_file, 'r') as f:
+                    rig_data = f.read()
+                info_text += f"- Rig data: `{rig_file.name}` ({len(rig_data.splitlines())} lines)\n"
+            
+            if mesh_file.exists():
+                output_files.append(str(mesh_file))
+                info_text += f"- Normalized mesh: `{mesh_file.name}`\n"
+            
+            info_text += f"\n**Statistics:**\n"
+            info_text += f"- Joints: {results.get('num_joints', 'N/A')}\n"
+            info_text += f"- Bones: {results.get('num_bones', 'N/A')}\n"
+            info_text += f"- Processing time: {results.get('time', 'N/A'):.2f}s\n"
+            
+            progress(1.0, desc="Done!")
+            
+            # Memory cleanup
+            if DEVICE == "cuda":
+                torch.cuda.empty_cache()
+            gc.collect()
+            
+            return output_files, str(skeleton_file) if skeleton_file.exists() else None, \
+                   str(rig_file) if rig_file.exists() else None, info_text
+    
     except Exception as e:
+        error_msg = f"❌ **Error:** {str(e)}"
+        print(f"Error processing mesh: {e}")
         import traceback
-        error_msg = f"❌ Error processing mesh: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
-        print(error_msg)
+        traceback.print_exc()
         return None, None, None, error_msg
 
 # Create Gradio interface
 def create_interface():
-    """Create the Gradio interface"""
-
-    with gr.Blocks(title="MagicArticulate - 3D Model Rigging", theme=gr.themes.Soft()) as demo:
+    with gr.Blocks(title="MagicArticulate - 3D Mesh Rigging") as demo:
         gr.Markdown("""
-        # 🪄 MagicArticulate: Make Your 3D Models Articulation-Ready
-
-        Upload a 3D mesh (.obj, .ply, or .stl) to automatically generate skeletal rigging.
-
-        **Paper**: [CVPR 2025] MagicArticulate | **Project**: [chaoyuesong.github.io/MagicArticulate/](https://chaoyuesong.github.io/MagicArticulate/)
+        # 🎨 MagicArticulate: Make Your 3D Models Articulation-Ready
+        
+        **CVPR 2025** | Upload a 3D mesh and automatically generate skeletal rigging for animation.
+        
+        Supported formats: `.obj`, `.ply`, `.stl`
         """)
-
+        
         with gr.Row():
             with gr.Column(scale=1):
                 gr.Markdown("### 📤 Input")
@@ -293,120 +140,99 @@ def create_interface():
                     file_types=[".obj", ".ply", ".stl"],
                     type="filepath"
                 )
-
-                gr.Markdown("### ⚙️ Options")
-                apply_marching_cubes = gr.Checkbox(
-                    label="Apply Marching Cubes",
-                    value=False,
-                    info="More accurate mesh processing (slower, ~30s longer)"
-                )
-
-                hier_order = gr.Checkbox(
-                    label="Hierarchical Ordering",
-                    value=False,
-                    info="Use hierarchical skeleton structure (better for characters)"
-                )
-
-                octree_depth = gr.Slider(
-                    minimum=5,
-                    maximum=9,
-                    value=7,
-                    step=1,
-                    label="Octree Depth",
-                    info="Only used when Marching Cubes is enabled"
-                )
-
-                generate_btn = gr.Button("🚀 Generate Rigging", variant="primary", size="lg")
-
+                
+                with gr.Accordion("⚙️ Advanced Settings", open=False):
+                    num_points = gr.Slider(
+                        minimum=2048,
+                        maximum=16384,
+                        value=8192,
+                        step=1024,
+                        label="Point Cloud Size",
+                        info="Higher values = more accurate but slower"
+                    )
+                    
+                    use_marching_cubes = gr.Checkbox(
+                        label="Use Marching Cubes",
+                        value=False,
+                        info="Apply marching cubes for better point cloud (slower)"
+                    )
+                    
+                    sequence_type = gr.Radio(
+                        choices=["spatial", "hierarchical"],
+                        value="spatial",
+                        label="Sequence Ordering",
+                        info="Spatial is faster, hierarchical preserves hierarchy"
+                    )
+                
+                process_btn = gr.Button("🚀 Generate Rigging", variant="primary", size="lg")
+            
             with gr.Column(scale=1):
                 gr.Markdown("### 📥 Output")
-                status_text = gr.Textbox(
-                    label="Status",
-                    lines=4,
-                    interactive=False,
-                    placeholder="Upload a mesh and click Generate to start..."
-                )
-
-                skel_output = gr.File(
-                    label="📥 Skeleton (.obj)",
-                    interactive=False
-                )
-
-                rig_output = gr.File(
-                    label="📥 Rig Prediction (.txt)",
-                    interactive=False
-                )
-
-                mesh_output = gr.File(
-                    label="📥 Normalized Mesh (.obj)",
-                    interactive=False
-                )
-
+                info_output = gr.Markdown()
+                
+                with gr.Tabs():
+                    with gr.Tab("📦 Download Files"):
+                        output_files = gr.File(
+                            label="Generated Files",
+                            file_count="multiple"
+                        )
+                    
+                    with gr.Tab("🦴 Skeleton"):
+                        skeleton_viewer = gr.Model3D(
+                            label="Generated Skeleton",
+                            height=400
+                        )
+                    
+                    with gr.Tab("📄 Rig Data"):
+                        rig_text = gr.File(
+                            label="Rig Text File"
+                        )
+        
         gr.Markdown("""
-        ### 📖 About
-
-        MagicArticulate automatically generates skeletal structures for 3D models, making them animation-ready.
-        The system uses a transformer-based approach to predict joint positions and bone connections.
-
-        **Processing time**: ~5-15 seconds (spatial ordering) or ~10-20 seconds (hierarchical ordering)
-
-        **Outputs**:
-        - **Skeleton (.obj)**: 3D visualization of the generated skeleton
-        - **Rig Prediction (.txt)**: Detailed rigging information (joints, bones, hierarchy)
-        - **Normalized Mesh (.obj)**: The input mesh normalized to unit cube
-
-        **Tips**:
-        - For furniture/objects: Use spatial ordering (default)
-        - For characters/animals: Use hierarchical ordering
-        - Marching Cubes: Better for noisy meshes but slower
-
-        **Citation**:
-        ```bibtex
+        ### 📖 How to Use
+        1. Upload your 3D mesh file (`.obj`, `.ply`, or `.stl`)
+        2. Adjust settings if needed (default works well)
+        3. Click "Generate Rigging" and wait for processing
+        4. Download the generated skeleton and rigging files
+        
+        ### 💡 Tips
+        - **Memory limited?** Use fewer points (e.g., 4096) for large meshes
+        - **Better quality?** Enable Marching Cubes (slower but more accurate)
+        - **Preserve hierarchy?** Use hierarchical ordering
+        
+        ### 📚 Citation
+        ```
         @inproceedings{song2025magicarticulate,
           title={MagicArticulate: Make Your 3D Models Articulation-Ready},
-          author={Song, Chaoyue and others},
+          author={Chaoyue Song and Jianfeng Zhang and others},
           booktitle={CVPR},
           year={2025}
         }
         ```
+        
+        **Links:** [Paper](https://chaoyuesong.github.io/MagicArticulate/) | [GitHub](https://github.com/Seed3D/MagicArticulate)
         """)
-
-        # Connect the button to the processing function
-        generate_btn.click(
+        
+        # Connect components
+        process_btn.click(
             fn=process_mesh,
-            inputs=[input_file, apply_marching_cubes, hier_order, octree_depth],
-            outputs=[skel_output, rig_output, mesh_output, status_text]
+            inputs=[input_file, num_points, use_marching_cubes, sequence_type],
+            outputs=[output_files, skeleton_viewer, rig_text, info_output]
         )
-
+        
     return demo
 
 if __name__ == "__main__":
-    print("=" * 70)
-    print("🚀 Starting MagicArticulate Gradio Space")
-    print("=" * 70)
-
-    # Check CUDA availability
-    print(f"PyTorch version: {torch.__version__}")
-    print(f"CUDA available: {torch.cuda.is_available()}")
-    if torch.cuda.is_available():
-        print(f"CUDA device: {torch.cuda.get_device_name(0)}")
-        print(f"CUDA memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
-
-    # Initialize model at startup (spatial ordering by default)
-    try:
-        print("\nInitializing model...")
-        initialize_model(use_hierarchical=False)
-        print("✅ Model ready!")
-    except Exception as e:
-        print(f"⚠️ Warning: Could not initialize model at startup: {e}")
-        print("Model will be initialized on first request.")
-
-    # Launch Gradio app
-    print("\nLaunching Gradio interface...")
+    # Download models on first run
+    if not os.path.exists("skeleton_ckpt"):
+        print("Downloading model checkpoints...")
+        os.system("python download_models.py")
+    
+    # Create and launch interface
     demo = create_interface()
-    demo.queue(max_size=10)
+    demo.queue(max_size=5)  # Limit queue to manage memory
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
         share=False
-    )
+    )