app.py

import gradio as gr
import torch
import os
import sys
import tempfile
import shutil
import subprocess
import spaces
from typing import Any, Dict, Union, List

# --- Configuration ---
UNIRIG_REPO_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "UniRig"))
BLENDER_INSTALL_DIR = "/opt/blender-4.2.0-linux-x64"
BLENDER_PYTHON_VERSION_DIR = "4.2"
BLENDER_PYTHON_VERSION = "python3.11"

# Construct paths
BLENDER_PYTHON_DIR = os.path.join(BLENDER_INSTALL_DIR, BLENDER_PYTHON_VERSION_DIR, "python")
BLENDER_PYTHON_BIN_DIR = os.path.join(BLENDER_PYTHON_DIR, "bin")
BLENDER_PYTHON_EXEC = os.path.join(BLENDER_PYTHON_BIN_DIR, BLENDER_PYTHON_VERSION)
BLENDER_PYTHON_LIB_PATH = os.path.join(BLENDER_PYTHON_DIR, "lib", BLENDER_PYTHON_VERSION)
BLENDER_PYTHON_SITE_PACKAGES = os.path.join(BLENDER_PYTHON_LIB_PATH, "site-packages")
# Path to the main blender executable (assuming symlink exists or using direct path)
BLENDER_EXEC = os.path.join(BLENDER_INSTALL_DIR, "blender") # Use direct path first
BLENDER_EXEC_SYMLINK = "/usr/local/bin/blender" # Fallback symlink

SETUP_SCRIPT = os.path.join(os.path.dirname(__file__), "setup_blender.sh")

# --- Initial Checks ---
print("--- Environment Checks ---")

# Check if main Blender executable exists
blender_executable_to_use = None
if os.path.exists(BLENDER_EXEC):
    print(f"Blender executable found at direct path: {BLENDER_EXEC}")
    blender_executable_to_use = BLENDER_EXEC
elif os.path.exists(BLENDER_EXEC_SYMLINK):
    print(f"Blender executable found via symlink: {BLENDER_EXEC_SYMLINK}")
    blender_executable_to_use = BLENDER_EXEC_SYMLINK
else:
    # Try running setup if Blender executable is missing
    print(f"Blender executable not found at {BLENDER_EXEC} or {BLENDER_EXEC_SYMLINK}. Running setup script...")
    if os.path.exists(SETUP_SCRIPT):
        try:
            setup_result = subprocess.run(["bash", SETUP_SCRIPT], check=True, capture_output=True, text=True)
            print("Setup script executed successfully.")
            print(f"Setup STDOUT:\n{setup_result.stdout}")
            if setup_result.stderr: print(f"Setup STDERR:\n{setup_result.stderr}")

            # Re-check for executable
            if os.path.exists(BLENDER_EXEC):
                blender_executable_to_use = BLENDER_EXEC
            elif os.path.exists(BLENDER_EXEC_SYMLINK):
                blender_executable_to_use = BLENDER_EXEC_SYMLINK

            if not blender_executable_to_use:
                 raise RuntimeError(f"Setup script ran but Blender executable still not found at {BLENDER_EXEC} or {BLENDER_EXEC_SYMLINK}.")

        except subprocess.CalledProcessError as e:
             print(f"ERROR running setup script: {SETUP_SCRIPT}\nStderr: {e.stderr}")
             raise gr.Error(f"Failed to execute setup script. Check logs. Stderr: {e.stderr[-500:]}")
        except Exception as e:
             raise gr.Error(f"Unexpected error running setup script '{SETUP_SCRIPT}': {e}")
    else:
        raise gr.Error(f"Blender executable not found and setup script missing: {SETUP_SCRIPT}")

# Check Python executable (still useful for checks)
if not os.path.exists(BLENDER_PYTHON_EXEC):
     print(f"WARNING: Blender Python executable not found at {BLENDER_PYTHON_EXEC}, though main executable exists.")


