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()