zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-19 07:48:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add support for taeh3. (#15695)

Browse Source
This commit is contained in:
comfyanonymous
2026-08-17 20:20:28 -07:00
committed by GitHub
parent ff6c8a8af1
commit 8e869efc87
5 changed files with 37 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
comfy/latent_formats.py
View File

@@ -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],

9
comfy/sd.py
View File

@@ -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

27
comfy/taesd/taehv.py
View File

@@ -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)

3
latent_preview.py
View File

@@ -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

2
nodes.py
View File

@@ -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