Update app.py

app.py

@@ -9,82 +9,96 @@ import spaces
 from typing import Any, Dict, Union, List
 
 # --- Configuration ---
-# Path to the cloned UniRig repository directory within the Space
-# Ensure this path is correct relative to app.py
-UNIRIG_REPO_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "UniRig")) # Get absolute path
-
-# Absolute path to the Blender installation provided in the Space environment
+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" # Corresponds to Blender 4.2.x
-BLENDER_PYTHON_VERSION = "python3.11" # Blender 4.2 uses Python 3.11
+BLENDER_PYTHON_VERSION_DIR = "4.2"
+BLENDER_PYTHON_VERSION = "python3.11"
 
-# Construct paths dynamically based on the above
+# 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") # Directory containing python executable
-BLENDER_PYTHON_EXEC = os.path.join(BLENDER_PYTHON_BIN_DIR, BLENDER_PYTHON_VERSION) # Full path to executable
+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
 
-# Path to the setup script (executed only if Blender isn't found initially)
 SETUP_SCRIPT = os.path.join(os.path.dirname(__file__), "setup_blender.sh")
 
 # --- Initial Checks ---
 print("--- Environment Checks ---")
 
-# Check if Blender Python executable exists (needed for environment setup)
-# Run setup script if Blender isn't found (assuming setup script handles installation)
-if not os.path.exists(BLENDER_PYTHON_EXEC):
-    print(f"Blender Python executable not found at {BLENDER_PYTHON_EXEC}. Running setup script...")
+# 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:
-            # Run the setup script using bash
             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 if the executable exists after running the script
-            if not os.path.exists(BLENDER_PYTHON_EXEC):
-                 raise RuntimeError(f"Setup script ran but Blender Python still not found at {BLENDER_PYTHON_EXEC}.")
+            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}")
-             print(f"Return Code: {e.returncode}")
-             print(f"Stdout: {e.stdout}")
-             print(f"Stderr: {e.stderr}")
-             # Raise a more informative error for Gradio if setup fails
+             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:
-        # If setup script is missing, raise a Gradio error
-        raise gr.Error(f"Blender Python not found and setup script missing: {SETUP_SCRIPT}")
-else:
-    print(f"Blender Python executable found: {BLENDER_PYTHON_EXEC}")
+        raise gr.Error(f"Blender executable not found and setup script missing: {SETUP_SCRIPT}")
 
-# Verify Blender Python site-packages path and bpy module presence
-bpy_found = False
+# 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}")
-    # Check 1: 'bpy' directory with __init__.py
-    bpy_module_dir = os.path.join(BLENDER_PYTHON_SITE_PACKAGES, "bpy")
-    if os.path.isdir(bpy_module_dir) and os.path.exists(os.path.join(bpy_module_dir, "__init__.py")):
-        print("Found 'bpy' module directory in site-packages.")
-        bpy_found = True
-    # Check 2: 'bpy.so' file (less common structure)
-    elif os.path.exists(os.path.join(BLENDER_PYTHON_SITE_PACKAGES, "bpy.so")):
-         print("Found 'bpy.so' in site-packages.")
-         bpy_found = True
-
-    if not bpy_found:
-        print("WARNING: Blender Python 'bpy' module indicator not found in site-packages. Imports might fail.")
+    # 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):
-    print(f"ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}.")
-    # Raise Gradio error if critical component is missing
     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}")
@@ -99,8 +113,7 @@ if DEVICE.type == 'cuda':
     except Exception as e:
         print(f"Could not get CUDA device details: {e}")
 else:
-    print("Warning: Gradio environment CUDA not available or not detected by PyTorch.")
-    print("UniRig subprocess will attempt to use GPU via Blender's Python environment (invoked by .sh scripts).")
+    print("Warning: Gradio environment CUDA not available.")
 
 print("--- End Environment Checks ---")
 
@@ -108,107 +121,101 @@ print("--- End Environment Checks ---")
 
 def patch_asset_py():
     """Temporary patch to fix type hinting error in UniRig's asset.py"""
-    # This patch might still be needed if the .sh scripts call Python code that uses asset.py
     asset_py_path = os.path.join(UNIRIG_REPO_DIR, "src", "data", "asset.py")
     try:
