Update app.py

app.py

@@ -6,170 +6,353 @@ import tempfile
 import shutil
 import subprocess
 import spaces
-from typing import Any, Dict, Union
+from typing import Any, Dict, Union, List
 
 # --- Configuration ---
 # Path to the cloned UniRig repository directory within the Space
-UNIRIG_REPO_DIR = os.path.join(os.path.dirname(__file__), "UniRig")
+# 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
 
-# Path to the Blender Python site-packages
-BLENDER_PYTHON_PATH = "/opt/blender-4.2.0-linux-x64/4.2/python/lib/python3.11/site-packages"
+# Absolute path to the Blender installation provided in the Space environment
+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_EXEC = "/opt/blender-4.2.0-linux-x64/4.2/python/bin/python3.11"
+# Construct paths dynamically based on the above
+BLENDER_PYTHON_DIR = os.path.join(BLENDER_INSTALL_DIR, BLENDER_PYTHON_VERSION_DIR, "python")
+BLENDER_PYTHON_EXEC = os.path.join(BLENDER_PYTHON_DIR, "bin", BLENDER_PYTHON_VERSION) # Still needed for PYTHONPATH
+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 setup script
+# 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")
 
-# Check if Blender is installed
-if not os.path.exists("/usr/local/bin/blender"):
-    print("Blender not found. Installing...")
-    subprocess.run(["bash", SETUP_SCRIPT], check=True)
+# --- 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...")
+    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}.")
+        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
+             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("Blender is already installed.")
+    print(f"Blender Python executable found: {BLENDER_PYTHON_EXEC}")
 
-# Verify Blender Python path
-if os.path.exists(BLENDER_PYTHON_PATH):
-    print(f"Blender Python modules found at {BLENDER_PYTHON_PATH}")
-    if os.path.exists(os.path.join(BLENDER_PYTHON_PATH, "bpy.so")):
-        print("bpy.so found - Blender Python integration should work")
+# Verify Blender Python site-packages path and bpy module presence
+if os.path.exists(BLENDER_PYTHON_SITE_PACKAGES):
+    print(f"Blender Python site-packages found at: {BLENDER_PYTHON_SITE_PACKAGES}")
+    # Check for 'bpy' directory which contains __init__.py and modules
+    bpy_module_path = os.path.join(BLENDER_PYTHON_SITE_PACKAGES, "bpy")
+    if os.path.isdir(bpy_module_path) and os.path.exists(os.path.join(bpy_module_path, "__init__.py")):
+        print("Blender Python 'bpy' module directory found.")
+    # Fallback check for bpy.so if directory structure is different
+    elif os.path.exists(os.path.join(BLENDER_PYTHON_SITE_PACKAGES, "bpy.so")):
+         print("Blender Python 'bpy.so' found.")
     else:
-        print("bpy.so not found - Check Blender installation or path")
+        print("WARNING: Blender Python 'bpy' module not found in site-packages. Check Blender installation or paths.")
 else:
-    print(f"Blender Python modules not found at {BLENDER_PYTHON_PATH} - Update BLENDER_PYTHON_PATH")
+    print(f"WARNING: Blender Python site-packages 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}. Please clone it there.")
-    # Consider raising an error or displaying it in the UI if UniRig is critical for startup
+    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}")
 
+# Check PyTorch and CUDA
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {DEVICE}")
+print(f"Gradio environment using device: {DEVICE}")
 if DEVICE.type == 'cuda':
-    print(f"CUDA Device Name: {torch.cuda.get_device_name(0)}")
-    print(f"CUDA Version: {torch.version.cuda}")
+    try:
+        print(f"Gradio CUDA Device Name: {torch.cuda.get_device_name(0)}")
+        print(f"Gradio PyTorch CUDA Version: {torch.version.cuda}")
+    except Exception as e:
+        print(f"Could not get CUDA device details: {e}")
 else:
-    print("Warning: CUDA not available or not detected by PyTorch. UniRig performance will be severely impacted.")
+    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("--- End Environment Checks ---")
+
+# --- Helper Functions ---
 
 def patch_asset_py():
