app.py

import gradio as gr
import torch
import os
import sys
import tempfile
import shutil
import subprocess

# --- 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.")

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_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=<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"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()