jkorstad commited on
Commit
b5736ea
·
verified ·
1 Parent(s): ae0a4fe

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +97 -153
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
- UNIRIG_REPO_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "UniRig"))
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
- # Construct paths
 
 
 
 
 
 
 
 
 
 
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
- SETUP_SCRIPT = os.path.join(os.path.dirname(__file__), "setup_blender.sh")
25
- # ** INCREASED TIMEOUT FOR SETUP SCRIPT **
26
- SETUP_SCRIPT_TIMEOUT = 800
 
 
27
 
28
 
 
 
 
29
  # --- Initial Checks ---
30
  print("--- Environment Checks ---")
31
- blender_executable_to_use = None # Use the main blender executable
 
 
 
 
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(BLENDER_EXEC_SYMLINK):
36
- print(f"Blender executable found via symlink: {BLENDER_EXEC_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 # Use the increased 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
- # If still not found after setup, raise a clear error
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 # Add timeout
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] # Add both
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 (for CUDA) and add Blender's libraries.
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
- "--" # Separator between blender args and script args
224
  ] + script_args
225
-
226
  print(f"\n--- Running UniRig Step: {step_name} ---")
227
- print(f"Command: {' '.join(cmd)}") # Log the actual command
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, # CWD is set to UniRig repo root
235
  capture_output=True,
236
  text=True,
237
- check=True, # Raises CalledProcessError on non-zero exit codes
238
- env=process_env, # Pass the modified environment
239
- timeout=1800 # 30 minutes timeout for UniRig steps
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: # General file read error
276
  specific_error = f"Blender could not read an input/output file. Check paths. Details: {last_lines}"
277
- elif "Error:" in e.stderr: # Generic Blender error
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() # Apply patch if needed
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()}}") # Should be UNIRIG_REPO_DIR
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") # Use temp dir
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 above for details. Error: {str(e_diag)[:500]}")
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", # Hydra main 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. Check logs.")
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
- "with",
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. Check logs.")
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
- "with",
453
- f"source_path={abs_skin_output_path}", # Merge skinned skeleton
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", # Specify merge mode
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. Check logs.")
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) # Update Gradio output
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
- # Ensure share=False for security unless public link is needed
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