modular cleanup

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/framepack/framepack_api.py

@@ -37,7 +37,6 @@ class ReqFramepack(BaseModel):
     mp4_opt: str | None = Field(default="crf=16", title="Options", description="Options for the video codec")
     mp4_ext: str | None = Field(default="mp4", title="Format", description="Format for the video")
     mp4_interpolate: int | None = Field(default=0, title="Interpolation", description="Interpolation for the video")
-    attention: str | None = Field(default="Default", title="Attention", description="Attention type for the model")
     vae_type: str | None = Field(default="Local", title="VAE", description="VAE type for the model")
     vlm_enhance: bool | None = Field(default=False, title="VLM enhance", description="Enable VLM enhance")
     vlm_model: str | None = Field(default=None, title="VLM model", description="VLM model to use")
@@ -114,7 +113,6 @@ def framepack_post(request: ReqFramepack):
         mp4_opt=request.mp4_opt,
         mp4_ext=request.mp4_ext,
         mp4_interpolate=request.mp4_interpolate,
-        attention=request.attention,
         vae_type=request.vae_type,
         vlm_enhance=request.vlm_enhance,
         vlm_model=request.vlm_model,

modules/framepack/framepack_install.py

@@ -16,7 +16,7 @@ def rename(src:str, dst:str):
             raise e
 
 
-def install_requirements(attention:str='SDPA'):
+def install_requirements():
     install('av')
 
 

modules/framepack/framepack_ui.py

