Update app.py

app.py

@@ -15,6 +15,8 @@ UNIRIG_REPO_DIR = os.path.join(os.path.dirname(__file__), "UniRig")
 # Path to the Blender Python site-packages
 BLENDER_PYTHON_PATH = "/opt/blender-4.2.0-linux-x64/4.2/python/lib/python3.11/site-packages"
 
+BLENDER_PYTHON_EXEC = "/opt/blender-4.2.0-linux-x64/4.2/python/bin/python3.11"
+
 # Path to the setup script
 SETUP_SCRIPT = os.path.join(os.path.dirname(__file__), "setup_blender.sh")
 
@@ -71,11 +73,8 @@ def patch_asset_py():
 
 @spaces.GPU # Decorator for ZeroGPU
 def run_unirig_command(command_list, step_name):
-    """
-    Helper function to run UniRig commands (now expecting bash scripts) using subprocess.
-    command_list: The full command and its arguments, e.g., ["bash", "script.sh", "--arg", "value"]
-    """
-    cmd = command_list
+    """Run UniRig commands using Blender's Python interpreter."""
+    cmd = [BLENDER_PYTHON_EXEC] + command_list[1:]  # Use Blender's Python instead of bash!!
     
     print(f"Running {step_name}: {' '.join(cmd)}")
     
@@ -86,9 +85,9 @@ def run_unirig_command(command_list, step_name):
 
     # Explicitly add BLENDER_PYTHON_PATH, UNIRIG_REPO_DIR/src, and UNIRIG_REPO_DIR to PYTHONPATH
     new_pythonpath_parts = [BLENDER_PYTHON_PATH, unirig_src_dir, UNIRIG_REPO_DIR]
-    existing_pythonpath = process_env.get('PYTHONPATH', '')
-    if existing_pythonpath:
-        new_pythonpath_parts.extend(existing_pythonpath.split(os.pathsep))
+    #existing_pythonpath = process_env.get('PYTHONPATH', '')
+    #if existing_pythonpath:
+    #    new_pythonpath_parts.extend(existing_pythonpath.split(os.pathsep))
     process_env["PYTHONPATH"] = os.pathsep.join(filter(None, new_pythonpath_parts))
     print(f"Set PYTHONPATH for subprocess: {process_env['PYTHONPATH']}")
 
@@ -106,23 +105,19 @@ def run_unirig_command(command_list, step_name):
         error_summary = e.stderr.splitlines()[-5:]
         raise gr.Error(f"Error in UniRig {step_name}. Details: {' '.join(error_summary)}")
     except FileNotFoundError:
-        print(f"ERROR: Could not find executable or script for {step_name}. Command: {' '.join(cmd)}. Is UniRig cloned correctly and 'bash' available?")
-        raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs, UniRig directory structure, and script paths.")
+        print(f"ERROR: Could not find executable or script for {step_name}. Command: {' '.join(cmd)}")
+        raise gr.Error(f"Setup error for UniRig {step_name}. Check server logs and script paths.")
     except Exception as e_general:
         print(f"An unexpected Python exception occurred in run_unirig_command for {step_name}: {e_general}")
         raise gr.Error(f"Unexpected Python error during {step_name}: {str(e_general)[:500]}")
 
-@spaces.GPU # Decorator for ZeroGPU
+@spaces.GPU
 def rig_glb_mesh_multistep(input_glb_file_obj):
-    """
-    Takes an input GLB file object (from gr.File with type="filepath"),
-    rigs it using the new UniRig multi-step process by calling its bash scripts,
-    and returns the path to the final rigged GLB file.
-    """
-    patch_asset_py()  # Apply the patch before running commands
+    """Rig a GLB mesh using UniRig's multi-step process."""
+    patch_asset_py()
 
     if not os.path.isdir(UNIRIG_REPO_DIR):
-         raise gr.Error(f"UniRig repository not found at {UNIRIG_REPO_DIR}. Cannot proceed. Please check Space setup.")
+        raise gr.Error(f"UniRig repository not found at {UNIRIG_REPO_DIR}.")
 
     if input_glb_file_obj is None:
         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):
 
     try:
         base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
-        
         abs_skeleton_output_path = os.path.join(processing_temp_dir, f"{base_name}_skeleton.fbx")
         abs_skin_output_path = os.path.join(processing_temp_dir, f"{base_name}_skin.fbx")
         abs_final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")
 
         # Step 1: Skeleton Prediction
         print("Step 1: Predicting Skeleton...")