-    """Temporary patch to fix type hinting error in 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()
-        # Check if patch is already applied
-        if "meta: Union[Dict[str, Any], None]" in content:
+
+        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
-        # Replace the problematic line
-        content = content.replace("meta: Union[Dict[str, ...], None]=None", "meta: Union[Dict[str, Any], None]=None")
-        # Ensure 'Any' is imported
-        if "from typing import Any" not in content:
-            content = "from typing import Any\n" + content
+        if problematic_line not in content:
+             print("Problematic line not found in asset.py, patch might be unnecessary.")
+             return
+
+        print("Applying patch to asset.py...")
+        content = content.replace(problematic_line, corrected_line)
+        if typing_import not in content:
+            if "from typing import" in content:
+                 content = content.replace("from typing import", f"{typing_import}\nfrom typing import", 1)
+            else:
+                 content = f"{typing_import}\n{content}"
+
         with open(asset_py_path, "w") as f:
             f.write(content)
         print("Successfully patched asset.py")
+
     except Exception as e:
-        print(f"Failed to patch asset.py: {e}")
-        raise
-
-@spaces.GPU # Decorator for ZeroGPU
-def run_unirig_command(command_list, step_name):
-    """Run UniRig commands using Blender's Python interpreter."""
-    cmd = [BLENDER_PYTHON_EXEC] + command_list[1:]  # Use Blender's Python instead of bash!!
-    
-    print(f"Running {step_name}: {' '.join(cmd)}")
-    
+        # 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):
+    """
+    Runs a specific UniRig SHELL script (.sh) using bash in a subprocess,
+    ensuring the correct environment (PYTHONPATH, etc.) is set.
+
+    Args:
+        script_path: Absolute path to the .sh script to execute.
+        args: A list of command-line arguments for the script.
+        step_name: Name of the step for logging.
+    """
+    # Command is now bash + script path + arguments
+    cmd = ["bash", script_path] + args
+
+    print(f"\n--- Running UniRig Step: {step_name} ---")
+    # Use subprocess.list2cmdline for a more accurate representation if needed,
+    # but simple join is often fine for logging.
+    print(f"Command: {' '.join(cmd)}")
+
+    # Prepare the environment for the subprocess (shell script)
+    # The shell script will likely invoke python and needs the correct env
     process_env = os.environ.copy()
-    
-    # Determine the path to the 'src' directory within UniRig, where the 'unirig' package resides.
     unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, "src")
 
-    # Explicitly add BLENDER_PYTHON_PATH, UNIRIG_REPO_DIR/src, and UNIRIG_REPO_DIR to PYTHONPATH
-    new_pythonpath_parts = [BLENDER_PYTHON_PATH, unirig_src_dir, UNIRIG_REPO_DIR]
-    #existing_pythonpath = process_env.get('PYTHONPATH', '')
-    #if existing_pythonpath:
-    #    new_pythonpath_parts.extend(existing_pythonpath.split(os.pathsep))
-    process_env["PYTHONPATH"] = os.pathsep.join(filter(None, new_pythonpath_parts))
-    print(f"Set PYTHONPATH for subprocess: {process_env['PYTHONPATH']}")
+    # Crucial: Set PYTHONPATH so the python invoked by the shell script finds Blender's
+    # site-packages and the UniRig source code.
+    pythonpath_parts = [
+        BLENDER_PYTHON_SITE_PACKAGES,
+        unirig_src_dir,
+        UNIRIG_REPO_DIR  # Include base repo dir as well
+    ]
+    process_env["PYTHONPATH"] = os.pathsep.join(filter(None, pythonpath_parts))
+    print(f"Subprocess PYTHONPATH: {process_env['PYTHONPATH']}")
+
+    # Add Blender's library path and potentially the python bin path to PATH
+    # This helps the shell find the correct python if needed, and shared libs
+    blender_lib_path = os.path.join(BLENDER_PYTHON_DIR, "lib")
+    blender_bin_path = os.path.join(BLENDER_PYTHON_DIR, "bin")
+    process_env["LD_LIBRARY_PATH"] = f"{blender_lib_path}{os.pathsep}{process_env.get('LD_LIBRARY_PATH', '')}"
+    process_env["PATH"] = f"{blender_bin_path}{os.pathsep}{process_env.get('PATH', '')}" # Add blender python bin to PATH
+    print(f"Subprocess LD_LIBRARY_PATH: {process_env['LD_LIBRARY_PATH']}")
+    print(f"Subprocess PATH: {process_env['PATH']}")
+
 
     try:
-        result = subprocess.run(cmd, cwd=UNIRIG_REPO_DIR, capture_output=True, text=True, check=True, env=process_env)
+        # Execute the shell script.
+        # cwd=UNIRIG_REPO_DIR ensures the script runs from the repo's root,
+        # which is often necessary for finding relative paths within the script.
+        result = subprocess.run(
+            cmd,
+            cwd=UNIRIG_REPO_DIR,
+            capture_output=True,
+            text=True,
+            check=True, # Raises CalledProcessError on non-zero exit codes
+            env=process_env
+        )
         print(f"{step_name} STDOUT:\n{result.stdout}")
         if result.stderr:
