Update app.py

app.py

@@ -175,15 +175,14 @@ def run_unirig_command(python_script_path: str, script_args: List[str], step_nam
 
     # Environment variables might still be useful if the script imports other non-blender libs
     process_env = os.environ.copy()
-    # Ensure UniRig source is findable if scripts use relative imports from repo root
-    # or if they import modules from UniRig/src
     unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, "src")
+
+    # Set PYTHONPATH: Include UniRig source dir, UniRig base dir, and '.' for CWD.
     pythonpath_parts = [
-        # Blender's site-packages should be found automatically when run via Blender exec
         unirig_src_dir,
-        UNIRIG_REPO_DIR
+        UNIRIG_REPO_DIR,
+        '.' # Add current working directory explicitly
     ]
-    # Add existing PYTHONPATH if any
     existing_pythonpath = process_env.get('PYTHONPATH', '')
     if existing_pythonpath:
         pythonpath_parts.append(existing_pythonpath)
@@ -202,7 +201,7 @@ def run_unirig_command(python_script_path: str, script_args: List[str], step_nam
         # cwd=UNIRIG_REPO_DIR ensures script runs relative to repo root.
         result = subprocess.run(
             cmd,
-            cwd=UNIRIG_REPO_DIR,
+            cwd=UNIRIG_REPO_DIR, # CWD is set to UniRig repo root
             capture_output=True,
             text=True,
             check=True, # Raises CalledProcessError on non-zero exit codes
@@ -225,9 +224,10 @@ def run_unirig_command(python_script_path: str, script_args: List[str], step_nam
         print(f"--- {step_name} STDERR ---:\n{e.stderr}")
         error_summary = e.stderr.strip().splitlines()
         last_lines = "\n".join(error_summary[-5:]) if error_summary else "No stderr output."
-        # Check specifically for bpy/torch import errors within the subprocess stderr
-        # These *shouldn't* happen now, but check just in case
-        if "ModuleNotFoundError: No module named 'bpy'" in e.stderr:
+        # Check specifically for import errors within the subprocess stderr
+        if "ModuleNotFoundError: No module named 'src'" in e.stderr:
+             raise gr.Error(f"Error in UniRig '{step_name}': Script failed to import the 'src' module. Check PYTHONPATH and script CWD.")
+        elif "ModuleNotFoundError: No module named 'bpy'" in e.stderr:
              raise gr.Error(f"Error in UniRig '{step_name}': Blender failed to provide 'bpy' module internally.")
         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 even within Blender. Check installation.")
@@ -281,13 +281,7 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
         abs_final_rigged_glb_path = os.path.join(processing_temp_dir, f"{base_name}_rigged_final.glb")
 
         # --- Define Absolute Paths to UniRig PYTHON Scripts ---
-        # *** IMPORTANT: Assuming the .sh scripts primarily wrap these .py scripts. ***
-        # *** This might need adjustment based on the actual content of the .sh files. ***
-        # Use run.py as the main entry point, passing task-specific configs/args
-        # This requires inspecting run.py and the shell scripts to know the correct args/configs.
-        # For now, let's assume a structure like: python run.py --config <config_file> --input ... --output ...
-
-        # Example: Assuming run.py takes task name and args
+        # *** Assuming run.py is the correct entry point. VERIFY THIS. ***
         run_script_path = os.path.join(UNIRIG_REPO_DIR, "run.py") # Main script?
 
         # --- Execute UniRig Steps ---
@@ -295,17 +289,14 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
         # Step 1: Skeleton Prediction
         print("\nStarting Step 1: Predicting Skeleton...")
         # Arguments for the run.py script for skeleton task
-        # ** These are GUESSES based on typical structure - VERIFY from UniRig code **
+        # ** These are GUESSES - VERIFY from UniRig code / shell scripts **
         skeleton_args = [
-            # Might need a config file path or task name argument first
-            # e.g., "--task", "generate_skeleton",
+            # "--task", "generate_skeleton", # Example: if run.py needs a task identifier
             "--input", abs_input_glb_path,
             "--output", abs_skeleton_output_path
-            # Add other relevant args like model path if needed
         ]
-        if not os.path.exists(run_script_path): # Check if assumed script exists
+        if not os.path.exists(run_script_path):
             raise gr.Error(f"UniRig main script not found at: {run_script_path}")
-        # Execute run.py via Blender
         run_unirig_command(run_script_path, skeleton_args, "Skeleton Prediction")
         if not os.path.exists(abs_skeleton_output_path):
             raise gr.Error("Skeleton prediction failed. Output file not created. Check logs.")
@@ -314,11 +305,11 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
         # Step 2: Skinning Weight Prediction
         print("\nStarting Step 2: Predicting Skinning Weights...")
         # Arguments for the run.py script for skinning task
-        # ** GUESSES - VERIFY from UniRig code **
+        # ** GUESSES - VERIFY **
         skin_args = [
-            # e.g., "--task", "generate_skin",
-            "--input", abs_skeleton_output_path, # Input is the skeleton from step 1
-            "--source", abs_input_glb_path,    # Source mesh (if needed by skin script)
+            # "--task", "generate_skin",
+            "--input", abs_skeleton_output_path,
+            "--source", abs_input_glb_path,
             "--output", abs_skin_output_path
         ]
         run_unirig_command(run_script_path, skin_args, "Skinning Prediction")
@@ -329,11 +320,9 @@ 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 run.py script for merging task
-        # ** GUESSES - VERIFY from UniRig code **
-        # Alternatively, merge might use a different script like src/data/merge_fbx.py?
-        # Let's assume run.py handles it for now.
+        # ** GUESSES - VERIFY **
         merge_args = [
-            # e.g., "--task", "merge",
+            # "--task", "merge",
             "--source", abs_skin_output_path,
             "--target", abs_input_glb_path,
             "--output", abs_final_rigged_glb_path