Update app.py

app.py

@@ -1,172 +1,202 @@
 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.")
+import subprocess
 
 # --- 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")
+# Path to the cloned UniRig repository directory within the Space
+# IMPORTANT: You must clone the UniRig repository into this directory in your Hugging Face Space.
+UNIRIG_REPO_DIR = os.path.join(os.path.dirname(__file__), "UniRig")
+
+# Check if UniRig directory exists
+if not os.path.isdir(UNIRIG_REPO_DIR):
+    # This message will appear in logs, Gradio app might fail to start fully.
+    print(f"ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. Please clone it there.")
+    # Optionally, raise an error to make it more visible if the app starts
+    # raise RuntimeError(f"UniRig repository not found at {UNIRIG_REPO_DIR}. Please clone it there.")
 
-# 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}")
+    # Note: UniRig scripts might manage device internally or via Hydra configs.
 else:
-    print("CUDA not available, UniRig performance will be significantly slower on CPU.")
+    print("Warning: CUDA not available or not detected by PyTorch. UniRig performance will be severely impacted.")
+
+def run_unirig_command(command_args, step_name):
+    """Helper function to run UniRig commands using subprocess."""
+    python_exe = sys.executable  # Use the current python interpreter
+    cmd = [python_exe, "-m"] + command_args
+    
+    print(f"Running {step_name}: {' '.join(cmd)}")
+    
+    # UniRig scripts often expect to be run from the root of the UniRig repository
+    # because they use Hydra and relative paths for configs (e.g., conf/config.yaml)
+    process_env = os.environ.copy()
+    
+    # Add UniRig's parent directory to PYTHONPATH so `import unirig` works if needed,
+    # and the UniRig directory itself so its internal imports work.
+    # However, `python -m` typically handles package discovery well if CWD is correct.
+    # process_env["PYTHONPATH"] = f"{UNIRIG_REPO_DIR}{os.pathsep}{process_env.get('PYTHONPATH', '')}"
+
+    try:
+        # Execute the command
+        # It's crucial to set `cwd=UNIRIG_REPO_DIR` for Hydra to find its configs.
+        result = subprocess.run(cmd, cwd=UNIRIG_REPO_DIR, capture_output=True, text=True, check=True, env=process_env)
+        print(f"{step_name} output:\n{result.stdout}")
+        if result.stderr:
+            print(f"{step_name} errors (non-fatal if check=True passed):\n{result.stderr}")
+    except subprocess.CalledProcessError as e:
+        print(f"ERROR during {step_name}:")
+        print(f"Command: {' '.join(e.cmd)}")
+        print(f"Return code: {e.returncode}")
+        print(f"Stdout: {e.stdout}")
+        print(f"Stderr: {e.stderr}")
+        raise gr.Error(f"Error in UniRig {step_name}: {e.stderr[:500]}") # Show first 500 chars of error
+    except FileNotFoundError:
+        # This can happen if UNIRIG_REPO_DIR is not populated correctly or python_exe is wrong
+        print(f"ERROR: Could not find executable or script for {step_name}. Is UniRig cloned correctly in {UNIRIG_REPO_DIR}?")
+        raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs.")
 
 
 # --- Core Rigging Function ---
-def rig_glb_mesh(input_glb_file):
+def rig_glb_mesh_multistep(input_glb_file_obj):
     """
-    Takes an input GLB file, rigs it using UniRig, and returns the path to the rigged GLB file.
+    Takes an input GLB file object, rigs it using the new UniRig multi-step process,
+    and returns the path to the final rigged GLB file.
     """
-    if input_glb_file is None:
+    if not os.path.isdir(UNIRIG_REPO_DIR):
+         raise gr.Error(f"UniRig repository not found at {UNIRIG_REPO_DIR}. Cannot proceed.")
+
+    if input_glb_file_obj 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
+    input_glb_path = input_glb_file_obj.name  # Path to the temporary 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.")
+    # Create a dedicated temporary directory for all intermediate and final files for this run
+    # This helps in organization and cleanup.
+    processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_")
+    print(f"Using temporary processing directory: {processing_temp_dir}")
 
     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.")
+        # Define paths for intermediate and final files within the processing_temp_dir
+        # UniRig scripts expect output paths.
+        base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
         
