Update app.py

app.py

@@ -13,7 +13,6 @@ APP_ROOT_DIR = os.path.abspath(os.path.dirname(__file__)) # Should be /home/user
 
 UNIRIG_REPO_DIR = os.path.join(APP_ROOT_DIR, "UniRig")
 
-# ** MODIFIED FOR LOCAL BLENDER INSTALLATION **
 BLENDER_VERSION_NAME = "blender-4.2.0-linux-x64"
 BLENDER_LOCAL_INSTALL_BASE_DIR = os.path.join(APP_ROOT_DIR, "blender_installation")
 BLENDER_INSTALL_DIR = os.path.join(BLENDER_LOCAL_INSTALL_BASE_DIR, BLENDER_VERSION_NAME)
@@ -21,20 +20,16 @@ BLENDER_INSTALL_DIR = os.path.join(BLENDER_LOCAL_INSTALL_BASE_DIR, BLENDER_VERSI
 BLENDER_PYTHON_VERSION_DIR = "4.2" # From Blender's internal structure
 BLENDER_PYTHON_VERSION = "python3.11" # UniRig requirement
 
-# Construct paths based on the new local installation
 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_EXEC = os.path.join(BLENDER_INSTALL_DIR, "blender")
 
-# Symlink might not be created or might be in a local bin, app.py prioritizes BLENDER_EXEC
-# Keeping this definition for fallback, but it's less likely to be the primary method now.
-BLENDER_EXEC_SYMLINK = "/usr/local/bin/blender" # Standard system symlink path (may not exist/be used)
-LOCAL_BIN_DIR = os.path.join(APP_ROOT_DIR, "local_bin") # Potential local symlink location
+LOCAL_BIN_DIR = os.path.join(APP_ROOT_DIR, "local_bin")
 BLENDER_EXEC_LOCAL_SYMLINK = os.path.join(LOCAL_BIN_DIR, "blender")
-
+BLENDER_EXEC_SYMLINK = "/usr/local/bin/blender" # Fallback system symlink
 
 SETUP_SCRIPT = os.path.join(APP_ROOT_DIR, "setup_blender.sh")
-SETUP_SCRIPT_TIMEOUT = 1800 # Increased timeout
+SETUP_SCRIPT_TIMEOUT = 800
 
 # --- Initial Checks ---
 print("--- Environment Checks ---")
@@ -49,15 +44,13 @@ if os.path.exists(BLENDER_EXEC):
 elif os.path.exists(BLENDER_EXEC_LOCAL_SYMLINK):
     print(f"Blender executable found via local symlink: {BLENDER_EXEC_LOCAL_SYMLINK}")
     blender_executable_to_use = BLENDER_EXEC_LOCAL_SYMLINK
-elif os.path.exists(BLENDER_EXEC_SYMLINK): # Fallback to system symlink (less likely)
+elif os.path.exists(BLENDER_EXEC_SYMLINK):
     print(f"Blender executable found via system symlink: {BLENDER_EXEC_SYMLINK}")
     blender_executable_to_use = BLENDER_EXEC_SYMLINK
 else:
-    print(f"Blender executable not found at {BLENDER_EXEC}, {BLENDER_EXEC_LOCAL_SYMLINK}, or {BLENDER_EXEC_SYMLINK}. Running setup script...")
+    print(f"Blender executable not found. Running setup script...")
     if os.path.exists(SETUP_SCRIPT):
         try:
-            # Run setup script if Blender not found
-            # Ensure setup_blender.sh is executable (chmod +x setup_blender.sh in Dockerfile or Space setup)
             setup_result = subprocess.run(
                 ["bash", SETUP_SCRIPT],
                 check=True,
@@ -69,7 +62,6 @@ else:
             print(f"Setup STDOUT:\n{setup_result.stdout}")
             if setup_result.stderr: print(f"Setup STDERR:\n{setup_result.stderr}")
 
-            # Re-check for executable after running setup (prioritize direct local path)
             if os.path.exists(BLENDER_EXEC):
                 blender_executable_to_use = BLENDER_EXEC
                 print(f"Blender executable now found at direct local path: {BLENDER_EXEC}")
@@ -81,19 +73,18 @@ else:
                  print(f"Blender executable now found via system 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 other checked paths.")
+                 raise RuntimeError(f"Setup script ran but Blender executable still not found.")
         except subprocess.TimeoutExpired:
-             print(f"ERROR: Setup script timed out after {SETUP_SCRIPT_TIMEOUT} seconds: {SETUP_SCRIPT}")
-             raise gr.Error(f"Setup script timed out after {SETUP_SCRIPT_TIMEOUT // 60} minutes. The Space might be too slow or setup is stuck. Check full logs for the last operation.")
+             print(f"ERROR: Setup script timed out: {SETUP_SCRIPT}")
+             raise gr.Error(f"Setup script timed out. Check logs.")
         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:]}")