@@ -63,7 +63,6 @@ def create_ui(prompt, negative, styles, _overrides, script_inputs, mp4_fps, mp4_
                 optimized_prompt = gr.Checkbox(label='FP use optimized system prompt', value=True)
                 use_cfgzero = gr.Checkbox(label='FP enable CFGZero', value=False)
                 use_preview = gr.Checkbox(label='FP enable Preview', value=True)
-                attention = gr.Dropdown(label="FP attention", choices=['Default', 'Xformers', 'FlashAttention', 'SageAttention'], value='Default', type='value')
                 vae_type = gr.Dropdown(label="FP VAE", choices=['Full', 'Tiny', 'Remote'], value='Full', type='value')
 
         with gr.Column(elem_id='framepack-output-column', scale=2) as _column_output:
@@ -88,7 +87,7 @@ def create_ui(prompt, negative, styles, _overrides, script_inputs, mp4_fps, mp4_
     duration.change(fn=change_sections, inputs=[duration, mp4_fps, mp4_interpolate, latent_ws, variant], outputs=[section_html, section_prompt])
     mp4_fps.change(fn=change_sections, inputs=[duration, mp4_fps, mp4_interpolate, latent_ws, variant], outputs=[section_html, section_prompt])
     mp4_interpolate.change(fn=change_sections, inputs=[duration, mp4_fps, mp4_interpolate, latent_ws, variant], outputs=[section_html, section_prompt])
-    btn_load.click(fn=load_model, inputs=[variant, attention], outputs=framepack_outputs)
+    btn_load.click(fn=load_model, inputs=[variant], outputs=framepack_outputs)
     btn_unload.click(fn=unload_model, outputs=framepack_outputs)
     receipe_get.click(fn=framepack_load.get_model, inputs=[], outputs=receipe)
     receipe_set.click(fn=framepack_load.set_model, inputs=[receipe], outputs=[])
@@ -108,7 +107,7 @@ def create_ui(prompt, negative, styles, _overrides, script_inputs, mp4_fps, mp4_
         use_teacache, use_cfgzero, use_preview,
         mp4_fps, mp4_codec, mp4_sf, mp4_video, mp4_frames, mp4_thumb, mp4_opt, mp4_ext, mp4_interpolate,
         mp4_scale, mp4_upscaler,
-        attention, vae_type, variant,
+        vae_type, variant,
         vlm_enhance, vlm_model, vlm_system_prompt,
     ]
 

modules/framepack/framepack_wrappers.py

@@ -91,11 +91,11 @@ def prepare_prompts(p, init_image, prompt:str, section_prompt:str, num_sections:
     return generated_prompts
 
 
-def load_model(variant, attention):
+def load_model(variant):
     global loaded_variant # pylint: disable=global-statement
     if (shared.sd_model_type != 'hunyuanvideo') or (loaded_variant != variant):
         yield gr.update(), gr.update(), 'Verifying FramePack'
-        framepack_install.install_requirements(attention)
+        framepack_install.install_requirements()
         # framepack_install.git_clone(git_repo=git_repo, git_dir=git_dir, tmp_dir=tmp_dir)
         # framepack_install.git_update(git_dir=git_dir, git_commit=git_commit)
         # sys.path.append(git_dir)
@@ -114,7 +114,7 @@ def unload_model():
     yield gr.update(), gr.update(), 'Model unloaded'
 
 
-def run_framepack(task_id, _ui_state, init_image, end_image, start_weight, end_weight, vision_weight, prompt, system_prompt, optimized_prompt, section_prompt, negative_prompt, styles, seed, resolution, duration, latent_ws, steps, cfg_scale, cfg_distilled, cfg_rescale, shift, use_teacache, use_cfgzero, use_preview, mp4_fps, mp4_codec, mp4_sf, mp4_video, mp4_frames, mp4_thumb, mp4_opt, mp4_ext, mp4_interpolate, mp4_scale, mp4_upscaler, attention, vae_type, variant, vlm_enhance, vlm_model, vlm_system_prompt, *_args, **_kwargs):
+def run_framepack(task_id, _ui_state, init_image, end_image, start_weight, end_weight, vision_weight, prompt, system_prompt, optimized_prompt, section_prompt, negative_prompt, styles, seed, resolution, duration, latent_ws, steps, cfg_scale, cfg_distilled, cfg_rescale, shift, use_teacache, use_cfgzero, use_preview, mp4_fps, mp4_codec, mp4_sf, mp4_video, mp4_frames, mp4_thumb, mp4_opt, mp4_ext, mp4_interpolate, mp4_scale, mp4_upscaler, vae_type, variant, vlm_enhance, vlm_model, vlm_system_prompt, *_args, **_kwargs):
     variant = variant or 'Bi-Directional'
     if variant == 'None':
         log.error('FramePack: no model selected')
@@ -137,7 +137,7 @@ def run_framepack(task_id, _ui_state, init_image, end_image, start_weight, end_w
     with call_queue.get_lock():
         progress.start_task(task_id)
 
-        yield from load_model(variant, attention)
+        yield from load_model(variant)
         if shared.sd_model_type != 'hunyuanvideo':
             progress.finish_task(task_id)
             yield gr.update(), gr.update(), 'Model load failed'

modules/minimax/minimax_ui.py

@@ -25,16 +25,17 @@ def create_ui(prompt, _negative, styles, overrides, script_inputs, mp4_fps, mp4_
                 with gr.Row():
                     width, height = ui_sections.create_resolution_inputs('minimax', default_width=1024, default_height=576, step=32)
                 with gr.Row():
-                    steps = gr.Slider(minimum=2, maximum=100, step=1, label="MiniMax Steps", elem_id='minimax_steps', value=30)
-                    frames = gr.Slider(label='MiniMax Frames', minimum=22, maximum=362, step=17, value=124, elem_id='minimax_frames')
+                    steps = gr.Slider(minimum=2, maximum=100, step=1, label="MiniMax steps", elem_id='minimax_steps', value=30)
+                    frames = gr.Slider(label='MiniMax frames', minimum=22, maximum=362, step=17, value=124, elem_id='minimax_frames')
                 with gr.Row():
-                    video_shift = gr.Slider(minimum=8.0, maximum=16.0, step=0.1, label="MiniMax Video Shift", elem_id='minimax_video_shift', value=12)
-                    audio_shift = gr.Slider(minimum=1.5, maximum=6.0, step=0.1, label="MiniMax Audio Shift", elem_id='minimax_audio_shift', value=3)
+                    video_shift = gr.Slider(minimum=8.0, maximum=16.0, step=0.1, label="MiniMax video shift", elem_id='minimax_video_shift', value=12)
+                    audio_shift = gr.Slider(minimum=1.5, maximum=6.0, step=0.1, label="MiniMax audio shift", elem_id='minimax_audio_shift', value=3)
                 with gr.Row():
                     seed = gr.Number(label='Seed', value=-1, elem_id='minimax_seed', container=True)
                     random_seed = ToolButton(ui_symbols.random, elem_id='minimax_seed_random')
                     random_seed.click(fn=lambda: -1, show_progress='hidden', inputs=[], outputs=[seed])
-                    audio_enable = gr.Checkbox(label='Audio Enabled', value=True, elem_id="minimax_audio_enable")
+                with gr.Row():
+                    audio_enable = gr.Checkbox(label='Enable audio', value=True, elem_id="minimax_audio_enable")
             with gr.Accordion(open=False, label="Input media", elem_id='minimax_input_media_accordion', visible=True) as input_accordion:
                 with gr.Row():
                     init_image = gr.Image(label='Image', elem_id='minimax_init_image', type='pil', image_mode='RGB', width=256, height=256)

modules/modular.py

@@ -7,7 +7,7 @@ from modules.logger import log
 debug = os.environ.get('SD_MODULAR_DEBUG', None) is not None
 
 
-def get_modular_class_name(diffusion_pipeline: diffusers.DiffusionPipeline) -> str:
+def get_modular_class(diffusion_pipeline: diffusers.DiffusionPipeline):
     name = diffusion_pipeline.__class__.__name__
     name = name.replace('Pipeline', '').replace('Img2Img', '').replace('Inpaint', '').replace('ImageToVideo', '')
     name = f'{name}AutoBlocks'
@@ -20,7 +20,7 @@ def get_modular_class_name(diffusion_pipeline: diffusers.DiffusionPipeline) -> s
 def is_compatible(diffusion_pipeline: diffusers.DiffusionPipeline) -> bool:
     if not shared.opts.model_modular_enable:
         return False
-    compatible = get_modular_class_name(diffusion_pipeline) is not None
+    compatible = get_modular_class(diffusion_pipeline) is not None
     if not compatible:
         log.debug(f'Modular: source={diffusion_pipeline.__class__.__name__} incompatible pipeline')
     return compatible
@@ -34,7 +34,7 @@ def is_guider(diffusion_pipeline: diffusers.DiffusionPipeline) -> bool:
 def convert_to_modular(diffusion_pipeline: diffusers.DiffusionPipeline) -> diffusers.ModularPipeline:
     modular_pipe = None
     try:
-        modular_cls = get_modular_class_name(diffusion_pipeline)
+        modular_cls = get_modular_class(diffusion_pipeline)
         if modular_cls is None:
             raise ValueError(f'unknown: cls={diffusion_pipeline.__class__.__name__}')
         modular_blocks = modular_cls()

modules/sd_hijack_modular.py

@@ -19,8 +19,11 @@ def modular_step(state: diffusers.modular_pipelines.modular_pipeline.BlockState)
         shared.state.sampling_steps = state.num_inference_steps
     if 'latents' in keys:
         shared.state.step()
+        # TODO: MiniMax runs AfterDenoise only at the end of the loop and it does latents unpacking
         shared.state.current_latent = state.latents
         lora_stack.on_step(shared.state.sampling_step)
+        if debug:
+            log.trace(f'Modular step: step={shared.state.sampling_step} latent={state.latents.shape}')
     if shared.state.interrupted or shared.state.skipped:
         raise AssertionError('Interrupted...')
     if shared.state.paused:
@@ -60,7 +63,7 @@ def patch_blocks(blocks: diffusers.ModularPipelineBlocks):
                 block_cls._is_patched = True # pylint: disable=protected-access
                 intercepted.add(block_cls)
                 if debug:
-                    log.trace(f'Modular hijack: {block_cls.__name__}')
+                    log.trace(f'Modular hijack: block={block_cls.__name__}')
         for attr in ("sub_blocks", "blocks"): # recurse into child blocks if containers exist
             sub = getattr(current_block, attr, None)
             if isinstance(sub, dict):

modules/taesd/taehv.py

@@ -9,7 +9,6 @@ import torch.nn as nn
 import torch.nn.functional as F
 from tqdm.auto import tqdm
 
-DecoderResult = namedtuple("DecoderResult", ("frame", "memory"))
 TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index"))
 
 def conv(n_in, n_out, **kwargs):
@@ -28,6 +27,18 @@ class MemBlock(nn.Module):
     def forward(self, x, past):
         return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x))
 
+class SuperMemBlock(nn.Module):
+    """MemBlock variant used by the Super decoder (ConvNeXt-style: 7x7 depthwise conv + inverted bottleneck)."""
+    def __init__(self, n_f):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(n_f*2, n_f*2, 7, padding=3, groups=n_f*2, bias=False),
+            nn.Conv2d(n_f*2, n_f*4, 1), nn.ReLU(inplace=True),
+            nn.Conv2d(n_f*4, n_f, 1, bias=False),
+        )
+    def forward(self, x, past):
+        return self.conv(torch.cat([x, past], 1)) + x
+
 class TPool(nn.Module):
     def __init__(self, n_f, stride):
         super().__init__()
@@ -47,6 +58,106 @@ class TGrow(nn.Module):
         x = self.conv(x)
         return x.reshape(-1, C, H, W)
 
+def apply_model_with_memblocks_parallel(model, x, show_progress_bar):
+    """
+    Apply a sequential model with memblocks to the given input,
+    with parallelization over the time axis and iteration over blocks.
+
+    Args:
+    - model: nn.Sequential of blocks to apply
+    - x: input data, of dimensions NTCHW
+    - 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
+    x = x.reshape(N*T, C, H, W)
+
+    # parallel over input timesteps, iterate over blocks
+    for b in tqdm(model, disable=not show_progress_bar):
+        if isinstance(b, (MemBlock, SuperMemBlock)):
+            NT, C, H, W = x.shape
+            T = NT // N
+            _x = x.reshape(N, T, C, H, W)
+            # pad with zeros along time axis (i.e. empty memory), slice
+            block_memory = F.pad(_x, (0,0,0,0,0,0,1,0), value=0)[:,:T].reshape(x.shape)
+            x = b(x, block_memory)
+        else:
+            x = b(x)
+    NT, C, H, W = x.shape
+    T = NT // N
+    return x.view(N, T, C, H, W)
+
+def apply_model_with_memblocks_sequential_single_step(model, memory, work_queue, progress_bar=None):
+    """
+    Process the work queue (a graph traversal over blocks and timesteps)
+    until an output frame is produced or the queue is empty.
+    Mutates memory and work_queue in place.
+
+    Returns N1CHW output tensor, or None if the queue needs more input.
+    """
+    while work_queue:
+        xt, i = work_queue.pop(0)
+        if progress_bar is not None and i == 0:
+            progress_bar.update(1)
+        if i == len(model):
+            return xt.unsqueeze(1)
+        b = model[i]
+        if isinstance(b, (MemBlock, SuperMemBlock)):
+            # mem blocks are simple since we're visiting the graph in causal order
+            if memory[i] is None:
+                xt_new = b(xt, xt * 0)
+            else:
+                xt_new = b(xt, memory[i])
+            memory[i] = xt
+            work_queue.insert(0, TWorkItem(xt_new, i+1))
+        elif isinstance(b, TPool):
+            # pool blocks accumulate inputs until they have enough to pool
+            if memory[i] is None:
+                memory[i] = []
+            memory[i].append(xt)
+            if len(memory[i]) > b.stride:
+                raise ValueError(f"TPool memory overflow: {len(memory[i])} items for stride {b.stride}")
+            elif len(memory[i]) == b.stride:
+                N, C, H, W = xt.shape
+                xt = b(torch.cat(memory[i], 1).view(N*b.stride, C, H, W))
+                memory[i] = []
+                work_queue.insert(0, TWorkItem(xt, i+1))
+        elif isinstance(b, TGrow):
+            xt = b(xt)
+            NT, C, H, W = xt.shape
+            for xt_next in reversed(xt.view(NT//b.stride, b.stride*C, H, W).chunk(b.stride, 1)):
+                work_queue.insert(0, TWorkItem(xt_next, i+1))
+        else:
+            xt = b(xt)
+            work_queue.insert(0, TWorkItem(xt, i+1))
+    return None
+
+def apply_model_with_memblocks_sequential(model, x, show_progress_bar):
+    """
+    Apply a sequential model with memblocks to the given input,
+    with iteration over timesteps as well as blocks.
+
+    Args:
+    - model: nn.Sequential of blocks to apply
+    - x: input data, of dimensions NTCHW
+    - 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"
+    work_queue = [TWorkItem(xt, 0) for xt in x.unbind(1)]
+    memory = [None] * len(model)
+    progress_bar = tqdm(range(len(work_queue)), disable=not show_progress_bar)
+    out = []
+    while work_queue:
+        xt = apply_model_with_memblocks_sequential_single_step(model, memory, work_queue, progress_bar)
+        if xt is not None:
+            out.append(xt)
+    progress_bar.close()
+    return torch.cat(out, 1)
+
 def apply_model_with_memblocks(model, x, parallel, show_progress_bar):
     """
     Apply a sequential model with memblocks to the given input.
@@ -59,141 +170,72 @@ def apply_model_with_memblocks(model, x, parallel, show_progress_bar):
 
     Returns NTCHW tensor of output data.
     """
-    if x.ndim == 4:
-        x = x.unsqueeze(0)
-    assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor"
-    if x.shape[1] == 16 and x.shape[2] != 16:
-        x = x.transpose(1,2) # NCTHW to NTCHW
-    N, T, C, H, W = x.shape
     if parallel:
-        x = x.reshape(N*T, C, H, W)
-        # parallel over input timesteps, iterate over blocks
-        for b in tqdm(model, disable=not show_progress_bar):
-            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)
-                x = b(x, mem)
-            else:
-                x = b(x)
-        NT, C, H, W = x.shape
-        T = NT // N
-        x = x.view(N, T, C, H, W)
+        return apply_model_with_memblocks_parallel(model, x, show_progress_bar)
     else:
-        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))]
-        # in addition to manually managing our queue, we also need to manually manage our progressbar.
-        # we'll update it for every source node that we consume.
-        progress_bar = tqdm(range(T), disable=not show_progress_bar)
-        # 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)
-            if i == 0:
-                # new source node consumed
-                progress_bar.update(1)
-            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))
-        progress_bar.close()
-        x = torch.stack(out, 1)
-    return x
+        return apply_model_with_memblocks_sequential(model, x, show_progress_bar)
 
 class TAEHV(nn.Module):
