Ethanfel/ComfyUI-SelVA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: switch managed venv to jax[cuda13] for GPU feature extraction

Browse Source 
RTX 6000 Pro (Blackwell SM 10.0) fully supports CUDA 13. Switch from
jax[cpu]+jaxlib to jax[cuda13] which bundles jaxlib and uses
pip-managed CUDA libraries. Delete _extract_env to force a rebuild.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Ethanfel
2026-03-27 20:33:45 +01:00
parent 06f8dbbab4
commit 0f46e8359d
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/feature_extractor.py
+2 -1
View File
@@ -17,7 +17,8 @@ _EXTRACT_PACKAGES = [
"torch", "torchaudio", "torchvision", "torch", "torchaudio", "torchvision",
# TF 2.15 only supports Python <=3.11; use >=2.16 for Python 3.12+ # TF 2.15 only supports Python <=3.11; use >=2.16 for Python 3.12+
"tensorflow-cpu>=2.16.0", "tensorflow-cpu>=2.16.0",
"jax[cpu]", "jaxlib", # jax[cuda13] includes jaxlib; pip-managed CUDA libs (no local toolkit needed)
"jax[cuda13]",
"transformers", "decord", "einops", "numpy", "mediapy", "transformers", "decord", "einops", "numpy", "mediapy",
"git+https://github.com/google-deepmind/videoprism.git", "git+https://github.com/google-deepmind/videoprism.git",
] ]