Update app.py

app.py

@@ -5,32 +5,27 @@ import sys
 import tempfile
 import shutil
 import subprocess
-import spaces
-from typing import Any, Dict, Union, List
+import spaces # Keep this if you use @spaces.GPU
+from typing import Any, Dict, List, Union # Added Union
 
 # --- 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"
+BLENDER_INSTALL_DIR = "/opt/blender-4.2.0-linux-x64" # From your setup_blender.sh
+BLENDER_PYTHON_VERSION_DIR = "4.2" # From your app.py
+BLENDER_PYTHON_VERSION = "python3.11" # From your app.py and UniRig
 
 # 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
+# Use Blender's main executable to run Python scripts via --python flag
+BLENDER_EXEC = os.path.join(BLENDER_INSTALL_DIR, "blender")
 BLENDER_EXEC_SYMLINK = "/usr/local/bin/blender" # Fallback symlink
 
 SETUP_SCRIPT = os.path.join(os.path.dirname(__file__), "setup_blender.sh")
 
-# --- Initial Checks ---
+# --- Initial Checks (adapted from your script) ---
 print("--- Environment Checks ---")
-
-# Check if main Blender executable exists
-blender_executable_to_use = None
+blender_executable_to_use = None # Use the main blender executable
 if os.path.exists(BLENDER_EXEC):
     print(f"Blender executable found at direct path: {BLENDER_EXEC}")
     blender_executable_to_use = BLENDER_EXEC
@@ -38,7 +33,6 @@ 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:
@@ -50,30 +44,25 @@ else:
             # Re-check for executable
             if os.path.exists(BLENDER_EXEC):
                 blender_executable_to_use = BLENDER_EXEC
+                print(f"Blender executable now found at: {BLENDER_EXEC}")
             elif os.path.exists(BLENDER_EXEC_SYMLINK):
-                blender_executable_to_use = BLENDER_EXEC_SYMLINK
+                 blender_executable_to_use = BLENDER_EXEC_SYMLINK
+                 print(f"Blender executable now found via symlink: {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:]}")
+            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}")
+            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
+# Verify bpy import using the found Blender executable
 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
+if blender_executable_to_use:
     try:
         test_script_content = "import bpy; print('bpy imported successfully')"
         test_result = subprocess.run(
@@ -87,14 +76,13 @@ if os.path.exists(BLENDER_PYTHON_SITE_PACKAGES):
             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.")
+        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.")
+     print("WARNING: Cannot test bpy import as Blender executable was not found.")
 
 
 # Check for UniRig repository
@@ -102,14 +90,21 @@ 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 for UniRig's run.py
+    UNIRIG_RUN_PY = os.path.join(UNIRIG_REPO_DIR, "run.py")
+    if not os.path.exists(UNIRIG_RUN_PY):
+        raise gr.Error(f"UniRig's run.py not found at {UNIRIG_RUN_PY}. Check UniRig clone.")
 
-# Check PyTorch and CUDA
+
+# Check PyTorch and CUDA for Gradio environment
 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}")
+        # Note: torch.version.cuda might show the version PyTorch was *built* with,
+        # not necessarily the runtime driver version. Check nvidia-smi in Space terminal if needed.
+        print(f"Gradio PyTorch CUDA Built Version: {torch.version.cuda}")
     except Exception as e:
         print(f"Could not get CUDA device details: {e}")
 else:
@@ -117,8 +112,7 @@ else:
 
 print("--- End Environment Checks ---")
 
-# --- Helper Functions ---
-
+# --- Helper Functions (Keep patch_asset_py if needed) ---
 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")
@@ -128,106 +122,150 @@ def patch_asset_py():
             return
 
         with open(asset_py_path, "r") as f: content = f.read()
+        # Adjust problematic/corrected lines if UniRig code changes
         problematic_line = "meta: Union[Dict[str, ...], None]=None"
