diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 1f1f103529..9cdf9f5e7c 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -518,6 +518,14 @@ surfaces: true_cfg_scale: request.sampling.true_cfg_scale boundary_ratio: request.sampling.boundary_ratio sigmas: request.sampling.sigmas + pyramid_num_inference_steps_list: request.sampling.pyramid_num_inference_steps_list + history_sizes: request.sampling.history_sizes + num_latent_frames_per_chunk: request.sampling.num_latent_frames_per_chunk + keep_first_frame: request.sampling.keep_first_frame + is_skip_first_chunk: request.sampling.is_skip_first_chunk + use_zero_init: request.sampling.use_zero_init + zero_steps: request.sampling.zero_steps + is_amplify_first_chunk: request.sampling.is_amplify_first_chunk enable_teacache: request.runtime.enable_teacache save_video: request.output.save_video return_frames: request.output.return_frames diff --git a/examples/inference/basic/basic_helios_distilled_t2v.py b/examples/inference/basic/basic_helios_distilled_t2v.py new file mode 100644 index 0000000000..d37b2eec51 --- /dev/null +++ b/examples/inference/basic/basic_helios_distilled_t2v.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generate one Helios-Distilled T2V chunk through FastVideo's typed API. + +Set ``HELIOS_MODEL_PATH`` to a local snapshot to avoid downloading the public +checkpoint again. The 33-frame example is a short integration run; increase +``num_frames`` to 240 for the official eight-chunk default. +""" + +from __future__ import annotations + +import os + +from fastvideo import VideoGenerator +from fastvideo.api import ( + EngineConfig, + GenerationRequest, + GeneratorConfig, + OffloadConfig, + OutputConfig, + SamplingConfig, +) + +MODEL_PATH = os.getenv("HELIOS_MODEL_PATH", "BestWishYsh/Helios-Distilled") +OUTPUT_PATH = os.getenv( + "HELIOS_OUTPUT_PATH", + "outputs_video/helios/helios_distilled_t2v.mp4", +) +PROMPT = ("A vibrant tropical fish swims gracefully through a colorful coral reef " + "in clear turquoise water, cinematic close-up, fluid motion, vivid detail.") +NEGATIVE_PROMPT = ("Bright tones, overexposed, static, blurred details, subtitles, paintings, " + "images, overall gray, worst quality, low quality, JPEG artifacts, ugly, " + "deformed, disfigured, still picture, messy background.") + + +def main() -> None: + generator = VideoGenerator.from_config( + GeneratorConfig( + model_path=MODEL_PATH, + engine=EngineConfig( + num_gpus=1, + use_fsdp_inference=False, + offload=OffloadConfig( + dit=False, + dit_layerwise=True, + text_encoder=True, + vae=True, + pin_cpu_memory=False, + ), + ), + )) + request = GenerationRequest( + prompt=PROMPT, + negative_prompt=NEGATIVE_PROMPT, + sampling=SamplingConfig( + seed=42, + height=384, + width=640, + num_frames=33, + fps=24, + num_inference_steps=2, + guidance_scale=1.0, + pyramid_num_inference_steps_list=[2, 2, 2], + history_sizes=[16, 2, 1], + num_latent_frames_per_chunk=9, + keep_first_frame=True, + is_skip_first_chunk=False, + use_zero_init=True, + zero_steps=1, + is_amplify_first_chunk=True, + ), + output=OutputConfig( + output_path=OUTPUT_PATH, + save_video=True, + return_frames=False, + ), + ) + + try: + generator.generate(request=request) + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/fastvideo/api/sampling_param.py b/fastvideo/api/sampling_param.py index 464d7f0fa3..52434946a7 100644 --- a/fastvideo/api/sampling_param.py +++ b/fastvideo/api/sampling_param.py @@ -108,6 +108,16 @@ class SamplingParam: boundary_ratio: float | None = None sigmas: list[float] | None = None + # Helios autoregressive spatial-pyramid sampling. + pyramid_num_inference_steps_list: list[int] | None = None + history_sizes: list[int] | None = None + num_latent_frames_per_chunk: int = 9 + keep_first_frame: bool = True + is_skip_first_chunk: bool = False + use_zero_init: bool = True + zero_steps: int = 1 + is_amplify_first_chunk: bool = False + # TeaCache parameters enable_teacache: bool = False @@ -375,6 +385,56 @@ def add_cli_args(parser: Any) -> Any: default=SamplingParam.boundary_ratio, help="Boundary timestep ratio", ) + parser.add_argument( + "--pyramid-num-inference-steps-list", + nargs=3, + type=int, + default=SamplingParam.pyramid_num_inference_steps_list, + help="Denoising steps for the three Helios pyramid stages", + ) + parser.add_argument( + "--history-sizes", + nargs=3, + type=int, + default=SamplingParam.history_sizes, + help="Long, mid, and short Helios latent history sizes", + ) + parser.add_argument( + "--num-latent-frames-per-chunk", + type=int, + default=SamplingParam.num_latent_frames_per_chunk, + help="Helios autoregressive latent frames per chunk", + ) + parser.add_argument( + "--keep-first-frame", + action=StoreBoolean, + default=SamplingParam.keep_first_frame, + help="Keep the first Helios latent frame as prefix conditioning", + ) + parser.add_argument( + "--is-skip-first-chunk", + action=StoreBoolean, + default=SamplingParam.is_skip_first_chunk, + help="Skip the first Helios autoregressive chunk", + ) + parser.add_argument( + "--use-zero-init", + action=StoreBoolean, + default=SamplingParam.use_zero_init, + help="Enable Helios CFG zero initialization when supported", + ) + parser.add_argument( + "--zero-steps", + type=int, + default=SamplingParam.zero_steps, + help="Number of Helios CFG zero-initialization steps", + ) + parser.add_argument( + "--is-amplify-first-chunk", + action=StoreBoolean, + default=SamplingParam.is_amplify_first_chunk, + help="Use the amplified DMD schedule for the first Helios chunk", + ) parser.add_argument( "--save-video", action="store_true", diff --git a/fastvideo/api/schema.py b/fastvideo/api/schema.py index 7b75cbc290..40506b317c 100644 --- a/fastvideo/api/schema.py +++ b/fastvideo/api/schema.py @@ -166,6 +166,16 @@ class SamplingConfig: boundary_ratio: float | None = None sigmas: list[float] | None = None + # Helios autoregressive spatial-pyramid sampling. + pyramid_num_inference_steps_list: list[int] | None = None + history_sizes: list[int] | None = None + num_latent_frames_per_chunk: int = 9 + keep_first_frame: bool = True + is_skip_first_chunk: bool = False + use_zero_init: bool = True + zero_steps: int = 1 + is_amplify_first_chunk: bool = False + @dataclass class RequestRuntimeConfig: diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 0322e78d5b..c571eb31f0 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -4,6 +4,7 @@ from fastvideo.configs.models.dits.flux import FluxDiTConfig from fastvideo.configs.models.dits.flux_2 import Flux2Config from fastvideo.configs.models.dits.glm_image import GlmImageDiTConfig +from fastvideo.configs.models.dits.helios import HeliosConfig from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config @@ -24,6 +25,6 @@ "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig", "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", - "StableAudioConfig", "GlmImageDiTConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotVideoConfig", - "MiniMaxH3Config", "ZImageDiTConfig", "MMAudioArchConfig", "MMAudioTransformerConfig" + "StableAudioConfig", "GlmImageDiTConfig", "HeliosConfig", "LingBotWorld2CausalFastVideoConfig", + "LingBotVideoConfig", "MiniMaxH3Config", "ZImageDiTConfig", "MMAudioArchConfig", "MMAudioTransformerConfig" ] diff --git a/fastvideo/configs/models/dits/helios.py b/fastvideo/configs/models/dits/helios.py new file mode 100644 index 0000000000..42ad39bc9f --- /dev/null +++ b/fastvideo/configs/models/dits/helios.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig +from fastvideo.platforms import AttentionBackendEnum + + +def _is_transformer_block(name: str, module) -> bool: + del module + return name.startswith("blocks.") and name.split(".")[-1].isdigit() + + +@dataclass +class HeliosArchConfig(DiTArchConfig): + """Architecture fields for Helios-Distilled's history-aware DiT.""" + + _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_transformer_block]) + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ) + param_names_mapping: dict = field(default_factory=dict) + reverse_param_names_mapping: dict = field(default_factory=dict) + lora_param_names_mapping: dict = field(default_factory=dict) + + patch_size: tuple[int, int, int] = (1, 2, 2) + num_attention_heads: int = 40 + attention_head_dim: int = 128 + in_channels: int = 16 + out_channels: int = 16 + text_dim: int = 4096 + freq_dim: int = 256 + ffn_dim: int = 13824 + num_layers: int = 40 + cross_attn_norm: bool = True + qk_norm: str = "rms_norm_across_heads" + eps: float = 1e-6 + added_kv_proj_dim: int | None = None + rope_dim: tuple[int, int, int] = (44, 42, 42) + rope_theta: float = 10000.0 + guidance_cross_attn: bool = True + zero_history_timestep: bool = True + has_multi_term_memory_patch: bool = True + is_amplify_history: bool = False + history_scale_mode: str = "per_head" + + def __post_init__(self) -> None: + super().__post_init__() + self.out_channels = self.out_channels or self.in_channels + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.num_channels_latents = self.out_channels + if not self.cross_attn_norm: + raise ValueError("Helios currently requires cross_attn_norm=True") + if self.qk_norm != "rms_norm_across_heads": + raise ValueError("Helios currently requires qk_norm='rms_norm_across_heads'") + if self.added_kv_proj_dim is not None: + raise ValueError("Helios added_kv_proj_dim variants are not supported") + if not self.guidance_cross_attn: + raise ValueError("Helios currently requires guidance_cross_attn=True") + if not self.zero_history_timestep: + raise ValueError("Helios currently requires zero_history_timestep=True") + if not self.has_multi_term_memory_patch: + raise ValueError("Helios currently requires has_multi_term_memory_patch=True") + if self.is_amplify_history: + raise ValueError("Helios is_amplify_history variants are not supported") + if self.history_scale_mode != "per_head": + raise ValueError("Helios currently requires history_scale_mode='per_head'") + if sum(self.rope_dim) != self.attention_head_dim: + raise ValueError( + f"Helios rope_dim must sum to attention_head_dim, got {self.rope_dim} and {self.attention_head_dim}") + if any(dim % 2 for dim in self.rope_dim): + raise ValueError(f"Helios rope dimensions must be even: {self.rope_dim}") + + +@dataclass +class HeliosConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=HeliosArchConfig) + prefix: str = "Helios" diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index 7897e297c2..68b0b7a80d 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -5,6 +5,7 @@ from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig +from fastvideo.configs.pipelines.helios import HeliosPipelineConfig from fastvideo.configs.pipelines.hyworld import HYWorldConfig from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig, Kandinsky5I2VConfig, Kandinsky5T2VConfig from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig @@ -21,7 +22,8 @@ "HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig", "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig", "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", - "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig", - "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig", - "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig", "get_pipeline_config_cls_from_name" + "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HeliosPipelineConfig", "HYWorldConfig", + "Kandinsky5T2VConfig", "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "LingBotWorld2CausalFastI2V480PConfig", + "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig", + "get_pipeline_config_cls_from_name" ] diff --git a/fastvideo/configs/pipelines/helios.py b/fastvideo/configs/pipelines/helios.py new file mode 100644 index 0000000000..345f1c84bb --- /dev/null +++ b/fastvideo/configs/pipelines/helios.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Typed component wiring for Helios-Distilled T2V inference.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +import html + +import ftfy +import regex as re +import torch + +from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig +from fastvideo.configs.models.dits.helios import HeliosConfig +from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config +from fastvideo.configs.models.encoders.base import TextEncoderArchConfig +from fastvideo.configs.models.encoders.t5 import T5ArchConfig +from fastvideo.configs.models.vaes import WanVAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig + + +@dataclass +class HeliosT5ArchConfig(T5ArchConfig): + + def __post_init__(self) -> None: + super().__post_init__() + self.tokenizer_kwargs["padding"] = "max_length" + + +@dataclass +class HeliosT5Config(T5Config): + arch_config: TextEncoderArchConfig = field(default_factory=HeliosT5ArchConfig) + + +def helios_preprocess_text(prompt: str) -> str: + text = ftfy.fix_text(prompt) + text = html.unescape(html.unescape(text)).strip() + return re.sub(r"\s+", " ", text).strip() + + +def helios_postprocess_text(output: BaseEncoderOutput) -> torch.Tensor: + if output.last_hidden_state is None or output.attention_mask is None: + raise ValueError("Helios UMT5 output requires hidden states and attention mask") + hidden_states = output.last_hidden_state + sequence_lengths = output.attention_mask.gt(0).sum(dim=1).long() + if torch.isnan(hidden_states).any(): + raise ValueError("Helios UMT5 produced NaN hidden states") + trimmed = [hidden[:min(int(length), 512)] for hidden, length in zip(hidden_states, sequence_lengths, strict=True)] + return torch.stack( + [torch.cat([hidden, hidden.new_zeros(512 - hidden.shape[0], hidden.shape[1])]) for hidden in trimmed], + dim=0, + ) + + +def make_helios_text_encoder_config() -> T5Config: + return HeliosT5Config( + arch_config=HeliosT5ArchConfig( + architectures=["UMT5EncoderModel"], + vocab_size=256384, + d_model=4096, + d_kv=64, + d_ff=10240, + num_layers=24, + num_decoder_layers=24, + num_heads=64, + relative_attention_num_buckets=32, + relative_attention_max_distance=128, + dropout_rate=0.1, + layer_norm_epsilon=1e-6, + feed_forward_proj="gated-gelu", + is_encoder_decoder=True, + use_cache=True, + text_len=512, + ), + prefix="umt5", + ) + + +@dataclass +class HeliosPipelineConfig(PipelineConfig): + dit_config: DiTConfig = field(default_factory=HeliosConfig) + vae_config: VAEConfig = field(default_factory=WanVAEConfig) + text_encoder_configs: tuple[EncoderConfig, + ...] = field(default_factory=lambda: (make_helios_text_encoder_config(), )) + preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(default_factory=lambda: (helios_preprocess_text, )) + postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], + ...] = field(default_factory=lambda: (helios_postprocess_text, )) + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", )) + dit_precision: str = "bf16" + vae_precision: str = "fp32" + vae_decode_precision: str | None = "fp32" + vae_tiling: bool = False + vae_sp: bool = False + flow_shift: float | None = None + + def __post_init__(self) -> None: + self.vae_config.load_encoder = False + self.vae_config.load_decoder = True diff --git a/fastvideo/models/dits/helios.py b/fastvideo/models/dits/helios.py new file mode 100644 index 0000000000..5096128e40 --- /dev/null +++ b/fastvideo/models/dits/helios.py @@ -0,0 +1,770 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# Copyright 2025 The Helios Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Native FastVideo implementation of the Helios history-aware video DiT. + +The architecture follows the Apache-2.0 Diffusers Helios implementation while +using FastVideo linear, attention, sharding, and model-loader boundaries. +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.attention import DistributedAttention, LocalAttention +from fastvideo.configs.models.dits.helios import HeliosConfig +from fastvideo.distributed.communication_op import ( + sequence_model_parallel_all_gather_with_unpad, + sequence_model_parallel_shard, +) +from fastvideo.distributed.parallel_state import get_sp_world_size +from fastvideo.layers.layernorm import FP32LayerNorm, RMSNorm +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.visual_embedding import Timesteps +from fastvideo.models.dits.base import BaseDiT + + +def pad_for_3d_conv(value: torch.Tensor, kernel_size: tuple[int, int, int]) -> torch.Tensor: + _, _, frames, height, width = value.shape + patch_frames, patch_height, patch_width = kernel_size + pad_frames = (patch_frames - frames % patch_frames) % patch_frames + pad_height = (patch_height - height % patch_height) % patch_height + pad_width = (patch_width - width % patch_width) % patch_width + return F.pad( + value, + (0, pad_width, 0, pad_height, 0, pad_frames), + mode="replicate", + ) + + +def center_down_sample_3d(value: torch.Tensor, kernel_size: tuple[int, int, int]) -> torch.Tensor: + return F.avg_pool3d(value, kernel_size, stride=kernel_size) + + +def apply_rotary_emb_transposed(hidden_states: torch.Tensor, rotary_emb: torch.Tensor) -> torch.Tensor: + first, second = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + cos, sin = rotary_emb.unsqueeze(-2).chunk(2, dim=-1) + output = torch.empty_like(hidden_states) + output[..., 0::2] = first * cos[..., 0::2] - second * sin[..., 1::2] + output[..., 1::2] = first * sin[..., 1::2] + second * cos[..., 0::2] + return output.type_as(hidden_states) + + +class HeliosOutputNorm(nn.Module): + + def __init__(self, dim: int, eps: float = 1e-6) -> None: + super().__init__() + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + self.norm = FP32LayerNorm(dim, eps=eps, elementwise_affine=False) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + original_context_length: int, + ) -> torch.Tensor: + temb = temb[:, -original_context_length:, :] + shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2) + shift = shift.squeeze(2).to(hidden_states.device) + scale = scale.squeeze(2).to(hidden_states.device) + hidden_states = hidden_states[:, -original_context_length:, :] + return (self.norm(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states) + + +class HeliosAttention(nn.Module): + + def __init__( + self, + dim: int, + heads: int, + dim_head: int, + eps: float, + *, + is_cross_attention: bool, + supported_attention_backends, + quant_config=None, + prefix: str = "", + is_amplify_history: bool = False, + history_scale_mode: str = "per_head", + ) -> None: + super().__init__() + self.heads = heads + self.dim_head = dim_head + self.inner_dim = heads * dim_head + self.is_cross_attention = is_cross_attention + self.is_amplify_history = is_amplify_history + self.history_scale_mode = history_scale_mode + + self.to_q = ReplicatedLinear( + dim, + self.inner_dim, + quant_config=quant_config, + prefix=f"{prefix}.to_q", + ) + self.to_k = ReplicatedLinear( + dim, + self.inner_dim, + quant_config=quant_config, + prefix=f"{prefix}.to_k", + ) + self.to_v = ReplicatedLinear( + dim, + self.inner_dim, + quant_config=quant_config, + prefix=f"{prefix}.to_v", + ) + self.to_out = nn.ModuleList([ + ReplicatedLinear( + self.inner_dim, + dim, + quant_config=quant_config, + prefix=f"{prefix}.to_out.0", + ), + nn.Dropout(0.0), + ]) + self.norm_q = RMSNorm(self.inner_dim, eps=eps) + self.norm_k = RMSNorm(self.inner_dim, eps=eps) + + attention_cls = LocalAttention + if not is_cross_attention and get_sp_world_size() > 1: + attention_cls = DistributedAttention + self.attn = attention_cls( + num_heads=heads, + head_size=dim_head, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.impl", + ) + + if is_amplify_history: + if history_scale_mode == "scalar": + self.history_key_scale = nn.Parameter(torch.ones(1)) + elif history_scale_mode == "per_head": + self.history_key_scale = nn.Parameter(torch.ones(heads)) + else: + raise ValueError(f"Unknown history_scale_mode: {history_scale_mode}") + self.max_scale = 10.0 + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + rotary_emb: torch.Tensor | None = None, + original_seq_len: int | None = None, + original_context_length: int | None = None, + ) -> torch.Tensor: + key_value_states = hidden_states if encoder_hidden_states is None else encoder_hidden_states + query = self.to_q(hidden_states)[0] + key = self.to_k(key_value_states)[0] + value = self.to_v(key_value_states)[0] + query = self.norm_q(query).unflatten(2, (self.heads, self.dim_head)) + key = self.norm_k(key).unflatten(2, (self.heads, self.dim_head)) + value = value.unflatten(2, (self.heads, self.dim_head)) + + if rotary_emb is not None: + query = apply_rotary_emb_transposed(query, rotary_emb) + key = apply_rotary_emb_transposed(key, rotary_emb) + + if not self.is_cross_attention and self.is_amplify_history and original_context_length is not None: + history_seq_len = hidden_states.shape[1] - original_context_length + if history_seq_len > 0: + scale = 1.0 + torch.sigmoid(self.history_key_scale) * (self.max_scale - 1.0) + if self.history_scale_mode == "per_head": + scale = scale.view(1, 1, -1, 1) + key = torch.cat( + [key[:, :history_seq_len] * scale, key[:, history_seq_len:]], + dim=1, + ) + + if isinstance(self.attn, DistributedAttention): + output = self.attn(query, key, value, original_seq_len)[0] + else: + output = self.attn(query, key, value) + output = output.flatten(2) + output = self.to_out[0](output)[0] + return self.to_out[1](output) + + +class HeliosTimestepEmbedding(nn.Module): + + def __init__(self, frequency_dim: int, dim: int, *, prefix: str = "") -> None: + super().__init__() + self.linear_1 = ReplicatedLinear(frequency_dim, dim, prefix=f"{prefix}.linear_1") + self.act = nn.SiLU() + self.linear_2 = ReplicatedLinear(dim, dim, prefix=f"{prefix}.linear_2") + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + sample = self.linear_1(sample)[0] + sample = self.act(sample) + return self.linear_2(sample)[0] + + +class HeliosTextProjection(nn.Module): + + def __init__(self, input_dim: int, dim: int, *, prefix: str = "") -> None: + super().__init__() + self.linear_1 = ReplicatedLinear(input_dim, dim, prefix=f"{prefix}.linear_1") + self.act_1 = nn.GELU(approximate="tanh") + self.linear_2 = ReplicatedLinear(dim, dim, prefix=f"{prefix}.linear_2") + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.linear_1(hidden_states)[0] + hidden_states = self.act_1(hidden_states) + return self.linear_2(hidden_states)[0] + + +class HeliosTimeTextEmbedding(nn.Module): + + def __init__( + self, + dim: int, + time_freq_dim: int, + time_proj_dim: int, + text_embed_dim: int, + *, + prefix: str = "", + ) -> None: + super().__init__() + self.timesteps_proj = Timesteps( + num_channels=time_freq_dim, + flip_sin_to_cos=True, + downscale_freq_shift=0, + ) + self.time_embedder = HeliosTimestepEmbedding(time_freq_dim, dim, prefix=f"{prefix}.time_embedder") + self.act_fn = nn.SiLU() + self.time_proj = ReplicatedLinear(dim, time_proj_dim, prefix=f"{prefix}.time_proj") + self.text_embedder = HeliosTextProjection(text_embed_dim, dim, prefix=f"{prefix}.text_embedder") + + def forward( + self, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + is_return_encoder_hidden_states: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + timestep = self.timesteps_proj(timestep) + timestep = timestep.to(self.time_embedder.linear_1.weight.dtype) + temb = self.time_embedder(timestep) + if encoder_hidden_states is not None: + temb = temb.type_as(encoder_hidden_states) + timestep_proj = self.time_proj(self.act_fn(temb))[0] + if encoder_hidden_states is not None and is_return_encoder_hidden_states: + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + return temb, timestep_proj, encoder_hidden_states + + +class HeliosRotaryPosEmbed(nn.Module): + + def __init__(self, rope_dim: tuple[int, int, int], theta: float) -> None: + super().__init__() + self.dim_t, self.dim_y, self.dim_x = rope_dim + self.theta = theta + self.register_buffer("freqs_base_t", self._get_freqs_base(self.dim_t), persistent=False) + self.register_buffer("freqs_base_y", self._get_freqs_base(self.dim_y), persistent=False) + self.register_buffer("freqs_base_x", self._get_freqs_base(self.dim_x), persistent=False) + + def _get_freqs_base(self, dim: int) -> torch.Tensor: + exponent = torch.arange(0, dim, 2, dtype=torch.float32)[:dim // 2] / dim + return 1.0 / self.theta**exponent + + def materialize_buffers(self, device: torch.device) -> None: + self.freqs_base_t = self._get_freqs_base(self.dim_t).to(device) + self.freqs_base_y = self._get_freqs_base(self.dim_y).to(device) + self.freqs_base_x = self._get_freqs_base(self.dim_x).to(device) + + @staticmethod + def _get_frequency_batched(freqs_base: torch.Tensor, positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + with torch.autocast(device_type=positions.device.type, enabled=False): + freqs = torch.einsum("d,bthw->dbthw", freqs_base, positions) + freqs = freqs.repeat_interleave(2, dim=0) + return freqs.cos(), freqs.sin() + + @torch.no_grad() + def forward( + self, + frame_indices: torch.Tensor, + height: int, + width: int, + device: torch.device, + ) -> torch.Tensor: + batch_size, num_frames = frame_indices.shape + frame_indices = frame_indices.to(device=device, dtype=torch.float32) + y_coords = torch.arange(height, device=device, dtype=torch.float32) + x_coords = torch.arange(width, device=device, dtype=torch.float32) + grid_y, grid_x = torch.meshgrid(y_coords, x_coords, indexing="ij") + grid_t = frame_indices[:, :, None, None].expand(batch_size, num_frames, height, width) + grid_y = grid_y[None, None].expand(batch_size, num_frames, -1, -1) + grid_x = grid_x[None, None].expand(batch_size, num_frames, -1, -1) + cos_t, sin_t = self._get_frequency_batched(self.freqs_base_t, grid_t) + cos_y, sin_y = self._get_frequency_batched(self.freqs_base_y, grid_y) + cos_x, sin_x = self._get_frequency_batched(self.freqs_base_x, grid_x) + result = torch.cat([cos_t, cos_y, cos_x, sin_t, sin_y, sin_x], dim=0) + return result.permute(1, 0, 2, 3, 4) + + +class HeliosFeedForwardProject(nn.Module): + + def __init__(self, dim: int, ffn_dim: int, *, quant_config=None, prefix: str = "") -> None: + super().__init__() + self.proj = ReplicatedLinear(dim, ffn_dim, quant_config=quant_config, prefix=f"{prefix}.proj") + self.gelu = nn.GELU(approximate="tanh") + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.gelu(self.proj(hidden_states)[0]) + + +class HeliosFeedForward(nn.Module): + + def __init__(self, dim: int, ffn_dim: int, *, quant_config=None, prefix: str = "") -> None: + super().__init__() + self.net = nn.ModuleList([ + HeliosFeedForwardProject( + dim, + ffn_dim, + quant_config=quant_config, + prefix=f"{prefix}.net.0", + ), + nn.Dropout(0.0), + ReplicatedLinear( + ffn_dim, + dim, + quant_config=quant_config, + prefix=f"{prefix}.net.2", + ), + ]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.net[0](hidden_states) + hidden_states = self.net[1](hidden_states) + return self.net[2](hidden_states)[0] + + +class HeliosTransformerBlock(nn.Module): + + def __init__( + self, + dim: int, + ffn_dim: int, + num_heads: int, + qk_norm: str, + cross_attn_norm: bool, + eps: float, + guidance_cross_attn: bool, + supported_attention_backends, + *, + quant_config=None, + prefix: str = "", + is_amplify_history: bool = False, + history_scale_mode: str = "per_head", + ) -> None: + super().__init__() + if qk_norm != "rms_norm_across_heads": + raise ValueError(f"Unsupported Helios qk_norm: {qk_norm}") + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.attn1 = HeliosAttention( + dim, + num_heads, + dim // num_heads, + eps, + is_cross_attention=False, + supported_attention_backends=supported_attention_backends, + quant_config=quant_config, + prefix=f"{prefix}.attn1", + is_amplify_history=is_amplify_history, + history_scale_mode=history_scale_mode, + ) + self.attn2 = HeliosAttention( + dim, + num_heads, + dim // num_heads, + eps, + is_cross_attention=True, + supported_attention_backends=supported_attention_backends, + quant_config=quant_config, + prefix=f"{prefix}.attn2", + ) + self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.ffn = HeliosFeedForward(dim, ffn_dim, quant_config=quant_config, prefix=f"{prefix}.ffn") + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + self.guidance_cross_attn = guidance_cross_attn + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + rotary_emb: torch.Tensor, + current_token_mask: torch.Tensor, + original_seq_len: int, + original_context_length: int, + ) -> torch.Tensor: + ( + shift_msa, + scale_msa, + gate_msa, + c_shift_msa, + c_scale_msa, + c_gate_msa, + ) = (self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(6, dim=2) + shift_msa = shift_msa.squeeze(2) + scale_msa = scale_msa.squeeze(2) + gate_msa = gate_msa.squeeze(2) + c_shift_msa = c_shift_msa.squeeze(2) + c_scale_msa = c_scale_msa.squeeze(2) + c_gate_msa = c_gate_msa.squeeze(2) + + norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states) + attn_output = self.attn1( + norm_hidden_states, + rotary_emb=rotary_emb, + original_seq_len=original_seq_len, + original_context_length=(original_context_length if get_sp_world_size() == 1 else None), + ) + hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) + + if self.guidance_cross_attn: + if get_sp_world_size() == 1: + history_length = hidden_states.shape[1] - original_context_length + history_hidden_states, current_hidden_states = hidden_states.split( + [history_length, original_context_length], + dim=1, + ) + norm_hidden_states = self.norm2(current_hidden_states.float()).type_as(current_hidden_states) + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + ) + hidden_states = torch.cat( + [history_hidden_states, current_hidden_states + attn_output], + dim=1, + ) + else: + current_mask = current_token_mask.squeeze(-1).bool() + current_counts = current_mask.sum(dim=1) + if not torch.equal(current_counts, current_counts[:1].expand_as(current_counts)): + raise ValueError("Helios SP shards require equal current-token counts per batch item") + current_count = int(current_counts[0].item()) + if current_count > 0: + current_hidden_states = hidden_states[current_mask].reshape( + hidden_states.shape[0], + current_count, + hidden_states.shape[-1], + ) + norm_hidden_states = self.norm2(current_hidden_states.float()).type_as(current_hidden_states) + current_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + ) + attn_output = torch.zeros_like(hidden_states).masked_scatter( + current_mask.unsqueeze(-1).expand_as(hidden_states), + current_output, + ) + hidden_states = hidden_states + attn_output + else: + norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states) + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + ) + hidden_states = hidden_states + attn_output + + norm_hidden_states = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + + c_shift_msa).type_as(hidden_states) + ff_output = self.ffn(norm_hidden_states) + return (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + + +class HeliosTransformer3DModel(BaseDiT): + _fsdp_shard_conditions = HeliosConfig()._fsdp_shard_conditions + _compile_conditions = HeliosConfig()._compile_conditions + _supported_attention_backends = HeliosConfig()._supported_attention_backends + param_names_mapping = HeliosConfig().param_names_mapping + reverse_param_names_mapping = HeliosConfig().reverse_param_names_mapping + lora_param_names_mapping = HeliosConfig().lora_param_names_mapping + + def __init__(self, config: HeliosConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.in_channels = config.in_channels + self.out_channels = config.out_channels + self.num_channels_latents = config.num_channels_latents + self.patch_size = tuple(config.patch_size) + self.zero_history_timestep = config.zero_history_timestep + self.quant_config = config.quant_config + + if config.num_attention_heads % get_sp_world_size() != 0: + raise ValueError(f"Helios heads ({config.num_attention_heads}) must be divisible by " + f"sequence parallel size ({get_sp_world_size()})") + inner_dim = config.hidden_size + self.rope = HeliosRotaryPosEmbed(tuple(config.rope_dim), config.rope_theta) + self.patch_embedding = nn.Conv3d( + config.in_channels, + inner_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + ) + if config.has_multi_term_memory_patch: + self.patch_short = nn.Conv3d( + config.in_channels, + inner_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + ) + self.patch_mid = nn.Conv3d( + config.in_channels, + inner_dim, + kernel_size=tuple(2 * value for value in self.patch_size), + stride=tuple(2 * value for value in self.patch_size), + ) + self.patch_long = nn.Conv3d( + config.in_channels, + inner_dim, + kernel_size=tuple(4 * value for value in self.patch_size), + stride=tuple(4 * value for value in self.patch_size), + ) + self.condition_embedder = HeliosTimeTextEmbedding( + dim=inner_dim, + time_freq_dim=config.freq_dim, + time_proj_dim=inner_dim * 6, + text_embed_dim=config.text_dim, + prefix=f"{config.prefix}.condition_embedder", + ) + self.blocks = nn.ModuleList([ + HeliosTransformerBlock( + inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.guidance_cross_attn, + self._supported_attention_backends, + quant_config=config.quant_config, + prefix=f"{config.prefix}.blocks.{index}", + is_amplify_history=config.is_amplify_history, + history_scale_mode=config.history_scale_mode, + ) for index in range(config.num_layers) + ]) + self.norm_out = HeliosOutputNorm(inner_dim, config.eps) + self.proj_out = ReplicatedLinear( + inner_dim, + config.out_channels * math.prod(self.patch_size), + quant_config=config.quant_config, + prefix=f"{config.prefix}.proj_out", + ) + self.gradient_checkpointing = False + self.__post_init__() + + def materialize_non_persistent_buffers(self, device: torch.device, dtype: torch.dtype | None = None) -> None: + del dtype + self.rope.materialize_buffers(device) + + @staticmethod + def _validate_history_pair( + history: torch.Tensor | None, + history_indices: torch.Tensor | None, + history_name: str, + history_indices_name: str, + ) -> None: + if (history is None) != (history_indices is None): + raise ValueError(f"{history_name} and {history_indices_name} must be provided together") + + def _patch_history( + self, + hidden_states: torch.Tensor, + rotary_emb: torch.Tensor, + history: torch.Tensor | None, + history_indices: torch.Tensor | None, + patch: nn.Conv3d, + rope_height: int | None, + rope_width: int | None, + downsample: tuple[int, int, int] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, int | None, int | None]: + if history is None or history_indices is None: + return hidden_states, rotary_emb, rope_height, rope_width + if downsample is not None: + history = pad_for_3d_conv(history, tuple(patch.kernel_size)) + history = patch(history) + _, _, _, patched_height, patched_width = history.shape + history = history.flatten(2).transpose(1, 2) + rope_height = patched_height if rope_height is None else rope_height + rope_width = patched_width if rope_width is None else rope_width + history_rotary = self.rope( + history_indices, + height=rope_height, + width=rope_width, + device=history.device, + ) + if downsample is not None: + history_rotary = pad_for_3d_conv(history_rotary, downsample) + history_rotary = center_down_sample_3d(history_rotary, downsample) + history_rotary = history_rotary.flatten(2).transpose(1, 2) + return ( + torch.cat([history, hidden_states], dim=1), + torch.cat([history_rotary, rotary_emb], dim=1), + patched_height, + patched_width, + ) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.LongTensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + indices_hidden_states: torch.Tensor | None = None, + indices_latents_history_short: torch.Tensor | None = None, + indices_latents_history_mid: torch.Tensor | None = None, + indices_latents_history_long: torch.Tensor | None = None, + latents_history_short: torch.Tensor | None = None, + latents_history_mid: torch.Tensor | None = None, + latents_history_long: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + del kwargs + self._validate_history_pair( + latents_history_short, + indices_latents_history_short, + "latents_history_short", + "indices_latents_history_short", + ) + self._validate_history_pair( + latents_history_mid, + indices_latents_history_mid, + "latents_history_mid", + "indices_latents_history_mid", + ) + self._validate_history_pair( + latents_history_long, + indices_latents_history_long, + "latents_history_long", + "indices_latents_history_long", + ) + if isinstance(encoder_hidden_states, list): + encoder_hidden_states = encoder_hidden_states[0] + batch_size = hidden_states.shape[0] + patch_frames, patch_height, patch_width = self.patch_size + + hidden_states = self.patch_embedding(hidden_states) + _, _, post_frames, post_height, post_width = hidden_states.shape + if indices_hidden_states is None: + indices_hidden_states = torch.arange(post_frames).unsqueeze(0).expand(batch_size, -1) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + rotary_emb = (self.rope( + indices_hidden_states, + height=post_height, + width=post_width, + device=hidden_states.device, + ).flatten(2).transpose(1, 2)) + original_context_length = hidden_states.shape[1] + + hidden_states, rotary_emb, history_height, history_width = self._patch_history( + hidden_states, + rotary_emb, + latents_history_short, + indices_latents_history_short, + self.patch_short, + None, + None, + ) + if (latents_history_mid is not None or latents_history_long is not None) and (latents_history_short is None): + raise ValueError("Helios mid/long history requires short history geometry") + hidden_states, rotary_emb, _, _ = self._patch_history( + hidden_states, + rotary_emb, + latents_history_mid, + indices_latents_history_mid, + self.patch_mid, + history_height, + history_width, + downsample=(2, 2, 2), + ) + hidden_states, rotary_emb, _, _ = self._patch_history( + hidden_states, + rotary_emb, + latents_history_long, + indices_latents_history_long, + self.patch_long, + history_height, + history_width, + downsample=(4, 4, 4), + ) + + history_context_length = hidden_states.shape[1] - original_context_length + if self.zero_history_timestep: + timestep_zero = torch.zeros(1, dtype=timestep.dtype, device=timestep.device) + temb_zero, timestep_proj_zero, _ = self.condition_embedder( + timestep_zero, + encoder_hidden_states, + is_return_encoder_hidden_states=False, + ) + temb_zero = temb_zero.unsqueeze(1).expand(batch_size, history_context_length, -1) + timestep_proj_zero = (timestep_proj_zero.unflatten(-1, (6, -1)).view(1, 6, 1, -1).expand( + batch_size, -1, history_context_length, -1)) + + temb, timestep_proj, encoder_hidden_states = self.condition_embedder(timestep, encoder_hidden_states) + assert encoder_hidden_states is not None + timestep_proj = timestep_proj.unflatten(-1, (6, -1)) + temb = temb.view(batch_size, 1, -1).expand(batch_size, original_context_length, -1) + timestep_proj = timestep_proj.view(batch_size, 6, 1, -1).expand(batch_size, 6, original_context_length, -1) + if self.zero_history_timestep: + temb = torch.cat([temb_zero, temb], dim=1) + timestep_proj = torch.cat([timestep_proj_zero, timestep_proj], dim=2) + timestep_proj = timestep_proj.permute(0, 2, 1, 3) + + current_token_mask = hidden_states.new_zeros(batch_size, hidden_states.shape[1], 1) + current_token_mask[:, -original_context_length:] = 1 + original_seq_len = hidden_states.shape[1] + if get_sp_world_size() > 1: + hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1) + rotary_emb = sequence_model_parallel_shard(rotary_emb, dim=1)[0] + timestep_proj = sequence_model_parallel_shard(timestep_proj, dim=1)[0] + current_token_mask = sequence_model_parallel_shard(current_token_mask, dim=1)[0] + + for block in self.blocks: + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + rotary_emb, + current_token_mask, + original_seq_len, + original_context_length, + ) + + if get_sp_world_size() > 1: + hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) + hidden_states = self.norm_out(hidden_states, temb, original_context_length) + hidden_states = self.proj_out(hidden_states)[0] + hidden_states = hidden_states.reshape( + batch_size, + post_frames, + post_height, + post_width, + patch_frames, + patch_height, + patch_width, + -1, + ) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + return hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + +EntryClass = HeliosTransformer3DModel diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 4f9939a62d..b2141d1a8d 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -141,6 +141,7 @@ "SelfForcingFlowMatchScheduler": ("schedulers", "scheduling_self_forcing_flow_match", "SelfForcingFlowMatchScheduler"), "RCMScheduler": ("schedulers", "scheduling_rcm", "RCMScheduler"), + "HeliosDMDScheduler": ("schedulers", "scheduling_helios_dmd", "HeliosDMDScheduler"), } _UPSAMPLERS = { diff --git a/fastvideo/models/schedulers/scheduling_helios_dmd.py b/fastvideo/models/schedulers/scheduling_helios_dmd.py new file mode 100644 index 0000000000..8e9f8e43cb --- /dev/null +++ b/fastvideo/models/schedulers/scheduling_helios_dmd.py @@ -0,0 +1,318 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# Copyright 2025 The Helios Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Native scheduler for the distilled three-stage Helios pyramid.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Literal + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + +from fastvideo.models.schedulers.base import BaseScheduler + + +@dataclass +class HeliosDMDSchedulerOutput(BaseOutput): + prev_sample: torch.FloatTensor + model_outputs: torch.FloatTensor | None = None + last_sample: torch.FloatTensor | None = None + this_order: int | None = None + + +class HeliosDMDScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): + """DMD flow scheduler used by `BestWishYsh/Helios-Distilled`.""" + + _compatibles: list[Any] = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + stages: int = 3, + stage_range: list[float] | None = None, + gamma: float = 1 / 3, + prediction_type: str = "flow_prediction", + use_flow_sigmas: bool = True, + use_dynamic_shifting: bool = False, + time_shift_type: Literal["exponential", "linear"] = "linear", + scheduler_type: str = "dmd", + _diffusers_version: str | None = None, + **kwargs, + ) -> None: + del scheduler_type, _diffusers_version, kwargs + if stage_range is None: + stage_range = [0, 1 / 3, 2 / 3, 1] + self.register_to_config(stage_range=stage_range) + self.num_train_timesteps = num_train_timesteps + self.timestep_ratios: dict[int, tuple[float, float]] = {} + self.timesteps_per_stage: dict[int, torch.Tensor] = {} + self.sigmas_per_stage: dict[int, torch.Tensor] = {} + self.start_sigmas: dict[int, float] = {} + self.end_sigmas: dict[int, float] = {} + self.ori_start_sigmas: dict[int, float] = {} + + self.init_sigmas_for_each_stage() + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + self.gamma = gamma + self.last_sample = None + self._step_index = None + self._begin_index = None + BaseScheduler.__init__(self) + + def init_sigmas(self) -> None: + alphas = np.linspace( + 1, + 1 / self.config.num_train_timesteps, + self.config.num_train_timesteps + 1, + ) + sigmas = 1.0 - alphas + sigmas = np.flip(self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas))[:-1].copy() + self.sigmas = torch.from_numpy(sigmas) + self.timesteps = (self.sigmas * self.config.num_train_timesteps).clone() + self._step_index = None + self._begin_index = None + + def init_sigmas_for_each_stage(self) -> None: + self.init_sigmas() + stage_distance = [] + training_steps = self.config.num_train_timesteps + for stage_index in range(self.config.stages): + start_index = max(int(self.config.stage_range[stage_index] * training_steps), 0) + end_index = min( + int(self.config.stage_range[stage_index + 1] * training_steps), + training_steps, + ) + start_sigma = self.sigmas[start_index].item() + end_sigma = self.sigmas[end_index].item() if end_index < training_steps else 0.0 + self.ori_start_sigmas[stage_index] = start_sigma + if stage_index != 0: + original_sigma = 1 - start_sigma + corrected_sigma = (1 / (math.sqrt(1 + 1 / self.config.gamma) * (1 - original_sigma) + original_sigma) * + original_sigma) + start_sigma = 1 - corrected_sigma + stage_distance.append(start_sigma - end_sigma) + self.start_sigmas[stage_index] = start_sigma + self.end_sigmas[stage_index] = end_sigma + + total_distance = sum(stage_distance) + for stage_index in range(self.config.stages): + start_ratio = 0.0 if stage_index == 0 else sum(stage_distance[:stage_index]) / total_distance + end_ratio = (0.9999999999999999 if stage_index == self.config.stages - + 1 else sum(stage_distance[:stage_index + 1]) / total_distance) + self.timestep_ratios[stage_index] = (start_ratio, end_ratio) + + for stage_index in range(self.config.stages): + start_ratio, end_ratio = self.timestep_ratios[stage_index] + timestep_max = min(self.timesteps[int(start_ratio * training_steps)], 999) + timestep_min = self.timesteps[min(int(end_ratio * training_steps), training_steps - 1)] + timesteps = np.linspace(timestep_max, timestep_min, training_steps + 1) + self.timesteps_per_stage[stage_index] = (timesteps[:-1] if isinstance(timesteps, torch.Tensor) else + torch.from_numpy(timesteps[:-1])) + stage_sigmas = np.linspace(0.999, 0, training_steps + 1) + self.sigmas_per_stage[stage_index] = torch.from_numpy(stage_sigmas[:-1]) + + @property + def step_index(self): + return self._step_index + + @property + def begin_index(self): + return self._begin_index + + def set_begin_index(self, begin_index: int = 0) -> None: + self._begin_index = begin_index + + def set_shift(self, shift: float) -> None: + self.register_to_config(shift=shift) + self.timestep_ratios.clear() + self.timesteps_per_stage.clear() + self.sigmas_per_stage.clear() + self.start_sigmas.clear() + self.end_sigmas.clear() + self.ori_start_sigmas.clear() + self.init_sigmas_for_each_stage() + + def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor: + del timestep + return sample + + def set_timesteps( + self, + num_inference_steps: int, + stage_index: int | None = None, + device: str | torch.device | None = None, + sigmas: np.ndarray | None = None, + mu: float | None = None, + is_amplify_first_chunk: bool = False, + ) -> None: + num_inference_steps = num_inference_steps * 2 + 1 if is_amplify_first_chunk else num_inference_steps + 1 + self.num_inference_steps = num_inference_steps + self.init_sigmas() + + if self.config.stages == 1: + if sigmas is None: + sigmas = np.linspace( + 1, + 1 / self.config.num_train_timesteps, + num_inference_steps + 1, + )[:-1].astype(np.float32) + if self.config.shift != 1.0: + if self.config.use_dynamic_shifting: + raise ValueError("Fixed shift and dynamic shifting cannot both be active") + sigmas = self.time_shift(self.config.shift, 1.0, sigmas) + timesteps = (sigmas * self.config.num_train_timesteps).copy() + sigma_tensor = torch.from_numpy(sigmas) + else: + if stage_index is None: + raise ValueError("stage_index is required for multi-stage Helios") + stage_timesteps = self.timesteps_per_stage[stage_index] + timesteps = np.linspace( + stage_timesteps[0].item(), + stage_timesteps[-1].item(), + num_inference_steps, + ) + stage_sigmas = self.sigmas_per_stage[stage_index] + ratios = np.linspace( + stage_sigmas[0].item(), + stage_sigmas[-1].item(), + num_inference_steps, + ) + sigma_tensor = torch.from_numpy(ratios) + + self.timesteps = torch.from_numpy(timesteps).to(device=device) + self.sigmas = torch.cat([sigma_tensor, torch.zeros(1)]).to(device=device) + self._step_index = None + self.reset_scheduler_history() + self.timesteps = self.timesteps[:-1] + self.sigmas = torch.cat([self.sigmas[:-2], self.sigmas[-1:]]) + + if self.config.use_dynamic_shifting: + if self.config.shift != 1.0: + raise ValueError("Dynamic shifting requires shift=1.0") + if mu is None: + raise ValueError("Dynamic shifting requires mu") + self.sigmas = self.time_shift(mu, 1.0, self.sigmas) + if self.config.stages == 1: + self.timesteps = self.sigmas[:-1] * self.config.num_train_timesteps + else: + assert stage_index is not None + stage_timesteps = self.timesteps_per_stage[stage_index] + self.timesteps = stage_timesteps.min() + self.sigmas[:-1] * (stage_timesteps.max() - + stage_timesteps.min()) + + def time_shift(self, mu: float, sigma: float, timesteps): + if self.config.time_shift_type == "exponential": + return math.exp(mu) / (math.exp(mu) + (1 / timesteps - 1)**sigma) + if self.config.time_shift_type == "linear": + return mu / (mu + (1 / timesteps - 1)**sigma) + raise ValueError(f"Unknown time_shift_type: {self.config.time_shift_type}") + + @staticmethod + def add_noise( + original_samples: torch.Tensor, + noise: torch.Tensor, + timestep: torch.Tensor, + sigmas: torch.Tensor, + timesteps: torch.Tensor, + ) -> torch.Tensor: + sigmas = sigmas.to(noise.device) + timesteps = timesteps.to(noise.device) + timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1) + return ((1 - sigma) * original_samples + sigma * noise).type_as(noise) + + @staticmethod + def convert_flow_pred_to_x0( + flow_pred: torch.Tensor, + sample: torch.Tensor, + timestep: torch.Tensor, + sigmas: torch.Tensor, + timesteps: torch.Tensor, + ) -> torch.Tensor: + original_dtype = flow_pred.dtype + device = flow_pred.device + flow_pred, sample, sigmas, timesteps = (value.double().to(device) + for value in (flow_pred, sample, sigmas, timesteps)) + timestep_id = torch.argmin((timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1) + return (sample - sigma * flow_pred).to(original_dtype) + + def step( + self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor | None = None, + sample: torch.FloatTensor | None = None, + generator: torch.Generator | None = None, + return_dict: bool = True, + cur_sampling_step: int = 0, + dmd_noisy_tensor: torch.FloatTensor | None = None, + dmd_sigmas: torch.FloatTensor | None = None, + dmd_timesteps: torch.FloatTensor | None = None, + all_timesteps: torch.FloatTensor | None = None, + ) -> HeliosDMDSchedulerOutput | tuple[torch.Tensor]: + del generator + if (timestep is None or sample is None or dmd_noisy_tensor is None or dmd_sigmas is None + or dmd_timesteps is None or all_timesteps is None): + raise ValueError("Helios DMD step requires all stage-local tensors") + predicted_x0 = self.convert_flow_pred_to_x0( + model_output, + sample, + torch.full( + (model_output.shape[0], ), + timestep, + dtype=torch.long, + device=model_output.device, + ), + dmd_sigmas, + dmd_timesteps, + ) + if cur_sampling_step < len(all_timesteps) - 1: + prev_sample = self.add_noise( + predicted_x0, + dmd_noisy_tensor, + torch.full( + (model_output.shape[0], ), + all_timesteps[cur_sampling_step + 1], + dtype=torch.long, + device=model_output.device, + ), + dmd_sigmas, + dmd_timesteps, + ) + else: + prev_sample = predicted_x0 + if not return_dict: + return (prev_sample, ) + return HeliosDMDSchedulerOutput(prev_sample=prev_sample) + + def reset_scheduler_history(self) -> None: + self._step_index = None + self._begin_index = None + + def __len__(self) -> int: + return self.config.num_train_timesteps + + +EntryClass = HeliosDMDScheduler diff --git a/fastvideo/pipelines/basic/helios/__init__.py b/fastvideo/pipelines/basic/helios/__init__.py new file mode 100644 index 0000000000..9881313609 --- /dev/null +++ b/fastvideo/pipelines/basic/helios/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/fastvideo/pipelines/basic/helios/helios_pipeline.py b/fastvideo/pipelines/basic/helios/helios_pipeline.py new file mode 100644 index 0000000000..64211eabe2 --- /dev/null +++ b/fastvideo/pipelines/basic/helios/helios_pipeline.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastVideo pipeline for Helios-Distilled text-to-video generation.""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline +from fastvideo.pipelines.basic.helios.stages import ( + HeliosChunkDecodingStage, + HeliosInputValidationStage, + HeliosPyramidDenoisingStage, +) +from fastvideo.pipelines.stages import ConditioningStage, TextEncodingStage + + +class HeliosPyramidPipeline(LoRAPipeline, ComposedPipelineBase): + """Autoregressive three-level spatial-pyramid pipeline for Helios.""" + + _required_config_modules = [ + "text_encoder", + "tokenizer", + "vae", + "transformer", + "scheduler", + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + self.add_stage( + stage_name="input_validation_stage", + stage=HeliosInputValidationStage(), + ) + self.add_stage( + stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + ) + self.add_stage( + stage_name="conditioning_stage", + stage=ConditioningStage(), + ) + self.add_stage( + stage_name="pyramid_denoising_stage", + stage=HeliosPyramidDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + pipeline=self, + ), + ) + self.add_stage( + stage_name="chunk_decoding_stage", + stage=HeliosChunkDecodingStage( + vae=self.get_module("vae"), + pipeline=self, + ), + ) + + +EntryClass = HeliosPyramidPipeline diff --git a/fastvideo/pipelines/basic/helios/pipeline_utils.py b/fastvideo/pipelines/basic/helios/pipeline_utils.py new file mode 100644 index 0000000000..02f3a2c685 --- /dev/null +++ b/fastvideo/pipelines/basic/helios/pipeline_utils.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Deterministic math and RNG helpers for Helios pyramid sampling.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence + +import torch +import torch.nn.functional as F + + +def calculate_shift( + image_seq_len: int, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +) -> float: + """Calculate the dynamic flow shift for a pyramid stage.""" + slope = (max_shift - base_shift) / (max_seq_len - base_seq_len) + intercept = base_shift - slope * base_seq_len + return image_seq_len * slope + intercept + + +def get_num_latent_chunks( + num_frames: int, + num_latent_frames_per_chunk: int, + temporal_scale_factor: int, +) -> int: + pixel_frames_per_chunk = (num_latent_frames_per_chunk - 1) * temporal_scale_factor + 1 + return max(1, math.ceil(num_frames / pixel_frames_per_chunk)) + + +def get_generated_pixel_frames(num_latent_frames: int, temporal_scale_factor: int) -> int: + return (num_latent_frames - 1) // temporal_scale_factor * temporal_scale_factor + 1 + + +def build_helios_frame_indices( + history_sizes: Sequence[int], + num_latent_frames_per_chunk: int, + keep_first_frame: bool, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if len(history_sizes) != 3: + raise ValueError(f"Helios requires three history sizes, got {list(history_sizes)}") + if not keep_first_frame: + raise ValueError("The initial FastVideo Helios T2V port requires keep_first_frame=True") + + history_long, history_mid, history_one = history_sizes + indices = torch.arange( + 1 + sum(history_sizes) + num_latent_frames_per_chunk, + device=device, + ) + prefix, long_indices, mid_indices, one_indices, current_indices = indices.split([ + 1, + history_long, + history_mid, + history_one, + num_latent_frames_per_chunk, + ]) + short_indices = torch.cat([prefix, one_indices]) + return ( + current_indices.unsqueeze(0), + short_indices.unsqueeze(0), + mid_indices.unsqueeze(0), + long_indices.unsqueeze(0), + ) + + +def downsample_to_pyramid_base(latents: torch.Tensor, num_stages: int) -> torch.Tensor: + batch_size, channels, num_frames, height, width = latents.shape + flattened = latents.permute(0, 2, 1, 3, 4).reshape( + batch_size * num_frames, + channels, + height, + width, + ) + for _ in range(num_stages - 1): + height //= 2 + width //= 2 + flattened = F.interpolate(flattened, size=(height, width), mode="bilinear") * 2 + return flattened.reshape( + batch_size, + num_frames, + channels, + height, + width, + ).permute(0, 2, 1, 3, 4) + + +def _randn_tensor( + shape: tuple[int, ...], + generator: torch.Generator | list[torch.Generator], + device: torch.device, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + if isinstance(generator, list): + if len(generator) != shape[0]: + raise ValueError(f"Expected one generator per batch item ({shape[0]}), got {len(generator)}") + samples = [ + torch.randn( + (1, *shape[1:]), + generator=item, + device=item.device, + dtype=dtype, + ).to(device) for item in generator + ] + return torch.cat(samples, dim=0) + return torch.randn( + shape, + generator=generator, + device=generator.device, + dtype=dtype, + ).to(device) + + +def sample_block_noise( + scheduler, + shape: tuple[int, int, int, int, int], + patch_size: tuple[int, int, int], + device: torch.device, + generator: torch.Generator | list[torch.Generator], +) -> torch.Tensor: + if isinstance(generator, list): + generator = generator[0] + batch_size, channels, num_frames, height, width = shape + _, patch_height, patch_width = patch_size + if height % patch_height or width % patch_width: + raise ValueError(f"Noise shape {(height, width)} must be divisible by patch {(patch_height, patch_width)}") + + block_size = patch_height * patch_width + gamma = scheduler.config.gamma + covariance = (torch.eye(block_size, device=device) * (1 + gamma) - + torch.ones(block_size, block_size, device=device) * gamma) + covariance += torch.eye(block_size, device=device) * 1e-8 + cholesky = torch.linalg.cholesky(covariance.float()) + + block_count = batch_size * channels * num_frames * (height // patch_height) * (width // patch_width) + standard_noise = torch.randn( + block_count, + block_size, + generator=generator, + device=generator.device, + ).to(device) + noise = standard_noise @ cholesky.T + noise = noise.view( + batch_size, + channels, + num_frames, + height // patch_height, + width // patch_width, + patch_height, + patch_width, + ) + return noise.permute(0, 1, 2, 3, 5, 4, 6).reshape(shape) diff --git a/fastvideo/pipelines/basic/helios/presets.py b/fastvideo/pipelines/basic/helios/presets.py new file mode 100644 index 0000000000..babc9a23ba --- /dev/null +++ b/fastvideo/pipelines/basic/helios/presets.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Inference preset for the official Helios-Distilled checkpoint.""" + +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +HELIOS_DISTILLED_NEGATIVE_PROMPT = ("Bright tones, overexposed, static, blurred details, subtitles, style, works, " + "paintings, images, static, overall gray, worst quality, low quality, JPEG " + "compression residue, ugly, incomplete, extra fingers, poorly drawn hands, " + "poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, " + "still picture, messy background, three legs, many people in the background, " + "walking backwards") + +_PYRAMID_STAGE = PresetStageSpec( + name="pyramid_denoise", + kind="denoising", + description="Three-level autoregressive DMD denoising", + allowed_overrides=frozenset({ + "guidance_scale", + "pyramid_num_inference_steps_list", + "is_amplify_first_chunk", + }), +) + +HELIOS_DISTILLED_T2V = InferencePreset( + name="helios_distilled_t2v", + version=1, + model_family="helios", + description="Helios-Distilled autoregressive text-to-video", + workload_type="t2v", + stage_schemas=(_PYRAMID_STAGE, ), + defaults={ + "height": 384, + "width": 640, + "num_frames": 240, + "fps": 24, + "guidance_scale": 1.0, + "num_inference_steps": 2, + "pyramid_num_inference_steps_list": [2, 2, 2], + "history_sizes": [16, 2, 1], + "num_latent_frames_per_chunk": 9, + "keep_first_frame": True, + "is_skip_first_chunk": False, + "use_zero_init": True, + "zero_steps": 1, + "is_amplify_first_chunk": True, + "max_sequence_length": 512, + "negative_prompt": HELIOS_DISTILLED_NEGATIVE_PROMPT, + }, +) + +ALL_PRESETS = (HELIOS_DISTILLED_T2V, ) diff --git a/fastvideo/pipelines/basic/helios/stages.py b/fastvideo/pipelines/basic/helios/stages.py new file mode 100644 index 0000000000..94b14d0a3b --- /dev/null +++ b/fastvideo/pipelines/basic/helios/stages.py @@ -0,0 +1,409 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Model-specific stages for Helios-Distilled T2V inference.""" + +import math +import weakref + +import torch +import torch.nn.functional as F + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.hooks.activation_trace import trace_step +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import TransformerLoader, VAELoader +from fastvideo.pipelines.basic.helios.pipeline_utils import ( + _randn_tensor, + build_helios_frame_indices, + calculate_shift, + downsample_to_pyramid_base, + get_generated_pixel_frames, + get_num_latent_chunks, + sample_block_noise, +) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.decoding import DecodingStage +from fastvideo.pipelines.stages.input_validation import InputValidationStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult +from fastvideo.utils import PRECISION_TO_TYPE + +logger = init_logger(__name__) + + +class HeliosInputValidationStage(InputValidationStage): + """Validate the intentionally narrow first Helios contribution: T2V.""" + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + batch = super().forward(batch, fastvideo_args) + if not isinstance(batch.height, int) or not isinstance(batch.width, int): + raise ValueError("Helios T2V expects scalar height and width") + if batch.height % 64 or batch.width % 64: + raise ValueError("Helios height and width must be divisible by 64 for two pyramid downsamples") + if not isinstance(batch.num_frames, int) or batch.num_frames <= 0: + raise ValueError(f"Helios num_frames must be a positive integer, got {batch.num_frames}") + if batch.num_videos_per_prompt != 1: + raise ValueError("Helios currently supports num_videos_per_prompt=1") + if isinstance(batch.prompt, list) and len(batch.prompt) != 1: + raise ValueError("Helios currently accepts one prompt per FastVideo request") + if batch.image_path is not None or batch.video_path is not None or batch.pil_image is not None: + raise ValueError("This initial Helios contribution supports T2V only") + if batch.latents is not None: + raise ValueError("Pre-generated Helios chunk latents are not supported yet") + + steps = batch.pyramid_num_inference_steps_list + if steps is None: + steps = [batch.num_inference_steps] * 3 + if len(steps) != 3 or any(not isinstance(value, int) or value <= 0 for value in steps): + raise ValueError(f"Helios requires three positive pyramid step counts, got {steps}") + batch.pyramid_num_inference_steps_list = list(steps) + + history_sizes = batch.history_sizes or [16, 2, 1] + if len(history_sizes) != 3 or any(not isinstance(value, int) or value <= 0 for value in history_sizes): + raise ValueError(f"Helios requires three positive history sizes, got {history_sizes}") + batch.history_sizes = sorted(history_sizes, reverse=True) + if batch.num_latent_frames_per_chunk != 9: + raise ValueError("Helios-Distilled requires num_latent_frames_per_chunk=9") + if not batch.keep_first_frame: + raise ValueError("The initial Helios T2V port requires keep_first_frame=True") + if batch.is_skip_first_chunk: + raise ValueError("is_skip_first_chunk is only meaningful for conditioned Helios modes") + if batch.zero_steps < 0: + raise ValueError("Helios zero_steps must be non-negative") + return batch + + +class HeliosPyramidDenoisingStage(PipelineStage): + """Generate autoregressive latent chunks with the three-stage DMD sampler.""" + + performance_component_metric = "dit_time_s" + + def __init__(self, transformer, scheduler, pipeline=None) -> None: + super().__init__() + self.transformer = transformer + self.scheduler = scheduler + self.pipeline = weakref.ref(pipeline) if pipeline else None + + @staticmethod + def _scheduler_value(scheduler, name: str, default): + value = getattr(scheduler.config, name, None) + if value is None and hasattr(scheduler.config, "get"): + value = scheduler.config.get(name, default) + return default if value is None else value + + def _load_or_move_transformer(self, fastvideo_args: FastVideoArgs): + pipeline = self.pipeline() if self.pipeline else None + if not fastvideo_args.model_loaded["transformer"]: + self.transformer = TransformerLoader().load( + fastvideo_args.model_paths["transformer"], + fastvideo_args, + ) + if pipeline is not None: + pipeline.add_module("transformer", self.transformer) + fastvideo_args.model_loaded["transformer"] = True + + if (fastvideo_args.dit_cpu_offload and not fastvideo_args.dit_layerwise_offload + and not fastvideo_args.use_fsdp_inference and next(self.transformer.parameters()).device.type == "cpu"): + self.transformer.to(get_local_torch_device()) + + def _transformer_forward( + self, + latents: torch.Tensor, + timestep: torch.Tensor, + prompt_embeds: torch.Tensor, + histories: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + indices: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + batch: ForwardBatch, + target_dtype: torch.dtype, + ) -> torch.Tensor: + history_short, history_mid, history_long = histories + current_indices, short_indices, mid_indices, long_indices = indices + with set_forward_context( + current_timestep=int(timestep[0].item()), + attn_metadata=None, + forward_batch=batch, + ): + return self.transformer( + hidden_states=latents.to(target_dtype), + timestep=timestep, + encoder_hidden_states=prompt_embeds, + indices_hidden_states=current_indices, + indices_latents_history_short=short_indices, + indices_latents_history_mid=mid_indices, + indices_latents_history_long=long_indices, + latents_history_short=history_short.to(target_dtype), + latents_history_mid=history_mid.to(target_dtype), + latents_history_long=history_long.to(target_dtype), + ) + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + self._load_or_move_transformer(fastvideo_args) + device = get_local_torch_device() + target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision] + prompt_embeds = batch.prompt_embeds[0].to(device=device, dtype=target_dtype) + negative_prompt_embeds = None + if batch.do_classifier_free_guidance: + if not batch.negative_prompt_embeds: + raise ValueError("Helios CFG requires negative prompt embeddings") + negative_prompt_embeds = batch.negative_prompt_embeds[0].to( + device=device, + dtype=target_dtype, + ) + + batch_size = prompt_embeds.shape[0] + assert isinstance(batch.height, int) and isinstance(batch.width, int) + assert isinstance(batch.num_frames, int) + assert batch.generator is not None + assert batch.history_sizes is not None + assert batch.pyramid_num_inference_steps_list is not None + + vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config + spatial_scale = vae_arch.scale_factor_spatial + temporal_scale = vae_arch.scale_factor_temporal + latent_height = batch.height // spatial_scale + latent_width = batch.width // spatial_scale + num_channels = self.transformer.in_channels + chunk_size = batch.num_latent_frames_per_chunk + num_chunks = get_num_latent_chunks( + batch.num_frames, + chunk_size, + temporal_scale, + ) + history_sizes = list(batch.history_sizes) + history_frame_count = sum(history_sizes) + history_latents = torch.zeros( + batch_size, + num_channels, + history_frame_count, + latent_height, + latent_width, + device=device, + dtype=torch.float32, + ) + indices = build_helios_frame_indices( + history_sizes, + chunk_size, + batch.keep_first_frame, + device, + ) + latent_chunks: list[torch.Tensor] = [] + first_frame_latent: torch.Tensor | None = None + patch_size = tuple(self.transformer.patch_size) + num_stages = len(batch.pyramid_num_inference_steps_list) + global_step_index = 0 + + logger.info( + "Helios sampling %d chunk(s), %d latent frames each, pyramid steps=%s", + num_chunks, + chunk_size, + batch.pyramid_num_inference_steps_list, + ) + for chunk_index in range(num_chunks): + history_long, history_mid, history_one = history_latents[:, :, -history_frame_count:].split(history_sizes, + dim=2) + if first_frame_latent is None: + prefix = torch.zeros( + batch_size, + num_channels, + 1, + latent_height, + latent_width, + device=device, + dtype=history_one.dtype, + ) + else: + prefix = first_frame_latent + history_short = torch.cat([prefix, history_one], dim=2) + histories = (history_short, history_mid, history_long) + + latents = _randn_tensor( + ( + batch_size, + num_channels, + chunk_size, + latent_height, + latent_width, + ), + generator=batch.generator, + device=device, + ) + latents = downsample_to_pyramid_base(latents, num_stages) + start_points = [latents] + + for stage_index, stage_steps in enumerate(batch.pyramid_num_inference_steps_list): + image_seq_len = math.prod(latents.shape[-3:]) // math.prod(patch_size) + mu = calculate_shift( + image_seq_len, + self._scheduler_value(self.scheduler, "base_image_seq_len", 256), + self._scheduler_value(self.scheduler, "max_image_seq_len", 4096), + self._scheduler_value(self.scheduler, "base_shift", 0.5), + self._scheduler_value(self.scheduler, "max_shift", 1.15), + ) + self.scheduler.set_timesteps( + stage_steps, + stage_index, + device=device, + mu=mu, + is_amplify_first_chunk=(batch.is_amplify_first_chunk and chunk_index == 0), + ) + timesteps = self.scheduler.timesteps + + if stage_index > 0: + batch_count, channels, frames, height, width = latents.shape + flattened = latents.permute(0, 2, 1, 3, 4).reshape( + batch_count * frames, + channels, + height, + width, + ) + flattened = F.interpolate( + flattened, + size=(height * 2, width * 2), + mode="nearest", + ) + latents = flattened.reshape( + batch_count, + frames, + channels, + height * 2, + width * 2, + ).permute(0, 2, 1, 3, 4) + + original_signal = 1 - self.scheduler.ori_start_sigmas[stage_index] + gamma = self.scheduler.config.gamma + alpha = 1 / (math.sqrt(1 + 1 / gamma) * (1 - original_signal) + original_signal) + beta = alpha * (1 - original_signal) / math.sqrt(gamma) + noise = sample_block_noise( + self.scheduler, + tuple(latents.shape), + patch_size, + device, + batch.generator, + ).to(dtype=target_dtype) + latents = alpha * latents + beta * noise + start_points.append(latents) + + for step_index, timestep_value in enumerate(timesteps): + timestep = timestep_value.expand(batch_size).to(torch.int64) + with trace_step(global_step_index): + noise_pred = self._transformer_forward( + latents, + timestep, + prompt_embeds, + histories, + indices, + batch, + target_dtype, + ) + if batch.do_classifier_free_guidance: + assert negative_prompt_embeds is not None + noise_uncond = self._transformer_forward( + latents, + timestep, + negative_prompt_embeds, + histories, + indices, + batch, + target_dtype, + ) + noise_pred = noise_uncond + batch.guidance_scale * (noise_pred - noise_uncond) + + latents = self.scheduler.step( + noise_pred, + timestep_value, + latents, + generator=(batch.generator[0] if isinstance(batch.generator, list) else batch.generator), + return_dict=False, + cur_sampling_step=step_index, + dmd_noisy_tensor=start_points[stage_index], + dmd_sigmas=self.scheduler.sigmas, + dmd_timesteps=self.scheduler.timesteps, + all_timesteps=timesteps, + )[0] + global_step_index += 1 + + if first_frame_latent is None: + first_frame_latent = latents[:, :, :1] + history_latents = torch.cat([history_latents, latents], dim=2) + latent_chunks.append(latents) + logger.info("Helios completed latent chunk %d/%d", chunk_index + 1, num_chunks) + + batch.helios_latent_chunks = latent_chunks + batch.latents = torch.cat(latent_chunks, dim=2) + batch.timesteps = self.scheduler.timesteps + + if fastvideo_args.dit_layerwise_offload: + manager = getattr(self.transformer, "_layerwise_offload_manager", None) + if manager is not None and getattr(manager, "enabled", False): + manager.release_all() + elif (fastvideo_args.dit_cpu_offload and not fastvideo_args.use_fsdp_inference + and next(self.transformer.parameters()).device.type == "cuda"): + self.transformer.to("cpu") + return batch + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + result.add_check("generator", batch.generator, V.generator_or_list_generators) + return result + + def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) + return result + + +class HeliosChunkDecodingStage(DecodingStage): + """Decode each 9-latent chunk independently, matching official Helios.""" + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if fastvideo_args.output_type == "latent": + assert batch.latents is not None + batch.output = batch.latents.detach().to(dtype=torch.float32, device="cpu") + batch.latents = None + batch.helios_latent_chunks = None + return batch + + pipeline = self.pipeline() if self.pipeline else None + if not fastvideo_args.model_loaded["vae"]: + self.vae = VAELoader().load(fastvideo_args.model_paths["vae"], fastvideo_args) + if pipeline is not None: + pipeline.add_module("vae", self.vae) + fastvideo_args.model_loaded["vae"] = True + + latent_chunks = batch.helios_latent_chunks + if not isinstance(latent_chunks, list) or not latent_chunks: + raise ValueError("Helios decoding requires non-empty helios_latent_chunks") + decoded_chunks = [self.decode(chunk, fastvideo_args) for chunk in latent_chunks] + frames = torch.cat(decoded_chunks, dim=2) + + temporal_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_temporal + assert isinstance(batch.num_frames, int) + generated_frames = min( + batch.num_frames, + get_generated_pixel_frames(frames.shape[2], temporal_scale), + ) + batch.output = frames[:, :, :generated_frames].detach().to(dtype=torch.float32, device="cpu") + batch.latents = None + batch.helios_latent_chunks = None + + if fastvideo_args.vae_cpu_offload: + self.vae.to("cpu") + return batch + + +__all__ = [ + "HeliosChunkDecodingStage", + "HeliosInputValidationStage", + "HeliosPyramidDenoisingStage", + "build_helios_frame_indices", + "calculate_shift", + "downsample_to_pyramid_base", + "get_num_latent_chunks", + "get_generated_pixel_frames", + "sample_block_noise", +] diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 3a9ca14675..1f04008627 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -132,6 +132,7 @@ class ForwardBatch: raw_latent_shape: tuple[int, ...] | None = None noise_pred: torch.Tensor | None = None image_latent: torch.Tensor | None = None + helios_latent_chunks: list[torch.Tensor] | None = None # Action control inputs (Matrix-Game) mouse_cond: torch.Tensor | None = None # Shape: (B, T, 2) @@ -189,6 +190,16 @@ class ForwardBatch: eta: float = 0.0 sigmas: list[float] | None = None + # Helios autoregressive spatial-pyramid sampling. + pyramid_num_inference_steps_list: list[int] | None = None + history_sizes: list[int] | None = None + num_latent_frames_per_chunk: int = 9 + keep_first_frame: bool = True + is_skip_first_chunk: bool = False + use_zero_init: bool = True + zero_steps: int = 1 + is_amplify_first_chunk: bool = False + # TeaCache enable_teacache: bool = False diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 0837329014..f573536251 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -24,6 +24,7 @@ from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig from fastvideo.configs.pipelines.gen3c import Gen3CConfig +from fastvideo.configs.pipelines.helios import HeliosPipelineConfig from fastvideo.configs.pipelines.hunyuan15 import (Hunyuan15T2V480PConfig, Hunyuan15I2V480PStepDistilledConfig, Hunyuan15T2V720PConfig, Hunyuan15I2V720PConfig, Hunyuan15SR1080PConfig) @@ -216,6 +217,13 @@ def _get_config_info( config = maybe_download_model_index(model_path, revision=revision) pipeline_name = config.get("_class_name", "").lower() + scheduler = config.get("scheduler") + if (pipeline_name == "heliospyramidpipeline" and config.get("is_distilled") is True and isinstance(scheduler, list) + and len(scheduler) >= 2 and scheduler[1] == "HeliosDMDScheduler"): + helios_model_id = _MODEL_HF_PATH_TO_NAME.get("BestWishYsh/Helios-Distilled") + if helios_model_id is not None: + logger.debug("Resolved Helios-Distilled from authoritative model index metadata.") + return _CONFIG_REGISTRY.get(helios_model_id) matched_model_names: list[str] = [] for model_id, detector in _MODEL_NAME_DETECTORS: @@ -240,6 +248,17 @@ def _get_config_info( def _register_configs() -> None: + register_configs( + sampling_param_cls=None, + pipeline_config_cls=HeliosPipelineConfig, + workload_types=(WorkloadType.T2V, ), + hf_model_paths=["BestWishYsh/Helios-Distilled"], + model_detectors=[], + model_family="helios", + default_preset="helios_distilled_t2v", + pipeline_cls_name="HeliosPyramidPipeline", + ) + # MMAudio large-44k-v2 (video/text-to-audio). The checkpoint is converted # into standard per-component FastVideo/Diffusers-style directories by # scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py. @@ -1306,6 +1325,8 @@ def _register_presets() -> None: ALL_PRESETS as HUNYUAN_PRESETS, ) from fastvideo.pipelines.basic.hunyuan15.presets import ( ALL_PRESETS as HUNYUAN15_PRESETS, ) + from fastvideo.pipelines.basic.helios.presets import ( + ALL_PRESETS as HELIOS_PRESETS, ) from fastvideo.pipelines.basic.hyworld.presets import ( ALL_PRESETS as HYWORLD_PRESETS, ) from fastvideo.pipelines.basic.kandinsky5.presets import ( @@ -1349,6 +1370,7 @@ def _register_presets() -> None: GEN3C_PRESETS, HUNYUAN_PRESETS, HUNYUAN15_PRESETS, + HELIOS_PRESETS, HYWORLD_PRESETS, KANDINSKY5_PRESETS, LINGBOT_VIDEO_PRESETS, diff --git a/fastvideo/tests/api/test_parser.py b/fastvideo/tests/api/test_parser.py index e5370cc7bf..1caa1b406a 100644 --- a/fastvideo/tests/api/test_parser.py +++ b/fastvideo/tests/api/test_parser.py @@ -201,6 +201,14 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None: "use_embedded_guidance": None, "boundary_ratio": None, "sigmas": None, + "pyramid_num_inference_steps_list": None, + "history_sizes": None, + "num_latent_frames_per_chunk": 9, + "keep_first_frame": True, + "is_skip_first_chunk": False, + "use_zero_init": True, + "zero_steps": 1, + "is_amplify_first_chunk": False, }, "runtime": { "enable_teacache": False, diff --git a/tests/local_tests/encoders/test_helios_umt5_parity.py b/tests/local_tests/encoders/test_helios_umt5_parity.py new file mode 100644 index 0000000000..a24b4881d0 --- /dev/null +++ b/tests/local_tests/encoders/test_helios_umt5_parity.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Helios tokenizer and exact-checkpoint UMT5-XXL reuse parity. + +Coverage scope: both. The tokenizer test exercises FastVideo's production +third-party loader boundary. The encoder test runs the official and native +implementations sequentially so two UMT5-XXL copies do not occupy one GPU. +""" + +from __future__ import annotations + +import gc +import json +from pathlib import Path +import subprocess +import sys + +import pytest +import torch +from torch.testing import assert_close +from transformers import AutoTokenizer, UMT5EncoderModel as OfficialUMT5EncoderModel + +from fastvideo.configs.models.encoders import T5Config +from fastvideo.configs.models.encoders.t5 import T5ArchConfig +from fastvideo.models.encoders.t5 import UMT5EncoderModel +from fastvideo.models.loader.weight_utils import ( + resolve_safetensors_files, + safetensors_weights_iterator, +) + +REPO_ROOT = Path(__file__).resolve().parents[3] +MODEL_DIR = REPO_ROOT / "official_weights" / "helios" +TEXT_ENCODER_DIR = MODEL_DIR / "text_encoder" +TOKENIZER_DIR = MODEL_DIR / "tokenizer" +PARITY_SCOPE = "both" + + +def _require_assets() -> None: + required = ( + TEXT_ENCODER_DIR / "model.safetensors.index.json", + TOKENIZER_DIR / "tokenizer.json", + ) + missing = [str(path) for path in required if not path.exists()] + if missing: + pytest.skip(f"Helios text assets missing: {missing}") + + +def _native_config() -> T5Config: + return T5Config( + arch_config=T5ArchConfig( + architectures=["UMT5EncoderModel"], + vocab_size=256384, + d_model=4096, + d_kv=64, + d_ff=10240, + num_layers=24, + num_decoder_layers=24, + num_heads=64, + relative_attention_num_buckets=32, + relative_attention_max_distance=128, + dropout_rate=0.1, + layer_norm_epsilon=1e-6, + feed_forward_proj="gated-gelu", + is_encoder_decoder=True, + use_cache=True, + text_len=512, + ), + prefix="umt5", + ) + + +def _patch_single_process_text_parallel(monkeypatch) -> None: + import fastvideo.layers.linear as fastvideo_linear + import fastvideo.layers.vocab_parallel_embedding as fastvideo_embedding + import fastvideo.models.encoders.t5 as fastvideo_t5 + + for module in (fastvideo_t5, fastvideo_embedding, fastvideo_linear): + if hasattr(module, "get_tp_rank"): + monkeypatch.setattr(module, "get_tp_rank", lambda: 0) + if hasattr(module, "get_tp_world_size"): + monkeypatch.setattr(module, "get_tp_world_size", lambda: 1) + monkeypatch.setattr(fastvideo_embedding, "tensor_model_parallel_all_reduce", lambda value: value) + + +def test_helios_tokenizer_loader_matches_official_assets() -> None: + _require_assets() + official = AutoTokenizer.from_pretrained(str(TOKENIZER_DIR), local_files_only=True) + prompts = [ + "A glass sculpture turning slowly in a quiet studio.", + "海面上缓慢升起的清晨薄雾。", + ] + kwargs = { + "padding": "max_length", + "truncation": True, + "max_length": 32, + "return_tensors": "pt", + } + official_batch = official(prompts, **kwargs) + script = f""" +import json +from types import SimpleNamespace +from fastvideo.models.loader.component_loader import TokenizerLoader + +args = SimpleNamespace( + pipeline_config=SimpleNamespace(text_encoder_configs=()), + trust_remote_code=False, +) +tokenizer = TokenizerLoader().load({str(TOKENIZER_DIR)!r}, args) +batch = tokenizer({prompts!r}, padding='max_length', truncation=True, + max_length=32, return_tensors='pt') +print(json.dumps({{ + 'input_ids': batch.input_ids.tolist(), + 'attention_mask': batch.attention_mask.tolist(), +}})) +""" + result = subprocess.run( + [sys.executable, "-c", script], + cwd=REPO_ROOT, + check=True, + capture_output=True, + text=True, + ) + fastvideo_batch = json.loads(result.stdout.strip().splitlines()[-1]) + assert fastvideo_batch["input_ids"] == official_batch.input_ids.tolist() + assert fastvideo_batch["attention_mask"] == official_batch.attention_mask.tolist() + + +def test_helios_text_encoder_config_matches_umt5_xxl() -> None: + config = _native_config() + assert config.vocab_size == 256384 + assert config.d_model == 4096 + assert config.d_kv == 64 + assert config.d_ff == 10240 + assert config.num_layers == config.num_decoder_layers == 24 + assert config.num_heads == 64 + assert config.feed_forward_proj == "gated-gelu" + assert config.text_len == 512 + + +def test_helios_pipeline_config_reuses_verified_components() -> None: + try: + from fastvideo.configs.pipelines.helios import HeliosPipelineConfig + except ImportError as exc: + raise AssertionError("HeliosPipelineConfig has not been implemented yet") from exc + + config = HeliosPipelineConfig() + assert config.dit_config.__class__.__name__ == "HeliosConfig" + assert config.vae_config.__class__.__name__ == "WanVAEConfig" + assert config.vae_config.load_encoder is False + assert config.vae_config.load_decoder is True + assert config.text_encoder_precisions == ("bf16", ) + assert config.dit_precision == "bf16" + assert config.vae_decode_precision == "fp32" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for UMT5 parity.") +@pytest.mark.parametrize( + ("dtype", "atol", "rtol", "mean_limit"), + [ + pytest.param(torch.float32, 1e-4, 1e-4, 1e-5, id="fp32-math"), + pytest.param(torch.bfloat16, 3e-2, 3e-2, 3e-3, id="bf16-runtime"), + ], +) +def test_helios_umt5_hidden_state_parity( + monkeypatch, + dtype: torch.dtype, + atol: float, + rtol: float, + mean_limit: float, +) -> None: + _require_assets() + device = torch.device("cuda:0") + tokenizer = AutoTokenizer.from_pretrained(str(TOKENIZER_DIR), local_files_only=True) + batch = tokenizer( + ["A precise macro shot of frost forming on a red leaf."], + padding="max_length", + truncation=True, + max_length=32, + return_tensors="pt", + ) + input_ids = batch.input_ids.to(device) + attention_mask = batch.attention_mask.to(device) + + official = (OfficialUMT5EncoderModel.from_pretrained(str(TEXT_ENCODER_DIR), + local_files_only=True, + torch_dtype=dtype).to(device).eval()) + with torch.inference_mode(): + official_hidden = official(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state.float().cpu() + del official + gc.collect() + torch.cuda.empty_cache() + + _patch_single_process_text_parallel(monkeypatch) + native = UMT5EncoderModel(_native_config()).to(device=device, dtype=dtype) + files = resolve_safetensors_files(str(TEXT_ENCODER_DIR)) + loaded = native.load_weights(safetensors_weights_iterator(files, to_cpu=True)) + missing = {name for name, _ in native.named_parameters()} - loaded + assert missing == set() + native.eval() + with torch.inference_mode(): + fastvideo_hidden = native(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state.float().cpu() + + assert official_hidden.shape == fastvideo_hidden.shape + diff = (official_hidden - fastvideo_hidden).abs() + print(f"UMT5 dtype={dtype} diff_max={diff.max().item():.8f} diff_mean={diff.mean().item():.8f}") + assert diff.mean().item() < mean_limit + assert_close(fastvideo_hidden, official_hidden, atol=atol, rtol=rtol) diff --git a/tests/local_tests/helios/PORT_STATUS.md b/tests/local_tests/helios/PORT_STATUS.md new file mode 100644 index 0000000000..7b90e36662 --- /dev/null +++ b/tests/local_tests/helios/PORT_STATUS.md @@ -0,0 +1,110 @@ +# Helios Port Status + +## Summary + +- model_family: `helios` +- workload_types: `T2V` +- official_ref: Diffusers `0.39.0` `HeliosPyramidPipeline`; `PKU-YuanGroup/Helios@8f2a2faa` +- official_ref_dir: `../Helios` +- hf_weights_path: `BestWishYsh/Helios-Distilled` +- local_weights_dir: `official_weights/helios` +- source_layout: `diffusers` +- local_tests_readme: `tests/local_tests/helios/README.md` + +## Current Phase + +- phase: `final_verification` +- status: `complete` +- owner: `orchestrator` +- last_updated: `2026-08-26` + +## Component Matrix + +| Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | +| transformer | dit | ported | Diffusers `HeliosTransformer3DModel` | Exact HF transformer config; short/mid/long histories and frame indices | native Helios DiT/config | complete | identity 1101/1101 | non_skip_pass | none | +| scheduler | scheduler | ported | Diffusers `HeliosDMDScheduler` | Three stages, dynamic shift, stage-local re-noising | native Helios DMD scheduler | complete | not_needed | non_skip_pass, bit-exact | none | +| vae | vae | reused | Diffusers `AutoencoderKLWan` | Exact Helios VAE config/weights | native Wan VAE | complete | not_needed | non_skip_pass, decode exact | none | +| text_encoder | encoder | reused | Transformers `UMT5EncoderModel` | UMT5-XXL, max length 512 | native UMT5 | complete | not_needed | non_skip_pass FP32/BF16 | none | +| tokenizer | encoder | passthrough | `T5TokenizerFast` | Exact checkpoint assets | production tokenizer loader | complete | not_needed | exact IDs/masks | none | +| pipeline | pipeline | ported | Diffusers `HeliosPyramidPipeline` | 9-frame AR chunks, `[16,2,1]` history, 3 spatial stages | `HeliosPyramidPipeline` | complete | not_needed | non_skip_pass | T2V Distilled only | + +## Conversion State + +- conversion_script: `not_needed` +- converted_weights_dir: `not_needed` +- source_layout: `diffusers` +- strict_load_status: transformer 1101/1101; VAE strict; UMT5 all parameters loaded +- passthrough_components: tokenizer assets +- retry_history: none + +## Parity Commands + +| Scope | Command | Last Result | Notes | +| --- | --- | --- | --- | +| transformer | `pytest tests/local_tests/transformers/test_helios_transformer_parity.py -v -s` | current PR component evidence: non-skip PASS | strict load, tiny/full, SP=2, FA2/SDPA | +| scheduler | `pytest tests/local_tests/schedulers/test_helios_dmd_scheduler_parity.py -v -s` | 10 passed | registry plus bit-exact schedules/steps | +| VAE | `pytest tests/local_tests/vaes/test_helios_vae_parity.py -v -s` | 2 passed | decode diff max/mean 0/0 | +| UMT5/tokenizer | `pytest tests/local_tests/encoders/test_helios_umt5_parity.py -v -s` | 5 passed | exact tokenizer; FP32/BF16 encoder parity | +| pipeline math/stage/smoke | three `test_helios_pipeline_*` files | 29 passed | includes CUDA block noise and CPU-output regression | +| pipeline parity | `pytest tests/local_tests/pipelines/test_helios_pipeline_parity.py -v -s` | 1 passed in 82.85 s | cosine 0.979703, MAE 0.170017, RMSE 0.246508, drift 0.376% | +| typed API regression | smoke + parser/compat/config tests | 42 passed | registry, preset, class resolution, typed/CLI fields | +| typed example | `python examples/inference/basic/basic_helios_distilled_t2v.py` | PASS, generation 29.43 s | H.264 640×384, 33 frames, full decode | +| quality/container | `HELIOS_QUALITY_CANDIDATE=... pytest test_helios_quality_regression.py` | 1 passed, 1 skipped | reference comparison deferred pending upload approval | + +## Open Questions + +| ID | Question | Owner | Needed By Phase | Status | Resolution | +| --- | --- | --- | --- | --- | --- | +| Q001 | Is `transformer_ode` required for Distilled T2V? | orchestrator | Phase 1 | resolved | No model index or official T2V call loads it; excluded. | +| Q002 | Can the native Wan VAE be reused? | component:vae | Phase 3 | resolved | Exact config/weight strict load and decode parity are bit-exact. | +| Q003 | Can native UMT5/tokenizer be reused? | component:encoder | Phase 3 | resolved | Token IDs/masks exact; FP32 and BF16 output parity pass. | +| Q004 | Where do Helios DMD runtime fields live? | component:scheduler | Phase 1 | resolved | Scheduler owns sigma/timestep math; public per-call knobs live in SamplingParam/preset. | + +## Issues And Blockers + +| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | +| I001 | prep | environment | medium | No project venv initially. | system Python lacked ML deps | prep | resolved | Python 3.12 project venv installed. | +| I002 | prep | repository | low | Early branch was behind main. | rev-list showed upstream commits | prep | resolved | Final integration branch rebased on `a159b63c`. | +| I003 | prep | weights | low | Snapshot includes unused 58 GB ODE transformer. | model indexes omit `transformer_ode` | orchestrator | resolved | Pinned inference assets downloaded without ODE directory. | +| I004 | parity | tokenizer | low | Direct loader import had a repository-root import cycle in the older workspace. | collection traceback | orchestrator | resolved_for_test | Production loader runs in a clean subprocess. | +| I005 | official_reference | CUDA accounting | low | Peak reset initially used an invalid device after offload hooks. | runtime error before inference | orchestrator | resolved | Set current device before hooks; parity completes. | +| I006 | transformer | RoPE buffers | medium | Meta construction left non-persistent RoPE buffers on meta. | meta/cuda einsum mismatch | component:transformer | resolved | Loader materializes exact buffers. | +| I007 | encoder | BF16 drift | low | Fused-QKV and three-GEMM paths differ slightly. | max 0.0234375, mean 0.00160135 | component:encoder | resolved | FP32 proves math; separate bounded BF16 gate passes. | +| I008 | pipeline | history geometry | high | Coarse-stage history RoPE used current-grid positions. | 294 history tokens versus 72 positions | component:transformer | resolved | Short-grid positions are center-downsampled for mid/long histories. | +| I009 | pipeline | guidance cross-attention | medium | History and current queries were combined before masking. | activation trace diverged | component:transformer | resolved | Current/history split occurs before cross-attention. | +| I010 | tests | repository package root | low | A root `__init__.py` makes invalid worktree basenames fail mypy. | `helios-pr... is not a valid Python package name` | orchestrator | resolved | Worktree basename changed to `helios_pr1670_full_pipeline`. | +| I011 | registry | variant safety | high | Broad Helios detection would route Base/Mid to Distilled. | negative metadata probes | orchestrator | resolved | Require pipeline class, `is_distilled=true`, and Helios DMD scheduler. | +| I012 | quality | tiny smoke | medium | 128×192 is not a visual-quality target. | poor-detail local smoke | orchestrator | resolved_as_scope | Public example and integrity gate use 384×640. | +| I013 | distributed | nightly coverage | medium | SP/FSDP/repeated-run coverage is not a committed CI lane. | earlier local SP/FSDP smoke passed | orchestrator | open_nightly | Keep as follow-up; no absent runner is referenced by this PR. | +| I014 | production_validation | multiprocessing output | high | Returning GPU output through CUDA IPC OOMed while the 14.31B DiT remained resident. | v1 completed stages then failed in `_new_shared_cuda` | pipeline | resolved | Decode/latent outputs move to CPU inside worker; regression is GREEN and typed example v2 passes. | +| I015 | final_verification | upstream main | high | Integration base `6388db81` had eight unrelated unit-lane failures. | 909 passed and 8 failed before upstream CI/schema fixes. | upstream | resolved | Rebased through `b2062556` onto `a159b63c`; the current shared unit script passes all 1047 tests. | + +## Escape Hatches + +No escape hatch is open. CI SSIM reference upload remains deliberately +unauthorized and is recorded as a quality-regression deferral, not a blocker to +local pipeline parity. + +## Decisions + +| Date | Decision | Rationale | Impact | +| --- | --- | --- | --- | +| 2026-07-11 | First public scope is Helios-Distilled T2V. | Keep a reviewable, verifiable variant. | Base/Mid, ODE, training and conditioned modes remain out of scope. | +| 2026-07-11 | Use Diffusers 0.39.0 as executable parity reference. | It matches the published Diffusers-layout checkpoint. | No conflicting Helios research requirements are installed. | +| 2026-07-11 | Preserve 1101 transformer keys directly. | Official and native key surfaces match. | No conversion script. | +| 2026-07-11 | Reuse Wan VAE and UMT5 only after exact-asset parity. | Architecture resemblance is insufficient. | Both reused components have non-skip evidence. | +| 2026-08-25 | Use a dedicated Helios pyramid stage inside `ComposedPipelineBase`. | Generic denoising cannot express AR history, three spatial levels and stage-local DMD. | One FastVideo architecture, model-specific stages only where required. | +| 2026-08-25 | Keep zero-init call fields but do not apply zero-star math for the pinned Distilled checkpoint. | Its model index declares `is_cfg_zero_star=false`; Diffusers also takes the standard CFG branch. | Signature stays compatible without claiming an inactive feature changes output. | +| 2026-08-25 | Move worker output to CPU before multiprocessing return. | Prevent CUDA IPC allocation after high-memory inference. | Public typed example is robust on 48 GB cards. | + +## Handoff Notes + +- Required components and pipeline parity are green on the integration branch. +- `quality_regression=deferred_with_reason`: local real-video integrity passes; + publishing a CI reference needs separate approval. +- Repository-wide pre-commit is green. The current shared unit script passes + all 1047 collected tests; `I015` remains resolved on `a159b63c`. +- No weights, generated media, reference clone, private report, token, push, or + PR mutation is part of this state file. diff --git a/tests/local_tests/helios/README.md b/tests/local_tests/helios/README.md new file mode 100644 index 0000000000..b0d1909269 --- /dev/null +++ b/tests/local_tests/helios/README.md @@ -0,0 +1,175 @@ +# Helios Local Verification + +Reviewer-facing setup and verification record for the native +`BestWishYsh/Helios-Distilled` text-to-video port. These tests compare FastVideo +against the exact Diffusers Helios implementation and require local assets for +non-skip author verification. + +Port state and resolved issues live in +`tests/local_tests/helios/PORT_STATUS.md`. + +## Pinned Sources + +| Field | Value | +| --- | --- | +| Model family | `helios` | +| Public scope | Helios-Distilled T2V | +| Architecture source | `PKU-YuanGroup/Helios@8f2a2faab3298c8a7630a2c73aea37c01b5bab01` | +| Executable parity reference | Diffusers `0.39.0` `HeliosPyramidPipeline` | +| HF checkpoint | `BestWishYsh/Helios-Distilled` | +| HF revision | `1999182614cb08d3bdcc46b9827504af2914b87b` | +| Local assets | `official_weights/helios` | +| Source layout | Diffusers; no conversion required | +| Integration base | `upstream/main@a159b63c67a1a283ce55813b694524909ea67b15` | + +The official T2V model index loads one `transformer`; `transformer_ode` is not +declared or used and remains outside this PR. Base/Mid, training, ODE, I2V, and +other variants are not supported by this first pipeline. + +## Environment And Assets + +Run from the FastVideo repository root: + +```bash +uv venv --python 3.12 --seed +UV_TORCH_BACKEND=cu130 uv pip install -e ".[dev]" + +.venv/bin/python -c "from diffusers import AutoencoderKLWan, HeliosDMDScheduler, HeliosPyramidPipeline, HeliosTransformer3DModel; from transformers import AutoTokenizer, UMT5EncoderModel; print('imports ok')" +``` + +The verified environment uses Python 3.12.3, PyTorch `2.12.0+cu130`, CUDA 13, +Diffusers 0.39.0, Transformers 5.13.1, and `flash-attn==2.8.1`. + +Download the pinned inference snapshot without the unused ODE transformer: + +```bash +.venv/bin/python \ + .agents/skills/add-model-01-prep/scripts/download_hf_weights.py \ + BestWishYsh/Helios-Distilled \ + official_weights/helios \ + --revision 1999182614cb08d3bdcc46b9827504af2914b87b \ + --ignore-pattern "transformer_ode/*" +``` + +Use only `HF_TOKEN`, `HUGGINGFACE_HUB_TOKEN`, or `HF_API_KEY` when auth is +needed. Never put a token value in this file or a command log. + +## Component Matrix + +| Component | FastVideo target | Test | Current author result | +| --- | --- | --- | --- | +| Transformer | `fastvideo/models/dits/helios.py` | `tests/local_tests/transformers/test_helios_transformer_parity.py` | Non-skip PASS: 1101/1101 strict load, tiny exact normal/pyramid, full BF16 normal/pyramid, SP=2, FlashAttention-vs-SDPA | +| Scheduler | `fastvideo/models/schedulers/scheduling_helios_dmd.py` | `tests/local_tests/schedulers/test_helios_dmd_scheduler_parity.py` | 10 passed: registry resolution plus bit-exact three-stage/amplify/step parity | +| VAE | Existing native Wan VAE | `tests/local_tests/vaes/test_helios_vae_parity.py` | Non-skip PASS; decode max/mean diff `0/0` | +| UMT5 | Existing native UMT5 | `tests/local_tests/encoders/test_helios_umt5_parity.py` | FP32 max/mean `1.25e-6/1.0e-7`; BF16 `0.0234375/0.00160135` | +| Tokenizer | Production tokenizer loader | `tests/local_tests/encoders/test_helios_umt5_parity.py` | Exact IDs and attention masks | +| Pipeline | `fastvideo/pipelines/basic/helios/` | `tests/local_tests/pipelines/test_helios_pipeline_parity.py` | Non-skip final-latent PASS | + +## Mandatory Commands + +Run component gates before pipeline gates: + +```bash +PYTHONPATH="$PWD" .venv/bin/pytest \ + tests/local_tests/transformers/test_helios_transformer_parity.py -v -s + +PYTHONPATH="$PWD" .venv/bin/pytest \ + tests/local_tests/vaes/test_helios_vae_parity.py -v -s + +PYTHONPATH="$PWD" .venv/bin/pytest \ + tests/local_tests/encoders/test_helios_umt5_parity.py -v -s + +PYTHONPATH="$PWD" .venv/bin/pytest \ + tests/local_tests/schedulers/test_helios_dmd_scheduler_parity.py -v -s + +PYTHONPATH="$PWD" .venv/bin/pytest \ + tests/local_tests/pipelines/test_helios_pipeline_math.py \ + tests/local_tests/pipelines/test_helios_pipeline_stages.py \ + tests/local_tests/pipelines/test_helios_pipeline_smoke.py -v -s + +DISABLE_SP=1 PYTHONPATH="$PWD" .venv/bin/pytest \ + tests/local_tests/pipelines/test_helios_pipeline_parity.py -v -s +``` + +## End-To-End Latent Parity + +Both implementations load the same checkpoint from scratch and use CPU seed +42, 128×192, 33 pixel frames, guidance 1.0, three `[2,2,2]` DMD stages, +history `[16,2,1]`, and nine latent frames per chunk. + +```text +shape [1, 16, 9, 16, 24] +official abs mean 0.95911187 +FastVideo abs mean 0.96272093 +abs-mean relative drift 0.3763% +max absolute diff 1.89453125 +mean absolute diff 0.17001711 +RMSE 0.24650776 +cosine 0.97970295 +result PASS +``` + +Fixed gates are shape equality, cosine ≥0.95, abs-mean drift ≤5%, mean +absolute error ≤0.30, and RMSE ≤0.40. The managed current-head run completed in +82.85 seconds with one passed test and no skip. + +## Typed Public Example + +```bash +HELIOS_MODEL_PATH=official_weights/helios \ +HELIOS_OUTPUT_PATH=outputs/helios/helios_distilled_t2v_384x640_33f.mp4 \ +PYTHONPATH="$PWD" .venv/bin/python \ + examples/inference/basic/basic_helios_distilled_t2v.py +``` + +The example uses `VideoGenerator.from_config`, `GeneratorConfig`, +`GenerationRequest`, and typed sampling/output configs. The current-head run +completed generation in 29.43 seconds and produced: + +```text +codec H.264 +resolution 640x384 +fps 24 +frames 33 +duration 1.375 s +SHA-256 2e6d5099715256b16b865cb300479b328b2835f9f8792cb5427a98c17eb0b038 +``` + +Full ffmpeg decode and the non-black-frame/container quality gate pass. The +contact sheet contains the requested fish, coral, and underwater scene. + +Run the committed local quality gate with: + +```bash +HELIOS_QUALITY_CANDIDATE=outputs/helios/helios_distilled_t2v_384x640_33f.mp4 \ + PYTHONPATH="$PWD" .venv/bin/pytest \ + tests/local_tests/helios/test_helios_quality_regression.py -v -s +``` + +Expected author result is one passed integrity test and one skipped reference +comparison. A CI SSIM reference is deliberately `deferred_with_reason`: its +publication requires separate upload approval, which this PR does not have. + +## Review Notes + +- Production code does not import Diffusers or Transformers model classes; + those imports exist only in parity tests. The existing tokenizer boundary is + reused. +- The dedicated pyramid stage is required by the official nine-frame + autoregressive chunks, three spatial levels, history geometry, block-noise + covariance, stage-local DMD schedules, and chunk decode contract. +- `use_zero_init` and `zero_steps` remain in the public call surface for + Diffusers signature compatibility. The pinned Distilled model declares + `is_cfg_zero_star=false`, so both official and FastVideo standard-CFG paths + intentionally ignore zero-star initialization for this checkpoint. +- Final decoded and latent outputs are moved to CPU inside the worker before + multiprocessing return. A regression test prevents CUDA-IPC OOM from + returning a GPU tensor while the 14.31B transformer remains resident. +- Generated media, tensor artifacts, weights, reference clones, and private + reports remain ignored and must not be committed. + +## Latest-Main Unit-Lane Baseline + +The branch is rebased onto `a159b63c`. The current +`.buildkite/scripts/unit_test.sh` completes with `1047 passed` and 21 warnings +in 28.70 seconds. `pre-commit run --all-files` also passes completely. diff --git a/tests/local_tests/helios/run_helios_latent_parity.py b/tests/local_tests/helios/run_helios_latent_parity.py new file mode 100644 index 0000000000..28f4e16b99 --- /dev/null +++ b/tests/local_tests/helios/run_helios_latent_parity.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Compare official and FastVideo Helios latents with identical CPU RNG.""" + +from __future__ import annotations + +import argparse +import gc +import json +from pathlib import Path +import re +import time + +import torch +from diffusers import AutoModel, HeliosPyramidPipeline as OfficialHeliosPipeline +import torch.nn.functional as F + +PROMPT = ( + "A vibrant tropical fish swims gracefully through a colorful coral reef " + "in clear turquoise water, cinematic close-up, fluid motion, vivid detail." +) +NEGATIVE_PROMPT = ( + "Bright tones, overexposed, static, blurred details, subtitles, paintings, " + "images, overall gray, worst quality, low quality, JPEG artifacts, ugly, " + "deformed, disfigured, still picture, messy background." +) + + +def parse_args() -> argparse.Namespace: + repo_root = Path(__file__).resolve().parents[3] + parser = argparse.ArgumentParser() + parser.add_argument( + "--model-dir", + type=Path, + default=repo_root / "official_weights" / "helios", + ) + parser.add_argument("--output-dir", type=Path, default=repo_root / "outputs") + parser.add_argument("--height", type=int, default=128) + parser.add_argument("--width", type=int, default=192) + parser.add_argument("--num-frames", type=int, default=33) + parser.add_argument("--steps", type=int, default=2) + parser.add_argument("--gpu", type=int, default=0) + parser.add_argument("--share-official-prompt-embeds", action="store_true") + parser.add_argument("--reuse-saved-official-artifacts", action="store_true") + parser.add_argument("--official-trace-output", type=Path) + return parser.parse_args() + + +def attach_official_trace(model, output_path: Path): + pattern = re.compile(r"^(blocks\.(0|9|19|29|39)|proj_out)$") + output_path.parent.mkdir(parents=True, exist_ok=True) + sink = output_path.open("w", encoding="utf-8") + call_counts: dict[str, int] = {} + handles = [] + + def make_hook(module_name: str): + + def hook(module, inputs, output): + del module, inputs + tensor = output[0] if isinstance(output, tuple | list) else output + if not isinstance(tensor, torch.Tensor): + return + step = call_counts.get(module_name, 0) + call_counts[module_name] = step + 1 + value = tensor.detach().float() + sink.write( + json.dumps({ + "module": module_name, + "tensor": "out", + "step": step, + "abs_mean": value.abs().mean().item(), + "mean": value.mean().item(), + "std": value.std().item(), + "max": value.max().item(), + "shape": list(value.shape), + "dtype": str(tensor.dtype), + }) + "\n") + + return hook + + for name, module in model.named_modules(): + if pattern.fullmatch(name): + handles.append(module.register_forward_hook(make_hook(name))) + if len(handles) != 6: + sink.close() + raise RuntimeError(f"Expected 6 official trace modules, attached {len(handles)}") + return handles, sink + + +def main() -> None: + args = parse_args() + args.output_dir.mkdir(parents=True, exist_ok=True) + torch.cuda.set_device(args.gpu) + + official_path = args.output_dir / "helios_official_cpu_seed42_latents.pt" + prompt_embeds_path = args.output_dir / "helios_official_prompt_embeds.pt" + if args.reuse_saved_official_artifacts: + official_latents = torch.load(official_path, map_location="cpu") + shared_prompt_embeds_cpu = (torch.load(prompt_embeds_path, map_location="cpu").detach() + if args.share_official_prompt_embeds else None) + official_seconds = 0.0 + else: + official_started = time.perf_counter() + vae = AutoModel.from_pretrained( + str(args.model_dir), + subfolder="vae", + local_files_only=True, + torch_dtype=torch.float32, + ) + official_pipeline = OfficialHeliosPipeline.from_pretrained( + str(args.model_dir), + vae=vae, + local_files_only=True, + torch_dtype=torch.bfloat16, + ) + official_pipeline.enable_model_cpu_offload(gpu_id=args.gpu) + trace_handles = [] + trace_sink = None + if args.official_trace_output is not None: + trace_handles, trace_sink = attach_official_trace( + official_pipeline.transformer, + args.official_trace_output, + ) + shared_prompt_embeds = None + if args.share_official_prompt_embeds: + shared_prompt_embeds, _ = official_pipeline.encode_prompt( + prompt=PROMPT, + negative_prompt=None, + do_classifier_free_guidance=False, + num_videos_per_prompt=1, + max_sequence_length=512, + dtype=torch.bfloat16, + ) + try: + official_latents = (official_pipeline( + prompt=None if shared_prompt_embeds is not None else PROMPT, + negative_prompt=(None if shared_prompt_embeds is not None else NEGATIVE_PROMPT), + prompt_embeds=shared_prompt_embeds, + height=args.height, + width=args.width, + num_frames=args.num_frames, + pyramid_num_inference_steps_list=[args.steps] * 3, + guidance_scale=1.0, + is_amplify_first_chunk=True, + generator=torch.Generator("cpu").manual_seed(42), + output_type="latent", + ).frames.float().cpu()) + finally: + for handle in trace_handles: + handle.remove() + if trace_sink is not None: + trace_sink.close() + official_seconds = time.perf_counter() - official_started + torch.save(official_latents, official_path) + shared_prompt_embeds_cpu = shared_prompt_embeds.detach().cpu() if shared_prompt_embeds is not None else None + if shared_prompt_embeds_cpu is not None: + torch.save(shared_prompt_embeds_cpu, prompt_embeds_path) + + del official_pipeline, vae + gc.collect() + torch.cuda.empty_cache() + + from fastvideo import VideoGenerator + + fastvideo_load_started = time.perf_counter() + generator = VideoGenerator.from_pretrained( + str(args.model_dir), + num_gpus=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + dit_layerwise_offload=True, + text_encoder_cpu_offload=True, + vae_cpu_offload=True, + pin_cpu_memory=False, + enable_stage_verification=True, + output_type="latent", + ) + fastvideo_load_seconds = time.perf_counter() - fastvideo_load_started + fastvideo_started = time.perf_counter() + try: + generation_kwargs = { + "negative_prompt": NEGATIVE_PROMPT, + "output_path": str(args.output_dir), + "save_video": False, + "return_frames": True, + "height": args.height, + "width": args.width, + "num_frames": args.num_frames, + "fps": 24, + "num_inference_steps": args.steps, + "pyramid_num_inference_steps_list": [args.steps] * 3, + "history_sizes": [16, 2, 1], + "num_latent_frames_per_chunk": 9, + "keep_first_frame": True, + "is_skip_first_chunk": False, + "use_zero_init": True, + "zero_steps": 1, + "guidance_scale": 1.0, + "is_amplify_first_chunk": True, + "seed": 42, + } + if shared_prompt_embeds_cpu is None: + result = generator.generate_video(prompt=PROMPT, **generation_kwargs) + else: + result = generator._generate_video_impl( + prompt=PROMPT, + sampling_param=None, + fastvideo_args=generator.fastvideo_args, + prompt_embeds=[shared_prompt_embeds_cpu], + **generation_kwargs, + ) + finally: + generator.shutdown() + fastvideo_seconds = time.perf_counter() - fastvideo_started + fastvideo_latents = result["samples"].float().cpu() + fastvideo_path = args.output_dir / "helios_fastvideo_cpu_seed42_latents.pt" + torch.save(fastvideo_latents, fastvideo_path) + + if official_latents.shape != fastvideo_latents.shape: + raise AssertionError( + f"Latent shape mismatch: official={official_latents.shape}, FastVideo={fastvideo_latents.shape}") + difference = official_latents - fastvideo_latents + summary = { + "shape": + list(official_latents.shape), + "official_path": + str(official_path), + "fastvideo_path": + str(fastvideo_path), + "official_abs_mean": + official_latents.abs().mean().item(), + "fastvideo_abs_mean": + fastvideo_latents.abs().mean().item(), + "diff_max": + difference.abs().max().item(), + "diff_mean": + difference.abs().mean().item(), + "rmse": + difference.square().mean().sqrt().item(), + "cosine": + F.cosine_similarity( + official_latents.flatten().unsqueeze(0), + fastvideo_latents.flatten().unsqueeze(0), + ).item(), + "official_seconds": + round(official_seconds, 3), + "fastvideo_load_seconds": + round(fastvideo_load_seconds, 3), + "fastvideo_seconds": + round(fastvideo_seconds, 3), + "seed": + 42, + "generator_device": + "cpu", + "conditioning_source": + ("shared_official_prompt_embeds" if shared_prompt_embeds_cpu is not None else "each_pipeline_text_encoder"), + "steps": [args.steps] * 3, + } + print("HELIOS_LATENT_PARITY=" + json.dumps(summary, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/helios/test_helios_quality_regression.py b/tests/local_tests/helios/test_helios_quality_regression.py new file mode 100644 index 0000000000..fe26d4eeaa --- /dev/null +++ b/tests/local_tests/helios/test_helios_quality_regression.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Asset-optional video-level quality and container regression checks.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +import subprocess + +import numpy as np +import pytest + + +def _video_summary(path: Path) -> dict: + probe = subprocess.run( + [ + "ffprobe", "-v", "error", "-select_streams", "v:0", "-count_frames", + "-show_entries", "stream=codec_name,width,height,avg_frame_rate,nb_read_frames", + "-of", "json", str(path), + ], + check=True, + capture_output=True, + text=True, + ) + stream = json.loads(probe.stdout)["streams"][0] + width, height = int(stream["width"]), int(stream["height"]) + raw = subprocess.run( + ["ffmpeg", "-v", "error", "-i", str(path), "-f", "rawvideo", "-pix_fmt", "gray", "-"], + check=True, + capture_output=True, + ).stdout + frames = np.frombuffer(raw, dtype=np.uint8).reshape(-1, height, width) + means = frames.mean(axis=(1, 2)) + return { + "codec": stream["codec_name"], + "width": width, + "height": height, + "fps": stream["avg_frame_rate"], + "frames": int(stream["nb_read_frames"]), + "decoded_frames": len(frames), + "mean": float(frames.mean()), + "std": float(frames.std()), + "black_frame_count": int((means < 1.0).sum()), + "frame_mean_std": float(means.std()), + } + + +def test_helios_quality_candidate_has_valid_384x640_video(): + candidate = os.environ.get("HELIOS_QUALITY_CANDIDATE") + if not candidate: + pytest.skip("Set HELIOS_QUALITY_CANDIDATE to run the real-video quality gate") + path = Path(candidate) + if not path.is_file(): + pytest.skip(f"Quality candidate does not exist: {path}") + summary = _video_summary(path) + assert summary["codec"] in {"h264", "hevc", "av1"} + assert (summary["width"], summary["height"]) == (640, 384) + assert summary["fps"] == "24/1" + assert summary["frames"] == summary["decoded_frames"] == 33 + assert summary["std"] > 5.0 + assert summary["black_frame_count"] == 0 + + +def test_helios_quality_candidate_is_stable_against_reference(): + candidate = os.environ.get("HELIOS_QUALITY_CANDIDATE") + reference = os.environ.get("HELIOS_QUALITY_REFERENCE") + if not candidate or not reference: + pytest.skip("Set HELIOS_QUALITY_CANDIDATE and HELIOS_QUALITY_REFERENCE for comparison") + candidate_summary = _video_summary(Path(candidate)) + reference_summary = _video_summary(Path(reference)) + assert abs(candidate_summary["mean"] - reference_summary["mean"]) < 35.0 + assert abs(candidate_summary["std"] - reference_summary["std"]) < 35.0 + assert abs(candidate_summary["frame_mean_std"] - reference_summary["frame_mean_std"]) < 25.0 diff --git a/tests/local_tests/pipelines/test_helios_pipeline_math.py b/tests/local_tests/pipelines/test_helios_pipeline_math.py new file mode 100644 index 0000000000..30e42e903c --- /dev/null +++ b/tests/local_tests/pipelines/test_helios_pipeline_math.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Deterministic contracts for Helios pyramid sampling geometry.""" + +import math +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + + +def _helpers(): + try: + from fastvideo.pipelines.basic.helios.pipeline_utils import ( + build_helios_frame_indices, + calculate_shift, + downsample_to_pyramid_base, + get_generated_pixel_frames, + get_num_latent_chunks, + sample_block_noise, + ) + except ImportError as exc: + raise AssertionError("Helios pyramid helpers have not been implemented") from exc + return { + "build_indices": build_helios_frame_indices, + "calculate_shift": calculate_shift, + "downsample": downsample_to_pyramid_base, + "generated_frames": get_generated_pixel_frames, + "num_chunks": get_num_latent_chunks, + "block_noise": sample_block_noise, + } + + +def test_calculate_shift_matches_official_flux_formula() -> None: + calculate_shift = _helpers()["calculate_shift"] + expected_middle = 0.5 + (1024 - 256) * (1.15 - 0.5) / (4096 - 256) + + assert calculate_shift(256) == 0.5 + assert calculate_shift(1024) == expected_middle + assert calculate_shift(4096) == 1.15 + + +def test_history_frame_indices_match_official_keep_first_frame_layout() -> None: + current, short, mid, long = _helpers()["build_indices"]( + history_sizes=[16, 2, 1], + num_latent_frames_per_chunk=9, + keep_first_frame=True, + device=torch.device("cpu"), + ) + + assert current.tolist() == [list(range(20, 29))] + assert short.tolist() == [[0, 19]] + assert mid.tolist() == [[17, 18]] + assert long.tolist() == [list(range(1, 17))] + + +def test_downsample_to_pyramid_base_matches_official_bilinear_loop() -> None: + latents = torch.arange(1 * 2 * 3 * 8 * 12, dtype=torch.float32).reshape(1, 2, 3, 8, 12) + expected = latents.permute(0, 2, 1, 3, 4).reshape(3, 2, 8, 12) + expected = F.interpolate(expected, size=(4, 6), mode="bilinear") * 2 + expected = F.interpolate(expected, size=(2, 3), mode="bilinear") * 2 + expected = expected.reshape(1, 3, 2, 2, 3).permute(0, 2, 1, 3, 4) + + actual = _helpers()["downsample"](latents, num_stages=3) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_block_noise_matches_official_covariance_sampling() -> None: + if not torch.cuda.is_available(): + pytest.skip("The official FP32 block-noise Cholesky contract requires CUDA") + device = torch.device("cuda:0") + scheduler = SimpleNamespace(config=SimpleNamespace(gamma=1 / 3)) + actual = _helpers()["block_noise"]( + scheduler=scheduler, + shape=(1, 1, 2, 4, 4), + patch_size=(1, 2, 2), + device=device, + generator=torch.Generator("cpu").manual_seed(123), + ) + + block_size = 4 + covariance = (torch.eye(block_size, device=device) * (1 + 1 / 3) - + torch.ones(block_size, block_size, device=device) * (1 / 3) + + torch.eye(block_size, device=device) * 1e-8).float() + cholesky = torch.linalg.cholesky(covariance) + standard = torch.randn( + 1 * 1 * 2 * 2 * 2, + block_size, + generator=torch.Generator("cpu").manual_seed(123), + ).to(device) + expected = (standard @ cholesky.T).view(1, 1, 2, 2, 2, 2, 2) + expected = expected.permute(0, 1, 2, 3, 5, 4, 6).reshape(1, 1, 2, 4, 4) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_pixel_frame_count_rounds_up_to_complete_helios_chunks() -> None: + get_num_latent_chunks = _helpers()["num_chunks"] + + assert get_num_latent_chunks(1, 9, 4) == 1 + assert get_num_latent_chunks(33, 9, 4) == 1 + assert get_num_latent_chunks(34, 9, 4) == 2 + assert get_num_latent_chunks(240, 9, 4) == 8 + + +def test_pixel_frame_count_matches_official_chunk_decode_contract() -> None: + get_generated_pixel_frames = _helpers()["generated_frames"] + + assert get_generated_pixel_frames(33, 4) == 33 + assert get_generated_pixel_frames(66, 4) == 65 + assert get_generated_pixel_frames(264, 4) == 261 + + +def test_pyramid_shift_stays_linear_between_official_endpoints() -> None: + calculate_shift = _helpers()["calculate_shift"] + + assert math.isclose(calculate_shift(2176), 0.825, rel_tol=0, abs_tol=1e-12) diff --git a/tests/local_tests/pipelines/test_helios_pipeline_parity.py b/tests/local_tests/pipelines/test_helios_pipeline_parity.py new file mode 100644 index 0000000000..a4484ea1c5 --- /dev/null +++ b/tests/local_tests/pipelines/test_helios_pipeline_parity.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end latent parity for Helios-Distilled T2V. + +This local test is intentionally heavyweight: it executes the official +Diffusers pipeline and the public FastVideo pipeline with the same CPU RNG, +then compares their final denoised latents. It is not intended for CI without +the pinned local checkpoint. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest +import torch + +REPO_ROOT = Path(__file__).resolve().parents[3] +MODEL_DIR = REPO_ROOT / "official_weights" / "helios" +RUNNER = REPO_ROOT / "tests" / "local_tests" / "helios" / "run_helios_latent_parity.py" + + +def _parse_summary(stdout: str) -> dict: + prefix = "HELIOS_LATENT_PARITY=" + for line in reversed(stdout.splitlines()): + if line.startswith(prefix): + return json.loads(line.removeprefix(prefix)) + raise AssertionError(f"Parity runner did not emit {prefix!r}\n{stdout[-4000:]}") + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Helios pipeline parity requires CUDA.", +) +def test_helios_distilled_t2v_latent_parity() -> None: + if not (MODEL_DIR / "model_index.json").is_file(): + pytest.skip(f"Pinned Helios checkpoint not found at {MODEL_DIR}") + + environment = os.environ.copy() + environment.setdefault("CUDA_VISIBLE_DEVICES", "0") + completed = subprocess.run( + [ + sys.executable, + str(RUNNER), + "--model-dir", + str(MODEL_DIR), + "--gpu", + "0", + ], + cwd=REPO_ROOT, + env=environment, + check=True, + capture_output=True, + text=True, + timeout=900, + ) + summary = _parse_summary(completed.stdout) + + assert summary["shape"] == [1, 16, 9, 16, 24] + relative_abs_mean_drift = abs(summary["fastvideo_abs_mean"] - + summary["official_abs_mean"]) / summary["official_abs_mean"] + + # Component tests establish exact control flow and close single-forward + # BF16 parity. This full six-forward DMD loop uses tolerance metrics because + # fused-QKV and SDPA kernel rounding feeds back through later timesteps. + assert summary["cosine"] >= 0.95 + assert relative_abs_mean_drift <= 0.05 + assert summary["diff_mean"] <= 0.30 + assert summary["rmse"] <= 0.40 diff --git a/tests/local_tests/pipelines/test_helios_pipeline_smoke.py b/tests/local_tests/pipelines/test_helios_pipeline_smoke.py new file mode 100644 index 0000000000..0c537e3796 --- /dev/null +++ b/tests/local_tests/pipelines/test_helios_pipeline_smoke.py @@ -0,0 +1,285 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Public API contracts for the Helios-Distilled pipeline.""" + +from argparse import ArgumentParser +from dataclasses import fields +import json +from pathlib import Path +import subprocess +import sys +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.api.compat import normalize_generation_request, request_to_sampling_param +from fastvideo.api.sampling_param import SamplingParam +from fastvideo.api.schema import SamplingConfig +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + +REPO_ROOT = Path(__file__).resolve().parents[3] +MODEL_DIR = REPO_ROOT / "official_weights" / "helios" + + +HELIOS_SAMPLING_VALUES = { + "pyramid_num_inference_steps_list": [2, 2, 2], + "history_sizes": [16, 2, 1], + "num_latent_frames_per_chunk": 9, + "keep_first_frame": True, + "is_skip_first_chunk": False, + "use_zero_init": True, + "zero_steps": 1, + "is_amplify_first_chunk": True, +} + + +def _pipeline_symbols(): + try: + from fastvideo.configs.pipelines.helios import ( + HeliosPipelineConfig, + helios_postprocess_text, + helios_preprocess_text, + ) + from fastvideo.pipelines.basic.helios.helios_pipeline import HeliosPyramidPipeline + from fastvideo.pipelines.basic.helios.presets import HELIOS_DISTILLED_T2V + except ImportError as exc: + raise AssertionError("Helios pipeline/config/preset have not been implemented") from exc + return ( + HeliosPipelineConfig, + helios_preprocess_text, + helios_postprocess_text, + HeliosPyramidPipeline, + HELIOS_DISTILLED_T2V, + ) + + +def test_helios_public_sampling_fields_reach_forward_batch() -> None: + sampling_fields = {item.name for item in fields(SamplingParam)} + typed_fields = {item.name for item in fields(SamplingConfig)} + batch_fields = {item.name for item in fields(ForwardBatch)} + expected = set(HELIOS_SAMPLING_VALUES) + + assert expected <= sampling_fields + assert expected <= typed_fields + assert expected <= batch_fields + + +def test_helios_internal_chunk_state_is_declared_on_forward_batch() -> None: + batch_fields = {item.name for item in fields(ForwardBatch)} + + assert "helios_latent_chunks" in batch_fields + + +def test_helios_typed_request_maps_sampling_fields() -> None: + request = normalize_generation_request({"sampling": HELIOS_SAMPLING_VALUES}) + sampling = request_to_sampling_param( + request, + model_path="BestWishYsh/Helios-Distilled", + ) + + for name, expected in HELIOS_SAMPLING_VALUES.items(): + assert getattr(sampling, name) == expected + + +def test_helios_sampling_fields_are_available_from_cli() -> None: + parser = ArgumentParser() + SamplingParam.add_cli_args(parser) + + args = parser.parse_args([ + "--pyramid-num-inference-steps-list", + "3", + "2", + "1", + "--history-sizes", + "8", + "2", + "1", + "--num-latent-frames-per-chunk", + "9", + "--keep-first-frame", + "true", + "--is-skip-first-chunk", + "false", + "--use-zero-init", + "true", + "--zero-steps", + "2", + "--is-amplify-first-chunk", + "true", + ]) + + assert args.pyramid_num_inference_steps_list == [3, 2, 1] + assert args.history_sizes == [8, 2, 1] + assert args.num_latent_frames_per_chunk == 9 + assert args.keep_first_frame is True + assert args.is_skip_first_chunk is False + assert args.use_zero_init is True + assert args.zero_steps == 2 + assert args.is_amplify_first_chunk is True + + +def _valid_helios_batch(**overrides) -> ForwardBatch: + values = { + "data_type": "video", + "prompt": "A paper boat floats across a rain puddle.", + "seed": 42, + "height": 128, + "width": 192, + "num_frames": 33, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "pyramid_num_inference_steps_list": [2, 2, 2], + "history_sizes": [16, 2, 1], + "num_latent_frames_per_chunk": 9, + "keep_first_frame": True, + "is_skip_first_chunk": False, + "zero_steps": 1, + } + values.update(overrides) + return ForwardBatch(**values) + + +def _validation_args(): + return SimpleNamespace( + pipeline_config=SimpleNamespace( + ti2v_task=False, + is_causal=False, + ), ) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"height": 120}, "divisible by 64"), + ({"pyramid_num_inference_steps_list": [2, 2]}, "three positive pyramid step counts"), + ({"pyramid_num_inference_steps_list": [2, 0, 2]}, "three positive pyramid step counts"), + ({"history_sizes": [16, 2]}, "three positive history sizes"), + ({"history_sizes": [16, 0, 1]}, "three positive history sizes"), + ({"num_latent_frames_per_chunk": 8}, "requires num_latent_frames_per_chunk=9"), + ({"keep_first_frame": False}, "requires keep_first_frame=True"), + ({"is_skip_first_chunk": True}, "only meaningful for conditioned Helios modes"), + ({"zero_steps": -1}, "zero_steps must be non-negative"), + ], +) +def test_helios_input_validation_rejects_unverified_contracts(overrides, message) -> None: + try: + from fastvideo.pipelines.basic.helios.stages import HeliosInputValidationStage + except ImportError as exc: + raise AssertionError("HeliosInputValidationStage has not been implemented") from exc + + with pytest.raises(ValueError, match=message): + HeliosInputValidationStage().forward(_valid_helios_batch(**overrides), _validation_args()) + + +def test_helios_pipeline_config_and_text_contract() -> None: + from fastvideo.configs.models.encoders import BaseEncoderOutput + + HeliosPipelineConfig, preprocess, postprocess, _, _ = _pipeline_symbols() + config = HeliosPipelineConfig() + assert config.dit_config.__class__.__name__ == "HeliosConfig" + assert config.vae_config.__class__.__name__ == "WanVAEConfig" + assert config.vae_config.load_encoder is False + assert config.vae_config.load_decoder is True + assert config.text_encoder_configs[0].architectures == ["UMT5EncoderModel"] + assert config.text_encoder_configs[0].text_len == 512 + assert config.dit_precision == "bf16" + assert config.text_encoder_precisions == ("bf16", ) + assert config.vae_precision == "fp32" + assert config.flow_shift is None + assert preprocess(" A\n fish & reef ") == "A fish & reef" + + hidden = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) + mask = torch.tensor([[1, 1, 1, 0, 0, 0], [1, 1, 1, 1, 1, 0]]) + output = postprocess(BaseEncoderOutput(last_hidden_state=hidden, attention_mask=mask)) + assert output.shape == (2, 512, 4) + assert torch.equal(output[0, :3], hidden[0, :3]) + assert torch.count_nonzero(output[0, 3:]) == 0 + assert torch.equal(output[1, :5], hidden[1, :5]) + + long_hidden = torch.arange(520 * 4, dtype=torch.float32).reshape(1, 520, 4) + long_mask = torch.ones(1, 520, dtype=torch.long) + long_output = postprocess(BaseEncoderOutput(last_hidden_state=long_hidden, attention_mask=long_mask)) + assert long_output.shape == (1, 512, 4) + assert torch.equal(long_output[0], long_hidden[0, :512]) + + +def test_helios_preset_matches_official_distilled_defaults() -> None: + _, _, _, _, preset = _pipeline_symbols() + assert preset.model_family == "helios" + assert preset.workload_type == "t2v" + assert preset.defaults["height"] == 384 + assert preset.defaults["width"] == 640 + assert preset.defaults["num_frames"] == 240 + assert preset.defaults["fps"] == 24 + assert preset.defaults["guidance_scale"] == 1.0 + assert preset.defaults["pyramid_num_inference_steps_list"] == [2, 2, 2] + assert preset.defaults["history_sizes"] == [16, 2, 1] + assert preset.defaults["num_latent_frames_per_chunk"] == 9 + assert preset.defaults["is_amplify_first_chunk"] is True + + +def test_helios_pipeline_entry_class_is_discoverable() -> None: + from fastvideo.pipelines.pipeline_registry import PipelineType, import_pipeline_classes + + _, _, _, pipeline_class, _ = _pipeline_symbols() + assert pipeline_class.__name__ == "HeliosPyramidPipeline" + assert pipeline_class._required_config_modules == [ + "text_encoder", + "tokenizer", + "vae", + "transformer", + "scheduler", + ] + discovered = import_pipeline_classes(PipelineType.BASIC)["basic"] + assert discovered["HeliosPyramidPipeline"] is pipeline_class + + +def test_helios_local_model_registry_selection() -> None: + from fastvideo.registry import get_model_family, get_pipeline_config_cls_from_name + + if not MODEL_DIR.exists(): + pytest.skip(f"Pinned Helios checkpoint not found at {MODEL_DIR}") + assert get_model_family(str(MODEL_DIR)) == "helios" + assert get_pipeline_config_cls_from_name(str(MODEL_DIR)).__name__ == "HeliosPipelineConfig" + + +def test_helios_registry_requires_distilled_metadata() -> None: + script = r''' +import json +import tempfile +from pathlib import Path + +from fastvideo.registry import get_model_family + +def write_model_index(model_dir, is_distilled): + model_dir.mkdir() + for component in ("transformer", "scheduler", "text_encoder", "vae"): + (model_dir / component).mkdir() + (model_dir / "model_index.json").write_text(json.dumps({ + "_class_name": "HeliosPyramidPipeline", + "_diffusers_version": "0.39.0", + "is_distilled": is_distilled, + "scheduler": ["diffusers", "HeliosDMDScheduler"], + "transformer": ["diffusers", "HeliosTransformer3DModel"], + "text_encoder": ["transformers", "UMT5EncoderModel"], + "vae": ["diffusers", "AutoencoderKLWan"], + })) + +with tempfile.TemporaryDirectory() as root: + root = Path(root) + distilled = root / "renamed_helios_checkpoint" + write_model_index(distilled, True) + assert get_model_family(str(distilled)) == "helios" + for name in ("Helios-Base", "Helios-Mid", "unrelated-helios-experiment"): + model_dir = root / name + write_model_index(model_dir, False) + assert get_model_family(str(model_dir)) is None, name +''' + result = subprocess.run( + [sys.executable, "-c", script], + cwd=REPO_ROOT, + check=True, + capture_output=True, + text=True, + ) + assert result.returncode == 0 diff --git a/tests/local_tests/pipelines/test_helios_pipeline_stages.py b/tests/local_tests/pipelines/test_helios_pipeline_stages.py new file mode 100644 index 0000000000..06ff2c3bdf --- /dev/null +++ b/tests/local_tests/pipelines/test_helios_pipeline_stages.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tiny end-to-end contracts for Helios custom pipeline stages.""" + +from __future__ import annotations + +from functools import lru_cache +import json +from pathlib import Path +import subprocess +import sys + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[3] + + +@lru_cache +def _probe_stages() -> dict: + script = r""" +import json +import math +from types import SimpleNamespace + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.models.schedulers.scheduling_helios_dmd import HeliosDMDScheduler +from fastvideo.pipelines.basic.helios.stages import ( + HeliosChunkDecodingStage, + HeliosPyramidDenoisingStage, +) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + +if not torch.cuda.is_available(): + print(json.dumps({"cuda_available": False})) + raise SystemExit(0) + +device = torch.device("cuda") + + +class TinyTransformer(nn.Module): + in_channels = 2 + patch_size = (1, 2, 2) + + def __init__(self): + super().__init__() + self.anchor = nn.Parameter(torch.zeros((), device=device, dtype=torch.bfloat16)) + self.calls = [] + + def forward( + self, + hidden_states, + timestep, + encoder_hidden_states, + indices_hidden_states, + indices_latents_history_short, + indices_latents_history_mid, + indices_latents_history_long, + latents_history_short, + latents_history_mid, + latents_history_long, + **kwargs, + ): + del kwargs + self.calls.append({ + "latent_shape": list(hidden_states.shape), + "history_shapes": [ + list(latents_history_short.shape), + list(latents_history_mid.shape), + list(latents_history_long.shape), + ], + "indices": [ + indices_hidden_states.tolist(), + indices_latents_history_short.tolist(), + indices_latents_history_mid.tolist(), + indices_latents_history_long.tolist(), + ], + "history_means": [ + latents_history_short.float().mean().item(), + latents_history_mid.float().mean().item(), + latents_history_long.float().mean().item(), + ], + "short_prefix_mean": latents_history_short[:, :, :1].float().mean().item(), + }) + scalar = ( + timestep.float().view(-1, 1, 1, 1, 1) / 1000 + + encoder_hidden_states.float().mean(dim=(1, 2)).view(-1, 1, 1, 1, 1) * 0.01 + + latents_history_short.float().mean(dim=(1, 2, 3, 4)).view(-1, 1, 1, 1, 1) * 0.02 + + latents_history_mid.float().mean(dim=(1, 2, 3, 4)).view(-1, 1, 1, 1, 1) * 0.03 + + latents_history_long.float().mean(dim=(1, 2, 3, 4)).view(-1, 1, 1, 1, 1) * 0.04 + ) + return (hidden_states.float() * 0.125 + scalar).to(hidden_states.dtype) + + +def scheduler(): + return HeliosDMDScheduler( + stages=3, + stage_range=[0, 1 / 3, 2 / 3, 1], + gamma=1 / 3, + shift=1.0, + use_dynamic_shifting=True, + time_shift_type="linear", + ) + + +def block_noise(sched, shape, generator): + b, c, t, h, w = shape + block_size = 4 + gamma = sched.config.gamma + covariance = ( + torch.eye(block_size, device=device) * (1 + gamma) + - torch.ones(block_size, block_size, device=device) * gamma + ) + covariance += torch.eye(block_size, device=device) * 1e-8 + cholesky = torch.linalg.cholesky(covariance.float()) + z = torch.randn( + b * c * t * (h // 2) * (w // 2), + block_size, + generator=generator, + device=generator.device, + ).to(device) + noise = z @ cholesky.T + noise = noise.view(b, c, t, h // 2, w // 2, 2, 2) + return noise.permute(0, 1, 2, 3, 5, 4, 6).reshape(shape) + + +def reference_sample(model, sched, generator, prompt): + history_sizes = [16, 2, 1] + history = torch.zeros(1, 2, 19, 8, 8, device=device) + history_long, history_mid, history_one = history.split(history_sizes, dim=2) + history_short = torch.cat([torch.zeros(1, 2, 1, 8, 8, device=device), history_one], dim=2) + + all_indices = torch.arange(29, device=device) + prefix, long_idx, mid_idx, one_idx, current_idx = all_indices.split([1, 16, 2, 1, 9]) + indices = ( + current_idx.unsqueeze(0), + torch.cat([prefix, one_idx]).unsqueeze(0), + mid_idx.unsqueeze(0), + long_idx.unsqueeze(0), + ) + + latents = torch.randn((1, 2, 9, 8, 8), generator=generator).to(device) + flat = latents.permute(0, 2, 1, 3, 4).reshape(9, 2, 8, 8) + flat = F.interpolate(flat, size=(4, 4), mode="bilinear") * 2 + flat = F.interpolate(flat, size=(2, 2), mode="bilinear") * 2 + latents = flat.reshape(1, 9, 2, 2, 2).permute(0, 2, 1, 3, 4) + start_points = [latents] + + for stage_index in range(3): + image_seq_len = math.prod(latents.shape[-3:]) // 4 + mu = image_seq_len * ((1.15 - 0.5) / (4096 - 256)) + ( + 0.5 - ((1.15 - 0.5) / (4096 - 256)) * 256 + ) + sched.set_timesteps( + 1, + stage_index, + device=device, + mu=mu, + is_amplify_first_chunk=True, + ) + timesteps = sched.timesteps + if stage_index > 0: + b, c, t, h, w = latents.shape + flat = latents.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) + flat = F.interpolate(flat, size=(h * 2, w * 2), mode="nearest") + latents = flat.reshape(b, t, c, h * 2, w * 2).permute(0, 2, 1, 3, 4) + original_signal = 1 - sched.ori_start_sigmas[stage_index] + gamma = sched.config.gamma + alpha = 1 / (math.sqrt(1 + 1 / gamma) * (1 - original_signal) + original_signal) + beta = alpha * (1 - original_signal) / math.sqrt(gamma) + latents = alpha * latents + beta * block_noise(sched, tuple(latents.shape), generator).to(torch.bfloat16) + start_points.append(latents) + + for step_index, timestep_value in enumerate(timesteps): + timestep = timestep_value.expand(1).to(torch.int64) + prediction = model( + hidden_states=latents.to(torch.bfloat16), + timestep=timestep, + encoder_hidden_states=prompt, + indices_hidden_states=indices[0], + indices_latents_history_short=indices[1], + indices_latents_history_mid=indices[2], + indices_latents_history_long=indices[3], + latents_history_short=history_short.to(torch.bfloat16), + latents_history_mid=history_mid.to(torch.bfloat16), + latents_history_long=history_long.to(torch.bfloat16), + ) + latents = sched.step( + prediction, + timestep_value, + latents, + generator=generator, + return_dict=False, + cur_sampling_step=step_index, + dmd_noisy_tensor=start_points[stage_index], + dmd_sigmas=sched.sigmas, + dmd_timesteps=sched.timesteps, + all_timesteps=timesteps, + )[0] + return latents + + +prompt = torch.ones(1, 5, 4, device=device, dtype=torch.bfloat16) +actual_model = TinyTransformer() +actual_scheduler = scheduler() +batch = ForwardBatch( + data_type="video", + prompt_embeds=[prompt], + generator=[torch.Generator("cpu").manual_seed(321)], + height=64, + width=64, + num_frames=33, + guidance_scale=1.0, + pyramid_num_inference_steps_list=[1, 1, 1], + history_sizes=[16, 2, 1], + num_latent_frames_per_chunk=9, + keep_first_frame=True, + is_amplify_first_chunk=True, +) +args = SimpleNamespace( + pipeline_config=SimpleNamespace( + dit_precision="bf16", + vae_config=SimpleNamespace( + arch_config=SimpleNamespace(scale_factor_spatial=8, scale_factor_temporal=4) + ), + ), + model_loaded={"transformer": True, "vae": True}, + dit_cpu_offload=False, + dit_layerwise_offload=False, + use_fsdp_inference=False, +) +actual_batch = HeliosPyramidDenoisingStage(actual_model, actual_scheduler).forward(batch, args) + +autoregressive_model = TinyTransformer() +autoregressive_batch = ForwardBatch( + data_type="video", + prompt_embeds=[prompt], + generator=[torch.Generator("cpu").manual_seed(321)], + height=64, + width=64, + num_frames=65, + guidance_scale=1.0, + pyramid_num_inference_steps_list=[1, 1, 1], + history_sizes=[16, 2, 1], + num_latent_frames_per_chunk=9, + keep_first_frame=True, + is_amplify_first_chunk=True, +) +autoregressive_batch = HeliosPyramidDenoisingStage(autoregressive_model, scheduler()).forward( + autoregressive_batch, args +) + +reference_model = TinyTransformer() +expected = reference_sample( + reference_model, + scheduler(), + torch.Generator("cpu").manual_seed(321), + prompt, +) + + +class TinyVAE(nn.Module): + handles_latent_denorm = True + + def __init__(self): + super().__init__() + self.calls = [] + + def decode(self, latent): + self.calls.append(list(latent.shape)) + output_frames = (latent.shape[2] - 1) * 4 + 1 + value = -1.0 if len(self.calls) == 1 else 1.0 + return torch.full( + (latent.shape[0], 3, output_frames, latent.shape[3] * 8, latent.shape[4] * 8), + value, + device=latent.device, + ) + + +vae = TinyVAE().to(device) +decode_batch = ForwardBatch(data_type="video", num_frames=34) +decode_batch.latents = torch.cat([expected, expected], dim=2) +decode_batch.helios_latent_chunks = [expected, expected] +decode_args = SimpleNamespace( + output_type="video", + model_loaded={"vae": True}, + pipeline_config=SimpleNamespace( + vae_decode_precision="fp32", + vae_precision="fp32", + vae_tiling=False, + vae_config=SimpleNamespace(arch_config=SimpleNamespace(scale_factor_temporal=4)), + ), + disable_autocast=False, + vae_cpu_offload=False, +) +decoded = HeliosChunkDecodingStage(vae).forward(decode_batch, decode_args).output + +latent_batch = ForwardBatch(data_type="video", latents=expected) +latent_args = SimpleNamespace(output_type="latent") +latent_output = HeliosChunkDecodingStage(vae).forward(latent_batch, latent_args).output + +print(json.dumps({ + "cuda_available": True, + "latent_max_diff": (actual_batch.latents - expected).abs().max().item(), + "call_shapes": [item["latent_shape"] for item in actual_model.calls], + "history_shapes": actual_model.calls[0]["history_shapes"], + "indices": actual_model.calls[0]["indices"], + "autoregressive_latent_shape": list(autoregressive_batch.latents.shape), + "autoregressive_call_count": len(autoregressive_model.calls), + "autoregressive_second_history_means": autoregressive_model.calls[6]["history_means"], + "autoregressive_second_short_prefix_mean": autoregressive_model.calls[6]["short_prefix_mean"], + "vae_calls": vae.calls, + "decoded_shape": list(decoded.shape), + "decoded_device": decoded.device.type, + "decoded_first_mean": decoded[:, :, :33].mean().item(), + "decoded_second_mean": decoded[:, :, 33:].mean().item(), + "latent_output_device": latent_output.device.type, +})) +""" + result = subprocess.run( + [sys.executable, "-c", script], + cwd=REPO_ROOT, + check=True, + capture_output=True, + text=True, + ) + return json.loads(result.stdout.strip().splitlines()[-1]) + + +def _results() -> dict: + result = _probe_stages() + if not result["cuda_available"]: + pytest.skip("Helios tiny stage parity requires CUDA") + return result + + +def test_tiny_pyramid_stage_matches_independent_official_loop(): + result = _results() + assert result["latent_max_diff"] == 0 + assert result["call_shapes"] == [ + [1, 2, 9, 2, 2], + [1, 2, 9, 2, 2], + [1, 2, 9, 4, 4], + [1, 2, 9, 4, 4], + [1, 2, 9, 8, 8], + [1, 2, 9, 8, 8], + ] + + +def test_tiny_pyramid_stage_passes_exact_history_and_indices(): + result = _results() + assert result["history_shapes"] == [ + [1, 2, 2, 8, 8], + [1, 2, 2, 8, 8], + [1, 2, 16, 8, 8], + ] + current, short, mid, long = result["indices"] + assert current == [list(range(20, 29))] + assert short == [[0, 19]] + assert mid == [[17, 18]] + assert long == [list(range(1, 17))] + + +def test_tiny_pyramid_stage_uses_history_on_second_chunk(): + result = _results() + assert result["autoregressive_latent_shape"] == [1, 2, 18, 8, 8] + assert result["autoregressive_call_count"] == 9 + assert all(abs(value) > 1e-5 for value in result["autoregressive_second_history_means"]) + assert abs(result["autoregressive_second_short_prefix_mean"]) > 1e-5 + + +def test_chunk_decoder_calls_vae_per_chunk_and_trims_to_requested_frames(): + result = _results() + assert result["vae_calls"] == [[1, 2, 9, 8, 8], [1, 2, 9, 8, 8]] + assert result["decoded_shape"] == [1, 3, 34, 64, 64] + assert result["decoded_device"] == "cpu" + assert result["decoded_first_mean"] == 0 + assert result["decoded_second_mean"] == 1 + assert result["latent_output_device"] == "cpu" diff --git a/tests/local_tests/schedulers/test_helios_dmd_scheduler_parity.py b/tests/local_tests/schedulers/test_helios_dmd_scheduler_parity.py new file mode 100644 index 0000000000..762992965d --- /dev/null +++ b/tests/local_tests/schedulers/test_helios_dmd_scheduler_parity.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Helios-Distilled DMD scheduler parity. + +Coverage scope: implementation_subcomponent. The test compares the native +FastVideo scheduler against the exact Diffusers class declared by the pinned +Helios-Distilled scheduler config. It covers every pyramid stage, dynamic time +shift, first-chunk amplification, and both branches of the DMD step. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import torch +from diffusers import HeliosDMDScheduler as OfficialHeliosDMDScheduler +from torch.testing import assert_close + +REPO_ROOT = Path(__file__).resolve().parents[3] +SCHEDULER_DIR = REPO_ROOT / "official_weights" / "helios" / "scheduler" +PARITY_SCOPE = "implementation_subcomponent" + + +def _scheduler_kwargs() -> dict: + config_path = SCHEDULER_DIR / "scheduler_config.json" + if not config_path.exists(): + pytest.skip(f"Helios scheduler config missing: {config_path}") + config = json.loads(config_path.read_text(encoding="utf-8")) + for key in ("_class_name", "_diffusers_version", "scheduler_type"): + config.pop(key, None) + return config + + +def _fastvideo_scheduler_class(): + try: + from fastvideo.models.schedulers.scheduling_helios_dmd import ( + HeliosDMDScheduler, ) + except ImportError as exc: + raise AssertionError("Native FastVideo HeliosDMDScheduler has not been implemented yet") from exc + return HeliosDMDScheduler + + +def _make_pair(): + kwargs = _scheduler_kwargs() + return ( + OfficialHeliosDMDScheduler(**kwargs), + _fastvideo_scheduler_class()(**kwargs), + ) + + +def test_helios_dmd_scheduler_resolves_through_production_registry(): + from fastvideo.models.registry import ModelRegistry + + scheduler_cls, architecture = ModelRegistry.resolve_model_cls("HeliosDMDScheduler") + + assert architecture == "HeliosDMDScheduler" + assert scheduler_cls is _fastvideo_scheduler_class() + + +@pytest.mark.parametrize("stage_index", [0, 1, 2]) +@pytest.mark.parametrize("amplify", [False, True]) +def test_helios_dmd_scheduler_stage_schedule_parity(stage_index: int, amplify: bool): + official, fastvideo = _make_pair() + call_kwargs = { + "num_inference_steps": 2, + "stage_index": stage_index, + "device": "cpu", + "mu": 1.07, + "is_amplify_first_chunk": amplify, + } + official.set_timesteps(**call_kwargs) + fastvideo.set_timesteps(**call_kwargs) + + assert_close(fastvideo.timesteps, official.timesteps, atol=0, rtol=0) + assert_close(fastvideo.sigmas, official.sigmas, atol=0, rtol=0) + assert fastvideo.timestep_ratios == official.timestep_ratios + assert fastvideo.start_sigmas == official.start_sigmas + assert fastvideo.end_sigmas == official.end_sigmas + assert fastvideo.ori_start_sigmas == official.ori_start_sigmas + + +@pytest.mark.parametrize("stage_index", [0, 1, 2]) +def test_helios_dmd_scheduler_step_parity(stage_index: int): + official, fastvideo = _make_pair() + call_kwargs = { + "num_inference_steps": 2, + "stage_index": stage_index, + "device": "cpu", + "mu": 1.07, + } + official.set_timesteps(**call_kwargs) + fastvideo.set_timesteps(**call_kwargs) + + generator = torch.Generator(device="cpu").manual_seed(20260711 + stage_index) + sample = torch.randn(2, 4, 3, 4, 6, generator=generator) + noisy_start = torch.randn(sample.shape, generator=generator) + + for step_index, (official_t, fastvideo_t) in enumerate(zip(official.timesteps, fastvideo.timesteps, strict=True)): + model_output = torch.randn(sample.shape, generator=generator) + official_sample = official.step( + model_output=model_output, + timestep=official_t, + sample=sample, + cur_sampling_step=step_index, + dmd_noisy_tensor=noisy_start, + dmd_sigmas=official.sigmas, + dmd_timesteps=official.timesteps, + all_timesteps=official.timesteps, + return_dict=False, + )[0] + fastvideo_sample = fastvideo.step( + model_output=model_output, + timestep=fastvideo_t, + sample=sample, + cur_sampling_step=step_index, + dmd_noisy_tensor=noisy_start, + dmd_sigmas=fastvideo.sigmas, + dmd_timesteps=fastvideo.timesteps, + all_timesteps=fastvideo.timesteps, + return_dict=False, + )[0] + assert_close(fastvideo_sample, official_sample, atol=0, rtol=0) + sample = official_sample diff --git a/tests/local_tests/transformers/test_helios_transformer_parity.py b/tests/local_tests/transformers/test_helios_transformer_parity.py new file mode 100644 index 0000000000..4f15d4e279 --- /dev/null +++ b/tests/local_tests/transformers/test_helios_transformer_parity.py @@ -0,0 +1,640 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Helios transformer config, weight-schema, and tiny forward parity. + +Coverage scope: both. The small model exercises all three history terms and +frame-indexed 3D RoPE without allocating the 40-layer checkpoint. Official +Diffusers weights are loaded strictly into the native FastVideo key schema +before comparing deterministic float32 outputs. +""" + +from __future__ import annotations + +import argparse +import gc +import json +import os +import socket +import subprocess +import sys +from dataclasses import fields +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +from diffusers import HeliosTransformer3DModel as OfficialHeliosTransformer3DModel +from torch.testing import assert_close + +from fastvideo.forward_context import set_forward_context +from fastvideo.models.loader.fsdp_load import load_model_from_full_model_state_dict +from fastvideo.models.loader.utils import get_param_names_mapping, set_default_torch_dtype +from fastvideo.models.loader.weight_utils import ( + resolve_safetensors_files, + safetensors_weights_iterator, +) + +os.environ.setdefault("DISABLE_SP", "1") +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") + +REPO_ROOT = Path(__file__).resolve().parents[3] +TRANSFORMER_DIR = Path( + os.getenv( + "HELIOS_TRANSFORMER_DIR", + REPO_ROOT / "official_weights" / "helios" / "transformer", + )) +HF_REVISION = "1999182614cb08d3bdcc46b9827504af2914b87b" +PARITY_SCOPE = "both" +SP_WORLD_SIZE = 2 + + +def _native_types(): + try: + from fastvideo.configs.models.dits.helios import ( + HeliosArchConfig, + HeliosConfig, + ) + from fastvideo.models.dits.helios import HeliosTransformer3DModel + except ImportError as exc: + raise AssertionError("Native FastVideo Helios transformer/config have not been implemented yet") from exc + return HeliosArchConfig, HeliosConfig, HeliosTransformer3DModel + + +def _tiny_kwargs() -> dict: + return { + "patch_size": (1, 2, 2), + "num_attention_heads": 2, + "attention_head_dim": 32, + "in_channels": 4, + "out_channels": 4, + "text_dim": 48, + "freq_dim": 32, + "ffn_dim": 128, + "num_layers": 1, + "cross_attn_norm": True, + "qk_norm": "rms_norm_across_heads", + "eps": 1e-6, + "added_kv_proj_dim": None, + "rope_dim": (12, 10, 10), + "rope_theta": 10000.0, + "guidance_cross_attn": True, + "zero_history_timestep": True, + "has_multi_term_memory_patch": True, + "is_amplify_history": False, + "history_scale_mode": "per_head", + } + + +def _make_inputs() -> dict[str, torch.Tensor]: + generator = torch.Generator(device="cpu").manual_seed(3535) + return { + "hidden_states": torch.randn(1, 4, 2, 8, 8, generator=generator), + "timestep": torch.tensor([517], dtype=torch.long), + "encoder_hidden_states": torch.randn(1, 5, 48, generator=generator), + "indices_hidden_states": torch.tensor([[19, 20]]), + "indices_latents_history_short": torch.tensor([[17, 18]]), + "indices_latents_history_mid": torch.tensor([[15, 16]]), + "indices_latents_history_long": torch.tensor([[11, 12, 13, 14]]), + "latents_history_short": torch.randn(1, 4, 2, 8, 8, generator=generator), + "latents_history_mid": torch.randn(1, 4, 2, 8, 8, generator=generator), + "latents_history_long": torch.randn(1, 4, 4, 8, 8, generator=generator), + } + + +def _make_pyramid_inputs() -> dict[str, torch.Tensor]: + inputs = _make_inputs() + inputs["hidden_states"] = inputs["hidden_states"][:, :, :, :2, :2] + return inputs + + +def _make_real_inputs(device: torch.device) -> dict[str, torch.Tensor]: + generator = torch.Generator(device=device).manual_seed(1999) + dtype = torch.bfloat16 + return { + "hidden_states": torch.randn(1, 16, 1, 8, 8, device=device, dtype=dtype, generator=generator), + "timestep": torch.tensor([517], device=device, dtype=torch.long), + "encoder_hidden_states": torch.randn(1, 8, 4096, device=device, dtype=dtype, generator=generator), + "indices_hidden_states": torch.tensor([[20]], device=device), + "indices_latents_history_short": torch.tensor([[18, 19]], device=device), + "indices_latents_history_mid": torch.tensor([[16, 17]], device=device), + "indices_latents_history_long": torch.tensor([[12, 13, 14, 15]], device=device), + "latents_history_short": torch.randn(1, 16, 2, 8, 8, device=device, dtype=dtype, generator=generator), + "latents_history_mid": torch.randn(1, 16, 2, 8, 8, device=device, dtype=dtype, generator=generator), + "latents_history_long": torch.randn(1, 16, 4, 8, 8, device=device, dtype=dtype, generator=generator), + } + + +def _make_real_pyramid_inputs(device: torch.device) -> dict[str, torch.Tensor]: + generator = torch.Generator(device=device).manual_seed(2001) + dtype = torch.bfloat16 + return { + "hidden_states": torch.randn(1, 16, 9, 4, 6, device=device, dtype=dtype, generator=generator), + "timestep": torch.tensor([999], device=device, dtype=torch.long), + "encoder_hidden_states": torch.randn(1, 512, 4096, device=device, dtype=dtype, generator=generator), + "indices_hidden_states": torch.arange(20, 29, device=device).unsqueeze(0), + "indices_latents_history_short": torch.tensor([[0, 19]], device=device), + "indices_latents_history_mid": torch.tensor([[17, 18]], device=device), + "indices_latents_history_long": torch.arange(1, 17, device=device).unsqueeze(0), + "latents_history_short": torch.randn(1, 16, 2, 16, 24, device=device, dtype=dtype, generator=generator), + "latents_history_mid": torch.randn(1, 16, 2, 16, 24, device=device, dtype=dtype, generator=generator), + "latents_history_long": torch.randn(1, 16, 16, 16, 24, device=device, dtype=dtype, generator=generator), + } + + +def _move_inputs( + inputs: dict[str, torch.Tensor], + device: torch.device, + dtype: torch.dtype, +) -> dict[str, torch.Tensor]: + return { + name: value.to(device=device, dtype=dtype if value.is_floating_point() else value.dtype) + for name, value in inputs.items() + } + + +def _load_tiny_fastvideo( + backend, + device: torch.device, + dtype: torch.dtype, +): + from fastvideo.attention.selector import _component_attention_backend_scope + + HeliosArchConfig, HeliosConfig, FastVideoHeliosTransformer = _native_types() + kwargs = _tiny_kwargs() + torch.manual_seed(3535) + official = OfficialHeliosTransformer3DModel(**kwargs).to(dtype=dtype).eval() + with set_default_torch_dtype(dtype), _component_attention_backend_scope(backend, component="transformer"): + native = FastVideoHeliosTransformer( + config=HeliosConfig(arch_config=HeliosArchConfig(**kwargs)), + hf_config={}, + ).eval() + incompatible = load_model_from_full_model_state_dict( + native, + iter(official.state_dict().items()), + device=device, + param_dtype=dtype, + strict=True, + param_names_mapping=get_param_names_mapping(native.param_names_mapping), + training_mode=False, + ) + assert incompatible.missing_keys == [] + assert incompatible.unexpected_keys == [] + native.materialize_non_persistent_buffers(device, dtype) + return native + + +def _has_real_weights() -> bool: + return (TRANSFORMER_DIR / "diffusion_pytorch_model.safetensors.index.json").is_file() and any( + TRANSFORMER_DIR.glob("*.safetensors")) + + +def _load_fastvideo_production(): + from fastvideo.configs.models.dits.helios import HeliosConfig + from fastvideo.configs.pipelines.base import PipelineConfig + from fastvideo.fastvideo_args import FastVideoArgs + from fastvideo.models.dits.helios import HeliosTransformer3DModel + from fastvideo.models.loader.component_loader import TransformerLoader + + args = FastVideoArgs( + model_path=str(TRANSFORMER_DIR), + dit_cpu_offload=False, + dit_layerwise_offload=False, + use_fsdp_inference=False, + pipeline_config=PipelineConfig( + dit_config=HeliosConfig(), + dit_precision="bf16", + ), + ) + model = TransformerLoader().load(str(TRANSFORMER_DIR), args).eval() + assert isinstance(model, HeliosTransformer3DModel) + assert args.model_paths["transformer"] == str(TRANSFORMER_DIR) + assert next(model.parameters()).device.type == "cuda" + return model + + +def _assert_real_bf16_parity( + actual: torch.Tensor, + expected: torch.Tensor, + *, + scope: str, +) -> None: + diff = (expected - actual).abs() + expected_abs_mean = expected.abs().mean() + actual_abs_mean = actual.abs().mean() + abs_mean_drift = (actual_abs_mean - expected_abs_mean).abs() / expected_abs_mean.clamp_min(1e-6) + print(f"{scope} official_abs_mean={expected_abs_mean.item():.8f} " + f"fastvideo_abs_mean={actual_abs_mean.item():.8f} " + f"abs_mean_drift={abs_mean_drift.item():.4%} " + f"diff_max={diff.max().item():.8f} diff_mean={diff.mean().item():.8f}") + assert abs_mean_drift < 0.01 + assert diff.mean() < 0.01 + assert_close(actual, expected, atol=5e-2, rtol=5e-2) + + +def test_helios_distilled_config_defaults_match_distilled_variant(): + HeliosArchConfig, _, _ = _native_types() + config = HeliosArchConfig() + assert config.patch_size == (1, 2, 2) + assert config.num_attention_heads == 40 + assert config.attention_head_dim == 128 + assert config.hidden_size == 5120 + assert config.in_channels == config.out_channels == 16 + assert config.text_dim == 4096 + assert config.freq_dim == 256 + assert config.ffn_dim == 13824 + assert config.num_layers == 40 + assert config.rope_dim == (44, 42, 42) + assert config.zero_history_timestep is True + assert config.has_multi_term_memory_patch is True + assert config.guidance_cross_attn is True + + +def test_helios_distilled_config_matches_local_pinned_checkpoint(): + HeliosArchConfig, _, _ = _native_types() + config = HeliosArchConfig() + config_path = TRANSFORMER_DIR / "config.json" + if not config_path.is_file(): + pytest.skip( + "Pinned Helios transformer config is absent; set HELIOS_TRANSFORMER_DIR " + f"to BestWishYsh/Helios-Distilled@{HF_REVISION}/transformer") + checkpoint_config = json.loads(config_path.read_text(encoding="utf-8")) + assert checkpoint_config.pop("_class_name") == "HeliosTransformer3DModel" + checkpoint_config.pop("_diffusers_version", None) + arch_fields = {field.name for field in fields(HeliosArchConfig)} + assert set(checkpoint_config) <= arch_fields + for name, expected in checkpoint_config.items(): + actual = getattr(config, name) + if isinstance(actual, tuple): + expected = tuple(expected) + assert actual == expected, f"unexpected Helios config {name}={actual!r}" + + +@pytest.mark.parametrize( + ("history_name", "indices_name"), + [ + ("latents_history_short", "indices_latents_history_short"), + ("latents_history_mid", "indices_latents_history_mid"), + ("latents_history_long", "indices_latents_history_long"), + ], +) +@pytest.mark.parametrize("missing_input", ["history", "indices"]) +def test_helios_history_tensor_and_indices_must_be_paired( + history_name: str, + indices_name: str, + missing_input: str, +): + _, _, FastVideoHeliosTransformer = _native_types() + history = torch.empty(1) if missing_input == "indices" else None + indices = torch.empty(1, dtype=torch.long) if missing_input == "history" else None + + with pytest.raises(ValueError, match=rf"{history_name}.*{indices_name}"): + FastVideoHeliosTransformer._validate_history_pair(history, indices, history_name, indices_name) + + +@pytest.mark.parametrize( + ("unsupported_override", "message"), + [ + ({"cross_attn_norm": False}, "cross_attn_norm"), + ({"qk_norm": None}, "qk_norm"), + ({"added_kv_proj_dim": 64}, "added_kv_proj_dim"), + ({"guidance_cross_attn": False}, "guidance_cross_attn"), + ({"zero_history_timestep": False}, "zero_history_timestep"), + ({"has_multi_term_memory_patch": False}, "has_multi_term_memory_patch"), + ({"is_amplify_history": True}, "is_amplify_history"), + ({"history_scale_mode": "scalar"}, "history_scale_mode"), + ], +) +def test_helios_arch_config_rejects_unverified_variants( + unsupported_override: dict, + message: str, +): + """Variant knobs without parity evidence must fail instead of silently drifting.""" + HeliosArchConfig, _, _ = _native_types() + kwargs = _tiny_kwargs() + kwargs.update(unsupported_override) + with pytest.raises(ValueError, match=message): + HeliosArchConfig(**kwargs) + + +def test_helios_transformer_registry_resolves_native_class(): + from fastvideo.models.dits.helios import HeliosTransformer3DModel + from fastvideo.models.registry import ModelRegistry + + model_cls, architecture = ModelRegistry.resolve_model_cls("HeliosTransformer3DModel") + assert model_cls is HeliosTransformer3DModel + assert architecture == "HeliosTransformer3DModel" + + +def test_helios_real_checkpoint_uses_identity_key_mapping(monkeypatch): + HeliosArchConfig, HeliosConfig, FastVideoHeliosTransformer = _native_types() + del HeliosArchConfig + import fastvideo.models.dits.helios as fastvideo_helios + + monkeypatch.setattr(fastvideo_helios, "get_sp_world_size", lambda: 1) + index_path = TRANSFORMER_DIR / "diffusion_pytorch_model.safetensors.index.json" + if not index_path.exists(): + pytest.skip( + "Pinned Helios transformer index is absent; set HELIOS_TRANSFORMER_DIR " + f"to BestWishYsh/Helios-Distilled@{HF_REVISION}/transformer") + with torch.device("meta"): + native = FastVideoHeliosTransformer(config=HeliosConfig(), hf_config={}) + official_keys = set(json.loads(index_path.read_text(encoding="utf-8"))["weight_map"]) + native_keys = set(native.state_dict()) + assert native_keys == official_keys + assert native.param_names_mapping == {} + + +def test_helios_real_checkpoint_strict_loads(monkeypatch): + _, HeliosConfig, FastVideoHeliosTransformer = _native_types() + import fastvideo.models.dits.helios as fastvideo_helios + + monkeypatch.setattr(fastvideo_helios, "get_sp_world_size", lambda: 1) + if not _has_real_weights(): + pytest.skip(f"Pinned Helios transformer shards missing: {TRANSFORMER_DIR}") + files = resolve_safetensors_files(str(TRANSFORMER_DIR)) + with torch.device("meta"): + native = FastVideoHeliosTransformer(config=HeliosConfig(), hf_config={}) + incompatible = load_model_from_full_model_state_dict( + native, + safetensors_weights_iterator(files, to_cpu=True), + device=torch.device("cpu"), + param_dtype=torch.bfloat16, + strict=True, + param_names_mapping=get_param_names_mapping(native.param_names_mapping), + training_mode=False, + ) + assert incompatible.missing_keys == [] + assert incompatible.unexpected_keys == [] + native.materialize_non_persistent_buffers(torch.device("cpu"), torch.bfloat16) + assert not any(parameter.is_meta for parameter in native.parameters()) + assert not any(buffer.is_meta for buffer in native.buffers()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for full transformer parity.") +def test_helios_real_transformer_forward_parity(monkeypatch): + _, HeliosConfig, FastVideoHeliosTransformer = _native_types() + import fastvideo.models.dits.helios as fastvideo_helios + + monkeypatch.setattr(fastvideo_helios, "get_sp_world_size", lambda: 1) + if not _has_real_weights(): + pytest.skip( + "Pinned Helios transformer weights are absent; set HELIOS_TRANSFORMER_DIR " + f"to BestWishYsh/Helios-Distilled@{HF_REVISION}/transformer") + device = torch.device("cuda:0") + inputs = _make_real_inputs(device) + pyramid_inputs = _make_real_pyramid_inputs(device) + + official = (OfficialHeliosTransformer3DModel.from_pretrained( + str(TRANSFORMER_DIR), + local_files_only=True, + torch_dtype=torch.bfloat16, + ).to(device).eval()) + with torch.inference_mode(): + official_output = official(**inputs, return_dict=False)[0].float().cpu() + official_pyramid_output = official(**pyramid_inputs, return_dict=False)[0].float().cpu() + del official + gc.collect() + torch.cuda.empty_cache() + + del HeliosConfig, FastVideoHeliosTransformer + native = _load_fastvideo_production() + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): + fastvideo_output = native(**inputs).float().cpu() + fastvideo_pyramid_output = native(**pyramid_inputs).float().cpu() + + assert fastvideo_output.shape == official_output.shape == (1, 16, 1, 8, 8) + _assert_real_bf16_parity(fastvideo_output, official_output, scope="real transformer") + + assert (official_pyramid_output.shape == fastvideo_pyramid_output.shape == ( + 1, + 16, + 9, + 4, + 6, + )) + _assert_real_bf16_parity( + fastvideo_pyramid_output, + official_pyramid_output, + scope="real pyramid transformer", + ) + + +def test_helios_tiny_transformer_strict_load_and_forward_parity(monkeypatch): + HeliosArchConfig, HeliosConfig, FastVideoHeliosTransformer = _native_types() + import fastvideo.models.dits.helios as fastvideo_helios + + monkeypatch.setattr(fastvideo_helios, "get_sp_world_size", lambda: 1) + kwargs = _tiny_kwargs() + torch.manual_seed(17) + official = OfficialHeliosTransformer3DModel(**kwargs).float().eval() + fastvideo = (FastVideoHeliosTransformer( + config=HeliosConfig(arch_config=HeliosArchConfig(**kwargs)), + hf_config={}, + ).float().eval()) + + incompatible = load_model_from_full_model_state_dict( + fastvideo, + iter(official.state_dict().items()), + device=torch.device("cpu"), + param_dtype=torch.float32, + strict=True, + param_names_mapping=get_param_names_mapping(fastvideo.param_names_mapping), + training_mode=False, + ) + assert incompatible.missing_keys == [] + assert incompatible.unexpected_keys == [] + + inputs = _make_inputs() + with torch.inference_mode(): + official_output = official(**inputs, return_dict=False)[0] + with set_forward_context(current_timestep=0, attn_metadata=None): + fastvideo_output = fastvideo(**inputs) + + assert official_output.shape == fastvideo_output.shape == (1, 4, 2, 8, 8) + diff = (official_output - fastvideo_output).abs() + print(f"tiny transformer diff_max={diff.max().item():.8f} diff_mean={diff.mean().item():.8f}") + assert_close(fastvideo_output, official_output, atol=1e-5, rtol=1e-5) + + +def test_helios_tiny_transformer_pyramid_geometry_parity(monkeypatch): + """Current latents shrink per stage while history stays full resolution.""" + HeliosArchConfig, HeliosConfig, FastVideoHeliosTransformer = _native_types() + import fastvideo.models.dits.helios as fastvideo_helios + + monkeypatch.setattr(fastvideo_helios, "get_sp_world_size", lambda: 1) + kwargs = _tiny_kwargs() + torch.manual_seed(23) + official = OfficialHeliosTransformer3DModel(**kwargs).float().eval() + fastvideo = (FastVideoHeliosTransformer( + config=HeliosConfig(arch_config=HeliosArchConfig(**kwargs)), + hf_config={}, + ).float().eval()) + incompatible = load_model_from_full_model_state_dict( + fastvideo, + iter(official.state_dict().items()), + device=torch.device("cpu"), + param_dtype=torch.float32, + strict=True, + param_names_mapping=get_param_names_mapping(fastvideo.param_names_mapping), + training_mode=False, + ) + assert incompatible.missing_keys == [] + assert incompatible.unexpected_keys == [] + + inputs = _make_pyramid_inputs() + cross_attention_query_lengths = [] + + def record_cross_attention_query_length(module, args): + del module + cross_attention_query_lengths.append(args[0].shape[1]) + + handle = fastvideo.blocks[0].attn2.register_forward_pre_hook(record_cross_attention_query_length) + with torch.inference_mode(): + official_output = official(**inputs, return_dict=False)[0] + try: + with set_forward_context(current_timestep=0, attn_metadata=None): + fastvideo_output = fastvideo(**inputs) + finally: + handle.remove() + + assert official_output.shape == fastvideo_output.shape == (1, 4, 2, 2, 2) + assert cross_attention_query_lengths == [2] + diff = (official_output - fastvideo_output).abs() + print(f"pyramid transformer diff_max={diff.max().item():.8f} diff_mean={diff.mean().item():.8f}") + assert_close(fastvideo_output, official_output, atol=1e-5, rtol=1e-5) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for FlashAttention parity.") +def test_helios_tiny_flash_attention_matches_sdpa(monkeypatch): + """Both supported attention backends must execute with bounded BF16 drift.""" + pytest.importorskip("flash_attn", reason="Install optional flash-attn to verify the FLASH_ATTN backend.") + from fastvideo.platforms import AttentionBackendEnum + import fastvideo.models.dits.helios as fastvideo_helios + + monkeypatch.setattr(fastvideo_helios, "get_sp_world_size", lambda: 1) + device = torch.device("cuda:0") + inputs = _move_inputs(_make_inputs(), device, torch.bfloat16) + sdpa = _load_tiny_fastvideo(AttentionBackendEnum.TORCH_SDPA, device, torch.bfloat16) + flash = _load_tiny_fastvideo(AttentionBackendEnum.FLASH_ATTN, device, torch.bfloat16) + assert sdpa.blocks[0].attn1.attn.backend is AttentionBackendEnum.TORCH_SDPA + assert flash.blocks[0].attn1.attn.backend is AttentionBackendEnum.FLASH_ATTN + assert flash.blocks[0].attn2.attn.backend is AttentionBackendEnum.FLASH_ATTN + + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): + sdpa_output = sdpa(**inputs).float() + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): + flash_output = flash(**inputs).float() + + assert torch.isfinite(flash_output).all() + diff = (flash_output - sdpa_output).abs() + print(f"flash-vs-sdpa diff_max={diff.max().item():.8f} diff_mean={diff.mean().item():.8f}") + assert diff.mean() < 1e-2 + assert_close(flash_output, sdpa_output, atol=5e-2, rtol=5e-2) + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _run_sp_worker(mode: str, output_path: Path) -> None: + from fastvideo.distributed import ( + cleanup_dist_env_and_memory, + maybe_init_distributed_environment_and_model_parallel, + ) + from fastvideo.platforms import AttentionBackendEnum + + if mode not in {"single", "sp"}: + raise ValueError(f"Unsupported Helios SP worker mode: {mode}") + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + rank = int(os.environ.get("RANK", "0")) + sp_size = 1 if mode == "single" else SP_WORLD_SIZE + device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(device) + torch.manual_seed(3535) + torch.cuda.manual_seed_all(3535) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + + try: + maybe_init_distributed_environment_and_model_parallel(1, sp_size) + model = _load_tiny_fastvideo(AttentionBackendEnum.TORCH_SDPA, device, torch.float32) + inputs = _move_inputs(_make_inputs(), device, torch.float32) + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): + output = model(**inputs) + assert torch.isfinite(output).all() + if rank == 0: + torch.save({"output": output.detach().cpu()}, output_path) + dist.barrier() + finally: + cleanup_dist_env_and_memory() + + +def _run_torchrun( + script_path: Path, + mode: str, + nproc_per_node: int, + output_path: Path, +) -> None: + command = [ + sys.executable, + "-m", + "torch.distributed.run", + "--nnodes", + "1", + "--nproc_per_node", + str(nproc_per_node), + "--master_port", + str(_free_port()), + str(script_path), + "--helios-sp-worker", + "--mode", + mode, + "--output", + str(output_path), + ] + environment = os.environ.copy() + environment["DISABLE_SP"] = "0" + environment["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA" + process = subprocess.run(command, capture_output=True, text=True, env=environment) + if process.returncode != 0: + raise RuntimeError(f"{mode} worker failed with code {process.returncode}\n" + f"STDOUT:\n{process.stdout}\n" + f"STDERR:\n{process.stderr}") + + +def test_helios_tiny_sp2_matches_single_rank(tmp_path: Path): + """SP=2 must preserve the unpadded full output for unequal history geometry.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA required for Helios SP parity.") + if torch.cuda.device_count() < SP_WORLD_SIZE: + pytest.skip(f"Helios SP parity requires at least {SP_WORLD_SIZE} CUDA devices.") + script_path = Path(__file__).resolve() + single_path = tmp_path / "helios_single.pt" + sp_path = tmp_path / "helios_sp2.pt" + _run_torchrun(script_path, "single", 1, single_path) + _run_torchrun(script_path, "sp", SP_WORLD_SIZE, sp_path) + + single_output = torch.load(single_path, map_location="cpu", weights_only=True)["output"] + sp_output = torch.load(sp_path, map_location="cpu", weights_only=True)["output"] + assert single_output.shape == sp_output.shape == (1, 4, 2, 8, 8) + diff = (sp_output - single_output).abs() + print(f"sp2-vs-single diff_max={diff.max().item():.8f} diff_mean={diff.mean().item():.8f}") + assert_close(sp_output, single_output, atol=1e-5, rtol=1e-5) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--helios-sp-worker", action="store_true") + parser.add_argument("--mode", choices=["single", "sp"], default=None) + parser.add_argument("--output", type=str, default=None) + return parser.parse_args() + + +if __name__ == "__main__": + args = _parse_args() + if not args.helios_sp_worker: + raise SystemExit("This module is intended to be run by pytest.") + if args.mode is None or args.output is None: + raise SystemExit("--mode and --output are required in worker mode.") + _run_sp_worker(args.mode, Path(args.output)) diff --git a/tests/local_tests/vaes/test_helios_vae_parity.py b/tests/local_tests/vaes/test_helios_vae_parity.py new file mode 100644 index 0000000000..4f906b070e --- /dev/null +++ b/tests/local_tests/vaes/test_helios_vae_parity.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Helios exact-checkpoint Wan VAE reuse parity. + +Coverage scope: implementation_subcomponent. This is intentionally a real +weight test: matching config fields alone is insufficient evidence that the +existing FastVideo Wan VAE can safely serve Helios. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch +from diffusers import AutoencoderKLWan as OfficialAutoencoderKLWan +from safetensors.torch import load_file +from torch.testing import assert_close + +from fastvideo.configs.models.vaes.wanvae import WanVAEArchConfig, WanVAEConfig +from fastvideo.models.vaes.wanvae import AutoencoderKLWan + +REPO_ROOT = Path(__file__).resolve().parents[3] +VAE_DIR = REPO_ROOT / "official_weights" / "helios" / "vae" +PARITY_SCOPE = "implementation_subcomponent" + + +def _require_weights() -> Path: + weight_path = VAE_DIR / "diffusion_pytorch_model.safetensors" + if not weight_path.exists(): + pytest.skip(f"Helios VAE weights missing: {weight_path}") + return weight_path + + +def _fastvideo_vae(device: torch.device) -> AutoencoderKLWan: + config = WanVAEConfig(arch_config=WanVAEArchConfig()) + config.load_encoder = True + config.load_decoder = True + model = AutoencoderKLWan(config).to(device=device, dtype=torch.float32) + incompatible = model.load_state_dict(load_file(_require_weights()), strict=True) + assert incompatible.missing_keys == [] + assert incompatible.unexpected_keys == [] + return model.eval() + + +def test_helios_vae_config_matches_native_wan_candidate() -> None: + config = WanVAEArchConfig() + assert config.base_dim == 96 + assert config.z_dim == 16 + assert config.dim_mult == (1, 2, 4, 4) + assert config.temperal_downsample == (False, True, True) + assert config.scale_factor_temporal == 4 + assert config.scale_factor_spatial == 8 + assert len(config.latents_mean) == len(config.latents_std) == 16 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for VAE parity.") +def test_helios_vae_decode_parity() -> None: + device = torch.device("cuda:0") + official = (OfficialAutoencoderKLWan.from_pretrained(str(VAE_DIR), local_files_only=True, + torch_dtype=torch.float32).to(device).eval()) + fastvideo = _fastvideo_vae(device) + + generator = torch.Generator(device=device).manual_seed(4242) + latents = torch.randn(1, 16, 2, 8, 8, device=device, generator=generator) + with torch.inference_mode(): + official_output = official.decode(latents, return_dict=False)[0].float().cpu() + fastvideo_output = fastvideo.decode(latents).float().cpu() + + assert official_output.shape == fastvideo_output.shape + diff = (official_output - fastvideo_output).abs() + print(f"VAE decode diff_max={diff.max().item():.8f} diff_mean={diff.mean().item():.8f}") + assert_close(fastvideo_output, official_output, atol=1e-4, rtol=1e-4)