wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-26 14:20:27 +08:00
Code Issues Actions 10 Packages Projects Releases Wiki Activity

ROCm 6.0 seems to require get_device_name to be called before memory methods in order to return valid data

Browse Source
This commit is contained in:
doctorpangloss 2024-05-22 22:09:07 -07:00
parent bb159a219e
commit 35cf996b68
2 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
comfy/model_management.py
View File

@ -132,13 +132,15 @@ def get_total_memory(dev=None, torch_total_too=False):
_, mem_total_cuda = torch.cuda.mem_get_info(dev) _, mem_total_cuda = torch.cuda.mem_get_info(dev)
mem_total_torch = mem_reserved mem_total_torch = mem_reserved
mem_total = mem_total_cuda mem_total = mem_total_cuda
if torch_total_too: if torch_total_too:
return (mem_total, mem_total_torch) return mem_total, mem_total_torch
else: else:
return mem_total return mem_total
# we're required to call get_device_name early on to initialize the methods get_total_memory will call
if torch.cuda.is_available() and hasattr(torch.version, "hip") and torch.version.hip is not None:
logging.info(f"Detected HIP device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
total_vram = get_total_memory(get_torch_device()) / (1024 * 1024) total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
total_ram = psutil.virtual_memory().total / (1024 * 1024) total_ram = psutil.virtual_memory().total / (1024 * 1024)
logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram)) logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))

2
setup.py
View File

@ -162,7 +162,7 @@ def dependencies(force_nightly: bool = False) -> List[str]:
index_urls_selected = [stable for (stable, _) in index_urls] index_urls_selected = [stable for (stable, _) in index_urls]
_alternative_indices_selected = [stable for (stable, _) in _alternative_indices] _alternative_indices_selected = [stable for (stable, _) in _alternative_indices]
try: try:
# pip 23 # pip 23, 24
finder = PackageFinder.create(LinkCollector(session, SearchScope([], index_urls_selected, no_index=False)), finder = PackageFinder.create(LinkCollector(session, SearchScope([], index_urls_selected, no_index=False)),
SelectionPreferences(allow_yanked=False, prefer_binary=False, SelectionPreferences(allow_yanked=False, prefer_binary=False,
allow_all_prereleases=True)) allow_all_prereleases=True))