Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
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 |
-
|
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 |
-
|
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)}
|
110 |
-
raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs
|
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
|
116 |
def rig_glb_mesh_multistep(input_glb_file_obj):
|
117 |
-
"""
|
118 |
-
|
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 |
-
|
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.
|
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.
|
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}.
|
202 |
|
203 |
iface = gr.Interface(
|
204 |
fn=rig_glb_mesh_multistep,
|
205 |
-
inputs=gr.File(
|
206 |
-
|
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
|
216 |
-
"
|
217 |
-
"
|
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}.
|
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()
|