-    latent_channels = 16
-    image_channels = 3
-    def __init__(self, checkpoint_path="taehv.pth", decoder_time_upscale=(True, True), decoder_space_upscale=(True, True, True)):
+    def __init__(self, checkpoint_path="taehv.pth", encoder_time_downscale=(True, True, False), decoder_time_upscale=(False, True, True), decoder_space_upscale=(True, True, True), patch_size=1, latent_channels=16, arch_variant=None):
         """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.
+            encoder_time_downscale: whether temporal downsampling is enabled for each block.
             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.
+            patch_size: input/output pixelshuffle patch-size for this model.
+            latent_channels: number of latent channels (z dim) for this model.
+            arch_variant: decoder architecture variant. None (base) or "super" (higher-quality, ~2x decoder params). Autodetected from filename if None.
         """
         super().__init__()
-        from modules import shared
+        self.patch_size = patch_size
+        self.latent_channels = latent_channels
+        self.image_channels = 3
+        if len(decoder_time_upscale) == 2:
+            decoder_time_upscale = (False, *decoder_time_upscale)
+        self.is_cogvideox = checkpoint_path is not None and "taecvx" in checkpoint_path
+        self.is_h3 = checkpoint_path is not None and "taeh3" in checkpoint_path
+        if checkpoint_path is not None and "taew2_2" in checkpoint_path:
+            self.patch_size, self.latent_channels = 2, 48
+        if checkpoint_path is not None and "taehv1_5" in checkpoint_path:
+            self.patch_size, self.latent_channels = 2, 32
+        if self.is_h3:
+            self.patch_size, self.latent_channels, encoder_time_downscale = 2, 24, (True, True, False)
+        if checkpoint_path is not None and "taeltx" in checkpoint_path: # same for both 2 and 2.3
+            self.patch_size, self.latent_channels, encoder_time_downscale, decoder_time_upscale = 4, 128, (True, True, True), (True, True, True)
+        if arch_variant is None and checkpoint_path is not None and "_super" in checkpoint_path:
+            arch_variant = "super"
+        assert arch_variant in (None, "super"), f"unrecognized arch_variant {arch_variant!r}"
         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),
