jkorstad commited on
Commit
25e014c
·
verified ·
1 Parent(s): e843bc1

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +53 -37
app.py CHANGED
@@ -25,6 +25,16 @@ if not os.path.exists("/usr/local/bin/blender"):
25
  else:
26
  print("Blender is already installed.")
27
 
 
 
 
 
 
 
 
 
 
 
28
  if not os.path.isdir(UNIRIG_REPO_DIR):
29
  print(f"ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. Please clone it there.")
30
  # Consider raising an error or displaying it in the UI if UniRig is critical for startup
@@ -37,13 +47,34 @@ if DEVICE.type == 'cuda':
37
  else:
38
  print("Warning: CUDA not available or not detected by PyTorch. UniRig performance will be severely impacted.")
39
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
40
  @spaces.GPU # Decorator for ZeroGPU
41
  def run_unirig_command(command_list, step_name):
42
  """
43
  Helper function to run UniRig commands (now expecting bash scripts) using subprocess.
44
  command_list: The full command and its arguments, e.g., ["bash", "script.sh", "--arg", "value"]
45
  """
46
- # The command_list is now expected to be the full command, e.g., starting with "bash"
47
  cmd = command_list
48
 
49
  print(f"Running {step_name}: {' '.join(cmd)}")
@@ -53,23 +84,15 @@ def run_unirig_command(command_list, step_name):
53
  # Determine the path to the 'src' directory within UniRig, where the 'unirig' package resides.
54
  unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, "src")
55
 
56
- # Explicitly add UNIRIG_REPO_DIR/src to PYTHONPATH for the subprocess.
57
- # The bash scripts will internally call Python, which needs to find the 'unirig' package.
58
- # Also, keep UNIRIG_REPO_DIR itself in case some scripts or modules there are run directly
59
- # or expect the project root to be in PYTHONPATH.
60
  existing_pythonpath = process_env.get('PYTHONPATH', '')
61
- new_pythonpath_parts = [unirig_src_dir, UNIRIG_REPO_DIR] # UniRig/src first, then UniRig/
62
  if existing_pythonpath:
63
- # Prepend our paths to existing PYTHONPATH
64
  new_pythonpath_parts.extend(existing_pythonpath.split(os.pathsep))
65
-
66
- process_env["PYTHONPATH"] = os.pathsep.join(filter(None, new_pythonpath_parts)) # filter(None,...) handles empty existing_pythonpath
67
  print(f"Set PYTHONPATH for subprocess: {process_env['PYTHONPATH']}")
68
 
69
  try:
70
- # Execute the command from the UniRig directory (UNIRIG_REPO_DIR)
71
- # This is crucial for the bash scripts to find their relative paths (e.g., to Python scripts)
72
- # and for any underlying Python/Hydra calls to find configurations (e.g., in UniRig/configs/)
73
  result = subprocess.run(cmd, cwd=UNIRIG_REPO_DIR, capture_output=True, text=True, check=True, env=process_env)
74
  print(f"{step_name} STDOUT:\n{result.stdout}")
75
  if result.stderr:
@@ -80,11 +103,9 @@ def run_unirig_command(command_list, step_name):
80
  print(f"Return code: {e.returncode}")
81
  print(f"Stdout: {e.stdout}")
82
  print(f"Stderr: {e.stderr}")
83
- # Provide a more user-friendly error, potentially masking long tracebacks
84
- error_summary = e.stderr.splitlines()[-5:] # Last 5 lines of stderr
85
  raise gr.Error(f"Error in UniRig {step_name}. Details: {' '.join(error_summary)}")
86
  except FileNotFoundError:
87
- # This error means the executable (e.g., "bash" or the script itself) was not found.
88
  print(f"ERROR: Could not find executable or script for {step_name}. Command: {' '.join(cmd)}. Is UniRig cloned correctly and 'bash' available?")
89
  raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs, UniRig directory structure, and script paths.")
90
  except Exception as e_general:
@@ -98,61 +119,56 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
98
  rigs it using the new UniRig multi-step process by calling its bash scripts,
99
  and returns the path to the final rigged GLB file.