# Verify Blender Python site-packages path and attempt bpy import test
bpy_import_ok = False
if os.path.exists(BLENDER_PYTHON_SITE_PACKAGES):
    print(f"Blender Python site-packages found at: {BLENDER_PYTHON_SITE_PACKAGES}")
    # Try importing bpy using blender --background --python
    try:
        test_script_content = "import bpy; print('bpy imported successfully')"
        test_result = subprocess.run(
            [blender_executable_to_use, "--background", "--python-expr", test_script_content],
            capture_output=True, text=True, check=True, timeout=30 # Add timeout
        )
        if "bpy imported successfully" in test_result.stdout:
            print("Successfully imported 'bpy' using Blender executable.")
            bpy_import_ok = True
        else:
            print(f"WARNING: 'bpy' import test via Blender returned unexpected output:\nSTDOUT:{test_result.stdout}\nSTDERR:{test_result.stderr}")

    except subprocess.TimeoutExpired:
         print("WARNING: 'bpy' import test via Blender timed out.")
    except subprocess.CalledProcessError as e:
        print(f"WARNING: Failed to import 'bpy' using Blender executable:\nSTDOUT:{e.stdout}\nSTDERR:{e.stderr}")
    except Exception as e:
        print(f"WARNING: Unexpected error during 'bpy' import test: {e}")

else:
    print(f"WARNING: Blender Python site-packages directory not found at {BLENDER_PYTHON_SITE_PACKAGES}. Check paths.")


# Check for UniRig repository
if not os.path.isdir(UNIRIG_REPO_DIR):
    raise gr.Error(f"UniRig repository missing at: {UNIRIG_REPO_DIR}. Ensure it's cloned correctly.")
else:
    print(f"UniRig repository found at: {UNIRIG_REPO_DIR}")

# Check PyTorch and CUDA
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Gradio environment using device: {DEVICE}")
if DEVICE.type == 'cuda':
    try:
        print(f"Gradio CUDA Device Name: {torch.cuda.get_device_name(0)}")
        print(f"Gradio PyTorch CUDA Version: {torch.version.cuda}")
    except Exception as e:
        print(f"Could not get CUDA device details: {e}")
else:
    print("Warning: Gradio environment CUDA not available.")

print("--- End Environment Checks ---")

# --- Helper Functions ---

def patch_asset_py():
    """Temporary patch to fix type hinting error in UniRig's asset.py"""
    asset_py_path = os.path.join(UNIRIG_REPO_DIR, "src", "data", "asset.py")
    try:
        if not os.path.exists(asset_py_path):
            print(f"Warning: asset.py not found at {asset_py_path}, skipping patch.")
            return

        with open(asset_py_path, "r") as f: content = f.read()
        problematic_line = "meta: Union[Dict[str, ...], None]=None"
        corrected_line = "meta: Union[Dict[str, Any], None]=None"
        typing_import = "from typing import Any"

        if corrected_line in content:
            print("Patch already applied to asset.py"); return
        if problematic_line not in content:
             print("Problematic line not found in asset.py, patch might be unnecessary."); return

        print("Applying patch to asset.py...")
        content = content.replace(problematic_line, corrected_line)
        if typing_import not in content:
            if "from typing import" in content: content = content.replace("from typing import", f"{typing_import}\nfrom typing import", 1)
            else: content = f"{typing_import}\n{content}"
        with open(asset_py_path, "w") as f: f.write(content)
        print("Successfully patched asset.py")

    except Exception as e:
        print(f"ERROR: Failed to patch asset.py: {e}. Proceeding cautiously.")