-        corrected_line = "meta: Union[Dict[str, Any], None]=None"
+        corrected_line = "meta: Union[Dict[str, Any], None]=None" # Requires `from typing import Any`
         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("Problematic line not found in asset.py, patch might be unnecessary or file changed."); return
 
         print("Applying patch to asset.py...")
         content = content.replace(problematic_line, corrected_line)
+        # Ensure 'Any' is imported
         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}"
+            if "from typing import" in content:
+                 # Add 'Any' to an existing import line
+                 content = content.replace("from typing import", f"{typing_import}\nfrom typing import", 1)
+            else:
+                 # Add the import line at the top
+                 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.")
+# --- End Helper Functions ---
 
+
+# Decorator for ZeroGPU if needed
 @spaces.GPU
-def run_unirig_command(python_script_path: str, script_args: List[str], step_name: str):
+def run_unirig_command(unirig_run_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.
+        unirig_run_script_path: Absolute path to the UniRig run.py script.
         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
+        raise gr.Error("Blender executable path could not be determined. Cannot run UniRig step.")
 
-    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
+    # --- Environment Setup for Subprocess ---
     process_env = os.environ.copy()
-    unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, "src")
 
-    # Set PYTHONPATH: Include UniRig source dir, UniRig base dir, and '.' for CWD.
-    pythonpath_parts = [
-        unirig_src_dir,
-        UNIRIG_REPO_DIR,
-        '.' # Add current working directory explicitly
-    ]
+    # PYTHONPATH: Only need UniRig's root directory. Python automatically checks CWD.
+    # Add UniRig root to the start of PYTHONPATH.
+    pythonpath_parts = [UNIRIG_REPO_DIR]
     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")
+    # LD_LIBRARY_PATH: Inherit system path (for CUDA) and add Blender's libraries.
+    blender_main_lib_path = os.path.join(BLENDER_INSTALL_DIR, "lib")
+    blender_python_lib_path = os.path.join(BLENDER_PYTHON_DIR, "lib") # Libs for the interpreter itself
+    ld_path_parts = []
+    # Add Blender's libs first so they are preferred if needed
+    if os.path.exists(blender_main_lib_path): ld_path_parts.append(blender_main_lib_path)
+    if os.path.exists(blender_python_lib_path): ld_path_parts.append(blender_python_lib_path)
+    # Inherit existing LD_LIBRARY_PATH (should include CUDA from ZeroGPU)
     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']}")