100
  """
 
 
101
  if not os.path.isdir(UNIRIG_REPO_DIR):
102
  raise gr.Error(f"UniRig repository not found at {UNIRIG_REPO_DIR}. Cannot proceed. Please check Space setup.")
103
 
104
  if input_glb_file_obj is None:
105
  raise gr.Error("No input file provided. Please upload a .glb mesh.")
106
 
107
- input_glb_path = input_glb_file_obj # This is the absolute path from gr.File(type="filepath")
108
  print(f"Input GLB path received: {input_glb_path}")
109
 
110
- # Create a dedicated temporary directory for all intermediate and final files
111
- # The output paths for UniRig scripts will point into this directory.
112
  processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_")
113
  print(f"Using temporary processing directory: {processing_temp_dir}")
114
 
115
  try:
116
  base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
117
 
118
- # Define absolute paths for intermediate files within the processing_temp_dir
119
  abs_skeleton_output_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
120
  abs_skin_output_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx")
121
  abs_final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")
122
 
123
- # Step 1: Skeleton Prediction using generate_skeleton.sh
124
  print("Step 1: Predicting Skeleton...")
125
  skeleton_cmd = [
126
  "bash", "launch/inference/generate_skeleton.sh",
127
- "--input", input_glb_path, # Input is the original GLB
128
  "--output", abs_skeleton_output_path
129
  ]
130
  run_unirig_command(skeleton_cmd, "Skeleton Prediction")
131
  if not os.path.exists(abs_skeleton_output_path):
132
  raise gr.Error("Skeleton prediction failed to produce an output file. Check logs for UniRig errors.")
133
 
134
- # Step 2: Skinning Weight Prediction using generate_skin.sh
135
  print("Step 2: Predicting Skinning Weights...")
136
- # generate_skin.sh requires the skeleton from step 1 as --input,
137
- # and the original mesh as --source.
138
  skin_cmd = [
139
  "bash", "launch/inference/generate_skin.sh",
140
- "--input", abs_skeleton_output_path, # Input is the skeleton FBX from previous step
141
- "--source", input_glb_path, # Source is the original GLB mesh
142
  "--output", abs_skin_output_path
143
  ]
144
  run_unirig_command(skin_cmd, "Skinning Prediction")
145
  if not os.path.exists(abs_skin_output_path):
146
  raise gr.Error("Skinning prediction failed to produce an output file. Check logs for UniRig errors.")
147
 
148
- # Step 3: Merge Skeleton/Skin with Original Mesh using merge.sh
149
  print("Step 3: Merging Results...")
150
- # merge.sh requires the skinned FBX as --source (which contains skeleton and weights)
151
- # and the original GLB as --target.
152
  merge_cmd = [
153
  "bash", "launch/inference/merge.sh",
154
- "--source", abs_skin_output_path, # Source is the skinned FBX from previous step
155
- "--target", input_glb_path, # Target is the original GLB mesh
156
  "--output", abs_final_rigged_glb_path
157
  ]
158
  run_unirig_command(merge_cmd, "Merging")
@@ -161,14 +177,14 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
161
 
162
  return abs_final_rigged_glb_path
163
 
164
- except gr.Error:
165
- if os.path.exists(processing_temp_dir):
166
  shutil.rmtree(processing_temp_dir)
167
  print(f"Cleaned up temporary directory: {processing_temp_dir}")
168
  raise
169
  except Exception as e:
170
  print(f"An unexpected error occurred in rig_glb_mesh_multistep: {e}")
171
- if os.path.exists(processing_temp_dir):
172
  shutil.rmtree(processing_temp_dir)
173
  print(f"Cleaned up temporary directory: {processing_temp_dir}")
174
  raise gr.Error(f"An unexpected error occurred during processing: {str(e)[:500]}")
@@ -181,14 +197,14 @@ theme = gr.themes.Soft(
181
  font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
182
  )
183
 
184
- if not os.path.isdir(UNIRIG_REPO_DIR) and __name__ == "__main__":
185
  print(f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. The application will not work.")
186
 
187
  iface = gr.Interface(
188
- fn=rig_glb_mesh_multistep,
189
  inputs=gr.File(
190
  label="Upload .glb Mesh File",
191
- type="filepath"
192
  ),
193
  outputs=gr.Model3D(
194
  label="Rigged 3D Model (.glb)",
 
25
  else:
26
  print("Blender is already installed.")
27
 
28
+ # Verify Blender Python path
29
+ if os.path.exists(BLENDER_PYTHON_PATH):
30
+ print(f"Blender Python modules found at {BLENDER_PYTHON_PATH}")
31
+ if os.path.exists(os.path.join(BLENDER_PYTHON_PATH, "bpy.so")):
32
+ print("bpy.so found - Blender Python integration should work")
33
+ else:
34
+ print("bpy.so not found - Check Blender installation or path")
35
+ else:
36
+ print(f"Blender Python modules not found at {BLENDER_PYTHON_PATH} - Update BLENDER_PYTHON_PATH")
37
+
38
  if not os.path.isdir(UNIRIG_REPO_DIR):
39
  print(f"ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. Please clone it there.")
40
  # Consider raising an error or displaying it in the UI if UniRig is critical for startup
 
47
  else:
48
  print("Warning: CUDA not available or not detected by PyTorch. UniRig performance will be severely impacted.")
49
 
50
+ def patch_asset_py():
51
+ """Temporary patch to fix type hinting error in asset.py"""
52
+ asset_py_path = os.path.join(UNIRIG_REPO_DIR, "src", "data", "asset.py")
53
+ try:
54
+ with open(asset_py_path, "r") as f:
55
+ content = f.read()
56
+ # Check if patch is already applied
57
+ if "meta: Union[Dict[str, Any], None]" in content:
58
+ print("Patch already applied to asset.py")
59
+ return
60
+ # Replace the problematic line
61
+ content = content.replace("meta: Union[Dict[str, ...], None]=None", "meta: Union[Dict[str, Any], None]=None")
62
+ # Ensure 'Any' is imported
63
+ if "from typing import Any" not in content:
64
+ content = "from typing import Any\n" + content
65
+ with open(asset_py_path, "w") as f:
66
+ f.write(content)
67
+ print("Successfully patched asset.py")
68
+ except Exception as e:
69
+ print(f"Failed to patch asset.py: {e}")
70
+ raise
71
+
72
  @spaces.GPU # Decorator for ZeroGPU
73
  def run_unirig_command(command_list, step_name):
74
  """
