81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
# 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"
|
|
}
|