Update app.py

app.py

@@ -20,7 +20,8 @@ BLENDER_PYTHON_VERSION = "python3.11" # Blender 4.2 uses Python 3.11
 
 # Construct paths dynamically based on the above
 BLENDER_PYTHON_DIR = os.path.join(BLENDER_INSTALL_DIR, BLENDER_PYTHON_VERSION_DIR, "python")
-BLENDER_PYTHON_EXEC = os.path.join(BLENDER_PYTHON_DIR, "bin", BLENDER_PYTHON_VERSION) # Still needed for PYTHONPATH
+BLENDER_PYTHON_BIN_DIR = os.path.join(BLENDER_PYTHON_DIR, "bin") # Directory containing python executable
+BLENDER_PYTHON_EXEC = os.path.join(BLENDER_PYTHON_BIN_DIR, BLENDER_PYTHON_VERSION) # Full path to executable
 BLENDER_PYTHON_LIB_PATH = os.path.join(BLENDER_PYTHON_DIR, "lib", BLENDER_PYTHON_VERSION)
 BLENDER_PYTHON_SITE_PACKAGES = os.path.join(BLENDER_PYTHON_LIB_PATH, "site-packages")
 
@@ -61,19 +62,23 @@ else:
     print(f"Blender Python executable found: {BLENDER_PYTHON_EXEC}")
 
 # Verify Blender Python site-packages path and bpy module presence
+bpy_found = False
 if os.path.exists(BLENDER_PYTHON_SITE_PACKAGES):
     print(f"Blender Python site-packages found at: {BLENDER_PYTHON_SITE_PACKAGES}")
-    # Check for 'bpy' directory which contains __init__.py and modules
-    bpy_module_path = os.path.join(BLENDER_PYTHON_SITE_PACKAGES, "bpy")
-    if os.path.isdir(bpy_module_path) and os.path.exists(os.path.join(bpy_module_path, "__init__.py")):
-        print("Blender Python 'bpy' module directory found.")
-    # Fallback check for bpy.so if directory structure is different
+    # Check 1: 'bpy' directory with __init__.py
+    bpy_module_dir = os.path.join(BLENDER_PYTHON_SITE_PACKAGES, "bpy")
+    if os.path.isdir(bpy_module_dir) and os.path.exists(os.path.join(bpy_module_dir, "__init__.py")):
+        print("Found 'bpy' module directory in site-packages.")
+        bpy_found = True
+    # Check 2: 'bpy.so' file (less common structure)
     elif os.path.exists(os.path.join(BLENDER_PYTHON_SITE_PACKAGES, "bpy.so")):
-         print("Blender Python 'bpy.so' found.")
-    else:
-        print("WARNING: Blender Python 'bpy' module not found in site-packages. Check Blender installation or paths.")
+         print("Found 'bpy.so' in site-packages.")
+         bpy_found = True
+
+    if not bpy_found:
+        print("WARNING: Blender Python 'bpy' module indicator not found in site-packages. Imports might fail.")
 else:
-    print(f"WARNING: Blender Python site-packages not found at {BLENDER_PYTHON_SITE_PACKAGES}. Check paths.")
+    print(f"WARNING: Blender Python site-packages directory not found at {BLENDER_PYTHON_SITE_PACKAGES}. Check paths.")
 
 
 # Check for UniRig repository
@@ -157,63 +162,66 @@ def run_unirig_command(script_path: str, args: List[str], step_name: str):
     cmd = ["bash", script_path] + args
 
     print(f"\n--- Running UniRig Step: {step_name} ---")
-    # Use subprocess.list2cmdline for a more accurate representation if needed,
-    # but simple join is often fine for logging.
     print(f"Command: {' '.join(cmd)}")
 
     # Prepare the environment for the subprocess (shell script)
-    # The shell script will likely invoke python and needs the correct env
     process_env = os.environ.copy()
     unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, "src")
 
-    # Crucial: Set PYTHONPATH so the python invoked by the shell script finds Blender's
-    # site-packages and the UniRig source code.
+    # 1. Set PYTHONPATH: Blender's site-packages + UniRig source
     pythonpath_parts = [
         BLENDER_PYTHON_SITE_PACKAGES,
         unirig_src_dir,
-        UNIRIG_REPO_DIR  # Include base repo dir as well
+        UNIRIG_REPO_DIR
     ]
     process_env["PYTHONPATH"] = os.pathsep.join(filter(None, pythonpath_parts))
     print(f"Subprocess PYTHONPATH: {process_env['PYTHONPATH']}")
 