75
  Helper function to run UniRig commands (now expecting bash scripts) using subprocess.
76
  command_list: The full command and its arguments, e.g., ["bash", "script.sh", "--arg", "value"]
77
  """
 
78
  cmd = command_list
79
 
80
  print(f"Running {step_name}: {' '.join(cmd)}")
 
84
  # Determine the path to the 'src' directory within UniRig, where the 'unirig' package resides.
85
  unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, "src")
86
 
87
+ # Explicitly add BLENDER_PYTHON_PATH, UNIRIG_REPO_DIR/src, and UNIRIG_REPO_DIR to PYTHONPATH
88
+ new_pythonpath_parts = [BLENDER_PYTHON_PATH, unirig_src_dir, UNIRIG_REPO_DIR]
 
 
89
  existing_pythonpath = process_env.get('PYTHONPATH', '')
 
90
  if existing_pythonpath:
 
91
  new_pythonpath_parts.extend(existing_pythonpath.split(os.pathsep))
92
+ process_env["PYTHONPATH"] = os.pathsep.join(filter(None, new_pythonpath_parts))
 
93
  print(f"Set PYTHONPATH for subprocess: {process_env['PYTHONPATH']}")
94
 
95
  try:
 
 
 
96
  result = subprocess.run(cmd, cwd=UNIRIG_REPO_DIR, capture_output=True, text=True, check=True, env=process_env)
97
  print(f"{step_name} STDOUT:\n{result.stdout}")
98
  if result.stderr:
 
103
  print(f"Return code: {e.returncode}")
104
  print(f"Stdout: {e.stdout}")
105
  print(f"Stderr: {e.stderr}")
106
+ error_summary = e.stderr.splitlines()[-5:]
 
107
  raise gr.Error(f"Error in UniRig {step_name}. Details: {' '.join(error_summary)}")
108
  except FileNotFoundError:
 
109
  print(f"ERROR: Could not find executable or script for {step_name}. Command: {' '.join(cmd)}. Is UniRig cloned correctly and 'bash' available?")
110
  raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs, UniRig directory structure, and script paths.")
111
  except Exception as e_general:
 
119
  rigs it using the new UniRig multi-step process by calling its bash scripts,
120
  and returns the path to the final rigged GLB file.
121
  """