@spaces.GPU
def run_unirig_command(python_script_path: str, script_args: List[str], step_name: str):
    """
    Runs a specific UniRig PYTHON script (.py) using the Blender executable
    in background mode (`blender --background --python script.py -- args`).

    Args:
        python_script_path: Absolute path to the .py script to execute within Blender.
        script_args: A list of command-line arguments FOR THE PYTHON SCRIPT.
        step_name: Name of the step for logging.
    """
    if not blender_executable_to_use:
         raise gr.Error("Blender executable path could not be determined. Cannot run UniRig step.")

    # Command structure: blender --background --python script.py -- script_arg1 script_arg2 ...
    cmd = [
        blender_executable_to_use,
        "--background",
        "--python", python_script_path,
        "--" # Separator for script arguments
    ] + script_args

    print(f"\n--- Running UniRig Step: {step_name} ---")
    print(f"Command: {' '.join(cmd)}") # Note: Simple join for logging

    # Environment variables might still be useful if the script imports other non-blender libs
    process_env = os.environ.copy()
    # Ensure UniRig source is findable if scripts use relative imports from repo root
    # or if they import modules from UniRig/src
    unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, "src")
    pythonpath_parts = [
        # Blender's site-packages should be found automatically when run via Blender exec
        unirig_src_dir,
        UNIRIG_REPO_DIR
    ]
    # Add existing PYTHONPATH if any
    existing_pythonpath = process_env.get('PYTHONPATH', '')
    if existing_pythonpath:
        pythonpath_parts.append(existing_pythonpath)
    process_env["PYTHONPATH"] = os.pathsep.join(filter(None, pythonpath_parts))
    print(f"Subprocess PYTHONPATH: {process_env['PYTHONPATH']}")

    # LD_LIBRARY_PATH might still be needed for non-python shared libs used by dependencies
    blender_lib_path = os.path.join(BLENDER_PYTHON_DIR, "lib")
    existing_ld_path = process_env.get('LD_LIBRARY_PATH', '')
    process_env["LD_LIBRARY_PATH"] = f"{blender_lib_path}{os.pathsep}{existing_ld_path}" if existing_ld_path else blender_lib_path
    print(f"Subprocess LD_LIBRARY_PATH: {process_env['LD_LIBRARY_PATH']}")


    try:
        # Execute Blender with the Python script.
        # cwd=UNIRIG_REPO_DIR ensures script runs relative to repo root.
        result = subprocess.run(
            cmd,
            cwd=UNIRIG_REPO_DIR,
            capture_output=True,
            text=True,
            check=True, # Raises CalledProcessError on non-zero exit codes
            env=process_env # Pass the modified environment
        )
        print(f"{step_name} STDOUT:\n{result.stdout}")
        # Blender often prints info/warnings to stderr even on success
        if result.stderr:
            print(f"{step_name} STDERR (Info/Warnings):\n{result.stderr}")
            # Check stderr for common Blender errors even if exit code was 0
            if "Error: Cannot read file" in result.stderr:
                 raise gr.Error(f"Blender reported an error reading an input file during {step_name}. Check paths and file integrity.")
            # Add other specific Blender error checks if needed

    except subprocess.CalledProcessError as e:
        print(f"ERROR during {step_name}: Subprocess failed!")
        print(f"Command: {' '.join(e.cmd)}")
        print(f"Return code: {e.returncode}")
        print(f"--- {step_name} STDOUT ---:\n{e.stdout}")
        print(f"--- {step_name} STDERR ---:\n{e.stderr}")
        error_summary = e.stderr.strip().splitlines()
        last_lines = "\n".join(error_summary[-5:]) if error_summary else "No stderr output."
        # Check specifically for bpy/torch import errors within the subprocess stderr
        # These *shouldn't* happen now, but check just in case
        if "ModuleNotFoundError: No module named 'bpy'" in e.stderr:
             raise gr.Error(f"Error in UniRig '{step_name}': Blender failed to provide 'bpy' module internally.")
        elif "ImportError: Failed to load PyTorch C extensions" in e.stderr:
             raise gr.Error(f"Error in UniRig '{step_name}': Script failed to load PyTorch extensions even within Blender. Check installation.")
        else:
             raise gr.Error(f"Error in UniRig '{step_name}'. Check logs. Last error lines:\n{last_lines}")


    except FileNotFoundError:
        # This error means blender executable or the python script wasn't found
        print(f"ERROR: Could not find Blender executable '{blender_executable_to_use}' or script '{python_script_path}' for {step_name}.")
        print(f"Attempted command: {' '.join(cmd)}")
        raise gr.Error(f"Setup error for UniRig '{step_name}'. Blender or Python script not found.")

    except Exception as e_general:
        print(f"An unexpected Python exception occurred in run_unirig_command for {step_name}: {e_general}")
        import traceback
        traceback.print_exc()
        raise gr.Error(f"Unexpected Python error during '{step_name}' execution: {str(e_general)[:500]}")

    print(f"--- Finished UniRig Step: {step_name} ---")


