Update app.py

app.py

@@ -158,17 +158,16 @@ def run_unirig_command(script_path: str, args: List[str], step_name: str):
         args: A list of command-line arguments for the script.
         step_name: Name of the step for logging.
     """
-    # Command is now bash + script path + arguments
     cmd = ["bash", script_path] + args
 
     print(f"\n--- Running UniRig Step: {step_name} ---")
     print(f"Command: {' '.join(cmd)}")
 
-    # Prepare the environment for the subprocess (shell script)
+    # Prepare the environment for the subprocess
     process_env = os.environ.copy()
     unirig_src_dir = os.path.join(UNIRIG_REPO_DIR, "src")
 
-    # 1. Set PYTHONPATH: Blender's site-packages + UniRig source
+    # 1. Set PYTHONPATH: Blender’s site-packages + UniRig source
     pythonpath_parts = [
         BLENDER_PYTHON_SITE_PACKAGES,
         unirig_src_dir,
@@ -177,30 +176,29 @@ def run_unirig_command(script_path: str, args: List[str], step_name: str):
     process_env["PYTHONPATH"] = os.pathsep.join(filter(None, pythonpath_parts))
     print(f"Subprocess PYTHONPATH: {process_env['PYTHONPATH']}")
 
-    # 2. Set LD_LIBRARY_PATH: Include Blender's Python library directory
+    # 2. Set LD_LIBRARY_PATH: Include Blender’s Python library directory
     blender_lib_path = os.path.join(BLENDER_PYTHON_DIR, "lib")
-    # 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']}")
-
-
+    # 3. Set PATH with a temporary 'python' symlink to Blender’s Python
+    temp_dir = tempfile.mkdtemp()
     try:
-        # Execute the shell script.
-        # cwd=UNIRIG_REPO_DIR ensures the script runs from the repo's root,
+        python_link = os.path.join(temp_dir, "python")
+        os.symlink(BLENDER_PYTHON_EXEC, python_link)  # Links 'python' to 'python3.11'
+        existing_path = process_env.get('PATH', '')
+        process_env["PATH"] = f"{temp_dir}{os.pathsep}{BLENDER_PYTHON_BIN_DIR}{os.pathsep}{existing_path}"
+        print(f"Subprocess PATH: {process_env['PATH']}")
+
+        # Execute the shell 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 # Pass the modified environment
+            check=True,
+            env=process_env
         )
         print(f"{step_name} STDOUT:\n{result.stdout}")
         if result.stderr:
@@ -214,19 +212,15 @@ def run_unirig_command(script_path: str, args: List[str], step_name: str):
         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
         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.")
+            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.")
+            raise gr.Error(f"Error in UniRig '{step_name}': Script failed to load PyTorch extensions. Check environment and PyTorch installation.")
         else:
-             raise gr.Error(f"Error in UniRig '{step_name}'. Check logs. Last error lines:\n{last_lines}")
-
+            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
         print(f"ERROR: Could not find executable 'bash' or script '{script_path}' for {step_name}.")
-        print(f"Attempted command: {' '.join(cmd)}")
         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:
@@ -235,6 +229,9 @@ def run_unirig_command(script_path: str, args: List[str], step_name: str):
         traceback.print_exc()
         raise gr.Error(f"Unexpected Python error during '{step_name}' execution: {str(e_general)[:500]}")
 
+    finally:
+        shutil.rmtree(temp_dir)  # Clean up the temporary directory
+
     print(f"--- Finished UniRig Step: {step_name} ---")
 
 
@@ -263,6 +260,19 @@ def rig_glb_mesh_multistep(input_glb_file_obj):
     if not input_glb_path.lower().endswith(".glb"):
          raise gr.Error("Invalid file type. Please upload a .glb file.")
 
+    # Test PyTorch import in Blender’s Python
+    try:
+        result = subprocess.run(
+            [BLENDER_PYTHON_EXEC, "-c", "import torch; print(torch.__version__)"],
+            capture_output=True,
+            text=True,
+            check=True
+        )
+        print(f"PyTorch version in Blender's Python: {result.stdout.strip()}")
+    except subprocess.CalledProcessError as e:
+        print(f"Failed to import torch in Blender's Python: {e.stderr}")
+        raise gr.Error("PyTorch import failed in Blender's Python environment. Check installation.")
+
     # --- Setup Temporary Directory ---
     processing_temp_dir = tempfile.mkdtemp(prefix="unirig_processing_")
     print(f"Using temporary processing directory: {processing_temp_dir}")