122
+ patch_asset_py() # Apply the patch before running commands
123
+
124
  if not os.path.isdir(UNIRIG_REPO_DIR):
125
  raise gr.Error(f"UniRig repository not found at {UNIRIG_REPO_DIR}. Cannot proceed. Please check Space setup.")
126
 
127
  if input_glb_file_obj is None:
128
  raise gr.Error("No input file provided. Please upload a .glb mesh.")
129
 
130
+ input_glb_path = input_glb_file_obj
131
  print(f"Input GLB path received: {input_glb_path}")
132
 
 
 
133
  processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_")
134
  print(f"Using temporary processing directory: {processing_temp_dir}")
135
 
136
  try:
137
  base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
138
 
 
139
  abs_skeleton_output_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
140
  abs_skin_output_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx")
141
  abs_final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")
142
 
143
+ # Step 1: Skeleton Prediction
144
  print("Step 1: Predicting Skeleton...")
145
  skeleton_cmd = [
146
  "bash", "launch/inference/generate_skeleton.sh",
147
+ "--input", input_glb_path,
148
  "--output", abs_skeleton_output_path
149
  ]
150
  run_unirig_command(skeleton_cmd, "Skeleton Prediction")
151
  if not os.path.exists(abs_skeleton_output_path):
152
  raise gr.Error("Skeleton prediction failed to produce an output file. Check logs for UniRig errors.")
153
 
154
+ # Step 2: Skinning Weight Prediction
155
  print("Step 2: Predicting Skinning Weights...")
 
 
156
  skin_cmd = [
157
  "bash", "launch/inference/generate_skin.sh",
158
+ "--input", abs_skeleton_output_path,
159
+ "--source", input_glb_path,
160
  "--output", abs_skin_output_path
161
  ]
162
  run_unirig_command(skin_cmd, "Skinning Prediction")
163
  if not os.path.exists(abs_skin_output_path):
164
  raise gr.Error("Skinning prediction failed to produce an output file. Check logs for UniRig errors.")
165
 
166
+ # Step 3: Merge Skeleton/Skin with Original Mesh
167
  print("Step 3: Merging Results...")
 
 
168
  merge_cmd = [
169
  "bash", "launch/inference/merge.sh",
170
+ "--source", abs_skin_output_path,
171
+ "--target", input_glb_path,
172
  "--output", abs_final_rigged_glb_path
173
  ]
174
  run_unirig_command(merge_cmd, "Merging")
 
177
 
178
  return abs_final_rigged_glb_path
179
 
180
+ except gr.Error:
181
+ if os.path.exists(processing_temp_dir):
182
  shutil.rmtree(processing_temp_dir)
183
  print(f"Cleaned up temporary directory: {processing_temp_dir}")
184
  raise
185
  except Exception as e:
186
  print(f"An unexpected error occurred in rig_glb_mesh_multistep: {e}")
187
+ if os.path.exists(processing_temp_dir):
188
  shutil.rmtree(processing_temp_dir)
189
  print(f"Cleaned up temporary directory: {processing_temp_dir}")
190
  raise gr.Error(f"An unexpected error occurred during processing: {str(e)[:500]}")
 
197
  font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
198
  )
199
 
200
+ if not os.path.isdir(UNIRIG_REPO_DIR) and __name__ == "__main__":
201
  print(f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. The application will not work.")
202
 
203
  iface = gr.Interface(
204
+ fn=rig_glb_mesh_multistep,
205
  inputs=gr.File(
206
  label="Upload .glb Mesh File",
207
+ type="filepath"
208
  ),
209
  outputs=gr.Model3D(
210
  label="Rigged 3D Model (.glb)",