ComfyUI/cfz/cfz_cudnn.toggle.py
2025-06-02 17:45:49 +03:00

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"}