-
-    # Test import of 'src' module
-    test_cmd = [blender_executable_to_use, "--background", "--python", "-c", "import src; print('src module imported successfully')"]
-    
-    test_result = subprocess.run(
-        test_cmd,
-        cwd=UNIRIG_REPO_DIR,
-        capture_output=True,
-        text=True,
-        env=process_env
-    )
-    if test_result.returncode != 0:
-        print(f"Test import of 'src' failed with return code {test_result.returncode}")
-        print(f"Test STDOUT: {test_result.stdout}")
-        print(f"Test STDERR: {test_result.stderr}")
-        raise gr.Error("Failed to import 'src' module in Blender environment. Check PYTHONPATH and module structure.")
-    elif "src module imported successfully" not in test_result.stdout:
-        print(f"Test import of 'src' did not produce expected output.")
-        print(f"Test STDOUT: {test_result.stdout}")
-        print(f"Test STDERR: {test_result.stderr}")
-        raise gr.Error("Unexpected output from 'src' import test. Check script behavior.")
-    else:
-        print("Test import of 'src' succeeded.")
+    if existing_ld_path: ld_path_parts.append(existing_ld_path)
+
+    if ld_path_parts:
+        process_env["LD_LIBRARY_PATH"] = os.pathsep.join(filter(None, ld_path_parts))
+    print(f"Subprocess LD_LIBRARY_PATH: {process_env.get('LD_LIBRARY_PATH', 'Not set')}")
+
+    # --- Test Import of 'src' module ---
+    # *** FIX: Use '--python-expr' which is equivalent to '--python -c' but safer with args ***
+    test_cmd = [
+        blender_executable_to_use,
+        "--background",
+        "--python-expr", "import src; print('src module imported successfully')" # Use --python-expr
+    ]
+    print(f"Running test import command: {' '.join(test_cmd)}")
+    try:
+        test_result = subprocess.run(
+            test_cmd,
+            cwd=UNIRIG_REPO_DIR, # Run test from UniRig root
+            capture_output=True,
+            text=True,
+            check=True, # Raise error on non-zero exit
+            env=process_env,
+            timeout=30
+        )
+        if "src module imported successfully" not in test_result.stdout:
+            print(f"Test import of 'src' did not produce expected output.")
+            print(f"Test STDOUT: {test_result.stdout}")
+            print(f"Test STDERR: {test_result.stderr}")
+            # Don't raise here, maybe it prints elsewhere, but log warning
+            print("WARNING: Unexpected output from 'src' import test. Proceeding with main command.")
+        else:
+            print("Test import of 'src' succeeded.")
+
+    except subprocess.TimeoutExpired:
+         print("ERROR: Test import of 'src' timed out.")
+         raise gr.Error("Failed to verify 'src' module import within Blender environment (timeout). Check setup.")
+    except subprocess.CalledProcessError as e:
+        print(f"Test import of 'src' failed with return code {e.returncode}")
+        print(f"Test Command: {' '.join(e.cmd)}")
+        print(f"Test CWD: {UNIRIG_REPO_DIR}")
+        print(f"Test STDOUT: {e.stdout}")
+        print(f"Test STDERR: {e.stderr}")
+        # Check common errors
+        if "ModuleNotFoundError: No module named 'src'" in e.stderr:
+             error_msg = "Failed to import 'src' module in Blender environment. Check PYTHONPATH and that UniRig repo structure is correct."
+        elif "could not get a list of mounted file-systems" in e.stderr:
+             error_msg = "Blender environment issue ('could not get mounted file-systems'). Check Space permissions or Blender compatibility."
+        else:
+             error_msg = f"Failed to import 'src' module in Blender environment. Error: {e.stderr[:500]}"
+        raise gr.Error(error_msg)
+    except Exception as e_test:
+         print(f"Unexpected error during test import: {e_test}")
+         raise gr.Error(f"Unexpected error during 'src' import test: {str(e_test)[:500]}")
+
+    # --- Execute Main UniRig Command ---
+    # Command structure: blender --background --python script.py -- script_arg1 script_arg2 ...
+    # Note: UniRig's run.py uses Hydra, which parses args differently. It doesn't need '--'.
+    # The arguments are passed directly after the script name.
+    cmd = [
+        blender_executable_to_use,
+        "--background",
+        "--python", unirig_run_script_path,
+        # "--" # REMOVED: Hydra parses args directly after script name
+    ] + script_args
+
+    print(f"\n--- Running UniRig Step: {step_name} ---")
+    print(f"Command: {' '.join(cmd)}") # Log the actual command
 
     try:
         # Execute Blender with the Python script.
-        # cwd=UNIRIG_REPO_DIR ensures script runs relative to repo root.
+        # cwd=UNIRIG_REPO_DIR ensures script runs relative to repo root and finds configs.
         result = subprocess.run(
             cmd,
             cwd=UNIRIG_REPO_DIR, # CWD is set to UniRig repo root
             capture_output=True,
             text=True,
             check=True, # Raises CalledProcessError on non-zero exit codes
-            env=process_env # Pass the modified environment
+            env=process_env, # Pass the modified environment
+            timeout=1800 # 30 minutes timeout, adjust as needed
         )
         print(f"{step_name} STDOUT:\n{result.stdout}")
         # Blender often prints info/warnings to stderr even on success
