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

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +25 -54
app.py CHANGED
@@ -15,6 +15,8 @@ UNIRIG_REPO_DIR = os.path.join(os.path.dirname(__file__), "UniRig")
15
  # Path to the Blender Python site-packages
16
  BLENDER_PYTHON_PATH = "/opt/blender-4.2.0-linux-x64/4.2/python/lib/python3.11/site-packages"
17
 
 
 
18
  # Path to the setup script
19
  SETUP_SCRIPT = os.path.join(os.path.dirname(__file__), "setup_blender.sh")
20
 
@@ -71,11 +73,8 @@ def patch_asset_py():
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)}")
81
 
@@ -86,9 +85,9 @@ def run_unirig_command(command_list, step_name):
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
 
@@ -106,23 +105,19 @@ def run_unirig_command(command_list, step_name):
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:
112
  print(f"An unexpected Python exception occurred in run_unirig_command for {step_name}: {e_general}")
113
  raise gr.Error(f"Unexpected Python error during {step_name}: {str(e_general)[:500]}")
114
 
115
- @spaces.GPU # Decorator for ZeroGPU
116
  def rig_glb_mesh_multistep(input_glb_file_obj):
117
- """
118
- Takes an input GLB file object (from gr.File with type="filepath"),
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.")
@@ -135,45 +130,30 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
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")
175
  if not os.path.exists(abs_final_rigged_glb_path):
176
- raise gr.Error("Merging process failed to produce the final rigged GLB file. Check logs for UniRig errors.")
177
 
178
  return abs_final_rigged_glb_path
179
 
@@ -198,25 +178,17 @@ theme = gr.themes.Soft(
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)",
211
- clear_color=[0.8, 0.8, 0.8, 1.0],
212
- ),
213
  title="UniRig Auto-Rigger (Python 3.11 / PyTorch 2.3+)",
214
  description=(
215
- "Upload a 3D mesh in `.glb` format. This application uses the latest UniRig to automatically rig the mesh by calling its provided bash scripts.\n"
216
- "The process involves: 1. Skeleton Prediction, 2. Skinning Weight Prediction, 3. Merging.\n"
217
- "This may take several minutes. Ensure your GLB has clean geometry.\n"
218
- f"Running on: {str(DEVICE).upper()}. UniRig repo expected at: '{os.path.basename(UNIRIG_REPO_DIR)}'.\n"
219
- f"UniRig Source: https://github.com/VAST-AI-Research/UniRig"
220
  ),
221
  cache_examples=False,
222
  theme=theme
@@ -224,6 +196,5 @@ iface = gr.Interface(
224
 
225
  if __name__ == "__main__":
226
  if not os.path.isdir(UNIRIG_REPO_DIR):
227
- print(f"CRITICAL: UniRig repository not found at {UNIRIG_REPO_DIR}. Ensure it's cloned in the Space's root.")
228
-
229
  iface.launch()
 
15
  # Path to the Blender Python site-packages
16
  BLENDER_PYTHON_PATH = "/opt/blender-4.2.0-linux-x64/4.2/python/lib/python3.11/site-packages"
17
 
18
+ BLENDER_PYTHON_EXEC = "/opt/blender-4.2.0-linux-x64/4.2/python/bin/python3.11"
19
+
20
  # Path to the setup script
21
  SETUP_SCRIPT = os.path.join(os.path.dirname(__file__), "setup_blender.sh")
22
 
 
73
 
74
  @spaces.GPU # Decorator for ZeroGPU
75
  def run_unirig_command(command_list, step_name):
76
+ """Run UniRig commands using Blender's Python interpreter."""
77
+ cmd = [BLENDER_PYTHON_EXEC] + command_list[1:] # Use Blender's Python instead of bash!!
 
 
 
78
 
79
  print(f"Running {step_name}: {' '.join(cmd)}")
80
 
 
85
 
86
  # Explicitly add BLENDER_PYTHON_PATH, UNIRIG_REPO_DIR/src, and UNIRIG_REPO_DIR to PYTHONPATH
87
  new_pythonpath_parts = [BLENDER_PYTHON_PATH, unirig_src_dir, UNIRIG_REPO_DIR]
88
+ #existing_pythonpath = process_env.get('PYTHONPATH', '')
89
+ #if existing_pythonpath:
90
+ # new_pythonpath_parts.extend(existing_pythonpath.split(os.pathsep))
91
  process_env["PYTHONPATH"] = os.pathsep.join(filter(None, new_pythonpath_parts))
