Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
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
|
57 |
-
|
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 |
-
|
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
|
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
|
124 |
print("Step 1: Predicting Skeleton...")
|
125 |
skeleton_cmd = [
|
126 |
"bash", "launch/inference/generate_skeleton.sh",
|
127 |
-
"--input", input_glb_path,
|
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
|
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,
|
141 |
-
"--source", input_glb_path,
|
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
|
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,
|
155 |
-
"--target", input_glb_path,
|
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)",
|