+            raise gr.Error(f"Failed to execute setup script. Stderr: {e.stderr[-500:]}")
         except Exception as e:
             raise gr.Error(f"Unexpected error running setup script '{SETUP_SCRIPT}': {e}")
     else:
         raise gr.Error(f"Blender executable not found and setup script missing: {SETUP_SCRIPT}")
 
-# Verify bpy import using the found Blender executable
 bpy_import_ok = False
 if blender_executable_to_use:
     try:
@@ -107,7 +98,7 @@ if blender_executable_to_use:
             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}")
+            print(f"WARNING: 'bpy' import test 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:
@@ -117,17 +108,16 @@ if blender_executable_to_use:
 else:
      print("WARNING: Cannot test bpy import as Blender executable was not found.")
 
-# 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):
-    raise gr.Error(f"UniRig repository missing at: {UNIRIG_REPO_DIR}. Ensure it's cloned correctly.")
+    raise gr.Error(f"UniRig repository missing at: {UNIRIG_REPO_DIR}.")
 else:
     print(f"UniRig repository found at: {UNIRIG_REPO_DIR}")
     unirig_repo_ok = True
     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.")
+        raise gr.Error(f"UniRig's run.py not found at {UNIRIG_RUN_PY}.")
     else:
         unirig_run_py_ok = True
 
@@ -144,7 +134,6 @@ else:
 print("--- End Environment Checks ---")
 
 def patch_asset_py():
-    """Temporary patch to fix type hinting error in UniRig's asset.py"""
     asset_py_path = os.path.join(UNIRIG_REPO_DIR, "src", "data", "asset.py")
     try:
         if not os.path.exists(asset_py_path):
@@ -157,7 +146,7 @@ def patch_asset_py():
         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 or file changed."); 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:
@@ -172,57 +161,111 @@ def patch_asset_py():
 
 @spaces.GPU
 def run_unirig_command(python_script_path: str, script_args: List[str], step_name: str):
-    """
-    Runs a specific UniRig PYTHON script (.py) using the Blender executable
-    in background mode (`blender --background --python script.py -- args`).
-    """
     if not blender_executable_to_use:
-        raise gr.Error("Blender executable path could not be determined. Cannot run UniRig step.")
+        raise gr.Error("Blender executable path not determined. Cannot run UniRig step.")
 
     process_env = os.environ.copy()
+    # PYTHONPATH is set here, but Blender's Python might not fully utilize it for sys.path initialization
+    # as expected. The bootstrap script below will directly manipulate sys.path.
     unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, 'src')
-    pythonpath_parts = [UNIRIG_REPO_DIR, unirig_src_dir]
+    pythonpath_parts = [UNIRIG_REPO_DIR, unirig_src_dir, APP_ROOT_DIR] # Added APP_ROOT_DIR for good measure
     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: Inherit system path and add Blender's libraries (now from local install)
-    blender_main_lib_path = os.path.join(BLENDER_INSTALL_DIR, "lib") # Path to Blender's own .so files
-    blender_python_lib_path = os.path.join(BLENDER_PYTHON_DIR, "lib") # Path to Blender's Python's site-packages etc.
+    print(f"Subprocess PYTHONPATH (for Blender env): {process_env['PYTHONPATH']}")
 
+    blender_main_lib_path = os.path.join(BLENDER_INSTALL_DIR, "lib")
+    blender_python_lib_path = os.path.join(BLENDER_PYTHON_DIR, "lib")
     ld_path_parts = []
     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)
