import gradio as gr import torch import os import sys import tempfile import shutil import subprocess # from huggingface_hub import HfApi, snapshot_download # For future model management if needed # import spaces # For @spaces.GPU decorator if you add it # --- Configuration --- # Path to the cloned UniRig repository directory within the Space UNIRIG_REPO_DIR = os.path.join(os.path.dirname(__file__), "UniRig") if not os.path.isdir(UNIRIG_REPO_DIR): print(f"ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. Please clone it there.") # Consider raising an error or displaying it in the UI if UniRig is critical for startup 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("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 # Ensure the command starts with the python executable and '-m' for module execution cmd = [python_exe, "-m"] + command_args print(f"Running {step_name}: {' '.join(cmd)}") process_env = os.environ.copy() # It's generally better to ensure UniRig's internal scripts handle PYTHONPATH if needed, # or that it's installed in a way that `python -m` works correctly from its root. # Setting cwd=UNIRIG_REPO_DIR is often the key for Hydra. try: # Execute the command from the UniRig directory for Hydra to find configs result = subprocess.run(cmd, cwd=UNIRIG_REPO_DIR, capture_output=True, text=True, check=True, env=process_env) print(f"{step_name} STDOUT:\n{result.stdout}") if result.stderr: print(f"{step_name} STDERR (non-fatal or warnings):\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}") # Provide a more user-friendly error, potentially masking long tracebacks error_summary = e.stderr.splitlines()[-5:] # Last 5 lines of stderr raise gr.Error(f"Error in UniRig {step_name}. Details: {' '.join(error_summary)}") except FileNotFoundError: print(f"ERROR: Could not find executable or script for {step_name}. Is UniRig cloned correctly in {UNIRIG_REPO_DIR} and Python environment setup?") raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs and UniRig directory structure.") except Exception as e_general: print(f"An unexpected Python exception occurred in run_unirig_command for {step_name}: {e_general}") raise gr.Error(f"Unexpected Python error during {step_name}: {str(e_general)[:500]}") # If you are using @spaces.GPU, you would import it: # import spaces # @spaces.GPU # You can specify type like @spaces.GPU(type="t4") or count def rig_glb_mesh_multistep(input_glb_file_obj): """ Takes an input GLB file object (from gr.File with type="filepath"), 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. Please check Space setup.") if input_glb_file_obj is None: # This case should ideally be handled by Gradio's input validation if `allow_none=False` (default) raise gr.Error("No input file provided. Please upload a .glb mesh.") # When type="filepath", input_glb_file_obj is the path string directly input_glb_path = input_glb_file_obj print(f"Input GLB path received: {input_glb_path}") # Create a dedicated temporary directory for all intermediate and final files processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_") print(f"Using temporary processing directory: {processing_temp_dir}") try: base_name = os.path.splitext(os.path.basename(input_glb_path))[0] # Step 1: Skeleton Prediction temp_skeleton_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx") print("Step 1: Predicting Skeleton...") run_unirig_command([ "unirig.predict_skeleton", f"input.path={os.path.abspath(input_glb_path)}", # Use absolute path for robustness f"output.path={os.path.abspath(temp_skeleton_path)}", # f"device={str(DEVICE)}" # If UniRig's script accepts this override and handles it ], "Skeleton Prediction") if not os.path.exists(temp_skeleton_path): raise gr.Error("Skeleton prediction failed to produce an output file. Check logs for UniRig errors.") # Step 2: Skinning Weight Prediction temp_skin_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx") print("Step 2: Predicting Skinning Weights...") run_unirig_command([ "unirig.predict_skin", f"input.skeleton_path={os.path.abspath(temp_skeleton_path)}", f"input.source_mesh_path={os.path.abspath(input_glb_path)}", f"output.path={os.path.abspath(temp_skin_path)}", ], "Skinning Prediction") if not os.path.exists(temp_skin_path): raise gr.Error("Skinning prediction failed to produce an output file. Check logs for UniRig errors.") # Step 3: Merge Skeleton/Skin with Original Mesh final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb") print("Step 3: Merging Results...") run_unirig_command([ "unirig.merge_skeleton_skin", f"input.source_rig_path={os.path.abspath(temp_skin_path)}", f"input.target_mesh_path={os.path.abspath(input_glb_path)}", f"output.path={os.path.abspath(final_rigged_glb_path)}", ], "Merging") if not os.path.exists(final_rigged_glb_path): raise gr.Error("Merging process failed to produce the final rigged GLB file. Check logs for UniRig errors.") # final_rigged_glb_path is in processing_temp_dir. # Gradio's gr.Model3D output component will handle serving this file. return final_rigged_glb_path except gr.Error: # Re-raise Gradio errors directly if os.path.exists(processing_temp_dir): # Clean up on known Gradio error shutil.rmtree(processing_temp_dir) print(f"Cleaned up temporary directory: {processing_temp_dir}") raise except Exception as e: print(f"An unexpected error occurred in rig_glb_mesh_multistep: {e}") if os.path.exists(processing_temp_dir): # Clean up on unexpected error shutil.rmtree(processing_temp_dir) print(f"Cleaned up temporary directory: {processing_temp_dir}") raise gr.Error(f"An unexpected error occurred during processing: {str(e)[:500]}") # --- 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"], ) # Ensure UNIRIG_REPO_DIR check happens before interface is built if it's critical if not os.path.isdir(UNIRIG_REPO_DIR) and __name__ == "__main__": # Check only if running as main script print(f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. The application will not work.") # Define the interface # Note: The @spaces.GPU decorator would go above the function `rig_glb_mesh_multistep` iface = gr.Interface( fn=rig_glb_mesh_multistep, inputs=gr.File( label="Upload .glb Mesh File", type="filepath" # Corrected type for Gradio 5.x.x ), outputs=gr.Model3D( label="Rigged 3D Model (.glb)", clear_color=[0.8, 0.8, 0.8, 1.0], ), 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: '{os.path.basename(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__": if not os.path.isdir(UNIRIG_REPO_DIR): print(f"CRITICAL: UniRig repository not found at {UNIRIG_REPO_DIR}. Ensure it's cloned in the Space's root.") iface.launch()