@spaces.GPU
def rig_glb_mesh_multistep(input_glb_file_obj):
    """
    Main Gradio function to rig a GLB mesh using UniRig's multi-step process.
    Orchestrates calls to run_unirig_command for each step, executing .py scripts via Blender.
    """
    try:
        patch_asset_py() # Attempt patch
    except Exception as e:
         print(f"Ignoring patch error: {e}")

    # --- Input Validation ---
    if input_glb_file_obj is None: raise gr.Error("No input file provided.")
    input_glb_path = input_glb_file_obj
    print(f"Input GLB path received: {input_glb_path}")
    if not os.path.exists(input_glb_path): raise gr.Error(f"Input file path does not exist: {input_glb_path}")
    if not input_glb_path.lower().endswith(".glb"): raise gr.Error("Invalid file type. Please upload a .glb file.")

    # --- Setup Temporary Directory ---
    processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_")
    print(f"Using temporary processing directory: {processing_temp_dir}")

    try:
        # --- Define File Paths ---
        base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
        abs_input_glb_path = os.path.abspath(input_glb_path)
        abs_skeleton_output_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
        abs_skin_output_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx")
        abs_final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")

        # --- Define Absolute Paths to UniRig PYTHON Scripts ---
        # *** IMPORTANT: Assuming the .sh scripts primarily wrap these .py scripts. ***
        # *** This might need adjustment based on the actual content of the .sh files. ***
        # Use run.py as the main entry point, passing task-specific configs/args
        # This requires inspecting run.py and the shell scripts to know the correct args/configs.
        # For now, let's assume a structure like: python run.py --config <config_file> --input ... --output ...

        # Example: Assuming run.py takes task name and args
        run_script_path = os.path.join(UNIRIG_REPO_DIR, "run.py") # Main script?

        # --- Execute UniRig Steps ---

        # Step 1: Skeleton Prediction
        print("\nStarting Step 1: Predicting Skeleton...")
        # Arguments for the run.py script for skeleton task
        # ** These are GUESSES based on typical structure - VERIFY from UniRig code **
        skeleton_args = [
            # Might need a config file path or task name argument first
            # e.g., "--task", "generate_skeleton",
            "--input", abs_input_glb_path,
            "--output", abs_skeleton_output_path
            # Add other relevant args like model path if needed
        ]
        if not os.path.exists(run_script_path): # Check if assumed script exists
            raise gr.Error(f"UniRig main script not found at: {run_script_path}")
        # Execute run.py via Blender
        run_unirig_command(run_script_path, skeleton_args, "Skeleton Prediction")
        if not os.path.exists(abs_skeleton_output_path):
            raise gr.Error("Skeleton prediction failed. Output file not created. Check logs.")
        print("Step 1: Skeleton Prediction completed.")

        # Step 2: Skinning Weight Prediction
        print("\nStarting Step 2: Predicting Skinning Weights...")
        # Arguments for the run.py script for skinning task
        # ** GUESSES - VERIFY from UniRig code **
        skin_args = [
            # e.g., "--task", "generate_skin",
            "--input", abs_skeleton_output_path, # Input is the skeleton from step 1
            "--source", abs_input_glb_path,    # Source mesh (if needed by skin script)
            "--output", abs_skin_output_path
        ]
        run_unirig_command(run_script_path, skin_args, "Skinning Prediction")
        if not os.path.exists(abs_skin_output_path):
            raise gr.Error("Skinning prediction failed. Output file not created. Check logs.")
        print("Step 2: Skinning Prediction completed.")

        # Step 3: Merge Skeleton/Skin with Original Mesh
        print("\nStarting Step 3: Merging Results...")
        # Arguments for the run.py script for merging task
        # ** GUESSES - VERIFY from UniRig code **
        # Alternatively, merge might use a different script like src/data/merge_fbx.py?
        # Let's assume run.py handles it for now.
        merge_args = [
            # e.g., "--task", "merge",
            "--source", abs_skin_output_path,
            "--target", abs_input_glb_path,
            "--output", abs_final_rigged_glb_path
        ]
        run_unirig_command(run_script_path, merge_args, "Merging")
        if not os.path.exists(abs_final_rigged_glb_path):
            raise gr.Error("Merging process failed. Final rigged GLB file not created. Check logs.")
        print("Step 3: Merging completed.")

        # --- Return Result ---
        print(f"Successfully generated rigged model: {abs_final_rigged_glb_path}")
        return abs_final_rigged_glb_path

    except gr.Error as e:
        print(f"Gradio Error occurred: {e}")
        if os.path.exists(processing_temp_dir): shutil.rmtree(processing_temp_dir); print(f"Cleaned up temp dir: {processing_temp_dir}")
        raise e
    except Exception as e:
        print(f"An unexpected error occurred in rig_glb_mesh_multistep: {e}")
        import traceback; traceback.print_exc()
        if os.path.exists(processing_temp_dir): shutil.rmtree(processing_temp_dir); print(f"Cleaned up temp dir: {processing_temp_dir}")
        raise gr.Error(f"An unexpected error occurred during processing: {str(e)[:500]}")


