From 755a4e4470254337718dc234196f5b61044bc8cb Mon Sep 17 00:00:00 2001 From: Junda Su <90978028+Davids048@users.noreply.github.com> Date: Sat, 18 Jul 2026 20:18:16 -0700 Subject: [PATCH] [new-model] Add LingBot-Video Dense and MoE/refiner T2V inference (#1595) --- .../inference_schema_parity_inventory.yaml | 1 + .../inference/basic/basic_lingbot_video.py | 52 ++ fastvideo/api/compat.py | 2 + fastvideo/api/sampling_param.py | 7 + fastvideo/api/schema.py | 1 + fastvideo/configs/models/dits/__init__.py | 4 +- .../configs/models/dits/lingbot_video.py | 70 ++ fastvideo/configs/models/encoders/__init__.py | 3 +- .../configs/models/encoders/lingbot_video.py | 50 ++ fastvideo/configs/pipelines/__init__.py | 3 +- fastvideo/configs/pipelines/lingbot_video.py | 73 ++ fastvideo/entrypoints/video_generator.py | 46 +- fastvideo/models/dits/lingbot_video.py | 811 ++++++++++++++++++ fastvideo/models/encoders/lingbot_video.py | 221 +++++ fastvideo/models/loader/fsdp_load.py | 162 ++-- .../basic/glm_image/stages/decoding.py | 2 +- .../pipelines/basic/lingbot_video/__init__.py | 5 + .../lingbot_video/lingbot_video_pipeline.py | 106 +++ .../pipelines/basic/lingbot_video/presets.py | 92 ++ .../pipelines/basic/lingbot_video/stages.py | 345 ++++++++ fastvideo/pipelines/pipeline_batch_info.py | 1 + fastvideo/pipelines/stages/decoding.py | 24 +- fastvideo/registry.py | 26 + .../tests/api/test_compat_translation.py | 39 +- fastvideo/tests/api/test_parser.py | 61 +- fastvideo/tests/api/test_presets.py | 246 ++++-- .../distributed/test_sp_lingbot_video.py | 203 +++++ .../tests/entrypoints/test_video_generator.py | 73 +- .../lingbot_video_to_diffusers.py | 182 ++++ ...lingbot_video_dense_text_encoder_parity.py | 229 +++++ .../local_tests/lingbot_video/PORT_STATUS.md | 255 ++++++ tests/local_tests/lingbot_video/README.md | 369 ++++++++ tests/local_tests/lingbot_video/hf_assets.py | 80 ++ .../lingbot_video/t2v_example_1_prompt.json | 77 ++ .../lingbot_video/test_conversion_layout.py | 79 ++ .../models/test_fsdp_load_mixed_dtype.py | 172 ++++ .../test_lingbot_video_moe_pipeline_smoke.py | 65 ++ .../test_lingbot_video_pipeline_parity.py | 307 +++++++ .../test_lingbot_video_pipeline_smoke.py | 176 ++++ ...st_lingbot_video_refiner_pipeline_smoke.py | 70 ++ .../test_lingbot_video_refiner_stages.py | 246 ++++++ .../test_lingbot_video_scheduler_parity.py | 64 ++ .../test_lingbot_video_dense_dit_parity.py | 233 +++++ .../transformers/test_lingbot_video_moe.py | 229 +++++ .../test_lingbot_video_moe_block_parity.py | 136 +++ .../test_lingbot_video_moe_dit_parity.py | 167 ++++ .../vaes/test_lingbot_video_vae_parity.py | 86 ++ 47 files changed, 5747 insertions(+), 204 deletions(-) create mode 100644 examples/inference/basic/basic_lingbot_video.py create mode 100644 fastvideo/configs/models/dits/lingbot_video.py create mode 100644 fastvideo/configs/models/encoders/lingbot_video.py create mode 100644 fastvideo/configs/pipelines/lingbot_video.py create mode 100644 fastvideo/models/dits/lingbot_video.py create mode 100644 fastvideo/models/encoders/lingbot_video.py create mode 100644 fastvideo/pipelines/basic/lingbot_video/__init__.py create mode 100644 fastvideo/pipelines/basic/lingbot_video/lingbot_video_pipeline.py create mode 100644 fastvideo/pipelines/basic/lingbot_video/presets.py create mode 100644 fastvideo/pipelines/basic/lingbot_video/stages.py create mode 100644 fastvideo/tests/distributed/test_sp_lingbot_video.py create mode 100644 scripts/checkpoint_conversion/lingbot_video_to_diffusers.py create mode 100644 tests/local_tests/encoders/test_lingbot_video_dense_text_encoder_parity.py create mode 100644 tests/local_tests/lingbot_video/PORT_STATUS.md create mode 100644 tests/local_tests/lingbot_video/README.md create mode 100644 tests/local_tests/lingbot_video/hf_assets.py create mode 100644 tests/local_tests/lingbot_video/t2v_example_1_prompt.json create mode 100644 tests/local_tests/lingbot_video/test_conversion_layout.py create mode 100644 tests/local_tests/models/test_fsdp_load_mixed_dtype.py create mode 100644 tests/local_tests/pipelines/test_lingbot_video_moe_pipeline_smoke.py create mode 100644 tests/local_tests/pipelines/test_lingbot_video_pipeline_parity.py create mode 100644 tests/local_tests/pipelines/test_lingbot_video_pipeline_smoke.py create mode 100644 tests/local_tests/pipelines/test_lingbot_video_refiner_pipeline_smoke.py create mode 100644 tests/local_tests/pipelines/test_lingbot_video_refiner_stages.py create mode 100644 tests/local_tests/schedulers/test_lingbot_video_scheduler_parity.py create mode 100644 tests/local_tests/transformers/test_lingbot_video_dense_dit_parity.py create mode 100644 tests/local_tests/transformers/test_lingbot_video_moe.py create mode 100644 tests/local_tests/transformers/test_lingbot_video_moe_block_parity.py create mode 100644 tests/local_tests/transformers/test_lingbot_video_moe_dit_parity.py create mode 100644 tests/local_tests/vaes/test_lingbot_video_vae_parity.py diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 6def015f7..87d19ec24 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -470,6 +470,7 @@ surfaces: num_inference_steps: request.sampling.num_inference_steps num_inference_steps_sr: request.sampling.num_inference_steps_sr guidance_scale: request.sampling.guidance_scale + batch_cfg: request.sampling.batch_cfg guidance_scale_2: request.sampling.guidance_scale_2 cfg_normalization: request.sampling.cfg_normalization cfg_truncation: request.sampling.cfg_truncation diff --git a/examples/inference/basic/basic_lingbot_video.py b/examples/inference/basic/basic_lingbot_video.py new file mode 100644 index 000000000..a39fd0b4f --- /dev/null +++ b/examples/inference/basic/basic_lingbot_video.py @@ -0,0 +1,52 @@ +"""Generate a five-second Dense LingBot-Video clip with the official defaults.""" + +import argparse +from pathlib import Path + +from fastvideo import VideoGenerator + + +def parse_args() -> argparse.Namespace: + """Parse the converted checkpoint and output paths for the sample.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--model-path", + type=Path, + required=True, + help="Path to a converted Dense LingBot-Video checkpoint.", + ) + parser.add_argument( + "--output-path", + type=Path, + default=Path("outputs/lingbot-video/dense-t2v"), + help="Directory for the generated video.", + ) + return parser.parse_args() + + +def main() -> None: + """Load the converted Dense checkpoint and generate the default T2V sample.""" + args = parse_args() + generator = VideoGenerator.from_pretrained( + str(args.model_path), + num_gpus=1, + use_fsdp_inference=False, + text_encoder_cpu_offload=True, + vae_cpu_offload=False, + pin_cpu_memory=True, + ) + try: + generator.generate({ + "prompt": "A red fox runs through fresh snow at sunrise.", + "output": { + "output_path": str(args.output_path), + "save_video": True, + "return_frames": False, + }, + }) + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/fastvideo/api/compat.py b/fastvideo/api/compat.py index ea3ff10ac..eb44bb40a 100644 --- a/fastvideo/api/compat.py +++ b/fastvideo/api/compat.py @@ -304,6 +304,8 @@ def generator_config_to_fastvideo_args(config: GeneratorConfig | Mapping[str, An for key in _LTX2_REFINE_FLAT_KEYS: if key in refine: kwargs[f"ltx2_refine_{key}"] = refine[key] + if "enabled" in refine: + kwargs["refine_enabled"] = refine["enabled"] kwargs.update(preset_overrides) kwargs.update(deepcopy(normalized.pipeline.experimental)) return FastVideoArgs.from_kwargs(**kwargs) diff --git a/fastvideo/api/sampling_param.py b/fastvideo/api/sampling_param.py index 9db1de75e..d5d9d28b0 100644 --- a/fastvideo/api/sampling_param.py +++ b/fastvideo/api/sampling_param.py @@ -90,6 +90,7 @@ class SamplingParam: num_inference_steps: int = 50 num_inference_steps_sr: int = 50 guidance_scale: float = 1.0 + batch_cfg: bool = False guidance_scale_2: float | None = None # Z-Image CFG controls. ``cfg_normalization=True`` caps the guided # prediction norm at the positive-prediction norm; ``cfg_truncation`` @@ -341,6 +342,12 @@ class SamplingParam: default=SamplingParam.cfg_truncation, help="Disable Z-Image CFG above this normalized-noise threshold", ) + parser.add_argument( + "--batch-cfg", + action=StoreBoolean, + default=SamplingParam.batch_cfg, + help="Evaluate conditional and unconditional CFG branches in one batch", + ) parser.add_argument( "--guidance-rescale", type=float, diff --git a/fastvideo/api/schema.py b/fastvideo/api/schema.py index 825e096cf..8eede2fdc 100644 --- a/fastvideo/api/schema.py +++ b/fastvideo/api/schema.py @@ -149,6 +149,7 @@ class SamplingConfig: num_inference_steps: int = 50 num_inference_steps_sr: int = 50 guidance_scale: float = 1.0 + batch_cfg: bool = False guidance_scale_2: float | None = None cfg_normalization: bool = False cfg_truncation: float | None = 1.0 diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index cc6b7322c..1c2f97a1a 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -16,10 +16,12 @@ from fastvideo.configs.models.dits.zimage import ZImageDiTConfig from fastvideo.configs.models.dits.hyworld import HYWorldConfig from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig from fastvideo.configs.models.dits.lingbotworld2 import LingBotWorld2CausalFastVideoConfig +from fastvideo.configs.models.dits.lingbot_video import LingBotVideoConfig __all__ = [ "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig", "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "FluxDiTConfig", "Flux2Config", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", - "StableAudioConfig", "GlmImageDiTConfig", "LingBotWorld2CausalFastVideoConfig", "ZImageDiTConfig" + "StableAudioConfig", "GlmImageDiTConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotVideoConfig", + "ZImageDiTConfig" ] diff --git a/fastvideo/configs/models/dits/lingbot_video.py b/fastvideo/configs/models/dits/lingbot_video.py new file mode 100644 index 000000000..a9d883758 --- /dev/null +++ b/fastvideo/configs/models/dits/lingbot_video.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Architecture configuration for LingBot-Video Dense and MoE DiTs.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig +from fastvideo.platforms import AttentionBackendEnum + + +def _is_lingbot_video_block(name: str, module: object) -> bool: + """Select top-level transformer blocks for FSDP and compilation.""" + del module + parts = name.split(".") + return len(parts) == 2 and parts[0] == "blocks" and parts[1].isdigit() + + +@dataclass +class LingBotVideoArchConfig(DiTArchConfig): + """One-to-one representation of the released transformer config JSON.""" + + _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_lingbot_video_block]) + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + AttentionBackendEnum.TORCH_SDPA, + AttentionBackendEnum.FLASH_ATTN, + ) + param_names_mapping: dict = field(default_factory=lambda: {r"^(.*)$": r"\1"}) + + patch_size: tuple[int, int, int] = (1, 2, 2) + in_channels: int = 16 + out_channels: int = 16 + hidden_size: int = 2048 + num_attention_heads: int = 16 + depth: int = 24 + intermediate_size: int = 6144 + text_dim: int = 2560 + freq_dim: int = 256 + norm_eps: float = 1e-6 + rope_theta: float = 256.0 + axes_dims: tuple[int, int, int] = (32, 48, 48) + axes_lens: tuple[int, int, int] = (8192, 1024, 1024) + qkv_bias: bool = False + out_bias: bool = True + patch_embed_bias: bool = True + timestep_mlp_bias: bool = True + num_experts: int = 0 + num_experts_per_tok: int = 8 + moe_intermediate_size: int = 512 + decoder_sparse_step: int = 1 + mlp_only_layers: tuple[int, ...] = () + n_shared_experts: int | None = None + score_func: str = "sigmoid" + norm_topk_prob: bool = True + n_group: int | None = None + topk_group: int | None = None + routed_scaling_factor: float = 1.0 + + def __post_init__(self) -> None: + """Populate FastVideo loader fields from the released architecture.""" + super().__post_init__() + self.num_channels_latents = self.in_channels + self.attention_head_dim = self.hidden_size // self.num_attention_heads + + +@dataclass +class LingBotVideoConfig(DiTConfig): + """FastVideo component configuration for LingBot-Video transformers.""" + + arch_config: DiTArchConfig = field(default_factory=LingBotVideoArchConfig) diff --git a/fastvideo/configs/models/encoders/__init__.py b/fastvideo/configs/models/encoders/__init__.py index 644726e96..26c43ee0c 100644 --- a/fastvideo/configs/models/encoders/__init__.py +++ b/fastvideo/configs/models/encoders/__init__.py @@ -10,6 +10,7 @@ from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1 from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig +from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig from fastvideo.configs.models.encoders.stable_audio_conditioner import (StableAudioConditionerArchConfig, StableAudioConditionerConfig) from fastvideo.configs.models.encoders.t5gemma import T5GemmaEncoderConfig @@ -19,5 +20,5 @@ __all__ = [ "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig", "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig", "StableAudioConditionerArchConfig", "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig", - "LingBotWorld2UMT5ArchConfig", "LingBotWorld2UMT5Config" + "LingBotWorld2UMT5ArchConfig", "LingBotWorld2UMT5Config", "LingBotVideoQwen3VLTextConfig" ] diff --git a/fastvideo/configs/models/encoders/lingbot_video.py b/fastvideo/configs/models/encoders/lingbot_video.py new file mode 100644 index 000000000..a2ed28443 --- /dev/null +++ b/fastvideo/configs/models/encoders/lingbot_video.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen3-VL text-only encoder configuration used by LingBot-Video.""" + +from dataclasses import dataclass, field + +from fastvideo.configs.models.encoders.base import TextEncoderArchConfig +from fastvideo.configs.models.encoders.qwen3 import Qwen3TextArchConfig, Qwen3TextConfig + + +@dataclass +class LingBotVideoQwen3VLTextArchConfig(Qwen3TextArchConfig): + """Exact Qwen3-VL language-model architecture released with LingBot-Video.""" + + architectures: list[str] = field(default_factory=lambda: ["LingBotVideoQwen3VLTextModel"]) + vocab_size: int = 151936 + hidden_size: int = 2560 + intermediate_size: int = 9728 + num_hidden_layers: int = 36 + num_attention_heads: int = 32 + num_key_value_heads: int = 8 + max_position_embeddings: int = 262144 + rms_norm_eps: float = 1e-6 + rope_theta: float = 5000000.0 + rope_scaling: dict | None = None + mrope_interleaved: bool = True + mrope_section: tuple[int, int, int] = (24, 20, 20) + bos_token_id: int = 151643 + eos_token_id: int = 151645 + pad_token_id: int = 151643 + text_len: int = 37698 + output_hidden_states: bool = True + require_processor: bool = True + + def __post_init__(self) -> None: + """Match the official processor call used by LingBotVideoPipeline.""" + self.tokenizer_kwargs = { + "truncation": True, + "max_length": self.text_len, + "padding": "longest", + "return_tensors": "pt", + } + + +@dataclass +class LingBotVideoQwen3VLTextConfig(Qwen3TextConfig): + """FastVideo loader config for the LingBot-Video text-only Qwen3-VL path.""" + + arch_config: TextEncoderArchConfig = field(default_factory=LingBotVideoQwen3VLTextArchConfig) + prefix: str = "language_model" + is_chat_model: bool = False diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index d05fb5541..374cb1931 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -8,6 +8,7 @@ from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelin from fastvideo.configs.pipelines.hyworld import HYWorldConfig from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig +from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig @@ -20,6 +21,6 @@ __all__ = [ "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig", "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig", - "Kandinsky5I2VConfig", "LingBotWorld2CausalFastI2V480PConfig", "MatrixGame2I2V480PConfig", + "Kandinsky5I2VConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name" ] diff --git a/fastvideo/configs/pipelines/lingbot_video.py b/fastvideo/configs/pipelines/lingbot_video.py new file mode 100644 index 000000000..2a18a0fa7 --- /dev/null +++ b/fastvideo/configs/pipelines/lingbot_video.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Dense LingBot-Video T2V pipeline configuration.""" + +from collections.abc import Callable +from dataclasses import dataclass, field + +import torch + +from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig +from fastvideo.configs.models.dits.lingbot_video import LingBotVideoConfig +from fastvideo.configs.models.encoders import BaseEncoderOutput +from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig +from fastvideo.configs.models.vaes import WanVAEConfig +from fastvideo.configs.pipelines.base import PipelineConfig + +PROMPT_CROP_START = 140 +PROMPT_TEMPLATE = ("<|im_start|>system\nGiven a user input that may include a text prompt alone, " + "a text prompt with an image reference, or a text prompt with a video reference " + 'or a video reference alone, generate an "Enhanced prompt" that provides detailed ' + "visual descriptions suitable for video generation. Evaluate the level of detail " + "in the user's input: if it is simple, enrich it by adding specifics about colors, " + "shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal " + "progression, and spatial relationships to create vivid, concrete, and temporally " + "coherent scenes to create vivid and concrete scenes. Please generate only the " + "enhanced description for the prompt below and avoid including any additional " + "commentary or evaluations:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n" + "<|im_start|>assistant\n") + + +def preprocess_lingbot_video_prompt(prompt: str) -> str: + """Apply the released T2V system/user/assistant prompt template.""" + return PROMPT_TEMPLATE.format(prompt) + + +def postprocess_lingbot_video_text( + outputs: BaseEncoderOutput, + attention_mask: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Select the final hidden state, crop the template, and trim batch-one padding.""" + if outputs.hidden_states is None: + raise ValueError("LingBot-Video requires text-encoder hidden states") + prompt_embeds = outputs.hidden_states[-1][:, PROMPT_CROP_START:] + prompt_mask = attention_mask[:, PROMPT_CROP_START:] + if prompt_embeds.shape[0] == 1: + true_length = int(prompt_mask[0].sum().item()) + prompt_embeds = prompt_embeds[:, :true_length] + prompt_mask = prompt_mask[:, :true_length] + return prompt_embeds, prompt_mask + + +@dataclass +class LingBotVideoT2VConfig(PipelineConfig): + """Released Dense T2V component wiring and numerical precision policy.""" + + dit_config: DiTConfig = field(default_factory=LingBotVideoConfig) + vae_config: VAEConfig = field(default_factory=WanVAEConfig) + text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (LingBotVideoQwen3VLTextConfig(), )) + preprocess_text_funcs: tuple[Callable, ...] = field(default_factory=lambda: (preprocess_lingbot_video_prompt, )) + postprocess_text_funcs: tuple[Callable, ...] = field(default_factory=lambda: (postprocess_lingbot_video_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 = 3.0 + embedded_cfg_scale: float | None = None + scheduler_step_in_fp32: bool = True + + def __post_init__(self) -> None: + """Load only the VAE decoder for the T2V workload.""" + self.vae_config.load_encoder = False + self.vae_config.load_decoder = True diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 40203dd00..777e0ef08 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -112,6 +112,40 @@ def _infer_latent_batch_size(batch: ForwardBatch) -> int: return latent_batch_size +def _resolve_output_size( + samples: torch.Tensor, + fallback: tuple[int, int, int], + *, + pixel_output: bool, +) -> tuple[int, int, int]: + """Report the final decoded video's `(height, width, frames)`. + + Refiner stages can produce a different resolution from the base request, so + pixel outputs use the final `[batch, channels, frames, height, width]` tensor. + Latent and audio outputs keep the requested fallback because their tensor + dimensions do not describe decoded pixels. + """ + if pixel_output and samples.ndim == 5: + return (int(samples.shape[-2]), int(samples.shape[-1]), int(samples.shape[-3])) + return fallback + + +def _validate_request_stage_overrides(model_path: str, request: GenerationRequest) -> None: + """Validate typed stage overrides against the model's registered preset.""" + if not request.stage_overrides: + return + from fastvideo.api.presets import validate_preset_selection + from fastvideo.registry import get_preset_selection + preset_name, model_family = get_preset_selection(model_path) + if preset_name is None or model_family is None: + raise ValueError(f"Model {model_path!r} has no preset for stage override validation") + validate_preset_selection( + preset_name, + model_family, + stage_overrides=request.stage_overrides, + ) + + class VideoGenerator: """ A unified class for generating videos using diffusion models. @@ -443,6 +477,7 @@ class VideoGenerator: self, request: GenerationRequest, ) -> GenerationResult | list[GenerationResult]: + _validate_request_stage_overrides(self.fastvideo_args.model_path, request) if isinstance(request.prompt, list): if request.inputs.prompt_path is not None: raise ValueError("request.prompt list cannot be combined with request.inputs.prompt_path") @@ -808,6 +843,15 @@ class VideoGenerator: # 2. Audio-only workload — `samples` is a 1×3×1×8×8 placeholder # no caller will use; skip the grid loop and save a `.wav`. # 3. Pixel video / image — the historical happy path. + # `GenerationResult.size` describes the produced media, not only the + # base-stage request. Refiner pipelines can change the final pixel + # dimensions, so derive this result metadata from the decoded output. + output_size = _resolve_output_size( + samples, + (target_height, target_width, batch.num_frames), + pixel_output=not is_latent_output and not audio_only, + ) + postprocess_start = time.perf_counter() frames: list[np.ndarray] | None if is_latent_output or audio_only: @@ -916,7 +960,7 @@ class VideoGenerator: "audio": output_batch.extra.get("audio"), "audio_sample_rate": output_batch.extra.get("audio_sample_rate"), "ltx2_audio_latents": output_batch.extra.get("ltx2_audio_latents"), - "size": (target_height, target_width, batch.num_frames), + "size": output_size, "generation_time": gen_time, "e2e_latency": e2e_time, "logging_info": logging_info, diff --git a/fastvideo/models/dits/lingbot_video.py b/fastvideo/models/dits/lingbot_video.py new file mode 100644 index 000000000..6aff7cd56 --- /dev/null +++ b/fastvideo/models/dits/lingbot_video.py @@ -0,0 +1,811 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastVideo-native LingBot-Video Dense and MoE diffusion transformers.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.configs.models.dits.lingbot_video import LingBotVideoConfig, _is_lingbot_video_block +from fastvideo.distributed.communication_op import ( + sequence_model_parallel_all_gather_with_unpad, + sequence_model_parallel_all_to_all_4D, + sequence_model_parallel_shard, +) +from fastvideo.distributed.parallel_state import get_sp_world_size, model_parallel_is_initialized +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.visual_embedding import Timesteps +from fastvideo.models.dits.base import BaseDiT +from fastvideo.platforms import AttentionBackendEnum + + +@dataclass +class LingBotVideoTransformerOutput: + """Output container matching the released transformer contract.""" + + sample: torch.Tensor + + +_FP32_MODULE_NAMES = ( + "time_embedder", + "time_modulation", + "scale_shift_table", + "norm", + "norm1", + "norm2", + "norm_q", + "norm_k", + "norm_post_attn", + "norm_post_ffn", + "norm_out", + "norm_out_modulation", + "router", +) + + +def _keep_in_fp32(name: str) -> bool: + """Return whether a released checkpoint module keeps fp32 parameters.""" + return any(module_name in name.split(".") for module_name in _FP32_MODULE_NAMES) + + +def _sequence_parallel_world_size() -> int: + """Use standalone single-rank behavior before distributed initialization.""" + return get_sp_world_size() if model_parallel_is_initialized() else 1 + + +class LingBotVideoLinear(ReplicatedLinear): + """Replicated FastVideo linear with the tensor-only official call contract.""" + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Return the projected tensor while preserving the normal weight surface.""" + output, _ = super().forward(hidden_states) + return output + + +class LingBotVideoRMSNorm(nn.Module): + """RMSNorm with fp32 accumulation and input-dtype output.""" + + def __init__(self, dim: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Normalize the last dimension using the official accumulation order.""" + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return (self.weight * hidden_states).to(input_dtype) + + +def _apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + """Apply complex 3D rotary embeddings to `(B, S, H, D)` tensors.""" + with torch.amp.autocast("cuda", enabled=False): + x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + output = torch.view_as_real(x_complex * freqs_cis.unsqueeze(2)).flatten(3) + return output.type_as(x) + + +class LingBotVideoRotaryEmbedding(nn.Module): + """Complex64 rotary table indexed by temporal and spatial positions.""" + + def __init__(self, axes_dims: tuple[int, ...], axes_lens: tuple[int, ...], theta: float) -> None: + super().__init__() + self.axes_dims = tuple(axes_dims) + self.axes_lens = list(axes_lens) + self.theta = theta + self.freqs_cis: list[torch.Tensor] | None = None + + @staticmethod + def _precompute(dims: tuple[int, ...], lengths: tuple[int, ...], theta: float) -> list[torch.Tensor]: + """Build the per-axis complex frequency tables on CPU.""" + tables: list[torch.Tensor] = [] + for dim, length in zip(dims, lengths, strict=True): + frequencies = 1.0 / (theta**(torch.arange(0, dim, 2, dtype=torch.float64, device="cpu") / dim)) + positions = torch.arange(length, device=frequencies.device, dtype=torch.float64) + phases = torch.outer(positions, frequencies).float() + tables.append(torch.polar(torch.ones_like(phases), phases).to(torch.complex64)) + return tables + + def forward( + self, + position_ids: torch.Tensor, + maxima: tuple[int, ...] | None = None, + ) -> torch.Tensor: + """Gather and concatenate rotary frequencies for `(S, 3)` positions.""" + device = position_ids.device + if maxima is None: + maxima = tuple(int(value) for value in position_ids.max(dim=0).values.tolist()) + rebuild = self.freqs_cis is None or any(maximum >= length + for maximum, length in zip(maxima, self.axes_lens, strict=True)) + if rebuild: + for index, maximum in enumerate(maxima): + if maximum >= self.axes_lens[index]: + self.axes_lens[index] = int(maximum * 1.5) + 1 + self.freqs_cis = self._precompute(self.axes_dims, tuple(self.axes_lens), self.theta) + self.freqs_cis = [table.to(device) for table in self.freqs_cis] + elif self.freqs_cis[0].device != device: + self.freqs_cis = [table.to(device) for table in self.freqs_cis] + return torch.cat( + [self.freqs_cis[index][position_ids[:, index]] for index in range(len(self.axes_dims))], + dim=-1, + ) + + +def _make_joint_position_ids( + text_len: int, + grid_t: int, + grid_h: int, + grid_w: int, + device: torch.device, +) -> torch.Tensor: + """Create official `[video; text]` 3D positions for one sample.""" + temporal = torch.arange(grid_t, device=device, dtype=torch.int32) + text_len + 1 + height = torch.arange(grid_h, device=device, dtype=torch.int32) + width = torch.arange(grid_w, device=device, dtype=torch.int32) + video_positions = torch.stack(torch.meshgrid(temporal, height, width, indexing="ij"), dim=-1).flatten(0, 2) + text_temporal = torch.arange(text_len, device=device, dtype=torch.int32) + 1 + text_positions = torch.stack( + [text_temporal, torch.zeros_like(text_temporal), + torch.zeros_like(text_temporal)], dim=-1) + return torch.cat([video_positions, text_positions], dim=0) + + +class LingBotVideoTextEmbedder(nn.Module): + """Project Qwen3-VL hidden states into the DiT hidden dimension.""" + + def __init__(self, text_dim: int, hidden_size: int) -> None: + super().__init__() + self.norm = LingBotVideoRMSNorm(text_dim, eps=1e-6) + self.linear_1 = LingBotVideoLinear(text_dim, hidden_size, bias=True) + self.linear_2 = LingBotVideoLinear(hidden_size, hidden_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Apply RMSNorm followed by the released two-layer SiLU projection.""" + hidden_states = self.norm(hidden_states) + return self.linear_2(F.silu(self.linear_1(hidden_states))) + + +class LingBotVideoAttention(nn.Module): + """Joint video-text attention shared by the Dense and MoE variants.""" + + def __init__( + self, + hidden_size: int, + num_heads: int, + norm_eps: float, + qkv_bias: bool, + out_bias: bool, + ) -> None: + """Create released QKV projections, per-head norms, and output projection.""" + super().__init__() + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.to_q = LingBotVideoLinear(hidden_size, hidden_size, bias=qkv_bias) + self.to_k = LingBotVideoLinear(hidden_size, hidden_size, bias=qkv_bias) + self.to_v = LingBotVideoLinear(hidden_size, hidden_size, bias=qkv_bias) + self.norm_q = LingBotVideoRMSNorm(self.head_dim, norm_eps) + self.norm_k = LingBotVideoRMSNorm(self.head_dim, norm_eps) + self.to_out = LingBotVideoLinear(hidden_size, hidden_size, bias=out_bias) + + def forward( + self, + hidden_states: torch.Tensor, + rotary_emb: torch.Tensor, + attention_mask: torch.Tensor | None = None, + original_seq_len: int | None = None, + ) -> torch.Tensor: + """Project QKV, apply rotary embeddings, and run non-causal SDPA.""" + batch, sequence, _ = hidden_states.shape + query = self.to_q(hidden_states).unflatten(2, (self.num_heads, self.head_dim)) + key = self.to_k(hidden_states).unflatten(2, (self.num_heads, self.head_dim)) + value = self.to_v(hidden_states).unflatten(2, (self.num_heads, self.head_dim)) + query = _apply_rotary_emb(self.norm_q(query), rotary_emb) + key = _apply_rotary_emb(self.norm_k(key), rotary_emb) + + if original_seq_len is not None: + # Attention needs the full unpadded joint sequence while projections + # and residual blocks stay sharded over tokens. + qkv = torch.cat([query, key, value], dim=0) + qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1) + padded_seq_len = qkv.shape[1] + query, key, value = qkv[:, :original_seq_len].chunk(3, dim=0) + output = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2) + if original_seq_len is not None: + output = F.pad(output, (0, 0, 0, 0, 0, padded_seq_len - original_seq_len)) + output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2) + return self.to_out(output.reshape(batch, sequence, -1).type_as(hidden_states)) + + +class LingBotVideoMLP(nn.Module): + """Dense SwiGLU feed-forward network.""" + + def __init__(self, hidden_size: int, intermediate_size: int) -> None: + super().__init__() + self.gate_proj = LingBotVideoLinear(hidden_size, intermediate_size, bias=False) + self.up_proj = LingBotVideoLinear(hidden_size, intermediate_size, bias=False) + self.down_proj = LingBotVideoLinear(intermediate_size, hidden_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Apply the released SiLU-gated MLP ordering.""" + return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)) + + +class LingBotVideoRouter(nn.Module): + """Released token-choice router with bias-only expert selection correction.""" + + def __init__( + self, + hidden_size: int, + num_experts: int, + top_k: int, + score_func: str, + norm_topk_prob: bool, + n_group: int | None, + topk_group: int | None, + route_scale: float, + ) -> None: + """Create fp32-routed expert scores with the released persistent bias.""" + super().__init__() + self.num_experts = num_experts + self.top_k = top_k + self.score_func = score_func + self.norm_topk_prob = norm_topk_prob + self.n_group = n_group + self.topk_group = topk_group + self.route_scale = route_scale + self.weight = nn.Parameter(torch.empty(num_experts, hidden_size)) + self.register_buffer("e_score_correction_bias", torch.zeros(num_experts), persistent=True) + + def _group_limited_topk(self, scores_for_choice: torch.Tensor) -> torch.Tensor: + """Restrict token choices to groups with the two strongest expert scores.""" + sequence_length = scores_for_choice.shape[0] + experts_per_group = self.num_experts // self.n_group + grouped = scores_for_choice.view(sequence_length, self.n_group, experts_per_group) + group_scores = grouped.topk(2, dim=-1)[0].sum(dim=-1) + group_indices = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_indices, 1) + score_mask = (group_mask.unsqueeze(-1).expand(sequence_length, self.n_group, + experts_per_group).reshape(sequence_length, -1)) + masked_scores = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf")) + return torch.topk(masked_scores, k=self.top_k, dim=-1, sorted=False)[1] + + def forward(self, + tokens: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Score in fp32, select with correction bias, and weight without it.""" + with torch.amp.autocast(tokens.device.type, enabled=False): + logits = F.linear(tokens.float(), self.weight.float()) + scores = F.softmax(logits, dim=-1) if self.score_func == "softmax" else logits.sigmoid() + scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0) + if self.n_group is not None and self.n_group > 1: + top_indices = self._group_limited_topk(scores_for_choice) + else: + top_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1] + top_scores = scores.gather(1, top_indices) + if self.top_k > 1 and self.norm_topk_prob: + top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20) + top_scores = top_scores * self.route_scale + return top_indices, top_scores.to(tokens.dtype), logits, scores, scores_for_choice + + +class LingBotVideoGroupedExperts(nn.Module): + """Released grouped-expert parameter layout: w1/w3 `[E,I,H]`, w2 `[E,H,I]`.""" + + def __init__(self, num_experts: int, hidden_size: int, intermediate_size: int) -> None: + super().__init__() + self.num_experts = num_experts + self.w1 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size)) + self.w2 = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size)) + self.w3 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size)) + + +def _round_up_to_multiple(value: int, multiple: int) -> int: + """Round an integer up to the next multiple.""" + return ((value + multiple - 1) // multiple) * multiple + + +class LingBotVideoSparseMoeBlock(nn.Module): + """Token-choice sparse feed-forward block matching the released state surface.""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + num_experts: int, + top_k: int, + moe_intermediate_size: int, + score_func: str, + norm_topk_prob: bool, + n_group: int | None, + topk_group: int | None, + routed_scaling_factor: float, + n_shared_experts: int | None, + ) -> None: + """Create routed and optional shared experts with released parameter names.""" + super().__init__() + del intermediate_size + self.hidden_size = hidden_size + self.num_experts = num_experts + self.router = LingBotVideoRouter( + hidden_size, + num_experts, + top_k, + score_func, + norm_topk_prob, + n_group, + topk_group, + routed_scaling_factor, + ) + self.experts = LingBotVideoGroupedExperts(num_experts, hidden_size, moe_intermediate_size) + self.shared_experts = None + if n_shared_experts is not None and n_shared_experts > 0: + self.shared_experts = LingBotVideoMLP(hidden_size, moe_intermediate_size * n_shared_experts) + + @staticmethod + def _reorder_tokens( + tokens: torch.Tensor, + top_scores: torch.Tensor, + top_indices: torch.Tensor, + num_experts: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]: + """Pack active token choices into stable expert-major order.""" + num_tokens = tokens.shape[0] + top_k = top_indices.shape[1] + flat_scores = top_scores.reshape(-1) + flat_indices = top_indices.reshape(-1) + active_positions = torch.where(flat_scores != 0)[0] + active_experts = flat_indices[active_positions] + counts = torch.zeros(num_experts, device=tokens.device, dtype=torch.int64) + counts.scatter_add_(0, active_experts, torch.ones_like(active_experts, dtype=torch.int64)) + sort_order = torch.argsort(active_experts, stable=True) + sorted_positions = active_positions[sort_order] + sorted_scores = flat_scores[sorted_positions] + original_token_indices = sorted_positions // top_k + permuted_tokens = tokens[original_token_indices] + return permuted_tokens, counts, sorted_positions, sorted_scores, num_tokens, top_k + + @staticmethod + def _pad_grouped_tokens(tokens: torch.Tensor, + counts: torch.Tensor, + align: int = 8) -> tuple[torch.Size, torch.Tensor, torch.Tensor, torch.Tensor]: + """Align each expert segment for `torch._grouped_mm` and retain unpad indices.""" + num_tokens = tokens.shape[0] + num_experts = int(counts.shape[0]) + max_length = _round_up_to_multiple(num_tokens + num_experts * align, align) + # Build the small padding index on CPU after one synchronization instead + # of reading three CUDA scalars for every expert. + counts_cpu = counts.to(device="cpu", dtype=torch.int64) + total_per_expert = torch.clamp_min(counts_cpu, align) + aligned_counts_cpu = ((total_per_expert + align - 1) // align * align).to(torch.int32) + write_offsets = torch.cumsum(aligned_counts_cpu, dim=0) - aligned_counts_cpu + start_indices = torch.cumsum(counts_cpu, dim=0) - counts_cpu + permuted_indices_cpu = torch.full((max_length, ), num_tokens, dtype=torch.int64, device="cpu") + for expert_index in range(num_experts): + length = int(counts_cpu[expert_index]) + if length == 0: + continue + write_start = int(write_offsets[expert_index]) + start = int(start_indices[expert_index]) + permuted_indices_cpu[write_start:write_start + length] = torch.arange(start, + start + length, + device="cpu", + dtype=torch.int64) + permuted_indices = permuted_indices_cpu.to(tokens.device) + aligned_counts = aligned_counts_cpu.to(tokens.device) + tokens_with_pad = torch.vstack((tokens, tokens.new_zeros((tokens.shape[-1], )))) + input_shape = tokens_with_pad.shape + return input_shape, tokens_with_pad[permuted_indices], permuted_indices, aligned_counts + + @staticmethod + def _unpad_grouped_tokens(output: torch.Tensor, input_shape: torch.Size, + permuted_indices: torch.Tensor) -> torch.Tensor: + """Undo per-expert alignment while dropping the shared padding row.""" + unpermuted = output.new_empty(input_shape) + unpermuted[permuted_indices, :] = output + return unpermuted[:-1] + + def _run_grouped_experts(self, tokens: torch.Tensor, counts: torch.Tensor) -> torch.Tensor: + """Use the released bf16 grouped matmuls on CUDA and an eager CPU fallback.""" + if tokens.device.type == "cpu" or not hasattr(torch, "_grouped_mm"): + return self._run_experts_for_loop(tokens, counts) + input_shape, padded_tokens, permuted_indices, aligned_counts = self._pad_grouped_tokens(tokens, counts) + offsets = torch.cumsum(aligned_counts, dim=0, dtype=torch.int32) + hidden = F.silu( + torch._grouped_mm( + padded_tokens.bfloat16(), + self.experts.w1.bfloat16().transpose(-2, -1), + offs=offsets, + )) + hidden = hidden * torch._grouped_mm( + padded_tokens.bfloat16(), + self.experts.w3.bfloat16().transpose(-2, -1), + offs=offsets, + ) + output = torch._grouped_mm( + hidden, + self.experts.w2.bfloat16().transpose(-2, -1), + offs=offsets, + ).type_as(padded_tokens) + return self._unpad_grouped_tokens(output, input_shape, permuted_indices) + + def _run_experts_for_loop(self, tokens: torch.Tensor, counts: torch.Tensor) -> torch.Tensor: + """Evaluate contiguous expert segments eagerly for CPU correctness tests.""" + splits = torch.split(tokens, counts.tolist(), dim=0) + outputs: list[torch.Tensor] = [] + for expert_index, expert_tokens in enumerate(splits): + if expert_tokens.numel() == 0: + continue + hidden = F.silu(expert_tokens @ self.experts.w1[expert_index].transpose(-2, -1)) + hidden = hidden * (expert_tokens @ self.experts.w3[expert_index].transpose(-2, -1)) + outputs.append(hidden @ self.experts.w2[expert_index].transpose(-2, -1)) + if not outputs: + return tokens.new_zeros(tokens.shape) + return torch.cat(outputs, dim=0) + + @staticmethod + def _restore_tokens( + expert_output: torch.Tensor, + sorted_positions: torch.Tensor, + sorted_scores: torch.Tensor, + num_tokens: int, + top_k: int, + ) -> torch.Tensor: + """Restore token order and combine expert outputs with fp32 weighted sums.""" + hidden_size = expert_output.shape[-1] + unsorted = torch.zeros( + (num_tokens * top_k, hidden_size), + dtype=expert_output.dtype, + device=expert_output.device, + ) + unsorted[sorted_positions] = expert_output + unsorted = unsorted.reshape(num_tokens, top_k, hidden_size) + scores_unsorted = torch.zeros( + num_tokens * top_k, + dtype=sorted_scores.dtype, + device=sorted_scores.device, + ) + scores_unsorted[sorted_positions] = sorted_scores + scores_unsorted = scores_unsorted.reshape(num_tokens, top_k, 1) + return (unsorted.float() * scores_unsorted).sum(dim=1).to(expert_output.dtype) + + def _run_selected_experts( + self, + tokens: torch.Tensor, + top_scores: torch.Tensor, + top_indices: torch.Tensor, + ) -> torch.Tensor: + """Dispatch routed choices, execute experts, and restore token-major order.""" + permuted_tokens, counts, sorted_positions, sorted_scores, num_tokens, top_k = self._reorder_tokens( + tokens, top_scores, top_indices, self.router.num_experts) + expert_output = self._run_grouped_experts(permuted_tokens, counts) + return self._restore_tokens(expert_output, sorted_positions, sorted_scores, num_tokens, top_k) + + def forward(self, hidden_states: torch.Tensor, padding_mask: torch.Tensor | None = None) -> torch.Tensor: + """Route token choices, zero padded routes, and add optional shared experts.""" + batch = hidden_states.shape[0] + tokens = hidden_states.view(-1, self.hidden_size) + top_indices, top_scores, logits, scores, scores_for_choice = self.router(tokens) + del logits, scores, scores_for_choice + if padding_mask is not None: + mask = padding_mask.unsqueeze(-1).to(top_scores.dtype) + top_scores = top_scores * mask + top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-9) + top_scores = top_scores * self.router.route_scale + output = self._run_selected_experts(tokens, top_scores, top_indices) + output = output.view(batch, -1, self.hidden_size) + if self.shared_experts is not None: + output = output + self.shared_experts(hidden_states) + return output + + +class LingBotVideoBlock(nn.Module): + """One Dense or sparse LingBot-Video transformer block.""" + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + intermediate_size: int, + norm_eps: float, + qkv_bias: bool, + out_bias: bool, + num_experts: int, + num_experts_per_tok: int, + moe_intermediate_size: int, + decoder_sparse_step: int, + mlp_only_layers: tuple[int, ...] | list[int], + n_shared_experts: int | None, + score_func: str, + norm_topk_prob: bool, + n_group: int | None, + topk_group: int | None, + routed_scaling_factor: float, + layer_idx: int, + ) -> None: + """Select the released Dense or sparse feed-forward structure for one layer.""" + super().__init__() + self.layer_idx = layer_idx + self.scale_shift_table = nn.Parameter(torch.zeros(1, 6 * hidden_size)) + self.norm1 = LingBotVideoRMSNorm(hidden_size, norm_eps) + self.attn = LingBotVideoAttention(hidden_size, num_attention_heads, norm_eps, qkv_bias, out_bias) + self.norm_post_attn = LingBotVideoRMSNorm(hidden_size, norm_eps) + self.norm2 = LingBotVideoRMSNorm(hidden_size, norm_eps) + if layer_idx not in mlp_only_layers and (num_experts > 0 and (layer_idx + 1) % decoder_sparse_step == 0): + self.ffn = LingBotVideoSparseMoeBlock( + hidden_size, + intermediate_size, + num_experts, + num_experts_per_tok, + moe_intermediate_size, + score_func, + norm_topk_prob, + n_group, + topk_group, + routed_scaling_factor, + n_shared_experts, + ) + else: + self.ffn = LingBotVideoMLP(hidden_size, intermediate_size) + self.norm_post_ffn = LingBotVideoRMSNorm(hidden_size, norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + temb6: torch.Tensor, + rotary_emb: torch.Tensor, + attention_mask: torch.Tensor | None = None, + moe_padding_mask: torch.Tensor | None = None, + original_seq_len: int | None = None, + ) -> torch.Tensor: + """Run attention and configured feed-forward residual branches with fp32 AdaLN.""" + expected_tokens = hidden_states.shape[0] * hidden_states.shape[1] + if temb6.ndim != 2 or temb6.shape[0] != expected_tokens: + raise ValueError("LingBotVideoBlock expects token-level temb6 with shape " + f"(B*S, 6D); got {tuple(temb6.shape)} for {tuple(hidden_states.shape)}.") + modulation = temb6.view(*hidden_states.shape[:2], -1) + self.scale_shift_table.unsqueeze(0) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = modulation.chunk(6, dim=-1) + gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh() + scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp + bulk_dtype = self.attn.to_q.weight.dtype + + attention_input = (self.norm1(hidden_states) * scale_msa + shift_msa).to(bulk_dtype) + attention_output = self.attn(attention_input, rotary_emb, attention_mask, original_seq_len) + hidden_states = hidden_states + (gate_msa * self.norm_post_attn(attention_output)).to(hidden_states.dtype) + mlp_input = (self.norm2(hidden_states) * scale_mlp + shift_mlp).to(bulk_dtype) + if isinstance(self.ffn, LingBotVideoSparseMoeBlock): + mlp_output = self.ffn(mlp_input, padding_mask=moe_padding_mask) + else: + mlp_output = self.ffn(mlp_input) + mlp_output = self.norm_post_ffn(mlp_output) + return hidden_states + (gate_mlp * mlp_output).to(hidden_states.dtype) + + +class LingBotVideoTimestepEmbedding(nn.Module): + """Two-layer timestep embedding with released parameter names.""" + + def __init__(self, input_dim: int, hidden_size: int, bias: bool) -> None: + super().__init__() + self.linear_1 = LingBotVideoLinear(input_dim, hidden_size, bias=bias) + self.linear_2 = LingBotVideoLinear(hidden_size, hidden_size, bias=bias) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + """Apply the official linear-SiLU-linear timestep projection.""" + return self.linear_2(F.silu(self.linear_1(sample))) + + +class LingBotVideoTransformer3DModel(BaseDiT): + """LingBot-Video DiT with a source-compatible Dense or MoE state surface.""" + + _fsdp_shard_conditions = [_is_lingbot_video_block] + _compile_conditions = [_is_lingbot_video_block] + _supported_attention_backends = ( + AttentionBackendEnum.TORCH_SDPA, + AttentionBackendEnum.FLASH_ATTN, + ) + param_names_mapping = {r"^(.*)$": r"\1"} + reverse_param_names_mapping = {r"^(.*)$": r"\1"} + + def _get_parameter_dtype(self, name: str, default_dtype: torch.dtype) -> torch.dtype: + """Select the released mixed-precision dtype while loading each parameter.""" + return torch.float32 if _keep_in_fp32(name) else default_dtype + + def __init__(self, config: LingBotVideoConfig, hf_config: dict[str, Any]) -> None: + """Construct the Dense or MoE variant from its released transformer config.""" + config.update_model_arch(hf_config) + super().__init__(config, hf_config) + head_dim = config.hidden_size // config.num_attention_heads + if head_dim != sum(config.axes_dims): + raise ValueError(f"head_dim {head_dim} != sum(axes_dims) {sum(config.axes_dims)}") + sp_world_size = _sequence_parallel_world_size() + assert config.num_attention_heads % sp_world_size == 0, ( + f"The number of attention heads ({config.num_attention_heads}) must be divisible by " + f"the sequence parallel size ({sp_world_size})") + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_channels_latents = config.in_channels + self.patch_embedder = LingBotVideoLinear( + config.in_channels * math.prod(config.patch_size), + config.hidden_size, + bias=config.patch_embed_bias, + ) + self.time_proj = Timesteps(config.freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.time_embedder = LingBotVideoTimestepEmbedding(config.freq_dim, config.hidden_size, + config.timestep_mlp_bias) + self.time_modulation = nn.Sequential(nn.SiLU(), LingBotVideoLinear(config.hidden_size, 6 * config.hidden_size)) + self.text_embedder = LingBotVideoTextEmbedder(config.text_dim, config.hidden_size) + self.rope = LingBotVideoRotaryEmbedding(tuple(config.axes_dims), tuple(config.axes_lens), config.rope_theta) + self.blocks = nn.ModuleList([ + LingBotVideoBlock( + config.hidden_size, + config.num_attention_heads, + config.intermediate_size, + config.norm_eps, + config.qkv_bias, + config.out_bias, + config.num_experts, + config.num_experts_per_tok, + config.moe_intermediate_size, + config.decoder_sparse_step, + config.mlp_only_layers, + config.n_shared_experts, + config.score_func, + config.norm_topk_prob, + config.n_group, + config.topk_group, + config.routed_scaling_factor, + layer_index, + ) for layer_index in range(config.depth) + ]) + self.norm_out = nn.LayerNorm(config.hidden_size, elementwise_affine=False, eps=config.norm_eps) + self.norm_out_modulation = nn.Sequential(nn.SiLU(), + LingBotVideoLinear(config.hidden_size, 2 * config.hidden_size)) + self.proj_out = LingBotVideoLinear(config.hidden_size, math.prod(config.patch_size) * config.out_channels) + self.__post_init__() + + def to(self, *args: Any, **kwargs: Any) -> LingBotVideoTransformer3DModel: + """Cast bulk weights while retaining the released fp32-sensitive modules.""" + device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs) + if dtype is None or dtype == torch.float32: + return super().to(*args, **kwargs) + if not torch.is_floating_point(torch.empty((), dtype=dtype)): + return super().to(*args, **kwargs) + if device is not None: + super().to(device=device, non_blocking=non_blocking) + for name, parameter in self.named_parameters(): + if torch.is_floating_point(parameter): + target_dtype = torch.float32 if _keep_in_fp32(name) else dtype + parameter.data = parameter.data.to(target_dtype, non_blocking=non_blocking) + if parameter.grad is not None: + parameter.grad.data = parameter.grad.data.to(target_dtype, non_blocking=non_blocking) + for name, buffer in self.named_buffers(): + if torch.is_floating_point(buffer): + target_dtype = torch.float32 if _keep_in_fp32(name) else dtype + buffer.data = buffer.data.to(target_dtype, non_blocking=non_blocking) + return self + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + guidance: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + return_dict: bool = True, + **kwargs: Any, + ) -> LingBotVideoTransformerOutput | tuple[torch.Tensor]: + """Denoise video latents with joint video-text attention.""" + del encoder_hidden_states_image, guidance, kwargs + if isinstance(encoder_hidden_states, list): + if len(encoder_hidden_states) != 1: + raise ValueError("LingBot-Video expects one text-encoder output tensor.") + encoder_hidden_states = encoder_hidden_states[0] + batch, channels, frames, height, width = hidden_states.shape + patch_t, patch_h, patch_w = self.config.patch_size + grid_t, grid_h, grid_w = frames // patch_t, height // patch_h, width // patch_w + video_tokens = grid_t * grid_h * grid_w + text_tokens = encoder_hidden_states.shape[1] + device = hidden_states.device + if encoder_attention_mask is None: + encoder_attention_mask = torch.ones((batch, text_tokens), device=device, dtype=torch.bool) + text_lengths = encoder_attention_mask.sum(dim=-1).long() + + patches = hidden_states.reshape(batch, channels, grid_t, patch_t, grid_h, patch_h, grid_w, patch_w) + patches = patches.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(batch, video_tokens, + patch_t * patch_h * patch_w * channels) + video_hidden = self.patch_embedder(patches) + text_hidden = self.text_embedder(encoder_hidden_states) + joint = torch.cat([video_hidden, text_hidden], dim=1) + + rotary_parts: list[torch.Tensor] = [] + for index in range(batch): + real_text_length = int(text_lengths[index].item()) + positions = _make_joint_position_ids(real_text_length, grid_t, grid_h, grid_w, device) + maxima = (real_text_length + grid_t, grid_h - 1, grid_w - 1) + rotary = self.rope(positions, maxima=maxima) + if real_text_length < text_tokens: + padding = torch.zeros( + text_tokens - real_text_length, + rotary.shape[-1], + device=device, + dtype=rotary.dtype, + ) + rotary = torch.cat([rotary, padding], dim=0) + rotary_parts.append(rotary) + rotary_emb = torch.stack(rotary_parts, dim=0) + + attention_mask = None + moe_padding_mask = None + if bool((text_lengths < text_tokens).any()): + key_mask = torch.cat( + [ + torch.ones(batch, video_tokens, dtype=torch.bool, device=device), + encoder_attention_mask.bool(), + ], + dim=1, + ) + attention_mask = key_mask[:, None, None, :] + moe_padding_mask = key_mask + timestep_projection = self.time_proj(timestep.float()) + timestep_embedding = self.time_embedder(timestep_projection) + token_embedding = timestep_embedding.unsqueeze(1).expand(batch, joint.shape[1], -1) + original_joint_length: int | None = None + if _sequence_parallel_world_size() > 1: + # Match the official CP order: project token modulation before + # placing every token-aligned tensor on the same padded shard. + temb6 = self.time_modulation(token_embedding.reshape(-1, + self.hidden_size)).reshape(batch, joint.shape[1], -1) + joint, original_joint_length = sequence_model_parallel_shard(joint, dim=1) + rotary_emb, _ = sequence_model_parallel_shard(rotary_emb, dim=1) + token_embedding, _ = sequence_model_parallel_shard(token_embedding, dim=1) + temb6, _ = sequence_model_parallel_shard(temb6, dim=1) + if moe_padding_mask is None: + moe_padding_mask = torch.ones(batch, original_joint_length, dtype=torch.bool, device=device) + moe_padding_mask, _ = sequence_model_parallel_shard(moe_padding_mask, dim=1) + moe_padding_mask = moe_padding_mask.reshape(-1) + temb6 = temb6.reshape(-1, 6 * self.hidden_size) + else: + temb6 = self.time_modulation(token_embedding.reshape(-1, self.hidden_size)) + if moe_padding_mask is not None: + moe_padding_mask = moe_padding_mask.reshape(-1) + + for block in self.blocks: + joint = block( + joint, + temb6, + rotary_emb, + attention_mask, + moe_padding_mask, + original_joint_length, + ) + + final_modulation = self.norm_out_modulation(token_embedding.reshape(-1, self.hidden_size)) + shift, scale = final_modulation.reshape(*joint.shape[:2], -1).chunk(2, dim=-1) + final_hidden = self.norm_out(joint) * (1.0 + scale) + shift + projected = self.proj_out(final_hidden.to(self.proj_out.weight.dtype)) + if original_joint_length is not None: + projected = sequence_model_parallel_all_gather_with_unpad(projected, original_joint_length, dim=1) + projected = projected[:, :video_tokens] + output_channels = self.config.out_channels + output = projected.reshape(batch, grid_t, grid_h, grid_w, patch_t, patch_h, patch_w, output_channels) + output = output.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(batch, output_channels, frames, height, width) + if not return_dict: + return (output, ) + return LingBotVideoTransformerOutput(sample=output) + + +EntryClass = LingBotVideoTransformer3DModel diff --git a/fastvideo/models/encoders/lingbot_video.py b/fastvideo/models/encoders/lingbot_video.py new file mode 100644 index 000000000..3dddb87c4 --- /dev/null +++ b/fastvideo/models/encoders/lingbot_video.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Native LingBot-Video Qwen3-VL language model for text-only conditioning.""" + +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn + +from fastvideo.configs.models.encoders import BaseEncoderOutput +from fastvideo.layers.layernorm import RMSNorm +from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding +from fastvideo.models.encoders.base import TextEncoder +from fastvideo.models.encoders.qwen3 import ( + Qwen3Attention, + Qwen3DecoderLayer, + Qwen3ForCausalLM, + Qwen3MLP, +) + + +class LingBotVideoQwen3VLAttention(Qwen3Attention): + """Qwen3-VL attention with the official masked repeat-K/V SDPA path.""" + + def _apply_qwen3_vl_rope( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Apply NeoX RoPE with Qwen3-VL's input-dtype multiply-add ordering.""" + flat_positions = positions.flatten() + cos_sin = self.rotary_emb.cos_sin_cache.index_select(0, flat_positions) + cos_half, sin_half = cos_sin.chunk(2, dim=-1) + cos = torch.cat((cos_half, cos_half), dim=-1).to(query.dtype) + sin = torch.cat((sin_half, sin_half), dim=-1).to(query.dtype) + if flat_positions.numel() == query.shape[1]: + cos = cos.view(1, query.shape[1], 1, self.head_dim) + sin = sin.view(1, query.shape[1], 1, self.head_dim) + else: + cos = cos.view(*query.shape[:2], 1, self.head_dim) + sin = sin.view(*query.shape[:2], 1, self.head_dim) + + def rotate(tensor: torch.Tensor) -> torch.Tensor: + first, second = tensor.chunk(2, dim=-1) + rotated = torch.cat((-second, first), dim=-1) + return tensor * cos + rotated * sin + + return rotate(query), rotate(key) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Apply fused projections, QK norm, RoPE, and causal grouped attention.""" + qkv, _ = self.qkv_proj(hidden_states) + query, key, value = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + batch_size, sequence_length = query.shape[:2] + query = query.reshape(batch_size, sequence_length, self.num_heads, self.head_dim) + key = key.reshape(batch_size, sequence_length, self.num_kv_heads, self.head_dim) + value = value.reshape(batch_size, sequence_length, self.num_kv_heads, self.head_dim) + query = self.q_norm(query) + key = self.k_norm(key) + query, key = self._apply_qwen3_vl_rope(positions, query, key) + no_padding = attention_mask is None + if no_padding: + attention_output = torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + dropout_p=0.0, + is_causal=sequence_length > 1, + scale=self.scaling, + enable_gqa=self.num_heads != self.num_kv_heads, + ).transpose(1, 2) + else: + groups = self.num_heads // self.num_kv_heads + key = (key[:, :, :, None, :].expand(-1, -1, -1, groups, -1).reshape(batch_size, sequence_length, + self.num_heads, self.head_dim)) + value = (value[:, :, :, None, :].expand(-1, -1, -1, groups, -1).reshape(batch_size, sequence_length, + self.num_heads, self.head_dim)) + causal_mask = torch.ones(sequence_length, sequence_length, device=query.device, dtype=torch.bool).tril() + key_mask = attention_mask.to(device=query.device, dtype=torch.bool) + sdpa_mask = causal_mask[None, None, :, :] & key_mask[:, None, None, :] + attention_output = torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=sdpa_mask, + dropout_p=0.0, + is_causal=False, + scale=self.scaling, + ).transpose(1, 2) + output, _ = self.o_proj(attention_output.reshape(batch_size, sequence_length, -1)) + return output + + +class LingBotVideoQwen3VLDecoderLayer(Qwen3DecoderLayer): + """Qwen3-VL decoder layer with explicit official residual rounding order.""" + + def __init__(self, config: Any, prefix: str) -> None: + """Build the final Qwen3-VL attention once to avoid orphan parameters.""" + nn.Module.__init__(self) + self.hidden_size = config.hidden_size + quant_config = getattr(config, "quant_config", None) + self.self_attn = LingBotVideoQwen3VLAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + rope_theta=config.rope_theta, + rope_scaling=config.rope_scaling, + max_position_embeddings=config.max_position_embeddings, + quant_config=quant_config, + bias=config.attention_bias, + prefix=f"{prefix}.self_attn", + ) + self.mlp = Qwen3MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + bias=getattr(config, "mlp_bias", False), + prefix=f"{prefix}.mlp", + ) + self.input_layernorm = RMSNorm(self.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(self.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run attention and MLP with each residual sum rounded before normalization.""" + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(positions, hidden_states, attention_mask) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + +class LingBotVideoQwen3VLTextModel(Qwen3ForCausalLM): + """Load the Qwen3-VL language-model subset without its vision tower or LM head.""" + + supports_hf_from_pretrained = False + + def __init__(self, config) -> None: + """Construct the exact Qwen3-VL module graph without replacing base layers.""" + TextEncoder.__init__(self, config) + self.quant_config = getattr(config, "quant_config", None) + if getattr(config, "lora_config", None) is not None: + max_loras = getattr(config.lora_config, "max_loras", 1) + lora_vocab_size = getattr(config.lora_config, "lora_extra_vocab_size", 1) + lora_vocab = lora_vocab_size * max_loras + else: + lora_vocab = 0 + self.vocab_size = config.vocab_size + lora_vocab + self.org_vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + quant_config=self.quant_config, + ) + self.layers = nn.ModuleList( + LingBotVideoQwen3VLDecoderLayer(config, prefix=f"{config.prefix}.layers.{index}") + for index in range(config.num_hidden_layers)) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs: Any, + ) -> BaseEncoderOutput: + """Run explicit Qwen3-VL layers and return the requested hidden-state tuple.""" + del kwargs + output_hidden_states = (output_hidden_states + if output_hidden_states is not None else self.config.output_hidden_states) + if inputs_embeds is None: + if input_ids is None: + raise ValueError("input_ids or inputs_embeds is required") + hidden_states = self.get_input_embeddings(input_ids) + else: + hidden_states = inputs_embeds + if position_ids is None: + position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0) + if attention_mask is not None and bool(attention_mask.to(torch.bool).all()): + attention_mask = None + all_hidden_states: tuple[torch.Tensor, ...] | None = () if output_hidden_states else None + for layer in self.layers: + if all_hidden_states is not None: + all_hidden_states += (hidden_states, ) + hidden_states = layer(position_ids, hidden_states, attention_mask) + hidden_states = self.norm(hidden_states) + if all_hidden_states is not None: + all_hidden_states += (hidden_states, ) + return BaseEncoderOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Accept either official compound keys or converted native keys.""" + prefix = "model.language_model." + language_weights = ((name[len(prefix):] if name.startswith(prefix) else name, tensor) + for name, tensor in weights + if name.startswith(prefix) or not name.startswith(("model.", "lm_head."))) + return super().load_weights(language_weights) + + +EntryClass = LingBotVideoQwen3VLTextModel diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 29ddca0ee..e7045ff3d 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -15,13 +15,11 @@ import torch from torch import nn from torch.distributed import DeviceMesh, init_device_mesh from torch.distributed._tensor import distribute_tensor -from torch.distributed.fsdp import (CPUOffloadPolicy, FSDPModule, - MixedPrecisionPolicy, fully_shard) +from torch.distributed.fsdp import (CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, fully_shard) from torch.nn.modules.module import _IncompatibleKeys from fastvideo.logger import init_logger -from fastvideo.models.loader.utils import (get_param_names_mapping, - hf_to_custom_state_dict) +from fastvideo.models.loader.utils import (get_param_names_mapping, hf_to_custom_state_dict) from fastvideo.models.loader.weight_utils import safetensors_weights_iterator from fastvideo.utils import set_mixed_precision_policy, is_pin_memory_available @@ -43,13 +41,16 @@ def _maybe_quantize_model(model: nn.Module) -> None: """ # Defer imports: these modules pull in heavy symbols at module-load time. from fastvideo.layers.quantization.nvfp4_config import ( - NVFP4QuantizeMethod, convert_model_to_nvfp4, + NVFP4QuantizeMethod, + convert_model_to_nvfp4, ) from fastvideo.layers.quantization.nvfp4_qat_config import ( - NVFP4QATQuantizeMethod, convert_model_to_fp4, + NVFP4QATQuantizeMethod, + convert_model_to_fp4, ) from fastvideo.layers.quantization.fp8_config import ( - FP8QuantizeMethod, convert_model_to_fp8, + FP8QuantizeMethod, + convert_model_to_fp8, ) for mod in model.modules(): @@ -121,10 +122,7 @@ def maybe_load_fsdp_model( """ # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are # manually casting the inputs to the model - mp_policy = MixedPrecisionPolicy(param_dtype, - reduce_dtype, - output_dtype, - cast_forward_inputs=False) + mp_policy = MixedPrecisionPolicy(param_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False) set_mixed_precision_policy( param_dtype=param_dtype, @@ -137,6 +135,13 @@ def maybe_load_fsdp_model( with set_default_dtype(default_dtype), torch.device("meta"): model = model_cls(**init_params) + dtype_selector = getattr(model, "_get_parameter_dtype", None) + has_mixed_parameter_dtypes = callable(dtype_selector) and any( + dtype_selector(name, param_dtype) != param_dtype for name, _ in model.named_parameters()) + if training_mode and has_mixed_parameter_dtypes: + raise NotImplementedError("FSDP training with model-selected mixed parameter dtypes requires " + "separate gradient synchronization for replicated parameters.") + # Check if we should use FSDP use_fsdp = training_mode or fsdp_inference @@ -152,7 +157,7 @@ def maybe_load_fsdp_model( if not training_mode and not fsdp_inference: hsdp_replicate_dim = world_size hsdp_shard_dim = 1 - + if current_platform.is_npu(): with torch.device("cpu"): device_mesh = init_device_mesh( @@ -163,11 +168,11 @@ def maybe_load_fsdp_model( ) else: device_mesh = init_device_mesh( - "cuda", - # (Replicate(), Shard(dim=0)) - mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim), - mesh_dim_names=("replicate", "shard"), - ) + "cuda", + # (Replicate(), Shard(dim=0)) + mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim), + mesh_dim_names=("replicate", "shard"), + ) shard_model(model, cpu_offload=cpu_offload, reshard_after_forward=True, @@ -188,12 +193,10 @@ def maybe_load_fsdp_model( param_names_mapping=param_names_mapping_fn, ) if hasattr(model, "materialize_non_persistent_buffers"): - model.materialize_non_persistent_buffers( - device=device, dtype=default_dtype) + model.materialize_non_persistent_buffers(device=device, dtype=default_dtype) for n, p in chain(model.named_parameters(), model.named_buffers()): if p.is_meta: - raise RuntimeError( - f"Unexpected param or buffer {n} on meta device.") + raise RuntimeError(f"Unexpected param or buffer {n} on meta device.") # Avoid unintended computation graph accumulation during inference if isinstance(p, torch.nn.Parameter): p.requires_grad = False @@ -209,8 +212,7 @@ def maybe_load_fsdp_model( compile_in_loader = enable_torch_compile and training_mode if compile_in_loader: compile_kwargs = torch_compile_kwargs or {} - logger.info("Enabling torch.compile for FSDP training module with kwargs=%s", - compile_kwargs) + logger.info("Enabling torch.compile for FSDP training module with kwargs=%s", compile_kwargs) model = torch.compile(model, **compile_kwargs) logger.info("torch.compile enabled for %s", type(model).__name__) return model @@ -254,58 +256,81 @@ def shard_model( """ # Check if we should use size-based filtering use_size_filtering = os.environ.get("FASTVIDEO_FSDP2_AUTOWRAP", "0") == "1" - + if not fsdp_shard_conditions: logger.warning("No FSDP shard conditions provided; nothing will be sharded.") return + default_param_dtype = getattr(mp_policy, "param_dtype", None) + dtype_selector = getattr(model, "_get_parameter_dtype", None) + ignored_params: set[nn.Parameter] = set() + if callable(dtype_selector) and default_param_dtype is not None: + ignored_params = { + parameter + for name, parameter in model.named_parameters() + if dtype_selector(name, default_param_dtype) != default_param_dtype + } + named_modules = list(model.named_modules()) + ignored_params_by_module = { + id(module): ignored_params.intersection(set(module.parameters())) + for _, module in named_modules + } + fsdp_kwargs = { "reshard_after_forward": reshard_after_forward, "mesh": mesh, "mp_policy": mp_policy, } if cpu_offload: - fsdp_kwargs["offload_policy"] = CPUOffloadPolicy( - pin_memory=pin_cpu_memory) + fsdp_kwargs["offload_policy"] = CPUOffloadPolicy(pin_memory=pin_cpu_memory) # iterating in reverse to start with # lowest-level modules first num_layers_sharded = 0 - + if use_size_filtering: # Size-based filtering mode min_params = int(os.environ.get("FASTVIDEO_FSDP2_MIN_PARAMS", "10000000")) logger.info("Using size-based filtering with threshold: %.2fM", min_params / 1e6) - - for n, m in reversed(list(model.named_modules())): + + for n, m in reversed(named_modules): if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]): # Count all parameters param_count = sum(p.numel() for p in m.parameters(recurse=True)) - + # Skip small modules if param_count < min_params: - logger.info("Skipping module %s (%.2fM params < %.2fM threshold)", - n, param_count / 1e6, min_params / 1e6) + logger.info("Skipping module %s (%.2fM params < %.2fM threshold)", n, param_count / 1e6, + min_params / 1e6) continue - + # Shard this module logger.info("Sharding module %s (%.2fM params)", n, param_count / 1e6) - fully_shard(m, **fsdp_kwargs) + module_kwargs = fsdp_kwargs + local_ignored_params = ignored_params_by_module[id(m)] + if local_ignored_params: + module_kwargs = {**fsdp_kwargs, "ignored_params": local_ignored_params} + fully_shard(m, **module_kwargs) num_layers_sharded += 1 else: - # Shard all modules matching conditions - for n, m in reversed(list(model.named_modules())): + # Shard all modules matching conditions + for n, m in reversed(named_modules): if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]): - fully_shard(m, **fsdp_kwargs) + module_kwargs = fsdp_kwargs + local_ignored_params = ignored_params_by_module[id(m)] + if local_ignored_params: + module_kwargs = {**fsdp_kwargs, "ignored_params": local_ignored_params} + fully_shard(m, **module_kwargs) num_layers_sharded += 1 - + if num_layers_sharded == 0: - raise ValueError( - "No layer modules were sharded. Please check if shard conditions are working as expected." - ) + raise ValueError("No layer modules were sharded. Please check if shard conditions are working as expected.") # Finally shard the entire model to account for any stragglers - fully_shard(model, **fsdp_kwargs) + root_kwargs = fsdp_kwargs + if ignored_params: + root_kwargs = {**fsdp_kwargs, "ignored_params": ignored_params} + fully_shard(model, **root_kwargs) # TODO(PY): device mesh for cfg parallel @@ -341,17 +366,17 @@ def load_model_from_full_model_state_dict( """ meta_sd = model.state_dict() named_parameters = dict(model.named_parameters()) + named_buffers = dict(model.named_buffers()) sharded_sd = {} - custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict( - full_sd_iterator, param_names_mapping) # type: ignore + custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(full_sd_iterator, + param_names_mapping) # type: ignore for target_param_name, full_tensor in custom_param_sd.items(): meta_sharded_param = meta_sd.get(target_param_name) if meta_sharded_param is None: # Some checkpoints include extra entries that are not part of the # instantiated model's state_dict (e.g. `_extra_state` keys from # some FSDP checkpoint formats). These can be safely skipped. - if (target_param_name.endswith("._extra_state") - or target_param_name.endswith("_extra_state")): + if (target_param_name.endswith("._extra_state") or target_param_name.endswith("_extra_state")): logger.warning( "Skipping non-parameter checkpoint key: %s", target_param_name, @@ -370,8 +395,12 @@ def load_model_from_full_model_state_dict( raise ValueError( f"Parameter {target_param_name} not found in custom model state dict. The hf to custom mapping may be incorrect." ) + target_dtype = param_dtype + dtype_selector = getattr(model, "_get_parameter_dtype", None) + if callable(dtype_selector): + target_dtype = dtype_selector(target_param_name, param_dtype) if not hasattr(meta_sharded_param, "device_mesh"): - full_tensor = full_tensor.to(device=device, dtype=param_dtype) + full_tensor = full_tensor.to(device=device, dtype=target_dtype) target_param = named_parameters.get(target_param_name) weight_loader = getattr(target_param, "weight_loader", None) # Gated on a shape mismatch: only fused/stacked params with a custom @@ -380,9 +409,7 @@ def load_model_from_full_model_state_dict( # fall through to the original `sharded_tensor = full_tensor` below. if target_param is not None and callable(weight_loader) and tuple(target_param.shape) != tuple( full_tensor.shape): - loaded_param = nn.Parameter(torch.empty(tuple(target_param.shape), - device=device, - dtype=param_dtype), + loaded_param = nn.Parameter(torch.empty(tuple(target_param.shape), device=device, dtype=target_dtype), requires_grad=False) for attr_name, attr_value in vars(target_param).items(): setattr(loaded_param, attr_name, attr_value) @@ -392,7 +419,7 @@ def load_model_from_full_model_state_dict( # In cases where parts of the model aren't sharded, some parameters will be plain tensors. sharded_tensor = full_tensor else: - full_tensor = full_tensor.to(device=device, dtype=param_dtype) + full_tensor = full_tensor.to(device=device, dtype=target_dtype) sharded_tensor = distribute_tensor( full_tensor, meta_sharded_param.device_mesh, @@ -400,36 +427,35 @@ def load_model_from_full_model_state_dict( ) if cpu_offload: sharded_tensor = sharded_tensor.cpu() - sharded_sd[target_param_name] = nn.Parameter(sharded_tensor) + if target_param_name in named_buffers: + sharded_sd[target_param_name] = sharded_tensor + else: + sharded_sd[target_param_name] = nn.Parameter(sharded_tensor) model.reverse_param_names_mapping = reverse_param_names_mapping unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys()) if unused_keys: - logger.warning("Found unloaded parameters in meta state dict: %s", - unused_keys) + logger.warning("Found unloaded parameters in meta state dict: %s", unused_keys) # List of allowed parameter name patterns ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress", "proj_l"] # Can be extended as needed for new_param_name in unused_keys: - if not any(pattern in new_param_name - for pattern in ALLOWED_NEW_PARAM_PATTERNS): - logger.error("Unsupported new parameter: %s. Allowed patterns: %s", - new_param_name, ALLOWED_NEW_PARAM_PATTERNS) - raise ValueError( - f"New parameter '{new_param_name}' is not supported. " - f"Currently only parameters containing {ALLOWED_NEW_PARAM_PATTERNS} are allowed." - ) + if not any(pattern in new_param_name for pattern in ALLOWED_NEW_PARAM_PATTERNS): + logger.error("Unsupported new parameter: %s. Allowed patterns: %s", new_param_name, + ALLOWED_NEW_PARAM_PATTERNS) + raise ValueError(f"New parameter '{new_param_name}' is not supported. " + f"Currently only parameters containing {ALLOWED_NEW_PARAM_PATTERNS} are allowed.") meta_sharded_param = meta_sd.get(new_param_name) + target_dtype = param_dtype + dtype_selector = getattr(model, "_get_parameter_dtype", None) + if callable(dtype_selector): + target_dtype = dtype_selector(new_param_name, param_dtype) if not hasattr(meta_sharded_param, "device_mesh"): # Initialize with zeros - sharded_tensor = torch.zeros_like(meta_sharded_param, - device=device, - dtype=param_dtype) + sharded_tensor = torch.zeros_like(meta_sharded_param, device=device, dtype=target_dtype) else: # Initialize with zeros and distribute - full_tensor = torch.zeros_like(meta_sharded_param, - device=device, - dtype=param_dtype) + full_tensor = torch.zeros_like(meta_sharded_param, device=device, dtype=target_dtype) sharded_tensor = distribute_tensor( full_tensor, meta_sharded_param.device_mesh, diff --git a/fastvideo/pipelines/basic/glm_image/stages/decoding.py b/fastvideo/pipelines/basic/glm_image/stages/decoding.py index b34d4c195..3b78255c8 100644 --- a/fastvideo/pipelines/basic/glm_image/stages/decoding.py +++ b/fastvideo/pipelines/basic/glm_image/stages/decoding.py @@ -19,7 +19,7 @@ class GlmImageDecodingStage(DecodingStage): vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] vae_autocast = (vae_dtype != torch.float32 and not fastvideo_args.disable_autocast) - latents = self._denormalize_latents(latents) + latents = self._denormalize_latents(latents, fastvideo_args) if latents.dim() == 5: latents = latents.squeeze(2) diff --git a/fastvideo/pipelines/basic/lingbot_video/__init__.py b/fastvideo/pipelines/basic/lingbot_video/__init__.py new file mode 100644 index 000000000..09320d2e1 --- /dev/null +++ b/fastvideo/pipelines/basic/lingbot_video/__init__.py @@ -0,0 +1,5 @@ +"""Dense LingBot-Video inference pipeline.""" + +from fastvideo.pipelines.basic.lingbot_video.lingbot_video_pipeline import LingBotVideoPipeline + +__all__ = ["LingBotVideoPipeline"] diff --git a/fastvideo/pipelines/basic/lingbot_video/lingbot_video_pipeline.py b/fastvideo/pipelines/basic/lingbot_video/lingbot_video_pipeline.py new file mode 100644 index 000000000..f8d3fd778 --- /dev/null +++ b/fastvideo/pipelines/basic/lingbot_video/lingbot_video_pipeline.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Stage-composed LingBot-Video Dense and MoE/refiner T2V pipeline.""" + +from typing import Any + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline +from fastvideo.pipelines.basic.lingbot_video.stages import ( + LingBotVideoDenoisingStage, + LingBotVideoInputValidationStage, + LingBotVideoLatentPreparationStage, + LingBotVideoRefinerPreparationStage, +) +from fastvideo.pipelines.stages import ( + ConditioningStage, + DecodingStage, + TextEncodingStage, + TimestepPreparationStage, +) + + +class LingBotVideoPipeline(LoRAPipeline, ComposedPipelineBase): + """T2V pipeline with optional released MoE pixel-space refinement.""" + + is_video_pipeline = True + _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"] + + def load_modules( + self, + fastvideo_args: FastVideoArgs, + loaded_modules: dict[str, Any] | None = None, + ) -> dict[str, Any]: + """Load the optional refiner DiT and the VAE encoder only when declared.""" + model_index = self._load_config(self.model_path) + required = list(type(self)._required_config_modules) + load_refiner = "transformer_2" in model_index and getattr(fastvideo_args, "refine_enabled", None) is not False + if load_refiner: + required.append("transformer_2") + fastvideo_args.pipeline_config.vae_config.load_encoder = True + self._required_config_modules = required + return super().load_modules(fastvideo_args, loaded_modules) + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: + """Apply the released runtime flow shift to the loaded scheduler.""" + shift = fastvideo_args.pipeline_config.flow_shift + if shift is None: + raise ValueError("LingBot-Video requires a flow shift") + self.get_module("scheduler").set_shift(float(shift)) + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + """Create base generation and the optional decoded-video refiner stages.""" + refiner = self.get_module("transformer_2") + self.add_stage( + "input_validation_stage", + LingBotVideoInputValidationStage(refiner_enabled=refiner is not None), + ) + self.add_stage( + "prompt_encoding_stage", + TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + ) + self.add_stage("conditioning_stage", ConditioningStage()) + self.add_stage( + "timestep_preparation_stage", + TimestepPreparationStage(scheduler=self.get_module("scheduler")), + ) + self.add_stage( + "latent_preparation_stage", + LingBotVideoLatentPreparationStage(transformer=self.get_module("transformer")), + ) + self.add_stage( + "denoising_stage", + LingBotVideoDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ) + self.add_stage( + "decoding_stage", + DecodingStage(vae=self.get_module("vae"), pipeline=self), + ) + if refiner is not None: + self.add_stage( + "refiner_preparation_stage", + LingBotVideoRefinerPreparationStage( + vae=self.get_module("vae"), + scheduler=self.get_module("scheduler"), + ), + ) + self.add_stage( + "refiner_denoising_stage", + LingBotVideoDenoisingStage( + transformer=refiner, + scheduler=self.get_module("scheduler"), + refiner=True, + ), + ) + self.add_stage( + "refiner_decoding_stage", + DecodingStage(vae=self.get_module("vae"), pipeline=self), + ) + + +EntryClass = LingBotVideoPipeline diff --git a/fastvideo/pipelines/basic/lingbot_video/presets.py b/fastvideo/pipelines/basic/lingbot_video/presets.py new file mode 100644 index 000000000..6e460b375 --- /dev/null +++ b/fastvideo/pipelines/basic/lingbot_video/presets.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Official LingBot-Video T2V inference presets.""" + +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +DEFAULT_NEGATIVE_PROMPT = ('{"universal_negative": {"visual_quality": ["low quality", "worst quality", "blurry", ' + '"pixelated", "jpeg artifacts", "low resolution", "unstable color", "color flicker", ' + '"underexposed", "overexposed", "invisible subject", "subject hidden in darkness"], ' + '"artistic_style": ["painting", "illustration", "drawing", "cartoon", "3d render", ' + '"cgi", "sketch", "digital art"], "composition_and_content": ["text", "watermark", ' + '"signature", "logo", "subtitles", "pillarboxed", "side bars", "portrait image in ' + 'landscape frame"], "temporal_and_motion_stability": ["flickering", "jittery", ' + '"motion blur", "temporal inconsistency", "warping", "morphing", "incoherent motion", ' + '"unnatural movement", "static object with sudden jump", "frame-to-frame inconsistency"], ' + '"material_and_structure": ["plastic-like glass", "unrealistic texture", "deformed ' + 'bottle", "liquid freezing improperly", "distorted reflections"]}}') + +_DENOISE_STAGE = PresetStageSpec( + name="denoise", + kind="denoising", + description="LingBot-Video batched-CFG denoising", + allowed_overrides=frozenset({"num_inference_steps", "guidance_scale"}), +) + +_REFINE_STAGE = PresetStageSpec( + name="refine", + kind="refinement", + description="LingBot-Video pixel-space resize, VAE re-encode, and refiner denoising", + allowed_overrides=frozenset({ + "height_sr", + "width_sr", + "num_inference_steps_sr", + "guidance_scale_2", + "t_thresh", + }), +) + +LINGBOT_VIDEO_DENSE_T2V = InferencePreset( + name="lingbot_video_dense_t2v", + version=1, + model_family="lingbot_video", + description="LingBot-Video Dense 1.3B text-to-video at 480p", + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "height": 480, + "width": 832, + "num_frames": 121, + "fps": 24, + "num_inference_steps": 40, + "guidance_scale": 3.0, + "batch_cfg": True, + "seed": 42, + "negative_prompt": DEFAULT_NEGATIVE_PROMPT, + }, +) + +LINGBOT_VIDEO_MOE_REFINER_T2V = InferencePreset( + name="lingbot_video_moe_refiner_t2v", + version=1, + model_family="lingbot_video", + description="LingBot-Video MoE 30B-A3B T2V with 1080p refiner", + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, _REFINE_STAGE), + defaults={ + "height": 480, + "width": 832, + "height_sr": 1088, + "width_sr": 1920, + "num_frames": 121, + "fps": 24, + "num_inference_steps": 40, + "num_inference_steps_sr": 8, + "guidance_scale": 3.0, + "guidance_scale_2": 3.0, + "batch_cfg": True, + "t_thresh": 0.85, + "seed": 42, + "negative_prompt": DEFAULT_NEGATIVE_PROMPT, + }, + stage_defaults={ + "refine": { + "height_sr": 1088, + "width_sr": 1920, + "num_inference_steps_sr": 8, + "guidance_scale_2": 3.0, + "t_thresh": 0.85, + }, + }, +) + +ALL_PRESETS = (LINGBOT_VIDEO_DENSE_T2V, LINGBOT_VIDEO_MOE_REFINER_T2V) diff --git a/fastvideo/pipelines/basic/lingbot_video/stages.py b/fastvideo/pipelines/basic/lingbot_video/stages.py new file mode 100644 index 000000000..50be9c577 --- /dev/null +++ b/fastvideo/pipelines/basic/lingbot_video/stages.py @@ -0,0 +1,345 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LingBot-Video stages whose contracts differ from shared Wan behavior.""" + +from __future__ import annotations + +import numpy as np +import torch +import torch.nn.functional as F + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.input_validation import InputValidationStage + +LINGBOT_VIDEO_REFINER_TAIL_STEPS = 2 + + +def _compute_refiner_sigmas( + sigma_max: float, + sigma_min: float, + num_inference_steps: int, + shift: float, + t_thresh: float, +) -> np.ndarray: + """Build the released truncated schedule plus its two-step low-noise tail.""" + if not 0.0 < t_thresh <= 1.0: + raise ValueError(f"LingBot-Video refiner t_thresh must be in (0, 1], got {t_thresh}") + if num_inference_steps < 1: + raise ValueError("LingBot-Video refiner requires at least one inference step") + base = np.linspace(sigma_max, sigma_min, num_inference_steps + 1).copy()[:-1] + shifted = shift * base / (1.0 + (shift - 1.0) * base) + sigmas = shifted[shifted <= t_thresh + 1e-6] + if sigmas.size == 0 or abs(float(sigmas[0]) - t_thresh) > 1e-6: + sigmas = np.concatenate(([t_thresh], sigmas)) + tail = np.linspace( + float(sigmas[-1]), + min(sigma_min, float(sigmas[-1])), + LINGBOT_VIDEO_REFINER_TAIL_STEPS + 2, + )[1:-1] + sigmas = np.concatenate((sigmas, tail)) + if sigmas.size > 1 and not np.all(np.diff(sigmas) < 0.0): + raise ValueError(f"LingBot-Video refiner sigmas must descend strictly, got {sigmas.tolist()}") + return sigmas.astype(np.float32) + + +class LingBotVideoInputValidationStage(InputValidationStage): + """Validate released shape constraints and construct the official CUDA RNG.""" + + def __init__(self, refiner_enabled: bool = False) -> None: + self.refiner_enabled = refiner_enabled + + def _generate_seeds(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> None: + """Use one device-local generator, matching the official production runner.""" + del fastvideo_args + if batch.seed is None: + raise ValueError("LingBot-Video requires a seed") + if batch.num_videos_per_prompt != 1: + raise ValueError("LingBot-Video currently supports one video per prompt") + batch.seeds = [batch.seed] + batch.generator = torch.Generator(device=get_local_torch_device()).manual_seed(batch.seed) + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + """Run shared validation, then enforce LingBot temporal and spatial geometry.""" + batch = super().forward(batch, fastvideo_args) + if not isinstance(batch.num_frames, int): + raise TypeError("LingBot-Video num_frames must be an integer") + if batch.num_frames != 1 and (batch.num_frames - 1) % 4 != 0: + raise ValueError(f"num_frames must be 1 or 4n+1, got {batch.num_frames}") + if not isinstance(batch.height, int) or not isinstance(batch.width, int): + raise TypeError("LingBot-Video height and width must be integers") + if batch.height % 16 != 0 or batch.width % 16 != 0: + raise ValueError(f"height and width must be divisible by 16, got {batch.height}x{batch.width}") + if isinstance(batch.prompt, list) and len(batch.prompt) != 1: + raise ValueError("LingBot-Video currently supports prompt batch size one") + if self.refiner_enabled and fastvideo_args.output_type == "latent": + raise ValueError("LingBot-Video refinement requires decoded pixel output") + return batch + + +class LingBotVideoLatentPreparationStage(PipelineStage): + """Prepare fp32 latents in the released 4x temporal and 8x spatial geometry.""" + + def __init__(self, transformer) -> None: + self.transformer = transformer + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + """Generate or validate one normalized fp32 latent video.""" + del fastvideo_args + if not all(isinstance(value, int) for value in (batch.num_frames, batch.height, batch.width)): + raise TypeError("latent geometry must contain integer frames, height, and width") + shape = ( + 1, + self.transformer.num_channels_latents, + (batch.num_frames - 1) // 4 + 1, + batch.height // 8, + batch.width // 8, + ) + device = get_local_torch_device() + if batch.latents is None: + batch.latents = torch.randn( + shape, + generator=batch.generator, + device=device, + dtype=torch.float32, + ) + else: + if tuple(batch.latents.shape) != shape: + raise ValueError(f"supplied latent shape {tuple(batch.latents.shape)} does not match {shape}") + batch.latents = batch.latents.to(device=device, dtype=torch.float32) + batch.raw_latent_shape = shape + return batch + + +class LingBotVideoRefinerPreparationStage(PipelineStage): + """Resize and encode the base video, then initialize the released refiner state.""" + + performance_component_metric = "vae_encode_time_s" + + def __init__(self, vae, scheduler) -> None: + self.vae = vae + self.scheduler = scheduler + + @staticmethod + def _resize_video(video: torch.Tensor, height: int, width: int) -> torch.Tensor: + """Bicubic-resize every decoded frame using the released tensor layout.""" + batch, channels, frames, source_height, source_width = video.shape + flat = video.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, source_height, source_width) + resized = F.interpolate(flat, size=(height, width), mode="bicubic", align_corners=False).clamp(0.0, 1.0) + return resized.reshape(batch, frames, channels, height, width).permute(0, 2, 1, 3, 4).contiguous() + + def _encode_video( + self, + video: torch.Tensor, + generator: torch.Generator, + device: torch.device, + ) -> torch.Tensor: + """Encode `[0,1]` pixels and convert Wan VAE latents to normalized DiT space.""" + video = video.to(device=device, dtype=torch.float32).mul(2.0).sub(1.0) + with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"): + encoded = self.vae.encode(video) + if hasattr(encoded, "latent_dist"): + latents = encoded.latent_dist.sample(generator) + elif hasattr(encoded, "sample") and callable(encoded.sample): + latents = encoded.sample(generator) + elif isinstance(encoded, tuple | list): + latents = encoded[0] + else: + latents = encoded + mean = torch.tensor(self.vae.config.latents_mean, device=device, dtype=torch.float32).view(1, -1, 1, 1, 1) + std = torch.tensor(self.vae.config.latents_std, device=device, dtype=torch.float32).view(1, -1, 1, 1, 1) + return ((latents.float() - mean) / std).to(latents) + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + """Prepare high-resolution refiner latents and its exact truncated sigma schedule.""" + if batch.output is None or batch.output.ndim != 5: + raise ValueError("LingBot-Video refinement requires a decoded base video") + if not isinstance(batch.height_sr, int) or not isinstance(batch.width_sr, int): + raise TypeError("LingBot-Video refinement requires integer height_sr and width_sr") + if batch.height_sr % 16 != 0 or batch.width_sr % 16 != 0: + raise ValueError("LingBot-Video refiner height_sr and width_sr must be divisible by 16") + if batch.seed is None: + raise ValueError("LingBot-Video refinement requires a seed") + device = get_local_torch_device() + if isinstance(self.vae, torch.nn.Module): + self.vae.to(device) + generator = torch.Generator(device=device).manual_seed(batch.seed) + resized = self._resize_video(batch.output, batch.height_sr, batch.width_sr) + encoded = self._encode_video(resized, generator, device) + noise = torch.randn(encoded.shape, generator=generator, device=device, dtype=encoded.dtype) + batch.latents = ((1.0 - batch.t_thresh) * encoded + batch.t_thresh * noise).float() + batch.generator = generator + batch.height = batch.height_sr + batch.width = batch.width_sr + batch.raw_latent_shape = tuple(batch.latents.shape) + batch.extra["lingbot_video_base_shape"] = tuple(batch.output.shape) + batch.output = None + + shift = fastvideo_args.pipeline_config.flow_shift + if shift is None: + raise ValueError("LingBot-Video refinement requires a flow shift") + sigmas = _compute_refiner_sigmas( + float(self.scheduler.sigma_max), + float(self.scheduler.sigma_min), + batch.num_inference_steps_sr, + float(shift), + float(batch.t_thresh), + ) + self.scheduler.set_timesteps(len(sigmas), device=device, sigmas=sigmas, shift=1.0) + batch.timesteps = self.scheduler.timesteps + if getattr(fastvideo_args, "vae_cpu_offload", False): + self.vae.to("cpu") + return batch + + +class LingBotVideoDenoisingStage(PipelineStage): + """Run the released batched-CFG bf16 DiT loop with fp32 scheduler state.""" + + performance_component_metric = "dit_time_s" + + def __init__(self, transformer, scheduler, refiner: bool = False) -> None: + self.transformer = transformer + self.scheduler = scheduler + self.refiner = refiner + + @staticmethod + def _pad_condition( + embeds: torch.Tensor, + mask: torch.Tensor, + length: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Right-pad one condition stream to a shared batched-CFG length.""" + pad_length = length - embeds.shape[1] + if pad_length < 0 or embeds.shape[:2] != mask.shape: + raise ValueError("invalid LingBot-Video prompt embedding/mask shapes") + if pad_length == 0: + return embeds, mask + embed_padding = embeds.new_zeros(embeds.shape[0], pad_length, embeds.shape[2]) + mask_padding = mask.new_zeros(mask.shape[0], pad_length) + return ( + torch.cat((embeds, embed_padding), dim=1), + torch.cat((mask, mask_padding), dim=1), + ) + + @staticmethod + def _transformer_timestep(timestep: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """Reproduce the official divide-cast-multiply timestep rounding.""" + sigma = timestep.float() / 1000.0 + if dtype in (torch.bfloat16, torch.float16): + sigma = sigma.to(dtype) + return (sigma * 1000.0).float() + + def _prepare_conditions( + self, + batch: ForwardBatch, + dtype: torch.dtype, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Pack conditional then unconditional text streams for one batched CFG call.""" + prompt = batch.prompt_embeds[0].to(device=device, dtype=dtype) + if batch.prompt_attention_mask is None or not batch.prompt_attention_mask: + raise ValueError("LingBot-Video requires a prompt attention mask") + prompt_mask = batch.prompt_attention_mask[0].to(device=device) + if not self._uses_cfg(batch) or not batch.batch_cfg: + return prompt, prompt_mask + negative, negative_mask = self._negative_condition(batch, prompt, prompt_mask, dtype, device) + target_length = max(prompt.shape[1], negative.shape[1]) + prompt, prompt_mask = self._pad_condition(prompt, prompt_mask, target_length) + negative, negative_mask = self._pad_condition(negative, negative_mask, target_length) + return ( + torch.cat((prompt, negative), dim=0), + torch.cat((prompt_mask, negative_mask), dim=0), + ) + + def _negative_condition( + self, + batch: ForwardBatch, + prompt: torch.Tensor, + prompt_mask: torch.Tensor, + dtype: torch.dtype, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Use the refiner's zero-cloned null condition or the encoded negative prompt.""" + if self.refiner: + return torch.zeros_like(prompt), prompt_mask.clone() + if batch.negative_prompt_embeds is None or not batch.negative_prompt_embeds: + raise ValueError("LingBot-Video CFG requires negative prompt embeddings") + if batch.negative_attention_mask is None or not batch.negative_attention_mask: + raise ValueError("LingBot-Video CFG requires a negative prompt mask") + return ( + batch.negative_prompt_embeds[0].to(device=device, dtype=dtype), + batch.negative_attention_mask[0].to(device=device), + ) + + def _uses_cfg(self, batch: ForwardBatch) -> bool: + """Enable guidance independently for the base or refiner scale.""" + scale = batch.guidance_scale_2 if self.refiner else batch.guidance_scale + return scale is not None and scale > 1.0 + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + """Denoise latents while keeping scheduler samples and predictions in fp32.""" + if batch.latents is None or batch.timesteps is None: + raise ValueError("LingBot-Video denoising requires latents and timesteps") + device = get_local_torch_device() + transformer_dtype = next(self.transformer.parameters()).dtype + condition, condition_mask = self._prepare_conditions(batch, transformer_dtype, device) + latents = batch.latents.to(device=device, dtype=torch.float32) + do_cfg = self._uses_cfg(batch) + negative = negative_mask = None + if do_cfg and not batch.batch_cfg: + negative, negative_mask = self._negative_condition(batch, condition, condition_mask, transformer_dtype, + device) + trajectory: list[torch.Tensor] = [] + trajectory_timesteps: list[torch.Tensor] = [] + for timestep in batch.timesteps: + timestep_batch = self._transformer_timestep(timestep, transformer_dtype).expand(1).to(device) + latent_input = latents + if do_cfg and batch.batch_cfg: + latent_input = torch.cat((latents, latents), dim=0) + timestep_batch = torch.cat((timestep_batch, timestep_batch), dim=0) + autocast_enabled = device.type == "cuda" and transformer_dtype != torch.float32 + with torch.autocast(device_type=device.type, dtype=transformer_dtype, enabled=autocast_enabled): + prediction = self.transformer( + latent_input, + timestep_batch, + condition, + encoder_attention_mask=condition_mask, + return_dict=False, + )[0].float() + if do_cfg: + if batch.batch_cfg: + conditional, unconditional = prediction.chunk(2, dim=0) + else: + with torch.autocast( + device_type=device.type, + dtype=transformer_dtype, + enabled=autocast_enabled, + ): + unconditional = self.transformer( + latents, + timestep_batch, + negative, + encoder_attention_mask=negative_mask, + return_dict=False, + )[0].float() + conditional = prediction + guidance_scale = batch.guidance_scale_2 if self.refiner else batch.guidance_scale + if guidance_scale is None: + raise ValueError("LingBot-Video CFG requires a guidance scale") + prediction = unconditional + guidance_scale * (conditional - unconditional) + latents = self.scheduler.step( + prediction, + timestep, + latents, + return_dict=False, + generator=batch.generator, + )[0].float() + if batch.return_trajectory_latents: + trajectory.append(latents.detach().cpu()) + trajectory_timesteps.append(timestep.detach().cpu()) + batch.latents = latents + if trajectory: + batch.trajectory_latents = torch.stack(trajectory, dim=1) + batch.trajectory_timesteps = trajectory_timesteps + return batch diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 4859f7929..b8faf4983 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -178,6 +178,7 @@ class ForwardBatch: num_inference_steps: int = 50 num_inference_steps_sr: int = 50 guidance_scale: float = 1.0 + batch_cfg: bool = False guidance_scale_2: float | None = None cfg_normalization: bool = False cfg_truncation: float | None = 1.0 diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index 15a8e27a4..b95abd287 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/pipelines/stages/decoding.py @@ -65,7 +65,11 @@ class DecodingStage(PipelineStage): ch_dim = 1 if latents.ndim >= 4 else -1 return latents.shape[ch_dim] == packed_ch - def _denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor: + def _denormalize_latents( + self, + latents: torch.Tensor, + fastvideo_args: FastVideoArgs, + ) -> torch.Tensor: """Convert normalized latents into the VAE's expected latent space.""" # Some VAEs handle latent (de)normalization internally. if bool(getattr(self.vae, "handles_latent_denorm", False)): @@ -79,6 +83,15 @@ class DecodingStage(PipelineStage): latents_mean_value = getattr(cfg, "latents_mean", None) if cfg is not None else None latents_std_value = getattr(cfg, "latents_std", None) if cfg is not None else None if latents_mean_value is not None and latents_std_value is not None: + # Preserve the released LingBot reciprocal arithmetic and its float32 rounding for numeric alignment + if type(fastvideo_args.pipeline_config).__name__.startswith("LingBotVideo"): + latents_mean = torch.tensor(latents_mean_value, device=latents.device, + dtype=torch.float32).view(1, -1, 1, 1, 1) + latents_std_inv = 1.0 / torch.tensor(latents_std_value, device=latents.device, + dtype=torch.float32).view(1, -1, 1, 1, 1) + # See Lingbot: https://github.com/Robbyant/lingbot-video/blob/a638721cf2271804d02738b69f2ad788c4a559fc/lingbot_video/pipeline_lingbot_video.py#L282 + return latents.float() / latents_std_inv + latents_mean + latents_mean = torch.tensor(latents_mean_value, device=latents.device, dtype=latents.dtype).view(1, -1, 1, 1, 1) latents_std = torch.tensor(latents_std_value, device=latents.device, @@ -156,7 +169,12 @@ class DecodingStage(PipelineStage): # Flux2: skip denormalize on packed latents; BN denorm runs below instead if not (latents.ndim == 5 and self._is_flux2_packed(latents)): - latents = self._denormalize_latents(latents) + latents = self._denormalize_latents(latents, fastvideo_args) + + # The released LingBot decoder enters Wan VAE in channels-last-3d + # layout; matching that boundary avoids selecting a different Conv3d path. + if (type(fastvideo_args.pipeline_config).__name__.startswith("LingBotVideo") and latents.ndim == 5): + latents = latents.contiguous(memory_format=torch.channels_last_3d) # Decode latents with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): @@ -221,7 +239,7 @@ class DecodingStage(PipelineStage): vae_dtype = PRECISION_TO_TYPE[decode_precision] vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast - latents = self._denormalize_latents(latents) + latents = self._denormalize_latents(latents, fastvideo_args) # Initialize cache if needed if cache is None: diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 0db2cc6ec..1093d984d 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -29,6 +29,7 @@ from fastvideo.configs.pipelines.hunyuan15 import (Hunyuan15T2V480PConfig, Hunyu Hunyuan15SR1080PConfig) from fastvideo.configs.pipelines.hyworld import HYWorldConfig from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5I2VConfig, Kandinsky5T2VConfig +from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig @@ -521,6 +522,28 @@ def _register_configs() -> None: model_family="lingbotworld", default_preset="lingbotworld_i2v", ) + # LingBot-Video MoE T2V with the released second-stage refiner. + register_configs( + sampling_param_cls=None, + pipeline_config_cls=LingBotVideoT2VConfig, + workload_types=(WorkloadType.T2V, ), + hf_model_paths=["FastVideo/LingBot-Video-MoE-30B-A3B-Diffusers"], + model_detectors=[lambda path: "lingbotvideomoepipeline" in path.lower()], + model_family="lingbot_video", + default_preset="lingbot_video_moe_refiner_t2v", + pipeline_cls_name="LingBotVideoPipeline", + ) + # LingBot-Video Dense T2V + register_configs( + sampling_param_cls=None, + pipeline_config_cls=LingBotVideoT2VConfig, + workload_types=(WorkloadType.T2V, ), + hf_model_paths=["FastVideo/LingBot-Video-Dense-1.3B-Diffusers"], + model_detectors=[lambda path: "lingbotvideodensepipeline" in path.lower()], + model_family="lingbot_video", + default_preset="lingbot_video_dense_t2v", + pipeline_cls_name="LingBotVideoPipeline", + ) def _kandinsky5_detector(require: tuple[str, ...] = (), exclude: tuple[str, ...] = ()) -> Callable[[str], bool]: @@ -1246,6 +1269,8 @@ def _register_presets() -> None: ALL_PRESETS as LINGBOTWORLD_PRESETS, ) from fastvideo.pipelines.basic.lingbotworld2.presets import ( ALL_PRESETS as LINGBOTWORLD2_PRESETS, ) + from fastvideo.pipelines.basic.lingbot_video.presets import ( + ALL_PRESETS as LINGBOT_VIDEO_PRESETS, ) from fastvideo.pipelines.basic.longcat.presets import ( ALL_PRESETS as LONGCAT_PRESETS, ) from fastvideo.pipelines.basic.ltx2.presets import ( @@ -1277,6 +1302,7 @@ def _register_presets() -> None: HUNYUAN15_PRESETS, HYWORLD_PRESETS, KANDINSKY5_PRESETS, + LINGBOT_VIDEO_PRESETS, LINGBOTWORLD_PRESETS, LINGBOTWORLD2_PRESETS, LONGCAT_PRESETS, diff --git a/fastvideo/tests/api/test_compat_translation.py b/fastvideo/tests/api/test_compat_translation.py index 6100588ef..2c8002eda 100644 --- a/fastvideo/tests/api/test_compat_translation.py +++ b/fastvideo/tests/api/test_compat_translation.py @@ -8,8 +8,11 @@ from __future__ import annotations from fastvideo.api.compat import ( generator_config_to_fastvideo_args, legacy_from_pretrained_to_config, + request_to_sampling_param, ) -from fastvideo.api.schema import CompileConfig, GeneratorConfig +from fastvideo.api.parser import parse_config +from fastvideo.api.schema import CompileConfig, GenerationRequest, GeneratorConfig +from fastvideo.api.sampling_param import SamplingParam class TestLegacyTorchCompileKwargsTranslation: @@ -35,8 +38,7 @@ class TestCompileConfigRoundTrip: _stub_fastvideo_args_from_kwargs(monkeypatch) config = GeneratorConfig( model_path="/models/ltx2", - engine=_engine_with_compile( - CompileConfig(enabled=True, backend="inductor", fullgraph=True)), + engine=_engine_with_compile(CompileConfig(enabled=True, backend="inductor", fullgraph=True)), ) args = generator_config_to_fastvideo_args(config) assert args.kwargs["enable_torch_compile"] is True @@ -53,13 +55,17 @@ class TestCompileConfigRoundTrip: CompileConfig( enabled=True, mode="reduce-overhead", - extras={"options": {"triton.cudagraphs": False}}, + extras={"options": { + "triton.cudagraphs": False + }}, )), ) args = generator_config_to_fastvideo_args(config) assert args.kwargs["torch_compile_kwargs"] == { "mode": "reduce-overhead", - "options": {"triton.cudagraphs": False}, + "options": { + "triton.cudagraphs": False + }, } def test_none_fields_suppressed(self, monkeypatch) -> None: @@ -137,8 +143,7 @@ class TestLegacyTextEncoderCompileTranslation: _stub_fastvideo_args_from_kwargs(monkeypatch) config = GeneratorConfig( model_path="/models/ltx2", - engine=_engine_with_compile( - CompileConfig(text_encoder_enabled=True)), + engine=_engine_with_compile(CompileConfig(text_encoder_enabled=True)), ) args = generator_config_to_fastvideo_args(config) assert args.kwargs["enable_torch_compile_text_encoder"] is True @@ -153,6 +158,26 @@ class TestLegacyTextEncoderCompileTranslation: assert "enable_torch_compile_text_encoder" not in args.kwargs +def test_batch_cfg_typed_request_reaches_sampling_param(monkeypatch) -> None: + """Preserve explicit batched CFG through the canonical request adapter.""" + monkeypatch.setattr( + SamplingParam, + "from_pretrained", + classmethod(lambda cls, model_path: cls()), + ) + request = parse_config( + GenerationRequest, + { + "prompt": "fox", + "sampling": { + "batch_cfg": True + } + }, + ) + sampling_param = request_to_sampling_param(request, model_path="test-model") + assert sampling_param.batch_cfg is True + + # ------------------------------------------------------------------- # Helpers # ------------------------------------------------------------------- diff --git a/fastvideo/tests/api/test_parser.py b/fastvideo/tests/api/test_parser.py index 53968a6ed..b812584f2 100644 --- a/fastvideo/tests/api/test_parser.py +++ b/fastvideo/tests/api/test_parser.py @@ -41,15 +41,15 @@ def test_parse_config_builds_nested_typed_config() -> None: }, "state": { "kind": "ltx2_continuation", - "payload": {"segment_index": 1}, + "payload": { + "segment_index": 1 + }, }, "plan": { - "stages": [ - { - "name": "base", - "kind": "sample", - } - ] + "stages": [{ + "name": "base", + "kind": "sample", + }] }, }, } @@ -77,10 +77,14 @@ def test_parse_config_accepts_existing_typed_instance() -> None: def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None: raw = { - "generator": {"model_path": "/models/wan"}, + "generator": { + "model_path": "/models/wan" + }, "request": { "prompt": "hello", - "sampling": {"num_frames": 16}, + "sampling": { + "num_frames": 16 + }, }, } path = tmp_path / "run.yaml" @@ -182,6 +186,7 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None: "num_inference_steps": 50, "num_inference_steps_sr": 50, "guidance_scale": 1.0, + "batch_cfg": False, "guidance_scale_2": None, "cfg_normalization": False, "cfg_truncation": 1.0, @@ -213,9 +218,15 @@ def test_load_run_config_supports_yaml_roundtrip(tmp_path) -> None: def test_load_serve_config_supports_json_roundtrip(tmp_path) -> None: raw = { - "generator": {"model_path": "/models/server"}, - "server": {"port": 9000}, - "default_request": {"prompt": "serve default"}, + "generator": { + "model_path": "/models/server" + }, + "server": { + "port": 9000 + }, + "default_request": { + "prompt": "serve default" + }, } path = tmp_path / "serve.json" path.write_text(json.dumps(raw), encoding="utf-8") @@ -235,7 +246,9 @@ def test_serve_config_streaming_defaults_to_none() -> None: def test_serve_config_parses_streaming_block() -> None: raw = { - "generator": {"model_path": "/models/server"}, + "generator": { + "model_path": "/models/server" + }, "streaming": { "session_timeout_seconds": 120, "generation_segment_cap": 4, @@ -270,8 +283,7 @@ def test_serve_config_parses_streaming_block() -> None: assert loaded.streaming.session_timeout_seconds == 120 assert loaded.streaming.generation_segment_cap == 4 assert loaded.streaming.stream_mode == "legacy_jpeg" - assert loaded.streaming.warmup == WarmupConfig( - enabled=False, prompt="warmup prompt", timeout_seconds=600) + assert loaded.streaming.warmup == WarmupConfig(enabled=False, prompt="warmup prompt", timeout_seconds=600) assert loaded.streaming.pool == GpuPoolConfig( num_workers=2, enable_audio_reencode=False, @@ -284,14 +296,17 @@ def test_serve_config_parses_streaming_block() -> None: timeout_ms=10000, system_prompt_dir="/opt/prompts", ) - assert loaded.streaming.safety == PromptSafetyConfig( - enabled=True, classifier_path="/opt/safety.pt") + assert loaded.streaming.safety == PromptSafetyConfig(enabled=True, classifier_path="/opt/safety.pt") def test_serve_config_streaming_round_trip_through_config_to_dict() -> None: raw = { - "generator": {"model_path": "/models/server"}, - "streaming": {"session_timeout_seconds": 600}, + "generator": { + "model_path": "/models/server" + }, + "streaming": { + "session_timeout_seconds": 600 + }, } loaded = parse_config(ServeConfig, raw) dumped = config_to_dict(loaded) @@ -304,10 +319,14 @@ def test_serve_config_streaming_round_trip_through_config_to_dict() -> None: def test_load_serve_config_with_streaming_from_yaml(tmp_path) -> None: raw = { - "generator": {"model_path": "/models/server"}, + "generator": { + "model_path": "/models/server" + }, "streaming": { "stream_mode": "av_fmp4", - "pool": {"num_workers": 4}, + "pool": { + "num_workers": 4 + }, }, } path = tmp_path / "serve.yaml" diff --git a/fastvideo/tests/api/test_presets.py b/fastvideo/tests/api/test_presets.py index 33563be79..66cf8be55 100644 --- a/fastvideo/tests/api/test_presets.py +++ b/fastvideo/tests/api/test_presets.py @@ -1,4 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 +import json + import pytest from fastvideo.api.errors import ConfigValidationError @@ -14,11 +16,11 @@ from fastvideo.api.presets import ( validate_stage_overrides, ) - # ------------------------------------------------------------------- # Fixtures # ------------------------------------------------------------------- + @pytest.fixture() def _isolated_registry(): """Run each test with an empty preset registry, restoring after.""" @@ -49,11 +51,11 @@ _NO_OVERRIDES_STAGE = PresetStageSpec( def _make_preset( - name: str = "test_preset", - version: int = 1, - model_family: str = "test", - stage_schemas: tuple[PresetStageSpec, ...] = (_SIMPLE_STAGE, ), - **kwargs, + name: str = "test_preset", + version: int = 1, + model_family: str = "test", + stage_schemas: tuple[PresetStageSpec, ...] = (_SIMPLE_STAGE, ), + **kwargs, ) -> InferencePreset: return InferencePreset( name=name, @@ -112,8 +114,7 @@ class TestRegistration: def test_get_presets_for_family(self) -> None: register_preset(_make_preset(name="a")) register_preset(_make_preset(name="b")) - register_preset(_make_preset( - name="c", model_family="other")) + register_preset(_make_preset(name="c", model_family="other")) result = get_presets_for_family("test") assert {p.name for p in result} == {"a", "b"} @@ -133,14 +134,11 @@ class TestStageNameValidation: def test_valid_stage_name_passes(self) -> None: preset = _make_preset(stage_schemas=(_SIMPLE_STAGE, _SR_STAGE)) - validate_stage_names( - preset, {"denoise": {}, "sr": {}}) + validate_stage_names(preset, {"denoise": {}, "sr": {}}) def test_unknown_stage_name_raises(self) -> None: preset = _make_preset(stage_schemas=(_SIMPLE_STAGE, )) - with pytest.raises( - ConfigValidationError, match="stage_overrides.bogus" - ): + with pytest.raises(ConfigValidationError, match="stage_overrides.bogus"): validate_stage_names(preset, {"bogus": {}}) def test_empty_overrides_passes(self) -> None: @@ -148,11 +146,8 @@ class TestStageNameValidation: validate_stage_names(preset, {}) def test_error_lists_valid_stages(self) -> None: - preset = _make_preset( - stage_schemas=(_SIMPLE_STAGE, _SR_STAGE)) - with pytest.raises( - ConfigValidationError, match="'denoise'" - ): + preset = _make_preset(stage_schemas=(_SIMPLE_STAGE, _SR_STAGE)) + with pytest.raises(ConfigValidationError, match="'denoise'"): validate_stage_names(preset, {"nope": {}}) @@ -167,29 +162,35 @@ class TestStageOverrideValidation: preset = _make_preset(stage_schemas=(_SIMPLE_STAGE, )) validate_stage_overrides( preset, - {"denoise": {"num_inference_steps": 25}}, + {"denoise": { + "num_inference_steps": 25 + }}, ) def test_disallowed_override_raises(self) -> None: preset = _make_preset(stage_schemas=(_SIMPLE_STAGE, )) with pytest.raises( - ConfigValidationError, - match="stage_overrides.denoise.height", + ConfigValidationError, + match="stage_overrides.denoise.height", ): validate_stage_overrides( preset, - {"denoise": {"height": 720}}, + {"denoise": { + "height": 720 + }}, ) def test_override_on_stage_with_no_allowed_raises(self) -> None: preset = _make_preset(stage_schemas=(_NO_OVERRIDES_STAGE, )) with pytest.raises( - ConfigValidationError, - match="does not accept overrides", + ConfigValidationError, + match="does not accept overrides", ): validate_stage_overrides( preset, - {"encode": {"some_key": 1}}, + {"encode": { + "some_key": 1 + }}, ) def test_empty_override_on_no_allowed_passes(self) -> None: @@ -199,23 +200,21 @@ class TestStageOverrideValidation: def test_non_mapping_override_raises(self) -> None: preset = _make_preset(stage_schemas=(_SIMPLE_STAGE, )) with pytest.raises(ConfigValidationError, match="mapping"): - validate_stage_overrides( - preset, {"denoise": "not a dict"}) + validate_stage_overrides(preset, {"denoise": "not a dict"}) def test_unknown_stage_still_caught(self) -> None: preset = _make_preset(stage_schemas=(_SIMPLE_STAGE, )) with pytest.raises(ConfigValidationError, match="unknown"): - validate_stage_overrides( - preset, {"missing_stage": {"a": 1}}) + validate_stage_overrides(preset, {"missing_stage": {"a": 1}}) def test_error_lists_allowed_overrides(self) -> None: preset = _make_preset(stage_schemas=(_SIMPLE_STAGE, )) - with pytest.raises( - ConfigValidationError, match="guidance_scale" - ): + with pytest.raises(ConfigValidationError, match="guidance_scale"): validate_stage_overrides( preset, - {"denoise": {"bad_key": 1}}, + {"denoise": { + "bad_key": 1 + }}, ) @@ -228,15 +227,13 @@ class TestValidatePresetSelection: @pytest.mark.usefixtures("_isolated_registry") def test_none_preset_returns_none(self) -> None: - assert validate_preset_selection( - None, "test") is None + assert validate_preset_selection(None, "test") is None @pytest.mark.usefixtures("_isolated_registry") def test_valid_preset_resolves(self) -> None: p = _make_preset() register_preset(p) - result = validate_preset_selection( - "test_preset", "test") + result = validate_preset_selection("test_preset", "test") assert result is p @pytest.mark.usefixtures("_isolated_registry") @@ -246,7 +243,9 @@ class TestValidatePresetSelection: result = validate_preset_selection( "test_preset", "test", - stage_overrides={"denoise": {"guidance_scale": 2.0}}, + stage_overrides={"denoise": { + "guidance_scale": 2.0 + }}, ) assert result is p @@ -262,7 +261,9 @@ class TestValidatePresetSelection: validate_preset_selection( "test_preset", "test", - stage_overrides={"denoise": {"bad": 1}}, + stage_overrides={"denoise": { + "bad": 1 + }}, ) @@ -307,19 +308,22 @@ class TestWanPresets: # Valid override. validate_stage_overrides( preset, - {"denoise": {"num_inference_steps": 25}}, + {"denoise": { + "num_inference_steps": 25 + }}, ) # Invalid override key. with pytest.raises(ConfigValidationError): validate_stage_overrides( preset, - {"denoise": {"height": 1080}}, + {"denoise": { + "height": 1080 + }}, ) def test_wan_model_family_in_registry(self) -> None: from fastvideo.registry import get_model_family - family = get_model_family( - "Wan-AI/Wan2.1-T2V-1.3B-Diffusers") + family = get_model_family("Wan-AI/Wan2.1-T2V-1.3B-Diffusers") assert family == "wan" @@ -369,22 +373,17 @@ class TestLtx2Presets: def test_ltx2_two_stage_refine_overrides_valid(self) -> None: import fastvideo.registry # noqa: F401 p = get_preset("ltx2_two_stage", "ltx2") - validate_stage_overrides( - p, {"refine": {"num_inference_steps": 3}}) - validate_stage_overrides( - p, {"refine": {"guidance_scale": 1.0}}) - validate_stage_overrides( - p, {"refine": {"image_crf": 18}}) - validate_stage_overrides( - p, {"refine": {"video_position_offset_sec": 2.5}}) + validate_stage_overrides(p, {"refine": {"num_inference_steps": 3}}) + validate_stage_overrides(p, {"refine": {"guidance_scale": 1.0}}) + validate_stage_overrides(p, {"refine": {"image_crf": 18}}) + validate_stage_overrides(p, {"refine": {"video_position_offset_sec": 2.5}}) def test_ltx2_two_stage_rejects_unknown_refine_override(self) -> None: import fastvideo.registry # noqa: F401 from fastvideo.api.errors import ConfigValidationError p = get_preset("ltx2_two_stage", "ltx2") with pytest.raises(ConfigValidationError): - validate_stage_overrides( - p, {"refine": {"bogus_field": 1}}) + validate_stage_overrides(p, {"refine": {"bogus_field": 1}}) # ------------------------------------------------------------------- @@ -437,12 +436,10 @@ class TestHunyuan15Presets: import fastvideo.registry # noqa: F401 p = get_preset("hunyuan15_sr_1080p", "hunyuan15") # Valid: override sr num_inference_steps. - validate_stage_overrides( - p, {"sr": {"num_inference_steps": 12}}) + validate_stage_overrides(p, {"sr": {"num_inference_steps": 12}}) # Invalid: height not in sr allowed_overrides. with pytest.raises(ConfigValidationError): - validate_stage_overrides( - p, {"sr": {"height": 1080}}) + validate_stage_overrides(p, {"sr": {"height": 1080}}) # ------------------------------------------------------------------- @@ -471,8 +468,7 @@ class TestCosmosPresets: """ import fastvideo.registry # noqa: F401 from fastvideo.api.sampling_param import SamplingParam - sp = SamplingParam.from_pretrained( - "KyleShao/Cosmos-Predict2.5-2B-Diffusers") + sp = SamplingParam.from_pretrained("KyleShao/Cosmos-Predict2.5-2B-Diffusers") assert sp.guidance_scale == 7.0 assert sp.num_inference_steps == 35 @@ -523,6 +519,116 @@ class TestSD35Presets: assert p.defaults["num_frames"] == 1 +# ------------------------------------------------------------------- +# LingBot-Video preset integration (Dense and MoE/refiner) +# ------------------------------------------------------------------- + + +class TestLingBotVideoPresets: + + @pytest.mark.parametrize( + ("directory_name", "pipeline_class", "expected_preset"), + ( + ("arbitrary-layout-a", "LingBotVideoDensePipeline", "lingbot_video_dense_t2v"), + ("arbitrary-layout-b", "LingBotVideoMoePipeline", "lingbot_video_moe_refiner_t2v"), + ), + ) + def test_local_layouts_select_one_variant( + self, + tmp_path, + caplog, + directory_name: str, + pipeline_class: str, + expected_preset: str, + ) -> None: + """Resolve canonical local Dense and MoE paths without detector overlap.""" + from fastvideo.registry import get_preset_selection + model_dir = tmp_path / directory_name + model_dir.mkdir() + (model_dir / "transformer").mkdir() + (model_dir / "model_index.json").write_text( + json.dumps({ + "_class_name": pipeline_class, + "_diffusers_version": "0.39.0" + }), + encoding="utf-8", + ) + assert get_preset_selection(str(model_dir)) == (expected_preset, "lingbot_video") + assert "Multiple models matched" not in caplog.text + + def test_lingbot_video_presets_registered(self) -> None: + """Register both released LingBot-Video inference layouts.""" + import fastvideo.registry # noqa: F401 + from fastvideo.registry import get_preset_selection + + presets = get_presets_for_family("lingbot_video") + assert {preset.name + for preset in presets} == { + "lingbot_video_dense_t2v", + "lingbot_video_moe_refiner_t2v", + } + assert get_preset_selection("FastVideo/LingBot-Video-Dense-1.3B-Diffusers") == ( + "lingbot_video_dense_t2v", + "lingbot_video", + ) + assert get_preset_selection("FastVideo/LingBot-Video-MoE-30B-A3B-Diffusers") == ( + "lingbot_video_moe_refiner_t2v", + "lingbot_video", + ) + + def test_moe_refiner_defaults(self) -> None: + """Expose the official base and 1080p refiner sampling defaults.""" + import fastvideo.registry # noqa: F401 + preset = get_preset("lingbot_video_moe_refiner_t2v", "lingbot_video") + expected = { + "height": 480, + "width": 832, + "height_sr": 1088, + "width_sr": 1920, + "num_frames": 121, + "num_inference_steps": 40, + "num_inference_steps_sr": 8, + "guidance_scale": 3.0, + "guidance_scale_2": 3.0, + "batch_cfg": True, + "t_thresh": 0.85, + "seed": 42, + } + assert {key: preset.defaults[key] for key in expected} == expected + assert preset.stage_defaults["refine"] == { + key: expected[key] + for key in ("height_sr", "width_sr", "num_inference_steps_sr", "guidance_scale_2", "t_thresh") + } + + def test_moe_refiner_stage_schema_uses_sampling_fields(self) -> None: + """Accept only declared ``SamplingParam`` fields as stage overrides.""" + import fastvideo.registry # noqa: F401 + from fastvideo.api.sampling_param import SamplingParam + preset = get_preset("lingbot_video_moe_refiner_t2v", "lingbot_video") + sampling_fields = set(SamplingParam.__dataclass_fields__) + assert [(stage.name, stage.kind) for stage in preset.stage_schemas] == [ + ("denoise", "denoising"), + ("refine", "refinement"), + ] + assert set(preset.defaults) <= sampling_fields + assert all(stage.allowed_overrides <= sampling_fields for stage in preset.stage_schemas) + validate_stage_overrides( + preset, + { + "denoise": { + "num_inference_steps": 32, + "guidance_scale": 2.5 + }, + "refine": { + "height_sr": 1088, + "t_thresh": 0.8 + }, + }, + ) + with pytest.raises(ConfigValidationError): + validate_stage_overrides(preset, {"refine": {"refiner_steps": 8}}) + + # ------------------------------------------------------------------- # LingBotWorld preset integration (dual guidance) # ------------------------------------------------------------------- @@ -540,11 +646,9 @@ class TestLingBotWorldPresets: def test_lingbotworld_override_validation(self) -> None: import fastvideo.registry # noqa: F401 p = get_preset("lingbotworld_i2v", "lingbotworld") - validate_stage_overrides( - p, {"denoise": {"boundary_ratio": 0.95}}) + validate_stage_overrides(p, {"denoise": {"boundary_ratio": 0.95}}) with pytest.raises(ConfigValidationError): - validate_stage_overrides( - p, {"denoise": {"height": 720}}) + validate_stage_overrides(p, {"denoise": {"height": 720}}) # ------------------------------------------------------------------- @@ -583,9 +687,7 @@ class TestSinglePresetFamilies: import fastvideo.registry # noqa: F401 presets = get_presets_for_family("longcat") names = {p.name for p in presets} - assert names == { - "longcat_t2v", "longcat_i2v", "longcat_vc" - } + assert names == {"longcat_t2v", "longcat_i2v", "longcat_vc"} # ------------------------------------------------------------------- @@ -628,12 +730,10 @@ class TestPresetDefaultTypes: import fastvideo.registry # noqa: F401 from fastvideo.api.presets import _PRESET_REGISTRY offenders = [ - f"{preset.model_family}/{preset.name}" - for preset in _PRESET_REGISTRY.values() + f"{preset.model_family}/{preset.name}" for preset in _PRESET_REGISTRY.values() if preset.defaults.get("negative_prompt", "") is None ] - assert not offenders, ( - "These presets set negative_prompt=None, which violates " - "SamplingParam.negative_prompt's typed str contract and " - "crashes the CFG path in text_encoding. Use \"\" instead:\n" - + "\n".join(f" - {p}" for p in offenders)) + assert not offenders, ("These presets set negative_prompt=None, which violates " + "SamplingParam.negative_prompt's typed str contract and " + "crashes the CFG path in text_encoding. Use \"\" instead:\n" + + "\n".join(f" - {p}" for p in offenders)) diff --git a/fastvideo/tests/distributed/test_sp_lingbot_video.py b/fastvideo/tests/distributed/test_sp_lingbot_video.py new file mode 100644 index 000000000..fcf8668d8 --- /dev/null +++ b/fastvideo/tests/distributed/test_sp_lingbot_video.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Forward parity coverage for LingBot-Video sequence parallelism.""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path +import socket +import subprocess +import sys + +import pytest +import torch +import torch.distributed as dist +from torch.testing import assert_close + +from fastvideo.configs.models.dits.lingbot_video import LingBotVideoArchConfig, LingBotVideoConfig +from fastvideo.distributed import ( + cleanup_dist_env_and_memory, + maybe_init_distributed_environment_and_model_parallel, +) +from fastvideo.models.dits.lingbot_video import LingBotVideoTransformer3DModel + +SP_WORLD_SIZE = 2 +SEED = 20260711 + + +def _free_port() -> int: + """Reserve a currently unused localhost port for one torchrun invocation.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _seed_everything(seed: int) -> None: + """Make model initialization and raw SDPA deterministic across runs.""" + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.backends.cuda.enable_cudnn_sdp(False) + torch.backends.cuda.enable_flash_sdp(False) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_math_sdp(True) + + +def _build_tiny_config() -> LingBotVideoConfig: + """Build a two-block Dense model whose four heads divide over two ranks.""" + arch_config = LingBotVideoArchConfig( + patch_size=(1, 1, 1), + in_channels=4, + out_channels=4, + hidden_size=32, + num_attention_heads=4, + depth=2, + intermediate_size=64, + text_dim=12, + freq_dim=16, + axes_dims=(2, 2, 4), + axes_lens=(16, 8, 8), + ) + return LingBotVideoConfig(arch_config=arch_config) + + +def _initialize_model_parameters(model: torch.nn.Module) -> None: + """Initialize empty replicated-linear parameters identically in each run.""" + torch.manual_seed(SEED + 1) + torch.cuda.manual_seed_all(SEED + 1) + with torch.no_grad(): + for name, parameter in model.named_parameters(): + if parameter.ndim <= 1: + if name.endswith("weight") and "norm" in name: + parameter.fill_(1.0) + else: + parameter.zero_() + continue + torch.nn.init.xavier_uniform_(parameter) + + +def _build_inputs(config: LingBotVideoConfig, device: torch.device) -> dict[str, torch.Tensor]: + """Create B=2 inputs with unequal masks and a seven-token joint sequence.""" + generator = torch.Generator(device="cpu").manual_seed(SEED + 2) + hidden_states = torch.randn( + 2, + config.in_channels, + 1, + 2, + 2, + generator=generator, + dtype=torch.float32, + ) + encoder_hidden_states = torch.randn( + 2, + 3, + config.text_dim, + generator=generator, + dtype=torch.float32, + ) + encoder_attention_mask = torch.tensor([[1, 1, 1], [1, 0, 0]], dtype=torch.long) + joint_length = 1 * 2 * 2 + encoder_hidden_states.shape[1] + assert joint_length == 7 and joint_length % SP_WORLD_SIZE != 0 + assert encoder_attention_mask.sum(dim=1).tolist() == [3, 1] + return { + "hidden_states": hidden_states.to(device), + "encoder_hidden_states": encoder_hidden_states.to(device), + "encoder_attention_mask": encoder_attention_mask.to(device), + "timestep": torch.tensor([10.0, 20.0], device=device), + } + + +def _run_worker(mode: str, output_path: Path) -> None: + """Run one deterministic single-rank or sequence-parallel model forward.""" + if mode not in {"single", "sp"}: + raise ValueError(f"Unsupported 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) + _seed_everything(SEED) + + try: + maybe_init_distributed_environment_and_model_parallel(1, sp_size) + config = _build_tiny_config() + model = LingBotVideoTransformer3DModel(config=config, hf_config={}).to(device=device, dtype=torch.float32) + _initialize_model_parameters(model) + model.eval() + + with torch.inference_mode(): + output = model(**_build_inputs(config, device)).sample + 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: + """Launch one isolated worker group and surface its complete failure output.""" + command = [ + sys.executable, + "-m", + "torch.distributed.run", + "--nnodes", + "1", + "--nproc_per_node", + str(nproc_per_node), + "--master_port", + str(_free_port()), + str(script_path), + "--sp-worker", + "--mode", + mode, + "--output", + str(output_path), + ] + environment = os.environ.copy() + 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_sp_forward_matches_single_rank(tmp_path: Path) -> None: + """Compare full model outputs from one rank and a padded two-rank shard.""" + if not torch.cuda.is_available(): + pytest.skip("This test requires CUDA.") + if torch.cuda.device_count() < SP_WORLD_SIZE: + pytest.skip(f"This test requires at least {SP_WORLD_SIZE} CUDA devices.") + + script_path = Path(__file__).resolve() + single_path = tmp_path / "single_rank_output.pt" + sp_path = tmp_path / f"sp{SP_WORLD_SIZE}_output.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 == (2, 4, 1, 2, 2) + assert_close(sp_output, single_output, atol=1e-5, rtol=1e-5) + + +def _parse_args() -> argparse.Namespace: + """Parse the private worker interface used by the parent pytest process.""" + parser = argparse.ArgumentParser() + parser.add_argument("--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.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_worker(mode=args.mode, output_path=Path(args.output)) diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index ab4cdd59d..6d64ced98 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -7,6 +7,7 @@ import torch import fastvideo.entrypoints.video_generator as video_generator_module from fastvideo.api import ( + ConfigValidationError, GenerationRequest, GenerationResult, GeneratorConfig, @@ -15,7 +16,7 @@ from fastvideo.api import ( load_run_config, ) from fastvideo.api.sampling_param import SamplingParam -from fastvideo.entrypoints.video_generator import VideoGenerator +from fastvideo.entrypoints.video_generator import VideoGenerator, _resolve_output_size from fastvideo.fastvideo_args import WorkloadType from fastvideo.pipelines import ForwardBatch from fastvideo.worker.gpu_worker import Worker @@ -81,6 +82,7 @@ def _patch_fastvideo_args_from_kwargs(monkeypatch): def _patch_sampling_param_from_pretrained(monkeypatch): + def fake_from_pretrained(cls, model_path): return cls() @@ -141,6 +143,13 @@ def _single_video_generator(output_batch, fastvideo_args): return generator +def test_resolve_output_size_uses_produced_pixel_geometry() -> None: + """Report refined dimensions instead of the base request geometry.""" + samples = torch.empty(1, 3, 5, 64, 96) + assert _resolve_output_size(samples, (32, 48, 1), pixel_output=True) == (64, 96, 5) + assert _resolve_output_size(samples, (32, 48, 1), pixel_output=False) == (32, 48, 1) + + def test_prepare_output_path_file_sanitization(tmp_path): vg = _new_video_generator() target_dir = tmp_path / "dir" @@ -523,8 +532,7 @@ def test_generate_uses_typed_request_path(monkeypatch): GenerationRequest( prompt="hello world", sampling=SamplingConfig(num_frames=81, height=480, width=832), - ) - ) + )) assert isinstance(result, GenerationResult) assert captured["prompt"] == "hello world" @@ -534,6 +542,24 @@ def test_generate_uses_typed_request_path(monkeypatch): assert result.video_path == "outputs/test.mp4" +def test_generate_rejects_stage_override_outside_registered_stage(monkeypatch) -> None: + """Validate typed stage overrides at the public generation entrypoint.""" + generator = _new_runtime_video_generator() + monkeypatch.setattr( + "fastvideo.registry.get_preset_selection", + lambda _model_path: ("lingbot_video_moe_refiner_t2v", "lingbot_video"), + ) + with pytest.raises(ConfigValidationError, match="stage_overrides.denoise.t_thresh"): + generator.generate({ + "prompt": "hello world", + "stage_overrides": { + "denoise": { + "t_thresh": 0.7 + } + }, + }) + + def test_generate_preserves_schema_defaults_for_dataclass_request(monkeypatch): generator = _new_runtime_video_generator() captured = {} @@ -558,8 +584,7 @@ def test_generate_preserves_schema_defaults_for_dataclass_request(monkeypatch): prompt="hello world", negative_prompt=None, sampling=SamplingConfig(num_frames=125, height=720, width=1280), - ) - ) + )) assert captured["sampling_param"].negative_prompt is None assert captured["sampling_param"].num_frames == 125 @@ -567,9 +592,7 @@ def test_generate_preserves_schema_defaults_for_dataclass_request(monkeypatch): assert captured["sampling_param"].width == 1280 -def test_generate_mapping_request_preserves_model_defaults_for_omitted_fields( - monkeypatch, -): +def test_generate_mapping_request_preserves_model_defaults_for_omitted_fields(monkeypatch, ): generator = _new_runtime_video_generator() captured = {} @@ -590,11 +613,9 @@ def test_generate_mapping_request_preserves_model_defaults_for_omitted_fields( monkeypatch.setattr(SamplingParam, "from_pretrained", classmethod(fake_from_pretrained)) monkeypatch.setattr(generator, "_generate_video_impl", fake_generate_video_impl) - generator.generate( - { - "prompt": "hello world", - } - ) + generator.generate({ + "prompt": "hello world", + }) assert captured["sampling_param"].negative_prompt == "model default" assert captured["sampling_param"].num_frames == 61 @@ -786,16 +807,16 @@ def test_generate_batch_prompt_file_returns_typed_results(tmp_path, monkeypatch) monkeypatch.setattr(generator, "_generate_single_video", fake_generate_single_video) - results = generator.generate( - { - "inputs": {"prompt_path": str(prompt_file)}, - "output": { - "output_path": str(output_dir), - "save_video": False, - "return_frames": False, - }, - } - ) + results = generator.generate({ + "inputs": { + "prompt_path": str(prompt_file) + }, + "output": { + "output_path": str(output_dir), + "save_video": False, + "return_frames": False, + }, + }) assert isinstance(results, list) assert [result.prompt for result in results] == ["first prompt", "second prompt"] @@ -821,8 +842,7 @@ def test_generate_batched_request_fans_out_media_inputs(monkeypatch): image_path=["first.png", "second.png"], video_path=["first.mp4", "second.mp4"], ), - ) - ) + )) assert [result.prompt for result in results] == ["first prompt", "second prompt"] assert captured == [ @@ -840,5 +860,4 @@ def test_generate_batched_request_rejects_mismatched_media_inputs(monkeypatch): GenerationRequest( prompt=["first prompt", "second prompt"], inputs=InputConfig(image_path=["first.png"]), - ) - ) + )) diff --git a/scripts/checkpoint_conversion/lingbot_video_to_diffusers.py b/scripts/checkpoint_conversion/lingbot_video_to_diffusers.py new file mode 100644 index 000000000..7140cd26c --- /dev/null +++ b/scripts/checkpoint_conversion/lingbot_video_to_diffusers.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Convert official LingBot-Video components into a native FastVideo layout.""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +from dataclasses import fields +from pathlib import Path + +import torch +from safetensors.torch import load_file, save_file + +from fastvideo.configs.models.dits.lingbot_video import LingBotVideoConfig +from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig + +LANGUAGE_PREFIX = "model.language_model." +PASSTHROUGH_COMPONENTS = ("transformer", "vae", "scheduler") +DENSE_PIPELINE_CLASS = "LingBotVideoDensePipeline" +MOE_PIPELINE_CLASS = "LingBotVideoMoePipeline" + + +def _copy_tree_with_weight_links(source: Path, destination: Path) -> None: + """Copy metadata and hard-link immutable weight files on the same filesystem.""" + destination.mkdir(parents=True, exist_ok=False) + for source_path in sorted(source.rglob("*")): + relative = source_path.relative_to(source) + destination_path = destination / relative + if source_path.is_dir(): + destination_path.mkdir(parents=True, exist_ok=True) + elif source_path.suffix == ".safetensors": + os.link(source_path, destination_path) + else: + shutil.copy2(source_path, destination_path) + + +def _fuse_language_shard(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Filter Qwen3-VL language tensors and fuse native QKV and gate/up weights.""" + converted = { + name[len(LANGUAGE_PREFIX):]: tensor + for name, tensor in state.items() if name.startswith(LANGUAGE_PREFIX) + } + layer_ids = {int(name.split(".")[1]) for name in converted if name.startswith("layers.")} + for layer_id in sorted(layer_ids): + attention = f"layers.{layer_id}.self_attn" + qkv_names = [f"{attention}.{kind}_proj.weight" for kind in ("q", "k", "v")] + present_qkv = [name in converted for name in qkv_names] + if any(present_qkv): + if not all(present_qkv): + raise KeyError(f"QKV tensors cross shard boundary for layer {layer_id}") + converted[f"{attention}.qkv_proj.weight"] = torch.cat([converted.pop(name) for name in qkv_names], dim=0) + mlp = f"layers.{layer_id}.mlp" + gate_up_names = [f"{mlp}.{kind}_proj.weight" for kind in ("gate", "up")] + present_gate_up = [name in converted for name in gate_up_names] + if any(present_gate_up): + if not all(present_gate_up): + raise KeyError(f"gate/up tensors cross shard boundary for layer {layer_id}") + converted[f"{mlp}.gate_up_proj.weight"] = torch.cat([converted.pop(name) for name in gate_up_names], dim=0) + return converted + + +def _convert_text_encoder(source: Path, destination: Path) -> dict[str, str]: + """Convert each source shard independently and emit a valid safetensors index.""" + destination.mkdir(parents=True, exist_ok=False) + source_index = json.loads((source / "model.safetensors.index.json").read_text()) + source_shards = sorted(set(source_index["weight_map"].values())) + weight_map: dict[str, str] = {} + total_size = 0 + for shard_name in source_shards: + source_state = load_file(source / shard_name, device="cpu") + converted = _fuse_language_shard(source_state) + if not converted: + continue + output_name = shard_name + save_file(converted, destination / output_name, metadata={"format": "pt"}) + for name, tensor in converted.items(): + if name in weight_map: + raise ValueError(f"duplicate converted text-encoder tensor: {name}") + weight_map[name] = output_name + total_size += tensor.numel() * tensor.element_size() + index = {"metadata": {"total_size": total_size}, "weight_map": dict(sorted(weight_map.items()))} + (destination / "model.safetensors.index.json").write_text(json.dumps(index, indent=2, sort_keys=True) + "\n") + return weight_map + + +def _write_text_encoder_config(source: Path, destination: Path) -> dict: + """Flatten the compound Qwen3-VL config into FastVideo's text-model bucket.""" + official = json.loads((source / "config.json").read_text()) + text = dict(official["text_config"]) + valid = {item.name for item in fields(LingBotVideoQwen3VLTextConfig().arch_config)} + config = {key: value for key, value in text.items() if key in valid} + config.update({ + "_class_name": "LingBotVideoQwen3VLTextModel", + "architectures": ["LingBotVideoQwen3VLTextModel"], + "pad_token_id": 151643, + "text_len": 37698, + "output_hidden_states": True, + "require_processor": True, + "rope_scaling": None, + "mrope_interleaved": True, + "mrope_section": [24, 20, 20], + }) + LingBotVideoQwen3VLTextConfig().update_model_arch(config) + (destination / "config.json").write_text(json.dumps(config, indent=2, sort_keys=True) + "\n") + for filename in ("generation_config.json", ): + candidate = source / filename + if candidate.is_file(): + shutil.copy2(candidate, destination / filename) + return config + + +def _validate_transformer_config(path: Path) -> None: + """Exercise the same architecture update used by the production transformer loader.""" + config = json.loads(path.read_text()) + if config.get("_class_name") != "LingBotVideoTransformer3DModel": + raise ValueError("converted transformer config has the wrong _class_name") + LingBotVideoConfig().update_model_arch(config) + + +def _write_model_index(destination: Path, *, has_refiner: bool) -> None: + """Declare the native T2V components, including an optional refiner DiT.""" + model_index = { + "_class_name": MOE_PIPELINE_CLASS if has_refiner else DENSE_PIPELINE_CLASS, + "_diffusers_version": "0.39.0", + "workload_type": "video-generation", + "transformer": ["diffusers", "LingBotVideoTransformer3DModel"], + "vae": ["diffusers", "AutoencoderKLWan"], + "text_encoder": ["transformers", "LingBotVideoQwen3VLTextModel"], + "tokenizer": ["transformers", "Qwen3VLProcessor"], + "scheduler": ["diffusers", "FlowUniPCMultistepScheduler"], + } + if has_refiner: + model_index["transformer_2"] = ["diffusers", "LingBotVideoTransformer3DModel"] + (destination / "model_index.json").write_text(json.dumps(model_index, indent=2, sort_keys=True) + "\n") + + +def convert(source: Path, destination: Path) -> None: + """Create a Dense or MoE/refiner T2V layout without altering source assets.""" + required = (*PASSTHROUGH_COMPONENTS, "text_encoder", "processor") + missing = [name for name in required if not (source / name).is_dir()] + if missing: + raise FileNotFoundError(f"missing official component directories: {missing}") + if destination.exists() and any(destination.iterdir()): + raise FileExistsError(f"refusing to overwrite non-empty destination: {destination}") + destination.mkdir(parents=True, exist_ok=True) + for component in PASSTHROUGH_COMPONENTS: + _copy_tree_with_weight_links(source / component, destination / component) + _validate_transformer_config(destination / "transformer" / "config.json") + refiner_source = source / "refiner" + has_refiner = refiner_source.is_dir() + if has_refiner: + _copy_tree_with_weight_links(refiner_source, destination / "transformer_2") + _validate_transformer_config(destination / "transformer_2" / "config.json") + text_destination = destination / "text_encoder" + weight_map = _convert_text_encoder(source / "text_encoder", text_destination) + text_config = _write_text_encoder_config(source / "text_encoder", text_destination) + _copy_tree_with_weight_links(source / "processor", destination / "tokenizer") + _write_model_index(destination, has_refiner=has_refiner) + print(f"converted text tensors: {len(weight_map)}") + print(f"text architecture: {text_config['architectures'][0]}") + print(f"output: {destination}") + + +def parse_args() -> argparse.Namespace: + """Parse explicit local source and destination paths.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--src", type=Path, required=True) + parser.add_argument("--dst", type=Path, required=True) + return parser.parse_args() + + +def main() -> None: + """Run the deterministic local conversion.""" + args = parse_args() + convert(args.src.resolve(), args.dst.resolve()) + + +if __name__ == "__main__": + main() diff --git a/tests/local_tests/encoders/test_lingbot_video_dense_text_encoder_parity.py b/tests/local_tests/encoders/test_lingbot_video_dense_text_encoder_parity.py new file mode 100644 index 000000000..edb631af4 --- /dev/null +++ b/tests/local_tests/encoders/test_lingbot_video_dense_text_encoder_parity.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical parity for Dense LingBot-Video T2V Qwen3-VL text conditioning.""" + +import os +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +from transformers import AutoProcessor, Qwen3VLForConditionalGeneration + +from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig +from fastvideo.distributed.parallel_state import init_distributed_environment, initialize_model_parallel +from fastvideo.models.encoders.lingbot_video import ( + LingBotVideoQwen3VLAttention, + LingBotVideoQwen3VLDecoderLayer, + LingBotVideoQwen3VLTextModel, +) +from fastvideo.models.loader.component_loader import TextEncoderLoader +from fastvideo.models.loader.utils import set_default_torch_dtype +from tests.local_tests.lingbot_video.hf_assets import ( + FASTVIDEO_DENSE, + OFFICIAL_DENSE, + download_components, +) + +PROMPT_TEMPLATE = ("<|im_start|>system\nGiven a user input that may include a text prompt alone, " + "a text prompt with an image reference, or a text prompt with a video reference " + 'or a video reference alone, generate an "Enhanced prompt" that provides detailed ' + "visual descriptions suitable for video generation. Evaluate the level of detail " + "in the user's input: if it is simple, enrich it by adding specifics about colors, " + "shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal " + "progression, and spatial relationships to create vivid, concrete, and temporally " + "coherent scenes to create vivid and concrete scenes. Please generate only the " + "enhanced description for the prompt below and avoid including any additional " + "commentary or evaluations:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n" + "<|im_start|>assistant\n") + + +def _require_gpu_test() -> torch.device: + """Keep heavyweight parity opt-in and off non-Slurm development nodes.""" + if os.environ.get("LINGBOT_VIDEO_RUN_GPU_TESTS") != "1": + pytest.skip("set LINGBOT_VIDEO_RUN_GPU_TESTS=1 on an allocated GPU") + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the released Qwen3-VL encoder") + return torch.device("cuda") + + +def _load_official(device: torch.device, checkpoint: Path) -> Qwen3VLForConditionalGeneration: + """Load the released full Qwen3-VL wrapper used by the official pipeline.""" + return Qwen3VLForConditionalGeneration.from_pretrained( + checkpoint / "text_encoder", + dtype=torch.bfloat16, + attn_implementation="sdpa", + local_files_only=True, + ).to(device).eval() + + +def _load_native(device: torch.device, checkpoint: Path) -> LingBotVideoQwen3VLTextModel: + """Load converted fused tensors through FastVideo's production text loader.""" + args = SimpleNamespace( + text_encoder_cpu_offload=True, + override_text_encoder_quant=None, + override_text_encoder_safetensors=None, + pin_cpu_memory=False, + ) + model = TextEncoderLoader().load_model( + str(checkpoint / "text_encoder"), + LingBotVideoQwen3VLTextConfig(), + device, + args, + dtype="bf16", + use_text_encoder_override=True, + ) + return model.eval() + + +def _capture_intermediates(model: torch.nn.Module) -> tuple[dict[str, torch.Tensor], list[Any]]: + """Capture shared Qwen3 checkpoints and materialize native deferred residuals.""" + captured: dict[str, torch.Tensor] = {} + handles: list[Any] = [] + suffixes = ( + "embed_tokens", + "input_layernorm", + "q_norm", + "k_norm", + "o_proj", + "post_attention_layernorm", + "down_proj", + "norm", + ) + + def save(name: str): + + def hook(_module, _inputs, output): + if name.count(".") == 1 and name.startswith("layers.") and isinstance(output, tuple): + tensor = output[0] if output[1] is None else output[0] + output[1] + else: + tensor = output[0] if isinstance(output, tuple) else output + if torch.is_tensor(tensor): + captured[name] = tensor.detach().float().cpu() + + return hook + + for name, module in model.named_modules(): + is_layer = name.startswith("layers.") and name.count(".") == 1 + if is_layer or name == "norm" or name.endswith(suffixes): + handles.append(module.register_forward_hook(save(name))) + return captured, handles + + +def _report_first_intermediate_drift(official: dict[str, torch.Tensor], native: dict[str, torch.Tensor]) -> None: + """Print the earliest shared module output with material numerical drift.""" + for name, expected in official.items(): + actual = native.get(name) + if actual is None or actual.shape != expected.shape: + continue + drift = (actual - expected).abs() + if drift.max().item() > 1e-3: + print( + f"first_intermediate_drift={name} max_abs={drift.max().item():.8f} mean_abs={drift.mean().item():.8f}") + return + print("first_intermediate_drift=not_found") + + +def test_lingbot_video_text_encoder_constructs_final_layers_once(monkeypatch: pytest.MonkeyPatch) -> None: + """Build a tiny meta encoder without invoking either replace-after-build constructor.""" + import fastvideo.layers.linear as linear + import fastvideo.layers.vocab_parallel_embedding as embedding + import fastvideo.models.encoders.qwen3 as qwen3 + + monkeypatch.setattr(linear, "get_tp_rank", lambda: 0) + monkeypatch.setattr(linear, "get_tp_world_size", lambda: 1) + monkeypatch.setattr(embedding, "get_tp_rank", lambda: 0) + monkeypatch.setattr(embedding, "get_tp_world_size", lambda: 1) + monkeypatch.setattr(qwen3, "get_tp_world_size", lambda: 1) + + def reject_base_constructor(*_args: Any, **_kwargs: Any) -> None: + raise AssertionError("replace-after-build base constructor was invoked") + + attention_types: list[type[torch.nn.Module]] = [] + original_attention_init = qwen3.Qwen3Attention.__init__ + + def record_attention_type(module: torch.nn.Module, *args: Any, **kwargs: Any) -> None: + attention_types.append(type(module)) + original_attention_init(module, *args, **kwargs) + + monkeypatch.setattr(qwen3.Qwen3ForCausalLM, "__init__", reject_base_constructor) + monkeypatch.setattr(qwen3.Qwen3DecoderLayer, "__init__", reject_base_constructor) + monkeypatch.setattr(qwen3.Qwen3Attention, "__init__", record_attention_type) + config = LingBotVideoQwen3VLTextConfig() + for name, value in { + "vocab_size": 32, + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + "max_position_embeddings": 32, + }.items(): + setattr(config.arch_config, name, value) + with set_default_torch_dtype(torch.bfloat16), torch.device("meta"): + model = LingBotVideoQwen3VLTextModel(config) + + layer_parameter_names = { + "self_attn.qkv_proj.weight", + "self_attn.o_proj.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + "mlp.gate_up_proj.weight", + "mlp.down_proj.weight", + "input_layernorm.weight", + "post_attention_layernorm.weight", + } + expected_names = {"embed_tokens.weight", "norm.weight" + } | {f"layers.{index}.{name}" + for index in range(2) + for name in layer_parameter_names} + assert set(model.state_dict()) == expected_names + assert all(parameter.dtype == torch.bfloat16 for parameter in model.parameters()) + assert all(isinstance(layer, LingBotVideoQwen3VLDecoderLayer) for layer in model.layers) + assert all(isinstance(layer.self_attn, LingBotVideoQwen3VLAttention) for layer in model.layers) + assert attention_types == [LingBotVideoQwen3VLAttention, LingBotVideoQwen3VLAttention] + + +def test_lingbot_video_dense_text_encoder_parity() -> None: + """Compare official and native final hidden states with identical T2V tokens.""" + device = _require_gpu_test() + official_checkpoint = download_components(OFFICIAL_DENSE, "text_encoder", "processor") + fastvideo_checkpoint = download_components(FASTVIDEO_DENSE, "text_encoder") + init_distributed_environment(world_size=1, rank=0, local_rank=0) + initialize_model_parallel(tensor_model_parallel_size=1, sequence_model_parallel_size=1) + processor = AutoProcessor.from_pretrained(official_checkpoint / "processor") + inputs = processor( + text=[PROMPT_TEMPLATE.format("A red fox runs through fresh snow at sunrise.")], + images=None, + videos=None, + do_resize=False, + truncation=True, + max_length=37698, + padding="longest", + return_tensors="pt", + ).to(device) + official = _load_official(device, official_checkpoint) + native = _load_native(device, fastvideo_checkpoint) + official_text_model = official.model.language_model + official_intermediates, official_handles = _capture_intermediates(official_text_model) + native_intermediates, native_handles = _capture_intermediates(native) + with torch.no_grad(): + expected = official( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + output_hidden_states=True, + use_cache=False, + ).hidden_states[-1] + native_outputs = native( + input_ids=inputs["input_ids"].cpu(), + attention_mask=inputs["attention_mask"].cpu(), + output_hidden_states=True, + ) + actual = native_outputs.hidden_states[-1] + for handle in official_handles + native_handles: + handle.remove() + drift = (actual.float() - expected.float()).abs() + print(f"max_abs={drift.max().item():.8f} mean_abs={drift.mean().item():.8f}") + _report_first_intermediate_drift(official_intermediates, native_intermediates) + assert torch.equal(actual, expected) diff --git a/tests/local_tests/lingbot_video/PORT_STATUS.md b/tests/local_tests/lingbot_video/PORT_STATUS.md new file mode 100644 index 000000000..c360abe2f --- /dev/null +++ b/tests/local_tests/lingbot_video/PORT_STATUS.md @@ -0,0 +1,255 @@ +# LingBot-Video Port Status + +## Summary + +- model_family: `lingbot_video` +- workload_types: T2V first, then T2I and TI2V (`WorkloadType.I2V`) +- official_ref: `https://github.com/robbyant/lingbot-video` at `a638721cf2271804d02738b69f2ad788c4a559fc` +- official_weights: `robbyant/lingbot-video-dense-1.3b`; `robbyant/lingbot-video-moe-30b-a3b` +- fastvideo_weights: `FastVideo/LingBot-Video-Dense-1.3B-Diffusers`; `FastVideo/LingBot-Video-MoE-30B-A3B-Diffusers` +- source_layout: `official custom-code Diffusers; converted FastVideo packages published on Hugging Face` +- converted_model_index_classes: `LingBotVideoDensePipeline`; `LingBotVideoMoePipeline` +- local_tests_readme: `tests/local_tests/lingbot_video/README.md` + +## Current Phase + +- phase: `t2v_complete` +- status: `complete` +- owner: `pipeline` +- last_updated: `2026-07-13` + +## Component Matrix + +| Component | Type | Reuse/Port | FastVideo Target | Prototype | Conversion | Parity | Open Issues | +| ---------------- | --------- | ---------- | ----------------------------------------------------- | --------- | ----------- | ------------------------------ | ---------------------- | +| Dense DiT | DiT | port | `fastvideo/models/dits/lingbot_video.py` | pass | pass | production exact pass | none | +| MoE DiT | DiT | port | `fastvideo/models/dits/lingbot_video.py` | pass | pass | full 48-block exact pass | none | +| MoE DiT refiner | DiT | port | `fastvideo/models/dits/lingbot_video.py` | pass | pass | full 48-block exact pass | none | +| Qwen3-VL text | encoder | reuse/port | `fastvideo/models/encoders/lingbot_video.py` | pass | pass | production exact pass | none | +| Qwen3-VL vision | encoder | deferred | `fastvideo/models/encoders/lingbot_video.py` | deferred | not_started | outside current T2V scope | TI2V phase | +| Wan VAE | VAE | reuse | `fastvideo/models/vaes/wanvae.py` | pass | passthrough | non-skip pass (`2/2`) | none | +| Flow-UniPC | scheduler | reuse | `fastvideo/models/schedulers/` | pass | not_needed | exact schedule/update pass | none | +| Base pipeline | pipeline | port | `fastvideo/pipelines/basic/lingbot_video/` | pass | pass | Dense exact; MoE smoke pass | none | +| Refiner pipeline | pipeline | port | `fastvideo/pipelines/basic/lingbot_video/` | pass | pass | FSDP plus SP=8 smoke pass | none | + +## Checkpoint Conversion + +The official checkpoints cannot be loaded directly by native FastVideo because +their folder layout and text-encoder format are different. The conversion +script creates FastVideo-compatible packages and does not modify the official +downloads. Local tests use the published packages below; developers do not need +to run conversion before testing. + +- Script: `scripts/checkpoint_conversion/lingbot_video_to_diffusers.py` + +| Variant | Official input repository | Published FastVideo repository | +| ------------- | ------------------------------------------ | ----------------------------------------------------------- | +| Dense 1.3B | `robbyant/lingbot-video-dense-1.3b` | `FastVideo/LingBot-Video-Dense-1.3B-Diffusers` | +| MoE + refiner | `robbyant/lingbot-video-moe-30b-a3b` | `FastVideo/LingBot-Video-MoE-30B-A3B-Diffusers` | + +What changes: + +- Text encoder conversion: `scripts/checkpoint_conversion/lingbot_video_to_diffusers.py` + converts the official Qwen3-VL checkpoint into the format loaded by FastVideo's + `LingBotVideoQwen3VLTextModel`. The official checkpoint is multimodal: it + contains both a vision encoder and a language model. The current port supports + text-to-video generation, so the converter keeps the language-model weights and + omits the unused vision-encoder weights. Within each language-model layer, + "projection weights" are the learned matrices used by the attention and MLP + computations. Qwen3-VL stores the attention matrices separately as `q_proj`, + `k_proj`, and `v_proj`; FastVideo reads the same values stacked into one + `qkv_proj` tensor. Qwen3-VL likewise stores `gate_proj` and `up_proj` separately; + FastVideo reads them stacked into one `gate_up_proj` tensor. This conversion + changes only the names and grouping of the retained weights. It does not + quantize, retrain, approximate, or otherwise change their numerical values. + +- The official `processor/` files are copied to `tokenizer/`. +- For MoE, the official `refiner/` is exposed as FastVideo's `transformer_2/`. +- A new `model_index.json` identifies the package as Dense or MoE/refiner. Both + packages use the shared `LingBotVideoPipeline` implementation. + +What does not change: + +- The Dense, MoE, and refiner transformer weights are reused without changing + their tensor values. +- The Wan VAE and scheduler are reused without weight conversion. +- Safetensors files are hard-linked where possible; configuration and tokenizer + files are copied. + +Related files: + +- Published package metadata: each `FastVideo/*-Diffusers` repository's `model_index.json` +- Dense key mappings: `checkpoints/lingbot-video/mapping/dense/` +- Conversion layout test: `tests/local_tests/lingbot_video/test_conversion_layout.py` +- Numerical loading and output checks: see **Validation Results** below. + +## Validation Results + +The results below separate three questions: + +1. Did the test run successfully? +2. Did FastVideo exactly match the original LingBot-Video implementation? +3. Which command and Slurm job provide the evidence? + +"Not tested" means that the check did not compare FastVideo with the original +implementation. It does not mean that the check failed. + +### Numerical parity with the original implementation + +These tests ran the same inputs through FastVideo and the original +LingBot-Video implementation. + +| What was tested | Successful? | Exact match with original? | Evidence | +| ----------------------------- | -------------- | --------------------------------------------------- | ----------------------------------------------------------------------------------------------- | +| Dense denoising model | Yes | Yes | `scripts/slurm/lingbot_video_dense_dit_parity.sbatch`, Slurm `1859061` | +| Dense VAE encode and decode | Yes: 2 tests | No exact claim; passed with `atol=rtol=0.05` | `scripts/slurm/lingbot_video_dense_vae_parity.sbatch`, Slurm `1859017` | +| Dense text encoder | Yes: 2 tests | Yes for encoder output; the other test checks setup | `scripts/slurm/lingbot_video_dense_text_encoder_parity.sbatch`, Slurm `1859116` | +| Scheduler | Yes: 2 tests | Yes | `tests/local_tests/schedulers/test_lingbot_video_scheduler_parity.py` | +| Dense pipeline, one step | Yes | Yes | `scripts/slurm/lingbot_video_dense_pipeline_parity.sbatch`, Slurm `1859069` | +| Dense pipeline, two steps | Yes | Yes | `scripts/slurm/lingbot_video_dense_pipeline_multistep_parity.sbatch`, Slurm `1859076` | +| Dense pipeline on two GPUs | Yes | Yes, using the same mathematical attention backend | `scripts/slurm/lingbot_video_dense_pipeline_sp_parity.sbatch`, Slurm `1859099` | +| One MoE transformer block | Yes | Yes | `scripts/slurm/lingbot_video_moe_block_parity.sbatch`, Slurm `1859113` | +| Complete MoE base transformer | Yes | Yes | `scripts/slurm/lingbot_video_moe_dit_parity.sbatch`, Slurm `1859114` | +| Complete MoE refiner | Yes | Yes | `scripts/slurm/lingbot_video_moe_refiner_dit_parity.sbatch`, Slurm `1859145` | + +### Functional tests without an original-output comparison + +These tests checked that FastVideo could load and run the port correctly. They +were not numerical comparisons with the original implementation. + +| What was tested | Successful? | Evidence | +| --------------------------------------- | ------------ | ------------------------------------------------------------------------------------------ | +| Dense pipeline on one GPU | Yes: 4 tests | `scripts/slurm/lingbot_video_dense_pipeline_smoke.sbatch`, Slurm `1859073` | +| Dense pipeline with sharded model state | Yes | `scripts/slurm/lingbot_video_dense_pipeline_fsdp_smoke.sbatch`, Slurm `1859075` | +| Mixed-precision sharded model loading | Yes: 2 tests | `scripts/slurm/lingbot_video_mixed_dtype_fsdp_test.sbatch`, Slurm `1859072` | +| Dense pipeline on two GPUs | Yes | `scripts/slurm/lingbot_video_dense_pipeline_sp_smoke.sbatch`, Slurm `1859089` | +| MoE base pipeline on one GPU | Yes | `scripts/slurm/lingbot_video_moe_pipeline_smoke.sbatch`, Slurm `1859121` | +| MoE base plus refiner on eight GPUs | Yes | `scripts/slurm/lingbot_video_refiner_pipeline_smoke.sbatch`, Slurm `1859122` | + +### Local API and CPU regression tests + +This was a selected local test suite, not an original-versus-FastVideo parity +test. It checked API parsing and presets, checkpoint conversion layout, +mixed-precision loading, pipeline setup, refiner stages, scheduler behavior, +and MoE routing. The result was `127 passed, 3 skipped`. + +The three skipped tests require allocated GPUs: Dense pipeline generation, MoE +pipeline generation, and mixed-precision sharded execution. Their corresponding +GPU runs are recorded in the functional-test table above. + +The exact local command was: + +```bash +python -m pytest -q \ + fastvideo/tests/api/test_compat_translation.py \ + fastvideo/tests/api/test_parser.py \ + fastvideo/tests/api/test_presets.py \ + fastvideo/tests/entrypoints/test_video_generator.py \ + tests/local_tests/lingbot_video/test_conversion_layout.py \ + tests/local_tests/models/test_fsdp_load_mixed_dtype.py \ + tests/local_tests/pipelines/test_lingbot_video_moe_pipeline_smoke.py \ + tests/local_tests/pipelines/test_lingbot_video_pipeline_smoke.py \ + tests/local_tests/pipelines/test_lingbot_video_refiner_stages.py \ + tests/local_tests/schedulers/test_lingbot_video_scheduler_parity.py \ + tests/local_tests/transformers/test_lingbot_video_moe.py +``` + +### Generation and performance runs + +These runs produced final videos or performance measurements. They are not +original-versus-FastVideo numerical parity tests. + +| What was run | Successful? | Evidence | +| ---------------------------- | ----------- | ----------------------------------------------------------------------------------------------- | +| Dense final video | Yes | `scripts/slurm/lingbot_video_dense_final_generation.sbatch`, Slurm `1859166` | +| MoE plus refiner final video | Yes | `scripts/slurm/lingbot_video_moe_refiner_final_generation.sbatch`, Slurm `1859167` | +| Dense performance benchmark | Yes | `scripts/slurm/lingbot_video_dense_benchmark.sbatch`, Slurm `1859117` | +| Dense performance profile | Yes | `scripts/slurm/lingbot_video_dense_nsys_profile.sbatch`, Slurm `1859106` | + +## Open Questions + +| ID | Question | Owner | Needed By Phase | Status | Resolution | +| ---- | ---------------------------------------------------------------------------------------- | ---------------- | --------------- | -------- | ----------------------------------------------------------- | +| Q001 | Do existing FastVideo Wan VAE and Flow-UniPC classes match exact official instantiation? | component owners | reuse gate | resolved | VAE encode/decode GPU parity and scheduler CPU parity pass. | + +## Issues And Blockers + +| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution | +| ---- | --------- | --------- | -------- | -------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------- | ---------- | -------- | ------------------------------------------------------------------------------------------------------------------------------ | +| I001 | prep | runtime | medium | Official recommends newer Torch/CUDA than shared venv. | All parity and production smoke jobs pass on Torch 2.11/cu128. | prep | resolved | No dependency changes were required. | +| I002 | parity | Dense DiT | medium | GPU numerical parity had not run. | Initial implementation had an exact 377-key/shape surface. | parity | resolved | Production-loaded output is exact in Slurm `1859061`. | +| I003 | parity | Dense DiT | medium | An all-true mask selected a different bf16 SDPA path. | Slurm `1859029` first drifted at block-1 attention. | components | resolved | Only build a mask for real padding; production parity is exact in Slurm `1859061`. | +| I004 | parity | Qwen3-VL | medium | Generic Qwen3 RoPE and causal-mask arithmetic differed from Qwen3-VL. | Slurm `1859030` first drifted at layer-0 `o_proj`. | components | resolved | LingBot specialization matches bf16 RoPE and causal GQA; Slurm `1859068` is exact. | +| I005 | pipeline | SDPA | medium | Official and native imports could select different cuDNN SDPA state. | The one-step full-pipeline comparison drifted before alignment. | pipeline | resolved | Match the backend state for both paths and restore it; Slurm `1859069` is exact. | +| I006 | inference | FSDP | medium | Mixed fp32/bf16 parameters cannot share one FSDP parameter group. | Nested child/root inference coverage passes in Slurm `1859072`. | loader | resolved | Preserve mixed dtypes and exclude fp32 parameters from inference FSDP; reject training use. | +| I007 | parity | SP SDPA | medium | Head sharding selected a different optimized bf16 Torch SDPA kernel. | Jobs `1859090`/`1859098`: drift `0.01412875`/`0.10981119`. | pipeline | resolved | Math SDPA on both paths is exactly zero in Slurm `1859099`, validating SP movement/masking. | +| I008 | benchmark | text | high | Temporary Qwen replacement retained 300 unreachable parameters totaling 7,424,306,688 bytes. | Slurm `1859109`: native steady allocation was 18,747,731,456 bytes, 6,451,255,808 above official. | encoder | resolved | Construct final layers directly; Slurm `1859117` native steady allocation is 11,319,211,520 bytes, 977,264,128 below official. | +| I009 | registry | discovery | low | One shared model-index identity made converted local MoE discovery ambiguous. | Pre-fix jobs `1859121`/`1859122` logged both registry IDs. | registry | resolved | Emit distinct Dense/MoE `_class_name` identities and map both to the shared native runtime. | +| I010 | refiner | workflow | low | Official CLI round-trips base pixels through MP4 before refinement. | `lingbot_video/runner.py` saves, reloads, then VAE-encodes. | pipeline | accepted | Native composition refines the decoded tensor in memory; document that pixels need not match. | + +> FastVideo disables PyTorch’s cuDNN attention backend during CUDA-platform initialization, while the original implementation may leave it enabled. The two implementations therefore used different bfloat16 attention kernels during the initial pipeline comparison, +> creating small rounding differences. The parity test now applies the same backend settings to both implementations and restores the previous process settings afterward. + +> The original implementation omits the attention mask when every token is valid. FastVideo initially passed an all-true mask, which is mathematically equivalent but caused PyTorch to choose a different optimized bfloat16 attention kernel. The different rounding broke +> exact parity. FastVideo now creates a mask only when the input contains real padding. + +## Escape Hatches + +An escape hatch is a documented exception used when the planned implementation +or validation cannot be completed as intended. Each entry records the blocker, +the proposed alternative or reduced scope, whether that alternative was +approved, and its consequences. An empty table means that this port did not use +any escape hatches. + +| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution | +| --- | ----- | ------------- | -------- | ------------------ | ------ | ---------- | + +## Decisions + +| Date | Decision | Rationale | Impact | +| ---------- | ----------------------------------------------------------- | -------------------------------------------------------------- | ------------------------------------------- | +| 2026-07-11 | Isolate this port from both prior worktrees. | User requirement. | Only current main and pinned official code. | +| 2026-07-11 | Port Dense completely before beginning MoE implementation. | User-requested order. | Dense parity is the MoE prerequisite. | +| 2026-07-11 | Exclude prompt rewriter and auto-negative inference. | User-selected scope. | Use fixed structured prompts/negatives. | +| 2026-07-11 | Keep every runtime write below the shared workspace root. | Home quota is small and can cause disk failures. | Launches source `launch_env.sh`. | +| 2026-07-11 | Reuse native Qwen3 for Dense/MoE text-only conditioning. | Pure-text MRoPE axes are identical; vision is not used by T2V. | TI2V vision remains separately gated. | +| 2026-07-11 | Use native grouped-mm for the released MoE expert layout. | It preserves the checkpoint surface and matches official CUDA. | Base and refiner DiTs are output-exact. | +| 2026-07-11 | Pass decoded base tensor directly to the refiner. | Original reloads an MP4; FastVideo keeps the tensor in memory. | No MP4 loss; not exact end-to-end parity. | +| 2026-07-11 | Give converted Dense and MoE layouts distinct identities. | Model discovery must not depend on an ambiguous shared name. | Both identities use one native runtime. | + +NOTE: Component-level numerical parity was achieved. Full MoE/refiner end-to-end parity was not tested or claimed because FastVideo intentionally omits the original MP4 round trip. + - Provenance: See fastvideo/pipelines/basic/lingbot_video/lingbot_video_pipeline.py and lingbot_video/runner.py for the difference. + +“Grouped-mm” means grouped matrix multiplication. +- Keep the checkpoint’s experts as separate expert weights. +- Preserve the original router’s expert selections. +- Use PyTorch’s built-in grouped-matrix operation to execute them. +- Do not merge the experts into a Dense model or implement a separate custom GPU kernel. +- Reference: fastvideo/models/dits/lingbot_video.py + +## Handoff Notes + +- Worktree base is `19a51a1fe630bcbeaf9fb6d864ad5ed3f31a3536` on branch `model/port-lingbot-video`. +- Official source imports successfully in the shared venv with no dependency changes. +- Dense official and FastVideo weights are pinned in `tests/local_tests/lingbot_video/hf_assets.py`. +- Dense DiT key manifests are under `checkpoints/lingbot-video/mapping/dense`; official and native surfaces both have 377 tensors with exact shapes. +- Dense Qwen3-VL text-only native config/class and parity test are present; conversion filters `model.language_model.*` and fuses QKV and gate/up through the existing Qwen3 loader. +- Scheduler, Dense DiT, VAE, text, and one/two-step sequential-CFG pipeline parity all pass without skips. +- Both official `robbyant` model repositories require conversion before native FastVideo use; the public `FastVideo` repositories provide that converted layout. +- Dense conversion is published at `FastVideo/LingBot-Video-Dense-1.3B-Diffusers`; its text surface has 290 fused tensors with exact native shapes, and its model index uses `_class_name: LingBotVideoDensePipeline`. +- Production inference preserves the Dense checkpoint's mixed fp32/bf16 parameter policy. Nested FSDP inference passes, while mixed-dtype FSDP training is intentionally rejected because the replicated fp32 parameters would not have synchronized gradients. +- Native batched CFG passes the non-FSDP and FSDP smoke tests. A full-size, 40-step MoE comparison also passes the decoded-frame semantic gate against the original packed FlashAttention-3 path; like-for-like sequential parity remains exact. +- Sequence-parallel validation is complete: SP=1 Dense DiT and full-pipeline regressions are exact in Slurm `1859083`/`1859084`; two-GPU B=2 unequal-mask odd-padding coverage passes in Slurm `1859085` and the official CP-order rerun `1859097`; production batched-CFG smoke passes in Slurm `1859089`. +- Official/native SP=2 parity is exactly zero with math SDPA at 192x320, 9 frames, and two steps in Slurm `1859099`. Normal optimized Torch SDPA is not exact: diagnostic Slurm jobs `1859090`/`1859098` produced mean/max latent drift of `0.01412875`/`0.10981119` because head sharding selected a different bf16 kernel. +- Canonical Dense final generation Slurm `1859166` produced a visually coherent 832x480, 121-frame, 24-FPS artifact in 330.31 seconds with a 26,738.52 MB rank-zero worker lifetime peak; its bytes exactly match Slurm `1859101`. +- The corrected Dense benchmark is under `validation/lingbot-video/benchmarks/job-1859117`: native/official mean latency is 0.214387/0.201173 seconds; native steady/peak allocation is 977,264,128/977,289,728 bytes lower. +- Nsight Systems artifacts are under `validation/lingbot-video/profiles/job-1859106`; launch and synchronization overhead dominate CUDA API time. Its pre-fix steady allocation is superseded by Slurm `1859117`. +- MoE conversion is published at `FastVideo/LingBot-Video-MoE-30B-A3B-Diffusers`; official base and refiner shards map to native `transformer/` and `transformer_2/` with exact 977-key surfaces, and its model index uses `_class_name: LingBotVideoMoePipeline`. +- Base and refiner 30.08B DiTs are exact across every block and final output in Slurm `1859114` and `1859145`. Production base-only and FSDP/SP=8 refiner smokes pass in Slurm `1859121` and `1859122`. +- Canonical Slurm `1859167` completed a 1920x1088, 121-frame, 24-FPS generation in 341.71 FastVideo seconds with a 63,450.65 MB rank-zero worker lifetime peak. Its bytes match superseded Slurm `1859144`. +- The final targeted CPU/API regression passed 127 tests with 3 expected GPU-gated skips; targeted Ruff and tracked/untracked whitespace validation pass. + +## Remaining T2V Gates + +- Attach reviewer-visible official and FastVideo comparison videos. +- T2I/TI2V and Qwen3-VL vision conditioning remain outside the completed T2V scope. diff --git a/tests/local_tests/lingbot_video/README.md b/tests/local_tests/lingbot_video/README.md new file mode 100644 index 000000000..8b3a7e2e4 --- /dev/null +++ b/tests/local_tests/lingbot_video/README.md @@ -0,0 +1,369 @@ +# LingBot-Video Local Tests + +This directory is the developer runbook for validating the LingBot-Video port in +FastVideo. The current port covers text-to-video (T2V) inference for both released +checkpoints: + +- Dense: a single Dense transformer generates the video. +- MoE/refiner: an MoE transformer generates a base video, then a second MoE + transformer refines it at a larger resolution. + +These are local tests, which means they are run manually in this worktree and are +not part of FastVideo's default CI suite. Prompt rewriting, automatic negative-prompt + generation, T2I, TI2V are outside the current port. + +## Current Testing Status + +Comparing official Lingbot-Video impl vs. FastVideo impl. + +| Area | Validation performed | Current result | +| ---------------------------- | ---------------------------------------------------------- | ------------------------------------------- | +| Qwen3-VL text-only encoder | Official versus FastVideo component comparison | Exact parity | +| Flow-UniPC scheduler | Timesteps, sigmas, and deterministic update comparison | Exact parity | +| Dense transformer | Full official versus FastVideo transformer comparison | Exact parity | +| Wan VAE | Official versus FastVideo encode and decode comparison | Passes with `atol=0.05`, `rtol=0.05` | +| Dense controlled pipeline | One- and two-step latent comparison | Exact parity | +| Dense sequence parallelism | Two-GPU pipeline latent comparison with math SDPA | Exact parity | +| MoE transformer block | Official versus FastVideo real-checkpoint block comparison | Exact parity | +| MoE base transformer | Full 48-block sequential comparison | Exact parity | +| MoE refiner transformer | Full 48-block sequential comparison | Exact parity | +| MoE production pipeline | Full-size, 40-step batched-CFG frame comparison | Mean SSIM exceeds 0.78 | +| MoE base pipeline parity | Full-size, 40-step sequential-CFG decoded-pixel comparison | Exact parity | +| MoE base pipeline | Batched CFG over five frames and two denoising steps | Pass | +| MoE plus refiner pipeline | Five-frame, two-step base-to-refiner handoff | Pass | +| Final generated MP4 | Production generation and decode | Generation pass; no exact end-to-end parity | +| T2I, TI2V, and vision branch | Not implemented by this T2V port | Outside current scope | + +The production batched classifier-free guidance (CFG) comparison uses the +original repository's packed FlashAttention-3 path and FastVideo's padded Torch +SDPA path. Because those kernels have different bfloat16 rounding, this test +compares decoded frames with a mean SSIM threshold of 0.78 instead of requiring +identical pixels. Like-for-like sequential component and pipeline parity remains +bit-exact. + +### One NOTE: +Parity is established by pairing the official and FastVideo implementations of +the text encoder, scheduler, Dense or MoE transformer, and VAE. A controlled +Dense pipeline test also compares the latent tensor after denoising. There is no +exact final-MP4 comparison for the complete MoE/refiner workflow because the two +pipelines hand the base video to the refiner differently: + +```text +Official repo: base decode -> save MP4 -> load MP4 -> resize and VAE-encode -> refine +FastVideo: base decode -> keep tensor in memory -> resize and VAE-encode -> refine +``` + +MP4 encoding is lossy, so the official workflow changes the pixels before the +refiner receives them. FastVideo intentionally avoids that loss. The component +parity tests therefore provide the meaningful numerical comparison, while the +complete MoE/refiner test is a production smoke test. + +## Setup: Required Workspace And Environment + +Complete the following steps before running the local tests. The example uses a +directory named `fv-hub`, but that directory can live anywhere: + +```bash +mkdir -p fv-hub +cd fv-hub +export FV_HUB="$PWD" +``` + +Run the remaining commands in the same shell so `FV_HUB` remains defined. Place +the FastVideo checkout being tested at +`$FV_HUB/fastvideo-port-lingbot-video`. + +### Step 1: Clone The Official Reference Implementation + +The parity tests import the official LingBot-Video implementation as the Python +package `lingbot_video`. Clone the official repository at the validated commit: + +```bash +cd "$FV_HUB/fastvideo-port-lingbot-video" +python3 \ + .agents/skills/add-model-01-prep/scripts/clone_reference_repo.py \ + https://github.com/robbyant/lingbot-video.git \ + "$FV_HUB/lingbot-video-reference" \ + --commit a638721cf2271804d02738b69f2ad788c4a559fc +``` + +### Step 2: Prepare The Virtual Environment + +The validated environment used Python 3.12, PyTorch 2.11.0+cu128, +Transformers 5.13.0, and Diffusers 0.39.0. Create the virtual environment, install +the normal FastVideo dependencies for the local platform, then register both +checkouts without changing dependency versions: + +```bash +python3.12 -m venv "$FV_HUB/.venv" +PY="$FV_HUB/.venv/bin/python" + +cd "$FV_HUB/fastvideo-port-lingbot-video" +$PY -m pip install --no-deps -e . +$PY -m pip install --no-deps -e "$FV_HUB/lingbot-video-reference" +``` + +The `--no-deps` commands do not populate an empty virtual environment; they only +register the two checkouts after the FastVideo dependencies are available. The +tests import the installed `lingbot_video` package; they do not modify +`sys.path`. Do not build FlashAttention as part of this setup. + +### Step 3: Configure The Hugging Face Cache + +The tests download their required model components from public Hugging Face +repositories after the corresponding opt-in test flag is set. No Hugging Face +token and no manually converted checkpoint directory are required. + +| Role | Repository | Pinned revision | +| --------------------- | ------------------------------------------------------------ | ------------------------------------------ | +| Official Dense | `robbyant/lingbot-video-dense-1.3b` | `f9789a7d9b4772a47aba62d4eb5282ddefd1da21` | +| FastVideo Dense | `FastVideo/LingBot-Video-Dense-1.3B-Diffusers` | `743ed04b96d77150d952eb08a59a56ee61b9bc95` | +| Official MoE/refiner | `robbyant/lingbot-video-moe-30b-a3b` | `f2e538f64afe00cc4ae674db2aeb52e2945edfd5` | +| FastVideo MoE/refiner | `FastVideo/LingBot-Video-MoE-30B-A3B-Diffusers` | `401dce84db5897cb950969e766410c8eadd4fbdf` | + +Keep the cache inside the shared `fv-hub` workspace by setting `HF_HOME` before +running pytest: + +```bash +export HF_HOME="$FV_HUB/.cache/huggingface" +mkdir -p "$HF_HOME" +``` + +The helper in `tests/local_tests/lingbot_video/hf_assets.py` pins the revisions +above and requests only the components used by each test. A complete run of both +MoE transformer variants and the refiner pipeline can retain roughly 260 GB of +official and FastVideo weights in the shared cache. + +### Step 4: Understand Checkpoint Conversion + +Conversion is already reflected in the two public `FastVideo/*-Diffusers` +repositories and is not a test setup step. Maintainers reproducing those +packages can use +`scripts/checkpoint_conversion/lingbot_video_to_diffusers.py`. The script keeps +the released transformer, VAE, and scheduler tensors; converts the text-only +Qwen3-VL weights to FastVideo's fused layout; and maps the official MoE +`refiner/` component to FastVideo's `transformer_2/` component. + +### Setup Result + +| Purpose | Location or identifier | +| ----------------------------- | -------------------------------------------------------------- | +| FastVideo checkout | `$FV_HUB/fastvideo-port-lingbot-video` | +| Official implementation | `$FV_HUB/lingbot-video-reference` at the pinned commit | +| Python executable | `$FV_HUB/.venv/bin/python` | +| Official Dense weights | `robbyant/lingbot-video-dense-1.3b` | +| Official MoE/refiner weights | `robbyant/lingbot-video-moe-30b-a3b` | +| FastVideo Dense weights | `FastVideo/LingBot-Video-Dense-1.3B-Diffusers` | +| FastVideo MoE/refiner weights | `FastVideo/LingBot-Video-MoE-30B-A3B-Diffusers` | + +### Compute Requirements For Tests + +Run GPU tests only inside an allocated GPU job. The activation flags documented +below are safety gates: without the required flag, pytest skips the expensive +test instead of consuming a GPU accidentally. + +| Test group | Required compute | +| ---------------------------------- | ------------------------------------------------------------------- | +| Routing, layout, and refiner logic | CPU | +| Scheduler parity | CPU; its small pinned scheduler component downloads automatically | +| Dense component and pipeline tests | 1 allocated CUDA GPU; prior acceptance runs used an H200 | +| Dense sequence-parallel test | 2 allocated CUDA GPUs | +| Full MoE DiT and base pipeline | 1 H200; official and FastVideo models are loaded sequentially | +| Complete MoE/refiner pipeline | 8 H200 GPUs by default; FSDP and sequence parallelism are exercised | + +Minimum GPU memory has not been established for the Dense tests. The H200 counts +above describe the configuration used for the retained acceptance results. + +--- + +## Test Catalog + +Tests are organized beside the FastVideo component they validate, so most files +are outside this `lingbot_video/` directory. + +### CPU And Preflight Tests + +| Test path | What it checks | +| --------------------------------------------------------------------- | ------------------------------------------------------------------- | +| `tests/local_tests/lingbot_video/test_conversion_layout.py` | Dense versus MoE converted directory and model-index layout | +| `tests/local_tests/transformers/test_lingbot_video_moe.py` | MoE routing, expert weighting, dtype policy, and checkpoint surface | +| `tests/local_tests/schedulers/test_lingbot_video_scheduler_parity.py` | Official and FastVideo scheduler timesteps, sigmas, and updates | +| `tests/local_tests/pipelines/test_lingbot_video_refiner_stages.py` | Refiner resize, VAE encode/noise preparation, schedule, and loading | +| `tests/local_tests/models/test_fsdp_load_mixed_dtype.py` | Checkpoint dtype restoration; its nested-FSDP case is GPU-gated | + +Run them from the FastVideo worktree: + +```bash +cd "$FV_HUB/fastvideo-port-lingbot-video" +PY="$FV_HUB/.venv/bin/python" + +$PY -m pytest -q \ + tests/local_tests/lingbot_video/test_conversion_layout.py \ + tests/local_tests/transformers/test_lingbot_video_moe.py \ + tests/local_tests/schedulers/test_lingbot_video_scheduler_parity.py \ + tests/local_tests/pipelines/test_lingbot_video_refiner_stages.py \ + tests/local_tests/models/test_fsdp_load_mixed_dtype.py +``` + +### Dense GPU Tests + +| Test path | What it checks | +| ---------------------------------------------------------------------------- | ---------------------------------------------------------- | +| `tests/local_tests/encoders/test_lingbot_video_dense_text_encoder_parity.py` | Text-only Qwen3-VL construction and official output parity | +| `tests/local_tests/transformers/test_lingbot_video_dense_dit_parity.py` | Full Dense transformer output parity | +| `tests/local_tests/vaes/test_lingbot_video_vae_parity.py` | Wan VAE encode and decode parity within tolerance | +| `tests/local_tests/pipelines/test_lingbot_video_pipeline_parity.py` | Controlled official and FastVideo denoised-latent parity | +| `tests/local_tests/pipelines/test_lingbot_video_pipeline_smoke.py` | Registry, input contracts, model loading, and generation | +| `tests/local_tests/models/test_fsdp_load_mixed_dtype.py` | Mixed-bf16/fp32 checkpoint loading under FSDP | + +Run the files individually so a large model is released before the next test: + +```bash +cd "$FV_HUB/fastvideo-port-lingbot-video" +PY="$FV_HUB/.venv/bin/python" + +LINGBOT_VIDEO_RUN_GPU_TESTS=1 $PY -m pytest -v -s \ + tests/local_tests/encoders/test_lingbot_video_dense_text_encoder_parity.py +LINGBOT_VIDEO_RUN_GPU_TESTS=1 $PY -m pytest -v -s \ + tests/local_tests/transformers/test_lingbot_video_dense_dit_parity.py +LINGBOT_VIDEO_RUN_GPU_TESTS=1 $PY -m pytest -v -s \ + tests/local_tests/vaes/test_lingbot_video_vae_parity.py +LINGBOT_VIDEO_RUN_GPU_TESTS=1 $PY -m pytest -v -s \ + tests/local_tests/pipelines/test_lingbot_video_pipeline_parity.py +LINGBOT_VIDEO_RUN_GPU_TESTS=1 $PY -m pytest -v -s \ + tests/local_tests/pipelines/test_lingbot_video_pipeline_smoke.py +LINGBOT_VIDEO_RUN_GPU_TESTS=1 $PY -m pytest -v -s \ + tests/local_tests/models/test_fsdp_load_mixed_dtype.py +``` + +The pipeline parity defaults to a small one-step fixture. The following optional +variables select its shape and distributed mode: + +```text +LINGBOT_VIDEO_PARITY_HEIGHT +LINGBOT_VIDEO_PARITY_WIDTH +LINGBOT_VIDEO_PARITY_NUM_FRAMES +LINGBOT_VIDEO_PARITY_NUM_INFERENCE_STEPS +LINGBOT_VIDEO_PARITY_NUM_GPUS +LINGBOT_VIDEO_PARITY_SP_SIZE +LINGBOT_VIDEO_PARITY_USE_FSDP=1 +LINGBOT_VIDEO_PARITY_FORCE_MATH_SDPA=1 +LINGBOT_VIDEO_PARITY_VARIANT=dense|moe +LINGBOT_VIDEO_PARITY_BATCH_CFG=1 +LINGBOT_VIDEO_PARITY_OUTPUT_TYPE=latent|np +LINGBOT_VIDEO_PARITY_DETERMINISTIC=1 +``` + +`LINGBOT_VIDEO_PARITY_FORCE_MATH_SDPA=1` is required for an exact two-GPU +sequence-parallel comparison. With optimized SDPA, head sharding can select a +different bf16 attention kernel; that kernel-choice difference creates small +numerical drift even when the sequence-parallel data movement is correct. + +Run the full-size MoE batched-CFG semantic gate with: + +```bash +LINGBOT_VIDEO_RUN_GPU_TESTS=1 \ +LINGBOT_VIDEO_PARITY_VARIANT=moe \ +LINGBOT_VIDEO_PARITY_BATCH_CFG=1 \ +LINGBOT_VIDEO_PARITY_HEIGHT=480 \ +LINGBOT_VIDEO_PARITY_WIDTH=832 \ +LINGBOT_VIDEO_PARITY_NUM_FRAMES=121 \ +LINGBOT_VIDEO_PARITY_NUM_INFERENCE_STEPS=40 \ +LINGBOT_VIDEO_PARITY_OUTPUT_TYPE=np \ +LINGBOT_VIDEO_PARITY_DETERMINISTIC=1 \ + $PY -m pytest -v -s \ + tests/local_tests/pipelines/test_lingbot_video_pipeline_parity.py +``` + +The original repository implements batched CFG with the optional +`flash_attn_interface.flash_attn_varlen_func` kernel. When that kernel is +available, both sides run batched CFG and the test compares decoded frames with +SSIM. When it is unavailable, the test prints the reason and uses the original +repository's sequential-CFG result as the numerical oracle while FastVideo +still runs batched CFG; that dependency-free fallback requires +`LINGBOT_VIDEO_PARITY_FORCE_MATH_SDPA=1` and uses tight pixel-drift thresholds. +The fallback exists only in the parity test and does not change either +production pipeline. + +For a bit-exact, full-size MoE base comparison, keep both sides on sequential +CFG: + +```bash +LINGBOT_VIDEO_RUN_GPU_TESTS=1 \ +LINGBOT_VIDEO_PARITY_VARIANT=moe \ +LINGBOT_VIDEO_PARITY_HEIGHT=480 \ +LINGBOT_VIDEO_PARITY_WIDTH=832 \ +LINGBOT_VIDEO_PARITY_NUM_FRAMES=121 \ +LINGBOT_VIDEO_PARITY_NUM_INFERENCE_STEPS=40 \ +LINGBOT_VIDEO_PARITY_OUTPUT_TYPE=np \ +LINGBOT_VIDEO_PARITY_DETERMINISTIC=1 \ + $PY -m pytest -v -s \ + tests/local_tests/pipelines/test_lingbot_video_pipeline_parity.py +``` + +### MoE And Refiner GPU Tests + +| Test path | What it checks | +| -------------------------------------------------------------------------- | ---------------------------------------------------------- | +| `tests/local_tests/transformers/test_lingbot_video_moe_block_parity.py` | One real-checkpoint MoE block against the official block | +| `tests/local_tests/transformers/test_lingbot_video_moe_dit_parity.py` | Full base or refiner 48-block transformer parity | +| `tests/local_tests/pipelines/test_lingbot_video_moe_pipeline_smoke.py` | Five-frame, two-step batched-CFG base-MoE generation | +| `tests/local_tests/pipelines/test_lingbot_video_refiner_pipeline_smoke.py` | Five-frame, two-step base-to-refiner in-memory workflow | + +Run each expensive test separately: + +```bash +cd "$FV_HUB/fastvideo-port-lingbot-video" +PY="$FV_HUB/.venv/bin/python" + +LINGBOT_VIDEO_RUN_GPU_TESTS=1 $PY -m pytest -v -s \ + tests/local_tests/transformers/test_lingbot_video_moe_block_parity.py + +LINGBOT_VIDEO_RUN_FULL_MOE_TESTS=1 \ +LINGBOT_VIDEO_MOE_PARITY_VARIANT=base \ + $PY -m pytest -v -s \ + tests/local_tests/transformers/test_lingbot_video_moe_dit_parity.py + +LINGBOT_VIDEO_RUN_FULL_MOE_TESTS=1 \ +LINGBOT_VIDEO_MOE_PARITY_VARIANT=refiner \ + $PY -m pytest -v -s \ + tests/local_tests/transformers/test_lingbot_video_moe_dit_parity.py + +LINGBOT_VIDEO_RUN_MOE_PIPELINE_TESTS=1 $PY -m pytest -v -s \ + tests/local_tests/pipelines/test_lingbot_video_moe_pipeline_smoke.py + +LINGBOT_VIDEO_RUN_REFINER_PIPELINE_TESTS=1 \ +LINGBOT_VIDEO_REFINER_NUM_GPUS=8 \ + $PY -m pytest -v -s \ + tests/local_tests/pipelines/test_lingbot_video_refiner_pipeline_smoke.py +``` + +### Two-GPU Sequence-Parallel Test + +`fastvideo/tests/distributed/test_sp_lingbot_video.py` starts its own two-process +`torchrun` worker and compares sequence-parallel output with the single-rank +reference. Run pytest once from a two-GPU allocation; do not wrap this command in +another `torchrun`: + +```bash +cd "$FV_HUB/fastvideo-port-lingbot-video" +"$FV_HUB/.venv/bin/python" -m pytest -v -s \ + fastvideo/tests/distributed/test_sp_lingbot_video.py +``` + +## Interpreting Skips And Failures + +- A skipped GPU test means no validation occurred. Re-run it in the required GPU + allocation with the documented activation flag. +- An exact-parity failure means the paired official and FastVideo tensors differ; + inspect the first reported mismatch rather than treating successful generation + as a substitute. +- A VAE failure means the error exceeded the explicit `0.05` absolute or relative + tolerance used by the VAE tests. +- A smoke-test failure means FastVideo could not complete that production path; it + does not by itself identify which component lost numerical parity. +- Missing files under the required workspace paths are setup failures, not model + correctness failures. + +This README intentionally uses direct pytest commands. There are no retained +LingBot-Video `launch_env.sh` or Slurm wrapper scripts in this worktree. diff --git a/tests/local_tests/lingbot_video/hf_assets.py b/tests/local_tests/lingbot_video/hf_assets.py new file mode 100644 index 000000000..fa499e520 --- /dev/null +++ b/tests/local_tests/lingbot_video/hf_assets.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pinned public Hugging Face assets used by LingBot-Video local tests.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +from huggingface_hub import snapshot_download + + +@dataclass(frozen=True, slots=True) +class HubModel: + """Identify one immutable public model snapshot.""" + + repo_id: str + revision: str + + +OFFICIAL_DENSE = HubModel( + "robbyant/lingbot-video-dense-1.3b", + "f9789a7d9b4772a47aba62d4eb5282ddefd1da21", +) +OFFICIAL_MOE = HubModel( + "robbyant/lingbot-video-moe-30b-a3b", + "f2e538f64afe00cc4ae674db2aeb52e2945edfd5", +) +FASTVIDEO_DENSE = HubModel( + "FastVideo/LingBot-Video-Dense-1.3B-Diffusers", + "743ed04b96d77150d952eb08a59a56ee61b9bc95", +) +FASTVIDEO_MOE = HubModel( + "FastVideo/LingBot-Video-MoE-30B-A3B-Diffusers", + "401dce84db5897cb950969e766410c8eadd4fbdf", +) + + +def download_patterns(model: HubModel, *patterns: str) -> Path: + """Download selected files from one pinned snapshot and return its root.""" + if not patterns: + raise ValueError("at least one Hugging Face allow-pattern is required") + return Path(snapshot_download( + repo_id=model.repo_id, + revision=model.revision, + allow_patterns=list(patterns), + )) + + +def download_components(model: HubModel, *components: str) -> Path: + """Download root metadata and the requested Diffusers component folders.""" + return download_patterns( + model, + "model_index.json", + *(f"{component}/**" for component in components), + ) + + +def materialize_component_view( + model: HubModel, + destination: Path, + *components: str, +) -> Path: + """Create a temporary valid model layout containing selected components.""" + snapshot = download_components(model, *components) + selected = set(components) + model_index = json.loads((snapshot / "model_index.json").read_text()) + filtered_index = { + key: value + for key, value in model_index.items() + if not (isinstance(value, list) and value and value[0] is not None and key not in selected) + } + destination.mkdir(parents=True) + (destination / "model_index.json").write_text( + json.dumps(filtered_index, indent=2) + "\n", + encoding="utf-8", + ) + for component in components: + (destination / component).symlink_to(snapshot / component, target_is_directory=True) + return destination diff --git a/tests/local_tests/lingbot_video/t2v_example_1_prompt.json b/tests/local_tests/lingbot_video/t2v_example_1_prompt.json new file mode 100644 index 000000000..92afdf69b --- /dev/null +++ b/tests/local_tests/lingbot_video/t2v_example_1_prompt.json @@ -0,0 +1,77 @@ +{ + "caption": { + "comprehensive_description": { + "scene_content_description": "A young woman with long, wavy brown hair is standing in a bright, modern apartment living room. She is wearing a stylish, oversized cream-colored knit cardigan over a white tank top, paired with high-waisted, wide-leg beige trousers. She holds a small, structured tan leather handbag in her left hand. The background features a neutral-toned interior with a beige sofa, a potted plant, and large windows that let in soft, natural light, creating a warm and inviting atmosphere. The woman is smiling and looking directly at the camera, showcasing her outfit with a confident and friendly demeanor.", + "camera_movement_description": "The camera is positioned at eye level and remains essentially stationary throughout the video, maintaining a medium shot that captures the subject from the waist up. There is a very shallow depth of field, keeping the woman in sharp focus while the background remains softly blurred." + }, + "camera_info": { + "color": "Warm", + "frame_size": "Medium", + "shot_type_angle": "Eye level", + "lens_size": "Medium", + "composition": "Center", + "lighting": "Soft light", + "lighting_type": "Daylight" + }, + "world_knowledge": [], + "prominent_elements": [ + { + "name": "young woman", + "description": "A woman with long, wavy brown hair and a friendly expression, modeling a fashion outfit.", + "actions": [ + { + "timestamp": "[0.0s - 0.5s]", + "action": "stands still, smiling at the camera" + }, + { + "timestamp": "[0.5s - 2.0s]", + "action": "shifts her weight and turns her body slightly to the right" + }, + { + "timestamp": "[2.0s - 3.5s]", + "action": "adjusts the collar of her cardigan with her right hand" + }, + { + "timestamp": "[3.5s - 5.0s]", + "action": "returns to a neutral pose, smiling at the camera" + } + ], + "location": "center of the frame", + "relative_size": "dominant", + "shape_and_color": "slender build; wearing cream, white, and beige", + "texture": "soft knit cardigan, smooth fabric trousers", + "appearance_details": "long wavy brown hair, gold hoop earrings, tan leather handbag", + "relationship": "the main subject of the video, standing in front of a blurred apartment background", + "orientation": "upright, facing the camera", + "pose": "standing, shifting weight, and adjusting clothing", + "expression": "smiling and confident", + "clothing": "oversized cream knit cardigan, white tank top, high-waisted wide-leg beige trousers", + "gender": "female", + "skin_tone_and_texture": "fair skin with a smooth texture" + }, + { + "name": "tan handbag", + "description": "A small, structured leather handbag with a top handle.", + "actions": [ + { + "timestamp": "[0.0s - 5.0s]", + "action": "held steady in the woman's left hand" + } + ], + "location": "held in the woman's left hand, lower center of the frame", + "relative_size": "small", + "shape_and_color": "rectangular, tan or light brown", + "texture": "smooth leather", + "appearance_details": "structured shape with a top handle", + "relationship": "held by the woman as an accessory", + "orientation": "upright", + "pose": "", + "expression": "", + "clothing": "", + "gender": "", + "skin_tone_and_texture": "" + } + ] + }, + "duration": 5 +} diff --git a/tests/local_tests/lingbot_video/test_conversion_layout.py b/tests/local_tests/lingbot_video/test_conversion_layout.py new file mode 100644 index 000000000..a48bdf42d --- /dev/null +++ b/tests/local_tests/lingbot_video/test_conversion_layout.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only layout tests for the LingBot-Video checkpoint converter.""" + +import json +import os +from pathlib import Path + +import pytest + +from scripts.checkpoint_conversion import lingbot_video_to_diffusers as converter + + +def _make_source(tmp_path: Path, has_refiner: bool) -> Path: + """Create the minimal official component tree exercised by ``convert``.""" + source = tmp_path / "source" + for index, component_name in enumerate(converter.PASSTHROUGH_COMPONENTS): + component = source / component_name + component.mkdir(parents=True) + (component / "config.json").write_text("{}\n") + (component / "weights.safetensors").write_bytes(bytes([index])) + for component_name in ("text_encoder", "processor"): + component = source / component_name + component.mkdir(parents=True) + (component / "config.json").write_text("{}\n") + if has_refiner: + refiner = source / "refiner" + refiner.mkdir() + (refiner / "config.json").write_text("{}\n") + (refiner / "weights.safetensors").write_bytes(b"refiner") + return source + + +def _fake_convert_text_encoder(_source: Path, destination: Path) -> dict[str, str]: + """Create a lightweight destination in place of tensor conversion.""" + destination.mkdir() + return {"weight": "model.safetensors"} + + +def _fake_write_text_encoder_config(_source: Path, _destination: Path) -> dict: + """Return the metadata consumed by the converter's completion message.""" + return {"architectures": ["LingBotVideoQwen3VLTextModel"]} + + +@pytest.mark.parametrize("has_refiner", [False, True]) +def test_convert_emits_optional_transformer_2_layout( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + has_refiner: bool, +) -> None: + """Keep Dense output unchanged and map official refiners to ``transformer_2``.""" + source = _make_source(tmp_path, has_refiner) + destination = tmp_path / "converted" + validated_configs: list[Path] = [] + monkeypatch.setattr(converter, "_convert_text_encoder", _fake_convert_text_encoder) + monkeypatch.setattr(converter, "_write_text_encoder_config", _fake_write_text_encoder_config) + monkeypatch.setattr(converter, "_validate_transformer_config", validated_configs.append) + + converter.convert(source, destination) + + model_index = json.loads((destination / "model_index.json").read_text()) + expected_pipeline_class = converter.MOE_PIPELINE_CLASS if has_refiner else converter.DENSE_PIPELINE_CLASS + assert model_index["_class_name"] == expected_pipeline_class + assert ("transformer_2" in model_index) is has_refiner + assert (destination / "transformer_2").is_dir() is has_refiner + assert os.path.samefile( + source / "transformer" / "weights.safetensors", + destination / "transformer" / "weights.safetensors", + ) + assert len(validated_configs) == (2 if has_refiner else 1) + if has_refiner: + assert model_index["transformer_2"] == ["diffusers", "LingBotVideoTransformer3DModel"] + assert os.path.samefile( + source / "refiner" / "weights.safetensors", + destination / "transformer_2" / "weights.safetensors", + ) + assert not os.path.samefile( + destination / "transformer" / "weights.safetensors", + destination / "transformer_2" / "weights.safetensors", + ) diff --git a/tests/local_tests/models/test_fsdp_load_mixed_dtype.py b/tests/local_tests/models/test_fsdp_load_mixed_dtype.py new file mode 100644 index 000000000..9ef5db503 --- /dev/null +++ b/tests/local_tests/models/test_fsdp_load_mixed_dtype.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Mixed-precision parameter loading coverage for native model dtype policies.""" + +import os + +import pytest +import torch +from torch.distributed._tensor import DTensor +from torch.distributed.fsdp import MixedPrecisionPolicy + +from fastvideo.distributed.parallel_state import init_distributed_environment +from fastvideo.models.dits.lingbot_video import LingBotVideoRouter +from fastvideo.models.loader.fsdp_load import ( + load_model_from_full_model_state_dict, + maybe_load_fsdp_model, + shard_model, +) + + +class _MixedDtypeModel(torch.nn.Module): + """Tiny model that keeps one checkpoint parameter in fp32.""" + + def __init__(self) -> None: + super().__init__() + self.bulk = torch.nn.Parameter(torch.zeros(2, dtype=torch.bfloat16)) + self.sensitive = torch.nn.Parameter(torch.zeros(2, dtype=torch.bfloat16)) + + def _get_parameter_dtype(self, name: str, default_dtype: torch.dtype) -> torch.dtype: + """Keep the sensitive test tensor in fp32.""" + return torch.float32 if name == "sensitive" else default_dtype + + +class _MixedDtypeBlock(torch.nn.Module): + """Small FSDP child with one managed and one replicated parameter.""" + + def __init__(self) -> None: + super().__init__() + self.bulk = torch.nn.Parameter(torch.eye(2, dtype=torch.bfloat16)) + self.sensitive = torch.nn.Parameter(torch.ones(2, dtype=torch.float32)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Apply both parameters so FSDP must initialize the mixed child.""" + return torch.nn.functional.linear(hidden_states, self.bulk) + self.sensitive.to(hidden_states.dtype) + + +class _NestedMixedDtypeModel(torch.nn.Module): + """Nested model exercising both child and root ignored-parameter wiring.""" + + def __init__(self) -> None: + super().__init__() + self.blocks = torch.nn.ModuleList([_MixedDtypeBlock()]) + self.root_bulk = torch.nn.Parameter(torch.ones(2, dtype=torch.bfloat16)) + self.root_sensitive = torch.nn.Parameter(torch.ones(2, dtype=torch.float32)) + + def _get_parameter_dtype(self, name: str, default_dtype: torch.dtype) -> torch.dtype: + """Keep every sensitive test parameter replicated in fp32.""" + return torch.float32 if "sensitive" in name else default_dtype + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Run the sharded child and consume both root parameters.""" + hidden_states = self.blocks[0](hidden_states) + return hidden_states * self.root_bulk + self.root_sensitive.to(hidden_states.dtype) + + +class _RouterBufferModel(torch.nn.Module): + """Wrap the released router to exercise its persistent correction buffer.""" + + def __init__(self) -> None: + super().__init__() + self.router = LingBotVideoRouter(2, 3, 1, "sigmoid", True, None, None, 1.0) + + def _get_parameter_dtype(self, name: str, default_dtype: torch.dtype) -> torch.dtype: + """Keep the released router state in fp32.""" + return torch.float32 if "router" in name else default_dtype + + +def test_full_state_dict_loader_honors_model_parameter_dtypes() -> None: + """Load exact fp32 values for selected tensors while casting ordinary weights.""" + model = _MixedDtypeModel() + bulk = torch.tensor([1.001, -2.003], dtype=torch.float32) + sensitive = torch.tensor([3.001, -4.003], dtype=torch.float32) + incompatible = load_model_from_full_model_state_dict( + model, + iter((("bulk", bulk), ("sensitive", sensitive))), + device=torch.device("cpu"), + param_dtype=torch.bfloat16, + strict=True, + param_names_mapping=lambda name: (name, None, None), + ) + assert incompatible.missing_keys == [] + assert incompatible.unexpected_keys == [] + assert model.bulk.dtype == torch.bfloat16 + assert model.sensitive.dtype == torch.float32 + torch.testing.assert_close(model.bulk, bulk.to(torch.bfloat16)) + torch.testing.assert_close(model.sensitive, sensitive) + + +def test_full_state_dict_loader_preserves_router_bias_buffer() -> None: + """Keep the MoE correction bias registered as a non-trainable buffer.""" + model = _RouterBufferModel() + weight = torch.arange(6, dtype=torch.float32).reshape(3, 2) + bias = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32) + + incompatible = load_model_from_full_model_state_dict( + model, + iter((("router.weight", weight), ("router.e_score_correction_bias", bias))), + device=torch.device("cpu"), + param_dtype=torch.bfloat16, + strict=True, + param_names_mapping=lambda name: (name, None, None), + ) + + assert incompatible.missing_keys == [] + assert incompatible.unexpected_keys == [] + assert "router.e_score_correction_bias" in dict(model.named_buffers()) + assert "router.e_score_correction_bias" not in dict(model.named_parameters()) + torch.testing.assert_close(model.router.e_score_correction_bias, bias) + + +def test_training_rejection_uses_fsdp_parameter_dtype() -> None: + """Reject replicated fp32 training state when construction defaults to fp32.""" + with pytest.raises(NotImplementedError, match="separate gradient synchronization"): + maybe_load_fsdp_model( + model_cls=_MixedDtypeModel, + init_params={}, + weight_dir_list=[], + device=torch.device("cpu"), + hsdp_replicate_dim=1, + hsdp_shard_dim=1, + default_dtype=torch.float32, + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + training_mode=True, + pin_cpu_memory=False, + ) + + +def test_nested_fsdp_ignores_selected_fp32_parameters() -> None: + """Run a CUDA forward with bf16 DTensors and replicated fp32 parameters.""" + if os.environ.get("LINGBOT_VIDEO_RUN_GPU_TESTS") != "1": + pytest.skip("set LINGBOT_VIDEO_RUN_GPU_TESTS=1 on an allocated GPU") + if not torch.cuda.is_available(): + raise RuntimeError("mixed-dtype FSDP coverage requires CUDA") + if not torch.distributed.is_initialized(): + init_distributed_environment(world_size=1, rank=0, local_rank=0) + model = _NestedMixedDtypeModel().cuda() + mesh = torch.distributed.init_device_mesh( + "cuda", + mesh_shape=(1, 1), + mesh_dim_names=("replicate", "shard"), + ) + shard_model( + model, + cpu_offload=False, + mp_policy=MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=None, + cast_forward_inputs=False, + ), + mesh=mesh, + fsdp_shard_conditions=[lambda name, module: isinstance(module, _MixedDtypeBlock)], + pin_cpu_memory=False, + ) + output = model(torch.ones(1, 2, device="cuda", dtype=torch.bfloat16)) + assert torch.isfinite(output).all() + assert isinstance(model.blocks[0].bulk, DTensor) + assert not isinstance(model.blocks[0].sensitive, DTensor) + assert model.blocks[0].sensitive.dtype == torch.float32 + assert isinstance(model.root_bulk, DTensor) + assert not isinstance(model.root_sensitive, DTensor) + assert model.root_sensitive.dtype == torch.float32 diff --git a/tests/local_tests/pipelines/test_lingbot_video_moe_pipeline_smoke.py b/tests/local_tests/pipelines/test_lingbot_video_moe_pipeline_smoke.py new file mode 100644 index 000000000..3ebf0790b --- /dev/null +++ b/tests/local_tests/pipelines/test_lingbot_video_moe_pipeline_smoke.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Production-loader smoke coverage for the base-only LingBot-Video MoE pipeline.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, cast + +import pytest +import torch + +from tests.local_tests.lingbot_video.hf_assets import FASTVIDEO_MOE, materialize_component_view + + +def test_lingbot_video_moe_base_pipeline_smoke(tmp_path: Path) -> None: + """Run two batched-CFG steps over two temporal latent frames with the 30B MoE.""" + if os.environ.get("LINGBOT_VIDEO_RUN_MOE_PIPELINE_TESTS") != "1": + pytest.skip("Set LINGBOT_VIDEO_RUN_MOE_PIPELINE_TESTS=1 on a scheduled H200.") + if not torch.cuda.is_available(): + pytest.skip("LingBot-Video MoE pipeline smoke requires CUDA.") + from fastvideo import VideoGenerator + + model_dir = materialize_component_view( + FASTVIDEO_MOE, + tmp_path / "base_model", + "scheduler", + "text_encoder", + "tokenizer", + "transformer", + "vae", + ) + generator = VideoGenerator.from_pretrained( + str(model_dir), + num_gpus=1, + sp_size=1, + use_fsdp_inference=False, + refine_enabled=False, + dit_cpu_offload=False, + dit_layerwise_offload=False, + vae_cpu_offload=True, + text_encoder_cpu_offload=True, + pin_cpu_memory=False, + output_type="latent", + ) + try: + result = generator.generate_video( + prompt="A red fox runs through fresh snow at sunrise.", + output_path=str(tmp_path), + save_video=False, + return_frames=True, + height=32, + width=32, + num_frames=5, + num_inference_steps=2, + guidance_scale=3.0, + batch_cfg=True, + seed=42, + ) + finally: + generator.shutdown() + samples = cast(dict[str, Any], result)["samples"] + assert torch.is_tensor(samples) + assert tuple(samples.shape) == (1, 16, 2, 4, 4) + assert torch.isfinite(samples).all() diff --git a/tests/local_tests/pipelines/test_lingbot_video_pipeline_parity.py b/tests/local_tests/pipelines/test_lingbot_video_pipeline_parity.py new file mode 100644 index 000000000..c5c50cee2 --- /dev/null +++ b/tests/local_tests/pipelines/test_lingbot_video_pipeline_parity.py @@ -0,0 +1,307 @@ +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end latent or decoded-pixel parity for LingBot-Video T2V.""" + +from __future__ import annotations + +import gc +import json +import os +from pathlib import Path +from typing import Any, cast + +import pytest +import torch + +from tests.local_tests.lingbot_video.hf_assets import ( + FASTVIDEO_DENSE, + FASTVIDEO_MOE, + OFFICIAL_DENSE, + OFFICIAL_MOE, + download_components, + materialize_component_view, +) + +PROMPT_PATH = Path(__file__).parents[1] / "lingbot_video" / "t2v_example_1_prompt.json" +HEIGHT = int(os.environ.get("LINGBOT_VIDEO_PARITY_HEIGHT", "32")) +WIDTH = int(os.environ.get("LINGBOT_VIDEO_PARITY_WIDTH", "32")) +NUM_FRAMES = int(os.environ.get("LINGBOT_VIDEO_PARITY_NUM_FRAMES", "1")) +NUM_INFERENCE_STEPS = int(os.environ.get("LINGBOT_VIDEO_PARITY_NUM_INFERENCE_STEPS", "1")) +NUM_GPUS = int(os.environ.get("LINGBOT_VIDEO_PARITY_NUM_GPUS", "1")) +SP_SIZE = int(os.environ.get("LINGBOT_VIDEO_PARITY_SP_SIZE", "1")) +USE_FSDP = os.environ.get("LINGBOT_VIDEO_PARITY_USE_FSDP") == "1" +FORCE_MATH_SDPA = os.environ.get("LINGBOT_VIDEO_PARITY_FORCE_MATH_SDPA") == "1" +PARITY_VARIANT = os.environ.get("LINGBOT_VIDEO_PARITY_VARIANT", "dense") +BATCH_CFG = os.environ.get("LINGBOT_VIDEO_PARITY_BATCH_CFG") == "1" +PARITY_OUTPUT_TYPE = os.environ.get("LINGBOT_VIDEO_PARITY_OUTPUT_TYPE", "latent") +DETERMINISTIC = os.environ.get("LINGBOT_VIDEO_PARITY_DETERMINISTIC") == "1" +BATCH_CFG_MIN_SSIM = 0.78 +BATCH_CFG_PIXEL_MAX_ABS = 3e-3 +BATCH_CFG_PIXEL_MEAN_ABS = 6e-4 + + +def _load_structured_prompt() -> str: + """Load the released LingBot T2V example in the format expected by Qwen3-VL.""" + sample = json.loads(PROMPT_PATH.read_text(encoding="utf-8")) + return str(sample["caption"]["comprehensive_description"]) + + +PROMPT = _load_structured_prompt() + + +def _as_channel_first_tensor(value: Any) -> torch.Tensor: + """Normalize latent or decoded pipeline output to a CPU tensor.""" + tensor = value if torch.is_tensor(value) else torch.as_tensor(value) + if tensor.ndim == 5 and tensor.shape[-1] == 3: + tensor = tensor.permute(0, 4, 1, 2, 3) + return tensor.detach().float().cpu() + + +def _configure_worker_backends(_worker: Any) -> dict[str, bool]: + """Mirror the parent process's deterministic and SDPA settings in a worker.""" + torch.backends.cuda.enable_cudnn_sdp(False) + if DETERMINISTIC: + torch.use_deterministic_algorithms(True) + torch.backends.cudnn.enabled = False + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + if FORCE_MATH_SDPA: + torch.backends.cuda.enable_flash_sdp(False) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_math_sdp(True) + return { + "deterministic": torch.are_deterministic_algorithms_enabled(), + "cudnn": torch.backends.cudnn.enabled, + "matmul_tf32": torch.backends.cuda.matmul.allow_tf32, + "cudnn_tf32": torch.backends.cudnn.allow_tf32, + "cudnn_sdp": torch.backends.cuda.cudnn_sdp_enabled(), + "flash": torch.backends.cuda.flash_sdp_enabled(), + "mem_efficient": torch.backends.cuda.mem_efficient_sdp_enabled(), + "math": torch.backends.cuda.math_sdp_enabled(), + } + + +def _require_gpu_test() -> None: + """Require an explicitly allocated CUDA worker for heavyweight pipeline parity.""" + if os.environ.get("LINGBOT_VIDEO_RUN_GPU_TESTS") != "1": + pytest.skip("set LINGBOT_VIDEO_RUN_GPU_TESTS=1 on an allocated GPU") + if not torch.cuda.is_available(): + pytest.skip("LingBot-Video pipeline parity requires CUDA") + if torch.cuda.device_count() < NUM_GPUS: + pytest.skip(f"LingBot-Video pipeline parity requires {NUM_GPUS} CUDA devices") + + +def _mean_frame_ssim(actual: torch.Tensor, expected: torch.Tensor) -> float: + """Measure visible frame similarity without introducing MP4 encoding loss.""" + from pytorch_msssim import ssim + + actual_frames = actual.permute(0, 2, 1, 3, 4).flatten(0, 1).clamp(0, 1) + expected_frames = expected.permute(0, 2, 1, 3, 4).flatten(0, 1).clamp(0, 1) + scores = [] + for start in range(0, actual_frames.shape[0], 8): + frame_slice = slice(start, start + 8) + scores.append( + ssim( + actual_frames[frame_slice].cuda(), + expected_frames[frame_slice].cuda(), + data_range=1.0, + size_average=False, + ).cpu()) + return torch.cat(scores).mean().item() + + +def _run_official(latents: torch.Tensor, model_dir: Path) -> tuple[torch.Tensor, bool, str]: + """Run the released production loader with the configured parity dimensions.""" + from lingbot_video.pipeline_lingbot_video import DEFAULT_NEGATIVE_PROMPT + from lingbot_video.runner import _load_diffusers_pipe + from lingbot_video.transformer_lingbot_video import flash_attn_varlen_func_v3 + + dtype_map = { + "default": torch.bfloat16, + "transformer": torch.bfloat16, + "text_encoder": torch.bfloat16, + "vae": torch.float32, + } + pipe = _load_diffusers_pipe( + model_dir, + dtype_map, + mode="t2v", + transformer_subfolder="transformer", + ) + official_batch_cfg = BATCH_CFG and flash_attn_varlen_func_v3 is not None + if BATCH_CFG and not official_batch_cfg: + print("official_batch_cfg=sequential reason=flash_attn_interface.flash_attn_varlen_func unavailable") + try: + with torch.inference_mode(): + output = pipe( + prompt=PROMPT, + negative_prompt=DEFAULT_NEGATIVE_PROMPT, + height=HEIGHT, + width=WIDTH, + num_frames=NUM_FRAMES, + num_inference_steps=NUM_INFERENCE_STEPS, + guidance_scale=3.0, + shift=3.0, + generator=torch.Generator(device="cuda").manual_seed(42), + latents=latents.clone(), + output_type=PARITY_OUTPUT_TYPE, + batch_cfg=official_batch_cfg, + return_dict=True, + ) + return _as_channel_first_tensor(output.frames), official_batch_cfg, DEFAULT_NEGATIVE_PROMPT + finally: + del pipe + gc.collect() + torch.cuda.empty_cache() + + +def _run_fastvideo( + latents: torch.Tensor, + model_dir: Path, + output_dir: Path, + negative_prompt: str, +) -> torch.Tensor: + """Run the converted native pipeline with the same prompt, seed, shape, and schedule.""" + from fastvideo import VideoGenerator + + generator = VideoGenerator.from_pretrained( + str(model_dir), + num_gpus=NUM_GPUS, + sp_size=SP_SIZE, + use_fsdp_inference=USE_FSDP, + dit_cpu_offload=False, + dit_layerwise_offload=False, + vae_cpu_offload=True, + text_encoder_cpu_offload=True, + pin_cpu_memory=False, + output_type=PARITY_OUTPUT_TYPE, + refine_enabled=False, + ) + try: + worker_backends = generator.executor.collective_rpc(_configure_worker_backends) + for backend in worker_backends: + assert backend["cudnn_sdp"] is False + if DETERMINISTIC: + assert backend["deterministic"] is True + assert backend["cudnn"] is False + assert backend["matmul_tf32"] is False + assert backend["cudnn_tf32"] is False + if FORCE_MATH_SDPA: + assert backend["flash"] is False + assert backend["mem_efficient"] is False + assert backend["math"] is True + result = generator.generate_video( + prompt=PROMPT, + negative_prompt=negative_prompt, + output_path=str(output_dir), + save_video=False, + return_frames=True, + height=HEIGHT, + width=WIDTH, + num_frames=NUM_FRAMES, + num_inference_steps=NUM_INFERENCE_STEPS, + guidance_scale=3.0, + batch_cfg=BATCH_CFG, + seed=42, + latents=latents.clone(), + ) + samples = cast(dict[str, Any], result)["samples"] + return _as_channel_first_tensor(samples) + finally: + generator.shutdown() + + +def test_lingbot_video_pipeline_matches(tmp_path: Path) -> None: + """Compare official and FastVideo latent or decoded outputs.""" + _require_gpu_test() + if PARITY_VARIANT == "dense": + official_root = download_components( + OFFICIAL_DENSE, + "scheduler", + "text_encoder", + "processor", + "transformer", + "vae", + ) + fastvideo_root = download_components( + FASTVIDEO_DENSE, + "scheduler", + "text_encoder", + "tokenizer", + "transformer", + "vae", + ) + elif PARITY_VARIANT == "moe": + official_root = materialize_component_view( + OFFICIAL_MOE, + tmp_path / "official_moe_base", + "scheduler", + "text_encoder", + "processor", + "transformer", + "vae", + ) + fastvideo_root = materialize_component_view( + FASTVIDEO_MOE, + tmp_path / "fastvideo_moe_base", + "scheduler", + "text_encoder", + "tokenizer", + "transformer", + "vae", + ) + else: + raise ValueError(f"unsupported LingBot-Video parity variant: {PARITY_VARIANT}") + latents = torch.randn( + (1, 16, (NUM_FRAMES - 1) // 4 + 1, HEIGHT // 8, WIDTH // 8), + generator=torch.Generator(device="cpu").manual_seed(42), + dtype=torch.float32, + ) + original_sdp_backends = ( + torch.backends.cuda.cudnn_sdp_enabled(), + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + ) + original_deterministic = torch.are_deterministic_algorithms_enabled() + original_cudnn = torch.backends.cudnn.enabled + original_matmul_tf32 = torch.backends.cuda.matmul.allow_tf32 + original_cudnn_tf32 = torch.backends.cudnn.allow_tf32 + torch.backends.cuda.enable_cudnn_sdp(False) + if DETERMINISTIC: + torch.use_deterministic_algorithms(True) + torch.backends.cudnn.enabled = False + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + if FORCE_MATH_SDPA: + torch.backends.cuda.enable_flash_sdp(False) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_math_sdp(True) + try: + expected, official_batch_cfg, negative_prompt = _run_official(latents, official_root) + actual = _run_fastvideo(latents, fastvideo_root, tmp_path, negative_prompt) + finally: + torch.backends.cuda.enable_cudnn_sdp(original_sdp_backends[0]) + torch.backends.cuda.enable_flash_sdp(original_sdp_backends[1]) + torch.backends.cuda.enable_mem_efficient_sdp(original_sdp_backends[2]) + torch.backends.cuda.enable_math_sdp(original_sdp_backends[3]) + torch.use_deterministic_algorithms(original_deterministic) + torch.backends.cudnn.enabled = original_cudnn + torch.backends.cuda.matmul.allow_tf32 = original_matmul_tf32 + torch.backends.cudnn.allow_tf32 = original_cudnn_tf32 + assert actual.shape == expected.shape + drift = (actual - expected).abs() + print(f"pipeline_max_abs={drift.max().item():.8f} pipeline_mean_abs={drift.mean().item():.8f}") + if not BATCH_CFG: + assert torch.equal(actual, expected) + return + if PARITY_OUTPUT_TYPE != "np": + raise ValueError("the batched-CFG semantic gate requires decoded pixels") + if official_batch_cfg: + mean_ssim = _mean_frame_ssim(actual, expected) + print(f"pipeline_mean_ssim={mean_ssim:.8f}") + assert mean_ssim >= BATCH_CFG_MIN_SSIM + return + if not FORCE_MATH_SDPA: + raise ValueError("the dependency-free batched-CFG semantic gate requires math SDPA") + assert drift.max().item() <= BATCH_CFG_PIXEL_MAX_ABS + assert drift.mean().item() <= BATCH_CFG_PIXEL_MEAN_ABS diff --git a/tests/local_tests/pipelines/test_lingbot_video_pipeline_smoke.py b/tests/local_tests/pipelines/test_lingbot_video_pipeline_smoke.py new file mode 100644 index 000000000..8667908bf --- /dev/null +++ b/tests/local_tests/pipelines/test_lingbot_video_pipeline_smoke.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Import, registry, stage-contract, and optional real Dense T2V smoke tests.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch + +from tests.local_tests.lingbot_video.hf_assets import FASTVIDEO_DENSE, download_components + + +def test_lingbot_video_pipeline_registry_and_preset(tmp_path: Path) -> None: + """Verify exact class resolution, required modules, config, and official defaults.""" + from fastvideo.api.presets import get_preset + from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig + from fastvideo.fastvideo_args import WorkloadType + from fastvideo.pipelines.basic.lingbot_video.lingbot_video_pipeline import ( + EntryClass, + LingBotVideoPipeline, + ) + from fastvideo.registry import get_model_info, get_preset_selection + + model_dir = tmp_path / "arbitrary-dense-layout" + model_index = { + "_class_name": "LingBotVideoDensePipeline", + "_diffusers_version": "0.39.0", + "scheduler": ["diffusers", "FlowUniPCMultistepScheduler"], + "text_encoder": ["transformers", "LingBotVideoQwen3VLTextModel"], + "tokenizer": ["transformers", "Qwen3VLProcessor"], + "transformer": ["diffusers", "LingBotVideoTransformer3DModel"], + "vae": ["diffusers", "AutoencoderKLWan"], + } + model_dir.mkdir() + for component in LingBotVideoPipeline._required_config_modules: + (model_dir / component).mkdir() + (model_dir / "model_index.json").write_text(json.dumps(model_index), encoding="utf-8") + assert model_index["_class_name"] == "LingBotVideoDensePipeline" + assert EntryClass is LingBotVideoPipeline + assert set(model_index) >= set(LingBotVideoPipeline._required_config_modules) + preset_name, family = get_preset_selection(str(model_dir)) + assert (preset_name, family) == ("lingbot_video_dense_t2v", "lingbot_video") + preset = get_preset(preset_name, family) + expected_defaults = { + "height": 480, + "width": 832, + "num_frames": 121, + "fps": 24, + "num_inference_steps": 40, + "guidance_scale": 3.0, + "batch_cfg": True, + "seed": 42, + } + assert all(preset.defaults[key] == value for key, value in expected_defaults.items()) + info = get_model_info(str(model_dir), workload_type=WorkloadType.T2V) + assert info.pipeline_cls is LingBotVideoPipeline + assert info.pipeline_config_cls is LingBotVideoT2VConfig + + +def test_lingbot_video_prompt_crop_contract() -> None: + """Crop exactly 140 template tokens and trim right padding for batch one.""" + from fastvideo.configs.models.encoders import BaseEncoderOutput + from fastvideo.configs.pipelines.lingbot_video import ( + PROMPT_CROP_START, + postprocess_lingbot_video_text, + preprocess_lingbot_video_prompt, + ) + + assert PROMPT_CROP_START == 140 + assert "<|im_start|>user\nfox<|im_end|>" in preprocess_lingbot_video_prompt("fox") + hidden = torch.arange(146 * 2, dtype=torch.float32).reshape(1, 146, 2) + mask = torch.cat((torch.ones(1, 144), torch.zeros(1, 2)), dim=1) + embeds, cropped_mask = postprocess_lingbot_video_text(BaseEncoderOutput(hidden_states=(hidden, )), mask) + assert embeds.shape == (1, 4, 2) + assert cropped_mask.shape == (1, 4) + assert torch.equal(embeds, hidden[:, 140:144]) + + +def test_lingbot_video_latent_geometry_and_validation(monkeypatch: pytest.MonkeyPatch) -> None: + """Validate 4n+1 frames, multiples of 16, and fp32 latent geometry.""" + from fastvideo.pipelines.basic.lingbot_video import stages + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + monkeypatch.setattr(stages, "get_local_torch_device", lambda: torch.device("cpu")) + transformer = SimpleNamespace(num_channels_latents=16) + stage = stages.LingBotVideoLatentPreparationStage(transformer) + supplied = torch.zeros(1, 16, 3, 8, 10, dtype=torch.bfloat16) + batch = ForwardBatch( + data_type="video", + prompt="fox", + num_frames=9, + height=64, + width=80, + latents=supplied, + ) + result = stage.forward(batch, cast(Any, SimpleNamespace())) + assert result.raw_latent_shape == (1, 16, 3, 8, 10) + assert result.latents is not None and result.latents.dtype == torch.float32 + + +def test_lingbot_video_uses_released_vae_denormalization_arithmetic() -> None: + """Use reciprocal-then-divide for LingBot while preserving the generic formula.""" + from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig + from fastvideo.pipelines.stages.decoding import DecodingStage + + vae = SimpleNamespace(config=SimpleNamespace(latents_mean=(-0.7571, ), latents_std=(2.8184, )), ) + stage = DecodingStage(vae) + latents = torch.tensor([[[[[-0.7278813123703003]]]]], dtype=torch.float32) + mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1) + std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1) + released = latents / (1.0 / std) + mean + generic = latents * std + mean + assert not torch.equal(released, generic) + + lingbot_args = cast(Any, SimpleNamespace(pipeline_config=LingBotVideoT2VConfig())) + generic_args = cast(Any, SimpleNamespace(pipeline_config=SimpleNamespace())) + assert torch.equal(stage._denormalize_latents(latents, lingbot_args), released) + assert torch.equal(stage._denormalize_latents(latents, generic_args), generic) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="real pipeline smoke requires CUDA") +def test_lingbot_video_pipeline_load_generate_smoke(tmp_path: Path) -> None: + """Load every converted component and run one tiny latent denoising step.""" + if os.environ.get("LINGBOT_VIDEO_RUN_GPU_TESTS") != "1": + pytest.skip("set LINGBOT_VIDEO_RUN_GPU_TESTS=1 on an allocated GPU") + from fastvideo import VideoGenerator + + use_fsdp_inference = os.environ.get("LINGBOT_VIDEO_USE_FSDP") == "1" + num_gpus = int(os.environ.get("LINGBOT_VIDEO_NUM_GPUS", "1")) + sp_size = int(os.environ.get("LINGBOT_VIDEO_SP_SIZE", "1")) + if torch.cuda.device_count() < num_gpus: + pytest.skip(f"LingBot-Video pipeline smoke requires {num_gpus} CUDA devices") + model_dir = download_components( + FASTVIDEO_DENSE, + "scheduler", + "text_encoder", + "tokenizer", + "transformer", + "vae", + ) + generator = VideoGenerator.from_pretrained( + str(model_dir), + num_gpus=num_gpus, + sp_size=sp_size, + use_fsdp_inference=use_fsdp_inference, + dit_cpu_offload=False, + dit_layerwise_offload=False, + vae_cpu_offload=True, + text_encoder_cpu_offload=True, + pin_cpu_memory=False, + output_type="latent", + ) + try: + result = generator.generate_video( + prompt="A red fox runs through fresh snow at sunrise.", + output_path=str(tmp_path), + save_video=False, + return_frames=True, + height=32, + width=32, + num_frames=1, + num_inference_steps=1, + guidance_scale=3.0, + seed=42, + ) + finally: + generator.shutdown() + samples = cast(dict[str, Any], result)["samples"] + assert torch.is_tensor(samples) + assert tuple(samples.shape) == (1, 16, 1, 4, 4) + assert torch.isfinite(samples).all() diff --git a/tests/local_tests/pipelines/test_lingbot_video_refiner_pipeline_smoke.py b/tests/local_tests/pipelines/test_lingbot_video_refiner_pipeline_smoke.py new file mode 100644 index 000000000..586775e98 --- /dev/null +++ b/tests/local_tests/pipelines/test_lingbot_video_refiner_pipeline_smoke.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Multi-GPU production smoke for LingBot-Video base generation plus refinement.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, cast + +import pytest +import torch + +from tests.local_tests.lingbot_video.hf_assets import FASTVIDEO_MOE, download_components + + +def test_lingbot_video_moe_refiner_pipeline_smoke(tmp_path: Path) -> None: + """Run two base/refiner steps over a five-frame VAE handoff with both MoE DiTs.""" + if os.environ.get("LINGBOT_VIDEO_RUN_REFINER_PIPELINE_TESTS") != "1": + pytest.skip("Set LINGBOT_VIDEO_RUN_REFINER_PIPELINE_TESTS=1 on a scheduled H200 node.") + required_gpus = int(os.environ.get("LINGBOT_VIDEO_REFINER_NUM_GPUS", "8")) + if not torch.cuda.is_available() or torch.cuda.device_count() < required_gpus: + pytest.skip(f"LingBot-Video refiner smoke requires {required_gpus} CUDA devices.") + from fastvideo import VideoGenerator + + model_dir = download_components( + FASTVIDEO_MOE, + "scheduler", + "text_encoder", + "tokenizer", + "transformer", + "transformer_2", + "vae", + ) + generator = VideoGenerator.from_pretrained( + str(model_dir), + num_gpus=required_gpus, + sp_size=required_gpus, + use_fsdp_inference=True, + refine_enabled=True, + dit_cpu_offload=False, + dit_layerwise_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=True, + pin_cpu_memory=False, + ) + try: + result = generator.generate_video( + prompt="A red fox runs through fresh snow at sunrise.", + output_path=str(tmp_path), + save_video=False, + return_frames=True, + height=32, + width=32, + height_sr=64, + width_sr=64, + num_frames=5, + num_inference_steps=2, + num_inference_steps_sr=2, + guidance_scale=3.0, + guidance_scale_2=3.0, + t_thresh=0.85, + batch_cfg=True, + seed=42, + ) + finally: + generator.shutdown() + samples = cast(dict[str, Any], result)["samples"] + assert torch.is_tensor(samples) + assert tuple(samples.shape) == (1, 3, 5, 64, 64) + assert torch.isfinite(samples).all() diff --git a/tests/local_tests/pipelines/test_lingbot_video_refiner_stages.py b/tests/local_tests/pipelines/test_lingbot_video_refiner_stages.py new file mode 100644 index 000000000..77731148b --- /dev/null +++ b/tests/local_tests/pipelines/test_lingbot_video_refiner_stages.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU contracts for LingBot-Video's decoded-video MoE refiner stages.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any, cast + +import numpy as np +import pytest +import torch + + +class _ZeroDistribution: + """Return a fixed VAE sample without consuming the supplied generator.""" + + def __init__(self, sample: torch.Tensor) -> None: + self.value = sample + + def sample(self, generator: torch.Generator) -> torch.Tensor: + del generator + return self.value + + +class _FakeVAE: + """Record normalized pixels and emit zero latents in Wan geometry.""" + + def __init__(self) -> None: + self.config = SimpleNamespace(latents_mean=[0.0] * 16, latents_std=[1.0] * 16) + self.last_input: torch.Tensor | None = None + self.last_device: torch.device | None = None + + def to(self, device: str | torch.device) -> _FakeVAE: + """Record explicit offload moves used by the refiner stage.""" + self.last_device = torch.device(device) + return self + + def encode(self, video: torch.Tensor) -> _ZeroDistribution: + """Create one zero latent for each 4x8x8 video cell.""" + self.last_input = video + shape = (video.shape[0], 16, (video.shape[2] - 1) // 4 + 1, video.shape[3] // 8, video.shape[4] // 8) + return _ZeroDistribution(torch.zeros(shape, device=video.device, dtype=video.dtype)) + + +class _FakeScheduler: + """Capture the custom refiner schedule passed by the preparation stage.""" + + sigma_max = 0.999 + sigma_min = 0.0 + + def __init__(self) -> None: + self.sigmas: np.ndarray | None = None + self.shift: float | None = None + self.timesteps = torch.empty(0) + + def set_timesteps( + self, + num_inference_steps: int, + device: torch.device, + sigmas: np.ndarray, + shift: float, + ) -> None: + """Mirror the schedule surface needed by LingBot refiner preparation.""" + assert num_inference_steps == len(sigmas) + self.sigmas = sigmas + self.shift = shift + self.timesteps = torch.from_numpy(sigmas * 1000.0).to(device=device, dtype=torch.int64) + + +def test_lingbot_video_refiner_sigmas_match_released_schedule() -> None: + """Match the official eight-step, threshold-0.85 schedule including its tail.""" + from fastvideo.pipelines.basic.lingbot_video.stages import _compute_refiner_sigmas + + actual = _compute_refiner_sigmas(0.999, 0.0, 8, 3.0, 0.85) + expected = np.asarray( + [0.85, 0.83296275, 0.7496248, 0.6424896, 0.49966654, 0.29975995, 0.19983996, 0.09991998], + dtype=np.float32, + ) + np.testing.assert_array_equal(actual, expected) + + +def test_lingbot_video_refiner_preparation_resizes_encodes_and_noises(monkeypatch: pytest.MonkeyPatch, ) -> None: + """Exercise pixel resize, VAE normalization, seeded noise, and scheduler setup.""" + from fastvideo.pipelines.basic.lingbot_video import stages + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + monkeypatch.setattr(stages, "get_local_torch_device", lambda: torch.device("cpu")) + vae = _FakeVAE() + scheduler = _FakeScheduler() + stage = stages.LingBotVideoRefinerPreparationStage(vae, scheduler) + batch = ForwardBatch( + data_type="video", + output=torch.linspace(0.0, 1.0, 3 * 5 * 16 * 32).reshape(1, 3, 5, 16, 32), + height=16, + width=32, + height_sr=32, + width_sr=64, + num_frames=5, + num_inference_steps_sr=8, + t_thresh=0.85, + seed=42, + ) + args = cast( + Any, + SimpleNamespace( + pipeline_config=SimpleNamespace(flow_shift=3.0), + vae_cpu_offload=True, + ), + ) + result = stage.forward(batch, args) + + assert vae.last_input is not None and tuple(vae.last_input.shape) == (1, 3, 5, 32, 64) + assert float(vae.last_input.min()) >= -1.0 and float(vae.last_input.max()) <= 1.0 + assert result.output is None and result.raw_latent_shape == (1, 16, 2, 4, 8) + expected_noise = torch.randn(result.raw_latent_shape, generator=torch.Generator().manual_seed(42)) + torch.testing.assert_close(result.latents, expected_noise * 0.85) + assert scheduler.shift == 1.0 + assert scheduler.sigmas is not None and scheduler.sigmas.shape == (8, ) + assert result.extra["lingbot_video_base_shape"] == (1, 3, 5, 16, 32) + assert vae.last_device == torch.device("cpu") + + +def test_lingbot_video_refiner_uses_zero_cloned_null_condition() -> None: + """Ignore encoded negative text and clone the positive mask for refiner CFG.""" + from fastvideo.pipelines.basic.lingbot_video.stages import LingBotVideoDenoisingStage + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + stage = LingBotVideoDenoisingStage(torch.nn.Linear(2, 2), SimpleNamespace(), refiner=True) + prompt = torch.ones(1, 3, 2) + prompt_mask = torch.tensor([[1, 1, 0]]) + batch = ForwardBatch( + data_type="video", + prompt_embeds=[prompt], + prompt_attention_mask=[prompt_mask], + negative_prompt_embeds=[torch.full_like(prompt, 7.0)], + negative_attention_mask=[torch.zeros_like(prompt_mask)], + guidance_scale=1.0, + guidance_scale_2=3.0, + batch_cfg=True, + ) + condition, mask = stage._prepare_conditions(batch, torch.float32, torch.device("cpu")) + torch.testing.assert_close(condition[0], prompt[0]) + torch.testing.assert_close(condition[1], torch.zeros_like(prompt[0])) + torch.testing.assert_close(mask[0], prompt_mask[0]) + torch.testing.assert_close(mask[1], prompt_mask[0]) + + base_stage = LingBotVideoDenoisingStage(torch.nn.Linear(2, 2), SimpleNamespace()) + base_condition, base_mask = base_stage._prepare_conditions(batch, torch.float32, torch.device("cpu")) + assert base_condition.shape == prompt.shape + torch.testing.assert_close(base_condition, prompt) + torch.testing.assert_close(base_mask, prompt_mask) + + +def test_lingbot_video_sequential_cfg_prepares_negative_once(monkeypatch: pytest.MonkeyPatch, ) -> None: + """Reuse the invariant negative condition across all denoising steps.""" + from fastvideo.pipelines.basic.lingbot_video import stages + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + class _ZeroTransformer(torch.nn.Module): + """Return a zero prediction while exposing one parameter dtype.""" + + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(())) + + def forward(self, hidden_states, *_args, **_kwargs): + """Match the DiT tuple return contract.""" + return (torch.zeros_like(hidden_states), ) + + class _IdentityScheduler: + """Keep the latent unchanged across the focused CFG test.""" + + def step(self, _prediction, _timestep, latents, **_kwargs): + """Match the scheduler tuple return contract.""" + return (latents, ) + + monkeypatch.setattr(stages, "get_local_torch_device", lambda: torch.device("cpu")) + stage = stages.LingBotVideoDenoisingStage(_ZeroTransformer(), _IdentityScheduler()) + original = stage._negative_condition + calls = 0 + + def counted_negative(*args, **kwargs): + """Count condition preparation while preserving its output.""" + nonlocal calls + calls += 1 + return original(*args, **kwargs) + + monkeypatch.setattr(stage, "_negative_condition", counted_negative) + batch = ForwardBatch( + data_type="video", + latents=torch.zeros(1, 1, 1, 1, 1), + timesteps=torch.tensor([1000.0, 500.0]), + prompt_embeds=[torch.ones(1, 2, 3)], + prompt_attention_mask=[torch.ones(1, 2, dtype=torch.long)], + negative_prompt_embeds=[torch.zeros(1, 2, 3)], + negative_attention_mask=[torch.ones(1, 2, dtype=torch.long)], + guidance_scale=3.0, + batch_cfg=False, + ) + + stage.forward(batch, cast(Any, SimpleNamespace())) + + assert calls == 1 + + +def test_lingbot_video_pipeline_loads_refiner_and_vae_encoder(monkeypatch: pytest.MonkeyPatch, ) -> None: + """Honor typed refiner opt-out while loading the optional second transformer.""" + from fastvideo import fastvideo_args as fva + from fastvideo.api.compat import generator_config_to_fastvideo_args + from fastvideo.api.schema import GeneratorConfig + from fastvideo.pipelines import ComposedPipelineBase + from fastvideo.pipelines.basic.lingbot_video.lingbot_video_pipeline import LingBotVideoPipeline + + captured: dict[str, Any] = {} + + def fake_super_load(self, fastvideo_args, loaded_modules=None): + """Capture the required module list after LingBot-specific discovery.""" + del fastvideo_args, loaded_modules + captured["required"] = list(self.required_config_modules) + return {name: object() for name in self.required_config_modules} + + monkeypatch.setattr(ComposedPipelineBase, "load_modules", fake_super_load) + pipe = object.__new__(LingBotVideoPipeline) + pipe.model_path = "/model" + monkeypatch.setattr(pipe, "_load_config", lambda _path: {"transformer_2": ["diffusers", "model"]}) + vae_config = SimpleNamespace(load_encoder=False) + args = cast(Any, SimpleNamespace(pipeline_config=SimpleNamespace(vae_config=vae_config))) + modules = pipe.load_modules(args) + + assert "transformer_2" in captured["required"] + assert "transformer_2" in modules + assert vae_config.load_encoder is True + + base_pipe = object.__new__(LingBotVideoPipeline) + base_pipe.model_path = "/model" + monkeypatch.setattr(base_pipe, "_load_config", lambda _path: {"transformer_2": ["diffusers", "model"]}) + base_vae_config = SimpleNamespace(load_encoder=False) + monkeypatch.setattr(fva.FastVideoArgs, "from_kwargs", lambda **kwargs: SimpleNamespace(**kwargs)) + config = GeneratorConfig(model_path="/model") + config.pipeline.preset_overrides = {"refine": {"enabled": False}} + base_args = generator_config_to_fastvideo_args(config) + base_args.pipeline_config = SimpleNamespace(vae_config=base_vae_config) + base_modules = base_pipe.load_modules(base_args) + assert base_args.refine_enabled is False + assert "transformer_2" not in base_modules + assert base_vae_config.load_encoder is False diff --git a/tests/local_tests/schedulers/test_lingbot_video_scheduler_parity.py b/tests/local_tests/schedulers/test_lingbot_video_scheduler_parity.py new file mode 100644 index 000000000..a96e6f83c --- /dev/null +++ b/tests/local_tests/schedulers/test_lingbot_video_scheduler_parity.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LingBot-Video Flow-UniPC scheduler parity. + +Coverage scope: implementation_subcomponent. The official repository vendors +the scheduler used by every Dense, MoE, and refiner denoising loop. These tests +compare its schedule and concrete multistep outputs with FastVideo's native +scheduler under the released checkpoint configuration. +""" +from __future__ import annotations + +import json + +import torch +from torch.testing import assert_close + +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler as FastVideoFlowUniPCMultistepScheduler, ) +from lingbot_video.scheduling_flow_unipc import ( + FlowUniPCMultistepScheduler as OfficialFlowUniPCMultistepScheduler, ) +from tests.local_tests.lingbot_video.hf_assets import OFFICIAL_DENSE, download_components + +PARITY_SCOPE = "implementation_subcomponent" + + +def _scheduler_kwargs() -> dict[str, object]: + """Load the exact scheduler constructor arguments released with Dense.""" + model_path = download_components(OFFICIAL_DENSE, "scheduler") + config_path = model_path / "scheduler" / "scheduler_config.json" + config = json.loads(config_path.read_text()) + return {key: value for key, value in config.items() if not key.startswith("_")} + + +def _scheduler_pair() -> tuple[OfficialFlowUniPCMultistepScheduler, FastVideoFlowUniPCMultistepScheduler]: + """Construct independent official and FastVideo schedulers from one config.""" + kwargs = _scheduler_kwargs() + return OfficialFlowUniPCMultistepScheduler(**kwargs), FastVideoFlowUniPCMultistepScheduler(**kwargs) + + +def test_lingbot_video_scheduler_timesteps_and_sigmas_match() -> None: + """Require exact released 40-step timestep and sigma schedules.""" + official, fastvideo = _scheduler_pair() + official.set_timesteps(40, device="cpu", shift=3.0) + fastvideo.set_timesteps(40, device="cpu", shift=3.0) + + assert_close(fastvideo.timesteps, official.timesteps, atol=0, rtol=0) + assert_close(fastvideo.sigmas, official.sigmas, atol=0, rtol=0) + + +def test_lingbot_video_scheduler_multistep_outputs_match() -> None: + """Compare four deterministic updates so UniPC history is exercised.""" + official, fastvideo = _scheduler_pair() + official.set_timesteps(4, device="cpu", shift=3.0) + fastvideo.set_timesteps(4, device="cpu", shift=3.0) + generator = torch.Generator(device="cpu").manual_seed(42) + official_sample = torch.randn((1, 16, 3, 4, 4), generator=generator) + fastvideo_sample = official_sample.clone() + + for index in range(4): + model_generator = torch.Generator(device="cpu").manual_seed(100 + index) + model_output = torch.randn(official_sample.shape, generator=model_generator) + official_sample = official.step(model_output, official.timesteps[index], official_sample, return_dict=False)[0] + fastvideo_sample = fastvideo.step(model_output, fastvideo.timesteps[index], fastvideo_sample, + return_dict=False)[0] + assert_close(fastvideo_sample, official_sample, atol=0, rtol=0) diff --git a/tests/local_tests/transformers/test_lingbot_video_dense_dit_parity.py b/tests/local_tests/transformers/test_lingbot_video_dense_dit_parity.py new file mode 100644 index 000000000..a94da3271 --- /dev/null +++ b/tests/local_tests/transformers/test_lingbot_video_dense_dit_parity.py @@ -0,0 +1,233 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Dense LingBot-Video DiT parity scaffold. + +Coverage scope: implementation_subcomponent. The official side loads the real +Dense 1.3B transformer checkpoint. The FastVideo side uses the planned native +DiT config/class path and strict-loads the same checkpoint tensor surface. +""" + +from __future__ import annotations + +import importlib +import json +import os +from pathlib import Path +from typing import Any + +import pytest +import torch +from torch.testing import assert_close + +from fastvideo.models.loader.fsdp_load import maybe_load_fsdp_model +from tests.local_tests.lingbot_video.hf_assets import ( + FASTVIDEO_DENSE, + OFFICIAL_DENSE, + download_components, +) + +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29519") +os.environ.setdefault("DISABLE_SP", "1") +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") +os.environ.setdefault("DIFFUSERS_ATTN_BACKEND", "native") + +REPO_ROOT = Path(__file__).resolve().parents[3] +FASTVIDEO_CONFIG_MODULE = "fastvideo.configs.models.dits.lingbot_video" +FASTVIDEO_CONFIG_CLASS = "LingBotVideoConfig" +FASTVIDEO_MODEL_MODULE = "fastvideo.models.dits.lingbot_video" +FASTVIDEO_MODEL_CLASS = "LingBotVideoTransformer3DModel" +FASTVIDEO_MODEL_FILE = REPO_ROOT / "fastvideo" / "models" / "dits" / "lingbot_video.py" +PARITY_SCOPE = "implementation_subcomponent" + + +def _load_fastvideo_types_or_skip(): + """Resolve the planned native class/config, skipping only until the class exists.""" + if not FASTVIDEO_MODEL_FILE.is_file(): + pytest.skip(f"FastVideo Dense DiT class missing: {FASTVIDEO_MODEL_MODULE}.{FASTVIDEO_MODEL_CLASS}") + model_module = importlib.import_module(FASTVIDEO_MODEL_MODULE) + try: + model_class = getattr(model_module, FASTVIDEO_MODEL_CLASS) + except AttributeError: + pytest.skip(f"FastVideo Dense DiT class missing: {FASTVIDEO_MODEL_MODULE}.{FASTVIDEO_MODEL_CLASS}") + + config_module = importlib.import_module(FASTVIDEO_CONFIG_MODULE) + return getattr(config_module, FASTVIDEO_CONFIG_CLASS), model_class + + +def _require_official_assets(model_dir: Path) -> None: + """Fail clearly when the pinned official Dense snapshot is incomplete.""" + transformer_dir = model_dir / "transformer" + required = ( + transformer_dir / "config.json", + transformer_dir / "diffusion_pytorch_model.safetensors", + ) + missing = [str(path) for path in required if not path.is_file()] + assert not missing, f"Missing required official Dense DiT assets: {missing}" + + +def _load_official_model(model_dir: Path, device: torch.device, dtype: torch.dtype) -> torch.nn.Module: + """Load the official Dense transformer through its production Diffusers path.""" + _require_official_assets(model_dir) + from lingbot_video.transformer_lingbot_video import LingBotVideoTransformer3DModel + + model = LingBotVideoTransformer3DModel.from_pretrained( + str(model_dir), + subfolder="transformer", + torch_dtype=dtype, + ) + return model.to(device=device).eval() + + +def _load_fastvideo_model(transformer_dir: Path, device: torch.device, dtype: torch.dtype) -> torch.nn.Module: + """Load the native Dense DiT through the production mixed-precision loader.""" + FastVideoConfig, FastVideoModel = _load_fastvideo_types_or_skip() + with (transformer_dir / "config.json").open(encoding="utf-8") as file: + hf_config = json.load(file) + + config = FastVideoConfig() + model = maybe_load_fsdp_model( + model_cls=FastVideoModel, + init_params={ + "config": config, + "hf_config": hf_config + }, + weight_dir_list=[str(transformer_dir / "diffusion_pytorch_model.safetensors")], + device=device, + hsdp_replicate_dim=1, + hsdp_shard_dim=1, + default_dtype=dtype, + param_dtype=dtype, + reduce_dtype=torch.float32, + strict=True, + cpu_offload=False, + fsdp_inference=False, + training_mode=False, + pin_cpu_memory=False, + ).eval() + assert model.patch_embedder.weight.dtype == torch.bfloat16 + assert model.blocks[0].scale_shift_table.dtype == torch.float32 + assert model.blocks[0].attn.to_q.weight.dtype == torch.bfloat16 + return model + + +def _make_inputs(device: torch.device, dtype: torch.dtype) -> dict[str, torch.Tensor]: + """Create a small deterministic upstream-format latent and text condition.""" + generator = torch.Generator(device="cpu").manual_seed(20260711) + hidden_states = torch.randn(1, 16, 1, 4, 4, generator=generator, dtype=torch.float32) + encoder_hidden_states = torch.randn(1, 22, 2560, generator=generator, dtype=torch.float32) + return { + "hidden_states": hidden_states.to(device=device), + "timestep": torch.tensor([500.0], device=device, dtype=torch.float32), + "encoder_hidden_states": encoder_hidden_states.to(device=device, dtype=dtype), + "encoder_attention_mask": torch.ones((1, 22), device=device, dtype=torch.long), + "return_dict": False, + } + + +def _run_model(model: torch.nn.Module, inputs: dict[str, torch.Tensor]) -> torch.Tensor: + """Execute the full Dense DiT and extract its denoised latent tensor.""" + with torch.inference_mode(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + output = model(**inputs) + if isinstance(output, tuple): + output = output[0] + elif hasattr(output, "sample"): + output = output.sample + assert torch.is_tensor(output), f"Dense DiT output is not a tensor: {type(output)}" + return output.detach().float().cpu() + + +def _capture_intermediates(model: torch.nn.Module) -> tuple[dict[str, torch.Tensor], list[Any]]: + """Capture ordered parity checkpoints without changing either implementation.""" + captured: dict[str, torch.Tensor] = {} + handles: list[Any] = [] + suffixes = ( + "patch_embedder", + "time_proj", + "time_embedder", + "time_modulation", + "text_embedder", + "norm1", + "to_q", + "norm_q", + "to_k", + "norm_k", + "to_v", + "to_out", + "norm_post_attn", + "norm2", + "gate_proj", + "up_proj", + "down_proj", + "norm_post_ffn", + "norm_out", + "norm_out_modulation", + "proj_out", + ) + + def save(name: str): + + def hook(_module, _inputs, output): + tensor = output[0] if isinstance(output, tuple) else output + if torch.is_tensor(tensor): + captured[name] = tensor.detach().float().cpu() + + return hook + + for name, module in model.named_modules(): + if name == "" or not (name.startswith("blocks.") or name in suffixes): + continue + if name.startswith("blocks.") and not name.endswith(suffixes): + parts = name.split(".") + if len(parts) != 2: + continue + handles.append(module.register_forward_hook(save(name))) + return captured, handles + + +def _report_first_intermediate_drift(official: dict[str, torch.Tensor], fastvideo: dict[str, torch.Tensor]) -> None: + """Print the first execution-ordered checkpoint whose values diverge materially.""" + for name, expected in official.items(): + actual = fastvideo.get(name) + if actual is None or actual.shape != expected.shape: + print(f"first_intermediate_drift={name} shape={getattr(actual, 'shape', None)} vs {expected.shape}") + return + drift = (actual - expected).abs() + if drift.max().item() > 1e-3: + print( + f"first_intermediate_drift={name} max_abs={drift.max().item():.8f} mean_abs={drift.mean().item():.8f}") + return + print("first_intermediate_drift=not_found") + + +def test_lingbot_video_dense_dit_outputs_match() -> None: + """Compare real official and native Dense DiT outputs for identical tensors.""" + if os.environ.get("LINGBOT_VIDEO_RUN_GPU_TESTS") != "1": + pytest.skip("Set LINGBOT_VIDEO_RUN_GPU_TESTS=1 on a scheduled GPU node.") + if not torch.cuda.is_available(): + pytest.skip("LingBot-Video Dense DiT parity requires CUDA.") + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 + official_model_dir = download_components(OFFICIAL_DENSE, "transformer") + fastvideo_model_dir = download_components(FASTVIDEO_DENSE, "transformer") + + fastvideo = _load_fastvideo_model(fastvideo_model_dir / "transformer", device, dtype) + official = _load_official_model(official_model_dir, device, dtype) + inputs = _make_inputs(device, dtype) + official_intermediates, official_handles = _capture_intermediates(official) + fastvideo_intermediates, fastvideo_handles = _capture_intermediates(fastvideo) + official_output = _run_model(official, inputs) + fastvideo_output = _run_model(fastvideo, inputs) + for handle in official_handles + fastvideo_handles: + handle.remove() + + assert official_output.shape == fastvideo_output.shape + difference = (official_output - fastvideo_output).abs() + official_abs_mean = official_output.abs().mean() + abs_mean_drift = (fastvideo_output.abs().mean() - official_abs_mean).abs() + print(f"official_abs_mean={official_abs_mean.item():.6f} " + f"fastvideo_abs_mean={fastvideo_output.abs().mean().item():.6f} " + f"diff_max={difference.max().item():.6f} diff_mean={difference.mean().item():.6f}") + _report_first_intermediate_drift(official_intermediates, fastvideo_intermediates) + + assert abs_mean_drift <= official_abs_mean * 0.05 + assert_close(fastvideo_output, official_output, atol=1e-2, rtol=1e-2) diff --git a/tests/local_tests/transformers/test_lingbot_video_moe.py b/tests/local_tests/transformers/test_lingbot_video_moe.py new file mode 100644 index 000000000..e06bc44b3 --- /dev/null +++ b/tests/local_tests/transformers/test_lingbot_video_moe.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Focused CPU coverage for the native LingBot-Video sparse MoE path.""" + +from __future__ import annotations + +import torch +from torch.nn import functional as F +from torch.testing import assert_close + +from fastvideo.configs.models.dits.lingbot_video import LingBotVideoConfig +from fastvideo.models.dits.lingbot_video import ( + LingBotVideoBlock, + LingBotVideoMLP, + LingBotVideoRotaryEmbedding, + LingBotVideoRouter, + LingBotVideoSparseMoeBlock, + LingBotVideoTransformer3DModel, + _make_joint_position_ids, +) +from fastvideo.models.loader.fsdp_load import set_default_dtype + + +def test_router_selection_bias_does_not_change_gating_weight() -> None: + """Select with the correction bias while gathering the bias-free sigmoid score.""" + router = LingBotVideoRouter(2, 3, 1, "sigmoid", True, None, None, 1.0) + with torch.no_grad(): + router.weight.zero_() + router.e_score_correction_bias.copy_(torch.tensor([0.0, 0.0, 1.0])) + + indices, top_scores, logits, scores, scores_for_choice = router(torch.zeros(1, 2, dtype=torch.bfloat16)) + + assert indices.item() == 2 + assert top_scores.dtype == torch.bfloat16 + assert top_scores.item() == 0.5 + assert logits.dtype == torch.float32 + assert scores.dtype == torch.float32 + assert scores_for_choice[0, 2].item() == 1.5 + + +def test_router_group_limited_topk_uses_two_expert_group_score() -> None: + """Keep choices inside the group whose two best expert scores have the largest sum.""" + router = LingBotVideoRouter(1, 8, 2, "sigmoid", False, 4, 1, 1.0) + with torch.no_grad(): + router.weight[:, 0].copy_(torch.tensor([6.0, 5.0, 9.0, -9.0, 4.0, 4.0, 3.0, 3.0])) + + indices, _, _, _, _ = router(torch.ones(1, 1)) + + assert set(indices[0].tolist()) == {0, 1} + + +def test_router_normalizes_then_applies_routed_scale() -> None: + """Normalize selected bias-free scores to one before multiplying by route scale.""" + router = LingBotVideoRouter(1, 3, 2, "sigmoid", True, None, None, 2.5) + with torch.no_grad(): + router.weight[:, 0].copy_(torch.tensor([2.0, 1.0, -1.0])) + + indices, top_scores, _, scores, _ = router(torch.ones(1, 1)) + selected_scores = scores.gather(1, indices) + expected = selected_scores / selected_scores.sum(dim=-1, keepdim=True) * 2.5 + + assert_close(top_scores, expected) + assert_close(top_scores.sum(dim=-1), torch.tensor([2.5])) + + +def test_sparse_restore_uses_fp32_weighted_sum() -> None: + """Restore expert-major rows to token-major order and weight each selected expert.""" + expert_output = torch.tensor( + [[1.0, 2.0], [10.0, 20.0], [3.0, 4.0], [30.0, 40.0]], + dtype=torch.bfloat16, + ) + sorted_positions = torch.tensor([0, 2, 1, 3]) + sorted_scores = torch.tensor([0.25, 0.5, 0.75, 0.5]) + + restored = LingBotVideoSparseMoeBlock._restore_tokens(expert_output, + sorted_positions, + sorted_scores, + num_tokens=2, + top_k=2) + + assert restored.dtype == torch.bfloat16 + assert_close(restored.float(), torch.tensor([[2.5, 3.5], [20.0, 30.0]])) + + +def test_grouped_expert_padding_preserves_segments() -> None: + """Build aligned expert segments without per-expert device scalar reads.""" + tokens = torch.arange(12, dtype=torch.float32).reshape(3, 4) + counts = torch.tensor([2, 0, 1], dtype=torch.int64) + + input_shape, padded, indices, aligned_counts = LingBotVideoSparseMoeBlock._pad_grouped_tokens(tokens, + counts, + align=2) + + assert input_shape == torch.Size((4, 4)) + assert torch.equal(aligned_counts, torch.tensor([2, 2, 2], dtype=torch.int32)) + assert torch.equal(indices[:6], torch.tensor([0, 1, 3, 3, 2, 3])) + assert_close(padded[:6], torch.vstack((tokens[:2], torch.zeros(2, 4), tokens[2:], torch.zeros(1, 4)))) + + +def test_rotary_precomputed_maxima_matches_tensor_reduction() -> None: + """Use known grid maxima without changing the gathered rotary frequencies.""" + positions = _make_joint_position_ids(3, 2, 2, 2, torch.device("cpu")) + dynamic = LingBotVideoRotaryEmbedding((2, 2, 2), (2, 1, 1), 256.0) + precomputed = LingBotVideoRotaryEmbedding((2, 2, 2), (2, 1, 1), 256.0) + + expected = dynamic(positions) + actual = precomputed(positions, maxima=(5, 1, 1)) + + assert torch.equal(actual, expected) + assert precomputed.axes_lens == dynamic.axes_lens + + +def test_sparse_moe_cpu_fallback_matches_direct_expert_sum() -> None: + """Exercise CPU dispatch and compare routed expert output with a direct calculation.""" + torch.manual_seed(7) + block = LingBotVideoSparseMoeBlock( + hidden_size=4, + intermediate_size=8, + num_experts=3, + top_k=2, + moe_intermediate_size=5, + score_func="sigmoid", + norm_topk_prob=True, + n_group=None, + topk_group=None, + routed_scaling_factor=1.5, + n_shared_experts=None, + ) + with torch.no_grad(): + for parameter in block.parameters(): + parameter.uniform_(-0.25, 0.25) + hidden_states = torch.randn(2, 3, 4) + tokens = hidden_states.reshape(-1, 4) + top_indices, top_scores, _, _, _ = block.router(tokens) + expected = torch.zeros_like(tokens, dtype=torch.float32) + for token_index in range(tokens.shape[0]): + for choice_index in range(top_indices.shape[1]): + expert_index = int(top_indices[token_index, choice_index]) + expert_input = tokens[token_index] + expert_hidden = F.silu(expert_input @ block.experts.w1[expert_index].T) + expert_hidden = expert_hidden * (expert_input @ block.experts.w3[expert_index].T) + expert_output = expert_hidden @ block.experts.w2[expert_index].T + expected[token_index] += expert_output.float() * top_scores[token_index, choice_index].float() + + actual = block(hidden_states) + + assert_close(actual.reshape(-1, 4).float(), expected, atol=1e-6, rtol=1e-6) + + +def test_block_selection_preserves_dense_layers() -> None: + """Use Dense MLPs unless the configured sparse schedule selects the layer.""" + common = { + "hidden_size": 12, + "num_attention_heads": 2, + "intermediate_size": 16, + "norm_eps": 1e-6, + "qkv_bias": False, + "out_bias": True, + "num_experts_per_tok": 2, + "moe_intermediate_size": 6, + "decoder_sparse_step": 2, + "mlp_only_layers": (), + "n_shared_experts": 1, + "score_func": "sigmoid", + "norm_topk_prob": True, + "n_group": 2, + "topk_group": 1, + "routed_scaling_factor": 1.0, + } + dense_variant = LingBotVideoBlock(**common, num_experts=0, layer_idx=1) + scheduled_dense = LingBotVideoBlock(**common, num_experts=4, layer_idx=0) + scheduled_sparse = LingBotVideoBlock(**common, num_experts=4, layer_idx=1) + + assert isinstance(dense_variant.ffn, LingBotVideoMLP) + assert isinstance(scheduled_dense.ffn, LingBotVideoMLP) + assert isinstance(scheduled_sparse.ffn, LingBotVideoSparseMoeBlock) + assert set(dense_variant.ffn.state_dict()) == { + "gate_proj.weight", + "up_proj.weight", + "down_proj.weight", + } + + +def test_released_moe_meta_state_names_shapes_and_dtype_policy() -> None: + """Build the released 30B MoE state surface on meta without materializing weights.""" + released_config = { + "patch_size": [1, 2, 2], + "in_channels": 16, + "out_channels": 16, + "hidden_size": 2048, + "num_attention_heads": 16, + "depth": 48, + "intermediate_size": 6144, + "text_dim": 2560, + "freq_dim": 256, + "norm_eps": 1e-6, + "rope_theta": 256.0, + "axes_dims": [32, 48, 48], + "axes_lens": [4096, 512, 512], + "qkv_bias": False, + "out_bias": True, + "patch_embed_bias": True, + "timestep_mlp_bias": True, + "num_experts": 128, + "num_experts_per_tok": 8, + "moe_intermediate_size": 768, + "decoder_sparse_step": 1, + "mlp_only_layers": [], + "n_shared_experts": 1, + "score_func": "sigmoid", + "norm_topk_prob": True, + "n_group": 4, + "topk_group": 2, + "routed_scaling_factor": 2.5, + } + with set_default_dtype(torch.bfloat16), torch.device("meta"): + model = LingBotVideoTransformer3DModel(LingBotVideoConfig(), released_config) + state = model.state_dict() + + assert state["blocks.0.ffn.router.weight"].shape == (128, 2048) + assert state["blocks.0.ffn.router.e_score_correction_bias"].shape == (128, ) + assert state["blocks.0.ffn.experts.w1"].shape == (128, 768, 2048) + assert state["blocks.0.ffn.experts.w2"].shape == (128, 2048, 768) + assert state["blocks.0.ffn.experts.w3"].shape == (128, 768, 2048) + assert state["blocks.0.ffn.shared_experts.gate_proj.weight"].shape == (768, 2048) + assert "blocks.47.ffn.router.weight" in state + assert state["blocks.0.ffn.router.weight"].is_meta + assert model._get_parameter_dtype("blocks.0.ffn.router.weight", torch.bfloat16) == torch.float32 + assert model._get_parameter_dtype("blocks.0.ffn.experts.w1", torch.bfloat16) == torch.bfloat16 + assert model._get_parameter_dtype("blocks.0.ffn.shared_experts.gate_proj.weight", torch.bfloat16) == torch.bfloat16 diff --git a/tests/local_tests/transformers/test_lingbot_video_moe_block_parity.py b/tests/local_tests/transformers/test_lingbot_video_moe_block_parity.py new file mode 100644 index 000000000..8d58381c3 --- /dev/null +++ b/tests/local_tests/transformers/test_lingbot_video_moe_block_parity.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Real-checkpoint CUDA parity for one released LingBot-Video MoE block.""" + +from __future__ import annotations + +import gc +import json +import os +from pathlib import Path + +import pytest +from safetensors import safe_open +import torch +from torch.testing import assert_close + +from fastvideo.models.loader.fsdp_load import set_default_dtype +from tests.local_tests.lingbot_video.hf_assets import OFFICIAL_MOE, download_patterns + +os.environ.setdefault("DIFFUSERS_ATTN_BACKEND", "native") +os.environ.setdefault("LINGBOT_MOE_EXPERT_BACKEND", "grouped_mm") +os.environ.setdefault("LINGBOT_MOE_PAD_BACKEND", "loop") +os.environ.setdefault("LINGBOT_MOE_REORDER_BACKEND", "sort") +os.environ.setdefault("LINGBOT_MOE_RESTORE_BACKEND", "scatter") + + +def _block_kwargs(config: dict, layer_index: int = 0) -> dict: + """Select the shared official/native constructor arguments for one block.""" + return { + "hidden_size": config["hidden_size"], + "num_attention_heads": config["num_attention_heads"], + "intermediate_size": config["intermediate_size"], + "norm_eps": config["norm_eps"], + "qkv_bias": config["qkv_bias"], + "out_bias": config["out_bias"], + "num_experts": config["num_experts"], + "num_experts_per_tok": config["num_experts_per_tok"], + "moe_intermediate_size": config["moe_intermediate_size"], + "decoder_sparse_step": config["decoder_sparse_step"], + "mlp_only_layers": config["mlp_only_layers"], + "n_shared_experts": config["n_shared_experts"], + "score_func": config["score_func"], + "norm_topk_prob": config["norm_topk_prob"], + "n_group": config["n_group"], + "topk_group": config["topk_group"], + "routed_scaling_factor": config["routed_scaling_factor"], + "layer_idx": layer_index, + } + + +def _download_transformer_for_block(layer_index: int = 0) -> Path: + """Download the config, index, and only the shards containing one block.""" + metadata = ( + "transformer/config.json", + "transformer/diffusion_pytorch_model.safetensors.index.json", + ) + model_dir = download_patterns(OFFICIAL_MOE, *metadata) + index = json.loads((model_dir / metadata[1]).read_text()) + prefix = f"blocks.{layer_index}." + shards = {shard for name, shard in index["weight_map"].items() if name.startswith(prefix)} + return (download_patterns( + OFFICIAL_MOE, + *metadata, + *(f"transformer/{shard}" for shard in sorted(shards)), + ) / "transformer") + + +def _load_block_state(transformer_dir: Path, layer_index: int = 0) -> dict[str, torch.Tensor]: + """Read only one block's tensors from the indexed 30B checkpoint shards.""" + index = json.loads((transformer_dir / "diffusion_pytorch_model.safetensors.index.json").read_text()) + prefix = f"blocks.{layer_index}." + keys_by_shard: dict[str, list[str]] = {} + for name, shard_name in index["weight_map"].items(): + if name.startswith(prefix): + keys_by_shard.setdefault(shard_name, []).append(name) + state: dict[str, torch.Tensor] = {} + for shard_name, names in keys_by_shard.items(): + with safe_open(transformer_dir / shard_name, framework="pt", device="cpu") as shard: + for name in names: + state[name.removeprefix(prefix)] = shard.get_tensor(name) + if not state: + raise RuntimeError(f"No checkpoint tensors found for {prefix}") + return state + + +def _make_block_inputs(device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Create deterministic mixed-precision block inputs and complex rotary phases.""" + generator = torch.Generator(device="cpu").manual_seed(20260711) + hidden_states = torch.randn(1, 17, 2048, generator=generator, dtype=torch.float32).to(device=device, + dtype=torch.bfloat16) + temb6 = torch.randn(17, 6 * 2048, generator=generator, dtype=torch.float32).to(device) + phases = torch.randn(1, 17, 64, generator=generator, dtype=torch.float32).to(device) + rotary = torch.polar(torch.ones_like(phases), phases).to(torch.complex64) + return hidden_states, temb6, rotary + + +def _run_block(block: torch.nn.Module, device: torch.device) -> torch.Tensor: + """Execute a block with the same Torch SDPA and grouped-mm runtime policy.""" + hidden_states, temb6, rotary = _make_block_inputs(device) + with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): + output = block(hidden_states, temb6, rotary) + return output.detach().float().cpu() + + +def test_lingbot_video_moe_block_matches_official_checkpoint() -> None: + """Compare official and native block zero sequentially on one scheduled GPU.""" + if os.environ.get("LINGBOT_VIDEO_RUN_GPU_TESTS") != "1": + pytest.skip("Set LINGBOT_VIDEO_RUN_GPU_TESTS=1 on a scheduled GPU node.") + if not torch.cuda.is_available(): + pytest.skip("LingBot-Video MoE block parity requires CUDA.") + from lingbot_video.transformer_lingbot_video import LingBotVideoBlock as OfficialBlock + from fastvideo.models.dits.lingbot_video import LingBotVideoBlock as NativeBlock + + transformer_dir = _download_transformer_for_block() + config = json.loads((transformer_dir / "config.json").read_text()) + state = _load_block_state(transformer_dir) + device = torch.device("cuda:0") + torch.backends.cuda.enable_cudnn_sdp(False) + + with set_default_dtype(torch.bfloat16), torch.device("meta"): + official = OfficialBlock(**_block_kwargs(config)) + official.load_state_dict(state, strict=True, assign=True) + official = official.to(device).eval() + official_output = _run_block(official, device) + del official + gc.collect() + torch.cuda.empty_cache() + + with set_default_dtype(torch.bfloat16), torch.device("meta"): + native = NativeBlock(**_block_kwargs(config)) + native.load_state_dict(state, strict=True, assign=True) + native = native.to(device).eval() + native_output = _run_block(native, device) + + difference = (native_output - official_output).abs() + print(f"diff_max={difference.max().item():.8f} diff_mean={difference.mean().item():.8f}") + assert_close(native_output, official_output, atol=0.0, rtol=0.0) diff --git a/tests/local_tests/transformers/test_lingbot_video_moe_dit_parity.py b/tests/local_tests/transformers/test_lingbot_video_moe_dit_parity.py new file mode 100644 index 000000000..4057bdfa2 --- /dev/null +++ b/tests/local_tests/transformers/test_lingbot_video_moe_dit_parity.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Sequential real-checkpoint parity for the released 30B-A3B MoE DiT.""" + +from __future__ import annotations + +import gc +import json +import os +from pathlib import Path +from typing import Any + +import pytest +import torch +from torch.testing import assert_close + +from fastvideo.configs.models.dits.lingbot_video import LingBotVideoConfig +from fastvideo.models.dits.lingbot_video import LingBotVideoTransformer3DModel +from fastvideo.models.loader.fsdp_load import maybe_load_fsdp_model +from tests.local_tests.lingbot_video.hf_assets import ( + FASTVIDEO_MOE, + OFFICIAL_MOE, + download_components, +) + +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29529") +os.environ.setdefault("DISABLE_SP", "1") +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") +os.environ.setdefault("DIFFUSERS_ATTN_BACKEND", "native") +os.environ.setdefault("LINGBOT_MOE_EXPERT_BACKEND", "grouped_mm") +os.environ.setdefault("LINGBOT_MOE_PAD_BACKEND", "loop") +os.environ.setdefault("LINGBOT_MOE_REORDER_BACKEND", "sort") +os.environ.setdefault("LINGBOT_MOE_RESTORE_BACKEND", "scatter") + +PARITY_VARIANT = os.environ.get("LINGBOT_VIDEO_MOE_PARITY_VARIANT", "base") +if PARITY_VARIANT not in {"base", "refiner"}: + raise ValueError(f"Unsupported LingBot-Video MoE parity variant: {PARITY_VARIANT}") +OFFICIAL_SUBFOLDER = "transformer" if PARITY_VARIANT == "base" else "refiner" +NATIVE_SUBFOLDER = "transformer" if PARITY_VARIANT == "base" else "transformer_2" + + +def _make_inputs(device: torch.device) -> dict[str, torch.Tensor | bool]: + """Create a deterministic small latent and single-branch text condition.""" + generator = torch.Generator(device="cpu").manual_seed(20260711) + return { + "hidden_states": torch.randn(1, 16, 1, 4, 4, generator=generator).to(device), + "timestep": torch.tensor([500.0], device=device), + "encoder_hidden_states": torch.randn(1, 22, 2560, generator=generator).to(device=device, dtype=torch.bfloat16), + "encoder_attention_mask": torch.ones(1, 22, device=device, dtype=torch.long), + "return_dict": False, + } + + +def _capture_blocks(model: torch.nn.Module) -> tuple[dict[str, torch.Tensor], list[Any]]: + """Capture each block output so failures identify the first divergent layer.""" + captured: dict[str, torch.Tensor] = {} + handles: list[Any] = [] + + def save(name: str): + + def hook(_module, _inputs, output): + tensor = output[0] if isinstance(output, tuple) else output + captured[name] = tensor.detach().float().cpu() + + return hook + + for name, module in model.named_modules(): + if name.startswith("blocks.") and len(name.split(".")) == 2: + handles.append(module.register_forward_hook(save(name))) + return captured, handles + + +def _run_model(model: torch.nn.Module, device: torch.device) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Run one full MoE DiT call and return final plus per-block CPU tensors.""" + captured, handles = _capture_blocks(model) + with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): + output = model(**_make_inputs(device)) + for handle in handles: + handle.remove() + if isinstance(output, tuple): + output = output[0] + elif hasattr(output, "sample"): + output = output.sample + return output.detach().float().cpu(), captured + + +def _load_official(model_dir: Path, device: torch.device) -> torch.nn.Module: + """Load the released MoE transformer through its official Diffusers class.""" + from lingbot_video.transformer_lingbot_video import LingBotVideoTransformer3DModel as OfficialModel + + model = OfficialModel.from_pretrained( + str(model_dir), + subfolder=OFFICIAL_SUBFOLDER, + torch_dtype=torch.bfloat16, + ) + return model.to(device).eval() + + +def _load_native(transformer_dir: Path, device: torch.device) -> torch.nn.Module: + """Load the native MoE transformer with FastVideo's mixed-dtype production loader.""" + hf_config = json.loads((transformer_dir / "config.json").read_text()) + weight_files = sorted(str(path) for path in transformer_dir.glob("*.safetensors")) + model = maybe_load_fsdp_model( + model_cls=LingBotVideoTransformer3DModel, + init_params={ + "config": LingBotVideoConfig(), + "hf_config": hf_config + }, + weight_dir_list=weight_files, + device=device, + hsdp_replicate_dim=1, + hsdp_shard_dim=1, + default_dtype=torch.bfloat16, + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + strict=True, + cpu_offload=False, + fsdp_inference=False, + training_mode=False, + pin_cpu_memory=False, + ).eval() + assert model.blocks[0].ffn.router.weight.dtype == torch.float32 + assert model.blocks[0].ffn.experts.w1.dtype == torch.bfloat16 + return model + + +def _report_first_block_drift(official: dict[str, torch.Tensor], native: dict[str, torch.Tensor]) -> None: + """Print the first block with nonzero drift in execution order.""" + for name, expected in official.items(): + actual = native[name] + difference = (actual - expected).abs() + if difference.max().item() != 0.0: + print(f"first_block_drift={name} max_abs={difference.max().item():.8f} " + f"mean_abs={difference.mean().item():.8f}") + return + print("first_block_drift=not_found") + + +def test_lingbot_video_moe_dit_matches_official_sequentially() -> None: + """Run one selected official/native 30B DiT pair sequentially on an H200.""" + if os.environ.get("LINGBOT_VIDEO_RUN_FULL_MOE_TESTS") != "1": + pytest.skip("Set LINGBOT_VIDEO_RUN_FULL_MOE_TESTS=1 on a scheduled H200.") + if not torch.cuda.is_available(): + pytest.skip("LingBot-Video MoE DiT parity requires CUDA.") + device = torch.device("cuda:0") + torch.backends.cuda.enable_cudnn_sdp(False) + official_model_dir = download_components(OFFICIAL_MOE, OFFICIAL_SUBFOLDER) + native_model_dir = download_components(FASTVIDEO_MOE, NATIVE_SUBFOLDER) + + official = _load_official(official_model_dir, device) + torch.cuda.reset_peak_memory_stats(device) + official_output, official_blocks = _run_model(official, device) + official_peak = torch.cuda.max_memory_allocated(device) + del official + gc.collect() + torch.cuda.empty_cache() + + native = _load_native(native_model_dir / NATIVE_SUBFOLDER, device) + torch.cuda.reset_peak_memory_stats(device) + native_output, native_blocks = _run_model(native, device) + native_peak = torch.cuda.max_memory_allocated(device) + difference = (native_output - official_output).abs() + print(f"variant={PARITY_VARIANT} diff_max={difference.max().item():.8f} " + f"diff_mean={difference.mean().item():.8f} " + f"official_peak_bytes={official_peak} native_peak_bytes={native_peak}") + _report_first_block_drift(official_blocks, native_blocks) + assert_close(native_output, official_output, atol=0.0, rtol=0.0) diff --git a/tests/local_tests/vaes/test_lingbot_video_vae_parity.py b/tests/local_tests/vaes/test_lingbot_video_vae_parity.py new file mode 100644 index 000000000..75b5753d3 --- /dev/null +++ b/tests/local_tests/vaes/test_lingbot_video_vae_parity.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LingBot-Video Wan VAE parity against the released Diffusers component. + +Coverage scope: implementation_subcomponent. The released Dense and MoE roots +share the same VAE checkpoint. The test uses the exact production architecture +config and compares deterministic encode-mode and decode tensors. +""" + +from __future__ import annotations + +import json +import os +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 WanVAEConfig +from fastvideo.models.vaes.wanvae import AutoencoderKLWan as FastVideoAutoencoderKLWan +from tests.local_tests.lingbot_video.hf_assets import ( + FASTVIDEO_DENSE, + OFFICIAL_DENSE, + download_components, +) + +PARITY_SCOPE = "implementation_subcomponent" + + +def _load_fastvideo_vae(vae_dir: Path, device: torch.device) -> FastVideoAutoencoderKLWan: + """Strict-load the released VAE weights into FastVideo's native class.""" + config_path = vae_dir / "config.json" + weights_path = vae_dir / "diffusion_pytorch_model.safetensors" + raw_config = json.loads(config_path.read_text()) + config = WanVAEConfig() + config.update_model_arch({key: value for key, value in raw_config.items() if not key.startswith("_")}) + config.load_encoder = True + config.load_decoder = True + model = FastVideoAutoencoderKLWan(config) + model.load_state_dict(load_file(weights_path), strict=True) + return model.to(device=device, dtype=torch.float32).eval() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="LingBot-Video VAE parity requires CUDA.") +def test_lingbot_video_vae_decode_matches_official() -> None: + """Decode the same normalized latent through both VAE implementations.""" + if os.environ.get("LINGBOT_VIDEO_RUN_GPU_TESTS") != "1": + pytest.skip("Set LINGBOT_VIDEO_RUN_GPU_TESTS=1 on a scheduled GPU node.") + device = torch.device("cuda:0") + official_root = download_components(OFFICIAL_DENSE, "vae") + fastvideo_root = download_components(FASTVIDEO_DENSE, "vae") + official = OfficialAutoencoderKLWan.from_pretrained(official_root / "vae", + torch_dtype=torch.float32).to(device).eval() + fastvideo = _load_fastvideo_vae(fastvideo_root / "vae", device) + generator = torch.Generator(device=device).manual_seed(42) + latents = torch.randn((1, 16, 3, 8, 8), generator=generator, device=device, dtype=torch.float32) + mean = torch.tensor(official.config.latents_mean, device=device).view(1, 16, 1, 1, 1) + std = torch.tensor(official.config.latents_std, device=device).view(1, 16, 1, 1, 1) + decoder_latents = latents * std + mean + + with torch.inference_mode(): + official_output = official.decode(decoder_latents, return_dict=False)[0].float() + fastvideo_output = fastvideo.decode(decoder_latents).float() + assert_close(fastvideo_output, official_output, atol=5e-2, rtol=5e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="LingBot-Video VAE parity requires CUDA.") +def test_lingbot_video_vae_encode_mode_matches_official() -> None: + """Compare deterministic posterior modes before pipeline normalization.""" + if os.environ.get("LINGBOT_VIDEO_RUN_GPU_TESTS") != "1": + pytest.skip("Set LINGBOT_VIDEO_RUN_GPU_TESTS=1 on a scheduled GPU node.") + device = torch.device("cuda:0") + official_root = download_components(OFFICIAL_DENSE, "vae") + fastvideo_root = download_components(FASTVIDEO_DENSE, "vae") + official = OfficialAutoencoderKLWan.from_pretrained(official_root / "vae", + torch_dtype=torch.float32).to(device).eval() + fastvideo = _load_fastvideo_vae(fastvideo_root / "vae", device) + generator = torch.Generator(device=device).manual_seed(43) + video = torch.randn((1, 3, 9, 64, 64), generator=generator, device=device, dtype=torch.float32) + + with torch.inference_mode(): + official_mode = official.encode(video).latent_dist.mode().float() + fastvideo_mode = fastvideo.encode(video).mode().float() + assert_close(fastvideo_mode, official_mode, atol=5e-2, rtol=5e-2)