Upload 2 files

7798da9 verified 7 months ago

19.7 kB

import torch
import torch.nn as nn
import torch.nn.functional as F

import comfy.utils
import comfy.ops
import comfy.model_management
import folder_paths


# ============================================================
# Layers (Comfy-style: disable_weight_init)
# ============================================================

def conv(n_in, n_out, **kwargs):
    return comfy.ops.disable_weight_init.Conv2d(n_in, n_out, 3, padding=1, **kwargs)


class Clamp(nn.Module):
    def forward(self, x):
        return torch.tanh(x / 3) * 3


class Block(nn.Module):
    def __init__(self, n_in, n_out, use_midblock_gn=False):
        super().__init__()
        self.conv = nn.Sequential(
            conv(n_in, n_out), nn.ReLU(),
            conv(n_out, n_out), nn.ReLU(),
            conv(n_out, n_out),
        )
        self.skip = comfy.ops.disable_weight_init.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
        self.fuse = nn.ReLU()

        self.pool = None
        if use_midblock_gn:
            conv1x1 = lambda a, b: comfy.ops.disable_weight_init.Conv2d(a, b, 1, bias=False)
            n_gn = n_in * 4
            self.pool = nn.Sequential(
                conv1x1(n_in, n_gn),
                comfy.ops.disable_weight_init.GroupNorm(4, n_gn),
                nn.ReLU(inplace=True),
                conv1x1(n_gn, n_in),
            )

    def forward(self, x):
        if self.pool is not None:
            x = x + self.pool(x)
        return self.fuse(self.conv(x) + self.skip(x))


# ============================================================
# TAESD scale8/scale16 builders
# Your file is scale8 (encoder final conv at layers.14)
# ============================================================

def build_encoder_scale8(latent_channels, pool_blocks):
    def B(idx): return Block(64, 64, use_midblock_gn=(idx in pool_blocks))
    return nn.Sequential(
        conv(3, 64),                # 0
        B(1),                       # 1
        conv(64, 64, stride=2, bias=False),  # 2
        B(3), B(4), B(5),           # 3-5
        conv(64, 64, stride=2, bias=False),  # 6
        B(7), B(8), B(9),           # 7-9
        conv(64, 64, stride=2, bias=False),  # 10
        B(11), B(12), B(13),        # 11-13
        conv(64, latent_channels),  # 14
    )


def build_decoder_scale8(latent_channels, pool_blocks):
    def B(idx): return Block(64, 64, use_midblock_gn=(idx in pool_blocks))
    return nn.Sequential(
        Clamp(),                    # 0 (no weights)
        conv(latent_channels, 64),  # 1
        nn.ReLU(),                  # 2
        B(3), B(4), B(5),           # 3-5
        nn.Upsample(scale_factor=2),            # 6
        conv(64, 64, bias=False),   # 7
        B(8), B(9), B(10),          # 8-10
        nn.Upsample(scale_factor=2),            # 11
        conv(64, 64, bias=False),   # 12
        B(13), B(14), B(15),        # 13-15
        nn.Upsample(scale_factor=2),            # 16
        conv(64, 64, bias=False),   # 17
        B(18),                      # 18
        conv(64, 3),                # 19
    )


def build_encoder_scale16(latent_channels, pool_blocks):
    def B(idx): return Block(64, 64, use_midblock_gn=(idx in pool_blocks))
    return nn.Sequential(
        conv(3, 64),                # 0
        B(1),                       # 1
        conv(64, 64, stride=2, bias=False),  # 2
        B(3), B(4), B(5),           # 3-5
        conv(64, 64, stride=2, bias=False),  # 6
        B(7), B(8), B(9),           # 7-9
        conv(64, 64, stride=2, bias=False),  # 10
        B(11), B(12), B(13),        # 11-13
        conv(64, 64, stride=2, bias=False),  # 14
        B(15), B(16), B(17),        # 15-17
        conv(64, latent_channels),  # 18
    )


def build_decoder_scale16(latent_channels, pool_blocks):
    def B(idx): return Block(64, 64, use_midblock_gn=(idx in pool_blocks))
    return nn.Sequential(
        Clamp(),                    # 0
        conv(latent_channels, 64),  # 1
        nn.ReLU(),                  # 2
        B(3), B(4), B(5),           # 3-5
        nn.Upsample(scale_factor=2),            # 6
        conv(64, 64, bias=False),   # 7
        B(8), B(9), B(10),          # 8-10
        nn.Upsample(scale_factor=2),            # 11
        conv(64, 64, bias=False),   # 12
        B(13), B(14), B(15),        # 13-15
        nn.Upsample(scale_factor=2),            # 16
        conv(64, 64, bias=False),   # 17
        B(18), B(19), B(20),        # 18-20
        nn.Upsample(scale_factor=2),            # 21
        conv(64, 64, bias=False),   # 22
        B(23),                      # 23
        conv(64, 3),                # 24
    )


