Upload 3 files

download_models.py

@@ -1,178 +1,73 @@
-#!/usr/bin/env python3
-"""
-Download large model checkpoints at runtime for Hugging Face Spaces
-This avoids the 1GB Git LFS repository limit
-"""
-
-import os
-import sys
-from pathlib import Path
-from huggingface_hub import hf_hub_download, snapshot_download
-import requests
-from tqdm import tqdm
-
-def download_file_with_progress(url, destination):
-    """Download a file with progress bar"""
-    response = requests.get(url, stream=True)
-    total_size = int(response.headers.get('content-length', 0))
-
-    destination = Path(destination)
-    destination.parent.mkdir(parents=True, exist_ok=True)
-
-    with open(destination, 'wb') as f, tqdm(
-        desc=destination.name,
-        total=total_size,
-        unit='iB',
-        unit_scale=True,
-        unit_divisor=1024,
-    ) as pbar:
-        for data in response.iter_content(chunk_size=1024):
-            size = f.write(data)
-            pbar.update(size)
-
-def download_from_huggingface_hub():
-    """Download models from Hugging Face Hub"""
-
-    print("=" * 70)
-    print("Downloading MagicArticulate checkpoints from Hugging Face Hub...")
-    print("=" * 70)
-
-    # Create checkpoints directory
-    checkpoints_dir = Path("checkpoints")
-    checkpoints_dir.mkdir(exist_ok=True)
-
-    try:
-        # Download MagicArticulate checkpoints
-        print("\n📥 Downloading MagicArticulate spatial_order checkpoint...")
-        spatial_order_path = hf_hub_download(
-            repo_id="Seed3D/MagicArticulate",
-            filename="checkpoints/spatial_order.pt",
-            local_dir=".",
-            local_dir_use_symlinks=False
-        )
-        print(f"✅ Downloaded: {spatial_order_path}")
-
-        print("\n📥 Downloading MagicArticulate hier_order checkpoint...")
-        hier_order_path = hf_hub_download(
-            repo_id="Seed3D/MagicArticulate",
-            filename="checkpoints/hier_order.pt",
-            local_dir=".",
-            local_dir_use_symlinks=False
-        )
-        print(f"✅ Downloaded: {hier_order_path}")
-
-    except Exception as e:
-        print(f"❌ Error downloading MagicArticulate checkpoints: {e}")
-        print("Falling back to direct download...")
-        download_from_direct_links()
-        return
-
-    # Download Michelangelo checkpoints
-    print("\n" + "=" * 70)
-    print("Downloading Michelangelo checkpoints...")
-    print("=" * 70)
-
-    try:
-        # Download aligned_shape_latents
-        print("\n📥 Downloading Michelangelo aligned_shape_latents...")
-        michelangelo_dir = Path("checkpoints/michelangelo")
-        michelangelo_dir.mkdir(parents=True, exist_ok=True)
-
-        # Download the full Michelangelo repo or specific files
-        snapshot_download(
-            repo_id="Maikou/Michelangelo",
-            allow_patterns=["checkpoints/aligned_shape_latents/*"],
-            local_dir=".",
-            local_dir_use_symlinks=False
-        )
-        print("✅ Michelangelo checkpoints downloaded")
-
-    except Exception as e:
-        print(f"⚠️ Warning: Could not download Michelangelo checkpoints: {e}")
-        print("You may need to download them manually or adjust the code.")
-
-def download_from_direct_links():
-    """Fallback: Download from direct links"""
-
-    print("\n📥 Using direct download links...")
-
-    # You can add direct download links here as fallback
-    # For example, from Google Drive, OneDrive, or other cloud storage
-
-    direct_links = {
-        # Add your direct links here if available
-        # "checkpoints/spatial_order.pt": "https://your-direct-link/spatial_order.pt",
-        # "checkpoints/hier_order.pt": "https://your-direct-link/hier_order.pt",
-    }
-
-    for filepath, url in direct_links.items():
-        print(f"Downloading {filepath}...")
-        download_file_with_progress(url, filepath)
-
-def verify_downloads():
-    """Verify that all required files are downloaded"""
-
-    required_files = [
-        "checkpoints/spatial_order.pt",
-        "checkpoints/hier_order.pt",
-    ]
-
-    optional_files = [
-        "checkpoints/aligned_shape_latents/model.safetensors",
-    ]
-
-    print("\n" + "=" * 70)
-    print("Verifying downloads...")
-    print("=" * 70)
-
-    all_good = True
-    for filepath in required_files:
-        if Path(filepath).exists():
-            size = Path(filepath).stat().st_size / (1024 * 1024)  # MB
-            print(f"✅ {filepath} ({size:.2f} MB)")
-        else:
-            print(f"❌ {filepath} - MISSING!")
-            all_good = False
-
-    for filepath in optional_files:
-        if Path(filepath).exists():
-            size = Path(filepath).stat().st_size / (1024 * 1024)  # MB
-            print(f"✅ {filepath} ({size:.2f} MB) [optional]")
-        else:
-            print(f"⚠️  {filepath} - Not found [optional]")
-
-    if all_good:
-        print("\n✅ All required checkpoints downloaded successfully!")
-    else:
-        print("\n❌ Some required files are missing!")
-        sys.exit(1)
-
-    # Calculate total size
-    total_size = 0
-    for filepath in required_files + optional_files:
-        if Path(filepath).exists():
-            total_size += Path(filepath).stat().st_size
-
-    print(f"\n📊 Total downloaded: {total_size / (1024 * 1024 * 1024):.2f} GB")
-
-def main():
-    """Main download function"""
-
-    print("\n" + "=" * 70)
-    print("🚀 MagicArticulate Model Downloader for Hugging Face Spaces")
-    print("=" * 70)
-    print("\nThis script downloads large model checkpoints at runtime")
-    print("to avoid Hugging Face Spaces 1GB Git LFS limit.\n")
-
-    # Download from Hugging Face Hub
-    download_from_huggingface_hub()
-
-    # Verify all downloads
-    verify_downloads()
-
-    print("\n" + "=" * 70)
-    print("✅ Setup complete! Models are ready to use.")
-    print("=" * 70)
-
-if __name__ == "__main__":
-    main()
+"""
+Download model checkpoints from Hugging Face Hub
+Optimized for Spaces environment with error handling
+"""
+import os
+from huggingface_hub import hf_hub_download
+from pathlib import Path
+
+def download_with_retry(repo_id, filename, local_dir, max_retries=3):
+    """Download file with retry logic"""
+    for attempt in range(max_retries):
+        try:
+            print(f"Downloading {filename} (attempt {attempt + 1}/{max_retries})...")
+            file_path = hf_hub_download(
+                repo_id=repo_id,
+                filename=filename,
+                local_dir=local_dir,
+                local_dir_use_symlinks=False  # Important for Spaces
+            )
+            print(f"✓ Successfully downloaded: {filename}")
+            return file_path
+        except Exception as e:
+            if attempt == max_retries - 1:
+                print(f"✗ Failed to download {filename}: {e}")
+                raise
+            print(f"Retry {attempt + 1} failed, trying again...")
+    return None
+
+def main():
+    print("=" * 50)
+    print("Downloading MagicArticulate Model Checkpoints")
+    print("=" * 50)
+    
+    # Create directories
+    Path("skeleton_ckpt").mkdir(exist_ok=True)
+    Path("third_partys/Michelangelo/checkpoints/aligned_shape_latents").mkdir(
+        parents=True, exist_ok=True
+    )
+    
+    # Download Michelangelo checkpoint (required dependency)
+    print("\n[1/3] Downloading Michelangelo checkpoint...")
+    try:
+        download_with_retry(
+            repo_id="Maikou/Michelangelo",
+            filename="checkpoints/aligned_shape_latents/shapevae-256.ckpt",
+            local_dir="third_partys/Michelangelo"
+        )
+    except Exception as e:
+        print(f"Warning: Michelangelo download failed: {e}")
+        print("This may affect some features.")
+    
+    # Download MagicArticulate spatial checkpoint
+    print("\n[2/3] Downloading MagicArticulate spatial checkpoint...")
+    download_with_retry(
+        repo_id="Seed3D/MagicArticulate",
+        filename="skeleton_ckpt/checkpoint_trainonv2_spatial.pth",
+        local_dir="."
+    )
+    
+    # Download MagicArticulate hierarchical checkpoint
+    print("\n[3/3] Downloading MagicArticulate hierarchical checkpoint...")
+    download_with_retry(
+        repo_id="Seed3D/MagicArticulate",
+        filename="skeleton_ckpt/checkpoint_trainonv2_hier.pth",
+        local_dir="."
+    )
+    
+    print("\n" + "=" * 50)
+    print("✓ All downloads completed successfully!")
+    print("=" * 50)
+
+if __name__ == "__main__":
+    main()

