olivv-cs-ComfyUI-FunPack/core/torch_swap.py

"""Switch this ComfyUI's torch to the CUDA 13 build of the SAME versions -- or leave it as it was.

Why: on a Blackwell rental, torch 2.10+cu128 sampled int8 models ~2x slower than bf16; the
same torch from the cu130 index made int8 faster than bf16 (2026-10-10).

Order is the safety. Everything is DOWNLOADED first (a failed download changes nothing), then
installed from those files only, then a fresh interpreter must import all three and report
CUDA 13. If any of that fails the old build is put back. A torch that does not import means a
ComfyUI that does not start, and then there is no UI left to repair it from.

Pins carry the local tag (`torch==2.10.0+cu130`): a bare `torch==2.10.0` is already satisfied
by `2.10.0+cu128`, so pip does nothing -- which is exactly how the hand-typed version failed.
"""

import shutil
import subprocess
import sys
import tempfile
from importlib import metadata

PACKAGES = ("torch", "torchvision", "torchaudio")
INDEX = "https://download.pytorch.org/whl/{tag}"
PYPI = "https://pypi.org/simple"
TARGET = "cu130"
NEED_FREE = 12 * 1024 ** 3       # the cu13 runtime wheels alone are several GB
# Compiled against one torch: worth checking after the swap, not touched by it.
COMPILED = ("xformers", "sageattn3", "flash_attn", "flash-attn")

step = None          # what is happening now, for the UI's progress card


def _version(name):
    try:
        return metadata.version(name)
    except metadata.PackageNotFoundError:
        return None


def _pip(args, timeout):
    return subprocess.run([sys.executable, "-m", "pip", *args, "--disable-pip-version-check"],
                          capture_output=True, text=True, timeout=timeout)


def _tail(proc):
    return "\n".join(((proc.stderr or "") + (proc.stdout or "")).strip().splitlines()[-6:])


def _pins(have, tag):
    """`torch==2.10.0+cu130` for each installed package; no tag means PyPI's own build."""
    return [f"{p}=={v.split('+')[0]}{'+' + tag if tag else ''}" for p, v in have.items()]


def _constraints(path):
    """Everything else stays at the version it is now. torch's own family and the CUDA runtime
    wheels are left free: they are what changes."""
    proc = _pip(["freeze", "--exclude-editable"], 120)
    free = PACKAGES + ("triton",)
    lines = [l for l in proc.stdout.splitlines()
             if "==" in l and not l.split("==")[0].lower().startswith(free + ("nvidia-",))]
    with open(path, "w") as f:
        f.write("\n".join(lines) + "\n")


def _imports(names):
    """CUDA version a FRESH interpreter's torch reports, or the error. This process keeps the old
    torch in memory until the restart, so asking here would always say the old answer."""
    proc = subprocess.run([sys.executable, "-c", f"import {', '.join(names)}; print(torch.version.cuda)"],
                          capture_output=True, text=True, timeout=300)
    return (proc.stdout.strip().splitlines() or [""])[-1] if proc.returncode == 0 else None, _tail(proc)


def swap(timeout=3600):
    """-> {"installed": bool (the new build is on disk), "message": str, "rebuild": [names]}."""
    global step
    have = {p: _version(p) for p in PACKAGES if _version(p)}
    if "torch" not in have:
        return {"installed": False, "message": "torch is not installed in this Python, so there is nothing to switch."}
    import torch
    if not getattr(torch.version, "cuda", None):
        return {"installed": False, "message": "This torch has no CUDA (a Mac, or a CPU build): there is no CUDA 13 build to switch to."}
    torch_v = have["torch"]
    old_tag = torch_v.split("+")[1] if "+" in torch_v else None
    if old_tag == TARGET or (old_tag or "").startswith("cu13"):
        return {"installed": False, "message": f"torch {torch_v} is already a CUDA 13 build."}
    if old_tag and not old_tag.startswith("cu"):
        return {"installed": False, "message": f"torch {torch_v} is not a CUDA build ({old_tag}); switching it to CUDA 13 is not something this can judge."}
    work = tempfile.mkdtemp(prefix="funpack_torch_")
    try:
        if shutil.disk_usage(work).free < NEED_FREE:
            return {"installed": False, "message": f"Not enough free disk for the download (needs ~{NEED_FREE // 1024 ** 3} GB). Nothing was changed."}
        constraints = f"{work}/constraints.txt"
        _constraints(constraints)
        new = _pins(have, TARGET)
        step = f"Downloading {', '.join(new)} (several GB)…"
        proc = _pip(["download", "-d", f"{work}/wheels", "--index-url", INDEX.format(tag=TARGET),
                     "--extra-index-url", PYPI, "-c", constraints, *new], timeout)
        if proc.returncode != 0:
            return {"installed": False, "message": f"The download failed, so nothing was changed.\n{_tail(proc)}"}
        step = "Installing the downloaded build…"
        proc = _pip(["install", "--no-index", "--find-links", f"{work}/wheels", "-c", constraints, *new], timeout)
        cuda, err = _imports(list(have)) if proc.returncode == 0 else (None, _tail(proc))
        if cuda and cuda.startswith("13"):
            rebuild = [n for n in COMPILED if _version(n)]
            return {"installed": True, "rebuild": rebuild,
                    "message": f"torch {', '.join(new)} installed and imports with CUDA {cuda}."
                               + (f" Check these after the restart (built for the old torch): {', '.join(rebuild)}." if rebuild else "")}
        step = "That did not work; putting the old build back…"
        back = _pins(have, old_tag) if old_tag else [f"{p}=={v}" for p, v in have.items()]
        index = ["--index-url", INDEX.format(tag=old_tag), "--extra-index-url", PYPI] if old_tag else []
        undo = _pip(["install", *index, "-c", constraints, *back], timeout)
        restored, _ = _imports(list(have))
        if undo.returncode == 0 and restored:
            return {"installed": False, "message": f"The CUDA 13 build failed to install or import, so the old one ({torch_v}) is back.\n{err}"}
        return {"installed": False, "broken": True,
                "message": (f"The CUDA 13 build failed AND putting {torch_v} back failed. ComfyUI will probably not start "
                            f"after a restart. Fix it in a terminal:\n  {sys.executable} -m pip install --force-reinstall "
                            f"{' '.join(back)} {' '.join(index)}\n{err}\n{_tail(undo)}")}
    except subprocess.TimeoutExpired:
        return {"installed": False, "message": "pip took too long and was stopped. Check the terminal: torch may be half-installed."}
    finally:
        step = None
        shutil.rmtree(work, ignore_errors=True)