# ============================================================
# Packed latents (auto-pad so it never errors)
# ============================================================

def unpack_packed_latents(x, latent_channels):
    # [B, C*4, H, W] -> [B, C, H*2, W*2]
    if x.ndim == 4 and x.shape[1] == latent_channels * 4:
        return (
            x.reshape(x.shape[0], latent_channels, 2, 2, x.shape[-2], x.shape[-1])
             .permute(0, 1, 4, 2, 5, 3)
             .reshape(x.shape[0], latent_channels, x.shape[-2] * 2, x.shape[-1] * 2)
        )
    return x


def pack_packed_latents(z, latent_channels):
    # [B, C, H, W] -> [B, C*4, H//2, W//2]
    if z.ndim == 4 and z.shape[1] == latent_channels:
        h, w = z.shape[-2], z.shape[-1]
        pad_h = h & 1
        pad_w = w & 1
        if pad_h or pad_w:
            z = F.pad(z, (0, pad_w, 0, pad_h), mode="replicate")
            h, w = z.shape[-2], z.shape[-1]

        return (
            z.reshape(z.shape[0], latent_channels, h // 2, 2, w // 2, 2)
             .permute(0, 1, 3, 5, 2, 4)
             .reshape(z.shape[0], latent_channels * 4, h // 2, w // 2)
        )
    return z


def pad_nchw_to_multiple(x, multiple):
    # replicate pad right/bottom so any size works
    _, _, h, w = x.shape
    pad_h = (multiple - (h % multiple)) % multiple
    pad_w = (multiple - (w % multiple)) % multiple
    if pad_h or pad_w:
        x = F.pad(x, (0, pad_w, 0, pad_h), mode="replicate")
    return x


# ============================================================
# Key conversion for your file format:
# encoder.layers.N.* and decoder.layers.N.*
# decoder layers must shift +1 because our decoder has Clamp() at index 0.
# ============================================================

def normalize_state_dict(sd_raw):
    keys = list(sd_raw.keys())

    # Already comfy split format?
    if any(k.startswith("taesd_encoder.") for k in keys) or any(k.startswith("taesd_decoder.") for k in keys):
        return sd_raw

    out = {}

    # Diffusers "encoder.layers.* / decoder.layers.*"
    if any(k.startswith("encoder.layers.") for k in keys) or any(k.startswith("decoder.layers.") for k in keys):
        for k, v in sd_raw.items():
            if k.startswith("encoder.layers."):
                # encoder.layers.N.xxx -> taesd_encoder.N.xxx
                out["taesd_encoder." + k[len("encoder.layers."):]] = v
            elif k.startswith("decoder.layers."):
                # decoder.layers.N.xxx -> taesd_decoder.(N+1).xxx  (Clamp at 0)
                rest = k[len("decoder.layers."):]
                parts = rest.split(".", 1)
                try:
                    n = int(parts[0])
                    n2 = n + 1
                    tail = parts[1] if len(parts) > 1 else ""
                    out_key = f"taesd_decoder.{n2}" + (("." + tail) if tail else "")
                    out[out_key] = v
                except Exception:
                    # fallback, keep
                    out[k] = v
            else:
                out[k] = v
        return out

    # Fallback: encoder./decoder. (numeric) — if decoder.0.weight looks like [64,C,3,3], offset it too
    if any(k.startswith("encoder.") for k in keys) or any(k.startswith("decoder.") for k in keys):
        decoder_needs_offset = False
        w0 = sd_raw.get("decoder.0.weight", None)
        if isinstance(w0, torch.Tensor) and w0.ndim == 4 and w0.shape[0] == 64 and w0.shape[2:] == (3, 3):
            decoder_needs_offset = True

        for k, v in sd_raw.items():
            if k.startswith("encoder."):
                out["taesd_encoder." + k[len("encoder."):]] = v
            elif k.startswith("decoder."):
                rest = k[len("decoder."):]
                if decoder_needs_offset:
                    parts = rest.split(".", 1)
                    if parts[0].isdigit():
                        n = int(parts[0]) + 1
                        tail = parts[1] if len(parts) > 1 else ""
                        out_key = f"taesd_decoder.{n}" + (("." + tail) if tail else "")
                        out[out_key] = v
                    else:
                        out["taesd_decoder." + rest] = v
                else:
                    out["taesd_decoder." + rest] = v
            else:
                out[k] = v
        return out

    # Unknown layout: return as-is (Dump node will show keys)
    return sd_raw


def split_encoder_decoder(sd):
    enc = {k[len("taesd_encoder."):]: v for k, v in sd.items() if k.startswith("taesd_encoder.")}
    dec = {k[len("taesd_decoder."):]: v for k, v in sd.items() if k.startswith("taesd_decoder.")}
    return enc, dec


def pool_blocks_from_sd(part_sd):
    blocks = set()
    for k in part_sd.keys():
        if ".pool.0.weight" in k or ".pool.0.bias" in k:
            head = k.split(".", 1)[0]
            if head.isdigit():
                blocks.add(int(head))
    return blocks


def infer_latent_channels_from_decoder(dec_sd):
    # Find smallest-index conv weight that looks like decoder input conv: [64, C, 3, 3]
    candidates = []
    for k, v in dec_sd.items():
        if not isinstance(v, torch.Tensor) or v.ndim != 4:
            continue
        head = k.split(".", 1)[0]
        if head.isdigit() and v.shape[0] == 64 and v.shape[2:] == (3, 3):
            candidates.append((int(head), int(v.shape[1])))
    if not candidates:
        raise RuntimeError("Could not infer latent_channels from decoder weights.")
    candidates.sort(key=lambda t: t[0])
    return candidates[0][1]


def detect_layout(enc_sd, latent_channels):
    # Your file has encoder.layers.14.* -> after normalize it's "14.weight"
    if "14.weight" in enc_sd:
        w = enc_sd["14.weight"]
        if isinstance(w, torch.Tensor) and w.ndim == 4 and w.shape[0] == latent_channels and w.shape[1] == 64:
            return "scale8"
    if "18.weight" in enc_sd:
        w = enc_sd["18.weight"]
        if isinstance(w, torch.Tensor) and w.ndim == 4 and w.shape[0] == latent_channels and w.shape[1] == 64:
            return "scale16"

    # Fallback: find earliest [C,64,3,3] conv in encoder
    best = None
    for k, v in enc_sd.items():
        if not isinstance(v, torch.Tensor) or v.ndim != 4:
            continue
        head = k.split(".", 1)[0]
        if head.isdigit() and v.shape[0] == latent_channels and v.shape[1] == 64 and v.shape[2:] == (3, 3):
            idx = int(head)
            best = idx if best is None else min(best, idx)
    if best is None:
        raise RuntimeError("Could not detect encoder layout (scale8 vs scale16).")
    return "scale8" if best <= 14 else "scale16"


# ============================================================
# Core model (PR behavior: decode -> [-1,1], encode -> packed for taef2)
# ============================================================

class TAESDCore(nn.Module):
    def __init__(self, encoder, decoder, latent_channels, is_taef2):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.latent_channels = int(latent_channels)
        self.is_taef2 = bool(is_taef2)

        self.vae_scale = nn.Parameter(torch.tensor(1.0))
        self.vae_shift = nn.Parameter(torch.tensor(0.0))

    @torch.inference_mode()
    def decode(self, x):
        x = unpack_packed_latents(x, self.latent_channels)
        x = (x - self.vae_shift) * self.vae_scale
        x_sample = self.decoder(x)
        # decoder output in [0,1] -> [-1,1]
        return x_sample.sub(0.5).mul(2.0)

    @torch.inference_mode()
    def encode(self, x):
        # x is [-1,1] -> encoder expects [0,1]
        z = (self.encoder(x * 0.5 + 0.5) / self.vae_scale) + self.vae_shift
        if self.is_taef2:
            z = pack_packed_latents(z, self.latent_channels)
        return z


def load_core(path, device, dtype):
    sd_raw = comfy.utils.load_torch_file(path, safe_load=True)
    sd = normalize_state_dict(sd_raw)
    enc_sd, dec_sd = split_encoder_decoder(sd)

    if not enc_sd or not dec_sd:
        sample = list(sd_raw.keys())[:40]
        raise RuntimeError(
            "Could not split encoder/decoder weights.\n"
            "Use Dump VAE Keys node and paste first ~40 keys.\n"
            f"First keys: {sample}"
        )

    enc_pool = pool_blocks_from_sd(enc_sd)
    dec_pool = pool_blocks_from_sd(dec_sd)

    latent_channels = infer_latent_channels_from_decoder(dec_sd)
    layout = detect_layout(enc_sd, latent_channels)

    # Flux2 taef2 packed-latents heuristic (matches your file):
    has_midblock_gn = (len(enc_pool) > 0) or (len(dec_pool) > 0)
    is_taef2 = (latent_channels == 32) and has_midblock_gn

    if layout == "scale8":
        encoder = build_encoder_scale8(latent_channels, enc_pool)
        decoder = build_decoder_scale8(latent_channels, dec_pool)
        base_downscale = 8
    else:
        encoder = build_encoder_scale16(latent_channels, enc_pool)
        decoder = build_decoder_scale16(latent_channels, dec_pool)
        base_downscale = 16

    # Load in fp32 first, then cast (more robust)
    core = TAESDCore(encoder, decoder, latent_channels, is_taef2)
    core.encoder.load_state_dict(enc_sd, strict=False)
    core.decoder.load_state_dict(dec_sd, strict=False)

    core = core.to(device=device, dtype=dtype).eval()
    for p in core.parameters():
        p.requires_grad_(False)

    core._base_downscale = base_downscale
    return core


# ============================================================
# Comfy VAE interface object
# ============================================================

class TAEF2VAE:
    def __init__(self, weights_path, device, dtype):
        self.device = device
        self.dtype = dtype
        self.core = load_core(weights_path, device=device, dtype=dtype)

        # packed latents halves latent H/W again -> effective downscale doubles
        self.downscale_ratio = self.core._base_downscale * (2 if self.core.is_taef2 else 1)

        print(
            f"[TAEF2] Loaded: {weights_path} | latent_channels={self.core.latent_channels} "
            f"| is_taef2={self.core.is_taef2} | base_downscale={self.core._base_downscale} "
            f"| effective_downscale={self.downscale_ratio}"
        )

    @torch.inference_mode()
    def decode(self, latents):
        x = latents.to(device=self.device, dtype=self.dtype)
        img = self.core.decode(x)                    # NCHW in [-1,1]
        img = img.clamp(-1, 1).add(1.0).mul(0.5)     # -> [0,1]
        return img.to(torch.float32).permute(0, 2, 3, 1).contiguous()  # NHWC float32

    @torch.inference_mode()
    def encode(self, pixels):
        # pixels NHWC [0,1]
        x = pixels[..., :3].permute(0, 3, 1, 2).contiguous()
        x = x.to(device=self.device, dtype=self.dtype).clamp(0, 1).mul(2.0).sub(1.0)  # -> [-1,1]

        # Make it behave like base VAE: pad to required multiple so any size works
        x = pad_nchw_to_multiple(x, self.downscale_ratio)

        z = self.core.encode(x)                      # packed if taef2
        return z.to(torch.float32)

    def decode_tiled(self, latents, **kwargs):
        return self.decode(latents)

    def encode_tiled(self, pixels, **kwargs):
        return self.encode(pixels)

    def spacial_compression_decode(self):
        return self.downscale_ratio

    def spacial_compression_encode(self):
        return self.downscale_ratio

    def temporal_compression_decode(self):
        return None

    def temporal_compression_encode(self):
        return None


# ============================================================
# Nodes
# ============================================================

def _list_vae_files():
    vae_files = folder_paths.get_filename_list("vae")
    approx_files = folder_paths.get_filename_list("vae_approx")
    return sorted(set(vae_files + approx_files))

def _resolve_vae_path(fname):
    path = folder_paths.get_full_path("vae_approx", fname)
    if path is None:
        path = folder_paths.get_full_path("vae", fname)
    return path


class LoadTAEF2VAE:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "weights": (_list_vae_files(),),
                "dtype": (["bf16", "fp16", "fp32"], {"default": "bf16"}),
            }
        }

    RETURN_TYPES = ("VAE",)
    FUNCTION = "load"
    CATEGORY = "latent/vae"

    def load(self, weights, dtype):
        path = _resolve_vae_path(weights)
        if path is None:
            raise FileNotFoundError(f"Could not find weights file: {weights}")

        device = comfy.model_management.get_torch_device()
        if dtype == "bf16":
            tdtype = torch.bfloat16
        elif dtype == "fp16":
            tdtype = torch.float16
        else:
            tdtype = torch.float32

        return (TAEF2VAE(path, device=device, dtype=tdtype),)


