import gradio as gr import torch import os import sys import tempfile import shutil import subprocess import spaces # --- Configuration --- # 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.") # Determine processing device (CUDA if available, otherwise CPU) 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("Warning: CUDA not available or not detected by PyTorch. UniRig performance will be severely impacted.") @spaces.GPU 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 --- @spaces.GPU def rig_glb_mesh_multistep(input_glb_file_obj): """ 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 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_obj.name # Path to the temporary uploaded file # 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: # 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] # 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= +output_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= +output_path= +source_mesh_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= +target_path= +output_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"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 --- theme = gr.themes.Soft( 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"], ) iface = gr.Interface( 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 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" ), cache_examples=False, theme=theme, allow_flagging="never" ) if __name__ == "__main__": # 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()