-    # Add Blender's library path and potentially the python bin path to PATH
-    # This helps the shell find the correct python if needed, and shared libs
+    # 2. Set LD_LIBRARY_PATH: Include Blender's Python library directory
     blender_lib_path = os.path.join(BLENDER_PYTHON_DIR, "lib")
-    blender_bin_path = os.path.join(BLENDER_PYTHON_DIR, "bin")
-    process_env["LD_LIBRARY_PATH"] = f"{blender_lib_path}{os.pathsep}{process_env.get('LD_LIBRARY_PATH', '')}"
-    process_env["PATH"] = f"{blender_bin_path}{os.pathsep}{process_env.get('PATH', '')}" # Add blender python bin to PATH
+    # Prepend Blender lib path to existing LD_LIBRARY_PATH if it exists
+    existing_ld_path = process_env.get('LD_LIBRARY_PATH', '')
+    process_env["LD_LIBRARY_PATH"] = f"{blender_lib_path}{os.pathsep}{existing_ld_path}" if existing_ld_path else blender_lib_path
     print(f"Subprocess LD_LIBRARY_PATH: {process_env['LD_LIBRARY_PATH']}")
+
+    # 3. Set PATH: *Prepend* Blender's Python bin directory to the system PATH.
+    # This makes it the *first* place the shell looks for 'python' or 'python3.11'.
+    existing_path = process_env.get('PATH', '')
+    process_env["PATH"] = f"{BLENDER_PYTHON_BIN_DIR}{os.pathsep}{existing_path}"
     print(f"Subprocess PATH: {process_env['PATH']}")
 
 
     try:
         # Execute the shell script.
         # cwd=UNIRIG_REPO_DIR ensures the script runs from the repo's root,
-        # which is often necessary for finding relative paths within the script.
         result = subprocess.run(
             cmd,
             cwd=UNIRIG_REPO_DIR,
             capture_output=True,
             text=True,
             check=True, # Raises CalledProcessError on non-zero exit codes
-            env=process_env
+            env=process_env # Pass the modified environment
         )
         print(f"{step_name} STDOUT:\n{result.stdout}")
         if result.stderr:
-            # Treat stderr as potential warnings unless check=True failed
             print(f"{step_name} STDERR (Warnings/Info):\n{result.stderr}")
 
     except subprocess.CalledProcessError as e:
-        # Error occurred within the shell script (non-zero exit code)
         print(f"ERROR during {step_name}: Subprocess failed!")
         print(f"Command: {' '.join(e.cmd)}")
         print(f"Return code: {e.returncode}")
         print(f"--- {step_name} STDOUT ---:\n{e.stdout}")
         print(f"--- {step_name} STDERR ---:\n{e.stderr}")
-        # Provide a concise error message to Gradio
         error_summary = e.stderr.strip().splitlines()
         last_lines = "\n".join(error_summary[-5:]) if error_summary else "No stderr output."
-        raise gr.Error(f"Error in UniRig '{step_name}'. Check logs. Last error lines:\n{last_lines}")
+        # Check specifically for bpy/torch import errors within the subprocess stderr
+        if "ModuleNotFoundError: No module named 'bpy'" in e.stderr:
+             raise gr.Error(f"Error in UniRig '{step_name}': Script failed to import Blender's 'bpy' module. Check environment setup.")
+        elif "ImportError: Failed to load PyTorch C extensions" in e.stderr:
+             raise gr.Error(f"Error in UniRig '{step_name}': Script failed to load PyTorch extensions. Check environment and PyTorch installation within Blender's Python.")
+        else:
+             raise gr.Error(f"Error in UniRig '{step_name}'. Check logs. Last error lines:\n{last_lines}")
+
 
     except FileNotFoundError:
         # This error means 'bash' or the script_path wasn't found
@@ -222,7 +230,6 @@ def run_unirig_command(script_path: str, args: List[str], step_name: str):
         raise gr.Error(f"Setup error for UniRig '{step_name}'. 'bash' or script '{os.path.basename(script_path)}' not found.")
 
     except Exception as e_general:
-        # Catch any other unexpected Python errors during subprocess handling
         print(f"An unexpected Python exception occurred in run_unirig_command for {step_name}: {e_general}")
         import traceback
         traceback.print_exc()