+            conv(self.image_channels*self.patch_size**2, 64), nn.ReLU(inplace=True),
+            TPool(64, 2 if encoder_time_downscale[0] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
+            TPool(64, 2 if encoder_time_downscale[1] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
+            TPool(64, 2 if encoder_time_downscale[2] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64),
+            conv(64, self.latent_channels),
         )
-        n_f = [256, 128, 64, 64]
-        self.frames_to_trim = 2**sum(decoder_time_upscale) - 1
-
-        if shared.opts.taesd_layers == 1:
+        if arch_variant == "super":
+            n_f = [512, 256, 128, 64]
             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.Identity(), 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.Identity(), 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),
-            )
-        elif shared.opts.taesd_layers == 2:
-            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.Identity(), 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),
+                nn.Conv2d(self.latent_channels, n_f[0], 1, bias=False),
+                SuperMemBlock(n_f[0]), SuperMemBlock(n_f[0]), SuperMemBlock(n_f[0]), conv(n_f[0], n_f[1]*(2 if decoder_space_upscale[0] else 1)**2), nn.ReLU(inplace=True), nn.PixelShuffle(2 if decoder_space_upscale[0] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1),
+                SuperMemBlock(n_f[1]), SuperMemBlock(n_f[1]), SuperMemBlock(n_f[1]), conv(n_f[1], n_f[2]*(2 if decoder_space_upscale[1] else 1)**2), nn.ReLU(inplace=True), nn.PixelShuffle(2 if decoder_space_upscale[1] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1),
+                SuperMemBlock(n_f[2]), SuperMemBlock(n_f[2]), SuperMemBlock(n_f[2]), conv(n_f[2], n_f[3]*(2 if decoder_space_upscale[2] else 1)**2), nn.ReLU(inplace=True), nn.PixelShuffle(2 if decoder_space_upscale[2] else 1), TGrow(n_f[3], 2 if decoder_time_upscale[2] else 1),
+                conv(n_f[3], self.image_channels*self.patch_size**2),
             )
         else:
