app.py

import gradio as gr
import torch
import os
import sys
import tempfile
import shutil
import subprocess
import spaces # Keep this if you use @spaces.GPU
from typing import Any, Dict, List, Union # Added Union

# --- Configuration ---
UNIRIG_REPO_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "UniRig"))
BLENDER_INSTALL_DIR = "/opt/blender-4.2.0-linux-x64" # From your setup_blender.sh
BLENDER_PYTHON_VERSION_DIR = "4.2" # From your app.py
BLENDER_PYTHON_VERSION = "python3.11" # From your app.py and UniRig

# Construct paths
BLENDER_PYTHON_DIR = os.path.join(BLENDER_INSTALL_DIR, BLENDER_PYTHON_VERSION_DIR, "python")
BLENDER_PYTHON_BIN_DIR = os.path.join(BLENDER_PYTHON_DIR, "bin")
# Use Blender's main executable to run Python scripts via --python flag
BLENDER_EXEC = os.path.join(BLENDER_INSTALL_DIR, "blender")
BLENDER_EXEC_SYMLINK = "/usr/local/bin/blender" # Fallback symlink

SETUP_SCRIPT = os.path.join(os.path.dirname(__file__), "setup_blender.sh")

# --- Initial Checks ---
print("--- Environment Checks ---")
blender_executable_to_use = None # Use the main blender executable
if os.path.exists(BLENDER_EXEC):
    print(f"Blender executable found at direct path: {BLENDER_EXEC}")
    blender_executable_to_use = BLENDER_EXEC
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:
    print(f"Blender executable not found at {BLENDER_EXEC} or {BLENDER_EXEC_SYMLINK}. Running setup script...")
    if os.path.exists(SETUP_SCRIPT):
        try:
            # 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 after running setup
            if os.path.exists(BLENDER_EXEC):
                blender_executable_to_use = BLENDER_EXEC
                print(f"Blender executable now found at: {BLENDER_EXEC}")
            elif os.path.exists(BLENDER_EXEC_SYMLINK):
                 blender_executable_to_use = BLENDER_EXEC_SYMLINK
                 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
        )
        if "bpy imported successfully" in test_result.stdout:
            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 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}")
    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 (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)}")
        print(f"Gradio PyTorch CUDA Built Version: {torch.version.cuda}")
    except Exception as 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 ---
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):
            print(f"Warning: asset.py not found at {asset_py_path}, skipping patch.")
            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
        if problematic_line not in content:
            print("Problematic line not found in asset.py, patch might be unnecessary or file changed."); return

        print("Applying patch to asset.py...")
        content = content.replace(problematic_line, corrected_line)
        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"ERROR: Failed to patch asset.py: {e}. Proceeding cautiously.")
# --- End Helper Functions ---


# Decorator for ZeroGPU if needed
@spaces.GPU
def run_unirig_command(python_script_path: str, script_args: List[str], step_name: str):
    """
    Runs a specific UniRig PYTHON script (.py) using the Blender executable
    in background mode (`blender --background --python script.py -- args`).

    Args:
        python_script_path: Absolute path to the 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.
    """
    if not blender_executable_to_use:
        raise gr.Error("Blender executable path could not be determined. Cannot run UniRig step.")

    # --- Environment Setup for Subprocess ---
    process_env = os.environ.copy()

    # 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)
    process_env["PYTHONPATH"] = os.pathsep.join(filter(None, pythonpath_parts))
    print(f"Subprocess PYTHONPATH: {process_env['PYTHONPATH']}")

    # 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")
    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)
    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')}")


    # --- Execute Command ---
    # Command structure: blender --background --python script.py -- args
    cmd = [
        blender_executable_to_use,
        "--background",
        "--python", python_script_path,
        "--" # Separator between blender args and script args
    ] + script_args

    print(f"\n--- Running UniRig Step: {step_name} ---")
    print(f"Command: {' '.join(cmd)}") # Log the actual command

    try:
        # Execute Blender with the Python script.
        # CWD is crucial for relative imports and finding config files.
        result = subprocess.run(
            cmd,
            cwd=UNIRIG_REPO_DIR, # CWD is set to UniRig repo root
            capture_output=True,
            text=True,
            check=True, # Raises CalledProcessError on non-zero exit codes
            env=process_env, # Pass the modified environment
            timeout=1800 # 30 minutes timeout
        )
        print(f"{step_name} STDOUT:\n{result.stdout}")
        if result.stderr:
            print(f"{step_name} STDERR (Info/Warnings):\n{result.stderr}")
            # 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.")
        raise gr.Error(f"Processing step '{step_name}' timed out. Please try with a simpler model or check logs.")
    except subprocess.CalledProcessError as e:
        print(f"ERROR during {step_name}: Subprocess failed!")
        print(f"Command: {' '.join(e.cmd)}")
        print(f"Return code: {e.returncode}")
        print(f"--- {step_name} STDOUT ---:\n{e.stdout}")
        print(f"--- {step_name} STDERR ---:\n{e.stderr}")
        error_summary = e.stderr.strip().splitlines()
        last_lines = "\n".join(error_summary[-15:]) if error_summary else "No stderr output."

        specific_error = "Unknown error."
        # 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" 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: # 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:
            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.")
    except Exception as e_general:
        print(f"An unexpected Python exception occurred 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} ---")


