Merge pull request #1 from Pratik-Doshi-99/pratik-taehv2

TAE for HV works with VAELoaders

comfy/latent_formats.py

@@ -382,6 +382,7 @@ class HunyuanVideo(LatentFormat):
     ]
 
     latent_rgb_factors_bias = [ 0.0259, -0.0192, -0.0761]
+    taesd_decoder_name = "taehv"
 
 class Cosmos1CV8x8x8(LatentFormat):
     latent_channels = 16
@@ -445,7 +446,7 @@ class Wan21(LatentFormat):
         ]).view(1, self.latent_channels, 1, 1, 1)
 
 
-        self.taesd_decoder_name = None #TODO
+        self.taesd_decoder_name = "taew2_1"
 
     def process_in(self, latent):
         latents_mean = self.latents_mean.to(latent.device, latent.dtype)

comfy/ldm/models/autoencoder.py

@@ -200,6 +200,8 @@ class AutoencodingEngineLegacy(AutoencodingEngine):
         return z
 
     def decode(self, z: torch.Tensor, **decoder_kwargs) -> torch.Tensor:
+        print('Decoding hunyuan latent. Received tensor:',z.shape)
+        
         if self.max_batch_size is None:
             dec = self.post_quant_conv(z)
             dec = self.decoder(dec, **decoder_kwargs)

comfy/sd.py

@@ -52,6 +52,7 @@ import comfy.lora_convert
 import comfy.hooks
 import comfy.t2i_adapter.adapter
 import comfy.taesd.taesd
+import comfy.taesd.taehv
 
 import comfy.ldm.flux.redux
 
@@ -297,6 +298,13 @@ class VAE:
             elif "taesd_decoder.1.weight" in sd:
                 self.latent_channels = sd["taesd_decoder.1.weight"].shape[1]
                 self.first_stage_model = comfy.taesd.taesd.TAESD(latent_channels=self.latent_channels)
+            elif "taehv_flag" in sd:
+                self.first_stage_model = comfy.taesd.taehv.TAEHV()
+                self.memory_used_decode = lambda shape, dtype: (1000 * shape[2] * shape[3] * shape[4] * 64) * model_management.dtype_size(dtype)
+                self.memory_used_encode = lambda shape, dtype: (1000 * shape[2] * shape[3] * shape[4]) * model_management.dtype_size(dtype)
+                self.latent_channels = 16
+                self.latent_dim = 3
+                sd.pop('taehv_flag',None)
             elif "vquantizer.codebook.weight" in sd: #VQGan: stage a of stable cascade
                 self.first_stage_model = StageA()
                 self.downscale_ratio = 4
@@ -401,6 +409,7 @@ class VAE:
                 self.downscale_index_formula = (4, 8, 8)
                 self.latent_dim = 3
                 self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.conv.weight"].shape[1]
+                print('Loading Hunyuan VAE. Latent channels = ',self.latent_channels)
                 self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=sd['post_quant_conv.weight'].shape[1])
                 self.memory_used_decode = lambda shape, dtype: (1500 * shape[2] * shape[3] * shape[4] * (4 * 8 * 8)) * model_management.dtype_size(dtype)
                 self.memory_used_encode = lambda shape, dtype: (900 * max(shape[2], 2) * shape[3] * shape[4]) * model_management.dtype_size(dtype)

comfy/taesd/taehv.py

