Spaces:
Running
on
Zero
Running
on
Zero
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 --- | |
APP_ROOT_DIR = os.path.abspath(os.path.dirname(__file__)) # Should be /home/user/app | |
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) | |
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 | |
BLENDER_EXEC_LOCAL_SYMLINK = os.path.join(LOCAL_BIN_DIR, "blender") | |
SETUP_SCRIPT = os.path.join(APP_ROOT_DIR, "setup_blender.sh") | |
SETUP_SCRIPT_TIMEOUT = 800 # Increased timeout | |
# --- Initial Checks --- | |
print("--- Environment Checks ---") | |
print(f"APP_ROOT_DIR: {APP_ROOT_DIR}") | |
print(f"Expected Blender Install Dir: {BLENDER_INSTALL_DIR}") | |
print(f"Expected Blender Executable: {BLENDER_EXEC}") | |
blender_executable_to_use = None | |
if os.path.exists(BLENDER_EXEC): | |
print(f"Blender executable found at direct local path: {BLENDER_EXEC}") | |
blender_executable_to_use = 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) | |
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...") | |
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, | |
capture_output=True, | |
text=True, | |
timeout=SETUP_SCRIPT_TIMEOUT | |
) | |
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 (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}") | |
elif os.path.exists(BLENDER_EXEC_LOCAL_SYMLINK): | |
blender_executable_to_use = BLENDER_EXEC_LOCAL_SYMLINK | |
print(f"Blender executable now found via local symlink: {BLENDER_EXEC_LOCAL_SYMLINK}") | |
elif os.path.exists(BLENDER_EXEC_SYMLINK): | |
blender_executable_to_use = BLENDER_EXEC_SYMLINK | |
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.") | |
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.") | |
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:]}") | |
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: | |
print("Testing bpy import via Blender...") | |
test_script_content = "import bpy; print('bpy imported successfully')" | |
test_result = subprocess.run( | |
[blender_executable_to_use, "--background", "--python-expr", test_script_content], | |
capture_output=True, text=True, check=True, timeout=30 | |
) | |
if "bpy imported successfully" in test_result.stdout: | |
print("Successfully imported 'bpy' using Blender executable.") | |
bpy_import_ok = True | |
else: | |
print(f"WARNING: 'bpy' import test via Blender returned unexpected output:\nSTDOUT:{test_result.stdout}\nSTDERR:{test_result.stderr}") | |
except subprocess.TimeoutExpired: | |
print("WARNING: 'bpy' import test via Blender timed out.") | |
except subprocess.CalledProcessError as e: | |
print(f"WARNING: Failed to import 'bpy' using Blender executable:\nSTDOUT:{e.stdout}\nSTDERR:{e.stderr}") | |
except Exception as e: | |
print(f"WARNING: Unexpected error during 'bpy' import test: {e}") | |
else: | |
print("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.") | |
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.") | |
else: | |
unirig_run_py_ok = True | |
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 ---") | |
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.") | |
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.") | |
process_env = os.environ.copy() | |
unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, 'src') | |
pythonpath_parts = [UNIRIG_REPO_DIR, unirig_src_dir] | |
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. | |
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')}") | |
print(f"Subprocess PATH: {process_env.get('PATH', 'Not set')}") | |
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)}") | |
try: | |
result = subprocess.run( | |
cmd, | |
cwd=UNIRIG_REPO_DIR, | |
capture_output=True, | |
text=True, | |
check=True, | |
env=process_env, | |
timeout=1800 | |
) | |
print(f"{step_name} STDOUT:\n{result.stdout}") | |
if result.stderr: | |
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.") | |
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." | |
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: | |
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}" | |
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} ---") | |
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 | |
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 | |
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() | |
except Exception as e: | |
print(f"Ignoring patch error: {e}") | |
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 | |
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}") | |
if not input_glb_path.lower().endswith(".glb"): | |
raise gr.Error("Invalid file type. Please upload a .glb file.") | |
processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_") | |
print(f"Using temporary processing directory: {processing_temp_dir}") | |
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 | |
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()}}") | |
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("\\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}}") | |
traceback.print_exc() | |
try: | |
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.") | |
print("\\nAttempting: from src.inference.download import download") | |
from src.inference.download import download # Example import | |
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") | |
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) | |
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}") | |
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.") | |
print("Step 1: Skeleton Prediction completed.") | |
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.") | |
print("Step 2: Skinning Prediction completed.") | |
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.") | |
print("Step 3: Merging completed.") | |
print(f"Successfully generated rigged model: {abs_final_rigged_glb_path}") | |
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 | |
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: | |
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}") | |
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"], | |
) | |
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}") | |
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: | |
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: | |
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}`. | |
* 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", # Returns a string path to a temporary copy | |
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], | |
) | |
submit_button.click( | |
fn=rig_glb_mesh_multistep, | |
inputs=[input_model], | |
outputs=[output_model] | |
) | |
if __name__ == "__main__": | |
if 'iface' in locals(): | |
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.") | |