Update app.py

app.py

@@ -23,7 +23,7 @@ BLENDER_EXEC_SYMLINK = "/usr/local/bin/blender" # Fallback symlink
 
 SETUP_SCRIPT = os.path.join(os.path.dirname(__file__), "setup_blender.sh")
 
-# --- Initial Checks (adapted from your script) ---
+# --- Initial Checks ---
 print("--- Environment Checks ---")
 blender_executable_to_use = None # Use the main blender executable
 if os.path.exists(BLENDER_EXEC):
@@ -36,12 +36,13 @@ else:
     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)
+            # Run setup script if Blender not found
+            setup_result = subprocess.run(["bash", SETUP_SCRIPT], check=True, capture_output=True, text=True, timeout=600) # 10 min timeout for setup
             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
+            # Re-check for executable after running setup
             if os.path.exists(BLENDER_EXEC):
                 blender_executable_to_use = BLENDER_EXEC
                 print(f"Blender executable now found at: {BLENDER_EXEC}")
@@ -50,21 +51,30 @@ else:
                  print(f"Blender executable now found via symlink: {BLENDER_EXEC_SYMLINK}")
 
             if not blender_executable_to_use:
+                 # If still not found after setup, raise a clear error
                  raise RuntimeError(f"Setup script ran but Blender executable still not found at {BLENDER_EXEC} or {BLENDER_EXEC_SYMLINK}.")
 
+        except subprocess.TimeoutExpired:
+             print(f"ERROR: Setup script timed out: {SETUP_SCRIPT}")
+             raise gr.Error("Setup script timed out. The Space might be too slow or setup is stuck.")
         except subprocess.CalledProcessError as e:
             print(f"ERROR running setup script: {SETUP_SCRIPT}\nStderr: {e.stderr}")
+            # Raise a Gradio error to notify the user
             raise gr.Error(f"Failed to execute setup script. Check logs. Stderr: {e.stderr[-500:]}")
         except Exception as e:
+            # Catch any other exceptions during setup
             raise gr.Error(f"Unexpected error running setup script '{SETUP_SCRIPT}': {e}")
     else:
+        # If setup script itself is missing
         raise gr.Error(f"Blender executable not found and setup script missing: {SETUP_SCRIPT}")
 
 # Verify bpy import using the found Blender executable
 bpy_import_ok = False
 if blender_executable_to_use:
     try:
+        print("Testing bpy import via Blender...")
         test_script_content = "import bpy; print('bpy imported successfully')"
+        # Use --python-expr for the bpy test
         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
@@ -73,46 +83,55 @@ if blender_executable_to_use:
             print("Successfully imported 'bpy' using Blender executable.")
             bpy_import_ok = True
         else:
+            # Log warning but don't necessarily stop startup
             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:
+        # Log specific error if bpy import fails
         print(f"WARNING: Failed to import 'bpy' using Blender executable:\nSTDOUT:{e.stdout}\nSTDERR:{e.stderr}")
     except Exception as e:
+        # Catch any other exception during the test
         print(f"WARNING: Unexpected error during 'bpy' import test: {e}")
 else:
+     # This case should ideally not be reached if setup logic is correct
      print("WARNING: Cannot test bpy import as Blender executable was not found.")
 
 
-# Check for UniRig repository
+# Check for UniRig repository and run.py
+unirig_repo_ok = False
+unirig_run_py_ok = False
+UNIRIG_RUN_PY = os.path.join(UNIRIG_REPO_DIR, "run.py")
 if not os.path.isdir(UNIRIG_REPO_DIR):
+    # Critical error if UniRig repo 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}")
-    # Check for UniRig's run.py
-    UNIRIG_RUN_PY = os.path.join(UNIRIG_REPO_DIR, "run.py")
+    unirig_repo_ok = True
+    # Check specifically for run.py within the repo
     if not os.path.exists(UNIRIG_RUN_PY):
+        # Critical error if run.py is missing
         raise gr.Error(f"UniRig's run.py not found at {UNIRIG_RUN_PY}. Check UniRig clone.")
+    else:
+        unirig_run_py_ok = True
 
 
-# Check PyTorch and CUDA for Gradio environment
+# Check PyTorch and CUDA for Gradio environment (less critical for startup)
 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)}")
-        # 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}")
+        print(f"Could not get Gradio CUDA device details: {e}")
 else:
     print("Warning: Gradio environment CUDA not available.")
 
 print("--- End Environment Checks ---")
 