@@ -235,9 +273,14 @@ def run_unirig_command(python_script_path: str, script_args: List[str], step_nam
             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.")
+                raise gr.Error(f"Blender reported an error reading an input file during {step_name}. Check paths and file integrity.")
+            elif "Error:" in result.stderr: # Catch generic Blender errors
+                 print(f"WARNING: Potential Blender error message found in STDERR for {step_name} despite success exit code.")
             # Add other specific Blender error checks if needed
 
+    except subprocess.TimeoutExpired:
+        print(f"ERROR: {step_name} timed out after 30 minutes.")
+        raise gr.Error(f"Processing step '{step_name}' timed out. Please try with a simpler model or check logs.")
     except subprocess.CalledProcessError as e:
         print(f"ERROR during {step_name}: Subprocess failed!")
         print(f"Command: {' '.join(e.cmd)}")
@@ -245,20 +288,33 @@ def run_unirig_command(python_script_path: str, script_args: List[str], step_nam
         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 import errors within the subprocess stderr
-        if "ModuleNotFoundError: No module named 'src'" in e.stderr:
-             raise gr.Error(f"Error in UniRig '{step_name}': Script failed to import the 'src' module. Check PYTHONPATH and script CWD.")
-        elif "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.")
+        last_lines = "\n".join(error_summary[-15:]) if error_summary else "No stderr output."
+
+        specific_error = "Unknown error."
+        # Check common Python/PyTorch/Blender errors in stderr
+        if "ModuleNotFoundError: No module named 'bpy'" in e.stderr:
+            specific_error = "The 'bpy' module could not be imported by the script. Check installation in Blender's Python (via setup_blender.sh)."
+        elif "ModuleNotFoundError: No module named 'flash_attn'" in e.stderr:
+            specific_error = "The 'flash_attn' module is missing or failed to import. It might have failed during installation. Check setup_blender.sh logs."
+        elif "ModuleNotFoundError: No module named 'src'" in e.stderr:
+             specific_error = "UniRig script failed to import its own 'src' module. Check PYTHONPATH and CWD for the subprocess."
+        elif "ModuleNotFoundError" in e.stderr or "ImportError" in e.stderr:
+            specific_error = f"An import error occurred. Check library installations. Details: {last_lines}"
+        elif "OutOfMemoryError" in e.stderr or "CUDA out of memory" in e.stderr:
+            specific_error = "CUDA out of memory. Try a smaller model or a Space with more GPU RAM."
+        elif "hydra.errors.ConfigCompositionException" in e.stderr:
+             specific_error = f"Hydra configuration error. Check the arguments passed to run.py: {script_args}. Details: {last_lines}"
+        elif "Error: Cannot read file" in e.stderr:
+             specific_error = f"Blender could not read an input file. Check paths: {last_lines}"
+        elif "Error:" in e.stderr: # Generic Blender error
+             specific_error = f"Blender reported an error. Details: {last_lines}"
         else:
-             raise gr.Error(f"Error in UniRig '{step_name}'. Check logs. Last error lines:\n{last_lines}")
+            specific_error = f"Check logs. Last error lines:\n{last_lines}"
+        raise gr.Error(f"Error in UniRig '{step_name}'. {specific_error}")
 
     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"ERROR: Could not find Blender executable '{blender_executable_to_use}' or script '{unirig_run_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.")
 