@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.
    """
    # Perform readiness checks at the start of the function call
    if not blender_executable_to_use:
        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() # Apply patch if needed
    except Exception as e:
        print(f"Ignoring patch error: {e}")

    # --- Input Validation ---
    if input_glb_file_obj is None:
        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 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_")
    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")

        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 ---
        unirig_device_arg = "device=cpu"
        if DEVICE.type == 'cuda':
             unirig_device_arg = "device=cuda:0"
        print(f"UniRig steps will attempt to use device argument: {unirig_device_arg}")

        # --- Execute UniRig Steps ---
        # Step 1: Skeleton Prediction
        print("\nStarting Step 1: Predicting Skeleton...")
        skeleton_args = [
            "--config-name=skeleton_config",
            "with",
            f"input={abs_input_glb_path}",
            f"output={abs_skeleton_output_path}",
            unirig_device_arg
        ]
        run_unirig_command(unirig_script_to_run, skeleton_args, "Skeleton Prediction")
        if not os.path.exists(abs_skeleton_output_path):
            raise gr.Error("Skeleton prediction failed. Output file not created. Check logs.")
        print("Step 1: Skeleton Prediction completed.")

        # Step 2: Skinning Weight Prediction
        print("\nStarting Step 2: Predicting Skinning Weights...")
        skin_args = [
            "--config-name=skin_config",
            "with",
            f"input={abs_skeleton_output_path}",
            f"output={abs_skin_output_path}",
            unirig_device_arg
        ]
        run_unirig_command(unirig_script_to_run, skin_args, "Skinning Prediction")
        if not os.path.exists(abs_skin_output_path):
            raise gr.Error("Skinning prediction failed. Output file not created. Check logs.")
        print("Step 2: Skinning Prediction completed.")

        # Step 3: Merge Results
        print("\nStarting Step 3: Merging Results...")
        merge_args = [
            "--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",
            unirig_device_arg
        ]
        run_unirig_command(unirig_script_to_run, merge_args, "Merging Results")
        if not os.path.exists(abs_final_rigged_glb_path):
            raise gr.Error("Merging process failed. Final rigged GLB file not created. Check logs.")
        print("Step 3: Merging completed.")

        # --- Return Result ---
        print(f"Successfully generated rigged model: {abs_final_rigged_glb_path}")
        return gr.update(value=abs_final_rigged_glb_path) # Update Gradio output

    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
    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.")
    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 ---
theme = gr.themes.Soft(
    primary_hue=gr.themes.colors.sky,
    secondary_hue=gr.themes.colors.blue,
    neutral_hue=gr.themes.colors.slate,
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
)

# 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 unirig_repo_ok:
     startup_error_message = (f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}.")
elif not unirig_run_py_ok:
      startup_error_message = (f"CRITICAL STARTUP ERROR: UniRig run.py not found at {UNIRIG_RUN_PY}.")


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 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)
             """
         )
        with gr.Row():
            with gr.Column(scale=1):
                input_model = gr.File(
                    label="Upload .glb Mesh File",
                    type="filepath",
                    file_types=[".glb"]
                )
                submit_button = gr.Button("Rig Model", variant="primary")
            with gr.Column(scale=2):
                output_model = gr.Model3D(
                    label="Rigged 3D Model (.glb)",
                    clear_color=[0.8, 0.8, 0.8, 1.0],
                )

        # Connect button click to the processing function
        submit_button.click(
            fn=rig_glb_mesh_multistep,
            inputs=[input_model],
            outputs=[output_model]
        )

# --- Launch the Application ---
if __name__ == "__main__":
    # Check if the interface was successfully created
    if 'iface' in locals():
        print("Launching Gradio interface...")
        # Ensure share=False for security unless public link is needed
        iface.launch(share=False, ssr_mode=False) # Disable SSR
    else:
        # This indicates a startup error occurred before interface creation
        print("ERROR: Gradio interface could not be created due to startup errors. Check logs above.")