92
  print(f"Set PYTHONPATH for subprocess: {process_env['PYTHONPATH']}")
93
 
 
105
  error_summary = e.stderr.splitlines()[-5:]
106
  raise gr.Error(f"Error in UniRig {step_name}. Details: {' '.join(error_summary)}")
107
  except FileNotFoundError:
108
+ print(f"ERROR: Could not find executable or script for {step_name}. Command: {' '.join(cmd)}")
109
+ raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs and script paths.")
110
  except Exception as e_general:
111
  print(f"An unexpected Python exception occurred in run_unirig_command for {step_name}: {e_general}")
112
  raise gr.Error(f"Unexpected Python error during {step_name}: {str(e_general)[:500]}")
113
 
114
+ @spaces.GPU
115
  def rig_glb_mesh_multistep(input_glb_file_obj):
116
+ """Rig a GLB mesh using UniRig's multi-step process."""
117
+ patch_asset_py()
 
 
 
 
118
 
119
  if not os.path.isdir(UNIRIG_REPO_DIR):
120
+ raise gr.Error(f"UniRig repository not found at {UNIRIG_REPO_DIR}.")
121
 
122
  if input_glb_file_obj is None:
123
  raise gr.Error("No input file provided. Please upload a .glb mesh.")
 
130
 
131
  try:
132
  base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
 
133
  abs_skeleton_output_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
134
  abs_skin_output_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx")
135
  abs_final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")
136
 
137
  # Step 1: Skeleton Prediction
138
  print("Step 1: Predicting Skeleton...")
139
+ skeleton_cmd = ["python", "launch/inference/generate_skeleton.py", "--input", input_glb_path, "--output", abs_skeleton_output_path]
 
 
 
 
140
  run_unirig_command(skeleton_cmd, "Skeleton Prediction")
141
  if not os.path.exists(abs_skeleton_output_path):
142
  raise gr.Error("Skeleton prediction failed to produce an output file. Check logs for UniRig errors.")
143
 
144
  # Step 2: Skinning Weight Prediction
145
  print("Step 2: Predicting Skinning Weights...")
146
+ skin_cmd = ["python", "launch/inference/generate_skin.py", "--input", abs_skeleton_output_path, "--source", input_glb_path, "--output", abs_skin_output_path]
 
 
 
 
 
147
  run_unirig_command(skin_cmd, "Skinning Prediction")
148
  if not os.path.exists(abs_skin_output_path):
149
+ raise gr.Error("Skinning prediction failed to produce an output file.")
150
 
151
  # Step 3: Merge Skeleton/Skin with Original Mesh
152
  print("Step 3: Merging Results...")
153
+ merge_cmd = ["python", "launch/inference/merge.py", "--source", abs_skin_output_path, "--target", input_glb_path, "--output", abs_final_rigged_glb_path]
 
 
 
 
 
154
  run_unirig_command(merge_cmd, "Merging")
155
  if not os.path.exists(abs_final_rigged_glb_path):
156
+ raise gr.Error("Merging process failed to produce the final rigged GLB file.")
157
 
158
  return abs_final_rigged_glb_path
159
 
 
178
  )
179
 
180
  if not os.path.isdir(UNIRIG_REPO_DIR) and __name__ == "__main__":
181
+ print(f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}.")
182
 
183
  iface = gr.Interface(
184
  fn=rig_glb_mesh_multistep,
185
+ inputs=gr.File(label="Upload .glb Mesh File", type="filepath"),
186
+ outputs=gr.Model3D(label="Rigged 3D Model (.glb)", clear_color=[0.8, 0.8, 0.8, 1.0]),
 
 
 
 
 
 
187
  title="UniRig Auto-Rigger (Python 3.11 / PyTorch 2.3+)",
188
  description=(
189
+ "Upload a 3D mesh in `.glb` format. This application uses the latest UniRig to automatically rig the mesh.\n"
190
+ "Running on: {DEVICE.upper()}. UniRig repo expected at: '{os.path.basename(UNIRIG_REPO_DIR)}'.\n"
191
+ "UniRig Source: https://github.com/VAST-AI-Research/UniRig"
 
 
192
  ),
193
  cache_examples=False,
194
  theme=theme
 
196
 
197
  if __name__ == "__main__":
198
  if not os.path.isdir(UNIRIG_REPO_DIR):
199
+ print(f"CRITICAL: UniRig repository not found at {UNIRIG_REPO_DIR}.")
 
200
  iface.launch()