diff --git a/fastvideo/configs/pipelines/open_h3_regenerate_2k.py b/fastvideo/configs/pipelines/open_h3_regenerate_2k.py new file mode 100644 index 000000000..83dd12452 --- /dev/null +++ b/fastvideo/configs/pipelines/open_h3_regenerate_2k.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pipeline configuration for Open-H3-Regenerate-2K (768P H3 draft -> 2560x1440).""" + +from __future__ import annotations + +from dataclasses import dataclass + +from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig + +DRAFT_SIZE_CHOICES = ("768p", "576p") + + +@dataclass +class OpenH3Regenerate2KPipelineConfig(MiniMaxH3PipelineConfig): + """MiniMax-H3 policy plus the refine-stage knobs. + + ``draft_size`` is the size the draft video is VAE-encoded at: 768p (1344x768) for MiniMax-H3 drafts, 576p + (1024x576) for drafts from other video models. ``refine_start_sigma`` re-noises the upscaled draft to another + level; the later steps then follow the shift-6 rule (None keeps the released grid). + """ + + draft_size: str = "768p" + refine_start_sigma: float | None = None + + def check_pipeline_config(self) -> None: + super().check_pipeline_config() + if self.draft_size not in DRAFT_SIZE_CHOICES: + raise ValueError(f"draft_size must be one of {DRAFT_SIZE_CHOICES}, got {self.draft_size!r}.") + if self.refine_start_sigma is not None and not 0.0 < self.refine_start_sigma < 1.0: + raise ValueError(f"refine_start_sigma must be in (0, 1), got {self.refine_start_sigma}.") diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 8cea9dde7..a9e3b801a 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -1655,6 +1655,8 @@ class PipelineComponentLoader: component_model_path, ) + # A component under a subfolder (e.g. "base_model/vae") dispatches on its own name. + module_name = module_name.rsplit("/", 1)[-1] # Get the appropriate loader for this module type loader = ComponentLoader.for_module_type(module_name, transformers_or_diffusers) diff --git a/fastvideo/models/upsamplers/minimax_h3_latent_upscaler.py b/fastvideo/models/upsamplers/minimax_h3_latent_upscaler.py new file mode 100644 index 000000000..9ed762431 --- /dev/null +++ b/fastvideo/models/upsamplers/minimax_h3_latent_upscaler.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-H3 3D-convolution latent upscaler (LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler, MIT). + +Port of the upstream node's ``LatentResizer3D`` and its ``target dimensions`` size logic, without the ComfyUI shims, +attention blocks (the node forces them off) and temporal chunking (Open-H3-Regenerate-2K runs without it). Module names +match the upstream ``latent_upsampler/*.safetensors`` checkpoint, so it loads without remapping. + +The network works on a latent that was normalized with the statistics below, not the VAE-normalized latent the DiT +sees: ``upscale_latent`` therefore re-normalizes its input and de-normalizes the result with the same statistics. +""" +from __future__ import annotations + +import re + +import torch +import torch.nn.functional as F +from safetensors.torch import load_file +from torch import nn + +VAE_DOWNSAMPLE = 16 +PIXEL_ALIGNMENT = 32 + +LATENTS_MEAN = ( + 0.858090341091156, -0.9606591463088989, 1.0661640167236328, -0.5090325474739075, -0.2727581858634949, + -1.3675414323806763, -0.2553254961967468, -0.26907554268836975, -0.5376840829849243, -0.0464097298681736, + 0.6657370328903198, 0.19690127670764923, -0.5460608005523682, -0.4035342037677765, -0.23683024942874908, + 0.25928452610969543, -0.30133944749832153, 0.211341992020607, -1.1206848621368408, 0.3581933379173279, + -0.04225143790245056, 0.2604829967021942, 0.22864092886447906, 0.7056031823158264) +LATENTS_STD = ( + 1.2223774194717407, 1.2767263650894165, 1.6831774711608887, 1.7549455165863037, 1.5636216402053833, + 2.194143533706665, 0.9653137922286987, 1.0569885969161987, 0.841948926448822, 0.7729952931404114, + 1.8955937623977661, 0.946841835975647, 0.7996809482574463, 0.44988900423049927, 0.7197399735450745, + 0.6936293244361877, 2.961095094680786, 2.7694199085235596, 3.0496184825897217, 2.1088054180145264, + 3.276226282119751, 3.1627357006073, 2.2816812992095947, 2.6127843856811523) + + +def _norm(channels: int) -> nn.GroupNorm: + return nn.GroupNorm(32, channels) + + +class ResBlockEmb3D(nn.Module): + + def __init__(self, channels: int, emb_channels: int) -> None: + super().__init__() + self.in_layers = nn.Sequential(_norm(channels), nn.SiLU(), nn.Conv3d(channels, channels, 3, padding=1)) + self.emb_layers = nn.Sequential(nn.SiLU(), nn.Linear(emb_channels, 2 * channels)) + self.out_norm = _norm(channels) + self.out_layers = nn.Sequential(nn.SiLU(), nn.Dropout(p=0.0), nn.Conv3d(channels, channels, 3, padding=1)) + + def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + h = self.in_layers(x) + scale, shift = self.emb_layers(emb).type(h.dtype)[..., None, None, None].chunk(2, dim=1) + h = self.out_norm(h) * (1 + scale) + shift + return x + self.out_layers(h) + + +class TemporalConv(nn.Module): + + def __init__(self, channels: int, kernel_size: int = 5) -> None: + super().__init__() + self.norm = _norm(channels) + self.dwconv = nn.Conv3d(channels, channels, (kernel_size, 1, 1), padding=(kernel_size // 2, 0, 0), + groups=channels) + self.pwconv = nn.Conv3d(channels, channels, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.pwconv(self.dwconv(F.silu(self.norm(x)))) + + +def _blocks(count: int, channels: int, emb_dim: int, temporal_every: int, kernel: int) -> nn.ModuleList: + blocks = nn.ModuleList() + for index in range(count): + blocks.append(ResBlockEmb3D(channels, emb_dim)) + if temporal_every > 0 and index % temporal_every == 0: + blocks.append(TemporalConv(channels, kernel)) + return blocks + + +class MiniMaxH3LatentUpscaler(nn.Module): + + def __init__(self, + in_channels: int = 24, + in_blocks: int = 12, + out_blocks: int = 12, + channels: int = 512, + temporal_every: int = 2, + temporal_kernel: int = 5) -> None: + super().__init__() + emb_dim = 64 + self.conv_in = nn.Conv3d(in_channels, channels, 3, padding=1) + self.embed = nn.Sequential(nn.Linear(1, emb_dim), nn.SiLU(), nn.Linear(emb_dim, emb_dim)) + self.in_blocks = _blocks(in_blocks, channels, emb_dim, temporal_every, temporal_kernel) + self.out_blocks = _blocks(out_blocks, channels, emb_dim, temporal_every, temporal_kernel) + self.norm_out = _norm(channels) + self.conv_out = nn.Conv3d(channels, in_channels, 3, padding=1) + + def _run(self, blocks: nn.ModuleList, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: + for block in blocks: + x = block(x, emb.expand(x.shape[0], -1)) if isinstance(block, ResBlockEmb3D) else block(x) + return x + + def forward(self, x: torch.Tensor, scale: float, target_size: tuple[int, int, int]) -> torch.Tensor: + """``x``: [B, C, T, h, w] normalized latent; returns [B, C, T, *target_size[1:]].""" + emb = self.embed(torch.tensor([[scale - 1.0]], dtype=x.dtype, device=x.device)) + x = self._run(self.in_blocks, self.conv_in(x), emb) + x = F.interpolate(x, size=target_size, mode="trilinear", align_corners=False) + x = self._run(self.out_blocks, x, emb) + return self.conv_out(F.silu(self.norm_out(x))) + + @classmethod + def from_checkpoint(cls, path: str) -> MiniMaxH3LatentUpscaler: + state = load_file(path, device="cpu") + model = cls(in_channels=state["conv_in.weight"].shape[1], + channels=state["conv_in.weight"].shape[0], + in_blocks=_count_blocks(state, "in_blocks"), + out_blocks=_count_blocks(state, "out_blocks"), + temporal_kernel=_temporal_kernel(state)) + model.load_state_dict(state, strict=True) + return model.eval().requires_grad_(False) + + +def _count_blocks(state: dict[str, torch.Tensor], prefix: str) -> int: + pattern = re.compile(rf"{prefix}\.(\d+)\.in_layers\.0\.weight") + return len({int(m.group(1)) for key in state if (m := pattern.match(key))}) + + +def _temporal_kernel(state: dict[str, torch.Tensor]) -> int: + return next(tensor.shape[2] for key, tensor in state.items() if key.endswith("dwconv.weight")) + + +def resolve_upscale_geometry(input_size: tuple[int, int], target_pixels: tuple[int, int]) -> tuple[tuple[int, int], float]: + """Latent (h, w) after aligning the target pixel size to 32, and the scale the network is conditioned on. + + ``input_size`` is the latent (h, w); ``target_pixels`` the requested (height, width). + """ + h_in, w_in = input_size + h_target, w_target = target_pixels + scale = (w_target / (w_in * VAE_DOWNSAMPLE) + h_target / (h_in * VAE_DOWNSAMPLE)) / 2.0 + aligned = (round(round(h_target / PIXEL_ALIGNMENT) * PIXEL_ALIGNMENT / VAE_DOWNSAMPLE), + round(round(w_target / PIXEL_ALIGNMENT) * PIXEL_ALIGNMENT / VAE_DOWNSAMPLE)) + return (max(1, aligned[0]), max(1, aligned[1])), scale + + +@torch.inference_mode() +def upscale_latent(model: MiniMaxH3LatentUpscaler, latent: torch.Tensor, target_pixels: tuple[int, int], + dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: + """Upscale a [C, T, h, w] latent to ``target_pixels`` (height, width); FP32 [C, T, H, W] out. + + The input is returned unchanged when it already has the target latent size. + """ + channels, frames, h_in, w_in = latent.shape + (h_out, w_out), scale = resolve_upscale_geometry((h_in, w_in), target_pixels) + if (h_out, w_out) == (h_in, w_in): + return latent.float() + if scale < 1.0: + raise ValueError("The MiniMax-H3 latent upscaler only upscales (effective scale >= 1.0).") + mean = torch.tensor(LATENTS_MEAN, dtype=dtype, device=latent.device).view(1, -1, 1, 1, 1) + std = torch.tensor(LATENTS_STD, dtype=dtype, device=latent.device).view(1, -1, 1, 1, 1) + source = (latent[None].to(dtype) - mean) / std + result = model(source, scale, (frames, h_out, w_out)) + return (result * std + mean).float()[0] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/__init__.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/__init__.py new file mode 100644 index 000000000..2d9dc85f5 --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.pipelines.basic.open_h3_regenerate_2k.open_h3_regenerate_2k_pipeline import ( + EntryClass, + OpenH3Regenerate2KPipeline, +) + +__all__ = ["EntryClass", "OpenH3Regenerate2KPipeline"] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/draft_video.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/draft_video.py new file mode 100644 index 000000000..ce5efb6a4 --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/draft_video.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Decode a draft video into the 124-frame pixel block the H3 VAE turns into 37 latent frames.""" + +from __future__ import annotations + +import subprocess + +import numpy as np +import torch + +import fastvideo.envs as envs + +DRAFT_SIZES: dict[str, tuple[int, int]] = {"768p": (1344, 768), "576p": (1024, 576)} +DRAFT_FRAMES = 121 +# The causal VAE maps 124 frames to 37 latent frames; the last frame is held for the three extra ones. +VAE_PADDING_FRAMES = 3 +# Seconds of held last frame ffmpeg appends so a short draft still yields DRAFT_FRAMES. +_TAIL_PAD_SECONDS = 6 +_DECODE_TIMEOUT_S = 600 + + +def ffmpeg_decode_command(path: str, width: int, height: int) -> list[str]: + """fps=24, cover-scale (bilinear) and centre-crop, last frame held, exactly DRAFT_FRAMES frames, no audio.""" + video_filter = (f"fps=24,scale={width}:{height}:force_original_aspect_ratio=increase:flags=bilinear," + f"crop={width}:{height},tpad=stop_mode=clone:stop_duration={_TAIL_PAD_SECONDS}") + return [ + envs.FASTVIDEO_FFMPEG_BIN.get(), "-v", "error", "-i", path, "-vf", video_filter, "-frames:v", + str(DRAFT_FRAMES), "-an", "-f", "rawvideo", "-pix_fmt", "rgb24", "pipe:1" + ] + + +def load_draft_pixels(path: str, size: str = "768p") -> torch.Tensor: + """uint8 pixels [1, 3, 124, H, W] on the CPU for the draft video at ``path``.""" + if size not in DRAFT_SIZES: + raise ValueError(f"draft size must be one of {sorted(DRAFT_SIZES)}, got {size!r}.") + width, height = DRAFT_SIZES[size] + try: + raw = subprocess.run(ffmpeg_decode_command(path, width, height), + capture_output=True, + check=True, + timeout=_DECODE_TIMEOUT_S).stdout + except FileNotFoundError as error: + raise RuntimeError(f"ffmpeg ({envs.FASTVIDEO_FFMPEG_BIN.get()!r}) is required to read the draft video; " + "set FASTVIDEO_FFMPEG_BIN.") from error + except subprocess.CalledProcessError as error: + raise ValueError( + f"Could not decode the draft video {path!r}: {error.stderr.decode(errors='replace')}") from error + expected = DRAFT_FRAMES * height * width * 3 + if len(raw) != expected: + raise ValueError(f"The draft video {path!r} decoded to {len(raw)} bytes, expected {expected} " + f"({DRAFT_FRAMES} frames at {width}x{height}).") + frames = torch.from_numpy(np.frombuffer(raw, np.uint8).reshape(DRAFT_FRAMES, height, width, 3).copy()) + pixels = frames.permute(3, 0, 1, 2)[None] + return torch.cat([pixels, pixels[:, :, -1:].expand(-1, -1, VAE_PADDING_FRAMES, -1, -1)], dim=2).contiguous() diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/open_h3_regenerate_2k_pipeline.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/open_h3_regenerate_2k_pipeline.py new file mode 100644 index 000000000..b5bf46360 --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/open_h3_regenerate_2k_pipeline.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Open-H3-Regenerate-2K: regenerate a 768P MiniMax-H3 draft video at 2560x1440. + +``draft video -> H3 VAE encode -> latent upscaler -> 3 fixed-sigma Euler steps (H3 DiT + merged LoRA + Sharp curve +AdaLN branches) -> H3 VAE decode``. The release repo holds the unmodified MiniMax-H3 components under ``base_model/`` +next to ``lora/`` and ``latent_upsampler/``; see ``docs/inference/open_h3_regenerate_2k.md``. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +import torch + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.upsamplers.minimax_h3_latent_upscaler import MiniMaxH3LatentUpscaler +from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3Pipeline +from fastvideo.pipelines.basic.minimax_h3.stages import MiniMaxH3VideoDecodingStage +from fastvideo.pipelines.basic.open_h3_regenerate_2k.sharp_curve import SharpCurveBranches +from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages import ( + Regenerate2KDenoisingStage, + Regenerate2KFinalizeStage, + Regenerate2KLatentPreparationStage, +) +from fastvideo.utils import maybe_download_model, verify_model_config_and_directory + +logger = init_logger(__name__) + +BASE_MODEL_DIR = "base_model" +LORA_FILE = "lora/open_h3_regenerate_2k_lora.safetensors" +UPSCALER_FILE = "latent_upsampler/minimax_h3_latent_upscaler_3d_conv_v1_bf16.safetensors" +_BASE_COMPONENTS = ("text_encoder", "tokenizer", "processor", "vae", "audio_vae", "transformer", "scheduler", + "audio_scheduler") + + +class OpenH3Regenerate2KPipeline(MiniMaxH3Pipeline): + """MiniMax-H3 T2VA layout with a draft-latent start, fixed sigmas and a silent output.""" + + _extra_config_module_map = {name: f"{BASE_MODEL_DIR}/{name}" for name in _BASE_COMPONENTS} + # The merged LoRA reaches only these linears; the AdaLN projections carry the Sharp curve branches instead. + lora_target_modules = ["attn.to_q", "attn.to_k", "attn.to_v", "attn.to_out", "ff.fc_in", "ff.fc_out"] + + @classmethod + def get_hf_download_allow_patterns(cls) -> list[str]: + components = [f"{BASE_MODEL_DIR}/{name}/**" for name in _BASE_COMPONENTS] + return [f"{BASE_MODEL_DIR}/model_index.json", *components, "lora/**", "latent_upsampler/*.safetensors"] + + def _load_config(self, model_path: str) -> dict[str, Any]: + revision = getattr(self.fastvideo_args, "revision", None) + self.model_path = maybe_download_model(model_path, + revision=revision, + allow_patterns=self.get_hf_download_allow_patterns()) + if self.fastvideo_args.lora_path is None: + # The merged refine LoRA is part of the release; it is a constructor-time adapter like any H3 LoRA. + self.fastvideo_args.lora_path = os.path.join(self.model_path, LORA_FILE) + return verify_model_config_and_directory(os.path.join(self.model_path, BASE_MODEL_DIR), + required_component_dirs=()) + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: + super().initialize_pipeline(fastvideo_args) + self._upscaler: MiniMaxH3LatentUpscaler | None = None + self._curve: SharpCurveBranches | None = None + + def _load_upscaler(self) -> MiniMaxH3LatentUpscaler: + if self._upscaler is None: + model = MiniMaxH3LatentUpscaler.from_checkpoint(os.path.join(self.model_path, UPSCALER_FILE)) + self._upscaler = model.to(device=self.device, dtype=torch.bfloat16) + return self._upscaler + + def _load_curve(self) -> SharpCurveBranches: + if self._curve is None: + self._curve = SharpCurveBranches(Path(self.model_path)) + return self._curve + + def _add_denoise_stages(self, *, ref2va: bool) -> None: + config = self.fastvideo_args.pipeline_config + vae = self.get_module("vae") + audio_vae = self.get_module("audio_vae") + transformer = self.get_module("transformer") + if vae is None or audio_vae is None or transformer is None: + raise RuntimeError("Open-H3-Regenerate-2K needs the video VAE, audio VAE and transformer loaded.") + self.add_stage( + "latent_preparation_stage", + Regenerate2KLatentPreparationStage(vae, + audio_vae, + self.get_module("scheduler"), + load_upscaler=self._load_upscaler, + draft_size=config.draft_size, + start_sigma=config.refine_start_sigma)) + self.add_stage("denoising_stage", + Regenerate2KDenoisingStage(transformer, self._load_curve(), config.refine_start_sigma)) + self.add_stage("video_decoding_stage", MiniMaxH3VideoDecodingStage(vae=vae)) + self.add_stage("audio_decoding_stage", Regenerate2KFinalizeStage()) + self._denoise_stages_ready = True + + +EntryClass = OpenH3Regenerate2KPipeline + +__all__ = ["EntryClass", "OpenH3Regenerate2KPipeline"] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/presets.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/presets.py new file mode 100644 index 000000000..1eb710e6a --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/presets.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Inference preset for Open-H3-Regenerate-2K.""" + +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +_REFINE_STAGE = PresetStageSpec( + name="refine", + kind="denoising", + description="Latent upscale, then three fixed-sigma Euler refine steps", + allowed_overrides=frozenset({"num_inference_steps"}), +) + +OPEN_H3_REGENERATE_2K = InferencePreset( + name="open_h3_regenerate_2k", + version=1, + model_family="open_h3_regenerate_2k", + description="Open-H3-Regenerate-2K: regenerate a 768P MiniMax-H3 draft video at 2560x1440 (silent output)", + workload_type=None, + stage_schemas=(_REFINE_STAGE, ), + defaults={ + "height": 1440, + "width": 2560, + "num_frames": 124, + "fps": 24, + "guidance_scale": 1.0, + "batch_cfg": False, + "negative_prompt": "", + "num_inference_steps": 3, + "seed": 1766580932, + }, +) + +ALL_PRESETS = (OPEN_H3_REGENERATE_2K, ) + +__all__ = ["ALL_PRESETS", "OPEN_H3_REGENERATE_2K"] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/sampling.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/sampling.py new file mode 100644 index 000000000..017415fb5 --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/sampling.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fixed Open-H3-Regenerate-2K refine schedule: Euler, CFG 1, model time ``1 - sigma``.""" + +from __future__ import annotations + +import torch + +REFINE_SHIFT = 6.0 +DEFAULT_REFINE_STEPS = 3 +# Start noise of the two-step variant; the three-step grid starts at 0.909375. +TWO_STEP_START_SIGMA = 0.7 + + +def shifted_grid(start: float, steps: int, shift: float = REFINE_SHIFT) -> list[float]: + """Euler grid from ``start`` to 0: even spacing in unshifted sigma, then ``s * x / (1 + (s - 1) * x)``.""" + if not 0.0 < start < 1.0: + raise ValueError(f"start sigma must be in (0, 1), got {start}.") + if steps < 1: + raise ValueError(f"steps must be positive, got {steps}.") + raw = start / (shift - (shift - 1) * start) + mapped = [shift * x / (1 + (shift - 1) * x) for x in (raw * (1 - i / steps) for i in range(1, steps))] + return [start, *mapped, 0.0] + + +REFINE_SCHEDULES: dict[int, tuple[float, ...]] = { + 3: (0.909375, 6 / 7, 2 / 3, 0.0), + 2: tuple(shifted_grid(TWO_STEP_START_SIGMA, 2)), +} + + +def refine_sigmas(steps: int = DEFAULT_REFINE_STEPS, start_sigma: float | None = None) -> list[float]: + """Sigma grid of ``steps`` Euler steps, ending at 0.""" + if start_sigma is not None: + return shifted_grid(start_sigma, steps) + if steps not in REFINE_SCHEDULES: + raise ValueError(f"Open-H3-Regenerate-2K supports {sorted(REFINE_SCHEDULES)} refine steps, got {steps}.") + return list(REFINE_SCHEDULES[steps]) + + +def mix_initial_state(upscaled: torch.Tensor, noise: torch.Tensor, sigma: float) -> torch.Tensor: + """Re-noise the upscaled draft latent to the first sigma: ``(1 - sigma) * upscaled + sigma * noise``.""" + return (1.0 - sigma) * upscaled + sigma * noise diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/sharp_curve.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/sharp_curve.py new file mode 100644 index 000000000..dc0645702 --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/sharp_curve.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LMS Sharp curve AdaLN branches. + +Fifty AdaLN projections get an extra low-rank term whose input is an 8-D curve interpolated from the model timestep +(``adaln_t_table``, 1025 x 8) instead of the layer input, so it cannot be folded into the LoRA. The term is added to the +projection output through a forward hook; ``set_timesteps`` must be called with the per-forward unique timesteps first. +""" + +from __future__ import annotations + +from pathlib import Path + +import torch +import torch.nn.functional as F +from safetensors.torch import load_file + +CURVE_BRANCHES_FILE = "lora/sharp_curve_branches.safetensors" +CURVE_TABLE_FILE = "lora/adaln_t_table.bin" +CURVE_TABLE_SHAPE = (1025, 8) +NUM_CURVE_BRANCHES = 50 +_KEY_PREFIX = "transformer." +_A_SUFFIX = ".curve_A.weight" +_B_SUFFIX = ".curve_B.weight" + + +def interpolate_curve(table: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor: + """Linear interpolation of ``table`` rows at model times in [0, 1]; ``timesteps`` [U] -> [U, 8].""" + t = timesteps.float() + if t.numel() and (t.min() < 0 or t.max() > 1): + raise ValueError("Sharp curve timesteps must lie in [0, 1].") + table = table.to(t.device) + position = t * (len(table) - 1) + low = position.floor().long().clamp(max=len(table) - 2) + return torch.lerp(table[low], table[low + 1], (position - low).unsqueeze(-1)) + + +class SharpCurveBranches: + + def __init__(self, weights_dir: str | Path) -> None: + root = Path(weights_dir) + raw = torch.frombuffer(bytearray((root / CURVE_TABLE_FILE).read_bytes()), dtype=torch.float32) + if raw.numel() != CURVE_TABLE_SHAPE[0] * CURVE_TABLE_SHAPE[1]: + raise ValueError(f"{root / CURVE_TABLE_FILE} must hold a {CURVE_TABLE_SHAPE} float32 table.") + self.table = raw.reshape(CURVE_TABLE_SHAPE).clone() + tensors = load_file(str(root / CURVE_BRANCHES_FILE), device="cpu") + self.factors: dict[str, tuple[torch.Tensor, torch.Tensor]] = { + key[len(_KEY_PREFIX):-len(_A_SUFFIX)]: (value.float(), tensors[key[:-len(_A_SUFFIX)] + _B_SUFFIX].float()) + for key, value in tensors.items() if key.endswith(_A_SUFFIX) + } + if len(self.factors) != NUM_CURVE_BRANCHES: + raise ValueError(f"Expected {NUM_CURVE_BRANCHES} Sharp curve branches, found {len(self.factors)}.") + self.coords: torch.Tensor | None = None + self._attached: set[int] = set() + self._device_factors: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + + def set_timesteps(self, unique_timesteps: torch.Tensor) -> None: + self.coords = interpolate_curve(self.table, unique_timesteps) + + def attach(self, transformer: torch.nn.Module) -> None: + """Register the branch hooks on ``transformer`` once.""" + if id(transformer) in self._attached: + return + modules = dict(transformer.named_modules()) + missing = sorted(set(self.factors) - set(modules)) + if missing: + raise ValueError(f"Sharp curve branches match no transformer module: {missing[:5]}") + for name in self.factors: + modules[name].register_forward_hook(self._hook(name)) + self._attached.add(id(transformer)) + + def _hook(self, name: str): + + def add_branch(module, inputs, output): + projected = output[0] if isinstance(output, tuple) else output + coords = self.coords + if coords is None or coords.shape[0] != projected.shape[0]: + raise RuntimeError("Sharp curve timesteps were not set for this forward.") + factors = self._device_factors.get(name) + if factors is None or factors[0].device != projected.device: + factors = tuple(f.to(projected.device) for f in self.factors[name]) + self._device_factors[name] = factors + delta = F.linear(F.linear(coords, factors[0]), factors[1]).to(projected.dtype) + projected = projected + delta + return (projected, *output[1:]) if isinstance(output, tuple) else projected + + return add_branch diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/__init__.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/__init__.py new file mode 100644 index 000000000..3c7a022bf --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages.denoising import Regenerate2KDenoisingStage +from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages.finalize import Regenerate2KFinalizeStage +from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages.latent_preparation import ( + Regenerate2KLatentPreparationStage, ) + +__all__ = [ + "Regenerate2KDenoisingStage", + "Regenerate2KFinalizeStage", + "Regenerate2KLatentPreparationStage", +] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py new file mode 100644 index 000000000..6dba2d58f --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fixed-sigma Euler refine loop of Open-H3-Regenerate-2K.""" + +from __future__ import annotations + +import contextlib +from typing import Any + +import torch +from torch.nn.attention import SDPBackend, sdpa_kernel + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.pipelines.basic.minimax_h3.packing import MiniMaxH3PackedLayout, build_row_timesteps +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY +from fastvideo.pipelines.basic.open_h3_regenerate_2k.sampling import refine_sigmas +from fastvideo.pipelines.basic.open_h3_regenerate_2k.sharp_curve import SharpCurveBranches +from fastvideo.pipelines.lazy_module import is_lazy_module +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + +logger = init_logger(__name__) + + +class Regenerate2KDenoisingStage(PipelineStage): + """Euler steps over the refine sigmas with the video clock at ``1 - sigma`` and the silent audio held clean.""" + + performance_component_metric = "dit_time_s" + + def __init__(self, transformer: Any, curve: SharpCurveBranches, start_sigma: float | None) -> None: + super().__init__() + self.transformer = transformer + self.curve = curve + self.start_sigma = start_sigma + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("layout", batch.extra.get(MINIMAX_H3_LAYOUT_KEY), V.not_none) + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_of_tensors_dims(3)) + result.add_check("latents", batch.latents, V.with_dims(2)) + result.add_check("audio_latents", batch.audio_latents, V.with_dims(2)) + return result + + def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, V.with_dims(2)) + return result + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + transformer = self.transformer + if is_lazy_module(transformer): + transformer.materialize() + self.curve.attach(transformer) + layout = batch.extra.get(MINIMAX_H3_LAYOUT_KEY) + if not isinstance(layout, MiniMaxH3PackedLayout): + raise ValueError("The MiniMax-H3 packed layout is missing before denoising.") + device = get_local_torch_device() + sigmas = refine_sigmas(batch.num_inference_steps, self.start_sigma) + state = batch.latents.to(device) + with self._attention_kernel(device): + state = self._euler_loop(batch, layout, state, sigmas, device, transformer) + batch.latents = state + return batch + + @staticmethod + def _attention_kernel(device: torch.device) -> contextlib.AbstractContextManager[None]: + """cuDNN attention: at 133k tokens PyTorch's default SDPA dispatch picks its flash kernel, 3x slower on B200.""" + return sdpa_kernel(SDPBackend.CUDNN_ATTENTION) if device.type == "cuda" else contextlib.nullcontext() + + def _euler_loop(self, batch: ForwardBatch, layout: MiniMaxH3PackedLayout, state: torch.Tensor, sigmas: list[float], + device: torch.device, transformer: Any) -> torch.Tensor: + position_ids = layout.position_ids.to(device) + token_tags = layout.token_tags.to(device) + video_indices = layout.video_indices.to(device) + audio_indices = layout.audio_indices.to(device) + text_indices = layout.text_indices.to(device) + prompt_embeds = batch.prompt_embeds[0].to(device) + audio = batch.audio_latents.to(device=device, dtype=torch.bfloat16)[None] + for index, (sigma, following) in enumerate(zip(sigmas[:-1], sigmas[1:], strict=True)): + unique, inverse = build_row_timesteps(layout, + video_timestep=1.0 - sigma, + audio_timestep=1.0, + condition_video_timestep=1.0, + condition_audio_timestep=1.0) + unique, inverse = unique.to(device), inverse.to(device) + self.curve.set_timesteps(unique) + torch.compiler.cudagraph_mark_step_begin() + with set_forward_context(current_timestep=index, attn_metadata=None, forward_batch=batch): + velocity, _ = transformer(hidden_states=state.to(torch.bfloat16)[None], + audio_hidden_states=audio, + encoder_hidden_states=prompt_embeds, + timestep=unique, + timestep_indices=inverse, + token_tags=token_tags, + position_ids=position_ids, + video_indices=video_indices, + audio_indices=audio_indices, + text_indices=text_indices) + state = state + (sigma - following) * velocity[0].float() + batch.step_index = index + return state + + +__all__ = ["Regenerate2KDenoisingStage"] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/finalize.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/finalize.py new file mode 100644 index 000000000..300caa4f8 --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/finalize.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Trim the decoded video to the draft length and drop the per-request H3 state.""" + +from __future__ import annotations + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.basic.open_h3_regenerate_2k.draft_video import DRAFT_FRAMES +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + + +class Regenerate2KFinalizeStage(PipelineStage): + """The released pipeline writes a silent video: the audio stream only conditions the DiT and is never decoded.""" + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("output", batch.output, V.with_dims(5)) + return result + + def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("output", batch.output, V.with_dims(5)) + return result + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + # The VAE decodes 124 frames (37 latent frames); the draft has 121. + batch.output = batch.output[:, :, :DRAFT_FRAMES] + batch.prompt_embeds = [] + batch.latents = None + batch.audio_latents = None + for key in tuple(batch.extra): + if key.startswith("minimax_h3_"): + batch.extra.pop(key) + return batch + + +__all__ = ["Regenerate2KFinalizeStage"] diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/latent_preparation.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/latent_preparation.py new file mode 100644 index 000000000..93f28431f --- /dev/null +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/latent_preparation.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Draft video -> upscaled, re-noised H3 latent plus the fixed silent-audio rows.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import torch + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models import pinned_offload +from fastvideo.models.upsamplers.minimax_h3_latent_upscaler import MiniMaxH3LatentUpscaler, upscale_latent +from fastvideo.pipelines.basic.minimax_h3.packing import ( + MINIMAX_H3_AUDIO_CHANNELS, + MINIMAX_H3_FPS, + MINIMAX_H3_KEYFRAME_ENCODE_SEED, + h3_dit_patch_size, + patchify_video_latents, +) +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import ( + MINIMAX_H3_LAYOUT_KEY, + MiniMaxH3LatentPreparationStage, +) +from fastvideo.pipelines.basic.open_h3_regenerate_2k.draft_video import load_draft_pixels +from fastvideo.pipelines.basic.open_h3_regenerate_2k.sampling import mix_initial_state, refine_sigmas +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.validators import VerificationResult + +logger = init_logger(__name__) + +DRAFT_FRAME_COUNT = 124 + + +class Regenerate2KLatentPreparationStage(MiniMaxH3LatentPreparationStage): + """Encode the draft with the H3 VAE, upscale its latent to the target grid and re-noise it to the first sigma.""" + + def __init__(self, vae: Any, audio_vae: Any, scheduler: Any, *, load_upscaler: Callable[[], + MiniMaxH3LatentUpscaler], + draft_size: str, start_sigma: float | None) -> None: + super().__init__(vae, audio_vae, scheduler, ref2va=False) + self.load_upscaler = load_upscaler + self.draft_size = draft_size + self.start_sigma = start_sigma + self._silence: tuple[tuple[str, int], torch.Tensor] | None = None + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = super().verify_input(batch, fastvideo_args) + result.add_check("video_path", batch.video_path, lambda value: isinstance(value, str) and bool(value)) + return result + + def _encode_draft(self, video_path: str, device: torch.device, fastvideo_args: FastVideoArgs) -> torch.Tensor: + """Normalized native latent [C, T, h, w] of the draft, with the VAE posterior sampled under a fixed seed.""" + pixels = load_draft_pixels(video_path, self.draft_size) + pinned_offload.load(self.vae, device, pin=fastvideo_args.pin_cpu_memory) + try: + posterior = self.vae.encode_pixels(pixels).latent_dist + sample = posterior.sample(generator=torch.Generator("cpu").manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED)) + return self.vae.normalize_latents(sample.float())[0].to(device) + finally: + if fastvideo_args.vae_cpu_offload: + pinned_offload.unload(self.vae) + + def _silent_audio_rows(self, num_frames: int, device: torch.device, fastvideo_args: FastVideoArgs) -> torch.Tensor: + """The encoded-silence audio stream, cached: it is the same for every request of this length.""" + key = (str(device), num_frames) + if self._silence is not None and self._silence[0] == key: + return self._silence[1] + pinned_offload.load(self.audio_vae, device, pin=fastvideo_args.pin_cpu_memory) + try: + samples = round(num_frames / MINIMAX_H3_FPS * int(self.audio_vae.sampling_rate)) + silence = torch.zeros(MINIMAX_H3_AUDIO_CHANNELS, 1, samples, device=device) + posterior = self.audio_vae.encode(silence).latent_dist + latents = self.audio_vae.normalize_latents(posterior.mode().float()).transpose(1, 2) + rows = latents.reshape(-1, self.audio_vae.latent_channels).to(torch.bfloat16) + finally: + if fastvideo_args.vae_cpu_offload: + pinned_offload.unload(self.audio_vae) + self._silence = (key, rows) + return rows + + @staticmethod + def _noise_rows(shape: tuple[int, int], device: torch.device, seed: int) -> torch.Tensor: + generator = torch.Generator(device=device).manual_seed(seed) + return torch.randn((1, *shape), device=device, dtype=torch.float32, generator=generator)[0] + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if batch.num_frames != DRAFT_FRAME_COUNT: + raise ValueError(f"Open-H3-Regenerate-2K refines {DRAFT_FRAME_COUNT} frames, got {batch.num_frames}.") + device = get_local_torch_device() + patch_size = h3_dit_patch_size(fastvideo_args) + _, channels, num_latent_frames, latent_height, latent_width = batch.raw_latent_shape + + draft = self._encode_draft(batch.video_path, device, fastvideo_args) + upscaler = self.load_upscaler() + upscaled = upscale_latent(upscaler, draft, (batch.height, batch.width)) + if tuple(upscaled.shape) != (channels, num_latent_frames, latent_height, latent_width): + raise ValueError(f"The upscaled latent has shape {tuple(upscaled.shape)}, expected " + f"{(channels, num_latent_frames, latent_height, latent_width)}.") + upscaled_rows = patchify_video_latents(upscaled[None], patch_size) + + if batch.latents is None: + noise_rows = self._noise_rows(tuple(upscaled_rows.shape), device, 0 if batch.seed is None else batch.seed) + else: + expected = (1, channels, num_latent_frames, latent_height, latent_width) + if tuple(batch.latents.shape) != expected: + raise ValueError(f"Injected noise must have shape {expected}, got {tuple(batch.latents.shape)}.") + noise_rows = patchify_video_latents(batch.latents.to(device=device, dtype=torch.float32), patch_size) + + layout = self._build_layout(batch, fastvideo_args) + audio_rows = self._silent_audio_rows(batch.num_frames, device, fastvideo_args) + if audio_rows.shape[0] != layout.audio_indices.numel(): + raise ValueError(f"The silent audio has {audio_rows.shape[0]} rows, the packed layout expects " + f"{layout.audio_indices.numel()} ({layout.num_audio_latents} latents).") + + batch.latents = mix_initial_state(upscaled_rows, noise_rows, + refine_sigmas(batch.num_inference_steps, self.start_sigma)[0]) + batch.audio_latents = audio_rows + batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout + for key in tuple(batch.extra): + if key.startswith("minimax_h3_") and key != MINIMAX_H3_LAYOUT_KEY: + batch.extra.pop(key) + return batch + + +__all__ = ["Regenerate2KLatentPreparationStage"] diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 414183892..c74d0d1f6 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -42,6 +42,7 @@ from fastvideo.configs.pipelines.flux_2 import ( from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig +from fastvideo.configs.pipelines.open_h3_regenerate_2k import OpenH3Regenerate2KPipelineConfig from fastvideo.configs.pipelines.mmaudio import ( MMAudioLarge44kV2PipelineConfig, MMAudioLarge44kV2AConfig, @@ -216,6 +217,13 @@ def _resolve_class_name(config: dict[str, Any]) -> str | None: return class_name +OPEN_H3_REGENERATE_2K_HF_PATH = "Efficient-Large-Model/Open-H3-Regenerate-2K" + + +def _is_open_h3_regenerate_2k_release(model_path: str) -> bool: + return os.path.isfile(os.path.join(model_path, "lora", "open_h3_regenerate_2k_lora.safetensors")) + + def _get_config_info( model_path: str, *, @@ -247,6 +255,9 @@ def _get_config_info( return _CONFIG_REGISTRY.get(model_id) # 3. Use detectors (path or model_index pipeline name). + if os.path.exists(model_path) and _is_open_h3_regenerate_2k_release(model_path): + # The release keeps its Diffusers manifest under base_model/, so a renamed local copy has no root manifest. + return _CONFIG_REGISTRY.get(_MODEL_HF_PATH_TO_NAME[OPEN_H3_REGENERATE_2K_HF_PATH]) if os.path.exists(model_path): config = verify_model_config_and_directory(model_path, required_component_dirs=[]) elif allow_remote: @@ -1152,6 +1163,22 @@ def _register_configs() -> None: _register_wan_configs(WAN_MODEL_DEFINITIONS) + # Open-H3-Regenerate-2K (video -> video). Registered before MiniMax H3 so the first matching detector wins. There is + # no V2V WorkloadType, so, like the Kandinsky6 SR entries, it declares no workload. + register_configs( + sampling_param_cls=None, + pipeline_config_cls=OpenH3Regenerate2KPipelineConfig, + workload_types=(), + hf_model_paths=[OPEN_H3_REGENERATE_2K_HF_PATH], + model_detectors=[ + lambda path: any(token in path.lower() + for token in ("open-h3-regenerate", "open_h3_regenerate", "openh3regenerate2kpipeline")), + ], + model_family="open_h3_regenerate_2k", + default_preset="open_h3_regenerate_2k", + pipeline_cls_name="OpenH3Regenerate2KPipeline", + ) + # MiniMax H3 register_configs( sampling_param_cls=None, @@ -1263,19 +1290,18 @@ def get_model_info( # is triggered just to log the replaced name). pipeline_name: str | None = override_pipeline_cls_name logger.info("Using override pipeline class name %s", pipeline_name) + elif config_info.pipeline_cls_name is not None: + # The resolved (path/detector-based) config pins the pipeline class, e.g. an I2V checkpoint whose + # `_class_name` would otherwise resolve to the T2V pipeline. The manifest is not read: a release such as + # Open-H3-Regenerate-2K keeps it in a subfolder. + pipeline_name = config_info.pipeline_cls_name + logger.info("Pinned pipeline class name %s", pipeline_name) else: if os.path.exists(model_path): config = verify_model_config_and_directory(model_path, required_component_dirs=[]) else: config = maybe_download_model_index(model_path, revision=revision) - pipeline_name = _resolve_class_name(config) - if config_info.pipeline_cls_name is not None: - # The resolved (path/detector-based) config pins the pipeline class, - # e.g. an I2V checkpoint whose `_class_name` would otherwise resolve - # to the T2V pipeline. - logger.info("Pinning pipeline class name from %s to %s", pipeline_name, config_info.pipeline_cls_name) - pipeline_name = config_info.pipeline_cls_name if pipeline_name is None: raise ValueError("Model config does not contain a _class_name attribute. " @@ -1350,6 +1376,8 @@ def _register_presets() -> None: ALL_PRESETS as MATRIXGAME3_PRESETS, ) from fastvideo.pipelines.basic.minimax_h3.presets import ( ALL_PRESETS as MINIMAX_H3_PRESETS, ) + from fastvideo.pipelines.basic.open_h3_regenerate_2k.presets import ( + ALL_PRESETS as OPEN_H3_REGENERATE_2K_PRESETS, ) from fastvideo.pipelines.basic.mmaudio.presets import ( ALL_PRESETS as MMAUDIO_PRESETS, ) from fastvideo.pipelines.basic.sd35.presets import ( @@ -1383,6 +1411,7 @@ def _register_presets() -> None: MATRIXGAME2_PRESETS, MATRIXGAME3_PRESETS, MINIMAX_H3_PRESETS, + OPEN_H3_REGENERATE_2K_PRESETS, MMAUDIO_PRESETS, SD35_PRESETS, STABLE_AUDIO_PRESETS, diff --git a/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py b/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py new file mode 100644 index 000000000..e2e3b250b --- /dev/null +++ b/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU unit tests for the Open-H3-Regenerate-2K building blocks.""" + +import numpy as np +import pytest +import torch +from safetensors.torch import save_file + +from fastvideo.models.upsamplers.minimax_h3_latent_upscaler import ( + MiniMaxH3LatentUpscaler, + resolve_upscale_geometry, + upscale_latent, +) +from fastvideo.pipelines.basic.open_h3_regenerate_2k import draft_video +from fastvideo.pipelines.basic.open_h3_regenerate_2k.sampling import ( + mix_initial_state, + refine_sigmas, + shifted_grid, +) +from fastvideo.pipelines.basic.open_h3_regenerate_2k.sharp_curve import ( + NUM_CURVE_BRANCHES, + SharpCurveBranches, + interpolate_curve, +) + + +def test_default_sigmas_match_the_released_grid(): + assert refine_sigmas(3) == pytest.approx([0.909375, 0.857143, 0.666667, 0.0], abs=1e-6) + + +def test_two_step_sigmas_start_at_0_7(): + assert refine_sigmas(2) == pytest.approx([0.7, 0.494118, 0.0], abs=1e-6) + + +def test_shifted_grid_from_the_released_start_is_the_original_four_step_grid(): + assert shifted_grid(0.909375, 4) == pytest.approx([0.909375, 0.841446, 0.732075, 0.526697, 0.0], abs=1e-6) + + +def test_custom_start_sigma_is_decreasing_and_ends_at_zero(): + grid = refine_sigmas(3, start_sigma=0.8) + assert grid[0] == 0.8 and grid[-1] == 0.0 + assert all(a > b for a, b in zip(grid, grid[1:])) + + +@pytest.mark.parametrize("steps", [0, 4]) +def test_unsupported_step_count_is_rejected(steps): + with pytest.raises(ValueError): + refine_sigmas(steps) + + +def test_initial_state_mix_weights_noise_by_sigma(): + upscaled, noise = torch.ones(2, 3), torch.zeros(2, 3) + assert torch.allclose(mix_initial_state(upscaled, noise, 0.909375), torch.full((2, 3), 0.090625)) + assert torch.equal(mix_initial_state(upscaled, noise, 0.0), upscaled) + + +def test_upscale_geometry_targets_2k_from_768p_and_576p(): + assert resolve_upscale_geometry((48, 84), (1440, 2560))[0] == (90, 160) + shape, scale = resolve_upscale_geometry((36, 64), (1440, 2560)) + assert shape == (90, 160) and scale == pytest.approx(2.5) + + +def test_tiny_upscaler_forward_shape_and_identity_passthrough(): + model = MiniMaxH3LatentUpscaler(channels=32, in_blocks=2, out_blocks=2).eval() + latent = torch.randn(24, 5, 4, 6) + out = upscale_latent(model, latent, (96, 160), dtype=torch.float32) + assert out.shape == (24, 5, 6, 10) and out.dtype == torch.float32 + assert torch.equal(upscale_latent(model, latent, (64, 96)), latent) + + +def test_upscaler_checkpoint_roundtrip(tmp_path): + model = MiniMaxH3LatentUpscaler(channels=32, in_blocks=3, out_blocks=2) + path = str(tmp_path / "upscaler.safetensors") + save_file(model.state_dict(), path) + loaded = MiniMaxH3LatentUpscaler.from_checkpoint(path) + assert len(loaded.in_blocks) == len(model.in_blocks) and len(loaded.out_blocks) == len(model.out_blocks) + for key, value in model.state_dict().items(): + assert torch.equal(loaded.state_dict()[key], value) + + +def test_curve_interpolation_hits_table_rows_and_midpoints(): + table = torch.arange(1025 * 8, dtype=torch.float32).reshape(1025, 8) + coords = interpolate_curve(table, torch.tensor([0.0, 1.0, 0.5, 1 / 2048])) + assert torch.equal(coords[0], table[0]) and torch.equal(coords[1], table[1024]) + assert torch.equal(coords[2], table[512]) + assert torch.allclose(coords[3], (table[0] + table[1]) / 2) + with pytest.raises(ValueError): + interpolate_curve(table, torch.tensor([1.5])) + + +class _Projection(torch.nn.Module): + + def __init__(self, out_features): + super().__init__() + self.linear = torch.nn.Linear(4, out_features) + + def forward(self, x): + return self.linear(x), None + + +def _write_curve_release(root, out_features, count): + (root / "lora").mkdir() + (root / "lora" / "adaln_t_table.bin").write_bytes(np.random.rand(1025 * 8).astype(np.float32).tobytes()) + tensors = {} + for index in range(count): + tensors[f"transformer.blocks.{index}.proj.curve_A.weight"] = torch.randn(3, 8) + tensors[f"transformer.blocks.{index}.proj.curve_B.weight"] = torch.randn(out_features, 3) + save_file(tensors, str(root / "lora" / "sharp_curve_branches.safetensors")) + + +def test_sharp_branches_require_all_fifty(tmp_path): + _write_curve_release(tmp_path, 6, 3) + with pytest.raises(ValueError, match="50"): + SharpCurveBranches(tmp_path) + + +def test_sharp_branch_adds_curve_term_to_projection_output(tmp_path): + _write_curve_release(tmp_path, 6, NUM_CURVE_BRANCHES) + curve = SharpCurveBranches(tmp_path) + model = torch.nn.Module() + model.blocks = torch.nn.ModuleList() + for _ in range(NUM_CURVE_BRANCHES): + block = torch.nn.Module() + block.proj = _Projection(6) + model.blocks.append(block) + curve.attach(model) + curve.attach(model) # idempotent: a second attach must not double the term + timesteps = torch.tensor([0.1, 0.9]) + x = torch.randn(2, 4) + base, _ = model.blocks[0].proj.linear(x), None + curve.set_timesteps(timesteps) + out, extra = model.blocks[0].proj(x) + a, b = curve.factors["blocks.0.proj"] + expected = base + torch.nn.functional.linear(torch.nn.functional.linear(curve.coords, a), b) + assert extra is None and torch.allclose(out, expected, atol=1e-6) + + +def test_sharp_branch_without_timesteps_fails_loudly(tmp_path): + _write_curve_release(tmp_path, 6, NUM_CURVE_BRANCHES) + curve = SharpCurveBranches(tmp_path) + model = torch.nn.Module() + model.blocks = torch.nn.ModuleList() + for _ in range(NUM_CURVE_BRANCHES): + block = torch.nn.Module() + block.proj = _Projection(6) + model.blocks.append(block) + curve.attach(model) + with pytest.raises(RuntimeError, match="timesteps"): + model.blocks[0].proj(torch.randn(2, 4)) + + +def test_ffmpeg_command_cover_scales_and_caps_at_121_frames(): + command = draft_video.ffmpeg_decode_command("clip.mp4", 1344, 768) + video_filter = command[command.index("-vf") + 1] + assert "fps=24" in video_filter and "scale=1344:768:force_original_aspect_ratio=increase" in video_filter + assert "crop=1344:768" in video_filter + assert command[command.index("-frames:v") + 1] == "121" + + +def test_draft_pixels_pad_the_last_frame_to_124(monkeypatch): + width, height = draft_video.DRAFT_SIZES["576p"] + frames = np.arange(121, dtype=np.uint8).reshape(121, 1, 1, 1) * np.ones((121, height, width, 3), np.uint8) + + class _Done: + stdout = frames.tobytes() + + monkeypatch.setattr(draft_video.subprocess, "run", lambda *a, **k: _Done()) + pixels = draft_video.load_draft_pixels("clip.mp4", "576p") + assert pixels.shape == (1, 3, 124, height, width) and pixels.dtype == torch.uint8 + assert torch.equal(pixels[:, :, 120], pixels[:, :, 123]) + + +def test_draft_pixels_reject_truncated_decode(monkeypatch): + + class _Short: + stdout = b"\x00" * 10 + + monkeypatch.setattr(draft_video.subprocess, "run", lambda *a, **k: _Short()) + with pytest.raises(ValueError, match="decoded to 10 bytes"): + draft_video.load_draft_pixels("clip.mp4") + + +class _ConstantVelocityTransformer: + """Predicts a constant video velocity and records what the stage feeds it.""" + + def __init__(self, velocity): + self.velocity = velocity + self.calls = [] + + def __call__(self, **kwargs): + self.calls.append(kwargs) + rows = kwargs["hidden_states"].shape[1] + audio_rows = kwargs["audio_hidden_states"].shape[1] + return torch.full((1, rows, 8), self.velocity), torch.zeros(1, audio_rows, 4) + + +class _RecordingCurve: + + def __init__(self): + self.timesteps = [] + self.attached = 0 + + def attach(self, transformer): + self.attached += 1 + + def set_timesteps(self, unique): + self.timesteps.append(unique.clone()) + + +def _denoise(monkeypatch, steps, velocity=0.5): + from fastvideo.pipelines.basic.minimax_h3.packing import build_packed_sequence + from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY + from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages import denoising + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + monkeypatch.setattr(denoising, "get_local_torch_device", lambda: torch.device("cpu")) + layout = build_packed_sequence(torch.ones(3, dtype=torch.long), 2, 2, 2, 2, (1, 2, 2)) + batch = ForwardBatch(data_type="video") + batch.num_inference_steps = steps + batch.prompt_embeds = [torch.zeros(1, 3, 4)] + batch.latents = torch.zeros(layout.video_indices.numel(), 8) + batch.audio_latents = torch.ones(layout.audio_indices.numel(), 4) + batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout + transformer, curve = _ConstantVelocityTransformer(velocity), _RecordingCurve() + stage = denoising.Regenerate2KDenoisingStage(transformer, curve, None) + return stage.forward(batch, None), transformer, curve + + +def test_euler_steps_integrate_the_constant_velocity_over_sigma_zero_to_start(monkeypatch): + batch, transformer, curve = _denoise(monkeypatch, steps=3, velocity=0.5) + assert len(transformer.calls) == 3 and curve.attached == 1 + assert torch.allclose(batch.latents, torch.full_like(batch.latents, 0.909375 * 0.5)) + + +def test_video_clock_is_one_minus_sigma_and_audio_is_held_clean(monkeypatch): + _, transformer, curve = _denoise(monkeypatch, steps=3) + for call, unique, sigma in zip(transformer.calls, curve.timesteps, [0.909375, 6 / 7, 2 / 3]): + assert torch.allclose(unique.sort().values, torch.tensor([1.0 - sigma, 1.0]).sort().values.float(), atol=1e-6) + assert call["hidden_states"].dtype == torch.bfloat16 and call["audio_hidden_states"].dtype == torch.bfloat16 + assert torch.equal(call["audio_hidden_states"].float(), torch.ones_like(call["audio_hidden_states"]).float()) + + +def test_two_step_request_runs_two_forwards(monkeypatch): + batch, transformer, _ = _denoise(monkeypatch, steps=2, velocity=1.0) + assert len(transformer.calls) == 2 + assert torch.allclose(batch.latents, torch.full_like(batch.latents, 0.7), atol=1e-6) + + +def test_registry_resolves_the_release_to_the_regenerate_pipeline(tmp_path): + from fastvideo.configs.pipelines.open_h3_regenerate_2k import OpenH3Regenerate2KPipelineConfig + from fastvideo.registry import get_model_info, get_pipeline_config_cls_from_name, get_preset_selection + + hub_id = "Efficient-Large-Model/Open-H3-Regenerate-2K" + (tmp_path / "lora").mkdir() + (tmp_path / "lora" / "open_h3_regenerate_2k_lora.safetensors").write_bytes(b"") + for path in (hub_id, str(tmp_path)): + assert get_pipeline_config_cls_from_name(path) is OpenH3Regenerate2KPipelineConfig + assert get_model_info(path).pipeline_cls.__name__ == "OpenH3Regenerate2KPipeline" + preset, family = get_preset_selection(hub_id) + assert (preset, family) == ("open_h3_regenerate_2k", "open_h3_regenerate_2k") diff --git a/tests/local_tests/open_h3_regenerate_2k/test_upscaler_parity.py b/tests/local_tests/open_h3_regenerate_2k/test_upscaler_parity.py new file mode 100644 index 000000000..b4bf1b155 --- /dev/null +++ b/tests/local_tests/open_h3_regenerate_2k/test_upscaler_parity.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Opt-in parity of the ported latent upscaler against the author's ComfyUI node. + +Needs the release weights (OPEN_H3_REGEN_WEIGHTS) and a checkout of NVlabs/Sana on the ``sol-engine`` branch +(OPEN_H3_REGEN_REF, the models/Open-H3-Regenerate-2K directory). GPU only: the node runs the network in bf16. +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +import types +from pathlib import Path + +import pytest +import torch + +from fastvideo.models.upsamplers.minimax_h3_latent_upscaler import MiniMaxH3LatentUpscaler, upscale_latent + +WEIGHTS = os.environ.get("OPEN_H3_REGEN_WEIGHTS") +REFERENCE = os.environ.get("OPEN_H3_REGEN_REF") +CHECKPOINT = "latent_upsampler/minimax_h3_latent_upscaler_3d_conv_v1_bf16.safetensors" + +pytestmark = pytest.mark.skipif(not (WEIGHTS and REFERENCE and torch.cuda.is_available()), + reason="needs OPEN_H3_REGEN_WEIGHTS, OPEN_H3_REGEN_REF and a GPU") + + +def _load_author_node(checkpoint: str): + folder = types.ModuleType("folder_paths") + folder.folder_names_and_paths = {"latent_upscale_models": ([], set())} + folder.get_full_path_or_raise = lambda category, name: checkpoint + sys.modules["folder_paths"] = folder + sys.modules["comfy.model_management"] = None + sys.modules["comfy_api.latest"] = None + spec = importlib.util.spec_from_file_location("author_upscaler", + Path(REFERENCE) / "open_h3" / "upscaler_node.py") + node = importlib.util.module_from_spec(spec) + spec.loader.exec_module(node) + node.io.NodeOutput = lambda value: (value, ) + return node + + +@pytest.mark.parametrize("shape", [(24, 37, 48, 84), (24, 37, 36, 64)]) +def test_upscaler_matches_the_author_node(shape): + checkpoint = str(Path(WEIGHTS) / CHECKPOINT) + node = _load_author_node(checkpoint) + torch.manual_seed(0) + latent = torch.randn(shape) + mode = {"mode": node.UpscaleMode.TARGET_DIMENSIONS, "width": 2560, "height": 1440} + with torch.inference_mode(): + (expected, ) = node.MinimaxH3LatentUpscaler3D.execute({"samples": latent[None].float()}, + "upscaler", + mode, + align=32, + enable_temporal_chunking=False, + force_unload=False, + device="cuda", + precision="bf16") + model = MiniMaxH3LatentUpscaler.from_checkpoint(checkpoint).to("cuda", torch.bfloat16) + actual = upscale_latent(model, latent.cuda(), (1440, 2560)).cpu() + assert torch.equal(actual, expected["samples"].float()[0])