@@ -0,0 +1,287 @@
+"""
+Tiny AutoEncoder for Hunyuan Video
+(DNN for encoding / decoding videos to Hunyuan Video's latent space)
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from collections import namedtuple
+
+
+import comfy.utils
+import comfy.ops
+
+DecoderResult = namedtuple("DecoderResult", ("frame", "memory"))
+TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index"))
+
+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 MemBlock(nn.Module):
+    def __init__(self, n_in, n_out):
+        super().__init__()
+        self.conv = nn.Sequential(conv(n_in * 2, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), nn.ReLU(inplace=True), 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.act = nn.ReLU(inplace=True)
+    def forward(self, x, past):
+        return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x))
+
+class TPool(nn.Module):
+    def __init__(self, n_f, stride):
+        super().__init__()
+        self.stride = stride
+        self.conv = comfy.ops.disable_weight_init.Conv2d(n_f*stride, n_f, 1, bias=False)
+    def forward(self, x):
+        _NT, C, H, W = x.shape
+        return self.conv(x.reshape(-1, self.stride * C, H, W))
+
+class TGrow(nn.Module):
+    def __init__(self, n_f, stride):
+        super().__init__()
+        self.stride = stride
+        self.conv = comfy.ops.disable_weight_init.Conv2d(n_f, n_f*stride, 1, bias=False)
+    def forward(self, x):
+        _NT, C, H, W = x.shape
+        x = self.conv(x)
+        return x.reshape(-1, C, H, W)
+
+def apply_model_with_memblocks(model, x, parallel, show_progress_bar):
+    """
+    Apply a sequential model with memblocks to the given input.
+    Args:
+    - model: nn.Sequential of blocks to apply
+    - x: input data, of dimensions NTCHW
+    - parallel: if True, parallelize over timesteps (fast but uses O(T) memory)
+        if False, each timestep will be processed sequentially (slow but uses O(1) memory)
+    - show_progress_bar: if True, enables tqdm progressbar display
+    Returns NTCHW tensor of output data.
+    """
+    assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor"
+    N, T, C, H, W = x.shape
+    print('Received tensor of shape:',x.shape)
+    if parallel:
+        x = x.reshape(N*T, C, H, W)
+        # parallel over input timesteps, iterate over blocks
+        for b in model:
+            if isinstance(b, MemBlock):
+                NT, C, H, W = x.shape
+                T = NT // N
+                _x = x.reshape(N, T, C, H, W)
+                mem = F.pad(_x, (0,0,0,0,0,0,1,0), value=0)[:,:T].reshape(x.shape)
+                print('Intermediate shape:',x.shape)
+                x = b(x, mem)
+            else:
+                print('Intermediate shape:',x.shape)
+                x = b(x)
+        NT, C, H, W = x.shape
+        T = NT // N
+        x = x.view(N, T, C, H, W)
+    else:
+        # TODO(oboerbohan): at least on macos this still gradually uses more memory during decode...
+        # need to fix :(
+        out = []
+        # iterate over input timesteps and also iterate over blocks.
+        # because of the cursed TPool/TGrow blocks, this is not a nested loop,
+        # it's actually a ***graph traversal*** problem! so let's make a queue
+        work_queue = [TWorkItem(xt, 0) for t, xt in enumerate(x.reshape(N, T * C, H, W).chunk(T, dim=1))]
+        # we'll also need a separate addressable memory per node as well
+        mem = [None] * len(model)
+        while work_queue:
+            xt, i = work_queue.pop(0)
+            print('Intermediate shape:', xt.shape)
+            if i == len(model):
+                # reached end of the graph, append result to output list
+                out.append(xt)
+            else:
+                # fetch the block to process
+                b = model[i]
+                if isinstance(b, MemBlock):
+                    # mem blocks are simple since we're visiting the graph in causal order
+                    if mem[i] is None:
+                        xt_new = b(xt, xt * 0)
+                        mem[i] = xt
+                    else:
+                        xt_new = b(xt, mem[i])
+                        mem[i].copy_(xt) # inplace might reduce mysterious pytorch memory allocations? doesn't help though
+                    # add successor to work queue
+                    work_queue.insert(0, TWorkItem(xt_new, i+1))
+                elif isinstance(b, TPool):
+                    # pool blocks are miserable
+                    if mem[i] is None:
+                        mem[i] = [] # pool memory is itself a queue of inputs to pool
+                    mem[i].append(xt)
+                    if len(mem[i]) > b.stride:
+                        # pool mem is in invalid state, we should have pooled before this
+                        raise ValueError("???")
+                    elif len(mem[i]) < b.stride:
+                        # pool mem is not yet full, go back to processing the work queue
+                        pass
+                    else:
+                        # pool mem is ready, run the pool block
+                        N, C, H, W = xt.shape 
+                        xt = b(torch.cat(mem[i], 1).view(N*b.stride, C, H, W))
+                        # reset the pool mem
+                        mem[i] = []
+                        # add successor to work queue
+                        work_queue.insert(0, TWorkItem(xt, i+1))
+                elif isinstance(b, TGrow):
+                    xt = b(xt)
+                    NT, C, H, W = xt.shape
+                    # each tgrow has multiple successor nodes
+                    for xt_next in reversed(xt.view(N, b.stride*C, H, W).chunk(b.stride, 1)):
+                        # add successor to work queue
+                        work_queue.insert(0, TWorkItem(xt_next, i+1))
+                else:
+                    # normal block with no funny business
+                    xt = b(xt)
+                    # add successor to work queue
+                    work_queue.insert(0, TWorkItem(xt, i+1))
+        x = torch.stack(out, 1)
+    return x
+
+class TAEHV(nn.Module):
+    latent_channels = 16
+    image_channels = 3
+    def __init__(self, checkpoint_path=None, decoder_time_upscale=(True, True), decoder_space_upscale=(True, True, True)):
+        """Initialize pretrained TAEHV from the given checkpoint.
+        Arg:
+            checkpoint_path: path to weight file to load. taehv.pth for Hunyuan, taew2_1.pth for Wan 2.1.
+            decoder_time_upscale: whether temporal upsampling is enabled for each block. upsampling can be disabled for a cheaper preview.
+            decoder_space_upscale: whether spatial upsampling is enabled for each block. upsampling can be disabled for a cheaper preview.
+        """
+        super().__init__()
+        self.encoder = nn.Sequential(
+            conv(TAEHV.image_channels, 64), nn.ReLU(inplace=True),
+            TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
+            TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
+            TPool(64, 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
+            conv(64, TAEHV.latent_channels),
+        )
+        n_f = [256, 128, 64, 64]
+        self.frames_to_trim = 2**sum(decoder_time_upscale) - 1
+        self.decoder = nn.Sequential(
+            Clamp(), conv(TAEHV.latent_channels, n_f[0]), nn.ReLU(inplace=True),
+            MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 1), conv(n_f[0], n_f[1], bias=False),
+            MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), conv(n_f[1], n_f[2], bias=False),
+            MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), conv(n_f[2], n_f[3], bias=False),
+            nn.ReLU(inplace=True), conv(n_f[3], TAEHV.image_channels),
+        )
+        if checkpoint_path is not None:
+            self.load_state_dict(comfy.utils.load_torch_file(checkpoint_path, safe_load=True))
+
+    
+    def load_state_dict(self, state_dict, strict=True):
+        return super().load_state_dict(self.patch_tgrow_layers(state_dict), strict=strict)
+    
+    
+    @staticmethod
+    def from_comfy_state_dict(state_dict):
+        """Create TAEHV model from ComfyUI-formatted state dict.
+        
+        Args:
+            state_dict: State dict with taehv_decoder.* and taehv_encoder.* keys
+            
+        Returns:
+            TAEHV model with loaded weights
+        """
+        # Create model without loading checkpoint
+        model = TAEHV(checkpoint_path=None)
+
+        # Convert ComfyUI state dict format back to TAEHV format
+        taehv_sd = {}
+        for key, value in state_dict.items():
+            if key.startswith("taehv_decoder."):
+                new_key = key.replace("taehv_decoder.", "decoder.")
+                taehv_sd[new_key] = value
+            elif key.startswith("taehv_encoder."):
+                new_key = key.replace("taehv_encoder.", "encoder.")
+                taehv_sd[new_key] = value
+
+        # Load the converted state dict
+        if taehv_sd:
+            model.load_state_dict(taehv_sd, strict=False)
+
+        return model
+
+    def patch_tgrow_layers(self, sd):
+        """Patch TGrow layers to use a smaller kernel if needed.
+        Args:
+            sd: state dict to patch
+        """
+        new_sd = self.state_dict()
+        for i, layer in enumerate(self.decoder):
+            if isinstance(layer, TGrow):
+                key = f"decoder.{i}.conv.weight"
+                if sd[key].shape[0] > new_sd[key].shape[0]:
+                    # take the last-timestep output channels
+                    sd[key] = sd[key][-new_sd[key].shape[0]:]
+        return sd
+
+    def encode_video(self, x, parallel=False, show_progress_bar=False):
+        """Encode a sequence of frames.
+        Args:
+            x: input NTCHW RGB (C=3) tensor with values in [0, 1].
+            parallel: if True, all frames will be processed at once.
+              (this is faster but may require more memory).
+              if False, frames will be processed sequentially.
+        Returns NTCHW latent tensor with ~Gaussian values.
+        """
+        return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar)
+
+    def decode_video(self, x, parallel=False, show_progress_bar=False):
+        """Decode a sequence of frames.
+        Args:
+            x: input NCTHW latent (C=16) tensor with ~Gaussian values.
+            parallel: if True, all frames will be processed at once.
+              (this is faster but may require more memory).
+              if False, frames will be processed sequentially.
+        Returns NTCHW RGB tensor with ~[0, 1] values.
+        """
+        #converting NCTHW to NTCHW
+        x = x.permute(0,2,1,3,4)
+        x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar)
+
+        
+        x = x[:, self.frames_to_trim:] # trim the time dimension
+        
+        #converting NTCHW to NCTHW
+        x = x.permute(0,2,1,3,4)
+
+        return x
+
+    def decode(self, x):
+        """Decode a single frame or batch of frames for preview."""
+        if x.ndim == 4:
+            # Add temporal dimension for single frame
+            x = x.unsqueeze(1)
+
+        # For preview, we'll just take the first frame after decoding
+        decoded = self.decode_video(x, parallel=False, show_progress_bar=False)
+        print(decoded.shape)
+
+        #converting
+
+        return decoded
+        # Return single frame for preview
+        # if decoded.shape[1] > 0:
+        #     return decoded[:, 0]
+        # else:
+        #     return decoded.squeeze(1)
+
+    def encode(self, x):
+        """Encode a single frame or batch of frames."""
+        if x.ndim == 4:
+            # Add temporal dimension for single frame
+            x = x.unsqueeze(1)
+
+        encoded = self.encode_video(x, parallel=False, show_progress_bar=False)
+
+        # Return single frame
+        return encoded.squeeze(1)
+
+    def forward(self, x):
+        return self.decode(x)

latent_preview.py

@@ -2,6 +2,7 @@ import torch
 from PIL import Image
 from comfy.cli_args import args, LatentPreviewMethod
 from comfy.taesd.taesd import TAESD
+from comfy.taesd.taehv import TAEHV
 import comfy.model_management
 import folder_paths
 import comfy.utils
@@ -35,6 +36,29 @@ class TAESDPreviewerImpl(LatentPreviewer):
         x_sample = self.taesd.decode(x0[:1])[0].movedim(0, 2)
         return preview_to_image(x_sample)
 
+# TODO: add a video preview instead of image
+class TAEHVPreviewerImpl(LatentPreviewer):
+    def __init__(self, taehv):
+        self.taehv = taehv
+
+    def decode_latent_to_preview(self, x0):
+        # For video models, we take the first frame for preview
+        if x0.ndim == 5:
+            # Already NTCHW format
+            decoded = self.taehv.decode(x0[:1, :1])
+            if decoded.ndim == 4:
+                # NCHW output
+                x_sample = decoded[0].movedim(0, 2)
+            else:
+                # NTCHW output, take first frame
+                x_sample = decoded[0, 0].movedim(0, 2)
+        else:
+            # NCHW format, add temporal dimension
+            decoded = self.taehv.decode(x0[:1])
+            x_sample = decoded[0].movedim(0, 2)
+        return preview_to_image(x_sample)
+
+
 
 class Latent2RGBPreviewer(LatentPreviewer):
     def __init__(self, latent_rgb_factors, latent_rgb_factors_bias=None):
@@ -78,8 +102,12 @@ def get_previewer(device, latent_format):
 
         if method == LatentPreviewMethod.TAESD:
             if taesd_decoder_path:
-                taesd = TAESD(None, taesd_decoder_path, latent_channels=latent_format.latent_channels).to(device)
-                previewer = TAESDPreviewerImpl(taesd)
+                if latent_format.taesd_decoder_name in ['taehv', 'taew2_1']:
+                    taehv = TAEHV(checkpoint_path=taesd_decoder_path).to(device)
+                    previewer = TAEHVPreviewerImpl(taehv)
+                else:
+                    taesd = TAESD(None, taesd_decoder_path, latent_channels=latent_format.latent_channels).to(device)
+                    previewer = TAESDPreviewerImpl(taesd)
             else:
                 logging.warning("Warning: TAESD previews enabled, but could not find models/vae_approx/{}".format(latent_format.taesd_decoder_name))
 

nodes.py

@@ -699,7 +699,8 @@ class VAELoader:
         sd3_taesd_dec = False
         f1_taesd_enc = False
         f1_taesd_dec = False
-
+        taehv_available = False
+        taew2_1_available = False
         for v in approx_vaes:
             if v.startswith("taesd_decoder."):
                 sd1_taesd_dec = True
@@ -717,6 +718,10 @@ class VAELoader:
                 f1_taesd_dec = True
             elif v.startswith("taef1_decoder."):
                 f1_taesd_enc = True
+            elif v.startswith("taehv."):
+                taehv_available = True
+            elif v.startswith("taew2_1."):
+                taew2_1_available = True
         if sd1_taesd_dec and sd1_taesd_enc:
             vaes.append("taesd")
         if sdxl_taesd_dec and sdxl_taesd_enc:
@@ -725,6 +730,10 @@ class VAELoader:
             vaes.append("taesd3")
         if f1_taesd_dec and f1_taesd_enc:
             vaes.append("taef1")
+        if taehv_available:
+            vaes.append("taehv")
+        if taew2_1_available:
+            vaes.append("taew2_1")
         return vaes
 
     @staticmethod
@@ -757,6 +766,36 @@ class VAELoader:
             sd["vae_shift"] = torch.tensor(0.1159)
         return sd
 
+    @staticmethod
+    def load_tae_video(name):
+        sd = {}
+        approx_vaes = folder_paths.get_filename_list("vae_approx")
+
+        # name is either taehv for Hunyuan Video or taew2_1 for WAN2.1
+        tae_file = next(filter(lambda a: a.startswith("{}.".format(name)), approx_vaes))
+        tae_weights = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", tae_file))
+
+        #Confirmed model structure. Dont need this anymore
+        # for k in tae_weights:
+        #     if k.startswith("decoder."):
+        #         sd["taehv_decoder.{}".format(k)] = tae_weights[k]
+        #     elif k.startswith("encoder."):
+        #         sd["taehv_encoder.{}".format(k)] = tae_weights[k]
+        #     else:
+        #         # For weights without clear prefix, assume they're decoder weights
+        #         sd["taehv_decoder.{}".format(k)] = tae_weights[k]
+
+        sd.update(tae_weights)
+        sd['taehv_flag'] = True
+        #TODO: Confirm scale/shift params
+        if name == "taehv":
+            sd["vae_scale"] = torch.tensor(0.476986)  # HunyuanVideo scale
+            sd["vae_shift"] = torch.tensor(0.0)
+        elif name == "taew21":
+            sd["vae_scale"] = torch.tensor(1.0)  # Wan21 scale
+            sd["vae_shift"] = torch.tensor(0.0)
+        return sd
+
     @classmethod
     def INPUT_TYPES(s):
         return {"required": { "vae_name": (s.vae_list(), )}}
@@ -769,6 +808,8 @@ class VAELoader:
     def load_vae(self, vae_name):
         if vae_name in ["taesd", "taesdxl", "taesd3", "taef1"]:
             sd = self.load_taesd(vae_name)
+        elif vae_name in ["taehv", "taew2_1"]:
+            sd = self.load_tae_video(vae_name)
         else:
             vae_path = folder_paths.get_full_path_or_raise("vae", vae_name)
             sd = comfy.utils.load_torch_file(vae_path)