-            print(f"{step_name} STDERR (non-fatal or warnings):\n{result.stderr}")
+            # Treat stderr as potential warnings unless check=True failed
+            print(f"{step_name} STDERR (Warnings/Info):\n{result.stderr}")
+
     except subprocess.CalledProcessError as e:
-        print(f"ERROR during {step_name}:")
+        # Error occurred within the shell script (non-zero exit code)
+        print(f"ERROR during {step_name}: Subprocess failed!")
         print(f"Command: {' '.join(e.cmd)}")
         print(f"Return code: {e.returncode}")
-        print(f"Stdout: {e.stdout}")
-        print(f"Stderr: {e.stderr}")
-        error_summary = e.stderr.splitlines()[-5:]
-        raise gr.Error(f"Error in UniRig {step_name}. Details: {' '.join(error_summary)}")
+        print(f"--- {step_name} STDOUT ---:\n{e.stdout}")
+        print(f"--- {step_name} STDERR ---:\n{e.stderr}")
+        # Provide a concise error message to Gradio
+        error_summary = e.stderr.strip().splitlines()
+        last_lines = "\n".join(error_summary[-5:]) if error_summary else "No stderr output."
+        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 or script for {step_name}. Command: {' '.join(cmd)}")
-        raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs and script paths.")
+        # This error means 'bash' or the script_path wasn't found
+        print(f"ERROR: Could not find executable 'bash' or script '{script_path}' for {step_name}.")
+        print(f"Attempted command: {' '.join(cmd)}")
+        raise gr.Error(f"Setup error for UniRig '{step_name}'. 'bash' or script '{os.path.basename(script_path)}' not found.")
+
     except Exception as e_general:
+        # Catch any other unexpected Python errors during subprocess handling
         print(f"An unexpected Python exception occurred in run_unirig_command for {step_name}: {e_general}")
-        raise gr.Error(f"Unexpected Python error during {step_name}: {str(e_general)[:500]}")
+        import traceback
+        traceback.print_exc()
+        raise gr.Error(f"Unexpected Python error during '{step_name}' execution: {str(e_general)[:500]}")
+
+    print(f"--- Finished UniRig Step: {step_name} ---")
+
 
 @spaces.GPU
 def rig_glb_mesh_multistep(input_glb_file_obj):
-    """Rig a GLB mesh using UniRig's multi-step process."""
-    patch_asset_py()
-
-    if not os.path.isdir(UNIRIG_REPO_DIR):
-        raise gr.Error(f"UniRig repository not found at {UNIRIG_REPO_DIR}.")
+    """
+    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.
+    """
+    try:
+        patch_asset_py() # Attempt patch, might still be relevant
+    except gr.Error as e:
+         # If patching is critical and fails, raise the error
+         # raise e
+         print(f"Ignoring patch error: {e}") # Or just log it
+    except Exception as e:
+         print(f"Ignoring unexpected patch error: {e}")
+         # raise gr.Error(f"Failed to prepare UniRig environment: {e}")
 
+    # --- Input Validation ---
     if input_glb_file_obj is None:
         raise gr.Error("No input file provided. Please upload a .glb mesh.")
 
     input_glb_path = input_glb_file_obj
     print(f"Input GLB path received: {input_glb_path}")
 
+    if not os.path.exists(input_glb_path):
+         raise gr.Error(f"Input file path does not exist: {input_glb_path}")
+    if not input_glb_path.lower().endswith(".glb"):
+         raise gr.Error("Invalid file type. Please upload a .glb file.")
+
+    # --- Setup Temporary Directory ---
     processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_")
     print(f"Using temporary processing directory: {processing_temp_dir}")
 
     try:
+        # --- Define File Paths ---
         base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
+        abs_input_glb_path = os.path.abspath(input_glb_path)
         abs_skeleton_output_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
         abs_skin_output_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx")
         abs_final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")
 
+        # --- Define Absolute Paths to UniRig SHELL Scripts ---
+        # Construct absolute paths based on UNIRIG_REPO_DIR, using .sh extension
+        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")
+
+        # --- Execute UniRig Steps ---
+
         # Step 1: Skeleton Prediction
-        print("Step 1: Predicting Skeleton...")
-        skeleton_cmd = ["python", "launch/inference/generate_skeleton.py", "--input", input_glb_path, "--output", abs_skeleton_output_path]
-        run_unirig_command(skeleton_cmd, "Skeleton Prediction")
+        print("\nStarting Step 1: Predicting Skeleton...")
+        # Arguments for the generate_skeleton.sh script
+        skeleton_args = [
+            "--input", abs_input_glb_path,
+            "--output", abs_skeleton_output_path
+            # Add any other arguments the shell script expects
+        ]
+        # Check if script exists before running
+        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(abs_skeleton_output_path):
-            raise gr.Error("Skeleton prediction failed to produce an output file. Check logs for UniRig errors.")
+            raise gr.Error("Skeleton prediction failed. Output file not created. Check logs.")
+        print("Step 1: Skeleton Prediction completed.")
 
         # Step 2: Skinning Weight Prediction
