File size: 4,434 Bytes
f2b53cf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """Loader for TAESD-family tiny VAE decoders that ComfyUI core can't build yet.
Core's comfy.taesd.taesd.Decoder hardcodes a 64-wide stack with 3 upsample stages, so
checkpoints that differ fail to load — e.g. taeh3 (MiniMax H3): 24 latent channels, 96
wide for the first half, 4 upsample stages for the VAE's 16x spatial ratio.
These decoders are a flat nn.Sequential whose state dict keys are the positional module
indices, so the exact architecture is recoverable from the checkpoint: indices carrying
`N.conv.0.weight` are Blocks, indices carrying `N.weight` are convs (bias-less ones sit
right after an upsample), and the gaps are the parameterless modules.
"""
import logging
import torch
import torch.nn as nn
import comfy.model_management
import comfy.utils
from comfy.taesd.taesd import Block, Clamp, conv
def build_tae_decoder(sd):
by_index = {}
for k, v in sd.items():
head, _, rest = k.partition(".")
if not head.isdigit():
raise ValueError(f"not a flat TAE decoder state dict (unexpected key '{k}')")
by_index.setdefault(int(head), {})[rest] = v
modules = []
for i in range(max(by_index) + 1):
entry = by_index.get(i)
if entry is None:
# index 0 is the input Clamp, 2 the ReLU after the input conv, the rest are upsamples
modules.append(Clamp() if i == 0 else nn.ReLU() if i == 2 else nn.Upsample(scale_factor=2))
elif "conv.0.weight" in entry:
w = entry["conv.0.weight"]
# only pass the kwarg when it's needed — older ComfyUI has no midblock-GN variant
if "pool.0.weight" in entry:
modules.append(Block(w.shape[1], w.shape[0], use_midblock_gn=True))
else:
modules.append(Block(w.shape[1], w.shape[0]))
elif "weight" in entry:
w = entry["weight"]
modules.append(conv(w.shape[1], w.shape[0], bias="bias" in entry))
else:
raise ValueError(f"unrecognized TAE decoder module at index {i}: {sorted(entry)}")
return nn.Sequential(*modules)
class TinyVAEDecoder:
"""Decode-only tiny VAE. Output is [0, 1], matching the TAE family's convention."""
def __init__(self, sd, device=None, dtype=None):
# keys may carry a "taesd_decoder."/"decoder." prefix; strip whatever is common
prefix = ""
first = next(iter(sd))
if not first.split(".")[0].isdigit():
prefix = first.split(".")[0] + "."
sd = {k[len(prefix):]: v for k, v in sd.items() if k.startswith(prefix)}
self.device = device if device is not None else comfy.model_management.vae_device()
self.dtype = dtype if dtype is not None else comfy.model_management.vae_dtype(
self.device, [torch.float16, torch.bfloat16])
self.model = build_tae_decoder(sd)
self.model.load_state_dict(sd)
self.model.eval().to(device=self.device, dtype=self.dtype)
self.latent_channels = self.model[1].weight.shape[1]
self.upscale_ratio = 2 ** sum(isinstance(m, nn.Upsample) for m in self.model)
def decode(self, latent):
"""[B, C, H, W] -> [B, 3, H*ratio, W*ratio], float32 on the input device."""
out = self.model(latent.to(device=self.device, dtype=self.dtype))
return out.to(device=latent.device, dtype=torch.float32)
def decode_video(self, latent, frame_indices=None):
"""[B, C, T, H, W] -> [T, H*ratio, W*ratio, 3]. Decodes one frame at a time —
at 16x the full-resolution activations are the memory peak, not the weights."""
x = latent[0]
indices = range(x.shape[1]) if frame_indices is None else frame_indices
frames = [self.decode(x[:, t].unsqueeze(0))[0].movedim(0, -1) for t in indices]
return torch.stack(frames, dim=0)
def load_tiny_vae_decoder(name, device=None, dtype=None):
"""Load by vae_approx filename. Returns None (and logs) if it can't be used."""
import folder_paths
path = folder_paths.get_full_path("vae_approx", name)
if path is None:
logging.warning(f"[KJ TinyVAE] '{name}' not found in models/vae_approx.")
return None
try:
sd = comfy.utils.load_torch_file(path, safe_load=True)
return TinyVAEDecoder(sd, device=device, dtype=dtype)
except Exception as e:
logging.warning(f"[KJ TinyVAE] Could not load '{name}': {e}")
return None
|