# --- Gradio Interface Definition ---
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"],
)

# Check UniRig repo existence again before building the interface
if not os.path.isdir(UNIRIG_REPO_DIR):
     startup_error_message = (f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}.")
     print(startup_error_message)
     with gr.Blocks(theme=theme) as iface: gr.Markdown(f"# Application Error\n\n{startup_error_message}")
elif not blender_executable_to_use:
     startup_error_message = (f"CRITICAL STARTUP ERROR: Blender executable not found.")
     print(startup_error_message)
     with gr.Blocks(theme=theme) as iface: gr.Markdown(f"# Application Error\n\n{startup_error_message}")
else:
     # Build the normal interface if UniRig is found
     iface = gr.Interface(
         fn=rig_glb_mesh_multistep,
         inputs=gr.File(label="Upload .glb Mesh File", type="filepath", file_types=[".glb"]),
         outputs=gr.Model3D(label="Rigged 3D Model (.glb)", clear_color=[0.8, 0.8, 0.8, 1.0]),
         title=f"UniRig Auto-Rigger (Blender {BLENDER_PYTHON_VERSION_DIR} / Python {BLENDER_PYTHON_VERSION})",
         description=(
             "Upload a 3D mesh in `.glb` format. This application uses UniRig via Blender's Python interface.\n"
             f"* Running main app on Python {sys.version.split()[0]}, UniRig steps use Blender's Python {BLENDER_PYTHON_VERSION}.\n"
             f"* Utilizing device: **{DEVICE.type.upper()}** (via ZeroGPU if available).\n"
             f"* UniRig Source: https://github.com/VAST-AI-Research/UniRig"
         ),
         cache_examples=False, theme=theme, allow_flagging='never'
     )

# --- Launch the Application ---
if __name__ == "__main__":
    if 'iface' in locals():
        print("Launching Gradio interface...")
        iface.launch()
    else:
        print("ERROR: Gradio interface not created due to startup errors.")