-        skeleton_cmd = [
-            "bash", "launch/inference/generate_skeleton.sh",
-            "--input", input_glb_path,
-            "--output", abs_skeleton_output_path
-        ]
+        skeleton_cmd = ["python", "launch/inference/generate_skeleton.py", "--input", input_glb_path, "--output", abs_skeleton_output_path]
         run_unirig_command(skeleton_cmd, "Skeleton Prediction")
         if not os.path.exists(abs_skeleton_output_path):
             raise gr.Error("Skeleton prediction failed to produce an output file. Check logs for UniRig errors.")
 
         # Step 2: Skinning Weight Prediction
         print("Step 2: Predicting Skinning Weights...")
-        skin_cmd = [
-            "bash", "launch/inference/generate_skin.sh",
-            "--input", abs_skeleton_output_path,
-            "--source", input_glb_path,
-            "--output", abs_skin_output_path
-        ]
+        skin_cmd = ["python", "launch/inference/generate_skin.py", "--input", abs_skeleton_output_path, "--source", input_glb_path, "--output", abs_skin_output_path]
         run_unirig_command(skin_cmd, "Skinning Prediction")
         if not os.path.exists(abs_skin_output_path):
-            raise gr.Error("Skinning prediction failed to produce an output file. Check logs for UniRig errors.")
+            raise gr.Error("Skinning prediction failed to produce an output file.")
 
         # Step 3: Merge Skeleton/Skin with Original Mesh
         print("Step 3: Merging Results...")
-        merge_cmd = [
-            "bash", "launch/inference/merge.sh",
-            "--source", abs_skin_output_path,
-            "--target", input_glb_path,
-            "--output", abs_final_rigged_glb_path
-        ]
+        merge_cmd = ["python", "launch/inference/merge.py", "--source", abs_skin_output_path, "--target", input_glb_path, "--output", abs_final_rigged_glb_path]
         run_unirig_command(merge_cmd, "Merging")
         if not os.path.exists(abs_final_rigged_glb_path):
-            raise gr.Error("Merging process failed to produce the final rigged GLB file. Check logs for UniRig errors.")
+            raise gr.Error("Merging process failed to produce the final rigged GLB file.")
 
         return abs_final_rigged_glb_path
 
@@ -198,25 +178,17 @@ theme = gr.themes.Soft(
 )
 
 if not os.path.isdir(UNIRIG_REPO_DIR) and __name__ == "__main__":
-    print(f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}. The application will not work.")
+    print(f"CRITICAL STARTUP ERROR: UniRig repository not found at {UNIRIG_REPO_DIR}.")
 
 iface = gr.Interface(
     fn=rig_glb_mesh_multistep,
-    inputs=gr.File(
-        label="Upload .glb Mesh File",
-        type="filepath"
-    ),
-    outputs=gr.Model3D(
-        label="Rigged 3D Model (.glb)",
-        clear_color=[0.8, 0.8, 0.8, 1.0],
-    ),
+    inputs=gr.File(label="Upload .glb Mesh File", type="filepath"),
+    outputs=gr.Model3D(label="Rigged 3D Model (.glb)", clear_color=[0.8, 0.8, 0.8, 1.0]),
     title="UniRig Auto-Rigger (Python 3.11 / PyTorch 2.3+)",
     description=(
-        "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"
-        "The process involves: 1. Skeleton Prediction, 2. Skinning Weight Prediction, 3. Merging.\n"
-        "This may take several minutes. Ensure your GLB has clean geometry.\n"
-        f"Running on: {str(DEVICE).upper()}. UniRig repo expected at: '{os.path.basename(UNIRIG_REPO_DIR)}'.\n"
-        f"UniRig Source: https://github.com/VAST-AI-Research/UniRig"
+        "Upload a 3D mesh in `.glb` format. This application uses the latest UniRig to automatically rig the mesh.\n"
+        "Running on: {DEVICE.upper()}. UniRig repo expected at: '{os.path.basename(UNIRIG_REPO_DIR)}'.\n"
+        "UniRig Source: https://github.com/VAST-AI-Research/UniRig"
     ),
     cache_examples=False,
     theme=theme
@@ -224,6 +196,5 @@ iface = gr.Interface(
 
 if __name__ == "__main__":
     if not os.path.isdir(UNIRIG_REPO_DIR):
-        print(f"CRITICAL: UniRig repository not found at {UNIRIG_REPO_DIR}. Ensure it's cloned in the Space's root.")
-    
+        print(f"CRITICAL: UniRig repository not found at {UNIRIG_REPO_DIR}.")
     iface.launch()