-        print("Step 2: Predicting Skinning Weights...")
-        skin_cmd = ["python", "launch/inference/generate_skin.py", "--input", abs_skeleton_output_path, "--source", input_glb_path, "--output", abs_skin_output_path]
-        run_unirig_command(skin_cmd, "Skinning Prediction")
+        print("\nStarting Step 2: Predicting Skinning Weights...")
+        # Arguments for the generate_skin.sh script
+        skin_args = [
+            "--input", abs_skeleton_output_path, # Input is the skeleton from step 1
+            "--source", abs_input_glb_path,    # Source mesh
+            "--output", abs_skin_output_path
+            # Add any other arguments the shell script expects
+        ]
+        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")
         if not os.path.exists(abs_skin_output_path):
-            raise gr.Error("Skinning prediction failed to produce an output file.")
+            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("Step 3: Merging Results...")
-        merge_cmd = ["python", "launch/inference/merge.py", "--source", abs_skin_output_path, "--target", input_glb_path, "--output", abs_final_rigged_glb_path]
-        run_unirig_command(merge_cmd, "Merging")
+        print("\nStarting Step 3: Merging Results...")
+        # Arguments for the merge.sh script
+        merge_args = [
+            "--source", abs_skin_output_path, # Input is the skinned result from step 2
+            "--target", abs_input_glb_path,   # Target is the original mesh
+            "--output", abs_final_rigged_glb_path
+            # Add any other arguments the shell script expects
+        ]
+        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")
         if not os.path.exists(abs_final_rigged_glb_path):
-            raise gr.Error("Merging process failed to produce the final rigged GLB file.")
+            raise gr.Error("Merging process failed. Final rigged GLB file not created. Check logs.")
+        print("Step 3: Merging completed.")
 
+        # --- Return Result ---
+        print(f"Successfully generated rigged model: {abs_final_rigged_glb_path}")
         return abs_final_rigged_glb_path
 
-    except gr.Error:
+    except gr.Error as e:
+        # If a gr.Error was raised, clean up and re-raise
+        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}")
-        raise
+        raise e
+
     except Exception as e:
+        # Catch any other unexpected errors
         print(f"An unexpected error occurred in rig_glb_mesh_multistep: {e}")
+        import traceback
+        traceback.print_exc()
+        # Clean up temporary files before raising the error
         if os.path.exists(processing_temp_dir):
             shutil.rmtree(processing_temp_dir)
             print(f"Cleaned up temporary directory: {processing_temp_dir}")
+        # Raise a generic Gradio error
         raise gr.Error(f"An unexpected error occurred during processing: {str(e)[:500]}")
 
-# --- Gradio Interface ---
+
+# --- Gradio Interface Definition ---
 theme = gr.themes.Soft(
     primary_hue=gr.themes.colors.sky,
     secondary_hue=gr.themes.colors.blue,
@@ -177,24 +360,49 @@ theme = gr.themes.Soft(
     font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
 )
 
-if not os.path.isdir(UNIRIG_REPO_DIR) and __name__ == "__main__":
-    print(f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}.")
-
-iface = gr.Interface(
-    fn=rig_glb_mesh_multistep,
-    inputs=gr.File(label="Upload .glb Mesh File", type="filepath"),
-    outputs=gr.Model3D(label="Rigged 3D Model (.glb)", clear_color=[0.8, 0.8, 0.8, 1.0]),
-    title="UniRig Auto-Rigger (Python 3.11 / PyTorch 2.3+)",
-    description=(
-        "Upload a 3D mesh in `.glb` format. This application uses the latest UniRig to automatically rig the mesh.\n"
-        "Running on: {DEVICE.upper()}. UniRig repo expected at: '{os.path.basename(UNIRIG_REPO_DIR)}'.\n"
-        "UniRig Source: https://github.com/VAST-AI-Research/UniRig"
-    ),
-    cache_examples=False,
-    theme=theme
-)
+# 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."
+     )
+     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}")
+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
+         ),
+         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"
+             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
+     )
 
+# --- Launch the Application ---
 if __name__ == "__main__":
-    if not os.path.isdir(UNIRIG_REPO_DIR):
-        print(f"CRITICAL: UniRig repository not found at {UNIRIG_REPO_DIR}.")
-    iface.launch()
+    # 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.")
+