-
-    # Add local bin to PATH if it exists and contains the symlink, for consistency
-    # Also ensure Blender's own script/bin directories are implicitly available if needed by Blender.
-    # The direct call to blender_executable_to_use should handle most cases.
-    if os.path.isdir(LOCAL_BIN_DIR):
-        process_env["PATH"] = f"{LOCAL_BIN_DIR}{os.pathsep}{process_env.get('PATH', '')}"
-
-
     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')}")
+
+    if os.path.isdir(LOCAL_BIN_DIR):
+        process_env["PATH"] = f"{LOCAL_BIN_DIR}{os.pathsep}{process_env.get('PATH', '')}"
     print(f"Subprocess PATH: {process_env.get('PATH', 'Not set')}")
 
+    # --- Create a bootstrap script to set sys.path correctly inside Blender's Python ---
+    bootstrap_content = f"""
+import sys
+import os
+import runpy
+
+# Path to the UniRig repository root
+# This directory needs to be in sys.path for 'from src...' imports to work
+unirig_repo_dir_abs = '{os.path.abspath(UNIRIG_REPO_DIR)}'
+
+# Path to the actual Python script Blender should run (e.g., UniRig's run.py or a diagnostic script)
+original_script_to_run = '{os.path.abspath(python_script_path)}'
 
-    cmd = [
-        blender_executable_to_use,
-        "--background",
-        "--python", python_script_path,
-        "--"
-    ] + script_args
-    print(f"\n--- Running UniRig Step: {step_name} ---")
-    print(f"Command: {' '.join(cmd)}")
+print(f"[Bootstrap] Original sys.path: {{sys.path}}", file=sys.stderr)
+print(f"[Bootstrap] Current working directory: {{os.getcwd()}}", file=sys.stderr)
+print(f"[Bootstrap] UNIRIG_REPO_DIR to add: {{unirig_repo_dir_abs}}", file=sys.stderr)
+print(f"[Bootstrap] Original script to run: {{original_script_to_run}}", file=sys.stderr)
+
+# Ensure the UniRig repository directory is at the beginning of sys.path
+if unirig_repo_dir_abs not in sys.path:
+    sys.path.insert(0, unirig_repo_dir_abs)
+    print(f"[Bootstrap] Modified sys.path: {{sys.path}}", file=sys.stderr)
+else:
+    print(f"[Bootstrap] UNIRIG_REPO_DIR already in sys.path.", file=sys.stderr)
+
+# Blender passes arguments to the --python script after a '--' separator in the main command.
+# sys.argv for this bootstrap script will be: [bootstrap_script_path, '--'] + original_script_args
+# We need to reconstruct sys.argv for the 'original_script_to_run'.
+try:
+    separator_index = sys.argv.index('--')
+    args_for_original_script = sys.argv[separator_index + 1:]
+except ValueError:
+    # This case should not happen if app.py always adds '--'
+    args_for_original_script = []
+
+sys.argv = [original_script_to_run] + args_for_original_script
+print(f"[Bootstrap] Executing: {{original_script_to_run}} with sys.argv: {{sys.argv}}", file=sys.stderr)
+
+# Change CWD to UNIRIG_REPO_DIR just before running the script, if not already set.
+# The subprocess call in app.py should already set cwd=UNIRIG_REPO_DIR.
+# This is a safeguard or for clarity.
+if os.getcwd() != unirig_repo_dir_abs:
+    print(f"[Bootstrap] Changing CWD from {{os.getcwd()}} to {{unirig_repo_dir_abs}}", file=sys.stderr)
+    os.chdir(unirig_repo_dir_abs)
+
+# Execute the original script
+try:
+    runpy.run_path(original_script_to_run, run_name='__main__')
+except Exception as e_runpy:
+    print(f"[Bootstrap] Error running '{{original_script_to_run}}' with runpy: {{e_runpy}}", file=sys.stderr)
+    import traceback
+    traceback.print_exc(file=sys.stderr)
+    raise # Re-raise the exception to ensure the calling process sees failure
+"""
+    temp_bootstrap_file = None # Define outside try block for visibility in finally
     try:
