vjumpkung-ComfyUI-STARWrapper/install.py

#!/usr/bin/env python3
"""
Installation script for STARVSRWrapper.

Resolves an xformers build that matches the PyTorch already installed in this
environment and installs it with uv (if available) or pip.

Resolution order:
  1. Ask the official PyTorch wheel index for the detected CUDA tag
     (https://download.pytorch.org/whl/cuXXX/xformers/) and pick the newest
     release whose declared `torch` pin is satisfied by the installed PyTorch.
  2. Fall back to a static torch -> xformers table if the index is unreachable.

PyTorch is never installed, upgraded or replaced: the install runs with
--no-deps so a mismatched xformers can't drag a different torch into a working
ComfyUI environment. Failures are reported and the script still exits 0 so it
can never block ComfyUI startup.
"""

import json
import platform
import re
import subprocess
import sys
import urllib.error
import urllib.request

PYTORCH_INDEX = "https://download.pytorch.org/whl/{tag}"
PYPI_RELEASE = "https://pypi.org/pypi/xformers/{version}/json"
NETWORK_TIMEOUT = 20

# Fallback table, newest first. Used only when the PyTorch index is unreachable.
# Each entry is the newest xformers release that pins that exact torch version.
FALLBACK_MAP = {
    "2.10.0": "0.0.34",
    "2.9.1": "0.0.33.post2",
    "2.9.0": "0.0.33.post1",
    "2.8.0": "0.0.32.post2",
    "2.7.1": "0.0.31.post1",
    "2.7.0": "0.0.30",
    "2.6.0": "0.0.29.post2",
}
# xformers releases that declare a floor (torch>=X) rather than an exact pin.
FALLBACK_FLOOR = [("2.10", "0.0.35")]


def run_command(cmd, check=True, capture_output=True):
    """Run a command and return the result."""
    result = subprocess.run(
        cmd, shell=True, capture_output=capture_output, text=True, check=False
    )
    if check and result.returncode != 0:
        raise RuntimeError(f"Command failed: {cmd}\n{result.stderr}")
    return result


def check_uv_available():
    """Check if uv is available in the system."""
    try:
        result = run_command("uv --version", check=False)
        return result.returncode == 0
    except Exception:
        return False


def get_pytorch_version():
    """Get the installed PyTorch version, e.g. '2.10.0+cu130'."""
    try:
        import torch

        return torch.__version__
    except ImportError:
        return None


def base_version(version_string):
    """'2.10.0+cu130' -> '2.10.0'. Returns None if unparseable."""
    if not version_string:
        return None
    match = re.match(r"(\d+)\.(\d+)\.(\d+)", version_string)
    return match.group(0) if match else None


def version_key(version_string):
    """Sortable key for a release like '0.0.33.post2' or '2.10'."""
    numbers = [int(n) for n in re.findall(r"\d+", version_string)]
    # A plain release outranks nothing, but ".postN" outranks the bare release.
    post = 1 if ".post" in version_string else 0
    return (numbers[:3] + [0, 0, 0])[:3] + [post] + numbers[3:]


def get_pytorch_cuda_tag():
    """Extract the wheel CUDA tag from PyTorch, e.g. 'cu130'. None if CPU/ROCm."""
    try:
        import torch
    except ImportError:
        return None

    local = re.search(r"\+(cu\d+)", torch.__version__)
    if local:
        return local.group(1)
    cuda = getattr(torch.version, "cuda", None)
    if cuda:
        major, _, minor = cuda.partition(".")
        return f"cu{major}{minor or '0'}"
    return None


def wheel_tags_match(filename):
    """Is this wheel installable on the running interpreter and platform?"""
    parts = filename[: -len(".whl")].split("-")
    if len(parts) < 5:
        return False
    python_tag, abi_tag, platform_tag = parts[-3], parts[-2], parts[-1]

    if sys.platform == "win32":
        want_platform = "win_amd64"
    elif sys.platform == "darwin":
        want_platform = "macosx"
    else:
        want_platform = "manylinux"
    if want_platform not in platform_tag:
        return False
    if (
        platform.machine().lower() not in ("amd64", "x86_64")
        and want_platform != "macosx"
    ):
        return False

    # abi3 / py3 wheels work on any newer CPython; cpXY wheels need an exact match.
    exact = f"cp{sys.version_info.major}{sys.version_info.minor}"
    if abi_tag in ("abi3", "none"):
        floor = re.match(r"(?:cp|py)(\d)(\d+)", python_tag)
        if not floor:
            return python_tag in (exact, "py3")
        return (sys.version_info.major, sys.version_info.minor) >= (
            int(floor.group(1)),
            int(floor.group(2)),
        )
    return python_tag == exact


def fetch_index_versions(cuda_tag):
    """Versions on the PyTorch index for this CUDA tag with an installable wheel."""
    url = PYTORCH_INDEX.format(tag=cuda_tag) + "/xformers/"
    try:
        with urllib.request.urlopen(url, timeout=NETWORK_TIMEOUT) as response:
            html = response.read().decode("utf-8", "ignore")
    except (urllib.error.URLError, OSError, TimeoutError) as exc:
        print(f"  Could not read {url}: {exc}")
        return []

    versions = set()
    for filename in re.findall(r"xformers-[^\"<>#]+?\.whl", html):
        filename = filename.split("/")[-1]
        if "%2B" in filename or "+" in filename:
            continue  # nightly / local-version build
        if wheel_tags_match(filename):
            versions.add(filename.split("-")[1])
    return sorted(versions, key=version_key, reverse=True)