class DumpVAEKeys:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "weights": (_list_vae_files(),),
                "include_shapes": ("BOOLEAN", {"default": True}),
                "sort_keys": ("BOOLEAN", {"default": True}),
                "max_lines": ("INT", {"default": 0, "min": 0, "max": 200000}),
            }
        }

    RETURN_TYPES = ("STRING",)
    FUNCTION = "dump"
    CATEGORY = "utils/debug"

    def dump(self, weights, include_shapes, sort_keys, max_lines):
        path = _resolve_vae_path(weights)
        if path is None:
            raise FileNotFoundError(f"Could not find weights file: {weights}")

        sd = comfy.utils.load_torch_file(path, safe_load=True)
        keys = list(sd.keys())
        if sort_keys:
            keys.sort()

        lines = []
        if include_shapes:
            for k in keys:
                v = sd[k]
                if isinstance(v, torch.Tensor):
                    lines.append(f"{k}\t{tuple(v.shape)}\t{str(v.dtype)}")
                else:
                    lines.append(f"{k}\t{type(v)}")
        else:
            lines = keys

        if max_lines and len(lines) > max_lines:
            head = lines[:max_lines]
            head.append(f"... TRUNCATED: total_keys={len(lines)} (showing first {max_lines}) ...")
            lines = head

        text = "\n".join(lines)
        return {"ui": {"text": [text]}, "result": (text,)}