inference.py

@@ -0,0 +1,176 @@
+import os
+import torch
+import trimesh
+import numpy as np
+from pathlib import Path
+import time
+
+# Import from MagicArticulate
+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
+)
+
+class SkeletonInferencer:
+    """Wrapper class for skeleton generation inference"""
+    
+    def __init__(self, pretrained_weights, device="cuda", precision="fp16"):
+        self.device = device
+        self.precision = precision
+        
+        # Create args object
+        class Args:
+            def __init__(self):
+                self.llm = "facebook/opt-350m"
+                self.pad_id = -1
+                self.n_discrete_size = 128
+                self.n_max_bones = 100
+                self.num_beams = 1
+                self.seed = 0
+        
+        self.args = Args()
+        
+        # Load model
+        print(f"Loading model from {pretrained_weights}...")
+        self.model = SkeletonGPT(self.args).to(device)
+        
+        pkg = torch.load(pretrained_weights, map_location=torch.device("cpu"))
+        self.model.load_state_dict(pkg["model"])
+        self.model.eval()
+        
+        # Set precision
+        if precision == "fp16" and device == "cuda":
+            self.model = self.model.half()
+        
+        print("Model loaded successfully!")
+    
+    @torch.no_grad()
+    def infer(
+        self,
+        input_path,
+        output_dir,
+        input_pc_num=8192,
+        apply_marching_cubes=False,
+        octree_depth=7,
+        sequence_type="spatial"
+    ):
+        """
+        Run inference on a single mesh file
+        
+        Returns:
+            dict: Results including paths and statistics
+        """
+        start_time = time.time()
+        
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Load mesh
+        mesh = trimesh.load(input_path, force='mesh')
+        
+        # Convert to point cloud
+        if apply_marching_cubes:
+            pc_list = MeshProcessor.convert_meshes_to_point_clouds(
+                [mesh], input_pc_num, 
+                apply_marching_cubes=True,
+                octree_depth=octree_depth
+            )
+            pc_normal = pc_list[0]
+        else:
+            # Simple sampling
+            points, face_indices = trimesh.sample.sample_surface(mesh, input_pc_num)
+            normals = mesh.face_normals[face_indices]
+            pc_normal = np.concatenate([points, normals], axis=-1)
+        
+        # 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)
+        
+        # Prepare transform parameters
+        bounds = np.array([pc_coor.min(axis=0), pc_coor.max(axis=0)])
+        pc_center = (bounds[0] + bounds[1]) / 2
+        pc_scale = (bounds[1] - bounds[0]).max() + 1e-5
+        
+        transform_params = torch.tensor([
+            center[0], center[1], center[2], scale,
+            pc_center[0], pc_center[1], pc_center[2], pc_scale
+        ], dtype=torch.float32)
+        
+        # Prepare batch data
+        pc_normal_normalized = np.concatenate([pc_coor, normals], axis=-1)
+        batch_data = {
+            'pc_normal': torch.from_numpy(pc_normal_normalized).half().unsqueeze(0).to(self.device),
+            'transform_params': transform_params.unsqueeze(0),
+            'vertices': torch.from_numpy(mesh.vertices).unsqueeze(0),
+            'faces': torch.from_numpy(mesh.faces).unsqueeze(0),
+            'file_name': [Path(input_path).stem]
+        }
+        
+        # Generate skeleton
+        pred_bone_coords = self.model.generate(batch_data)
+        
+        # Process results
+        file_name = Path(input_path).stem
+        skeleton = pred_bone_coords[0].cpu().numpy()
+        pred_joints, pred_bones = pred_joints_and_bones(skeleton.squeeze())
+        
+        # Post-process
+        hier_order = (sequence_type == "hierarchical")
+        if hier_order and len(pred_bones) > 0:
+            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
+            )
+        else:
+            pred_joints, pred_bones = merge_duplicate_joints_and_fix_bones(
+                pred_joints, pred_bones
+            )
+            pred_root_index = None
+        
+        # Denormalize for saving
+        trans = transform_params[:3].numpy()
+        scale_val = transform_params[3].item()
+        pc_trans = transform_params[4:7].numpy()
+        pc_scale_val = transform_params[7].item()
+        
+        pred_joints_denorm = pred_joints * pc_scale_val + pc_trans
+        pred_joints_denorm = pred_joints_denorm / scale_val + trans
+        
+        # Save files
+        pred_rig_filename = output_dir / f"{file_name}_pred.txt"
+        pred_skel_filename = output_dir / f"{file_name}_skel.obj"
+        mesh_filename = output_dir / f"{file_name}_mesh.obj"
+        
+        save_skeleton_to_txt(
+            pred_joints_denorm, pred_bones, pred_root_index,
+            hier_order, mesh.vertices, str(pred_rig_filename)
+        )
+        
+        save_skeleton_obj(
+            pred_joints, pred_bones, str(pred_skel_filename),
+            pred_root_index if hier_order else None,
+            use_cone=hier_order
+        )
+        
+        # Save normalized mesh
+        vertices_norm = (mesh.vertices - trans) * scale_val
+        vertices_norm = (vertices_norm - pc_trans) / pc_scale_val
+        save_mesh(vertices_norm, mesh.faces, str(mesh_filename))
+        
+        elapsed_time = time.time() - start_time
+        
+        return {
+            'skeleton_file': str(pred_skel_filename),
+            'rig_file': str(pred_rig_filename),
+            'mesh_file': str(mesh_filename),
+            'num_joints': len(pred_joints),
+            'num_bones': len(pred_bones),
+            'time': elapsed_time
+        }

packages.txt

@@ -0,0 +1,6 @@
+libgl1-mesa-glx
+libglib2.0-0
+libsm6
+libxext6
+libxrender-dev
+libgomp1