diff --git a/cfz/cfz_cudnn.toggle.py b/cfz/cfz_cudnn.toggle.py new file mode 100644 index 000000000..419d17ede --- /dev/null +++ b/cfz/cfz_cudnn.toggle.py @@ -0,0 +1,36 @@ +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"}