Create app.py

app.py

@@ -0,0 +1,172 @@
+import gradio as gr
+import torch
+import trimesh
+import os
+import sys
+import tempfile
+import shutil
+
+# Add UniRig source directory to Python path
+# Assuming UniRig files are in a subdirectory named 'UniRig_src'
+sys.path.append(os.path.join(os.path.dirname(__file__), 'UniRig_src'))
+
+# Conditional import for AutoRigger and setup_source_mesh
+# This helps in providing a clearer error if UniRig_src is not found
+try:
+    from autorig import AutoRigger
+    from utils import setup_source_mesh
+except ImportError as e:
+    print("Error importing from UniRig_src. Make sure the UniRig source files are in the 'UniRig_src' directory.")
+    print(f"Details: {e}")
+    # Define dummy functions if import fails, so Gradio can still load with an error message
+    def AutoRigger(*args, **kwargs):
+        raise RuntimeError("UniRig AutoRigger could not be loaded. Check UniRig_src setup.")
+    def setup_source_mesh(mesh, *args, **kwargs):
+        raise RuntimeError("UniRig setup_source_mesh could not be loaded. Check UniRig_src setup.")
+
+# --- Configuration ---
+# Define paths to the UniRig model files
+# These files should be placed in the 'model_files' directory in your Hugging Face Space
+MODEL_DIR = os.path.join(os.path.dirname(__file__), "model_files")
+SMPL_SKELETON_PATH = os.path.join(MODEL_DIR, "smpl_skeleton.pkl")
+SKIN_KPS_PREDICTOR_PATH = os.path.join(MODEL_DIR, "skin_kps_predictor.pkl")
+
+# Check if model files exist
+if not os.path.exists(SMPL_SKELETON_PATH) or not os.path.exists(SKIN_KPS_PREDICTOR_PATH):
+    print(f"Warning: Model files not found at {MODEL_DIR}. Please ensure smpl_skeleton.pkl and skin_kps_predictor.pkl are present.")
+
+# Determine processing device (CUDA if available, otherwise CPU)
+# ZeroGPU on Hugging Face Spaces should provide CUDA
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {DEVICE}")
+if DEVICE.type == 'cuda':
+    print(f"CUDA Device Name: {torch.cuda.get_device_name(0)}")
+    print(f"CUDA Version: {torch.version.cuda}")
+else:
+    print("CUDA not available, UniRig performance will be significantly slower on CPU.")
+
+
+# --- Core Rigging Function ---
+def rig_glb_mesh(input_glb_file):
+    """
+    Takes an input GLB file, rigs it using UniRig, and returns the path to the rigged GLB file.
+    """
+    if input_glb_file is None:
+        raise gr.Error("No input file provided. Please upload a .glb mesh.")
+
+    input_glb_path = input_glb_file.name  # Get the path of the uploaded file
+
+    # Ensure UniRig components are loaded (they might be dummy if import failed)
+    if not callable(getattr(AutoRigger, '__init__', None)) or not callable(setup_source_mesh):
+         raise gr.Error("UniRig components are not correctly loaded. Please check the server logs and UniRig_src setup.")
+
+    try:
+        # Create a temporary directory for output
+        temp_dir = tempfile.mkdtemp()
+        output_glb_filename = "rigged_output.glb"
+        output_glb_path = os.path.join(temp_dir, output_glb_filename)
+
+        # 1. Load the mesh using trimesh
+        print(f"Loading mesh from: {input_glb_path}")
+        mesh = trimesh.load_mesh(input_glb_path, force='mesh', process=False)
+        
+        if not isinstance(mesh, trimesh.Trimesh):
+            # If it's a Scene object, try to get a single geometry
+            if isinstance(mesh, trimesh.Scene):
+                if len(mesh.geometry) == 0:
+                    raise gr.Error("Input GLB file contains no mesh geometry.")
+                # Concatenate all meshes in the scene into a single mesh
+                # This is a common approach, but might not be ideal for all GLB files
+                print(f"Input is a scene with {len(mesh.geometry)} geometries. Attempting to merge.")
+                mesh = trimesh.util.concatenate(list(mesh.geometry.values()))
+                if not isinstance(mesh, trimesh.Trimesh):
+                     raise gr.Error(f"Could not extract a valid mesh from the GLB scene. Found type: {type(mesh)}")
+            else:
+                raise gr.Error(f"Failed to load a valid mesh from the input file. Loaded type: {type(mesh)}")
+
+        print("Mesh loaded successfully.")
+
+        # 2. Preprocess the mesh (as per UniRig's example)
+        # This step is crucial for UniRig to work correctly.
+        # It involves canonicalization and remeshing.
+        print("Preprocessing mesh...")
+        mesh = setup_source_mesh(mesh, device=DEVICE)
+        print("Mesh preprocessing complete.")
+
+        # 3. Initialize the AutoRigger
+        # Ensure model files are accessible
+        if not os.path.exists(SMPL_SKELETON_PATH) or not os.path.exists(SKIN_KPS_PREDICTOR_PATH):
+            raise gr.Error(f"UniRig model files not found. Searched in {MODEL_DIR}. Please check your Space's file structure.")
+        
+        print("Initializing AutoRigger...")
+        autorigger = AutoRigger(SMPL_SKELETON_PATH, SKIN_KPS_PREDICTOR_PATH, device=DEVICE)
+        print("AutoRigger initialized.")
+
+        # 4. Perform rigging
+        print("Starting rigging process...")
+        # The `rig` method might require specific verts, faces, and normals if not handled by `setup_source_mesh`
+        # Assuming `setup_source_mesh` prepares it adequately.
+        output_dict = autorigger.rig(mesh)
+        print("Rigging process complete.")
+
+        # 5. Extract the rigged mesh
+        rigged_mesh = output_dict['rigged_mesh'] # This should be a trimesh.Trimesh object
+        print("Rigged mesh extracted.")
+
+        # 6. Export the rigged mesh to GLB format
+        print(f"Exporting rigged mesh to: {output_glb_path}")
+        rigged_mesh.export(output_glb_path)
+        print("Export complete.")
+
+        return output_glb_path
+
+    except Exception as e:
+        print(f"Error during rigging: {e}")
+        # Clean up temp dir in case of error
+        if 'temp_dir' in locals() and os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir)
+        # Re-raise as Gradio error to display to user
+        raise gr.Error(f"An error occurred during processing: {str(e)}")
+    # No finally block for shutil.rmtree(temp_dir) here,
+    # because Gradio needs the file path to serve it.
+    # Gradio handles cleanup of temporary files created by gr.File.
+
+# --- Gradio Interface ---
+# Define a custom theme (Blue and Charcoal Gray)
+# Using Soft theme with sky blue and slate gray
+theme = gr.themes.Soft(
+    primary_hue=gr.themes.colors.sky,    # A nice blue
+    secondary_hue=gr.themes.colors.blue, # Can be same as primary or a complementary blue
+    neutral_hue=gr.themes.colors.slate,  # Charcoal gray
+    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
+).set(
+    # Further fine-tuning if needed
+    # button_primary_background_fill="*primary_500",
+    # button_primary_text_color="white",
+)
+
+# Interface definition
+iface = gr.Interface(
+    fn=rig_glb_mesh,
+    inputs=gr.File(label="Upload .glb Mesh File", type="file"), # 'file' gives a NamedTemporaryFile object
+    outputs=gr.Model3D(label="Rigged 3D Model (.glb)", clear_color=[0.8, 0.8, 0.8, 1.0]), # Model3D can display .glb
+    title="UniRig Auto-Rigger for 3D Meshes",
+    description=(
+        "Upload a 3D mesh in `.glb` format. This application uses UniRig to automatically rig the mesh.\n"
+        "The process may take a few minutes, especially for complex meshes. Ensure your GLB has clean geometry.\n"
+        f"Running on: {str(DEVICE).upper()}. Model files expected in '{MODEL_DIR}'.\n"
+        f"UniRig Source: https://github.com/VAST-AI-Research/UniRig"
+    ),
+    examples=[
+        # Add paths to example GLB files if you include them in your Space
+        # e.g., [os.path.join(os.path.dirname(__file__), "examples/sample_mesh.glb")]
+    ],
+    cache_examples=False, # Set to True if you have static examples and want to pre-process them
+    theme=theme,
+    allow_flagging="never"
+)
+
+if __name__ == "__main__":
+    if not os.path.exists(os.path.join(os.path.dirname(__file__), 'UniRig_src')):
+        print("CRITICAL: 'UniRig_src' directory not found. Please ensure UniRig source files are correctly placed.")
+    iface.launch()