@@ -275,16 +331,27 @@ def run_unirig_command(python_script_path: str, script_args: List[str], step_nam
 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.
+    Orchestrates calls to run_unirig_command for each step, executing run.py via Blender's Python.
     """
+    # Check readiness before proceeding
+    if not blender_executable_to_use:
+        # Use gr.Warning or gr.Info for non-blocking user feedback
+        gr.Info("System setup might still be in progress (Blender not found yet). Please wait a moment or check logs if this persists.")
+        return None # Return None for the output component
+    if not bpy_import_ok:
+         gr.Info("System setup might still be in progress (bpy import test failed). Please wait or check logs.")
+         return None
+
     try:
-        patch_asset_py() # Attempt patch
+        patch_asset_py() # If still needed
     except Exception as e:
-         print(f"Ignoring patch error: {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
+    if input_glb_file_obj is None:
+        gr.Warning("No input file provided.")
+        return None # Return None for the output component
+    input_glb_path = input_glb_file_obj # Gradio File(type="filepath") gives path string
     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.")
@@ -296,76 +363,92 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
     try:
         # --- Define File Paths ---
         base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
+        # Ensure paths are absolute for the subprocess
         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 ---
-        # *** Assuming run.py is the correct entry point. VERIFY THIS. ***
-        run_script_path = os.path.join(UNIRIG_REPO_DIR, "run.py") # Main script?
+        unirig_script_to_run = UNIRIG_RUN_PY # Path to UniRig/run.py
+
+        # --- Determine Device for UniRig ---
+        # Use CUDA if available in the Gradio env, otherwise CPU. Pass explicitly.
+        unirig_device_arg = "device=cpu"
+        if DEVICE.type == 'cuda':
+             # UniRig examples often use cuda:0 explicitly for Hydra config
+             unirig_device_arg = "device=cuda:0"
+        print(f"UniRig steps will attempt to use device argument: {unirig_device_arg}")
 
-        # --- Execute UniRig Steps ---
+        # --- Execute UniRig Steps using Hydra format ---
+        # UniRig uses Hydra (run.py --config-name=X with key=value ...)
 
         # Step 1: Skeleton Prediction
         print("\nStarting Step 1: Predicting Skeleton...")
-        # Arguments for the run.py script for skeleton task
-        # ** These are GUESSES - VERIFY from UniRig code / shell scripts **
         skeleton_args = [
-            # "--task", "generate_skeleton", # Example: if run.py needs a task identifier
-            "--input", abs_input_glb_path,
-            "--output", abs_skeleton_output_path
+            "--config-name", "skeleton_config", "with", # Hydra syntax
+            f"input={abs_input_glb_path}",
+            f"output={abs_skeleton_output_path}",
+            unirig_device_arg
+            # Add other optional args from generate_skeleton.sh if needed, e.g.:
+            # "seed=42", "verbose=True"
         ]
-        if not os.path.exists(run_script_path):
-            raise gr.Error(f"UniRig main script not found at: {run_script_path}")
-        run_unirig_command(run_script_path, skeleton_args, "Skeleton Prediction")
+        run_unirig_command(unirig_script_to_run, 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 **
         skin_args = [
-            # "--task", "generate_skin",
-            "--input", abs_skeleton_output_path,
-            "--source", abs_input_glb_path,
-            "--output", abs_skin_output_path
+            "--config-name", "skin_config", "with", # Hydra syntax
+            f"input={abs_skeleton_output_path}", # Uses the skeleton from previous step
+            f"output={abs_skin_output_path}",
+            unirig_device_arg
+            # Add other optional args from generate_skin.sh if needed, e.g.:
+            # "verbose=True"
         ]
-        run_unirig_command(run_script_path, skin_args, "Skinning Prediction")
+        run_unirig_command(unirig_script_to_run, 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
+        # Step 3: Merge Results
         print("\nStarting Step 3: Merging Results...")
-        # Arguments for the run.py script for merging task
-        # ** GUESSES - VERIFY **
         merge_args = [
-            # "--task", "merge",
-            "--source", abs_skin_output_path,
-            "--target", abs_input_glb_path,
-            "--output", abs_final_rigged_glb_path
+            "--config-name", "merge_config", "with", # Hydra syntax
+            f"source_path={abs_skin_output_path}", # Rigged data (skeleton + skin)
+            f"target_path={abs_input_glb_path}",   # Original mesh (for attributes, reference)
+            f"output_path={abs_final_rigged_glb_path}",
+            "mode=skin", # Explicitly set based on typical workflow
+            unirig_device_arg
+            # Add other optional args from merge.sh if needed
         ]
-        run_unirig_command(run_script_path, merge_args, "Merging")
+        run_unirig_command(unirig_script_to_run, merge_args, "Merging Results")
         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
+        # Use gr.update to explicitly update the Model3D component
+        return gr.update(value=abs_final_rigged_glb_path)
+
+    except gr.Error as e: # Catch Gradio errors specifically to re-raise
+        print(f"A Gradio Error occurred: {e}")
+        # Cleanup temp dir before re-raising
+        if os.path.exists(processing_temp_dir):
+            shutil.rmtree(processing_temp_dir)
+            print(f"Cleaned up temp dir: {processing_temp_dir}")
+        raise e # Re-raise the Gradio error
     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]}")
+        # Cleanup temp dir before raising Gradio error
+        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: {str(e)[:500]}. Check logs for details.")
+    # finally: # Cleanup is handled within except blocks to ensure it happens before raising
 
 
 # --- Gradio Interface Definition ---
@@ -376,36 +459,64 @@ theme = gr.themes.Soft(
     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):
+# Check critical components before building the interface
+startup_error_message = None
+if not blender_executable_to_use:
+     startup_error_message = (f"CRITICAL STARTUP ERROR: Blender executable could not be located or setup failed. Check logs.")
+elif 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}")
+elif not os.path.exists(UNIRIG_RUN_PY):
+      startup_error_message = (f"CRITICAL STARTUP ERROR: UniRig run.py not found at {UNIRIG_RUN_PY}.")
+# Add bpy_import_ok check? Maybe too strict for startup, handle in rig_glb_mesh_multistep
+# elif not bpy_import_ok:
+#      startup_error_message = (f"CRITICAL STARTUP ERROR: Failed initial 'bpy' import test using Blender.")
+
+
+if startup_error_message:
+    print(startup_error_message)
+    with gr.Blocks(theme=theme) as iface:
+        gr.Markdown(f"# Application Startup Error\n\n{startup_error_message}\n\nPlease check the Space logs for more details.")
 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'
-     )
+    # Build the normal interface if checks pass
+    with gr.Blocks(theme=theme) as iface:
+        gr.Markdown(
+            f"""
+            # UniRig Auto-Rigger (Blender {BLENDER_PYTHON_VERSION_DIR} / Python {BLENDER_PYTHON_VERSION})
+            Upload a 3D mesh in `.glb` format. This application uses UniRig via Blender's Python interface to predict skeleton and skinning weights.
+            * Running main app on Python `{sys.version.split()[0]}`, UniRig steps use Blender's Python `{BLENDER_PYTHON_VERSION}`.
+            * Utilizing device: **{DEVICE.type.upper()}** (via ZeroGPU if available).
+            * UniRig Source: [https://github.com/VAST-AI-Research/UniRig](https://github.com/VAST-AI-Research/UniRig)
+            """
+        )
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_model = gr.File(
+                    label="Upload .glb Mesh File",
+                    type="filepath", # Use filepath for direct path access
+                    file_types=[".glb"]
+                )
+                submit_button = gr.Button("Rig Model", variant="primary")
+            with gr.Column(scale=2):
+                output_model = gr.Model3D(
+                    label="Rigged 3D Model (.glb)",
+                    clear_color=[0.8, 0.8, 0.8, 1.0],
+                    # Add interaction settings if desired
+                    # camera_position=(0, 0, 5)
+                )
+
+        # Connect button click to the processing function
+        submit_button.click(
+            fn=rig_glb_mesh_multistep,
+            inputs=[input_model],
+            outputs=[output_model]
+        )
 
 # --- Launch the Application ---
 if __name__ == "__main__":
     if 'iface' in locals():
         print("Launching Gradio interface...")
-        iface.launch()
+        # Set share=True for public link if needed, debug=True for more logs locally
+        iface.launch(ssr_mode=False) # Disable SSR as it was experimental
     else:
         print("ERROR: Gradio interface not created due to startup errors.")