diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index 82d9f9bb8..1ab4cb280 100644 --- a/comfy/latent_formats.py +++ b/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) diff --git a/comfy/ldm/models/autoencoder.py b/comfy/ldm/models/autoencoder.py index 13bd6e16b..819df2437 100644 --- a/comfy/ldm/models/autoencoder.py +++ b/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) diff --git a/comfy/sd.py b/comfy/sd.py index 5b95cf75a..75c7b5973 100644 --- a/comfy/sd.py +++ b/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) diff --git a/comfy/taesd/taehv.py b/comfy/taesd/taehv.py new file mode 100644 index 000000000..2d019bec3 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/latent_preview.py b/latent_preview.py index 95d3cb733..2ba3a72a1 100644 --- a/latent_preview.py +++ b/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)) diff --git a/nodes.py b/nodes.py index 231d4d4de..e52b1dc43 100644 --- a/nodes.py +++ b/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)