-# --- Helper Functions (Keep patch_asset_py if needed) ---
+# --- 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")
@@ -122,9 +141,8 @@ 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" # Requires `from typing import Any`
+        corrected_line = "meta: Union[Dict[str, Any], None]=None"
         typing_import = "from typing import Any"
 
         if corrected_line in content:
@@ -134,15 +152,11 @@ def patch_asset_py():
 
         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:
-                 # 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")
 
@@ -153,13 +167,13 @@ def patch_asset_py():
 
 # Decorator for ZeroGPU if needed
 @spaces.GPU
-def run_unirig_command(unirig_run_script_path: str, script_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 PYTHON script (.py) using the Blender executable
     in background mode (`blender --background --python script.py -- args`).
 
     Args:
-        unirig_run_script_path: Absolute path to the UniRig run.py script.
+        python_script_path: Absolute path to the Python 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.
     """
@@ -169,9 +183,9 @@ def run_unirig_command(unirig_run_script_path: str, script_args: List[str], step
     # --- Environment Setup for Subprocess ---
     process_env = os.environ.copy()
 
-    # PYTHONPATH: Only need UniRig's root directory. Python automatically checks CWD.
-    # Add UniRig root to the start of PYTHONPATH.
-    pythonpath_parts = [UNIRIG_REPO_DIR]
+    # Explicitly add UniRig root AND its 'src' directory to PYTHONPATH
+    unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, 'src')
+    pythonpath_parts = [UNIRIG_REPO_DIR, unirig_src_dir] # Add both
     existing_pythonpath = process_env.get('PYTHONPATH', '')
     if existing_pythonpath:
         pythonpath_parts.append(existing_pythonpath)
@@ -180,76 +194,24 @@ def run_unirig_command(unirig_run_script_path: str, script_args: List[str], step
 
     # 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
+    blender_python_lib_path = os.path.join(BLENDER_PYTHON_DIR, "lib")
     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', '')
     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.
+    # --- Execute Command ---
+    # Command structure: blender --background --python script.py -- args
     cmd = [
         blender_executable_to_use,
         "--background",
-        "--python", unirig_run_script_path,
-        # "--" # REMOVED: Hydra parses args directly after script name
+        "--python", python_script_path,
+        "--" # Separator between blender args and script args
     ] + script_args
 
     print(f"\n--- Running UniRig Step: {step_name} ---")
@@ -257,7 +219,7 @@ def run_unirig_command(unirig_run_script_path: str, script_args: List[str], step
 
     try:
         # Execute Blender with the Python script.
-        # cwd=UNIRIG_REPO_DIR ensures script runs relative to repo root and finds configs.
+        # CWD is crucial for relative imports and finding config files.
         result = subprocess.run(
             cmd,
             cwd=UNIRIG_REPO_DIR, # CWD is set to UniRig repo root
@@ -265,18 +227,15 @@ def run_unirig_command(unirig_run_script_path: str, script_args: List[str], step
             text=True,
             check=True, # Raises CalledProcessError on non-zero exit codes
             env=process_env, # Pass the modified environment
-            timeout=1800 # 30 minutes timeout, adjust as needed
+            timeout=1800 # 30 minutes timeout
         )
         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.")
-            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
+            # More robust check for errors in stderr even on success
+            stderr_lower = result.stderr.lower()
+            if "error" in stderr_lower or "failed" in stderr_lower or "traceback" in stderr_lower:
+                 print(f"WARNING: Potential error messages found in STDERR for {step_name} despite success exit code.")
 
     except subprocess.TimeoutExpired:
         print(f"ERROR: {step_name} timed out after 30 minutes.")
@@ -291,21 +250,21 @@ def run_unirig_command(unirig_run_script_path: str, script_args: List[str], step
         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:
+        # Refine error checking based on common patterns
+        if "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. See diagnostic info above."
+        elif "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: Cannot read file" in e.stderr: # General file read error
+             specific_error = f"Blender could not read an input/output file. Check paths. Details: {last_lines}"
         elif "Error:" in e.stderr: # Generic Blender error
              specific_error = f"Blender reported an error. Details: {last_lines}"
         else:
@@ -313,11 +272,8 @@ def run_unirig_command(unirig_run_script_path: str, script_args: List[str], step
         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 '{unirig_run_script_path}' for {step_name}.")
-        print(f"Attempted command: {' '.join(cmd)}")
+        print(f"ERROR: Could not find Blender executable '{blender_executable_to_use}' or script '{python_script_path}' for {step_name}.")
         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
@@ -331,30 +287,32 @@ def run_unirig_command(unirig_run_script_path: str, script_args: List[str], step
 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 run.py via Blender's Python.
     """
-    # Check readiness before proceeding
+    # Perform readiness checks at the start of the function call
     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.")
+        gr.Warning("System not ready: Blender executable not found. Please wait or check logs.")
+        return None
+    if not unirig_repo_ok or not unirig_run_py_ok:
+         gr.Warning("System not ready: UniRig repository or run.py script not found. Check setup.")
          return None
+    if not bpy_import_ok:
+         gr.Warning("System warning: Initial 'bpy' import test failed. Attempting to proceed, but errors may occur.")
 
     try:
-        patch_asset_py() # If still needed
+        patch_asset_py() # Apply patch if needed
     except Exception as e:
         print(f"Ignoring patch error: {e}")
 
     # --- Input Validation ---
     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
+        gr.Info("Please upload a .glb file first.")
+        return None
+    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.")
+    if not isinstance(input_glb_path, str) or not os.path.exists(input_glb_path):
+         raise gr.Error(f"Invalid input file path received: {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_")
@@ -363,34 +321,99 @@ 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")
 
-        unirig_script_to_run = UNIRIG_RUN_PY # Path to UniRig/run.py
+        unirig_script_to_run = UNIRIG_RUN_PY
+
+        # --- Run Blender Python Environment Diagnostic Test ---
+        print("\n--- Running Blender Python Environment Diagnostic Test ---")
+        diagnostic_script_content = f"""
+import sys
+import os
+import traceback
+
+print("--- Diagnostic Info from Blender Python ---")
+print(f"Python Executable: {{sys.executable}}")
+print(f"Current Working Directory (inside script): {{os.getcwd()}}") # Should be UNIRIG_REPO_DIR
+print("sys.path:")
+for p in sys.path:
+    print(f"  {{p}}")
+print("\\nPYTHONPATH Environment Variable (as seen by script):")
+print(os.environ.get('PYTHONPATH', 'PYTHONPATH not set or empty'))
+print("\\n--- Attempting Imports ---")
+try:
+    import bpy
+    print("SUCCESS: 'bpy' imported.")
+except ImportError as e:
+    print(f"FAILED to import 'bpy': {{e}}")
+    traceback.print_exc()
+except Exception as e:
+    print(f"FAILED to import 'bpy' with other error: {{e}}")
+    traceback.print_exc()
+
+try:
+    # Check if CWD is correct and src is importable
+    print("\\nChecking for 'src' in CWD (should be UniRig repo root):")
+    if os.path.isdir('src'):
+        print("  'src' directory FOUND in CWD.")
+        if os.path.isfile(os.path.join('src', '__init__.py')):
+            print("  'src/__init__.py' FOUND.")
+        else:
+            print("  WARNING: 'src/__init__.py' NOT FOUND. 'src' may not be treated as a package.")
+    else:
+        print("  'src' directory NOT FOUND in CWD.")
+
+    # Attempt the import that failed previously
+    print("\\nAttempting: from src.inference.download import download")
+    from src.inference.download import download
+    print("SUCCESS: 'from src.inference.download import download' worked.")
+except ImportError as e:
+    print(f"FAILED: 'from src.inference.download import download': {{e}}")
+    traceback.print_exc()
+except Exception as e:
+    print(f"FAILED: 'from src.inference.download import download' with other error: {{e}}")
+    traceback.print_exc()
+print("--- End Diagnostic Info ---")
+"""
+        diagnostic_script_path = os.path.join(processing_temp_dir, "env_diagnostic_test.py") # Use temp dir
+        with open(diagnostic_script_path, "w") as f:
+            f.write(diagnostic_script_content)
+
+        try:
+            # Run the diagnostic script using the *exact same* command structure
+            # Pass an empty list for script_args as the script takes none
+            run_unirig_command(diagnostic_script_path, [], "Blender Env Diagnostic")
+            print("--- Finished Blender Python Environment Diagnostic Test ---\n")
+        except Exception as e_diag:
+            # If the diagnostic fails, raise an error - no point continuing
+            print(f"ERROR during diagnostic test execution: {e_diag}")
+            # Clean up the script file before raising
+            if os.path.exists(diagnostic_script_path): os.remove(diagnostic_script_path)
+            raise gr.Error(f"Blender environment diagnostic failed. Cannot proceed. Check logs above for details. Error: {str(e_diag)[:500]}")
+        finally:
+             # Ensure cleanup even if run_unirig_command doesn't raise but has issues
+            if os.path.exists(diagnostic_script_path): os.remove(diagnostic_script_path)
+        # --- END DIAGNOSTIC STEP ---
+
 
         # --- 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 using Hydra format ---
-        # UniRig uses Hydra (run.py --config-name=X with key=value ...)
-
+        # --- Execute UniRig Steps ---
         # Step 1: Skeleton Prediction
         print("\nStarting Step 1: Predicting Skeleton...")
         skeleton_args = [
-            "--config-name", "skeleton_config", "with", # Hydra syntax
+            "--config-name=skeleton_config",
+            "with",
             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"
         ]
         run_unirig_command(unirig_script_to_run, skeleton_args, "Skeleton Prediction")
         if not os.path.exists(abs_skeleton_output_path):
@@ -400,12 +423,11 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
         # Step 2: Skinning Weight Prediction
         print("\nStarting Step 2: Predicting Skinning Weights...")
         skin_args = [
-            "--config-name", "skin_config", "with", # Hydra syntax
-            f"input={abs_skeleton_output_path}", # Uses the skeleton from previous step
+            "--config-name=skin_config",
+            "with",
+            f"input={abs_skeleton_output_path}",
             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(unirig_script_to_run, skin_args, "Skinning Prediction")
         if not os.path.exists(abs_skin_output_path):
@@ -415,13 +437,13 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
         # Step 3: Merge Results
         print("\nStarting Step 3: Merging Results...")
         merge_args = [
-            "--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)
+            "--config-name=merge_config",
+            "with",
+            f"source_path={abs_skin_output_path}",
+            f"target_path={abs_input_glb_path}",
             f"output_path={abs_final_rigged_glb_path}",
-            "mode=skin", # Explicitly set based on typical workflow
+            "mode=skin",
             unirig_device_arg
-            # Add other optional args from merge.sh if needed
         ]
         run_unirig_command(unirig_script_to_run, merge_args, "Merging Results")
         if not os.path.exists(abs_final_rigged_glb_path):
@@ -430,25 +452,25 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
 
         # --- Return Result ---
         print(f"Successfully generated rigged model: {abs_final_rigged_glb_path}")
-        # Use gr.update to explicitly update the Model3D component
-        return gr.update(value=abs_final_rigged_glb_path)
+        return gr.update(value=abs_final_rigged_glb_path) # Update Gradio output
 
-    except gr.Error as e: # Catch Gradio errors specifically to re-raise
+    except gr.Error as e:
         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
+        if os.path.exists(processing_temp_dir): shutil.rmtree(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()
-        # 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}")
+        if os.path.exists(processing_temp_dir): shutil.rmtree(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
+    finally:
+        # General cleanup for the temp dir if it still exists (e.g., if no exception occurred)
+        if os.path.exists(processing_temp_dir):
+             try:
+                 shutil.rmtree(processing_temp_dir)
+                 print(f"Cleaned up temp dir: {processing_temp_dir}")
+             except Exception as cleanup_e:
+                 print(f"Error cleaning up temp dir {processing_temp_dir}: {cleanup_e}")
 
 
 # --- Gradio Interface Definition ---
@@ -459,40 +481,38 @@ theme = gr.themes.Soft(
     font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
 )
 
-# Check critical components before building the interface
+# Perform critical checks before defining 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):
+elif not unirig_repo_ok:
      startup_error_message = (f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}.")
-elif not os.path.exists(UNIRIG_RUN_PY):
+elif not unirig_run_py_ok:
       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:
+    # If critical error, display error message instead of interface
     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 checks pass
+    # Build the normal interface if essential 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)
-            """
-        )
+             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
+                    type="filepath",
                     file_types=[".glb"]
                 )
                 submit_button = gr.Button("Rig Model", variant="primary")
@@ -500,8 +520,6 @@ else:
                 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
@@ -513,10 +531,12 @@ else:
 
 # --- Launch the Application ---
 if __name__ == "__main__":
+    # Check if the interface was successfully created
     if 'iface' in locals():
         print("Launching Gradio interface...")
-        # 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
+        # Ensure share=False for security unless public link is needed
+        iface.launch(share=False, ssr_mode=False) # Disable SSR
     else:
-        print("ERROR: Gradio interface not created due to startup errors.")
+        # This indicates a startup error occurred before interface creation
+        print("ERROR: Gradio interface could not be created due to startup errors. Check logs above.")