Skip to content

Commit 523f46b

Browse files
committed
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 it happened to be available there, not just torch-ecosystem packages. Replace it with uv's --torch-backend flag which routes only torch, torchaudio, torchvision, and torchsde to the correct PyTorch index. Map NVIDIA to cu126, AMD to rocm6.1, and CPU to the cpu index so each backend gets the right torch variant. Pin uv>=0.6.0 in pyproject.toml as the minimum version supporting --torch-backend.
1 parent 8bab84d commit 523f46b

2 files changed

Lines changed: 6 additions & 2 deletions

File tree

comfy_cli/uv.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,9 +64,11 @@ def parse_req_file(rf: PathLike, skips: list[str] | None = None):
6464

6565

6666
class DependencyCompiler:
67+
cpuPytorchUrl = "https://download.pytorch.org/whl/cpu"
6768
rocmPytorchUrl = "https://download.pytorch.org/whl/rocm6.1"
6869
nvidiaPytorchUrl = "https://download.pytorch.org/whl/cu126"
6970

71+
cpuTorchBackend = "cpu"
7072
rocmTorchBackend = "rocm6.1"
7173
nvidiaTorchBackend = "cu126"
7274

@@ -387,11 +389,13 @@ def __init__(
387389
self.gpuUrl = (
388390
DependencyCompiler.nvidiaPytorchUrl if self.gpu == GPU_OPTION.NVIDIA else
389391
DependencyCompiler.rocmPytorchUrl if self.gpu == GPU_OPTION.AMD else
392+
DependencyCompiler.cpuPytorchUrl if self.gpu == GPU_OPTION.CPU else
390393
None
391394
) # fmt: skip
392395
self.torchBackend = (
393396
DependencyCompiler.nvidiaTorchBackend if self.gpu == GPU_OPTION.NVIDIA else
394397
DependencyCompiler.rocmTorchBackend if self.gpu == GPU_OPTION.AMD else
398+
DependencyCompiler.cpuTorchBackend if self.gpu == GPU_OPTION.CPU else
395399
None
396400
) # fmt: skip
397401
self.out: Path = self.outDir / outName

tests/uv/test_uv.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -94,8 +94,8 @@ def test_torch_backend_amd():
9494

9595
def test_torch_backend_cpu():
9696
depComp = DependencyCompiler(cwd=temp, gpu=GPU_OPTION.CPU, outDir=temp, reqFilesCore=[], reqFilesExt=[])
97-
assert depComp.torchBackend is None
98-
assert depComp.gpuUrl is None
97+
assert depComp.torchBackend == "cpu"
98+
assert depComp.gpuUrl == DependencyCompiler.cpuPytorchUrl
9999

100100

101101
def test_torch_backend_none():

0 commit comments

Comments
 (0)