From 8f2e2041465cbc601dbdc7ac5b186a18c8fb91b7 Mon Sep 17 00:00:00 2001
From: Ethanfel <ethan.fel@ts-pc.fr>
Date: Fri, 27 Mar 2026 19:55:55 +0100
Subject: [PATCH] fix: show pip output, handle incomplete venv, fix TF version
 for Python 3.12

- tensorflow-cpu==2.15.0 only supports Python <=3.11; relax to >=2.16.0
- capture_output=False so pip errors are visible in ComfyUI logs
- clean up incomplete venv dir before retrying install

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 nodes/feature_extractor.py     | 22 +++++++++++++++++++---
 scripts/install_extract_env.sh |  2 +-
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/nodes/feature_extractor.py b/nodes/feature_extractor.py
index f152fb1..99c738d 100644
--- a/nodes/feature_extractor.py
+++ b/nodes/feature_extractor.py
@@ -15,26 +15,42 @@ _MANAGED_PYTHON = os.path.join(_MANAGED_VENV, "bin", "python")
 
 _EXTRACT_PACKAGES = [
     "torch", "torchaudio", "torchvision",
-    "tensorflow-cpu==2.15.0",
+    # TF 2.15 only supports Python <=3.11; use >=2.16 for Python 3.12+
+    "tensorflow-cpu>=2.16.0",
     "jax[cpu]", "jaxlib",
     "transformers", "decord", "einops", "numpy", "mediapy",
     "git+https://github.com/google-deepmind/videoprism.git",
 ]
 
 
+def _pip_run(pip, *args):
+    """Run pip and stream output; raise with visible error on failure."""
+    result = subprocess.run([pip] + list(args), capture_output=False)
+    if result.returncode != 0:
+        raise RuntimeError(
+            f"[PrismAudio] pip {' '.join(args[:2])} failed (exit {result.returncode}). "
+            "Check the output above for details."
+        )
+
+
 def _ensure_extract_env():
     """Create and populate the managed venv on first use."""
     if os.path.exists(_MANAGED_PYTHON):
         return _MANAGED_PYTHON
 
+    import shutil
+    if os.path.exists(_MANAGED_VENV):
+        print("[PrismAudio] Removing incomplete venv and retrying...")
+        shutil.rmtree(_MANAGED_VENV)
+
     print("[PrismAudio] Feature-extraction env not found — creating venv at:", _MANAGED_VENV)
     subprocess.run([sys.executable, "-m", "venv", _MANAGED_VENV], check=True)
 
     pip = os.path.join(_MANAGED_VENV, "bin", "pip")
-    subprocess.run([pip, "install", "--upgrade", "pip"], check=True)
+    _pip_run(pip, "install", "--upgrade", "pip")
 
     print("[PrismAudio] Installing feature-extraction dependencies (this takes a few minutes)...")
-    subprocess.run([pip, "install"] + _EXTRACT_PACKAGES, check=True)
+    _pip_run(pip, "install", *_EXTRACT_PACKAGES)
 
     print("[PrismAudio] Feature-extraction env ready.")
     return _MANAGED_PYTHON
diff --git a/scripts/install_extract_env.sh b/scripts/install_extract_env.sh
index 621b6d3..1495e31 100755
--- a/scripts/install_extract_env.sh
+++ b/scripts/install_extract_env.sh
@@ -27,7 +27,7 @@ echo "[PrismAudio] Installing PyTorch stack..."
 
 echo "[PrismAudio] Installing feature-extraction dependencies..."
 "${PIP}" install \
-    "tensorflow-cpu==2.15.0" \
+    "tensorflow-cpu>=2.16.0" \
     "jax[cpu]" \
     "jaxlib" \
     "transformers" \