+            n_f = [256, 128, 64, 64]
             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),
+                Clamp(), conv(self.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], 2 if decoder_time_upscale[0] else 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[1] 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[2] else 1), conv(n_f[2], n_f[3], bias=False),
+                nn.ReLU(inplace=True), conv(n_f[3], self.image_channels*self.patch_size**2),
             )
+        # computed properties
+        self.t_downscale = 2**sum(t.stride == 2 for t in self.encoder if isinstance(t, TPool))
+        self.t_upscale = 2**sum(t.stride == 2 for t in self.decoder if isinstance(t, TGrow))
+        self.frames_to_trim = self.t_upscale - 1
 
         if checkpoint_path is not None:
             self.load_state_dict(self.patch_tgrow_layers(torch.load(checkpoint_path, map_location="cpu", weights_only=True)))
@@ -213,6 +255,38 @@ class TAEHV(nn.Module):
                     sd[key] = sd[key][-new_sd[key].shape[0]:]
         return sd
 
+    def preprocess_input_frames(self, x):
+        """Preprocess RGB input frames prior to the main encoder sequence."""
+        if self.patch_size > 1:
+            x = F.pixel_unshuffle(x, self.patch_size)
+        return x
+
+    def _encode_h3_video(self, x, parallel, show_progress_bar):
+        """Match H3's 17-frame chunks and three-token drop.
+        https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/fa9c8ab1eaa21c8ae25e7e40b83b2e6002f340af/FL2VA/video_vae/klvae.py#L461-L503
+        """
+        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))
+        x = self.preprocess_input_frames(x)
+        if parallel:
+            x = apply_model_with_memblocks(self.encoder, x.flatten(0, 1), True, show_progress_bar)
+            x = x.reshape(batch, -1, *x.shape[2:])
+        else:
+            x = torch.cat([apply_model_with_memblocks(self.encoder, chunk, False, False)
+                           for chunk in tqdm(x.unbind(1), disable=not show_progress_bar)], dim=1)
+        return x[:, :-3]
+
+    def _decode_h3_video(self, x, parallel, show_progress_bar):
+        """Match H3's five-token chunks and per-chunk prefix trim.
+        https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/fa9c8ab1eaa21c8ae25e7e40b83b2e6002f340af/FL2VA/video_vae/klvae.py#L678-L786
+        """
+        x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar)
+        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 self.postprocess_output_frames(x[:, :-3 * self.t_upscale])
+
     def encode_video(self, x, parallel=True, show_progress_bar=True):
         """Encode a sequence of frames.
 
@@ -223,23 +297,43 @@ class TAEHV(nn.Module):
               if False, frames will be processed sequentially.
         Returns NTCHW latent tensor with ~Gaussian values.
         """