def fetch_torch_requirement(xformers_version):
    """The `torch` specifier that this xformers release declares, or None."""
    try:
        url = PYPI_RELEASE.format(version=xformers_version)
        with urllib.request.urlopen(url, timeout=NETWORK_TIMEOUT) as response:
            metadata = json.load(response)
    except (urllib.error.URLError, OSError, TimeoutError, ValueError):
        return None
    for requirement in metadata.get("info", {}).get("requires_dist") or []:
        match = re.match(r"\s*torch\s*(==|>=)\s*([\d.]+)", requirement)
        if match:
            return match.group(1), match.group(2).rstrip(".")
    return None


def requirement_satisfied(requirement, torch_version):
    """Does the installed torch satisfy ('==' | '>=', version)?"""
    operator, wanted = requirement
    installed = version_key(torch_version)
    target = version_key(wanted)
    if operator == "==":
        # A pin of '2.10.0' matches 2.10.0 only; a pin of '2.10' matches 2.10.x.
        depth = len(re.findall(r"\d+", wanted))
        return installed[:depth] == target[:depth]
    return installed >= target


def resolve_xformers(torch_version, cuda_tag):
    """Pick (version, index_url) for the installed torch, or (None, None)."""
    torch_base = base_version(torch_version)
    if not torch_base:
        print(f"Could not parse PyTorch version: {torch_version}")
        return None, None

    if cuda_tag:
        index_url = PYTORCH_INDEX.format(tag=cuda_tag)
        candidates = fetch_index_versions(cuda_tag)
        if candidates:
            print(
                f"  {len(candidates)} candidate(s) on the {cuda_tag} index: {', '.join(candidates)}"
            )
            for candidate in candidates:
                requirement = fetch_torch_requirement(candidate)
                if requirement is None:
                    continue
                if requirement_satisfied(requirement, torch_base):
                    print(
                        f"  xformers {candidate} declares torch{requirement[0]}{requirement[1]}"
                        f" -> matches {torch_base}"
                    )
                    return candidate, index_url
            print(
                f"  No {cuda_tag} build declares compatibility with torch {torch_base}."
            )
        # Index unreachable or nothing compatible: try the static table below.
    else:
        index_url = None
        print("  PyTorch has no CUDA build tag (CPU or ROCm); using PyPI.")

    static = FALLBACK_MAP.get(torch_base)
    if not static:
        for floor, version in FALLBACK_FLOOR:
            if version_key(torch_base) >= version_key(floor):
                static = version
                break
    if static:
        print(f"  Falling back to the static table: xformers {static}")
        return static, index_url

    print(f"  No known xformers release for PyTorch {torch_base}.")
    return None, None


def get_installed_xformers():
    """(version, torch_requirement) for the installed xformers, or None."""
    try:
        import importlib.metadata as metadata

        version = metadata.version("xformers")
    except Exception:
        return None
    requirement = None
    try:
        for entry in metadata.requires("xformers") or []:
            match = re.match(r"\s*torch\s*(==|>=)\s*([\d.]+)", entry)
            if match:
                requirement = (match.group(1), match.group(2).rstrip("."))
                break
    except Exception:
        pass
    return version, requirement


def install_package(package_manager, package_spec, index_url=None):
    """Install one package without letting it touch the existing torch."""
    if package_manager == "uv":
        cmd = f'uv pip install --no-deps "{package_spec}"'
    else:
        cmd = f'"{sys.executable}" -m pip install --no-deps "{package_spec}"'
    if index_url:
        cmd += f" --index-url {index_url}"

    print(f"Running: {cmd}")
    result = run_command(cmd, check=False, capture_output=False)
    return result.returncode == 0


def main():
    """Main installation routine."""
    print("=" * 60)
    print("STARVSRWrapper Installation Script")
    print("=" * 60)

    print("\n[1/3] Checking package manager...")
    package_manager = "uv" if check_uv_available() else "pip"
    print(f"Using package manager: {package_manager}")

    print("\n[2/3] Inspecting the environment...")
    torch_version = get_pytorch_version()
    if not torch_version:
        print("PyTorch not found. Install PyTorch first, then re-run this script.")
        return 0

    cuda_tag = get_pytorch_cuda_tag()
    torch_base = base_version(torch_version)
    print(f"PyTorch version: {torch_version}")
    print(f"CUDA wheel tag:  {cuda_tag or 'none (CPU/ROCm build)'}")
    print(f"Python:          {sys.version.split()[0]} on {sys.platform}")

    installed = get_installed_xformers()
    if installed:
        version, requirement = installed
        if (
            requirement
            and torch_base
            and not requirement_satisfied(requirement, torch_base)
        ):
            print(
                f"xformers {version} is installed but was built for "
                f"torch{requirement[0]}{requirement[1]}, not {torch_base}. Reinstalling."
            )
        else:
            print(
                f"xformers {version} is installed and matches this PyTorch. Nothing to do."
            )
            return 0

    print("\n[3/3] Resolving a matching xformers build...")
    xformers_version, index_url = resolve_xformers(torch_version, cuda_tag)

    if not xformers_version:
        print("\n[!] Could not determine a compatible xformers version.")
        print("    Install it manually, e.g.:")
        hint = f" --index-url {PYTORCH_INDEX.format(tag=cuda_tag)}" if cuda_tag else ""
        print(f'      "{sys.executable}" -m pip install --no-deps xformers{hint}')
        print(
            "    Version matrix: https://github.com/facebookresearch/xformers#installing-xformers"
        )
        return 0

    if not install_package(package_manager, f"xformers=={xformers_version}", index_url):
        print("\n[!] xformers installation failed.")
        print("    Retry manually:")
        hint = f" --index-url {index_url}" if index_url else ""
        print(
            f'      "{sys.executable}" -m pip install --no-deps '
            f'"xformers=={xformers_version}"{hint}'
        )
        return 0

    print(f"\nInstalled xformers {xformers_version}.")
    return 0


if __name__ == "__main__":
    sys.exit(main())