From 6a3bcd3cd1bdcb454724bfffde65e04bc1fd0bec Mon Sep 17 00:00:00 2001 From: MasterDenis <61897713+MasterDenis@users.noreply.github.com> Date: Sat, 19 Jul 2025 14:59:25 +0300 Subject: [PATCH] Add files via upload --- __init__.py | 17 +++++++++ js/vae_decode_switcher.js | 72 +++++++++++++++++++++++++++++++++++ readme.md | 46 ++++++++++++++++++++++ vae_decode_switcher.py | 80 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 215 insertions(+) create mode 100644 __init__.py create mode 100644 js/vae_decode_switcher.js create mode 100644 readme.md create mode 100644 vae_decode_switcher.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..29c53cf --- /dev/null +++ b/__init__.py @@ -0,0 +1,17 @@ +# __init__.py +# +# Place this file inside your ComfyUI/custom_nodes/VAE-Decode-Switch/ directory. + +import os + +# Get the directory of the current file +NODE_DIR = os.path.dirname(__file__) +# Name of the folder containing our JS files +WEB_DIRECTORY = "js" + +# Import the node mappings from our python file. +from .vae_decode_switcher import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +# Tell ComfyUI that this custom node has web assets to load. +# The __all__ list is what ComfyUI's loader looks for. +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] diff --git a/js/vae_decode_switcher.js b/js/vae_decode_switcher.js new file mode 100644 index 0000000..c751b1a --- /dev/null +++ b/js/vae_decode_switcher.js @@ -0,0 +1,72 @@ +// vae_decode_switcher.js +// +// This file controls the UI for the VAE Decode Switch node. +// It uses a new strategy to truly remove widgets from the layout. + +import { app } from "/scripts/app.js"; + +app.registerExtension({ + name: "VAE-Decode-Switch.UI.NewStrategy", // A new name for a new approach. + async beforeRegisterNodeDef(nodeType, nodeData) { + // Check if this is the node we want to modify. + if (nodeData.name === "VAEDecodeSwitcher") { + console.log("[VAE-Decode-Switch] Applying new layout strategy to VAEDecodeSwitcher."); + + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + onNodeCreated?.apply(this, arguments); + + // Find all the widgets we need to control. + const decoderWidget = this.widgets.find((w) => w.name === "select_decoder"); + const tiledWidgets = [ + "tile_size", + "overlap", + "temporal_size", + "temporal_overlap", + ].map(name => this.widgets.find(w => w.name === name)); + + // Store the original properties of the widgets so we can restore them. + for (const widget of tiledWidgets) { + if (widget) { + // We only need to store the original computeSize. + widget._original_computeSize = widget.computeSize; + } + } + + const toggleVisibility = (isTiled) => { + for (const widget of tiledWidgets) { + if (!widget) continue; + + // Set the widget's hidden property. This prevents it from being drawn, + // which should fix the background flicker issue. + widget.hidden = !isTiled; + + if (isTiled) { + // If we want to show the widget, restore its original size calculation. + widget.computeSize = widget._original_computeSize; + } else { + // If we want to hide it, replace its size calculation with one + // that removes it from the layout, fixing the gap. + widget.computeSize = function() { return [0, -4]; } // -4 accounts for margin + } + } + + // Force the node to immediately re-evaluate its size and redraw. + // This is the key to preventing the gap when switching back to default. + this.size = this.computeSize(); + this.setDirtyCanvas(true, true); + }; + + // When the dropdown changes, run our new visibility logic. + decoderWidget.callback = (value) => { + toggleVisibility(value === "tiled"); + }; + + // Run once on creation to set the correct initial state. + setTimeout(() => { + decoderWidget.callback(decoderWidget.value); + }, 10); + }; + } + }, +}); diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..1d328fc --- /dev/null +++ b/readme.md @@ -0,0 +1,46 @@ +# VAE Decode Switch for ComfyUI + +A simple but powerful utility node for ComfyUI that allows you to switch between the standard `VAE Decode` and the `VAE Decode (Tiled)` nodes without rewiring your workflow. + +This node is designed to keep your graphs clean and make A/B testing different decoding methods quick and easy. + +![VAE Decode Switch Demo](https://i.imgur.com/aBcDeFg.gif) + +--- + +## Features + +* **Seamless Switching:** A simple dropdown menu lets you select either the "default" or "tiled" VAE decoder on the fly. +* **Dynamic UI:** The node's interface is dynamic. When "default" is selected, the extra settings for the tiled decoder (like `tile_size`, `overlap`, etc.) are completely hidden, keeping the node compact and clean. +* **Clean Workflows:** Avoid messy rerouting nodes or having to manually disconnect and reconnect pipelines when you want to try a different decoder. +* **Preserves Inputs:** All standard inputs (`samples`, `vae`) are maintained, and the node correctly passes the required parameters to the selected decoder. + +--- + +## Installation + +1. Navigate to your ComfyUI installation directory. +2. Go into the `custom_nodes` folder. +3. Download or clone this repository into the `custom_nodes` folder. Ensure the final folder structure looks like this: + ``` + ComfyUI/ + └── custom_nodes/ + └── VAE-Decode-Switch/ + ├── __init__.py + ├── vae_decode_switcher.py + └── js/ + └── vae_decode_switcher.js + ``` +4. Restart ComfyUI completely. + +--- + +## Usage + +1. After restarting ComfyUI, you can add the node by: + * Right-clicking on the graph and selecting `Add Node` -> `Logic/Switchers` -> `VAE Decode Switch`. + * Double-clicking the graph and searching for "VAE Decode Switch". +2. Connect your `samples` (LATENT) and `vae` (VAE) inputs as you would with a standard decoder. +3. Use the `select_decoder` dropdown to choose your desired method. + * **default:** Uses the standard `VAEDecode`. The node will shrink to hide the other settings. + * **tiled:** Uses the `VAEDecodeTiled`. The node will expand to show the tiled-specific settings. diff --git a/vae_decode_switcher.py b/vae_decode_switcher.py new file mode 100644 index 0000000..5fc49bb --- /dev/null +++ b/vae_decode_switcher.py @@ -0,0 +1,80 @@ +# vae_decode_switcher.py +# +# Place this file in your ComfyUI/custom_nodes/VAE-Decode-Switch/ directory. + +try: + import torch +except ImportError: + print("[Warning] VAE Decode Switcher: PyTorch is not installed. Error handling will be limited.") + torch = None + +from nodes import VAEDecode, VAEDecodeTiled + +class VAEDecodeSwitcher: + """ + A custom node that allows switching between the standard VAE decoder + and the tiled VAE decoder without changing workflow connections. + """ + + @classmethod + def INPUT_TYPES(s): + """ + Defines all possible input types for the node. + The frontend will hide/show them as needed. + """ + return { + "required": { + "samples": ("LATENT",), + "vae": ("VAE",), + "select_decoder": (["default", "tiled"],), + }, + "optional": { + "tile_size": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 64}), + "overlap": ("INT", {"default": 64, "min": 0, "max": 4096, "step": 8}), + "temporal_size": ("INT", {"default": 64, "min": 1, "max": 4096, "step": 1}), + "temporal_overlap": ("INT", {"default": 8, "min": 0, "max": 4096, "step": 1}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "switch_and_decode" + CATEGORY = "Logic/Switchers" + + def switch_and_decode(self, samples, vae, select_decoder, **kwargs): + print("\n--- VAE Decode Switcher ---") + print(f"Selected Decoder: {select_decoder}") + + if select_decoder == "tiled": + print("Executing Tiled VAE Decode.") + tiled_decoder = VAEDecodeTiled() + try: + result = tiled_decoder.decode( + vae, + samples, + tile_size=kwargs.get('tile_size', 512), + overlap=kwargs.get('overlap', 64), + temporal_size=kwargs.get('temporal_size', 64), + temporal_overlap=kwargs.get('temporal_overlap', 8) + ) + print("Tiled decoding successful.") + return result + except TypeError as e: + print(f"\n[ERROR] VAE Decode Switcher: 'tiled' decoder failed: {e}\n") + if torch: + return (torch.zeros(1, 64, 64, 3, dtype=torch.float32, device="cpu"), ) + else: + raise e + else: + print("Executing Default VAE Decode.") + default_decoder = VAEDecode() + result = default_decoder.decode(vae, samples) + print("Default decoding successful.") + return result + +NODE_CLASS_MAPPINGS = { + "VAEDecodeSwitcher": VAEDecodeSwitcher +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "VAEDecodeSwitcher": "VAE Decode Switch" +}