+        if self.is_h3:
+            return self._encode_h3_video(x, parallel, show_progress_bar)
+        x = self.preprocess_input_frames(x)
+        if x.shape[1] % self.t_downscale != 0:
+            # pad at end to multiple of self.t_downscale
+            n_pad = self.t_downscale - x.shape[1] % self.t_downscale
+            padding = x[:, -1:].repeat_interleave(n_pad, dim=1)
+            x = torch.cat([x, padding], 1)
         return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar)
 
+    def postprocess_output_frames(self, x):
+        """Postprocess RGB frames after the main decoder sequence."""
+        if self.patch_size > 1:
+            x = F.pixel_shuffle(x, self.patch_size)
+        return x.clamp_(0, 1)
+
     def decode_video(self, x, parallel=True, show_progress_bar=True):
         """Decode a sequence of frames.
 
         Args:
-            x: input NTCHW latent (C=12) tensor with ~Gaussian values.
+            x: input NTCHW latent (C=self.latent_channels) 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.
         """
+        if self.is_h3:
+            return self._decode_h3_video(x, parallel, show_progress_bar)
+        skip_trim = self.is_cogvideox and x.shape[1] % 2 == 0
         x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar)
-        return x[:, self.frames_to_trim:] * 2.0 - 1.0
-
-    def forward(self, x):
-        return self.c(x)
+        x = self.postprocess_output_frames(x)
+        if skip_trim:
+            # skip trimming for cogvideox to make frame counts match.
+            # this still doesn't have correct temporal alignment for certain frame counts
+            # (cogvideox seems to pad at the start?), but for multiple-of-4 it's fine.
+            return x
+        return x[:, self.frames_to_trim:]
 
     def decode(self, x, parallel=True, show_progress_bar=False, return_dict=False): # pylint: disable=unused-argument
         """Decode a sequence of frames."""
@@ -248,3 +342,193 @@ class TAEHV(nn.Module):
     def encode(self, x, parallel=True, show_progress_bar=False, return_dict=False): # pylint: disable=unused-argument
         """Encode a sequence of frames."""
         return self.encode_video(x, parallel=False, show_progress_bar=False)
