wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-18 20:38:15 +08:00
Code Issues Actions 27 Packages Projects Releases Wiki Activity

Preserve platform-specific PyTorch builds during update (#14902)

Browse Source 
When running `pip install -r requirements.txt` during an update, the bare
`torch`/`torchvision`/`torchaudio` entries can cause pip to overwrite
ROCm, DirectML, XPU, or other platform-specific torch builds with the
default PyPI (CUDA) packages.

This adds pip constraints to pin torch packages to their current versions
when a non-standard build variant (e.g. +rocm, +xpu) is detected,
preventing the update from breaking the user's GPU setup.

Fixes #14902

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Kaihui-AMD 2026-07-13 10:39:59 +08:00
parent c96fcddb81
commit b42cc81be1
1 changed files with 71 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

73
.ci/update_windows/update.py
View File

@ -4,6 +4,8 @@ import sys
import os
import shutil
import filecmp
import importlib.util
import tempfile
def pull(repo, remote_name='origin', branch='master'):
for remote in repo.remotes:
@ -144,10 +146,77 @@ if self_update and not files_equal(update_py_path, repo_update_py_path) and file
shutil.copy(repo_update_py_path, os.path.join(cur_path, "update_new.py"))
exit()
if not os.path.exists(req_path) or not files_equal(repo_req_path, req_path):
TORCH_PACKAGES = {'torch', 'torchvision', 'torchaudio'}
def get_installed_torch_versions():
"""Get current versions of torch packages without importing them.
Returns a dict of {package_name: version_string} for installed torch packages
that have a platform-specific build variant (e.g. +rocm, +xpu).
"""
result = {}
for pkg_name in TORCH_PACKAGES:
try:
spec = importlib.util.find_spec(pkg_name)
if spec is None:
continue
for folder in spec.submodule_search_locations:
ver_file = os.path.join(folder, "version.py")
if os.path.isfile(ver_file):
ver_spec = importlib.util.spec_from_file_location(f"{pkg_name}_version", ver_file)
module = importlib.util.module_from_spec(ver_spec)
ver_spec.loader.exec_module(module)
version = module.__version__
if '+' in version:
result[pkg_name] = version
break
except Exception:
pass
return result
def pip_install_requirements(requirements_path):
"""Install requirements while preserving platform-specific torch builds.
If torch is installed with a non-standard build variant (ROCm, DirectML, XPU, etc.),
uses pip constraints to pin torch packages and prevent them from being overwritten
by default PyPI builds.
"""
import subprocess
torch_versions = get_installed_torch_versions()
constraint_file = None
if torch_versions:
variant = list(torch_versions.values())[0].split('+')[1]
print("Detected platform-specific PyTorch build: +{}".format(variant)) # noqa: T201
print("Constraining torch packages to prevent overwrite: {}".format( # noqa: T201
', '.join('{}=={}'.format(k, v) for k, v in torch_versions.items())
))
constraint_file = tempfile.NamedTemporaryFile(
mode='w', suffix='.txt', delete=False, prefix='torch_constraints_'
)
for pkg, ver in torch_versions.items():
constraint_file.write("{}=={}\n".format(pkg, ver))
constraint_file.close()
try:
subprocess.check_call([sys.executable, '-s', '-m', 'pip', 'install', '-r', repo_req_path])
cmd = [sys.executable, '-s', '-m', 'pip', 'install', '-r', requirements_path]
if constraint_file:
cmd += ['-c', constraint_file.name]
subprocess.check_call(cmd)
finally:
if constraint_file:
try:
os.unlink(constraint_file.name)
except OSError:
pass
if not os.path.exists(req_path) or not files_equal(repo_req_path, req_path):
try:
pip_install_requirements(repo_req_path)
shutil.copy(repo_req_path, req_path)
except:
pass