From 8e869efc8764546415036e5fdac05fc287dbe926 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:20:28 -0700 Subject: [PATCH] Add support for taeh3. (#15695) --- comfy/latent_formats.py | 1 + comfy/sd.py | 9 ++++++++- comfy/taesd/taehv.py | 27 ++++++++++++++++++++++++++- latent_preview.py | 3 +-- nodes.py | 2 +- 5 files changed, 37 insertions(+), 5 deletions(-) diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index ed5167075..34c6d700c 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -577,6 +577,7 @@ class MiniMaxH3Video(LatentFormat): spacial_downscale_ratio = 16 temporal_downscale_ratio = 4 scale_factor = 1.0 + taesd_decoder_name = "taeh3" latent_rgb_factors = [ [-0.018555, 0.024344, -0.017536], diff --git a/comfy/sd.py b/comfy/sd.py index 4bdaa978c..06679c6fb 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -907,7 +907,14 @@ class VAE: self.upscale_index_formula = (4, 16, 16) self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 16, 16) self.downscale_index_formula = (4, 16, 16) - if self.latent_channels in [48, 128]: # Wan 2.2 and LTX2 + if self.latent_channels == 24 and sd["decoder.22.bias"].shape[0] == 12: # MiniMax H3 + self.first_stage_model = comfy.taesd.taehv.TAEHV(latent_channels=self.latent_channels, latent_format=None) + self.process_input = self.process_output = lambda image: image + self.upscale_ratio = (lambda a: max(1, (a - 2) // 5 * 17 + 5), 16, 16) + self.downscale_ratio = (lambda a: max(1, (a - 1) // 17 * 5 + 2) if a > 1 else 1, 16, 16) + self.memory_used_encode = lambda shape, dtype: (400 * ((shape[-3] + 16) // 17) * shape[-2] * shape[-1] * model_management.dtype_size(dtype)) + self.memory_used_decode = lambda shape, dtype: ((260 * 16 * 16 + shape[1] * shape[-3]) * shape[-2] * shape[-1] * model_management.dtype_size(dtype)) + elif self.latent_channels in [48, 128]: # Wan 2.2 and LTX2 self.first_stage_model = comfy.taesd.taehv.TAEHV(latent_channels=self.latent_channels, latent_format=None) # taehv doesn't need scaling self.process_input = self.process_output = lambda image: image self.process_output = lambda image: image diff --git a/comfy/taesd/taehv.py b/comfy/taesd/taehv.py index 696013200..ffa9f89d1 100644 --- a/comfy/taesd/taehv.py +++ b/comfy/taesd/taehv.py @@ -131,10 +131,11 @@ class TAEHV(nn.Module): self.latent_channels = latent_channels self.parallel = parallel self.latent_format = latent_format + self.is_h3 = self.latent_channels == 24 self.show_progress_bar = show_progress_bar self.process_in = latent_format().process_in if latent_format is not None else (lambda x: x) self.process_out = latent_format().process_out if latent_format is not None else (lambda x: x) - if self.latent_channels in [48, 32]: # Wan 2.2 and HunyuanVideo1.5 + if self.latent_channels in [48, 32, 24]: # Wan 2.2, HunyuanVideo1.5 and MiniMax H3 self.patch_size = 2 elif self.latent_channels == 128: # LTX2 self.patch_size, self.latent_channels, encoder_time_downscale, decoder_time_upscale = 4, 128, (True, True, True), (True, True, True) @@ -176,6 +177,21 @@ class TAEHV(nn.Module): def encode(self, x, **kwargs): x = x.movedim(2, 1) # [B, C, T, H, W] -> [B, T, C, H, W] + if self.is_h3: + single_frame = x.shape[1] == 1 + batch = x.shape[0] + x = torch.cat([x, x[:, -1:].expand(-1, -x.shape[1] % 17, -1, -1, -1)], dim=1) + x = F.pad(x.reshape(batch, -1, 17, *x.shape[2:]), (0, 0, 0, 0, 0, 0, 3, 0)) + if self.parallel: + x = apply_model_with_memblocks(self.encoder, x.flatten(0, 1), True, self.show_progress_bar, + patch_size=self.patch_size) + x = x.reshape(batch, -1, *x.shape[2:]) + else: + x = torch.cat([apply_model_with_memblocks(self.encoder, chunk, False, False, + patch_size=self.patch_size) + for chunk in tqdm(x.unbind(1), disable=not self.show_progress_bar)], dim=1) + x = x[:, :1] if single_frame else x[:, :-3] + return self.process_out(x.movedim(2, 1)) if x.shape[1] % self.t_downscale != 0: # pad at end to multiple of t_downscale n_pad = self.t_downscale - x.shape[1] % self.t_downscale @@ -189,7 +205,16 @@ class TAEHV(nn.Module): x = x.unsqueeze(0) if x.ndim == 4 else x # [T, C, H, W] -> [1, T, C, H, W] x = x.movedim(1, 2) if x.shape[1] != self.latent_channels else x # [B, T, C, H, W] or [B, C, T, H, W] x = self.process_in(x).movedim(2, 1) # [B, C, T, H, W] -> [B, T, C, H, W] + if self.is_h3: + single_frame = x.shape[1] == 1 x = apply_model_with_memblocks(self.decoder, x, self.parallel, self.show_progress_bar, output_device=comfy.model_management.intermediate_device(), patch_size=self.patch_size, decode=True) + if self.is_h3: + x.clamp_(0, 1) + if not single_frame: + chunk_frames = 5 * self.t_upscale + x = F.pad(x, (0, 0, 0, 0, 0, 0, 0, -x.shape[1] % chunk_frames)) + x = x.unflatten(1, (-1, chunk_frames))[:, :, self.frames_to_trim:].flatten(1, 2) + return x[:, :-3 * self.t_upscale].movedim(2, 1) return x[:, self.frames_to_trim:].movedim(2, 1) diff --git a/latent_preview.py b/latent_preview.py index 6bf2c1869..d98b70019 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -11,7 +11,7 @@ import logging default_preview_method = args.preview_method MAX_PREVIEW_RESOLUTION = args.preview_size -VIDEO_TAES = ["taehv", "lighttaew2_2", "lighttaew2_1", "lighttaehy1_5", "taeltx_2"] +VIDEO_TAES = ["taehv", "lighttaew2_2", "lighttaew2_1", "lighttaehy1_5", "taeltx_2", "taeh3"] def preview_to_image(latent_image, do_scale=True): if do_scale: @@ -136,4 +136,3 @@ def set_preview_method(override: str = None): args.preview_method = method return args.preview_method = default_preview_method - diff --git a/nodes.py b/nodes.py index 1a3dd3f48..fa3a77949 100644 --- a/nodes.py +++ b/nodes.py @@ -768,7 +768,7 @@ class LoraLoaderModelOnly(LoraLoader): return (self.load_lora(model, None, lora_name, strength_model, 0)[0],) class VAELoader: - video_taes = ["taehv", "lighttaew2_2", "lighttaew2_1", "lighttaehy1_5", "taeltx_2"] + video_taes = ["taehv", "lighttaew2_2", "lighttaew2_1", "lighttaehy1_5", "taeltx_2", "taeh3"] image_taes = ["taesd", "taesdxl", "taesd3", "taef1", "taef2"] @staticmethod