+
+class StreamingTAEHV(nn.Module):
+    def __init__(self, taehv):
+        """Streaming wrapper around TAEHV for real-time use-cases (where not all inputs are available immediately).
+
+        Encode-decode (video-to-video) usage:
+            streaming = StreamingTAEHV(taehv)
+            for frame in video_frames:
+                latent = streaming.encode(frame_tensor)
+                decoded = streaming.decode(latent)  # feeds latent if not None, then returns next frame
+                if decoded is not None:
+                    display(decoded)
+            for frame in streaming.flush():
+                display(frame)
+
+        Decode-only (world model) usage:
+            streaming = StreamingTAEHV(taehv)
+            while running:
+                latent = world_model.step()       # latent represents t_upscale frames
+                frame = streaming.decode(latent)   # returns first frame immediately
+                while frame is not None:           # retrieve remaining frames from this latent
+                    display(frame)
+                    frame = streaming.decode()
+        """
+        super().__init__()
+        self.taehv = taehv
+        self.reset()
+
+    def reset(self):
+        """Reset all internal state. Call this to start encoding/decoding a new stream."""
+        self.encoder_work_queue, self.encoder_memory = [], [None] * len(self.taehv.encoder)
+        self.decoder_work_queue, self.decoder_memory = [], [None] * len(self.taehv.decoder)
+        self.n_frames_encoded, self.n_frames_decoded = 0, 0
+        self._last_encoder_input_frame = None
+
+    def encode(self, x=None):
+        """Feed an input frame (optional) and try to produce an encoder output.
+
+        The encoder accumulates t_downscale input frames before producing one latent,
+        so most calls will return None. Use flush_encoder() at end-of-stream to pad and
+        drain any remaining latents.
+
+        Args:
+            x: NTCHW RGB frame tensor with values in [0, 1], or None to just process pending work.
+        Returns: N1CHW latent tensor, or None if not enough input has been accumulated.
+        """
+        if x is not None:
+            assert x.ndim == 5 and x.shape[2] == self.taehv.image_channels, f"Expected NTCHW frames but got {x.shape=}"
+            self._last_encoder_input_frame = x[:, -1:] # pylint: disable=attribute-defined-outside-init
+            x = self.taehv.preprocess_input_frames(x)
+            self.encoder_work_queue.extend(TWorkItem(xt, 0) for xt in x.unbind(1))
+            self.n_frames_encoded += x.shape[1]
+        xt = apply_model_with_memblocks_sequential_single_step(
+            self.taehv.encoder, self.encoder_memory, self.encoder_work_queue)
+        return xt
+
+    def decode(self, x=None):
+        """Feed a latent (optional) and try to produce a decoded frame.
+
+        Each latent produces t_upscale output frames due to temporal upscaling. The first
+        decode(latent) call returns the first of these frames; call decode() with no argument
+        to retrieve the rest, one at a time. Each call does the minimum decoder work needed to
+        produce one frame.
+
+        Startup frames (the first frames_to_trim raw decoder outputs, used for causal alignment
+        with the reference VAE) are consumed internally and never returned.
+
+        Args:
+            x: NTCHW latent tensor, or None to retrieve the next pending frame.
+        Returns: N1CHW decoded RGB frame tensor, or None if the queue needs more input.
+        """
+        if x is not None:
+            assert x.ndim == 5 and x.shape[2] == self.taehv.latent_channels, f"Expected NTCHW latents but got {x.shape=}"
+            self.decoder_work_queue.extend(TWorkItem(xt, 0) for xt in x.unbind(1))
+        while True:
+            xt = apply_model_with_memblocks_sequential_single_step(
+                self.taehv.decoder, self.decoder_memory, self.decoder_work_queue)
+            if xt is None:
+                return None
+            self.n_frames_decoded += 1
+            # skip startup frames (to match decode_video trim behavior)
+            if not self.taehv.is_cogvideox and self.n_frames_decoded <= self.taehv.frames_to_trim:
+                continue
+            return self.taehv.postprocess_output_frames(xt)
+
+    def flush_encoder(self):
+        """Pad (if needed) and drain all remaining latents from the encoder.
+
+        Returns list of N1CHW latent tensors.
+        """
+        latents = []
+        if self._last_encoder_input_frame is not None and self.n_frames_encoded % self.taehv.t_downscale != 0:
+            n_pad = self.taehv.t_downscale - self.n_frames_encoded % self.taehv.t_downscale
+            for _ in range(n_pad):
+                lat = self.encode(self._last_encoder_input_frame)
+                if lat is not None:
+                    latents.append(lat)
+        while (lat := self.encode()) is not None:
+            latents.append(lat)
+        return latents
+
+    def flush_decoder(self):
+        """Drain all remaining decoded frames from the decoder.
+
+        Returns list of N1CHW decoded RGB frame tensors.
+        """
+        frames = []
+        while (frame := self.decode()) is not None:
+            frames.append(frame)
+        return frames
+
+    def flush(self):
+        """Flush encoder (with padding) and decoder, returning all remaining decoded frames.
+
+        Returns list of N1CHW decoded RGB frame tensors.
+        """
+        frames = []
+        for latent in self.flush_encoder():
+            frame = self.decode(latent)
+            if frame is not None:
+                frames.append(frame)
+        frames.extend(self.flush_decoder())
+        return frames
+
+@torch.no_grad()
+def main():
+    """Run TAEHV roundtrip reconstruction on the given video paths."""
+    import os
+    import sys
+    import cv2 # no highly esteemed deed is commemorated here
+
+    class VideoTensorReader:
+        def __init__(self, video_file_path):
+            self.cap = cv2.VideoCapture(video_file_path)
+            assert self.cap.isOpened(), f"Could not load {video_file_path}"
+            self.fps = self.cap.get(cv2.CAP_PROP_FPS)
+        def __iter__(self):
+            return self
+        def __next__(self):
+            ret, frame = self.cap.read()
+            if not ret:
+                self.cap.release()
+                raise StopIteration  # End of video or error
+            return torch.from_numpy(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).permute(2, 0, 1) # BGR HWC -> RGB CHW
+
+    class VideoTensorWriter:
+        def __init__(self, video_file_path, width_height, fps=30):
+            self.writer = cv2.VideoWriter(video_file_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, width_height)
+            assert self.writer.isOpened(), f"Could not create writer for {video_file_path}"
+        def write(self, frame_tensor):
+            assert frame_tensor.ndim == 3 and frame_tensor.shape[0] == 3, f"{frame_tensor.shape}??"
+            self.writer.write(cv2.cvtColor(frame_tensor.permute(1, 2, 0).numpy(), cv2.COLOR_RGB2BGR)) # RGB CHW -> BGR HWC
+        def __del__(self):
+            if hasattr(self, 'writer'):
+                self.writer.release()
+
+    dev = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
+    dtype = torch.float16
+    checkpoint_path = os.getenv("TAEHV_CHECKPOINT_PATH", "taehv.pth")
+    checkpoint_name = os.path.splitext(os.path.basename(checkpoint_path))[0]
+    print(f"Using device \033[31m{dev}\033[0m, dtype \033[32m{dtype}\033[0m, checkpoint \033[34m{checkpoint_name}\033[0m ({checkpoint_path})")
+    taehv = TAEHV(checkpoint_path=checkpoint_path).to(dev, dtype)
+    for video_path in sys.argv[1:]:
+        print(f"Processing {video_path}...")
+        video_in = VideoTensorReader(video_path)
+        video = torch.stack(list(video_in), 0)[None]
+        vid_dev = video.to(dev, dtype).div_(255.0)
+        # convert to device tensor
+        if video.numel() < 100_000_000:
+            print(f"  {video_path} seems small enough, will process all frames in parallel")
+            # convert to device tensor
+            vid_enc = taehv.encode_video(vid_dev)
+            print(f"  Encoded {video_path} -> {vid_enc.shape}. Decoding...")
+            vid_dec = taehv.decode_video(vid_enc)
+            print(f"  Decoded {video_path} -> {vid_dec.shape}")
+        else:
+            print(f"  {video_path} seems large, will process each frame sequentially")
+            # convert to device tensor
+            vid_enc = taehv.encode_video(vid_dev, parallel=False)
+            print(f"  Encoded {video_path} -> {vid_enc.shape}. Decoding...")
+            vid_dec = taehv.decode_video(vid_enc, parallel=False)
+            print(f"  Decoded {video_path} -> {vid_dec.shape}")
+        video_out_path = video_path + f".reconstructed_by_{checkpoint_name}.mp4"
+        video_out = VideoTensorWriter(video_out_path, (vid_dec.shape[-1], vid_dec.shape[-2]), fps=int(round(video_in.fps)))
+        for frame in vid_dec.clamp_(0, 1).mul_(255).round_().byte().cpu()[0]:
+            video_out.write(frame)
+        print(f"  Saved to {video_out_path}")
+
+if __name__ == "__main__":
+    main()