-        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
-
+        # Step 1: Skeleton Prediction
+        # Output is typically an FBX file for the skeleton
+        temp_skeleton_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
+        print("Step 1: Predicting Skeleton...")
+        # Command: python -m unirig.predict_skeleton +input_path=<input_glb_path> +output_path=<temp_skeleton_path>
+        # Note: UniRig's scripts might have default output locations or require specific Hydra overrides.
+        # The `+` syntax is for Hydra overrides.
+        # Check UniRig's `conf/predict_skeleton.yaml` for default config values.
+        run_unirig_command([
+            "unirig.predict_skeleton",
+            f"input.path={input_glb_path}", # Use dot notation for Hydra parameters
+            f"output.path={temp_skeleton_path}",
+            # Add other necessary overrides, e.g., for device if not auto-detected well
+            # f"device={str(DEVICE)}" # If UniRig's script accepts this override
+        ], "Skeleton Prediction")
+        print(f"Skeleton predicted at: {temp_skeleton_path}")
+        if not os.path.exists(temp_skeleton_path):
+            raise gr.Error("Skeleton prediction failed to produce an output file.")
+
+        # Step 2: Skinning Weight Prediction
+        # Input: skeleton FBX and original GLB. Output: skinned FBX (or other format)
+        temp_skin_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx")
+        print("Step 2: Predicting Skinning Weights...")
+        # Command: python -m unirig.predict_skin +input_path=<temp_skeleton_path> +output_path=<temp_skin_path> +source_mesh_path=<input_glb_path>
+        run_unirig_command([
+            "unirig.predict_skin",
+            f"input.skeleton_path={temp_skeleton_path}", # Check exact Hydra param name in UniRig
+            f"input.source_mesh_path={input_glb_path}", # Check exact Hydra param name
+            f"output.path={temp_skin_path}",
+        ], "Skinning Prediction")
+        print(f"Skinning predicted at: {temp_skin_path}")
+        if not os.path.exists(temp_skin_path):
+            raise gr.Error("Skinning prediction failed to produce an output file.")
+
+        # Step 3: Merge Skeleton/Skin with Original Mesh
+        # Input: original GLB and the skin FBX (which contains skeleton + weights). Output: final rigged GLB
+        final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")
+        print("Step 3: Merging Results...")
+        # Command: python -m unirig.merge_skeleton_skin +source_path=<temp_skin_path> +target_path=<input_glb_path> +output_path=<final_rigged_glb_path>
+        run_unirig_command([
+            "unirig.merge_skeleton_skin",
+            f"input.source_rig_path={temp_skin_path}", # Path to the file with skeleton and skin weights
+            f"input.target_mesh_path={input_glb_path}", # Path to the original mesh
+            f"output.path={final_rigged_glb_path}",
+        ], "Merging")
+        print(f"Final rigged mesh at: {final_rigged_glb_path}")
+        if not os.path.exists(final_rigged_glb_path):
+            raise gr.Error("Merging process failed to produce the final rigged GLB file.")
+
+        # The final_rigged_glb_path needs to be accessible by Gradio to serve it.
+        # Gradio usually copies temp files it creates, but here we created it.
+        # We return the path, and Gradio should handle it.
+        # The processing_temp_dir will be cleaned up by Gradio if input_glb_file_obj is from gr.File
+        # or we can clean it up if we copy the final file to a Gradio managed temp location.
+        # For gr.Model3D, returning a path is fine.
+        return final_rigged_glb_path
+
+    except gr.Error: # Re-raise Gradio errors directly
+        raise
     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.
+        print(f"An unexpected error occurred: {e}")
+        # Clean up the processing directory in case of an unhandled error
+        if os.path.exists(processing_temp_dir):
+            shutil.rmtree(processing_temp_dir)
+        raise gr.Error(f"An unexpected error occurred during processing: {str(e)}")
+    # Note: Do not clean up processing_temp_dir in a `finally` block here if returning path from it,
+    # as Gradio needs the file to exist to serve it. Gradio's gr.File output type handles temp file cleanup.
+    # If outputting gr.File, copy the final file to a new tempfile managed by Gradio.
+    # For gr.Model3D, path is fine.
 
 # --- 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
+    primary_hue=gr.themes.colors.sky,
+    secondary_hue=gr.themes.colors.blue,
+    neutral_hue=gr.themes.colors.slate,
     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",
+    fn=rig_glb_mesh_multistep,
+    inputs=gr.File(label="Upload .glb Mesh File", type="file"),
+    outputs=gr.Model3D(
+        label="Rigged 3D Model (.glb)",
+        clear_color=[0.8, 0.8, 0.8, 1.0],
+        # Note: Model3D might have issues with complex GLBs or certain rigging structures.
+        # A gr.File output for download might be a safer fallback.
+        # outputs=[gr.Model3D(...), gr.File(label="Download Rigged GLB")]
+    ),
+    title="UniRig Auto-Rigger (Python 3.11 / PyTorch 2.3+)",
     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"
+        "Upload a 3D mesh in `.glb` format. This application uses the latest UniRig to automatically rig the mesh.\n"
+        "The process involves: 1. Skeleton Prediction, 2. Skinning Weight Prediction, 3. Merging.\n"
+        "This may take several minutes. Ensure your GLB has clean geometry.\n"
+        f"Running on: {str(DEVICE).upper()}. UniRig repo expected at: '{UNIRIG_REPO_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
+    cache_examples=False,
     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.")
+    # Perform a quick check for UniRig directory on launch
+    if not os.path.isdir(UNIRIG_REPO_DIR):
+        print(f"CRITICAL: UniRig repository not found at {UNIRIG_REPO_DIR}. The application will likely fail.")
+        # You could display this error in the Gradio interface itself using a dummy function or Markdown.
+    
+    # For local testing, you might need to set PYTHONPATH or ensure UniRig is installed.
+    # Example: os.environ["PYTHONPATH"] = f"{UNIRIG_REPO_DIR}{os.pathsep}{os.environ.get('PYTHONPATH', '')}"
+    
     iface.launch()