-        # Check if file exists before trying to open
         if not os.path.exists(asset_py_path):
             print(f"Warning: asset.py not found at {asset_py_path}, skipping patch.")
-            return # Don't raise error, maybe it's not needed
-
-        with open(asset_py_path, "r") as f:
-            content = f.read()
+            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
+            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."); 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)
+            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:
-        # Log error but don't necessarily stop the app, maybe patch isn't critical
         print(f"ERROR: Failed to patch asset.py: {e}. Proceeding cautiously.")
-        # raise gr.Error(f"Failed to apply necessary patch to UniRig code: {e}") # Optional: make it fatal
 
 @spaces.GPU
-def run_unirig_command(script_path: str, args: List[str], step_name: str):
+def run_unirig_command(python_script_path: str, script_args: List[str], step_name: str):
     """
-    Runs a specific UniRig SHELL script (.sh) using bash in a subprocess,
-    ensuring the correct environment (PYTHONPATH, etc.) is set within Blender's context.
+    Runs a specific UniRig PYTHON script (.py) using the Blender executable
+    in background mode (`blender --background --python script.py -- args`).
 
     Args:
-        script_path: Absolute path to the .sh script to execute.
-        args: A list of command-line arguments for the script.
+        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.
     """
-    # Use xvfb-run to enable headless execution of Blender
-    cmd = ["xvfb-run", "-a", "bash", script_path] + args
+    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)}")
+    print(f"Command: {' '.join(cmd)}") # Note: Simple join for logging
 
-    # Prepare the environment for the subprocess
+    # 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")
-
-    # Set PYTHONPATH: Blender’s site-packages + UniRig source
     pythonpath_parts = [
-        BLENDER_PYTHON_SITE_PACKAGES,
+        # 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']}")
 
-    # Set LD_LIBRARY_PATH: Include Blender’s Python library directory
+    # 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']}")
 
-    # Set PATH to prioritize Blender's Python bin directory
-    existing_path = process_env.get('PATH', '')
-    process_env["PATH"] = f"{BLENDER_PYTHON_BIN_DIR}{os.pathsep}{existing_path}"
-    print(f"Subprocess PATH: {process_env['PATH']}")
-
-    # Debug: Check which Python is used
-    python_bin = subprocess.run(["which", "python"], env=process_env, capture_output=True, text=True)
-    print(f"Python binary used: {python_bin.stdout.strip()}")
-
-    # Debug: Check Blender version
-    blender_check = subprocess.run([os.path.join(BLENDER_INSTALL_DIR, "blender"), "--version"], capture_output=True, text=True)
-    print(f"Blender version: {blender_check.stdout.strip()}")
 
     try:
-        # Execute the shell script with xvfb-run
+        # 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,
-            env=process_env
+            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 (Warnings/Info):\n{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!")
@@ -218,16 +225,21 @@ def run_unirig_command(script_path: str, args: List[str], step_name: str):
         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}': Script failed to import Blender's 'bpy' module. Check environment setup.")
+             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. Check environment and PyTorch installation.")
+             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}")
+             raise gr.Error(f"Error in UniRig '{step_name}'. Check logs. Last error lines:\n{last_lines}")
+
 
     except FileNotFoundError:
-        print(f"ERROR: Could not find executable 'bash' or script '{script_path}' for {step_name}.")
-        raise gr.Error(f"Setup error for UniRig '{step_name}'. 'bash' or script '{os.path.basename(script_path)}' not found.")
+        # 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}")
@@ -242,52 +254,19 @@ def run_unirig_command(script_path: str, args: List[str], step_name: str):
 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 .sh scripts.