modules/vae/sd_vae_taesd.py

@@ -27,6 +27,7 @@ TAESD_MODELS = {
     'TAE HunyuanVideo': { 'fn': 'taehv.pth', 'uri': 'https://github.com/madebyollin/taehv/raw/refs/heads/main/taehv.pth', 'model': None },
     'TAE WanVideo': { 'fn': 'taew1.pth', 'uri': 'https://github.com/madebyollin/taehv/raw/refs/heads/main/taew2_1.pth', 'model': None },
     'TAE MochiVideo': { 'fn': 'taem1.pth', 'uri': 'https://github.com/madebyollin/taem1/raw/refs/heads/main/taem1.pth', 'model': None },
+    'TAE MiniMax-H3': { 'fn': 'taeh3.pth', 'uri': 'https://github.com/madebyollin/taehv/raw/refs/heads/main/taeh3.pth', 'model': None },
 }
 CQYAN_MODELS = {
     'Hybrid-Tiny SD': {
@@ -79,6 +80,8 @@ def get_model(model_cls, variant=None):
         variant = 'TAE SD3'
     elif model_cls in {'wanai', 'qwen', 'chrono', 'cosmos', 'anima', 'fibo', 'joy', 'krea2'}:
         variant = 'TAE WanVideo'
+    elif model_cls in {'minimaxh3'}:
+        variant = 'TAE MiniMax-H3'
     else:
         warn_once(f'cls={shared.sd_model.__class__.__name__} type={shared.sd_model_type} unsuppported', variant=variant)
         return model_cls, None
@@ -122,10 +125,7 @@ def load_model(model_type = 'decoder', variant = None, vae_file: str | None = No
             log.print() # new line
             log.debug(f'Decode: type="taesd" variant="{variant}" fn="{fn}" layers={shared.opts.taesd_layers} load')
             vae = None
-            if 'TAE HunyuanVideo' in variant:
-                from modules.taesd.taehv import TAEHV
-                vae = TAEHV(checkpoint_path=fn)
-            elif 'TAE WanVideo' in variant:
+            if ('TAE HunyuanVideo' in variant) or ('TAE WanVideo' in variant) or ('TAE MiniMax-H3' in variant):
                 from modules.taesd.taehv import TAEHV
                 vae = TAEHV(checkpoint_path=fn)
             elif 'TAE MochiVideo' in variant:
@@ -194,7 +194,7 @@ def decode(latents, fast=False):
                 variant = prev_variant
             else:
                 vae, variant = load_model(model_type='decoder')
-                if vae is None or max(latents.shape) > 256: # safety check of large tensors
+                if vae is None or max(latents.shape) > 384: # safety check of large tensors
                     return latents
                 prev_model = vae
                 prev_variant = variant