NODE_CLASS_MAPPINGS = {
    "LoadTAEF2VAE": LoadTAEF2VAE,
    "DumpVAEKeys": DumpVAEKeys,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LoadTAEF2VAE": "Load TAEF2 (Flux2 Tiny VAE)",
    "DumpVAEKeys": "Dump VAE Keys (as String)",
}

	import torch
	import torch.nn as nn
	import torch.nn.functional as F

	import comfy.utils
	import comfy.ops
	import comfy.model_management
	import folder_paths


	# ============================================================
	# Layers (Comfy-style: disable_weight_init)
	# ============================================================

	def conv(n_in, n_out, **kwargs):
	return comfy.ops.disable_weight_init.Conv2d(n_in, n_out, 3, padding=1, **kwargs)


	class Clamp(nn.Module):
	def forward(self, x):
	return torch.tanh(x / 3) * 3


	class Block(nn.Module):
	def __init__(self, n_in, n_out, use_midblock_gn=False):
	super().__init__()
	self.conv = nn.Sequential(
	conv(n_in, n_out), nn.ReLU(),
	conv(n_out, n_out), nn.ReLU(),
	conv(n_out, n_out),
	)
	self.skip = comfy.ops.disable_weight_init.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
	self.fuse = nn.ReLU()

	self.pool = None
	if use_midblock_gn:
	conv1x1 = lambda a, b: comfy.ops.disable_weight_init.Conv2d(a, b, 1, bias=False)
	n_gn = n_in * 4
	self.pool = nn.Sequential(
	conv1x1(n_in, n_gn),
	comfy.ops.disable_weight_init.GroupNorm(4, n_gn),
	nn.ReLU(inplace=True),
	conv1x1(n_gn, n_in),
	)

	def forward(self, x):
	if self.pool is not None:
	x = x + self.pool(x)
	return self.fuse(self.conv(x) + self.skip(x))


	# ============================================================
	# TAESD scale8/scale16 builders
	# Your file is scale8 (encoder final conv at layers.14)
	# ============================================================

	def build_encoder_scale8(latent_channels, pool_blocks):
	def B(idx): return Block(64, 64, use_midblock_gn=(idx in pool_blocks))
	return nn.Sequential(
	conv(3, 64), # 0
	B(1), # 1
	conv(64, 64, stride=2, bias=False), # 2
	B(3), B(4), B(5), # 3-5
	conv(64, 64, stride=2, bias=False), # 6
	B(7), B(8), B(9), # 7-9
	conv(64, 64, stride=2, bias=False), # 10
	B(11), B(12), B(13), # 11-13
	conv(64, latent_channels), # 14
	)


	def build_decoder_scale8(latent_channels, pool_blocks):
	def B(idx): return Block(64, 64, use_midblock_gn=(idx in pool_blocks))
	return nn.Sequential(
	Clamp(), # 0 (no weights)
	conv(latent_channels, 64), # 1
	nn.ReLU(), # 2
	B(3), B(4), B(5), # 3-5
	nn.Upsample(scale_factor=2), # 6
	conv(64, 64, bias=False), # 7
	B(8), B(9), B(10), # 8-10
	nn.Upsample(scale_factor=2), # 11
	conv(64, 64, bias=False), # 12
	B(13), B(14), B(15), # 13-15
	nn.Upsample(scale_factor=2), # 16
	conv(64, 64, bias=False), # 17
	B(18), # 18
	conv(64, 3), # 19
	)


	def build_encoder_scale16(latent_channels, pool_blocks):
	def B(idx): return Block(64, 64, use_midblock_gn=(idx in pool_blocks))
	return nn.Sequential(
	conv(3, 64), # 0
	B(1), # 1
	conv(64, 64, stride=2, bias=False), # 2
	B(3), B(4), B(5), # 3-5
	conv(64, 64, stride=2, bias=False), # 6
	B(7), B(8), B(9), # 7-9
	conv(64, 64, stride=2, bias=False), # 10
	B(11), B(12), B(13), # 11-13
	conv(64, 64, stride=2, bias=False), # 14
	B(15), B(16), B(17), # 15-17
	conv(64, latent_channels), # 18
	)


	def build_decoder_scale16(latent_channels, pool_blocks):
	def B(idx): return Block(64, 64, use_midblock_gn=(idx in pool_blocks))
	return nn.Sequential(
	Clamp(), # 0
	conv(latent_channels, 64), # 1
	nn.ReLU(), # 2
	B(3), B(4), B(5), # 3-5
	nn.Upsample(scale_factor=2), # 6
	conv(64, 64, bias=False), # 7
	B(8), B(9), B(10), # 8-10
	nn.Upsample(scale_factor=2), # 11
	conv(64, 64, bias=False), # 12
	B(13), B(14), B(15), # 13-15
	nn.Upsample(scale_factor=2), # 16
	conv(64, 64, bias=False), # 17
	B(18), B(19), B(20), # 18-20
	nn.Upsample(scale_factor=2), # 21
	conv(64, 64, bias=False), # 22
	B(23), # 23
	conv(64, 3), # 24
	)


	# ============================================================
	# Packed latents (auto-pad so it never errors)
	# ============================================================

	def unpack_packed_latents(x, latent_channels):
	# [B, C4, H, W] -> [B, C, H2, W*2]
	if x.ndim == 4 and x.shape[1] == latent_channels * 4:
	return (
	x.reshape(x.shape[0], latent_channels, 2, 2, x.shape[-2], x.shape[-1])
	.permute(0, 1, 4, 2, 5, 3)
	.reshape(x.shape[0], latent_channels, x.shape[-2] * 2, x.shape[-1] * 2)
	)
	return x


	def pack_packed_latents(z, latent_channels):
	# [B, C, H, W] -> [B, C*4, H//2, W//2]
	if z.ndim == 4 and z.shape[1] == latent_channels:
	h, w = z.shape[-2], z.shape[-1]
	pad_h = h & 1
	pad_w = w & 1
	if pad_h or pad_w:
	z = F.pad(z, (0, pad_w, 0, pad_h), mode="replicate")
	h, w = z.shape[-2], z.shape[-1]

	return (
	z.reshape(z.shape[0], latent_channels, h // 2, 2, w // 2, 2)
	.permute(0, 1, 3, 5, 2, 4)
	.reshape(z.shape[0], latent_channels * 4, h // 2, w // 2)
	)
	return z


	def pad_nchw_to_multiple(x, multiple):
	# replicate pad right/bottom so any size works
	_, _, h, w = x.shape
	pad_h = (multiple - (h % multiple)) % multiple
	pad_w = (multiple - (w % multiple)) % multiple
	if pad_h or pad_w:
	x = F.pad(x, (0, pad_w, 0, pad_h), mode="replicate")
	return x


	# ============================================================
	# Key conversion for your file format:
	# encoder.layers.N.* and decoder.layers.N.*
	# decoder layers must shift +1 because our decoder has Clamp() at index 0.
	# ============================================================

	def normalize_state_dict(sd_raw):
	keys = list(sd_raw.keys())

	# Already comfy split format?
	if any(k.startswith("taesd_encoder.") for k in keys) or any(k.startswith("taesd_decoder.") for k in keys):
	return sd_raw

	out = {}

	# Diffusers "encoder.layers.* / decoder.layers.*"
	if any(k.startswith("encoder.layers.") for k in keys) or any(k.startswith("decoder.layers.") for k in keys):
	for k, v in sd_raw.items():
	if k.startswith("encoder.layers."):
	# encoder.layers.N.xxx -> taesd_encoder.N.xxx
	out["taesd_encoder." + k[len("encoder.layers."):]] = v
	elif k.startswith("decoder.layers."):
	# decoder.layers.N.xxx -> taesd_decoder.(N+1).xxx (Clamp at 0)
	rest = k[len("decoder.layers."):]
	parts = rest.split(".", 1)
	try:
	n = int(parts[0])
	n2 = n + 1
	tail = parts[1] if len(parts) > 1 else ""
	out_key = f"taesd_decoder.{n2}" + (("." + tail) if tail else "")
	out[out_key] = v
	except Exception:
	# fallback, keep
	out[k] = v
	else:
	out[k] = v
	return out

	# Fallback: encoder./decoder. (numeric) — if decoder.0.weight looks like [64,C,3,3], offset it too
	if any(k.startswith("encoder.") for k in keys) or any(k.startswith("decoder.") for k in keys):
	decoder_needs_offset = False
	w0 = sd_raw.get("decoder.0.weight", None)
	if isinstance(w0, torch.Tensor) and w0.ndim == 4 and w0.shape[0] == 64 and w0.shape[2:] == (3, 3):
	decoder_needs_offset = True

	for k, v in sd_raw.items():
	if k.startswith("encoder."):
	out["taesd_encoder." + k[len("encoder."):]] = v
	elif k.startswith("decoder."):
	rest = k[len("decoder."):]
	if decoder_needs_offset:
	parts = rest.split(".", 1)
	if parts[0].isdigit():
	n = int(parts[0]) + 1
	tail = parts[1] if len(parts) > 1 else ""
	out_key = f"taesd_decoder.{n}" + (("." + tail) if tail else "")
	out[out_key] = v
	else:
	out["taesd_decoder." + rest] = v
	else:
	out["taesd_decoder." + rest] = v
	else:
	out[k] = v
	return out

	# Unknown layout: return as-is (Dump node will show keys)
	return sd_raw


	def split_encoder_decoder(sd):
	enc = {k[len("taesd_encoder."):]: v for k, v in sd.items() if k.startswith("taesd_encoder.")}
	dec = {k[len("taesd_decoder."):]: v for k, v in sd.items() if k.startswith("taesd_decoder.")}
	return enc, dec


	def pool_blocks_from_sd(part_sd):
	blocks = set()
	for k in part_sd.keys():
	if ".pool.0.weight" in k or ".pool.0.bias" in k:
	head = k.split(".", 1)[0]
	if head.isdigit():
	blocks.add(int(head))
	return blocks


	def infer_latent_channels_from_decoder(dec_sd):
	# Find smallest-index conv weight that looks like decoder input conv: [64, C, 3, 3]
	candidates = []
	for k, v in dec_sd.items():
	if not isinstance(v, torch.Tensor) or v.ndim != 4:
	continue
	head = k.split(".", 1)[0]
	if head.isdigit() and v.shape[0] == 64 and v.shape[2:] == (3, 3):
	candidates.append((int(head), int(v.shape[1])))
	if not candidates:
	raise RuntimeError("Could not infer latent_channels from decoder weights.")
	candidates.sort(key=lambda t: t[0])
	return candidates[0][1]


	def detect_layout(enc_sd, latent_channels):
	# Your file has encoder.layers.14.* -> after normalize it's "14.weight"
	if "14.weight" in enc_sd:
	w = enc_sd["14.weight"]
	if isinstance(w, torch.Tensor) and w.ndim == 4 and w.shape[0] == latent_channels and w.shape[1] == 64:
	return "scale8"
	if "18.weight" in enc_sd:
	w = enc_sd["18.weight"]
	if isinstance(w, torch.Tensor) and w.ndim == 4 and w.shape[0] == latent_channels and w.shape[1] == 64:
	return "scale16"

	# Fallback: find earliest [C,64,3,3] conv in encoder
	best = None
	for k, v in enc_sd.items():
	if not isinstance(v, torch.Tensor) or v.ndim != 4:
	continue
	head = k.split(".", 1)[0]
	if head.isdigit() and v.shape[0] == latent_channels and v.shape[1] == 64 and v.shape[2:] == (3, 3):
	idx = int(head)
	best = idx if best is None else min(best, idx)
	if best is None:
	raise RuntimeError("Could not detect encoder layout (scale8 vs scale16).")
	return "scale8" if best <= 14 else "scale16"


	# ============================================================
	# Core model (PR behavior: decode -> [-1,1], encode -> packed for taef2)
	# ============================================================

	class TAESDCore(nn.Module):
	def __init__(self, encoder, decoder, latent_channels, is_taef2):
	super().__init__()
	self.encoder = encoder
	self.decoder = decoder
	self.latent_channels = int(latent_channels)
	self.is_taef2 = bool(is_taef2)

	self.vae_scale = nn.Parameter(torch.tensor(1.0))
	self.vae_shift = nn.Parameter(torch.tensor(0.0))

	@torch.inference_mode()
	def decode(self, x):
	x = unpack_packed_latents(x, self.latent_channels)
	x = (x - self.vae_shift) * self.vae_scale
	x_sample = self.decoder(x)
	# decoder output in [0,1] -> [-1,1]
	return x_sample.sub(0.5).mul(2.0)

	@torch.inference_mode()
	def encode(self, x):
	# x is [-1,1] -> encoder expects [0,1]
	z = (self.encoder(x * 0.5 + 0.5) / self.vae_scale) + self.vae_shift
	if self.is_taef2:
	z = pack_packed_latents(z, self.latent_channels)
	return z


	def load_core(path, device, dtype):
	sd_raw = comfy.utils.load_torch_file(path, safe_load=True)
	sd = normalize_state_dict(sd_raw)
	enc_sd, dec_sd = split_encoder_decoder(sd)

	if not enc_sd or not dec_sd:
	sample = list(sd_raw.keys())[:40]
	raise RuntimeError(
	"Could not split encoder/decoder weights.\n"
	"Use Dump VAE Keys node and paste first ~40 keys.\n"
	f"First keys: {sample}"
	)

	enc_pool = pool_blocks_from_sd(enc_sd)
	dec_pool = pool_blocks_from_sd(dec_sd)

	latent_channels = infer_latent_channels_from_decoder(dec_sd)
	layout = detect_layout(enc_sd, latent_channels)

	# Flux2 taef2 packed-latents heuristic (matches your file):
	has_midblock_gn = (len(enc_pool) > 0) or (len(dec_pool) > 0)
	is_taef2 = (latent_channels == 32) and has_midblock_gn

	if layout == "scale8":
	encoder = build_encoder_scale8(latent_channels, enc_pool)
	decoder = build_decoder_scale8(latent_channels, dec_pool)
	base_downscale = 8
	else:
	encoder = build_encoder_scale16(latent_channels, enc_pool)
	decoder = build_decoder_scale16(latent_channels, dec_pool)
	base_downscale = 16

	# Load in fp32 first, then cast (more robust)
	core = TAESDCore(encoder, decoder, latent_channels, is_taef2)
	core.encoder.load_state_dict(enc_sd, strict=False)
	core.decoder.load_state_dict(dec_sd, strict=False)

	core = core.to(device=device, dtype=dtype).eval()
	for p in core.parameters():
	p.requires_grad_(False)

	core._base_downscale = base_downscale
	return core


	# ============================================================
	# Comfy VAE interface object
	# ============================================================

	class TAEF2VAE:
	def __init__(self, weights_path, device, dtype):
	self.device = device
	self.dtype = dtype
	self.core = load_core(weights_path, device=device, dtype=dtype)

	# packed latents halves latent H/W again -> effective downscale doubles
	self.downscale_ratio = self.core._base_downscale * (2 if self.core.is_taef2 else 1)

	print(
	f"[TAEF2] Loaded: {weights_path} \| latent_channels={self.core.latent_channels} "
	f"\| is_taef2={self.core.is_taef2} \| base_downscale={self.core._base_downscale} "
	f"\| effective_downscale={self.downscale_ratio}"
	)

	@torch.inference_mode()
	def decode(self, latents):
	x = latents.to(device=self.device, dtype=self.dtype)
	img = self.core.decode(x) # NCHW in [-1,1]
	img = img.clamp(-1, 1).add(1.0).mul(0.5) # -> [0,1]
	return img.to(torch.float32).permute(0, 2, 3, 1).contiguous() # NHWC float32

	@torch.inference_mode()
	def encode(self, pixels):
	# pixels NHWC [0,1]
	x = pixels[..., :3].permute(0, 3, 1, 2).contiguous()
	x = x.to(device=self.device, dtype=self.dtype).clamp(0, 1).mul(2.0).sub(1.0) # -> [-1,1]

	# Make it behave like base VAE: pad to required multiple so any size works
	x = pad_nchw_to_multiple(x, self.downscale_ratio)

	z = self.core.encode(x) # packed if taef2
	return z.to(torch.float32)

	def decode_tiled(self, latents, **kwargs):
	return self.decode(latents)

	def encode_tiled(self, pixels, **kwargs):
	return self.encode(pixels)

	def spacial_compression_decode(self):
	return self.downscale_ratio

	def spacial_compression_encode(self):
	return self.downscale_ratio

	def temporal_compression_decode(self):
	return None

	def temporal_compression_encode(self):
	return None


	# ============================================================
	# Nodes
	# ============================================================

	def _list_vae_files():
	vae_files = folder_paths.get_filename_list("vae")
	approx_files = folder_paths.get_filename_list("vae_approx")
	return sorted(set(vae_files + approx_files))

	def _resolve_vae_path(fname):
	path = folder_paths.get_full_path("vae_approx", fname)
	if path is None:
	path = folder_paths.get_full_path("vae", fname)
	return path


	class LoadTAEF2VAE:
	@classmethod
	def INPUT_TYPES(cls):
	return {
	"required": {
	"weights": (_list_vae_files(),),
	"dtype": (["bf16", "fp16", "fp32"], {"default": "bf16"}),
	}
	}

	RETURN_TYPES = ("VAE",)
	FUNCTION = "load"
	CATEGORY = "latent/vae"

	def load(self, weights, dtype):
	path = _resolve_vae_path(weights)
	if path is None:
	raise FileNotFoundError(f"Could not find weights file: {weights}")

	device = comfy.model_management.get_torch_device()
	if dtype == "bf16":
	tdtype = torch.bfloat16
	elif dtype == "fp16":
	tdtype = torch.float16
	else:
	tdtype = torch.float32

	return (TAEF2VAE(path, device=device, dtype=tdtype),)


	class DumpVAEKeys:
	@classmethod
	def INPUT_TYPES(cls):
	return {
	"required": {
	"weights": (_list_vae_files(),),
	"include_shapes": ("BOOLEAN", {"default": True}),
	"sort_keys": ("BOOLEAN", {"default": True}),
	"max_lines": ("INT", {"default": 0, "min": 0, "max": 200000}),
	}
	}

	RETURN_TYPES = ("STRING",)
	FUNCTION = "dump"
	CATEGORY = "utils/debug"

	def dump(self, weights, include_shapes, sort_keys, max_lines):
	path = _resolve_vae_path(weights)
	if path is None:
	raise FileNotFoundError(f"Could not find weights file: {weights}")

	sd = comfy.utils.load_torch_file(path, safe_load=True)
	keys = list(sd.keys())
	if sort_keys:
	keys.sort()

	lines = []
	if include_shapes:
	for k in keys:
	v = sd[k]
	if isinstance(v, torch.Tensor):
	lines.append(f"{k}\t{tuple(v.shape)}\t{str(v.dtype)}")
	else:
	lines.append(f"{k}\t{type(v)}")
	else:
	lines = keys

	if max_lines and len(lines) > max_lines:
	head = lines[:max_lines]
	head.append(f"... TRUNCATED: total_keys={len(lines)} (showing first {max_lines}) ...")
	lines = head

	text = "\n".join(lines)
	return {"ui": {"text": [text]}, "result": (text,)}


	NODE_CLASS_MAPPINGS = {
	"LoadTAEF2VAE": LoadTAEF2VAE,
	"DumpVAEKeys": DumpVAEKeys,
	}

	NODE_DISPLAY_NAME_MAPPINGS = {
	"LoadTAEF2VAE": "Load TAEF2 (Flux2 Tiny VAE)",
	"DumpVAEKeys": "Dump VAE Keys (as String)",
	}