diff --git a/CHANGELOG.md b/CHANGELOG.md index e2fcfe1b1..95107a94f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,9 @@ to reduce memory usage and improve performance on large models - **Internal** - modular pipelines intercept and profiling hooks + - modular pipelines convert standard model on-demand + - modular pipelines previews + - modular pipelines basic guidance - attention mechanisms decision tree and apply method refactor ## Update for 2026-08-26 diff --git a/TODO.md b/TODO.md index e31f08986..ffa870508 100644 --- a/TODO.md +++ b/TODO.md @@ -3,7 +3,7 @@ ## Short-term - MiniMax LoRA: native loader for MiniMax-H3: fl2va, ref2va, pruned, @CalamitousFelicitousness -- MiniMax TAESD: +- MiniMax TAESD: need to unpack latents, @vladmandic - MiniMax: Create pre-quant for MiniMax-H3-Turbo - Benchmark tool productize: @CalamitousFelicitousness - Inpaint: https://discord.com/channels/1101998836328697867/1130536562422186044/1506850651035144322, @vladmandic diff --git a/modules/framepack/framepack_api.py b/modules/framepack/framepack_api.py index 9cc2ace1a..74cdbb2f0 100644 --- a/modules/framepack/framepack_api.py +++ b/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, diff --git a/modules/framepack/framepack_install.py b/modules/framepack/framepack_install.py index 404e89652..905010f9d 100644 --- a/modules/framepack/framepack_install.py +++ b/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') diff --git a/modules/framepack/framepack_ui.py b/modules/framepack/framepack_ui.py index c9fb51340..c2edbbb0f 100644 --- a/modules/framepack/framepack_ui.py +++ b/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, ] diff --git a/modules/framepack/framepack_wrappers.py b/modules/framepack/framepack_wrappers.py index 94eb5155d..1ac40aee3 100644 --- a/modules/framepack/framepack_wrappers.py +++ b/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' diff --git a/modules/minimax/minimax_ui.py b/modules/minimax/minimax_ui.py index 314628b06..56dbf0e65 100644 --- a/modules/minimax/minimax_ui.py +++ b/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) diff --git a/modules/modular.py b/modules/modular.py index 32a35d1a8..3916ca210 100644 --- a/modules/modular.py +++ b/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() diff --git a/modules/sd_hijack_modular.py b/modules/sd_hijack_modular.py index 2c06a0c14..2aa77fcdf 100644 --- a/modules/sd_hijack_modular.py +++ b/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): diff --git a/modules/taesd/taehv.py b/modules/taesd/taehv.py index 4e974deae..9239ac864 100644 --- a/modules/taesd/taehv.py +++ b/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() diff --git a/modules/vae/sd_vae_taesd.py b/modules/vae/sd_vae_taesd.py index f40d3b8cc..0436339af 100644 --- a/modules/vae/sd_vae_taesd.py +++ b/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