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

Update zluda.py

Browse Source
This commit is contained in:
patientx 2025-01-14 20:03:55 +03:00 committed by GitHub
parent c4861c74d4
commit 4f01f72bed
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 15 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
comfy/zluda.py
View File

@ -1,33 +1,16 @@
try: import torch
torch_device_name = get_torch_device_name(get_torch_device())
if "[ZLUDA]" in torch_device_name:
_torch_stft = torch.stft
_torch_istft = torch.istft
def z_stft(input: torch.Tensor, window: torch.Tensor, *args, **kwargs): if torch.cuda.get_device_name().endswith("[ZLUDA]"):
return _torch_stft(input=input.cpu(), window=window.cpu(), *args, **kwargs).to(input.device) print(" ")
print("***----------------------ZLUDA-----------------------------***")
def z_istft(input: torch.Tensor, window: torch.Tensor, *args, **kwargs): print(" :: ZLUDA detected, disabling non-supported functions.")
return _torch_istft(input=input.cpu(), window=window.cpu(), *args, **kwargs).to(input.device) torch.backends.cudnn.enabled = False
print(" :: CuDNN, flash_sdp, math_sdp, mem_efficient_sdp disabled) ")
def z_jit(f, *_, **__): torch.backends.cuda.enable_flash_sdp(False)
f.graph = torch._C.Graph() torch.backends.cuda.enable_math_sdp(True)
return f torch.backends.cuda.enable_mem_efficient_sdp(False)
print("***--------------------------------------------------------***")
# hijacks print(" :: Device:", torch.cuda.get_device_name())
torch.stft = z_stft print(" ")
torch.istft = z_istft else:
torch.jit.script = z_jit print(" :: ZLUDA isn't detected, please try patching it.")
print(" ")
print("***----------------------ZLUDA--------------------------***")
print(" :: ZLUDA detected, disabling non-supported functions.")
torch.backends.cudnn.enabled = False
print(" :: (cuDNN, flash_sdp, mem_efficient_sdp disabled) ")
torch.backends.cuda.enable_flash_sdp(False)
torch.backends.cuda.enable_math_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(False)
print("***-----------------------------------------------------***")
print(" :: Device:", torch_device_name)
print(" ")