@@ -240,12 +247,9 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
     try:
         patch_asset_py() # Attempt patch, might still be relevant
     except gr.Error as e:
-         # If patching is critical and fails, raise the error
-         # raise e
          print(f"Ignoring patch error: {e}") # Or just log it
     except Exception as e:
          print(f"Ignoring unexpected patch error: {e}")
-         # raise gr.Error(f"Failed to prepare UniRig environment: {e}")
 
     # --- Input Validation ---
     if input_glb_file_obj is None:
@@ -266,13 +270,13 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
     try:
         # --- Define File Paths ---
         base_name = os.path.splitext(os.path.basename(input_glb_path))[0]
+        # Ensure paths passed to scripts are absolute
         abs_input_glb_path = os.path.abspath(input_glb_path)
         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")
 
         # --- Define Absolute Paths to UniRig SHELL Scripts ---
-        # Construct absolute paths based on UNIRIG_REPO_DIR, using .sh extension
         skeleton_script_path = os.path.join(UNIRIG_REPO_DIR, "launch/inference/generate_skeleton.sh")
         skin_script_path = os.path.join(UNIRIG_REPO_DIR, "launch/inference/generate_skin.sh")
         merge_script_path = os.path.join(UNIRIG_REPO_DIR, "launch/inference/merge.sh")
@@ -281,28 +285,24 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
 
         # Step 1: Skeleton Prediction
         print("\nStarting Step 1: Predicting Skeleton...")
-        # Arguments for the generate_skeleton.sh script
         skeleton_args = [
             "--input", abs_input_glb_path,
             "--output", abs_skeleton_output_path
-            # Add any other arguments the shell script expects
         ]
-        # Check if script exists before running
         if not os.path.exists(skeleton_script_path):
             raise gr.Error(f"Skeleton script not found at: {skeleton_script_path}")
         run_unirig_command(skeleton_script_path, skeleton_args, "Skeleton Prediction")
         if not os.path.exists(abs_skeleton_output_path):
+            # Check if the error wasn't already raised by run_unirig_command
             raise gr.Error("Skeleton prediction failed. Output file not created. Check logs.")
         print("Step 1: Skeleton Prediction completed.")
 
         # Step 2: Skinning Weight Prediction
         print("\nStarting Step 2: Predicting Skinning Weights...")
-        # Arguments for the generate_skin.sh script
         skin_args = [
             "--input", abs_skeleton_output_path, # Input is the skeleton from step 1
             "--source", abs_input_glb_path,    # Source mesh
             "--output", abs_skin_output_path
-            # Add any other arguments the shell script expects
         ]
         if not os.path.exists(skin_script_path):
             raise gr.Error(f"Skinning script not found at: {skin_script_path}")
@@ -313,12 +313,12 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
 
         # Step 3: Merge Skeleton/Skin with Original Mesh
         print("\nStarting Step 3: Merging Results...")
-        # Arguments for the merge.sh script
         merge_args = [
-            "--source", abs_skin_output_path, # Input is the skinned result from step 2
-            "--target", abs_input_glb_path,   # Target is the original mesh
+            # Determine which source to use based on what exists or a user choice?
+            # Assuming skin output is the desired source if it exists.
+            "--source", abs_skin_output_path,
+            "--target", abs_input_glb_path,
             "--output", abs_final_rigged_glb_path
-            # Add any other arguments the shell script expects
         ]
         if not os.path.exists(merge_script_path):
             raise gr.Error(f"Merging script not found at: {merge_script_path}")
@@ -332,7 +332,6 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
         return abs_final_rigged_glb_path
 
     except gr.Error as e:
-        # If a gr.Error was raised, clean up and re-raise
         print(f"Gradio Error occurred: {e}")
         if os.path.exists(processing_temp_dir):
             shutil.rmtree(processing_temp_dir)
@@ -340,15 +339,12 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
         raise e
 
     except Exception as e:
-        # Catch any other unexpected errors
         print(f"An unexpected error occurred in rig_glb_mesh_multistep: {e}")
         import traceback
         traceback.print_exc()
-        # Clean up temporary files before raising the error
         if os.path.exists(processing_temp_dir):
             shutil.rmtree(processing_temp_dir)
             print(f"Cleaned up temporary directory: {processing_temp_dir}")
-        # Raise a generic Gradio error
         raise gr.Error(f"An unexpected error occurred during processing: {str(e)[:500]}")