mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-13 15:50:49 +08:00
37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
import torch
|
|
|
|
class AutoCUDNNToggle:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"enable_cudnn": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"latent": ("LATENT",),
|
|
"audio": ("AUDIO",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "AUDIO")
|
|
RETURN_NAMES = ("latent_out", "audio_out")
|
|
FUNCTION = "toggle"
|
|
CATEGORY = "advanced/utils"
|
|
|
|
def toggle(self, enable_cudnn, latent=None, audio=None):
|
|
# Set CuDNN state
|
|
torch.backends.cudnn.enabled = enable_cudnn
|
|
|
|
# Auto-detect active path
|
|
if latent is not None:
|
|
print(f"[CuDNN] Latent mode | Enabled: {enable_cudnn}")
|
|
return (latent, None)
|
|
elif audio is not None:
|
|
print(f"[CuDNN] Audio mode | Enabled: {enable_cudnn}")
|
|
return (None, audio)
|
|
else:
|
|
raise ValueError("No valid input connected - must connect either latent OR audio")
|
|
|
|
NODE_CLASS_MAPPINGS = {"CFZ-CUDNNToggle": AutoCUDNNToggle}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"CFZ-CUDNNToggle": "CFZ CuDNN Toggle"}
|