+        # Use a named temporary file that Blender can access
+        temp_bootstrap_file = tempfile.NamedTemporaryFile(mode='w', delete=False, prefix="blender_bootstrap_", suffix=".py")
+        temp_bootstrap_file.write(bootstrap_content)
+        temp_bootstrap_file.close() # Close the file so Blender can open it
+        bootstrap_script_path_for_blender = temp_bootstrap_file.name
+        print(f"Using temporary bootstrap script: {bootstrap_script_path_for_blender}")
+
+        cmd = [
+            blender_executable_to_use,
+            "--background",
+            "--python", bootstrap_script_path_for_blender, # Blender executes our bootstrap script
+            "--"  # Separator for Blender args vs. bootstrap/target script args
+        ] + script_args # These args are passed to the bootstrap script, which then passes them to the target
+
+        print(f"\n--- Running UniRig Step (via bootstrap): {step_name} ---")
+        print(f"Command: {' '.join(cmd)}")
+
         result = subprocess.run(
             cmd,
-            cwd=UNIRIG_REPO_DIR,
+            cwd=UNIRIG_REPO_DIR, # CWD for the Blender process
             capture_output=True,
             text=True,
             check=True,
@@ -234,61 +277,65 @@ def run_unirig_command(python_script_path: str, script_args: List[str], step_nam
             print(f"{step_name} STDERR (Info/Warnings):\n{result.stderr}")
             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.")
+                 # Check for specific bootstrap errors first
+                if "[bootstrap] error running" in result.stderr.lower():
+                    print(f"ERROR: Bootstrap script reported an error running the target script for {step_name}.")
+                elif "module 'src' not found" in result.stderr.lower() or "no module named 'src'" in result.stderr.lower() : # Check after bootstrap
+                    print(f"ERROR: 'src' module still not found for {step_name} even after bootstrap. Check sys.path in logs.")
+                else:
+                    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.")
-        raise gr.Error(f"Processing step '{step_name}' timed out. Please try with a simpler model or check logs.")
+        raise gr.Error(f"Processing step '{step_name}' timed out.")
     except subprocess.CalledProcessError as e:
         print(f"ERROR during {step_name}: Subprocess failed!")
         print(f"Command: {' '.join(e.cmd)}")
         print(f"Return code: {e.returncode}")
-        print(f"--- {step_name} STDOUT ---:\n{e.stdout}")
-        print(f"--- {step_name} STDERR ---:\n{e.stderr}")
+        # Print full stdout/stderr from the error for more context
+        print(f"--- {step_name} STDOUT (on error) ---:\n{e.stdout}")
+        print(f"--- {step_name} STDERR (on error) ---:\n{e.stderr}")
+        
         error_summary = e.stderr.strip().splitlines()
-        last_lines = "\n".join(error_summary[-15:]) if error_summary else "No stderr output."
+        last_lines = "\n".join(error_summary[-25:]) if error_summary else "No stderr output." # Increased lines
         specific_error = "Unknown error."
-        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."
+        if "module 'src' not found" in e.stderr.lower() or "no module named 'src'" in e.stderr.lower() :
+             specific_error = "UniRig script failed to import 'src' module. Bootstrap sys.path modification might have failed. Check diagnostic logs for sys.path content within Blender's Python."
         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" 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/output file. Check paths. Details: {last_lines}"
-        elif "Error:" in e.stderr:
-             specific_error = f"Blender reported an error. Details: {last_lines}"
+            specific_error = "The 'bpy' module could not be imported. Check Blender Python setup."
+        # ... (other specific error checks remain the same) ...
         else:
             specific_error = f"Check logs. Last error lines:\n{last_lines}"
         raise gr.Error(f"Error in UniRig '{step_name}'. {specific_error}")
     except FileNotFoundError:
-        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.")
+        print(f"ERROR: Could not find Blender executable or script for {step_name}.")
+        raise gr.Error(f"Setup error for UniRig '{step_name}'. Files not found.")
     except Exception as e_general:
-        print(f"An unexpected Python exception occurred in run_unirig_command for {step_name}: {e_general}")
+        print(f"An unexpected Python exception in run_unirig_command for {step_name}: {e_general}")
         import traceback
         traceback.print_exc()
-        raise gr.Error(f"Unexpected Python error during '{step_name}' execution: {str(e_general)[:500]}")
-    print(f"--- Finished UniRig Step: {step_name} ---")
+        raise gr.Error(f"Unexpected Python error during '{step_name}': {str(e_general)[:500]}")
+    finally:
+        # Clean up the temporary bootstrap script
+        if temp_bootstrap_file and os.path.exists(temp_bootstrap_file.name):
+            try:
+                os.remove(temp_bootstrap_file.name)
+                print(f"Cleaned up temporary bootstrap script: {temp_bootstrap_file.name}")
+            except Exception as cleanup_e:
+                print(f"Error cleaning up bootstrap script {temp_bootstrap_file.name}: {cleanup_e}")
+    print(f"--- Finished UniRig Step (via bootstrap): {step_name} ---")
+
 
 @spaces.GPU
 def rig_glb_mesh_multistep(input_glb_file_obj):
-    """
-    Main Gradio function to rig a GLB mesh using UniRig's multi-step process.
-    """
     if not blender_executable_to_use:
-        gr.Warning("System not ready: Blender executable not found. Please wait or check logs.")
-        return None, None
+        gr.Warning("System not ready: Blender executable not found.")
+        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, None
+         gr.Warning("System not ready: UniRig repository or run.py script not found.")
+         return None
     if not bpy_import_ok:
-         gr.Warning("System warning: Initial 'bpy' import test failed. Attempting to proceed, but errors may occur.")
+         gr.Warning("System warning: Initial 'bpy' import test failed. Proceeding cautiously.")
 
     try:
         patch_asset_py()
@@ -298,26 +345,33 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
     if input_glb_file_obj is None:
         gr.Info("Please upload a .glb file first.")
         return None
-    input_glb_path = input_glb_file_obj # Gradio File component with type="filepath" returns a string path
+    input_glb_path = input_glb_file_obj
     print(f"Input GLB path received: {input_glb_path}")
 
     if not isinstance(input_glb_path, str) or not os.path.exists(input_glb_path):
-         raise gr.Error(f"Invalid input file path received or file does not exist: {input_glb_path}")
+         raise gr.Error(f"Invalid input file path or file 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.")
 
+    # Use a single temporary directory for all processing for this run
+    # This directory will be cleaned up at the end.
+    # The bootstrap script will be created inside this if run_unirig_command doesn't make its own.
+    # For clarity, run_unirig_command now handles its own bootstrap temp file.
     processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_")
     print(f"Using temporary processing directory: {processing_temp_dir}")
+    
+    diagnostic_script_path = None # For cleanup in finally
+
     try:
         base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
-        # Ensure paths used by UniRig are absolute, especially if its CWD changes
         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
+        unirig_script_to_run = UNIRIG_RUN_PY # This is an absolute path
 
-        print("\n--- Running Blender Python Environment Diagnostic Test ---")
+        print("\n--- Running Blender Python Environment Diagnostic Test (via bootstrap) ---")
+        # Content of the diagnostic script remains the same
         diagnostic_script_content = f"""
 import sys
 import os
@@ -331,30 +385,25 @@ print("\\nPYTHONPATH Environment Variable (as seen by script):")
 print(os.environ.get('PYTHONPATH', 'PYTHONPATH not set or empty'))
 print("\\nLD_LIBRARY_PATH Environment Variable (as seen by script):")
 print(os.environ.get('LD_LIBRARY_PATH', 'LD_LIBRARY_PATH not set or empty'))
-print("\\nPATH Environment Variable (as seen by script):")
-print(os.environ.get('PATH', 'PATH 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}}")
+    print(f"FAILED to import 'bpy': {{e}}")
     traceback.print_exc()
 try:
     print("\\nChecking for 'src' in CWD (should be UniRig repo root):")
-    if os.path.isdir('src'):
+    if os.path.isdir('src'): # This check is relative to CWD
         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.")
+            print("  WARNING: 'src/__init__.py' NOT FOUND.")
     else:
         print("  'src' directory NOT FOUND in CWD.")
     print("\\nAttempting: from src.inference.download import download")
-    from src.inference.download import download # Example import
+    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}}")
@@ -364,17 +413,13 @@ except Exception as e:
     traceback.print_exc()
 print("--- End Diagnostic Info ---")
 """
+        # Save diagnostic script to the processing_temp_dir
         diagnostic_script_path = os.path.join(processing_temp_dir, "env_diagnostic_test.py")
         with open(diagnostic_script_path, "w") as f: f.write(diagnostic_script_content)
-        try:
-            run_unirig_command(diagnostic_script_path, [], "Blender Env Diagnostic")
-            print("--- Finished Blender Python Environment Diagnostic Test ---\n")
-        except Exception as e_diag:
-            print(f"ERROR during diagnostic test execution: {e_diag}")
-            if os.path.exists(diagnostic_script_path): os.remove(diagnostic_script_path)
-            raise gr.Error(f"Blender environment diagnostic failed. Cannot proceed. Check logs. Error: {str(e_diag)[:500]}")
-        finally:
-            if os.path.exists(diagnostic_script_path): os.remove(diagnostic_script_path)
+        
+        run_unirig_command(diagnostic_script_path, [], "Blender Env Diagnostic") # Args are empty for diagnostic
+        print("--- Finished Blender Python Environment Diagnostic Test ---\n")
+        # If the above didn't raise an error, sys.path was likely fixed by bootstrap for the diagnostic.
 
         unirig_device_arg = "device=cpu"
         if DEVICE.type == 'cuda':
@@ -423,14 +468,22 @@ print("--- End Diagnostic Info ---")
         return gr.update(value=abs_final_rigged_glb_path)
     except gr.Error as e:
         print(f"A Gradio Error occurred: {e}")
-        if os.path.exists(processing_temp_dir): shutil.rmtree(processing_temp_dir)
-        raise e # Re-raise Gradio errors to be displayed in UI
+        # No need to re-raise, run_unirig_command already does or it's an explicit raise here.
+        # Let Gradio handle displaying it.
+        raise # Re-raise to ensure Gradio UI shows the 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)
-        raise gr.Error(f"An unexpected error occurred: {str(e)[:500]}. Check logs for details.")
+        raise gr.Error(f"An unexpected error occurred: {str(e)[:500]}. Check logs.")
     finally:
+        # Cleanup diagnostic script if it was created
+        if diagnostic_script_path and os.path.exists(diagnostic_script_path):
+            try:
+                os.remove(diagnostic_script_path)
+                print(f"Cleaned up diagnostic script: {diagnostic_script_path}")
+            except Exception as cleanup_e:
+                print(f"Error cleaning up diagnostic script {diagnostic_script_path}: {cleanup_e}")
+        # Cleanup the main processing directory
         if os.path.exists(processing_temp_dir):
              try:
                  shutil.rmtree(processing_temp_dir)
@@ -447,7 +500,7 @@ theme = gr.themes.Soft(
 
 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. Expected at {BLENDER_EXEC}")
+     startup_error_message = (f"CRITICAL STARTUP ERROR: Blender executable not located. Expected at {BLENDER_EXEC}")
 elif not unirig_repo_ok:
      startup_error_message = (f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}.")
 elif not unirig_run_py_ok:
@@ -456,16 +509,16 @@ elif not unirig_run_py_ok:
 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.")
+        gr.Markdown(f"# Application Startup Error\n\n{startup_error_message}\n\nPlease check Space logs.")
 else:
     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).
-             * Blender executable: `{blender_executable_to_use}`.
+             Upload a `.glb` mesh. UniRig predicts skeleton and skinning weights via Blender's Python.
+             * App Python: `{sys.version.split()[0]}`, UniRig (Blender): `{BLENDER_PYTHON_VERSION}`.
+             * Device: **{DEVICE.type.upper()}**.
+             * Blender: `{blender_executable_to_use}`.
              * UniRig Source: [https://github.com/VAST-AI-Research/UniRig](https://github.com/VAST-AI-Research/UniRig)
              """
          )
@@ -473,7 +526,7 @@ else:
             with gr.Column(scale=1):
                 input_model = gr.File(
                     label="Upload .glb Mesh File",
-                    type="filepath", # Returns a string path to a temporary copy
+                    type="filepath",
                     file_types=[".glb"]
                 )
                 submit_button = gr.Button("Rig Model", variant="primary")
@@ -493,5 +546,5 @@ if __name__ == "__main__":
         print("Launching Gradio interface...")
         iface.launch(share=False, ssr_mode=False)
     else:
-        print("ERROR: Gradio interface could not be created due to startup errors. Check logs above.")
+        print("ERROR: Gradio interface not created due to startup errors. Check logs.")