fix: use --torch-backend instead of --extra-index-url for dependency compilation

The old --extra-index-url approach caused uv to resolve any package from
the PyTorch wheel index when available there, not just torch-ecosystem
packages. Replace it with --torch-backend which routes only torch/
torchaudio/torchvision/torchsde to the PyTorch index.

Pin uv>=0.6.0 as the minimum version supporting --torch-backend.

Author: bigcat88

comfy_cli/uv.py

@@ -67,10 +67,12 @@ class DependencyCompiler:
     rocmPytorchUrl = "https://download.pytorch.org/whl/rocm6.1"
     nvidiaPytorchUrl = "https://download.pytorch.org/whl/cu126"
 
+    rocmTorchBackend = "rocm6.1"
+    nvidiaTorchBackend = "cu126"
+
     overrideGpu = dedent(
         """
         # ensure usage of {gpu} version of pytorch
-        --extra-index-url {gpuUrl}
         torch
         torchaudio
         torchsde
@@ -118,6 +120,7 @@ def Compile(
         out: PathLike | None = None,
         override: PathLike | None = None,
         resolve_strategy: str | None = None,
+        torch_backend: str | None = None,
     ) -> subprocess.CompletedProcess[Any]:
         cmd = [
             str(executable),
@@ -136,8 +139,11 @@ def Compile(
         if emit_index_url:
             cmd.append("--emit-index-url")
 
+        if torch_backend is not None:
+            cmd.extend(["--torch-backend", torch_backend])
+
         # ensures that eg tqdm is latest version, even though an old tqdm is on the amd url
-        # see https://github.com/astral-sh/uv/blob/main/PIP_COMPATIBILITY.md#packages-that-exist-on-multiple-indexes and https://github.com/astral-sh/uv/issues/171
+        # see https://github.com/astral-sh/uv/blob/main/PIP_COMPATIBILITY.md#packages-that-exist-on-multiple-indexes
         if index_strategy is not None:
             cmd.extend(["--index-strategy", "unsafe-best-match"])
 
@@ -383,6 +389,11 @@ def __init__(
             DependencyCompiler.rocmPytorchUrl if self.gpu == GPU_OPTION.AMD else
             None
         )  # fmt: skip
+        self.torchBackend = (
+            DependencyCompiler.nvidiaTorchBackend if self.gpu == GPU_OPTION.NVIDIA else
+            DependencyCompiler.rocmTorchBackend if self.gpu == GPU_OPTION.AMD else
+            None
+        )  # fmt: skip
         self.out: Path = self.outDir / outName
         self.override = self.outDir / "override.txt"
 
@@ -401,8 +412,8 @@ def make_override(self):
         self.override.unlink(missing_ok=True)
 
         with open(self.override, "w") as f:
-            if self.gpu is not None and self.gpuUrl is not None:
-                f.write(DependencyCompiler.overrideGpu.format(gpu=self.gpu, gpuUrl=self.gpuUrl))
+            if self.gpu is not None:
+                f.write(DependencyCompiler.overrideGpu.format(gpu=self.gpu))
                 f.write("\n\n")
 
         completed = DependencyCompiler.Compile(
@@ -412,6 +423,7 @@ def make_override(self):
             emit_index_url=False,
             executable=self.executable,
             override=self.override,
+            torch_backend=self.torchBackend,
         )
 
         with open(self.override, "a") as f:
@@ -442,6 +454,7 @@ def compile_core_plus_ext(self):
                     override=self.override,
                     out=self.out,
                     resolve_strategy="ask",
+                    torch_backend=self.torchBackend,
                 )
 
                 break

pyproject.toml

@@ -48,7 +48,7 @@ dependencies = [
   "tomlkit",
   "typer>=0.12.5",
   "typing-extensions>=4.7",
-  "uv",
+  "uv>=0.6.0",
   "websocket-client",
 ]
 

tests/uv/mock_requirements/requirements.compiled

@@ -1,29 +1,27 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile /home/tel/git/comfy-cli/tests/uv/mock_requirements/core_reqs.txt /home/tel/git/comfy-cli/tests/uv/mock_requirements/x_reqs.txt /home/tel/git/comfy-cli/tests/uv/mock_requirements/y_reqs.txt --emit-index-annotation --emit-index-url --index-strategy unsafe-best-match -o /home/tel/git/comfy-cli/tests/temp/test_uv/requirements.compiled --override /home/tel/git/comfy-cli/tests/temp/test_uv/override.txt
+#    uv pip compile tests/uv/mock_requirements/core_reqs.txt tests/uv/mock_requirements/x_reqs.txt tests/uv/mock_requirements/y_reqs.txt --emit-index-annotation --emit-index-url --torch-backend rocm6.1 --index-strategy unsafe-best-match -o tests/temp/test_uv/requirements.compiled --override tests/temp/test_uv/override.txt
 --index-url https://pypi.org/simple
---extra-index-url https://download.pytorch.org/whl/rocm6.1
 
 mpmath==1.3.0
     # via
-    #   -r /home/tel/git/comfy-cli/tests/uv/mock_requirements/y_reqs.txt
+    #   -r tests/uv/mock_requirements/y_reqs.txt
     #   sympy
-    # from https://download.pytorch.org/whl/rocm6.1
+    # from https://pypi.org/simple
 numpy==2.0.2
     # via
-    #   --override override.txt
-    #   -r /home/tel/git/comfy-cli/tests/uv/mock_requirements/x_reqs.txt
-    #   -r /home/tel/git/comfy-cli/tests/uv/mock_requirements/y_reqs.txt
+    #   -r tests/uv/mock_requirements/x_reqs.txt
+    #   -r tests/uv/mock_requirements/y_reqs.txt
     # from https://pypi.org/simple
 sympy==1.13.0
     # via
     #   --override override.txt
-    #   -r /home/tel/git/comfy-cli/tests/uv/mock_requirements/x_reqs.txt
-    #   -r /home/tel/git/comfy-cli/tests/uv/mock_requirements/y_reqs.txt
+    #   -r tests/uv/mock_requirements/x_reqs.txt
+    #   -r tests/uv/mock_requirements/y_reqs.txt
     # from https://pypi.org/simple
 tqdm==4.66.4
     # via
     #   --override override.txt
-    #   -r /home/tel/git/comfy-cli/tests/uv/mock_requirements/core_reqs.txt
-    #   -r /home/tel/git/comfy-cli/tests/uv/mock_requirements/x_reqs.txt
-    #   -r /home/tel/git/comfy-cli/tests/uv/mock_requirements/y_reqs.txt
+    #   -r tests/uv/mock_requirements/core_reqs.txt
+    #   -r tests/uv/mock_requirements/x_reqs.txt
+    #   -r tests/uv/mock_requirements/y_reqs.txt
     # from https://pypi.org/simple

tests/uv/test_uv.py

@@ -1,5 +1,6 @@
 import shutil
 from pathlib import Path
+from unittest.mock import patch
 
 import pytest
 
@@ -77,3 +78,86 @@ def _filter_optional(lines: list[str]) -> list[str]:
         knownLines, testLines = [_filter_optional(lines) for lines in (knownLines, testLines)]
 
         assert knownLines == testLines
+
+
+def test_torch_backend_nvidia():
+    depComp = DependencyCompiler(cwd=temp, gpu=GPU_OPTION.NVIDIA, outDir=temp, reqFilesCore=[], reqFilesExt=[])
+    assert depComp.torchBackend == "cu126"
+    assert depComp.gpuUrl == DependencyCompiler.nvidiaPytorchUrl
+
+
+def test_torch_backend_amd():
+    depComp = DependencyCompiler(cwd=temp, gpu=GPU_OPTION.AMD, outDir=temp, reqFilesCore=[], reqFilesExt=[])
+    assert depComp.torchBackend == "rocm6.1"
+    assert depComp.gpuUrl == DependencyCompiler.rocmPytorchUrl
+
+
+def test_torch_backend_cpu():
+    depComp = DependencyCompiler(cwd=temp, gpu=GPU_OPTION.CPU, outDir=temp, reqFilesCore=[], reqFilesExt=[])
+    assert depComp.torchBackend is None
+    assert depComp.gpuUrl is None
+
+
+def test_torch_backend_none():
+    with patch.object(DependencyCompiler, "Resolve_Gpu", return_value=None):
+        depComp = DependencyCompiler(cwd=temp, gpu=None, outDir=temp, reqFilesCore=[], reqFilesExt=[])
+    assert depComp.torchBackend is None
+    assert depComp.gpuUrl is None
+
+
+def test_compile_passes_torch_backend():
+    """Verify that Compile() includes --torch-backend in the command when provided."""
+    with patch("comfy_cli.uv._run") as mock_run:
+        mock_run.return_value = type("R", (), {"stdout": "", "stderr": "", "returncode": 0})()
+        DependencyCompiler.Compile(
+            cwd=temp,
+            reqFiles=[mockReqsDir / "core_reqs.txt"],
+            torch_backend="cu126",
+        )
+    cmd = mock_run.call_args[0][0]
+    idx = cmd.index("--torch-backend")
+    assert cmd[idx + 1] == "cu126"
+
+
+def test_compile_omits_torch_backend_when_none():
+    """Verify that Compile() does not include --torch-backend when torch_backend is None."""
+    with patch("comfy_cli.uv._run") as mock_run:
+        mock_run.return_value = type("R", (), {"stdout": "", "stderr": "", "returncode": 0})()
+        DependencyCompiler.Compile(
+            cwd=temp,
+            reqFiles=[mockReqsDir / "core_reqs.txt"],
+            torch_backend=None,
+        )
+    cmd = mock_run.call_args[0][0]
+    assert "--torch-backend" not in cmd
+
+
+def test_compiled_output_has_no_extra_index_url(mock_prompt_select):
+    """The compiled output must not contain --extra-index-url (torch-backend handles routing)."""
+    depComp = DependencyCompiler(
+        cwd=temp,
+        gpu=GPU_OPTION.AMD,
+        outDir=temp,
+        reqFilesCore=[mockReqsDir / "core_reqs.txt"],
+        reqFilesExt=[mockReqsDir / "x_reqs.txt", mockReqsDir / "y_reqs.txt"],
+    )
+    depComp.make_override()
+    depComp.compile_core_plus_ext()
+
+    content = depComp.out.read_text()
+    assert "--extra-index-url" not in content
+
+
+def test_override_file_has_no_extra_index_url():
+    depComp = DependencyCompiler(
+        cwd=temp,
+        gpu=GPU_OPTION.AMD,
+        outDir=temp,
+        reqFilesCore=[mockReqsDir / "core_reqs.txt"],
+        reqFilesExt=[],
+    )
+    depComp.make_override()
+
+    content = depComp.override.read_text()
+    assert "--extra-index-url" not in content
+    assert "torch" in content