Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -9,72 +9,88 @@ import spaces # Keep this if you use @spaces.GPU
|
|
9 |
from typing import Any, Dict, List, Union # Added Union
|
10 |
|
11 |
# --- Configuration ---
|
12 |
-
|
13 |
-
BLENDER_INSTALL_DIR = "/opt/blender-4.2.0-linux-x64" # From your setup_blender.sh
|
14 |
-
BLENDER_PYTHON_VERSION_DIR = "4.2" # From your app.py
|
15 |
-
BLENDER_PYTHON_VERSION = "python3.11" # From your app.py and UniRig
|
16 |
|
17 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
18 |
BLENDER_PYTHON_DIR = os.path.join(BLENDER_INSTALL_DIR, BLENDER_PYTHON_VERSION_DIR, "python")
|
19 |
BLENDER_PYTHON_BIN_DIR = os.path.join(BLENDER_PYTHON_DIR, "bin")
|
20 |
-
# Use Blender's main executable to run Python scripts via --python flag
|
21 |
BLENDER_EXEC = os.path.join(BLENDER_INSTALL_DIR, "blender")
|
22 |
-
BLENDER_EXEC_SYMLINK = "/usr/local/bin/blender" # Fallback symlink
|
23 |
|
24 |
-
|
25 |
-
#
|
26 |
-
|
|
|
|
|
27 |
|
28 |
|
|
|
|
|
|
|
29 |
# --- Initial Checks ---
|
30 |
print("--- Environment Checks ---")
|
31 |
-
|
|
|
|
|
|
|
|
|
32 |
if os.path.exists(BLENDER_EXEC):
|
33 |
-
print(f"Blender executable found at direct path: {BLENDER_EXEC}")
|
34 |
blender_executable_to_use = BLENDER_EXEC
|
35 |
-
elif os.path.exists(
|
36 |
-
print(f"Blender executable found via symlink: {
|
|
|
|
|
|
|
37 |
blender_executable_to_use = BLENDER_EXEC_SYMLINK
|
38 |
else:
|
39 |
-
print(f"Blender executable not found at {BLENDER_EXEC} or {BLENDER_EXEC_SYMLINK}. Running setup script...")
|
40 |
if os.path.exists(SETUP_SCRIPT):
|
41 |
try:
|
42 |
# Run setup script if Blender not found
|
|
|
43 |
setup_result = subprocess.run(
|
44 |
["bash", SETUP_SCRIPT],
|
45 |
check=True,
|
46 |
capture_output=True,
|
47 |
text=True,
|
48 |
-
timeout=SETUP_SCRIPT_TIMEOUT
|
49 |
)
|
50 |
print("Setup script executed successfully.")
|
51 |
print(f"Setup STDOUT:\n{setup_result.stdout}")
|
52 |
if setup_result.stderr: print(f"Setup STDERR:\n{setup_result.stderr}")
|
53 |
|
54 |
-
# Re-check for executable after running setup
|
55 |
if os.path.exists(BLENDER_EXEC):
|
56 |
blender_executable_to_use = BLENDER_EXEC
|
57 |
-
print(f"Blender executable now found at: {BLENDER_EXEC}")
|
|
|
|
|
|
|
58 |
elif os.path.exists(BLENDER_EXEC_SYMLINK):
|
59 |
blender_executable_to_use = BLENDER_EXEC_SYMLINK
|
60 |
-
print(f"Blender executable now found via symlink: {BLENDER_EXEC_SYMLINK}")
|
61 |
|
62 |
if not blender_executable_to_use:
|
63 |
-
|
64 |
-
raise RuntimeError(f"Setup script ran but Blender executable still not found at {BLENDER_EXEC} or {BLENDER_EXEC_SYMLINK}.")
|
65 |
-
|
66 |
except subprocess.TimeoutExpired:
|
67 |
print(f"ERROR: Setup script timed out after {SETUP_SCRIPT_TIMEOUT} seconds: {SETUP_SCRIPT}")
|
68 |
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.")
|
69 |
except subprocess.CalledProcessError as e:
|
70 |
print(f"ERROR running setup script: {SETUP_SCRIPT}\nStderr: {e.stderr}")
|
71 |
-
# Raise a Gradio error to notify the user
|
72 |
raise gr.Error(f"Failed to execute setup script. Check logs. Stderr: {e.stderr[-500:]}")
|
73 |
except Exception as e:
|
74 |
-
# Catch any other exceptions during setup
|
75 |
raise gr.Error(f"Unexpected error running setup script '{SETUP_SCRIPT}': {e}")
|
76 |
else:
|
77 |
-
# If setup script itself is missing
|
78 |
raise gr.Error(f"Blender executable not found and setup script missing: {SETUP_SCRIPT}")
|
79 |
|
80 |
# Verify bpy import using the found Blender executable
|
@@ -83,50 +99,38 @@ if blender_executable_to_use:
|
|
83 |
try:
|
84 |
print("Testing bpy import via Blender...")
|
85 |
test_script_content = "import bpy; print('bpy imported successfully')"
|
86 |
-
# Use --python-expr for the bpy test
|
87 |
test_result = subprocess.run(
|
88 |
[blender_executable_to_use, "--background", "--python-expr", test_script_content],
|
89 |
-
capture_output=True, text=True, check=True, timeout=30
|
90 |
)
|
91 |
if "bpy imported successfully" in test_result.stdout:
|
92 |
print("Successfully imported 'bpy' using Blender executable.")
|
93 |
bpy_import_ok = True
|
94 |
else:
|
95 |
-
# Log warning but don't necessarily stop startup
|
96 |
print(f"WARNING: 'bpy' import test via Blender returned unexpected output:\nSTDOUT:{test_result.stdout}\nSTDERR:{test_result.stderr}")
|
97 |
-
|
98 |
except subprocess.TimeoutExpired:
|
99 |
print("WARNING: 'bpy' import test via Blender timed out.")
|
100 |
except subprocess.CalledProcessError as e:
|
101 |
-
# Log specific error if bpy import fails
|
102 |
print(f"WARNING: Failed to import 'bpy' using Blender executable:\nSTDOUT:{e.stdout}\nSTDERR:{e.stderr}")
|
103 |
except Exception as e:
|
104 |
-
# Catch any other exception during the test
|
105 |
print(f"WARNING: Unexpected error during 'bpy' import test: {e}")
|
106 |
else:
|
107 |
-
# This case should ideally not be reached if setup logic is correct
|
108 |
print("WARNING: Cannot test bpy import as Blender executable was not found.")
|
109 |
|
110 |
-
|
111 |
# Check for UniRig repository and run.py
|
112 |
unirig_repo_ok = False
|
113 |
unirig_run_py_ok = False
|
114 |
UNIRIG_RUN_PY = os.path.join(UNIRIG_REPO_DIR, "run.py")
|
115 |
if not os.path.isdir(UNIRIG_REPO_DIR):
|
116 |
-
# Critical error if UniRig repo is missing
|
117 |
raise gr.Error(f"UniRig repository missing at: {UNIRIG_REPO_DIR}. Ensure it's cloned correctly.")
|
118 |
else:
|
119 |
print(f"UniRig repository found at: {UNIRIG_REPO_DIR}")
|
120 |
unirig_repo_ok = True
|
121 |
-
# Check specifically for run.py within the repo
|
122 |
if not os.path.exists(UNIRIG_RUN_PY):
|
123 |
-
# Critical error if run.py is missing
|
124 |
raise gr.Error(f"UniRig's run.py not found at {UNIRIG_RUN_PY}. Check UniRig clone.")
|
125 |
else:
|
126 |
unirig_run_py_ok = True
|
127 |
|
128 |
-
|
129 |
-
# Check PyTorch and CUDA for Gradio environment (less critical for startup)
|
130 |
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
131 |
print(f"Gradio environment using device: {DEVICE}")
|
132 |
if DEVICE.type == 'cuda':
|
@@ -137,10 +141,8 @@ if DEVICE.type == 'cuda':
|
|
137 |
print(f"Could not get Gradio CUDA device details: {e}")
|
138 |
else:
|
139 |
print("Warning: Gradio environment CUDA not available.")
|
140 |
-
|
141 |
print("--- End Environment Checks ---")
|
142 |
|
143 |
-
# --- Helper Functions ---
|
144 |
def patch_asset_py():
|
145 |
"""Temporary patch to fix type hinting error in UniRig's asset.py"""
|
146 |
asset_py_path = os.path.join(UNIRIG_REPO_DIR, "src", "data", "asset.py")
|
@@ -148,17 +150,14 @@ def patch_asset_py():
|
|
148 |
if not os.path.exists(asset_py_path):
|
149 |
print(f"Warning: asset.py not found at {asset_py_path}, skipping patch.")
|
150 |
return
|
151 |
-
|
152 |
with open(asset_py_path, "r") as f: content = f.read()
|
153 |
problematic_line = "meta: Union[Dict[str, ...], None]=None"
|
154 |
corrected_line = "meta: Union[Dict[str, Any], None]=None"
|
155 |
typing_import = "from typing import Any"
|
156 |
-
|
157 |
if corrected_line in content:
|
158 |
print("Patch already applied to asset.py"); return
|
159 |
if problematic_line not in content:
|
160 |
print("Problematic line not found in asset.py, patch might be unnecessary or file changed."); return
|
161 |
-
|
162 |
print("Applying patch to asset.py...")
|
163 |
content = content.replace(problematic_line, corrected_line)
|
164 |
if typing_import not in content:
|
@@ -168,84 +167,74 @@ def patch_asset_py():
|
|
168 |
content = f"{typing_import}\n{content}"
|
169 |
with open(asset_py_path, "w") as f: f.write(content)
|
170 |
print("Successfully patched asset.py")
|
171 |
-
|
172 |
except Exception as e:
|
173 |
print(f"ERROR: Failed to patch asset.py: {e}. Proceeding cautiously.")
|
174 |
-
# --- End Helper Functions ---
|
175 |
|
176 |
-
|
177 |
-
# Decorator for ZeroGPU if needed
|
178 |
@spaces.GPU
|
179 |
def run_unirig_command(python_script_path: str, script_args: List[str], step_name: str):
|
180 |
"""
|
181 |
Runs a specific UniRig PYTHON script (.py) using the Blender executable
|
182 |
in background mode (`blender --background --python script.py -- args`).
|
183 |
-
|
184 |
-
Args:
|
185 |
-
python_script_path: Absolute path to the Python script to execute within Blender.
|
186 |
-
script_args: A list of command-line arguments FOR THE PYTHON SCRIPT.
|
187 |
-
step_name: Name of the step for logging.
|
188 |
"""
|
189 |
if not blender_executable_to_use:
|
190 |
raise gr.Error("Blender executable path could not be determined. Cannot run UniRig step.")
|
191 |
|
192 |
-
# --- Environment Setup for Subprocess ---
|
193 |
process_env = os.environ.copy()
|
194 |
-
|
195 |
-
# Explicitly add UniRig root AND its 'src' directory to PYTHONPATH
|
196 |
unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, 'src')
|
197 |
-
pythonpath_parts = [UNIRIG_REPO_DIR, unirig_src_dir]
|
198 |
existing_pythonpath = process_env.get('PYTHONPATH', '')
|
199 |
if existing_pythonpath:
|
200 |
pythonpath_parts.append(existing_pythonpath)
|
201 |
process_env["PYTHONPATH"] = os.pathsep.join(filter(None, pythonpath_parts))
|
202 |
print(f"Subprocess PYTHONPATH: {process_env['PYTHONPATH']}")
|
203 |
|
204 |
-
# LD_LIBRARY_PATH: Inherit system path
|
205 |
-
blender_main_lib_path = os.path.join(BLENDER_INSTALL_DIR, "lib")
|
206 |
-
blender_python_lib_path = os.path.join(BLENDER_PYTHON_DIR, "lib")
|
|
|
207 |
ld_path_parts = []
|
208 |
if os.path.exists(blender_main_lib_path): ld_path_parts.append(blender_main_lib_path)
|
209 |
if os.path.exists(blender_python_lib_path): ld_path_parts.append(blender_python_lib_path)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
210 |
existing_ld_path = process_env.get('LD_LIBRARY_PATH', '')
|
211 |
if existing_ld_path: ld_path_parts.append(existing_ld_path)
|
212 |
if ld_path_parts:
|
213 |
process_env["LD_LIBRARY_PATH"] = os.pathsep.join(filter(None, ld_path_parts))
|
214 |
print(f"Subprocess LD_LIBRARY_PATH: {process_env.get('LD_LIBRARY_PATH', 'Not set')}")
|
|
|
215 |
|
216 |
|
217 |
-
# --- Execute Command ---
|
218 |
-
# Command structure: blender --background --python script.py -- args
|
219 |
cmd = [
|
220 |
blender_executable_to_use,
|
221 |
"--background",
|
222 |
"--python", python_script_path,
|
223 |
-
"--"
|
224 |
] + script_args
|
225 |
-
|
226 |
print(f"\n--- Running UniRig Step: {step_name} ---")
|
227 |
-
print(f"Command: {' '.join(cmd)}")
|
228 |
-
|
229 |
try:
|
230 |
-
# Execute Blender with the Python script.
|
231 |
-
# CWD is crucial for relative imports and finding config files.
|
232 |
result = subprocess.run(
|
233 |
cmd,
|
234 |
-
cwd=UNIRIG_REPO_DIR,
|
235 |
capture_output=True,
|
236 |
text=True,
|
237 |
-
check=True,
|
238 |
-
env=process_env,
|
239 |
-
timeout=1800
|
240 |
)
|
241 |
print(f"{step_name} STDOUT:\n{result.stdout}")
|
242 |
if result.stderr:
|
243 |
print(f"{step_name} STDERR (Info/Warnings):\n{result.stderr}")
|
244 |
-
# More robust check for errors in stderr even on success
|
245 |
stderr_lower = result.stderr.lower()
|
246 |
if "error" in stderr_lower or "failed" in stderr_lower or "traceback" in stderr_lower:
|
247 |
print(f"WARNING: Potential error messages found in STDERR for {step_name} despite success exit code.")
|
248 |
-
|
249 |
except subprocess.TimeoutExpired:
|
250 |
print(f"ERROR: {step_name} timed out after 30 minutes.")
|
251 |
raise gr.Error(f"Processing step '{step_name}' timed out. Please try with a simpler model or check logs.")
|
@@ -257,9 +246,7 @@ def run_unirig_command(python_script_path: str, script_args: List[str], step_nam
|
|
257 |
print(f"--- {step_name} STDERR ---:\n{e.stderr}")
|
258 |
error_summary = e.stderr.strip().splitlines()
|
259 |
last_lines = "\n".join(error_summary[-15:]) if error_summary else "No stderr output."
|
260 |
-
|
261 |
specific_error = "Unknown error."
|
262 |
-
# Refine error checking based on common patterns
|
263 |
if "ModuleNotFoundError: No module named 'src'" in e.stderr:
|
264 |
specific_error = "UniRig script failed to import its own 'src' module. Check PYTHONPATH and CWD for the subprocess. See diagnostic info above."
|
265 |
elif "ModuleNotFoundError: No module named 'bpy'" in e.stderr:
|
@@ -272,14 +259,13 @@ def run_unirig_command(python_script_path: str, script_args: List[str], step_nam
|
|
272 |
specific_error = "CUDA out of memory. Try a smaller model or a Space with more GPU RAM."
|
273 |
elif "hydra.errors.ConfigCompositionException" in e.stderr:
|
274 |
specific_error = f"Hydra configuration error. Check the arguments passed to run.py: {script_args}. Details: {last_lines}"
|
275 |
-
elif "Error: Cannot read file" in e.stderr:
|
276 |
specific_error = f"Blender could not read an input/output file. Check paths. Details: {last_lines}"
|
277 |
-
elif "Error:" in e.stderr:
|
278 |
specific_error = f"Blender reported an error. Details: {last_lines}"
|
279 |
else:
|
280 |
specific_error = f"Check logs. Last error lines:\n{last_lines}"
|
281 |
raise gr.Error(f"Error in UniRig '{step_name}'. {specific_error}")
|
282 |
-
|
283 |
except FileNotFoundError:
|
284 |
print(f"ERROR: Could not find Blender executable '{blender_executable_to_use}' or script '{python_script_path}' for {step_name}.")
|
285 |
raise gr.Error(f"Setup error for UniRig '{step_name}'. Blender or Python script not found.")
|
@@ -288,70 +274,65 @@ def run_unirig_command(python_script_path: str, script_args: List[str], step_nam
|
|
288 |
import traceback
|
289 |
traceback.print_exc()
|
290 |
raise gr.Error(f"Unexpected Python error during '{step_name}' execution: {str(e_general)[:500]}")
|
291 |
-
|
292 |
print(f"--- Finished UniRig Step: {step_name} ---")
|
293 |
|
294 |
-
|
295 |
@spaces.GPU
|
296 |
def rig_glb_mesh_multistep(input_glb_file_obj):
|
297 |
"""
|
298 |
Main Gradio function to rig a GLB mesh using UniRig's multi-step process.
|
299 |
"""
|
300 |
-
# Perform readiness checks at the start of the function call
|
301 |
if not blender_executable_to_use:
|
302 |
gr.Warning("System not ready: Blender executable not found. Please wait or check logs.")
|
303 |
-
return None
|
304 |
if not unirig_repo_ok or not unirig_run_py_ok:
|
305 |
gr.Warning("System not ready: UniRig repository or run.py script not found. Check setup.")
|
306 |
-
return None
|
307 |
if not bpy_import_ok:
|
308 |
gr.Warning("System warning: Initial 'bpy' import test failed. Attempting to proceed, but errors may occur.")
|
309 |
|
310 |
try:
|
311 |
-
patch_asset_py()
|
312 |
except Exception as e:
|
313 |
print(f"Ignoring patch error: {e}")
|
314 |
|
315 |
-
# --- Input Validation ---
|
316 |
if input_glb_file_obj is None:
|
317 |
gr.Info("Please upload a .glb file first.")
|
318 |
return None
|
319 |
-
input_glb_path = input_glb_file_obj
|
320 |
print(f"Input GLB path received: {input_glb_path}")
|
|
|
321 |
if not isinstance(input_glb_path, str) or not os.path.exists(input_glb_path):
|
322 |
-
raise gr.Error(f"Invalid input file path received: {input_glb_path}")
|
323 |
if not input_glb_path.lower().endswith(".glb"):
|
324 |
raise gr.Error("Invalid file type. Please upload a .glb file.")
|
325 |
|
326 |
-
# --- Setup Temporary Directory ---
|
327 |
processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_")
|
328 |
print(f"Using temporary processing directory: {processing_temp_dir}")
|
329 |
-
|
330 |
try:
|
331 |
-
# --- Define File Paths ---
|
332 |
base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
|
|
|
333 |
abs_input_glb_path = os.path.abspath(input_glb_path)
|
334 |
abs_skeleton_output_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
|
335 |
abs_skin_output_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx")
|
336 |
abs_final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")
|
337 |
-
|
338 |
unirig_script_to_run = UNIRIG_RUN_PY
|
339 |
|
340 |
-
# --- Run Blender Python Environment Diagnostic Test ---
|
341 |
print("\n--- Running Blender Python Environment Diagnostic Test ---")
|
342 |
diagnostic_script_content = f"""
|
343 |
import sys
|
344 |
import os
|
345 |
import traceback
|
346 |
-
|
347 |
print("--- Diagnostic Info from Blender Python ---")
|
348 |
print(f"Python Executable: {{sys.executable}}")
|
349 |
-
print(f"Current Working Directory (inside script): {{os.getcwd()}}")
|
350 |
print("sys.path:")
|
351 |
-
for p in sys.path:
|
352 |
-
print(f" {{p}}")
|
353 |
print("\\nPYTHONPATH Environment Variable (as seen by script):")
|
354 |
print(os.environ.get('PYTHONPATH', 'PYTHONPATH not set or empty'))
|
|
|
|
|
|
|
|
|
355 |
print("\\n--- Attempting Imports ---")
|
356 |
try:
|
357 |
import bpy
|
@@ -362,9 +343,7 @@ except ImportError as e:
|
|
362 |
except Exception as e:
|
363 |
print(f"FAILED to import 'bpy' with other error: {{e}}")
|
364 |
traceback.print_exc()
|
365 |
-
|
366 |
try:
|
367 |
-
# Check if CWD is correct and src is importable
|
368 |
print("\\nChecking for 'src' in CWD (should be UniRig repo root):")
|
369 |
if os.path.isdir('src'):
|
370 |
print(" 'src' directory FOUND in CWD.")
|
@@ -374,10 +353,8 @@ try:
|
|
374 |
print(" WARNING: 'src/__init__.py' NOT FOUND. 'src' may not be treated as a package.")
|
375 |
else:
|
376 |
print(" 'src' directory NOT FOUND in CWD.")
|
377 |
-
|
378 |
-
# Attempt the import that failed previously
|
379 |
print("\\nAttempting: from src.inference.download import download")
|
380 |
-
from src.inference.download import download
|
381 |
print("SUCCESS: 'from src.inference.download import download' worked.")
|
382 |
except ImportError as e:
|
383 |
print(f"FAILED: 'from src.inference.download import download': {{e}}")
|
@@ -387,95 +364,73 @@ except Exception as e:
|
|
387 |
traceback.print_exc()
|
388 |
print("--- End Diagnostic Info ---")
|
389 |
"""
|
390 |
-
diagnostic_script_path = os.path.join(processing_temp_dir, "env_diagnostic_test.py")
|
391 |
-
with open(diagnostic_script_path, "w") as f:
|
392 |
-
f.write(diagnostic_script_content)
|
393 |
-
|
394 |
try:
|
395 |
-
# Run the diagnostic script using the *exact same* command structure
|
396 |
-
# Pass an empty list for script_args as the script takes none
|
397 |
run_unirig_command(diagnostic_script_path, [], "Blender Env Diagnostic")
|
398 |
print("--- Finished Blender Python Environment Diagnostic Test ---\n")
|
399 |
except Exception as e_diag:
|
400 |
-
# If the diagnostic fails, raise an error - no point continuing
|
401 |
print(f"ERROR during diagnostic test execution: {e_diag}")
|
402 |
-
# Clean up the script file before raising
|
403 |
if os.path.exists(diagnostic_script_path): os.remove(diagnostic_script_path)
|
404 |
-
raise gr.Error(f"Blender environment diagnostic failed. Cannot proceed. Check logs
|
405 |
finally:
|
406 |
-
# Ensure cleanup even if run_unirig_command doesn't raise but has issues
|
407 |
if os.path.exists(diagnostic_script_path): os.remove(diagnostic_script_path)
|
408 |
-
# --- END DIAGNOSTIC STEP ---
|
409 |
|
410 |
-
|
411 |
-
# --- Determine Device for UniRig ---
|
412 |
unirig_device_arg = "device=cpu"
|
413 |
if DEVICE.type == 'cuda':
|
414 |
unirig_device_arg = "device=cuda:0"
|
415 |
print(f"UniRig steps will attempt to use device argument: {unirig_device_arg}")
|
416 |
|
417 |
-
# --- Execute UniRig Steps ---
|
418 |
-
# Step 1: Skeleton Prediction
|
419 |
print("\nStarting Step 1: Predicting Skeleton...")
|
420 |
-
# Arguments for run.py using Hydra format (key=value)
|
421 |
skeleton_args = [
|
422 |
-
"--config-name=skeleton_config",
|
423 |
-
"with", # Hydra keyword for overrides
|
424 |
f"input={abs_input_glb_path}",
|
425 |
f"output={abs_skeleton_output_path}",
|
426 |
unirig_device_arg
|
427 |
-
# Add other args like "verbose=True" if needed for debugging
|
428 |
]
|
429 |
run_unirig_command(unirig_script_to_run, skeleton_args, "Skeleton Prediction")
|
430 |
if not os.path.exists(abs_skeleton_output_path):
|
431 |
-
raise gr.Error("Skeleton prediction failed. Output file not created.
|
432 |
print("Step 1: Skeleton Prediction completed.")
|
433 |
|
434 |
-
# Step 2: Skinning Weight Prediction
|
435 |
print("\nStarting Step 2: Predicting Skinning Weights...")
|
436 |
skin_args = [
|
437 |
-
"--config-name=skin_config",
|
438 |
-
"
|
439 |
-
f"input={abs_skeleton_output_path}", # Input is the predicted skeleton
|
440 |
f"output={abs_skin_output_path}",
|
441 |
unirig_device_arg
|
442 |
]
|
443 |
run_unirig_command(unirig_script_to_run, skin_args, "Skinning Prediction")
|
444 |
if not os.path.exists(abs_skin_output_path):
|
445 |
-
raise gr.Error("Skinning prediction failed. Output file not created.
|
446 |
print("Step 2: Skinning Prediction completed.")
|
447 |
|
448 |
-
# Step 3: Merge Results
|
449 |
print("\nStarting Step 3: Merging Results...")
|
450 |
merge_args = [
|
451 |
-
"--config-name=merge_config",
|
452 |
-
"
|
453 |
-
f"
|
454 |
-
f"target_path={abs_input_glb_path}", # With original mesh geometry/attrs
|
455 |
f"output_path={abs_final_rigged_glb_path}",
|
456 |
-
"mode=skin",
|
457 |
unirig_device_arg
|
458 |
]
|
459 |
run_unirig_command(unirig_script_to_run, merge_args, "Merging Results")
|
460 |
if not os.path.exists(abs_final_rigged_glb_path):
|
461 |
-
raise gr.Error("Merging process failed. Final rigged GLB file not created.
|
462 |
print("Step 3: Merging completed.")
|
463 |
|
464 |
-
# --- Return Result ---
|
465 |
print(f"Successfully generated rigged model: {abs_final_rigged_glb_path}")
|
466 |
-
return gr.update(value=abs_final_rigged_glb_path)
|
467 |
-
|
468 |
except gr.Error as e:
|
469 |
print(f"A Gradio Error occurred: {e}")
|
470 |
if os.path.exists(processing_temp_dir): shutil.rmtree(processing_temp_dir)
|
471 |
-
raise e
|
472 |
except Exception as e:
|
473 |
print(f"An unexpected error occurred in rig_glb_mesh_multistep: {e}")
|
474 |
import traceback; traceback.print_exc()
|
475 |
if os.path.exists(processing_temp_dir): shutil.rmtree(processing_temp_dir)
|
476 |
raise gr.Error(f"An unexpected error occurred: {str(e)[:500]}. Check logs for details.")
|
477 |
finally:
|
478 |
-
# General cleanup for the temp dir if it still exists (e.g., if no exception occurred)
|
479 |
if os.path.exists(processing_temp_dir):
|
480 |
try:
|
481 |
shutil.rmtree(processing_temp_dir)
|
@@ -483,8 +438,6 @@ print("--- End Diagnostic Info ---")
|
|
483 |
except Exception as cleanup_e:
|
484 |
print(f"Error cleaning up temp dir {processing_temp_dir}: {cleanup_e}")
|
485 |
|
486 |
-
|
487 |
-
# --- Gradio Interface Definition ---
|
488 |
theme = gr.themes.Soft(
|
489 |
primary_hue=gr.themes.colors.sky,
|
490 |
secondary_hue=gr.themes.colors.blue,
|
@@ -492,23 +445,19 @@ theme = gr.themes.Soft(
|
|
492 |
font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
|
493 |
)
|
494 |
|
495 |
-
# Perform critical checks before defining the interface
|
496 |
startup_error_message = None
|
497 |
if not blender_executable_to_use:
|
498 |
-
startup_error_message = (f"CRITICAL STARTUP ERROR: Blender executable could not be located or setup failed. Check logs.")
|
499 |
elif not unirig_repo_ok:
|
500 |
startup_error_message = (f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}.")
|
501 |
elif not unirig_run_py_ok:
|
502 |
startup_error_message = (f"CRITICAL STARTUP ERROR: UniRig run.py not found at {UNIRIG_RUN_PY}.")
|
503 |
|
504 |
-
|
505 |
if startup_error_message:
|
506 |
-
# If critical error, display error message instead of interface
|
507 |
print(startup_error_message)
|
508 |
with gr.Blocks(theme=theme) as iface:
|
509 |
gr.Markdown(f"# Application Startup Error\n\n{startup_error_message}\n\nPlease check the Space logs for more details.")
|
510 |
else:
|
511 |
-
# Build the normal interface if essential checks pass
|
512 |
with gr.Blocks(theme=theme) as iface:
|
513 |
gr.Markdown(
|
514 |
f"""
|
@@ -516,6 +465,7 @@ else:
|
|
516 |
Upload a 3D mesh in `.glb` format. This application uses UniRig via Blender's Python interface to predict skeleton and skinning weights.
|
517 |
* Running main app on Python `{sys.version.split()[0]}`, UniRig steps use Blender's Python `{BLENDER_PYTHON_VERSION}`.
|
518 |
* Utilizing device: **{DEVICE.type.upper()}** (via ZeroGPU if available).
|
|
|
519 |
* UniRig Source: [https://github.com/VAST-AI-Research/UniRig](https://github.com/VAST-AI-Research/UniRig)
|
520 |
"""
|
521 |
)
|
@@ -523,7 +473,7 @@ else:
|
|
523 |
with gr.Column(scale=1):
|
524 |
input_model = gr.File(
|
525 |
label="Upload .glb Mesh File",
|
526 |
-
type="filepath",
|
527 |
file_types=[".glb"]
|
528 |
)
|
529 |
submit_button = gr.Button("Rig Model", variant="primary")
|
@@ -532,22 +482,16 @@ else:
|
|
532 |
label="Rigged 3D Model (.glb)",
|
533 |
clear_color=[0.8, 0.8, 0.8, 1.0],
|
534 |
)
|
535 |
-
|
536 |
-
# Connect button click to the processing function
|
537 |
submit_button.click(
|
538 |
fn=rig_glb_mesh_multistep,
|
539 |
inputs=[input_model],
|
540 |
outputs=[output_model]
|
541 |
)
|
542 |
|
543 |
-
# --- Launch the Application ---
|
544 |
if __name__ == "__main__":
|
545 |
-
# Check if the interface was successfully created
|
546 |
if 'iface' in locals():
|
547 |
print("Launching Gradio interface...")
|
548 |
-
|
549 |
-
iface.launch(share=False, ssr_mode=False) # Disable SSR
|
550 |
else:
|
551 |
-
# This indicates a startup error occurred before interface creation
|
552 |
print("ERROR: Gradio interface could not be created due to startup errors. Check logs above.")
|
553 |
|
|
|
9 |
from typing import Any, Dict, List, Union # Added Union
|
10 |
|
11 |
# --- Configuration ---
|
12 |
+
APP_ROOT_DIR = os.path.abspath(os.path.dirname(__file__)) # Should be /home/user/app
|
|
|
|
|
|
|
13 |
|
14 |
+
UNIRIG_REPO_DIR = os.path.join(APP_ROOT_DIR, "UniRig")
|
15 |
+
|
16 |
+
# ** MODIFIED FOR LOCAL BLENDER INSTALLATION **
|
17 |
+
BLENDER_VERSION_NAME = "blender-4.2.0-linux-x64"
|
18 |
+
BLENDER_LOCAL_INSTALL_BASE_DIR = os.path.join(APP_ROOT_DIR, "blender_installation")
|
19 |
+
BLENDER_INSTALL_DIR = os.path.join(BLENDER_LOCAL_INSTALL_BASE_DIR, BLENDER_VERSION_NAME)
|
20 |
+
|
21 |
+
BLENDER_PYTHON_VERSION_DIR = "4.2" # From Blender's internal structure
|
22 |
+
BLENDER_PYTHON_VERSION = "python3.11" # UniRig requirement
|
23 |
+
|
24 |
+
# Construct paths based on the new local installation
|
25 |
BLENDER_PYTHON_DIR = os.path.join(BLENDER_INSTALL_DIR, BLENDER_PYTHON_VERSION_DIR, "python")
|
26 |
BLENDER_PYTHON_BIN_DIR = os.path.join(BLENDER_PYTHON_DIR, "bin")
|
|
|
27 |
BLENDER_EXEC = os.path.join(BLENDER_INSTALL_DIR, "blender")
|
|
|
28 |
|
29 |
+
# Symlink might not be created or might be in a local bin, app.py prioritizes BLENDER_EXEC
|
30 |
+
# Keeping this definition for fallback, but it's less likely to be the primary method now.
|
31 |
+
BLENDER_EXEC_SYMLINK = "/usr/local/bin/blender" # Standard system symlink path (may not exist/be used)
|
32 |
+
LOCAL_BIN_DIR = os.path.join(APP_ROOT_DIR, "local_bin") # Potential local symlink location
|
33 |
+
BLENDER_EXEC_LOCAL_SYMLINK = os.path.join(LOCAL_BIN_DIR, "blender")
|
34 |
|
35 |
|
36 |
+
SETUP_SCRIPT = os.path.join(APP_ROOT_DIR, "setup_blender.sh")
|
37 |
+
SETUP_SCRIPT_TIMEOUT = 800 # Increased timeout
|
38 |
+
|
39 |
# --- Initial Checks ---
|
40 |
print("--- Environment Checks ---")
|
41 |
+
print(f"APP_ROOT_DIR: {APP_ROOT_DIR}")
|
42 |
+
print(f"Expected Blender Install Dir: {BLENDER_INSTALL_DIR}")
|
43 |
+
print(f"Expected Blender Executable: {BLENDER_EXEC}")
|
44 |
+
|
45 |
+
blender_executable_to_use = None
|
46 |
if os.path.exists(BLENDER_EXEC):
|
47 |
+
print(f"Blender executable found at direct local path: {BLENDER_EXEC}")
|
48 |
blender_executable_to_use = BLENDER_EXEC
|
49 |
+
elif os.path.exists(BLENDER_EXEC_LOCAL_SYMLINK):
|
50 |
+
print(f"Blender executable found via local symlink: {BLENDER_EXEC_LOCAL_SYMLINK}")
|
51 |
+
blender_executable_to_use = BLENDER_EXEC_LOCAL_SYMLINK
|
52 |
+
elif os.path.exists(BLENDER_EXEC_SYMLINK): # Fallback to system symlink (less likely)
|
53 |
+
print(f"Blender executable found via system symlink: {BLENDER_EXEC_SYMLINK}")
|
54 |
blender_executable_to_use = BLENDER_EXEC_SYMLINK
|
55 |
else:
|
56 |
+
print(f"Blender executable not found at {BLENDER_EXEC}, {BLENDER_EXEC_LOCAL_SYMLINK}, or {BLENDER_EXEC_SYMLINK}. Running setup script...")
|
57 |
if os.path.exists(SETUP_SCRIPT):
|
58 |
try:
|
59 |
# Run setup script if Blender not found
|
60 |
+
# Ensure setup_blender.sh is executable (chmod +x setup_blender.sh in Dockerfile or Space setup)
|
61 |
setup_result = subprocess.run(
|
62 |
["bash", SETUP_SCRIPT],
|
63 |
check=True,
|
64 |
capture_output=True,
|
65 |
text=True,
|
66 |
+
timeout=SETUP_SCRIPT_TIMEOUT
|
67 |
)
|
68 |
print("Setup script executed successfully.")
|
69 |
print(f"Setup STDOUT:\n{setup_result.stdout}")
|
70 |
if setup_result.stderr: print(f"Setup STDERR:\n{setup_result.stderr}")
|
71 |
|
72 |
+
# Re-check for executable after running setup (prioritize direct local path)
|
73 |
if os.path.exists(BLENDER_EXEC):
|
74 |
blender_executable_to_use = BLENDER_EXEC
|
75 |
+
print(f"Blender executable now found at direct local path: {BLENDER_EXEC}")
|
76 |
+
elif os.path.exists(BLENDER_EXEC_LOCAL_SYMLINK):
|
77 |
+
blender_executable_to_use = BLENDER_EXEC_LOCAL_SYMLINK
|
78 |
+
print(f"Blender executable now found via local symlink: {BLENDER_EXEC_LOCAL_SYMLINK}")
|
79 |
elif os.path.exists(BLENDER_EXEC_SYMLINK):
|
80 |
blender_executable_to_use = BLENDER_EXEC_SYMLINK
|
81 |
+
print(f"Blender executable now found via system symlink: {BLENDER_EXEC_SYMLINK}")
|
82 |
|
83 |
if not blender_executable_to_use:
|
84 |
+
raise RuntimeError(f"Setup script ran but Blender executable still not found at {BLENDER_EXEC} or other checked paths.")
|
|
|
|
|
85 |
except subprocess.TimeoutExpired:
|
86 |
print(f"ERROR: Setup script timed out after {SETUP_SCRIPT_TIMEOUT} seconds: {SETUP_SCRIPT}")
|
87 |
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.")
|
88 |
except subprocess.CalledProcessError as e:
|
89 |
print(f"ERROR running setup script: {SETUP_SCRIPT}\nStderr: {e.stderr}")
|
|
|
90 |
raise gr.Error(f"Failed to execute setup script. Check logs. Stderr: {e.stderr[-500:]}")
|
91 |
except Exception as e:
|
|
|
92 |
raise gr.Error(f"Unexpected error running setup script '{SETUP_SCRIPT}': {e}")
|
93 |
else:
|
|
|
94 |
raise gr.Error(f"Blender executable not found and setup script missing: {SETUP_SCRIPT}")
|
95 |
|
96 |
# Verify bpy import using the found Blender executable
|
|
|
99 |
try:
|
100 |
print("Testing bpy import via Blender...")
|
101 |
test_script_content = "import bpy; print('bpy imported successfully')"
|
|
|
102 |
test_result = subprocess.run(
|
103 |
[blender_executable_to_use, "--background", "--python-expr", test_script_content],
|
104 |
+
capture_output=True, text=True, check=True, timeout=30
|
105 |
)
|
106 |
if "bpy imported successfully" in test_result.stdout:
|
107 |
print("Successfully imported 'bpy' using Blender executable.")
|
108 |
bpy_import_ok = True
|
109 |
else:
|
|
|
110 |
print(f"WARNING: 'bpy' import test via Blender returned unexpected output:\nSTDOUT:{test_result.stdout}\nSTDERR:{test_result.stderr}")
|
|
|
111 |
except subprocess.TimeoutExpired:
|
112 |
print("WARNING: 'bpy' import test via Blender timed out.")
|
113 |
except subprocess.CalledProcessError as e:
|
|
|
114 |
print(f"WARNING: Failed to import 'bpy' using Blender executable:\nSTDOUT:{e.stdout}\nSTDERR:{e.stderr}")
|
115 |
except Exception as e:
|
|
|
116 |
print(f"WARNING: Unexpected error during 'bpy' import test: {e}")
|
117 |
else:
|
|
|
118 |
print("WARNING: Cannot test bpy import as Blender executable was not found.")
|
119 |
|
|
|
120 |
# Check for UniRig repository and run.py
|
121 |
unirig_repo_ok = False
|
122 |
unirig_run_py_ok = False
|
123 |
UNIRIG_RUN_PY = os.path.join(UNIRIG_REPO_DIR, "run.py")
|
124 |
if not os.path.isdir(UNIRIG_REPO_DIR):
|
|
|
125 |
raise gr.Error(f"UniRig repository missing at: {UNIRIG_REPO_DIR}. Ensure it's cloned correctly.")
|
126 |
else:
|
127 |
print(f"UniRig repository found at: {UNIRIG_REPO_DIR}")
|
128 |
unirig_repo_ok = True
|
|
|
129 |
if not os.path.exists(UNIRIG_RUN_PY):
|
|
|
130 |
raise gr.Error(f"UniRig's run.py not found at {UNIRIG_RUN_PY}. Check UniRig clone.")
|
131 |
else:
|
132 |
unirig_run_py_ok = True
|
133 |
|
|
|
|
|
134 |
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
135 |
print(f"Gradio environment using device: {DEVICE}")
|
136 |
if DEVICE.type == 'cuda':
|
|
|
141 |
print(f"Could not get Gradio CUDA device details: {e}")
|
142 |
else:
|
143 |
print("Warning: Gradio environment CUDA not available.")
|
|
|
144 |
print("--- End Environment Checks ---")
|
145 |
|
|
|
146 |
def patch_asset_py():
|
147 |
"""Temporary patch to fix type hinting error in UniRig's asset.py"""
|
148 |
asset_py_path = os.path.join(UNIRIG_REPO_DIR, "src", "data", "asset.py")
|
|
|
150 |
if not os.path.exists(asset_py_path):
|
151 |
print(f"Warning: asset.py not found at {asset_py_path}, skipping patch.")
|
152 |
return
|
|
|
153 |
with open(asset_py_path, "r") as f: content = f.read()
|
154 |
problematic_line = "meta: Union[Dict[str, ...], None]=None"
|
155 |
corrected_line = "meta: Union[Dict[str, Any], None]=None"
|
156 |
typing_import = "from typing import Any"
|
|
|
157 |
if corrected_line in content:
|
158 |
print("Patch already applied to asset.py"); return
|
159 |
if problematic_line not in content:
|
160 |
print("Problematic line not found in asset.py, patch might be unnecessary or file changed."); return
|
|
|
161 |
print("Applying patch to asset.py...")
|
162 |
content = content.replace(problematic_line, corrected_line)
|
163 |
if typing_import not in content:
|
|
|
167 |
content = f"{typing_import}\n{content}"
|
168 |
with open(asset_py_path, "w") as f: f.write(content)
|
169 |
print("Successfully patched asset.py")
|
|
|
170 |
except Exception as e:
|
171 |
print(f"ERROR: Failed to patch asset.py: {e}. Proceeding cautiously.")
|
|
|
172 |
|
|
|
|
|
173 |
@spaces.GPU
|
174 |
def run_unirig_command(python_script_path: str, script_args: List[str], step_name: str):
|
175 |
"""
|
176 |
Runs a specific UniRig PYTHON script (.py) using the Blender executable
|
177 |
in background mode (`blender --background --python script.py -- args`).
|
|
|
|
|
|
|
|
|
|
|
178 |
"""
|
179 |
if not blender_executable_to_use:
|
180 |
raise gr.Error("Blender executable path could not be determined. Cannot run UniRig step.")
|
181 |
|
|
|
182 |
process_env = os.environ.copy()
|
|
|
|
|
183 |
unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, 'src')
|
184 |
+
pythonpath_parts = [UNIRIG_REPO_DIR, unirig_src_dir]
|
185 |
existing_pythonpath = process_env.get('PYTHONPATH', '')
|
186 |
if existing_pythonpath:
|
187 |
pythonpath_parts.append(existing_pythonpath)
|
188 |
process_env["PYTHONPATH"] = os.pathsep.join(filter(None, pythonpath_parts))
|
189 |
print(f"Subprocess PYTHONPATH: {process_env['PYTHONPATH']}")
|
190 |
|
191 |
+
# LD_LIBRARY_PATH: Inherit system path and add Blender's libraries (now from local install)
|
192 |
+
blender_main_lib_path = os.path.join(BLENDER_INSTALL_DIR, "lib") # Path to Blender's own .so files
|
193 |
+
blender_python_lib_path = os.path.join(BLENDER_PYTHON_DIR, "lib") # Path to Blender's Python's site-packages etc.
|
194 |
+
|
195 |
ld_path_parts = []
|
196 |
if os.path.exists(blender_main_lib_path): ld_path_parts.append(blender_main_lib_path)
|
197 |
if os.path.exists(blender_python_lib_path): ld_path_parts.append(blender_python_lib_path)
|
198 |
+
|
199 |
+
# Add local bin to PATH if it exists and contains the symlink, for consistency
|
200 |
+
# Also ensure Blender's own script/bin directories are implicitly available if needed by Blender.
|
201 |
+
# The direct call to blender_executable_to_use should handle most cases.
|
202 |
+
if os.path.isdir(LOCAL_BIN_DIR):
|
203 |
+
process_env["PATH"] = f"{LOCAL_BIN_DIR}{os.pathsep}{process_env.get('PATH', '')}"
|
204 |
+
|
205 |
+
|
206 |
existing_ld_path = process_env.get('LD_LIBRARY_PATH', '')
|
207 |
if existing_ld_path: ld_path_parts.append(existing_ld_path)
|
208 |
if ld_path_parts:
|
209 |
process_env["LD_LIBRARY_PATH"] = os.pathsep.join(filter(None, ld_path_parts))
|
210 |
print(f"Subprocess LD_LIBRARY_PATH: {process_env.get('LD_LIBRARY_PATH', 'Not set')}")
|
211 |
+
print(f"Subprocess PATH: {process_env.get('PATH', 'Not set')}")
|
212 |
|
213 |
|
|
|
|
|
214 |
cmd = [
|
215 |
blender_executable_to_use,
|
216 |
"--background",
|
217 |
"--python", python_script_path,
|
218 |
+
"--"
|
219 |
] + script_args
|
|
|
220 |
print(f"\n--- Running UniRig Step: {step_name} ---")
|
221 |
+
print(f"Command: {' '.join(cmd)}")
|
|
|
222 |
try:
|
|
|
|
|
223 |
result = subprocess.run(
|
224 |
cmd,
|
225 |
+
cwd=UNIRIG_REPO_DIR,
|
226 |
capture_output=True,
|
227 |
text=True,
|
228 |
+
check=True,
|
229 |
+
env=process_env,
|
230 |
+
timeout=1800
|
231 |
)
|
232 |
print(f"{step_name} STDOUT:\n{result.stdout}")
|
233 |
if result.stderr:
|
234 |
print(f"{step_name} STDERR (Info/Warnings):\n{result.stderr}")
|
|
|
235 |
stderr_lower = result.stderr.lower()
|
236 |
if "error" in stderr_lower or "failed" in stderr_lower or "traceback" in stderr_lower:
|
237 |
print(f"WARNING: Potential error messages found in STDERR for {step_name} despite success exit code.")
|
|
|
238 |
except subprocess.TimeoutExpired:
|
239 |
print(f"ERROR: {step_name} timed out after 30 minutes.")
|
240 |
raise gr.Error(f"Processing step '{step_name}' timed out. Please try with a simpler model or check logs.")
|
|
|
246 |
print(f"--- {step_name} STDERR ---:\n{e.stderr}")
|
247 |
error_summary = e.stderr.strip().splitlines()
|
248 |
last_lines = "\n".join(error_summary[-15:]) if error_summary else "No stderr output."
|
|
|
249 |
specific_error = "Unknown error."
|
|
|
250 |
if "ModuleNotFoundError: No module named 'src'" in e.stderr:
|
251 |
specific_error = "UniRig script failed to import its own 'src' module. Check PYTHONPATH and CWD for the subprocess. See diagnostic info above."
|
252 |
elif "ModuleNotFoundError: No module named 'bpy'" in e.stderr:
|
|
|
259 |
specific_error = "CUDA out of memory. Try a smaller model or a Space with more GPU RAM."
|
260 |
elif "hydra.errors.ConfigCompositionException" in e.stderr:
|
261 |
specific_error = f"Hydra configuration error. Check the arguments passed to run.py: {script_args}. Details: {last_lines}"
|
262 |
+
elif "Error: Cannot read file" in e.stderr:
|
263 |
specific_error = f"Blender could not read an input/output file. Check paths. Details: {last_lines}"
|
264 |
+
elif "Error:" in e.stderr:
|
265 |
specific_error = f"Blender reported an error. Details: {last_lines}"
|
266 |
else:
|
267 |
specific_error = f"Check logs. Last error lines:\n{last_lines}"
|
268 |
raise gr.Error(f"Error in UniRig '{step_name}'. {specific_error}")
|
|
|
269 |
except FileNotFoundError:
|
270 |
print(f"ERROR: Could not find Blender executable '{blender_executable_to_use}' or script '{python_script_path}' for {step_name}.")
|
271 |
raise gr.Error(f"Setup error for UniRig '{step_name}'. Blender or Python script not found.")
|
|
|
274 |
import traceback
|
275 |
traceback.print_exc()
|
276 |
raise gr.Error(f"Unexpected Python error during '{step_name}' execution: {str(e_general)[:500]}")
|
|
|
277 |
print(f"--- Finished UniRig Step: {step_name} ---")
|
278 |
|
|
|
279 |
@spaces.GPU
|
280 |
def rig_glb_mesh_multistep(input_glb_file_obj):
|
281 |
"""
|
282 |
Main Gradio function to rig a GLB mesh using UniRig's multi-step process.
|
283 |
"""
|
|
|
284 |
if not blender_executable_to_use:
|
285 |
gr.Warning("System not ready: Blender executable not found. Please wait or check logs.")
|
286 |
+
return None, None
|
287 |
if not unirig_repo_ok or not unirig_run_py_ok:
|
288 |
gr.Warning("System not ready: UniRig repository or run.py script not found. Check setup.")
|
289 |
+
return None, None
|
290 |
if not bpy_import_ok:
|
291 |
gr.Warning("System warning: Initial 'bpy' import test failed. Attempting to proceed, but errors may occur.")
|
292 |
|
293 |
try:
|
294 |
+
patch_asset_py()
|
295 |
except Exception as e:
|
296 |
print(f"Ignoring patch error: {e}")
|
297 |
|
|
|
298 |
if input_glb_file_obj is None:
|
299 |
gr.Info("Please upload a .glb file first.")
|
300 |
return None
|
301 |
+
input_glb_path = input_glb_file_obj # Gradio File component with type="filepath" returns a string path
|
302 |
print(f"Input GLB path received: {input_glb_path}")
|
303 |
+
|
304 |
if not isinstance(input_glb_path, str) or not os.path.exists(input_glb_path):
|
305 |
+
raise gr.Error(f"Invalid input file path received or file does not exist: {input_glb_path}")
|
306 |
if not input_glb_path.lower().endswith(".glb"):
|
307 |
raise gr.Error("Invalid file type. Please upload a .glb file.")
|
308 |
|
|
|
309 |
processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_")
|
310 |
print(f"Using temporary processing directory: {processing_temp_dir}")
|
|
|
311 |
try:
|
|
|
312 |
base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
|
313 |
+
# Ensure paths used by UniRig are absolute, especially if its CWD changes
|
314 |
abs_input_glb_path = os.path.abspath(input_glb_path)
|
315 |
abs_skeleton_output_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
|
316 |
abs_skin_output_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx")
|
317 |
abs_final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")
|
|
|
318 |
unirig_script_to_run = UNIRIG_RUN_PY
|
319 |
|
|
|
320 |
print("\n--- Running Blender Python Environment Diagnostic Test ---")
|
321 |
diagnostic_script_content = f"""
|
322 |
import sys
|
323 |
import os
|
324 |
import traceback
|
|
|
325 |
print("--- Diagnostic Info from Blender Python ---")
|
326 |
print(f"Python Executable: {{sys.executable}}")
|
327 |
+
print(f"Current Working Directory (inside script): {{os.getcwd()}}")
|
328 |
print("sys.path:")
|
329 |
+
for p in sys.path: print(f" {{p}}")
|
|
|
330 |
print("\\nPYTHONPATH Environment Variable (as seen by script):")
|
331 |
print(os.environ.get('PYTHONPATH', 'PYTHONPATH not set or empty'))
|
332 |
+
print("\\nLD_LIBRARY_PATH Environment Variable (as seen by script):")
|
333 |
+
print(os.environ.get('LD_LIBRARY_PATH', 'LD_LIBRARY_PATH not set or empty'))
|
334 |
+
print("\\nPATH Environment Variable (as seen by script):")
|
335 |
+
print(os.environ.get('PATH', 'PATH not set or empty'))
|
336 |
print("\\n--- Attempting Imports ---")
|
337 |
try:
|
338 |
import bpy
|
|
|
343 |
except Exception as e:
|
344 |
print(f"FAILED to import 'bpy' with other error: {{e}}")
|
345 |
traceback.print_exc()
|
|
|
346 |
try:
|
|
|
347 |
print("\\nChecking for 'src' in CWD (should be UniRig repo root):")
|
348 |
if os.path.isdir('src'):
|
349 |
print(" 'src' directory FOUND in CWD.")
|
|
|
353 |
print(" WARNING: 'src/__init__.py' NOT FOUND. 'src' may not be treated as a package.")
|
354 |
else:
|
355 |
print(" 'src' directory NOT FOUND in CWD.")
|
|
|
|
|
356 |
print("\\nAttempting: from src.inference.download import download")
|
357 |
+
from src.inference.download import download # Example import
|
358 |
print("SUCCESS: 'from src.inference.download import download' worked.")
|
359 |
except ImportError as e:
|
360 |
print(f"FAILED: 'from src.inference.download import download': {{e}}")
|
|
|
364 |
traceback.print_exc()
|
365 |
print("--- End Diagnostic Info ---")
|
366 |
"""
|
367 |
+
diagnostic_script_path = os.path.join(processing_temp_dir, "env_diagnostic_test.py")
|
368 |
+
with open(diagnostic_script_path, "w") as f: f.write(diagnostic_script_content)
|
|
|
|
|
369 |
try:
|
|
|
|
|
370 |
run_unirig_command(diagnostic_script_path, [], "Blender Env Diagnostic")
|
371 |
print("--- Finished Blender Python Environment Diagnostic Test ---\n")
|
372 |
except Exception as e_diag:
|
|
|
373 |
print(f"ERROR during diagnostic test execution: {e_diag}")
|
|
|
374 |
if os.path.exists(diagnostic_script_path): os.remove(diagnostic_script_path)
|
375 |
+
raise gr.Error(f"Blender environment diagnostic failed. Cannot proceed. Check logs. Error: {str(e_diag)[:500]}")
|
376 |
finally:
|
|
|
377 |
if os.path.exists(diagnostic_script_path): os.remove(diagnostic_script_path)
|
|
|
378 |
|
|
|
|
|
379 |
unirig_device_arg = "device=cpu"
|
380 |
if DEVICE.type == 'cuda':
|
381 |
unirig_device_arg = "device=cuda:0"
|
382 |
print(f"UniRig steps will attempt to use device argument: {unirig_device_arg}")
|
383 |
|
|
|
|
|
384 |
print("\nStarting Step 1: Predicting Skeleton...")
|
|
|
385 |
skeleton_args = [
|
386 |
+
"--config-name=skeleton_config", "with",
|
|
|
387 |
f"input={abs_input_glb_path}",
|
388 |
f"output={abs_skeleton_output_path}",
|
389 |
unirig_device_arg
|
|
|
390 |
]
|
391 |
run_unirig_command(unirig_script_to_run, skeleton_args, "Skeleton Prediction")
|
392 |
if not os.path.exists(abs_skeleton_output_path):
|
393 |
+
raise gr.Error("Skeleton prediction failed. Output file not created.")
|
394 |
print("Step 1: Skeleton Prediction completed.")
|
395 |
|
|
|
396 |
print("\nStarting Step 2: Predicting Skinning Weights...")
|
397 |
skin_args = [
|
398 |
+
"--config-name=skin_config", "with",
|
399 |
+
f"input={abs_skeleton_output_path}",
|
|
|
400 |
f"output={abs_skin_output_path}",
|
401 |
unirig_device_arg
|
402 |
]
|
403 |
run_unirig_command(unirig_script_to_run, skin_args, "Skinning Prediction")
|
404 |
if not os.path.exists(abs_skin_output_path):
|
405 |
+
raise gr.Error("Skinning prediction failed. Output file not created.")
|
406 |
print("Step 2: Skinning Prediction completed.")
|
407 |
|
|
|
408 |
print("\nStarting Step 3: Merging Results...")
|
409 |
merge_args = [
|
410 |
+
"--config-name=merge_config", "with",
|
411 |
+
f"source_path={abs_skin_output_path}",
|
412 |
+
f"target_path={abs_input_glb_path}",
|
|
|
413 |
f"output_path={abs_final_rigged_glb_path}",
|
414 |
+
"mode=skin",
|
415 |
unirig_device_arg
|
416 |
]
|
417 |
run_unirig_command(unirig_script_to_run, merge_args, "Merging Results")
|
418 |
if not os.path.exists(abs_final_rigged_glb_path):
|
419 |
+
raise gr.Error("Merging process failed. Final rigged GLB file not created.")
|
420 |
print("Step 3: Merging completed.")
|
421 |
|
|
|
422 |
print(f"Successfully generated rigged model: {abs_final_rigged_glb_path}")
|
423 |
+
return gr.update(value=abs_final_rigged_glb_path)
|
|
|
424 |
except gr.Error as e:
|
425 |
print(f"A Gradio Error occurred: {e}")
|
426 |
if os.path.exists(processing_temp_dir): shutil.rmtree(processing_temp_dir)
|
427 |
+
raise e # Re-raise Gradio errors to be displayed in UI
|
428 |
except Exception as e:
|
429 |
print(f"An unexpected error occurred in rig_glb_mesh_multistep: {e}")
|
430 |
import traceback; traceback.print_exc()
|
431 |
if os.path.exists(processing_temp_dir): shutil.rmtree(processing_temp_dir)
|
432 |
raise gr.Error(f"An unexpected error occurred: {str(e)[:500]}. Check logs for details.")
|
433 |
finally:
|
|
|
434 |
if os.path.exists(processing_temp_dir):
|
435 |
try:
|
436 |
shutil.rmtree(processing_temp_dir)
|
|
|
438 |
except Exception as cleanup_e:
|
439 |
print(f"Error cleaning up temp dir {processing_temp_dir}: {cleanup_e}")
|
440 |
|
|
|
|
|
441 |
theme = gr.themes.Soft(
|
442 |
primary_hue=gr.themes.colors.sky,
|
443 |
secondary_hue=gr.themes.colors.blue,
|
|
|
445 |
font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
|
446 |
)
|
447 |
|
|
|
448 |
startup_error_message = None
|
449 |
if not blender_executable_to_use:
|
450 |
+
startup_error_message = (f"CRITICAL STARTUP ERROR: Blender executable could not be located or setup failed. Check logs. Expected at {BLENDER_EXEC}")
|
451 |
elif not unirig_repo_ok:
|
452 |
startup_error_message = (f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}.")
|
453 |
elif not unirig_run_py_ok:
|
454 |
startup_error_message = (f"CRITICAL STARTUP ERROR: UniRig run.py not found at {UNIRIG_RUN_PY}.")
|
455 |
|
|
|
456 |
if startup_error_message:
|
|
|
457 |
print(startup_error_message)
|
458 |
with gr.Blocks(theme=theme) as iface:
|
459 |
gr.Markdown(f"# Application Startup Error\n\n{startup_error_message}\n\nPlease check the Space logs for more details.")
|
460 |
else:
|
|
|
461 |
with gr.Blocks(theme=theme) as iface:
|
462 |
gr.Markdown(
|
463 |
f"""
|
|
|
465 |
Upload a 3D mesh in `.glb` format. This application uses UniRig via Blender's Python interface to predict skeleton and skinning weights.
|
466 |
* Running main app on Python `{sys.version.split()[0]}`, UniRig steps use Blender's Python `{BLENDER_PYTHON_VERSION}`.
|
467 |
* Utilizing device: **{DEVICE.type.upper()}** (via ZeroGPU if available).
|
468 |
+
* Blender executable: `{blender_executable_to_use}`.
|
469 |
* UniRig Source: [https://github.com/VAST-AI-Research/UniRig](https://github.com/VAST-AI-Research/UniRig)
|
470 |
"""
|
471 |
)
|
|
|
473 |
with gr.Column(scale=1):
|
474 |
input_model = gr.File(
|
475 |
label="Upload .glb Mesh File",
|
476 |
+
type="filepath", # Returns a string path to a temporary copy
|
477 |
file_types=[".glb"]
|
478 |
)
|
479 |
submit_button = gr.Button("Rig Model", variant="primary")
|
|
|
482 |
label="Rigged 3D Model (.glb)",
|
483 |
clear_color=[0.8, 0.8, 0.8, 1.0],
|
484 |
)
|
|
|
|
|
485 |
submit_button.click(
|
486 |
fn=rig_glb_mesh_multistep,
|
487 |
inputs=[input_model],
|
488 |
outputs=[output_model]
|
489 |
)
|
490 |
|
|
|
491 |
if __name__ == "__main__":
|
|
|
492 |
if 'iface' in locals():
|
493 |
print("Launching Gradio interface...")
|
494 |
+
iface.launch(share=False, ssr_mode=False)
|
|
|
495 |
else:
|
|
|
496 |
print("ERROR: Gradio interface could not be created due to startup errors. Check logs above.")
|
497 |
|