+    Orchestrates calls to run_unirig_command for each step, executing .py scripts via Blender.
     """
     try:
-        patch_asset_py() # Attempt patch, might still be relevant
-    except gr.Error as e:
-         print(f"Ignoring patch error: {e}") # Or just log it
+        patch_asset_py() # Attempt patch
     except Exception as e:
-         print(f"Ignoring unexpected 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. Please upload a .glb mesh.")
-
+    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.")
-
-    # Test PyTorch import in Blender’s Python
-    try:
-        result = subprocess.run(
-            [BLENDER_PYTHON_EXEC, "-c", "import torch; print(torch.__version__)"],
-            capture_output=True,
-            text=True,
-            check=True
-        )
-        print(f"PyTorch version in Blender's Python: {result.stdout.strip()}")
-    except subprocess.CalledProcessError as e:
-        print(f"Failed to import torch in Blender's Python: {e.stderr}")
-        raise gr.Error("PyTorch import failed in Blender's Python environment. Check installation.")
-
-    # Test Blender execution
-    try:
-        blender_test = subprocess.run(
-            [os.path.join(BLENDER_INSTALL_DIR, "blender"), "--version"],
-            capture_output=True,
-            text=True,
-            check=True
-        )
-        print(f"Blender version: {blender_test.stdout.strip()}")
-    except subprocess.CalledProcessError as e:
-        print(f"Failed to run Blender: {e.stderr}")
-        raise gr.Error("Blender is not accessible. Check installation.")
+    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_")
@@ -296,59 +275,70 @@ 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 passed to scripts are absolute
         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 SHELL Scripts ---
-        skeleton_script_path = os.path.join(UNIRIG_REPO_DIR, "launch/inference/generate_skeleton.sh")
-        skin_script_path = os.path.join(UNIRIG_REPO_DIR, "launch/inference/generate_skin.sh")
-        merge_script_path = os.path.join(UNIRIG_REPO_DIR, "launch/inference/merge.sh")
+        # --- 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(skeleton_script_path):
-            raise gr.Error(f"Skeleton script not found at: {skeleton_script_path}")
-        run_unirig_command(skeleton_script_path, skeleton_args, "Skeleton Prediction")
+        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):
-            # Check if the error wasn't already raised by run_unirig_command
             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
+            "--source", abs_input_glb_path,    # Source mesh (if needed by skin script)
             "--output", abs_skin_output_path
         ]
-        if not os.path.exists(skin_script_path):
-            raise gr.Error(f"Skinning script not found at: {skin_script_path}")
-        run_unirig_command(skin_script_path, skin_args, "Skinning Prediction")
+        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 = [
-            # Determine which source to use based on what exists or a user choice?
-            # Assuming skin output is the desired source if it exists.
+            # e.g., "--task", "merge",
             "--source", abs_skin_output_path,
             "--target", abs_input_glb_path,
             "--output", abs_final_rigged_glb_path
         ]
-        if not os.path.exists(merge_script_path):
-            raise gr.Error(f"Merging script not found at: {merge_script_path}")
-        run_unirig_command(merge_script_path, merge_args, "Merging")
+        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.")
@@ -359,18 +349,12 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
 
     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 temporary directory: {processing_temp_dir}")
+        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 temporary directory: {processing_temp_dir}")
+        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]}")
 
 
@@ -384,47 +368,34 @@ theme = gr.themes.Soft(
 
 # 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}. "
-         "The application cannot start. Please ensure the repository is cloned correctly."
-     )
+     startup_error_message = (f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}.")
      print(startup_error_message)
-     # Display an error message if Gradio tries to build the UI without UniRig
-     with gr.Blocks(theme=theme) as iface:
-          gr.Markdown(f"# Application Error\n\n{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", # Provides the path to the uploaded file
-             file_types=[".glb"] # Restrict file types
-         ),
-         outputs=gr.Model3D(
-             label="Rigged 3D Model (.glb)",
-             clear_color=[0.8, 0.8, 0.8, 1.0] # Background color for the viewer
-         ),
+         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 (invoked through shell scripts) to automatically generate a skeleton and skinning weights.\n"
-             f"* Running main app on Python {sys.version.split()[0]}, UniRig steps use Blender's Python {BLENDER_PYTHON_VERSION} via .sh scripts.\n"
+             "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, # Disable caching if results depend heavily on GPU state or are large
-         theme=theme,
-         allow_flagging='never' # Disable flagging unless needed
+         cache_examples=False, theme=theme, allow_flagging='never'
      )
 
 # --- Launch the Application ---
 if __name__ == "__main__":
-    # Ensure the interface object exists before launching
     if 'iface' in locals():
         print("Launching Gradio interface...")
-        # Consider adding share=True for public link if needed, or server_name="0.0.0.0"
         iface.launch()
     else:
-        # This case should only happen if the UniRig repo check failed above
         print("ERROR: Gradio interface not created due to startup errors.")