wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-25 05:40:15 +08:00
Code Issues Actions 16 Packages Projects Releases Wiki Activity

Update model_management.py

Browse Source
This commit is contained in:
patientx 2025-02-14 12:30:19 +03:00 committed by GitHub
parent 4d66aa9709
commit 99d2824d5a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 18 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
comfy/model_management.py
View File

@ -25,7 +25,6 @@ import sys
import platform import platform
import weakref import weakref
import gc import gc
import zluda
class VRAMState(Enum): class VRAMState(Enum):
DISABLED = 0 #No vram present: no need to move models to vram DISABLED = 0 #No vram present: no need to move models to vram
@ -237,10 +236,23 @@ try:
except: except:
pass pass
#if ENABLE_PYTORCH_ATTENTION:
# torch.backends.cuda.enable_math_sdp(True) try:
# torch.backends.cuda.enable_flash_sdp(True) if is_amd():
# torch.backends.cuda.enable_mem_efficient_sdp(True) arch = torch.cuda.get_device_properties(get_torch_device()).gcnArchName
logging.info("AMD arch: {}".format(arch))
if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
if int(torch_version[0]) >= 2 and int(torch_version[2]) >= 7: # works on 2.6 but doesn't actually seem to improve much
if arch in ["gfx1100"]: #TODO: more arches
ENABLE_PYTORCH_ATTENTION = True
except:
pass
if ENABLE_PYTORCH_ATTENTION:
torch.backends.cuda.enable_math_sdp(True)
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
try: try:
if is_nvidia() and args.fast: if is_nvidia() and args.fast:
@ -307,6 +319,7 @@ try:
except: except:
logging.warning("Could not pick default device.") logging.warning("Could not pick default device.")
current_loaded_models = [] current_loaded_models = []
def module_size(module): def module_size(module):