[new-model]: Open-H3-Regenerate-2K pipeline (latent upscaler, Sharp curve branches, refine sampler)
This commit is contained in:
Aryan Kumar
@@ -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}.")
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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]
|
||||
@@ -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"]
|
||||
@@ -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()
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
@@ -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"]
|
||||
+36
-7
@@ -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,
|
||||
|
||||
@@ -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")
|
||||
@@ -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])
|
||||
Reference in New Issue
Block a user