From 4b3f99b224fb05b33df9783e5636f75d4ac3ac52 Mon Sep 17 00:00:00 2001 From: Ishan <34205320+Ishxn20@users.noreply.github.com> Date: Wed, 7 Oct 2026 03:32:38 -0700 Subject: [PATCH] [new-model] Added LingBot-World-Fast image-to-video support (#1665) Co-authored-by: Claude Sonnet 5 Co-authored-by: SolitaryThinker --- .buildkite/scripts/unit_test.sh | 1 + docs/inference/support_matrix.md | 1 + .../basic/basic_lingbotworld_fast.py | 38 +++ fastvideo/configs/models/dits/__init__.py | 2 + .../configs/models/dits/lingbotworld_fast.py | 34 +++ fastvideo/configs/pipelines/__init__.py | 4 +- .../configs/pipelines/lingbotworld_fast.py | 23 ++ .../models/dits/lingbotworld2/causal_fast.py | 24 +- fastvideo/models/registry.py | 7 + .../basic/lingbotworld_fast/__init__.py | 5 + .../basic/lingbotworld_fast/fast_pipeline.py | 38 +++ .../basic/lingbotworld_fast/presets.py | 31 +++ fastvideo/registry.py | 30 ++- fastvideo/tests/api/test_presets.py | 42 ++++ .../test_lingbotworld_fast_backend.py | 30 +++ .../ssim/test_lingbot_fast_similarity.py | 199 +++++++++++++++ .../lingbotworld_fast/PORT_STATUS.md | 64 +++++ tests/local_tests/lingbotworld_fast/README.md | 103 ++++++++ ...st_lingbotworld_fast_transformer_parity.py | 232 ++++++++++++++++++ 19 files changed, 901 insertions(+), 7 deletions(-) create mode 100644 examples/inference/basic/basic_lingbotworld_fast.py create mode 100644 fastvideo/configs/models/dits/lingbotworld_fast.py create mode 100644 fastvideo/configs/pipelines/lingbotworld_fast.py create mode 100644 fastvideo/pipelines/basic/lingbotworld_fast/__init__.py create mode 100644 fastvideo/pipelines/basic/lingbotworld_fast/fast_pipeline.py create mode 100644 fastvideo/pipelines/basic/lingbotworld_fast/presets.py create mode 100644 fastvideo/tests/attention/test_lingbotworld_fast_backend.py create mode 100644 fastvideo/tests/ssim/test_lingbot_fast_similarity.py create mode 100644 tests/local_tests/lingbotworld_fast/PORT_STATUS.md create mode 100644 tests/local_tests/lingbotworld_fast/README.md create mode 100644 tests/local_tests/lingbotworld_fast/test_lingbotworld_fast_transformer_parity.py diff --git a/.buildkite/scripts/unit_test.sh b/.buildkite/scripts/unit_test.sh index 35ada983f..09fed9e27 100644 --- a/.buildkite/scripts/unit_test.sh +++ b/.buildkite/scripts/unit_test.sh @@ -21,6 +21,7 @@ exec pytest \ ./fastvideo/tests/training/test_trackers.py \ ./fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py \ ./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py \ + ./fastvideo/tests/attention/test_lingbotworld_fast_backend.py \ ./fastvideo/tests/attention/test_vsa_h3_tile_grad_safety.py \ ./fastvideo/tests/attention/test_vsa_h3_metadata.py \ ./fastvideo/tests/attention/test_vsa_h3_ref2va_regions.py \ diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index 9eacb82b7..a4adcbeae 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -61,6 +61,7 @@ column links a runnable script in `examples/inference/basic/` where one exists. | lingbot_video | `FastVideo/LingBot-Video-Dense-1.3B-Diffusers` | T2V | [basic_lingbot_video.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbot_video.py) | | lingbotworld | `FastVideo/LingBot-World-Base-Cam-Diffusers` | I2V | [basic_lingbotworld_base_cam.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbotworld_base_cam.py) | | lingbotworld2 | `robbyant/lingbot-world-v2-14b-causal-fast` | I2V | [basic_lingbotworld2_causal_fast.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbotworld2_causal_fast.py) | +| lingbotworld_fast | `FastVideo/LingBot-World-Fast-Diffusers` | I2V | [basic_lingbotworld_fast.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_lingbotworld_fast.py) | | longcat | `FastVideo/LongCat-Video-T2V-Diffusers` | T2V | [basic_longcat_t2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_longcat_t2v.py) | | longcat | `FastVideo/LongCat-Video-I2V-Diffusers` | I2V | [basic_longcat_i2v.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_longcat_i2v.py) | | longcat | `FastVideo/LongCat-Video-VC-Diffusers` | — | [basic_longcat_vc.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_longcat_vc.py) | diff --git a/examples/inference/basic/basic_lingbotworld_fast.py b/examples/inference/basic/basic_lingbotworld_fast.py new file mode 100644 index 000000000..4d5260fbb --- /dev/null +++ b/examples/inference/basic/basic_lingbotworld_fast.py @@ -0,0 +1,38 @@ +from fastvideo import VideoGenerator + +OUTPUT_PATH = "video_samples_lingbotworld_fast" + + +def main(): + generator = VideoGenerator.from_pretrained( + "FastVideo/LingBot-World-Fast-Diffusers", + num_gpus=1, + use_fsdp_inference=False, # set to True if GPU is out of memory + dit_cpu_offload=True, + vae_cpu_offload=False, + text_encoder_cpu_offload=True, + pin_cpu_memory=True, + ) + + prompt = ( + "A serene lakeside scene with a lone tree standing in calm water, " + "surrounded by distant snow-capped mountains under a bright blue sky " + "with drifting white clouds — gentle ripples reflect the tree and sky, " + "creating a tranquil, meditative atmosphere.") + action_path = "examples/datasets/lingbotworld2" + image_path = f"{action_path}/image.jpg" + + generator.generate_video( + prompt, + image_path=image_path, + action_path=action_path, + output_path=OUTPUT_PATH, + save_video=True, + num_frames=81, + height=480, + width=832, + ) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 8dbc32690..50b13c299 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -21,6 +21,7 @@ from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig from fastvideo.configs.models.dits.kandinsky6 import Kandinsky6ArchConfig, Kandinsky6VideoAudioConfig from fastvideo.configs.models.dits.kandinsky6_sr import Kandinsky6SRArchConfig, Kandinsky6SRConfig from fastvideo.configs.models.dits.lingbotworld2 import LingBotWorld2CausalFastVideoConfig +from fastvideo.configs.models.dits.lingbotworld_fast import LingBotWorldFastVideoConfig from fastvideo.configs.models.dits.lingbot_video import LingBotVideoConfig __all__ = [ @@ -29,5 +30,6 @@ __all__ = [ "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "Kandinsky6ArchConfig", "Kandinsky6VideoAudioConfig", "Kandinsky6SRArchConfig", "Kandinsky6SRConfig", "MagiHumanVideoConfig", "StableAudioConfig", "GlmImageDiTConfig", "HeliosConfig", "LingBotWorld2CausalFastVideoConfig", + "LingBotWorldFastVideoConfig", "LingBotVideoConfig", "MiniMaxH3Config", "ZImageDiTConfig", "MMAudioArchConfig", "MMAudioTransformerConfig" ] diff --git a/fastvideo/configs/models/dits/lingbotworld_fast.py b/fastvideo/configs/models/dits/lingbotworld_fast.py new file mode 100644 index 000000000..5ca7f831a --- /dev/null +++ b/fastvideo/configs/models/dits/lingbotworld_fast.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig +from fastvideo.configs.models.dits.lingbotworld2 import ( + LingBotWorld2CausalFastArchConfig, ) + + +@dataclass +class LingBotWorldFastArchConfig(LingBotWorld2CausalFastArchConfig): + """Arch config for the released LingBot-World-Fast checkpoint. + + The tensor layout is identical to the LingBot World 2 causal-fast model, so + the parent's shapes and ``param_names_mapping`` are reused verbatim. Only + the sampling/attention-window values released with this checkpoint differ. + """ + + # The released `generate_fast.py` leaves `--local_attn_size` at -1, so + # self-attention stays global and the KV cache never evicts. `sink_size` is + # therefore unreachable, but it is pinned to the official constructor's + # value so the two implementations stay identical if the window is enabled. + local_attn_size: int = -1 + sink_size: int = 9 + # `chunk_size` and `timesteps_index` are absent from the checkpoint config, + # so these values are what the sampling loop actually runs with. + chunk_size: int = 3 + timesteps_index: tuple[int, int, int, int] = (0, 179, 358, 679) + + +@dataclass +class LingBotWorldFastVideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=LingBotWorldFastArchConfig) + + prefix: str = "Wan" diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index 0d586e573..aaf732a7f 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -11,6 +11,7 @@ from fastvideo.configs.pipelines.kandinsky5 import Kandinsky5DMDConfig, Kandinsk from fastvideo.configs.pipelines.kandinsky6 import Kandinsky6TI2VAConfig from fastvideo.configs.pipelines.kandinsky6_sr import Kandinsky6SRPipelineConfig from fastvideo.configs.pipelines.lingbotworld2 import LingBotWorld2CausalFastI2V480PConfig +from fastvideo.configs.pipelines.lingbotworld_fast import LingBotWorldFastI2V480PConfig from fastvideo.configs.pipelines.lingbot_video import LingBotVideoT2VConfig from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig @@ -26,6 +27,7 @@ __all__ = [ "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HeliosPipelineConfig", "HYWorldConfig", "Kandinsky5T2VConfig", "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "Kandinsky6TI2VAConfig", - "Kandinsky6SRPipelineConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotVideoT2VConfig", + "Kandinsky6SRPipelineConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotWorldFastI2V480PConfig", + "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig", "get_pipeline_config_cls_from_name" ] diff --git a/fastvideo/configs/pipelines/lingbotworld_fast.py b/fastvideo/configs/pipelines/lingbotworld_fast.py new file mode 100644 index 000000000..e4e10169c --- /dev/null +++ b/fastvideo/configs/pipelines/lingbotworld_fast.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models import DiTConfig, EncoderConfig +from fastvideo.configs.models.dits.lingbotworld_fast import ( + LingBotWorldFastVideoConfig, ) +from fastvideo.configs.models.encoders import T5Config +from fastvideo.configs.pipelines.lingbotworld2 import ( + LingBotWorld2CausalFastI2V480PConfig, ) + + +@dataclass +class LingBotWorldFastI2V480PConfig(LingBotWorld2CausalFastI2V480PConfig): + """Pipeline config for LingBot-World-Fast 480P image-to-video. + + Shares the LingBot World 2 causal-fast sampling loop and Wan VAE, but this + checkpoint ships the stock ``UMT5EncoderModel`` text encoder (``d_model`` + fields) rather than LingBot World 2's custom one (``dim`` fields), so the + standard ``T5Config`` is restored here alongside this DiT's arch config. + """ + + dit_config: DiTConfig = field(default_factory=LingBotWorldFastVideoConfig) + text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (T5Config(), )) diff --git a/fastvideo/models/dits/lingbotworld2/causal_fast.py b/fastvideo/models/dits/lingbotworld2/causal_fast.py index e3bac555b..b7d2fea65 100644 --- a/fastvideo/models/dits/lingbotworld2/causal_fast.py +++ b/fastvideo/models/dits/lingbotworld2/causal_fast.py @@ -10,6 +10,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from fastvideo.attention.selector import backend_name_to_enum, get_attn_backend from fastvideo.configs.models.dits.lingbotworld2 import ( LingBotWorld2CausalFastVideoConfig, ) from fastvideo.distributed.communication_op import ( @@ -19,6 +20,7 @@ from fastvideo.distributed.communication_op import ( from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size from fastvideo.models.dits.base import BaseDiT from fastvideo.platforms import AttentionBackendEnum +from fastvideo.utils import get_compute_dtype try: import flash_attn_interface @@ -151,6 +153,7 @@ def attention( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + backend: AttentionBackendEnum, q_lens: torch.Tensor | None = None, k_lens: torch.Tensor | None = None, dropout_p: float = 0.0, @@ -162,8 +165,8 @@ def attention( dtype: torch.dtype = torch.bfloat16, fa_version: int | None = None, ) -> torch.Tensor: - """Dispatch LingBot World 2 attention to FlashAttention when available.""" - if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE: + """Dispatch LingBot World 2 attention through the selected backend.""" + if backend == AttentionBackendEnum.FLASH_ATTN: return flash_attention( q=q, k=k, @@ -180,6 +183,7 @@ def attention( version=fa_version, ) + assert backend == AttentionBackendEnum.TORCH_SDPA if q_lens is not None or k_lens is not None: warnings.warn("Padding masks are disabled without FlashAttention.") q = q.transpose(1, 2).to(dtype) @@ -262,6 +266,14 @@ class CausalWanSelfAttention(nn.Module): self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads + attn_backend = get_attn_backend( + self.head_dim, + get_compute_dtype(), + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), + ) + backend_name = attn_backend.get_name() + self.backend = backend_name_to_enum(backend_name) + assert self.backend is not None, f"Unsupported attention backend: {backend_name}" self.local_attn_size = local_attn_size self.sink_size = sink_size self.qk_norm = qk_norm @@ -348,7 +360,7 @@ class CausalWanSelfAttention(nn.Module): k_cache = kv_cache["k"][:, max(0, local_end_index - max_attention_size):local_end_index] v_cache = kv_cache["v"][:, max(0, local_end_index - max_attention_size):local_end_index] - x = attention(roped_query, k_cache, v_cache) + x = attention(roped_query, k_cache, v_cache, self.backend) kv_cache["global_end_index"].fill_(current_end) kv_cache["local_end_index"].fill_(local_end_index) @@ -389,7 +401,11 @@ class WanCrossAttention(CausalWanSelfAttention): else: k = self.norm_k(self.k(context)).view(b, -1, n, d) v = self.v(context).view(b, -1, n, d) - x = flash_attention(q, k, v, k_lens=context_lens) + # Dispatch through `attention` so the TORCH_SDPA backend this model + # advertises actually works; self-attention already does the same. The + # SDPA path ignores `k_lens`, which is safe here because the caller + # always passes `context_lens=None`. + x = attention(q, k, v, self.backend, k_lens=context_lens) return self.o(x.flatten(2)) diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 31303f352..0f54accde 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -64,6 +64,13 @@ _IMAGE_TO_VIDEO_DIT_MODELS = { "lingbotworld2", "LingBotWorld2CausalFastTransformer3DModel", ), + # LingBot-World-Fast ships this class name; its tensor layout is identical + # to the LingBot World 2 causal-fast DiT. + "CausalLingBotWorldTransformer3DModel": ( + "dits", + "lingbotworld2", + "LingBotWorld2CausalFastTransformer3DModel", + ), "MatrixGame2WanModel": ("dits", "matrixgame2", "MatrixGame2WanModel"), "CausalMatrixGame2WanModel": ("dits", "matrixgame2", "CausalMatrixGame2WanModel"), # Legacy aliases for older HF model_index.json files diff --git a/fastvideo/pipelines/basic/lingbotworld_fast/__init__.py b/fastvideo/pipelines/basic/lingbotworld_fast/__init__.py new file mode 100644 index 000000000..bbc83685a --- /dev/null +++ b/fastvideo/pipelines/basic/lingbotworld_fast/__init__.py @@ -0,0 +1,5 @@ +from .fast_pipeline import LingBotWorldFastPipeline + +__all__ = ["LingBotWorldFastPipeline"] + +EntryClass = LingBotWorldFastPipeline diff --git a/fastvideo/pipelines/basic/lingbotworld_fast/fast_pipeline.py b/fastvideo/pipelines/basic/lingbotworld_fast/fast_pipeline.py new file mode 100644 index 000000000..037a6755b --- /dev/null +++ b/fastvideo/pipelines/basic/lingbotworld_fast/fast_pipeline.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LingBot-World-Fast causal image-to-video pipeline.""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler, ) +from fastvideo.pipelines.basic.lingbotworld2.causal_fast_pipeline import ( + LingBotWorld2CausalFastPipeline, ) + +logger = init_logger(__name__) + + +class LingBotWorldFastPipeline(LingBotWorld2CausalFastPipeline): + """LingBot-World-Fast I2V generation. + + The released checkpoint uses the same chunked causal sampling loop as + LingBot World 2 causal-fast; the chunk size, timestep indices, and + attention window come from ``LingBotWorldFastArchConfig``. + """ + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: + """Install the flow-matching scheduler the released model samples with. + + This checkpoint ships a stock ``UniPCMultistepScheduler``, whose + ``set_timesteps`` takes no ``shift``. The reference ``generate_fast.py`` + builds a ``FlowUniPCMultistepScheduler`` instead, so replace the loaded + one unconditionally rather than only when absent. + """ + arch_config = fastvideo_args.pipeline_config.dit_config.arch_config + self.modules["scheduler"] = FlowUniPCMultistepScheduler( + num_train_timesteps=arch_config.num_train_timesteps, + shift=1, + use_dynamic_shifting=False, + ) + + +EntryClass = LingBotWorldFastPipeline diff --git a/fastvideo/pipelines/basic/lingbotworld_fast/presets.py b/fastvideo/pipelines/basic/lingbotworld_fast/presets.py new file mode 100644 index 000000000..fbcfa88bb --- /dev/null +++ b/fastvideo/pipelines/basic/lingbotworld_fast/presets.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +"""LingBot-World-Fast pipeline preset.""" + +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +_DENOISE_STAGE = PresetStageSpec( + name="denoise", + kind="denoising", + description="Causal-fast denoising pass", +) + +LINGBOTWORLD_FAST_I2V = InferencePreset( + name="lingbotworld_fast_i2v", + version=1, + model_family="lingbotworld_fast", + description="LingBot-World-Fast 14B causal I2V", + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "guidance_scale": 1.0, + "num_inference_steps": 4, + "fps": 16, + "seed": 42, + "num_frames": 81, + "height": 480, + "width": 832, + "negative_prompt": "", + }, +) + +ALL_PRESETS = (LINGBOTWORLD_FAST_I2V, ) diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 5487d9aa4..096f5b4d0 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -35,6 +35,7 @@ from fastvideo.configs.pipelines.kandinsky6_sr import Kandinsky6SRPipelineConfig 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.lingbotworld_fast import LingBotWorldFastI2V480PConfig from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig from fastvideo.configs.pipelines.flux_2 import ( @@ -600,6 +601,27 @@ def _register_configs() -> None: default_preset="lingbotworld2_causal_fast_i2v", ) + # LingBotWorld-Fast — registered BEFORE the LingBotWorld base entry so its + # detector wins when both fire. The base detector only excludes the + # "causal-fast" spelling, so it also matches this checkpoint's path and its + # `lingbotworldcausaldmdpipeline` model_index name; the detector loop keeps + # the first match, which must be this more specific entry. + register_configs( + sampling_param_cls=None, + pipeline_config_cls=LingBotWorldFastI2V480PConfig, + workload_types=(WorkloadType.I2V, ), + hf_model_paths=[ + "FastVideo/LingBot-World-Fast-Diffusers", + ], + model_detectors=[ + lambda path: ("lingbot-world-fast" in path.lower() or "lingbotworldfast" in path.lower() or + "lingbotworldcausaldmdpipeline" in path.lower()) + ], + model_family="lingbotworld_fast", + default_preset="lingbotworld_fast_i2v", + pipeline_cls_name="LingBotWorldFastPipeline", + ) + # LingBotWorld register_configs( sampling_param_cls=None, @@ -609,8 +631,9 @@ def _register_configs() -> None: "FastVideo/LingBot-World-Base-Cam-Diffusers", ], model_detectors=[ - lambda path: (("lingbotworld" in path.lower() or "lingbot-world" in path.lower()) and "causal-fast" not in - path.lower() and "causalfast" not in path.lower()) + lambda path: + (("lingbotworld" in path.lower() or "lingbot-world" in path.lower()) and "causal-fast" not in path.lower() + and "causalfast" not in path.lower() and "-fast" not in path.lower() and "causaldmd" not in path.lower()) ], model_family="lingbotworld", default_preset="lingbotworld_i2v", @@ -1328,6 +1351,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.lingbotworld_fast.presets import ( + ALL_PRESETS as LINGBOTWORLD_FAST_PRESETS, ) from fastvideo.pipelines.basic.lingbot_video.presets import ( ALL_PRESETS as LINGBOT_VIDEO_PRESETS, ) from fastvideo.pipelines.basic.longcat.presets import ( @@ -1373,6 +1398,7 @@ def _register_presets() -> None: LINGBOT_VIDEO_PRESETS, LINGBOTWORLD_PRESETS, LINGBOTWORLD2_PRESETS, + LINGBOTWORLD_FAST_PRESETS, LONGCAT_PRESETS, LTX2_PRESETS, MAGI_HUMAN_PRESETS, diff --git a/fastvideo/tests/api/test_presets.py b/fastvideo/tests/api/test_presets.py index 39fb9edf2..75ee31113 100644 --- a/fastvideo/tests/api/test_presets.py +++ b/fastvideo/tests/api/test_presets.py @@ -715,6 +715,48 @@ class TestLingBotVideoPresets: validate_stage_overrides(preset, {"refine": {"refiner_steps": 8}}) +# ------------------------------------------------------------------- +# LingBot-World-Fast preset integration +# ------------------------------------------------------------------- + + +class TestLingBotWorldFastPresets: + + def test_registered_defaults_are_fixed(self) -> None: + import fastvideo.registry # noqa: F401 + from fastvideo.registry import get_preset_selection + + model_path = "FastVideo/LingBot-World-Fast-Diffusers" + assert get_preset_selection(model_path) == ("lingbotworld_fast_i2v", "lingbotworld_fast") + + preset = get_preset("lingbotworld_fast_i2v", "lingbotworld_fast") + assert preset.defaults["num_inference_steps"] == 4 + assert preset.defaults["guidance_scale"] == 1.0 + assert preset.stage_schemas[0].allowed_overrides == frozenset() + with pytest.raises(ConfigValidationError, match="does not accept overrides"): + validate_stage_overrides(preset, {"denoise": {"num_inference_steps": 8}}) + + def test_local_checkpoint_resolves_fast_pipeline(self, tmp_path) -> None: + from fastvideo.configs.pipelines.lingbotworld_fast import LingBotWorldFastI2V480PConfig + from fastvideo.fastvideo_args import WorkloadType + from fastvideo.pipelines.basic.lingbotworld_fast import LingBotWorldFastPipeline + from fastvideo.registry import get_model_info + + model_dir = tmp_path / "lingbot-world-fast" + (model_dir / "transformer").mkdir(parents=True) + (model_dir / "model_index.json").write_text( + json.dumps({ + "_class_name": "LingBotWorldCausalDMDPipeline", + "_diffusers_version": "0.36.0", + }), + encoding="utf-8", + ) + + info = get_model_info(str(model_dir), workload_type=WorkloadType.I2V) + assert info.pipeline_cls is LingBotWorldFastPipeline + assert info.pipeline_config_cls is LingBotWorldFastI2V480PConfig + + # ------------------------------------------------------------------- # LingBotWorld preset integration (dual guidance) # ------------------------------------------------------------------- diff --git a/fastvideo/tests/attention/test_lingbotworld_fast_backend.py b/fastvideo/tests/attention/test_lingbotworld_fast_backend.py new file mode 100644 index 000000000..a8f7516e3 --- /dev/null +++ b/fastvideo/tests/attention/test_lingbotworld_fast_backend.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +import torch.nn.functional as F + +from fastvideo.models.dits.lingbotworld2 import causal_fast +from fastvideo.platforms import AttentionBackendEnum + + +class _SDPABackend: + + @staticmethod + def get_name() -> str: + return "TORCH_SDPA" + + +def test_lingbot_attention_honors_selected_sdpa_backend(monkeypatch) -> None: + monkeypatch.setattr(causal_fast, "get_attn_backend", lambda *args, **kwargs: _SDPABackend) + layer = causal_fast.CausalWanSelfAttention(dim=8, num_heads=2) + assert layer.backend == AttentionBackendEnum.TORCH_SDPA + + monkeypatch.setattr(causal_fast, "flash_attention", lambda **kwargs: pytest.fail("FlashAttention was called")) + q = torch.arange(24, dtype=torch.float32).reshape(1, 3, 2, 4) / 24 + k = q.flip(1) + v = q + 1 + + actual = causal_fast.attention(q, k, v, layer.backend, dtype=torch.float32) + expected = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)) + torch.testing.assert_close(actual, expected.transpose(1, 2).contiguous()) diff --git a/fastvideo/tests/ssim/test_lingbot_fast_similarity.py b/fastvideo/tests/ssim/test_lingbot_fast_similarity.py new file mode 100644 index 000000000..27bc26df4 --- /dev/null +++ b/fastvideo/tests/ssim/test_lingbot_fast_similarity.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +"""SSIM-based similarity test for LingBot-World-Fast causal I2V. + +The camera trajectory and source image are the tracked copies of official +LingBot example 03, matching the released `run_fast.sh` workflow. + +Note: this checkpoint is 4-step distilled, so the step count is fixed by the +model rather than reduced for CI. +""" + +import os + +import pytest +import torch + +import fastvideo.envs as envs +from fastvideo import VideoGenerator +from fastvideo.logger import init_logger +from fastvideo.tests.ssim.bootstrap_references import ( + xfail_missing_reference_in_bootstrap_mode, ) +from fastvideo.tests.ssim.reference_utils import ( + build_generated_output_dir, + build_reference_folder_path, + get_cuda_device_name, + resolve_device_reference_folder, + select_ssim_params, +) +from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results + +logger = init_logger(__name__) + +REQUIRED_GPUS = 2 + +# The released checkpoint is 4-step distilled; see LingBotWorldFastArchConfig. +NUM_DISTILLED_STEPS = 4 + + +def _find_lingbotworld_action_path() -> str | None: + script_dir = os.path.dirname(os.path.abspath(__file__)) + repo_root = os.path.abspath(os.path.join(script_dir, "..", "..", "..")) + candidate = os.path.join(repo_root, "examples", "datasets", "lingbotworld2") + if (os.path.exists(os.path.join(candidate, "image.jpg")) + and os.path.exists(os.path.join(candidate, "poses.npy")) + and os.path.exists(os.path.join(candidate, "intrinsics.npy"))): + return os.path.abspath(candidate) + return None + + +device_name = get_cuda_device_name() +device_reference_folder = resolve_device_reference_folder( + ( + ("A40", "A40"), + ("L40S", "L40S"), + ("H100", "H100"), + # GB200 must precede any bare "B200" pattern: the lookup is a + # substring scan and "B200" is a substring of "NVIDIA GB200". + ("GB200", "GB200"), + ("H200", "H200"), + ), + device_name=device_name, + logger=logger, +) + +LINGBOT_FAST_PARAMS = { + "model_path": "FastVideo/LingBot-World-Fast-Diffusers", + "num_gpus": 2, + "height": 480, + "width": 832, + "num_frames": 25, # must be 4k+1; trimmed to a whole number of chunks + "seed": 42, + "fps": 16, +} +LINGBOT_FAST_FULL_QUALITY_PARAMS = { + **LINGBOT_FAST_PARAMS, + "num_frames": 81, +} + +TEST_PROMPTS = [ + "A serene lakeside scene with a lone tree standing in calm water, surrounded " + "by distant snow-capped mountains under a bright blue sky with drifting white " + "clouds — gentle ripples reflect the tree and sky, creating a tranquil, " + "meditative atmosphere.", +] + + +@pytest.mark.parametrize("prompt", TEST_PROMPTS) +@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) +def test_lingbot_fast_i2v_similarity(prompt: str, ATTENTION_BACKEND: str, env_overrides): + env_overrides.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override(ATTENTION_BACKEND)) + + params = select_ssim_params(LINGBOT_FAST_PARAMS, + LINGBOT_FAST_FULL_QUALITY_PARAMS) + + if device_reference_folder is None: + pytest.skip( + f"Unsupported device for LingBot-World-Fast SSIM test: {device_name}" + ) + if torch.cuda.device_count() < REQUIRED_GPUS: + pytest.skip( + f"LingBot-World-Fast SSIM test requires {REQUIRED_GPUS} GPUs, " + f"but only {torch.cuda.device_count()} detected.") + + action_path = _find_lingbotworld_action_path() + if action_path is None: + pytest.skip( + "LingBot camera assets not found under examples/datasets/lingbotworld2." + ) + image_path = os.path.join(action_path, "image.jpg") + + script_dir = os.path.dirname(os.path.abspath(__file__)) + model_id = "LingBot-World-Fast-Diffusers" + output_dir = build_generated_output_dir( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, + ) + output_video_name = f"{prompt[:100].strip()}.mp4" + os.makedirs(output_dir, exist_ok=True) + + init_kwargs = { + "num_gpus": params["num_gpus"], + "use_fsdp_inference": True, + "dit_cpu_offload": True, + "dit_layerwise_offload": False, + "text_encoder_cpu_offload": True, + "vae_cpu_offload": False, + "pin_cpu_memory": True, + } + generation_kwargs = { + "output_path": output_dir, + "image_path": image_path, + "action_path": action_path, + "height": params["height"], + "width": params["width"], + "num_frames": params["num_frames"], + "seed": params["seed"], + "fps": params["fps"], + } + + generator: VideoGenerator | None = None + try: + generator = VideoGenerator.from_pretrained( + model_path=params["model_path"], **init_kwargs) + generator.generate_video(prompt, **generation_kwargs) + finally: + if generator is not None: + generator.shutdown() + + generated_video_path = os.path.join(output_dir, output_video_name) + assert os.path.exists(generated_video_path), ( + f"Output video was not generated at {generated_video_path}") + + reference_folder = build_reference_folder_path( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, + ) + if not os.path.exists(reference_folder): + xfail_missing_reference_in_bootstrap_mode( + generated_artifact_path=generated_video_path, + reference_folder=reference_folder, + artifact_kind="video", + ) + raise FileNotFoundError( + f"Reference video folder does not exist: {reference_folder}") + + reference_video_name = None + for filename in os.listdir(reference_folder): + if filename.endswith(".mp4") and prompt[:100].strip() in filename: + reference_video_name = filename + break + if not reference_video_name: + xfail_missing_reference_in_bootstrap_mode( + generated_artifact_path=generated_video_path, + reference_folder=reference_folder, + artifact_kind="video", + ) + raise FileNotFoundError( + f"Reference video missing for prompt/backend under {reference_folder}" + ) + + reference_video_path = os.path.join(reference_folder, reference_video_name) + logger.info("Computing SSIM between %s and %s", reference_video_path, + generated_video_path) + ssim_values = compute_video_ssim_torchvision(reference_video_path, + generated_video_path, + use_ms_ssim=True) + mean_ssim = ssim_values[0] + logger.info("SSIM mean value: %s", mean_ssim) + + write_ssim_results(output_dir, ssim_values, reference_video_path, + generated_video_path, NUM_DISTILLED_STEPS, prompt) + + min_acceptable_ssim = 0.70 + assert mean_ssim >= min_acceptable_ssim, ( + f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " + f"for {model_id} with backend {ATTENTION_BACKEND}") diff --git a/tests/local_tests/lingbotworld_fast/PORT_STATUS.md b/tests/local_tests/lingbotworld_fast/PORT_STATUS.md new file mode 100644 index 000000000..c1c387677 --- /dev/null +++ b/tests/local_tests/lingbotworld_fast/PORT_STATUS.md @@ -0,0 +1,64 @@ +# LingBot-World-Fast port status + +## Summary + +- model family: `lingbotworld_fast` +- workload: causal I2V with camera poses +- official source: `https://github.com/Robbyant/lingbot-world` +- official checkpoint: `robbyant/lingbot-world-fast` +- FastVideo bundle: `FastVideo/LingBot-World-Fast-Diffusers` +- source layout: official DiT checkpoint plus shared Wan components +- current phase: SSIM seeding + +## Reused implementation evidence + +The FastVideo variant reuses +`LingBotWorld2CausalFastTransformer3DModel`. The released transformer index has +1,421 parameter keys, and the PR's initial validation found all names and +shapes identical to the reused implementation. + +Variant-specific runtime values: + +| Setting | Value | Source | +|---|---:|---| +| `local_attn_size` | `-1` | released checkpoint config / `generate_fast.py` default | +| `chunk_size` | `3` | official `generate_fast.py` call | +| fixed timestep indices | `(0, 179, 358, 679)` | released fast sampling path | +| `sample_shift` | `10.0` | official I2V config | + +## Verification matrix + +| Scope | Command | Latest result | +|---|---|---| +| Official-vs-FastVideo full DiT | See `README.md` | 2026-08-17: 2 passed; max/mean/relative-mean error all `0.0` | +| Backend selection | `pytest fastvideo/tests/attention/test_lingbotworld_fast_backend.py -v` | 2026-08-17: passed | +| Preset and exact pipeline resolution | `pytest fastvideo/tests/api/test_presets.py::TestLingBotWorldFastPresets -v` | 2026-08-17: 2 passed | +| SSIM collection on this GB10 host | `pytest fastvideo/tests/ssim/test_lingbot_fast_similarity.py -v -s` | 2026-08-17: 1 skipped; unsupported GB10 | +| L40S reference generation and SSIM | `seed-ssim-references` workflow | blocked: Modal CLI and HF write token unavailable locally | + +## Acceptance requirements + +- [x] Official and FastVideo implementations load real released weights. +- [x] Both execute a real full-transformer forward. +- [x] Full-transformer numerical comparison passes without skip. +- [x] Backend and preset/registry tests pass. +- [ ] L40S SSIM reference is generated and visually approved. +- [ ] SSIM reference is uploaded and the package test passes against it. + +## Environment notes + +Local verification host: + +- GPU: NVIDIA GB10 +- detected GPU count: 1 +- PyTorch: 2.12.0 + CUDA 13.0 +- FlashAttention: unavailable; parity aligns both implementations on SDPA + +The SSIM suite does not accept GB10 references and the test requires two GPUs. +This machine can validate transformer parity and CPU contract tests, but it +cannot produce the required L40S SSIM acceptance evidence. + +## Open blockers + +1. Run the new SSIM test on Modal L40S, inspect the generated video, and seed + `FastVideo/ssim-reference-videos`. diff --git a/tests/local_tests/lingbotworld_fast/README.md b/tests/local_tests/lingbotworld_fast/README.md new file mode 100644 index 000000000..6dc4ae79c --- /dev/null +++ b/tests/local_tests/lingbotworld_fast/README.md @@ -0,0 +1,103 @@ +# LingBot-World-Fast local verification + +Local-only numerical parity and port evidence for +`FastVideo/LingBot-World-Fast-Diffusers`. + +## Scope + +- Official implementation: `Robbyant/lingbot-world`, `wan/modules/model_fast.py` +- Released fast checkpoint: `robbyant/lingbot-world-fast` +- FastVideo bundle: `FastVideo/LingBot-World-Fast-Diffusers` +- Workload: causal image-to-video with camera poses +- Released sampling contract: + - global self-attention (`local_attn_size=-1`) + - three latent frames per chunk + - four fixed scheduler indices: `(0, 179, 358, 679)` + - flow shift `10.0` + +## Setup + +Use the current FastVideo environment. The official source needs `easydict` in +addition to FastVideo's dependencies. + +```bash +git clone https://github.com/Robbyant/lingbot-world.git \ + /tmp/lingbot-world-reference +uv pip install --python .venv/bin/python easydict +``` + +Resolve a local snapshot of the FastVideo bundle: + +```bash +LINGBOTWORLD_FAST_MODEL_DIR=$( + .venv/bin/python - <<'PY' +from huggingface_hub import snapshot_download +print(snapshot_download("FastVideo/LingBot-World-Fast-Diffusers")) +PY +) +export LINGBOTWORLD_FAST_MODEL_DIR +export LINGBOTWORLD_FAST_REFERENCE_DIR=/tmp/lingbot-world-reference +``` + +Do not store Hugging Face tokens in this directory. + +## Full transformer parity + +The parity test loads the same released transformer weights into both: + +1. the official `WanModelFast` implementation; and +2. FastVideo's production `TransformerLoader`. + +It runs the models sequentially to stay within memory limits. Both sides use +the mathematical SDPA backend. The official cross-attention helper is routed +through the official source's own SDPA-capable dispatcher because the direct +helper only supports FlashAttention. + +The deterministic input is one real three-frame causal chunk with: + +- 16 noisy latent channels; +- 20 image-conditioning channels; +- 4096-wide text context; +- packed 384-channel camera rays; +- real per-layer self- and cross-attention caches. + +Run: + +```bash +FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA \ +LINGBOTWORLD_FAST_MODEL_DIR="$LINGBOTWORLD_FAST_MODEL_DIR" \ +LINGBOTWORLD_FAST_REFERENCE_DIR="$LINGBOTWORLD_FAST_REFERENCE_DIR" \ +.venv/bin/python -m pytest \ + tests/local_tests/lingbotworld_fast/test_lingbotworld_fast_transformer_parity.py \ + -v -s +``` + +This is valid evidence only when the full-transformer test reports `PASSED`, +not `SKIPPED`. + +## Lightweight runtime-contract tests + +These tests cover backend dispatch, the fixed-step preset contract, and exact +registry/pipeline resolution without loading the 18.5B transformer: + +```bash +.venv/bin/python -m pytest \ + fastvideo/tests/attention/test_lingbotworld_fast_backend.py \ + fastvideo/tests/api/test_presets.py::TestLingBotWorldFastPresets \ + -v +``` + +## SSIM + +The package test is: + +```bash +pytest fastvideo/tests/ssim/test_lingbot_fast_similarity.py -v -s +``` + +It requires two GPUs and device-specific references. New L40S references must +be generated with the repository's `seed-ssim-references` workflow, reviewed +visually, and uploaded to `FastVideo/ssim-reference-videos`. A GB10 output is +not a valid replacement for the L40S CI reference. + +See `PORT_STATUS.md` for the latest measured results and remaining blockers. diff --git a/tests/local_tests/lingbotworld_fast/test_lingbotworld_fast_transformer_parity.py b/tests/local_tests/lingbotworld_fast/test_lingbotworld_fast_transformer_parity.py new file mode 100644 index 000000000..e58262d17 --- /dev/null +++ b/tests/local_tests/lingbotworld_fast/test_lingbotworld_fast_transformer_parity.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Official-vs-FastVideo full-DiT parity for LingBot-World-Fast. + +Coverage scope: both. The official ``WanModelFast`` implementation and the +FastVideo production ``TransformerLoader`` load the same released transformer +weights, then run one real cached causal forward with the released global +attention and three-frame chunk settings. +""" + +from __future__ import annotations + +import gc +import importlib +import json +import os +from pathlib import Path +import sys + +import pytest +import torch +from torch.testing import assert_close + +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29531") + +DEFAULT_MODEL_DIR = Path("official_weights/FastVideo/LingBot-World-Fast-Diffusers") +DEFAULT_REFERENCE_DIR = Path("official_repos/lingbot-world") + + +def _model_dir() -> Path: + return Path(os.getenv("LINGBOTWORLD_FAST_MODEL_DIR", str(DEFAULT_MODEL_DIR))).expanduser().resolve() + + +def _reference_dir() -> Path: + return Path(os.getenv("LINGBOTWORLD_FAST_REFERENCE_DIR", str(DEFAULT_REFERENCE_DIR))).expanduser().resolve() + + +def _transformer_dir() -> Path: + return _model_dir() / "transformer" + + +def _has_real_weights() -> bool: + transformer_dir = _transformer_dir() + return transformer_dir.exists() and any(transformer_dir.glob("*.safetensors")) + + +def _official_constructor_kwargs() -> dict: + return { + "model_type": "i2v", + "control_type": "cam", + "patch_size": (1, 2, 2), + "text_len": 512, + "in_dim": 36, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "local_attn_size": -1, + "sink_size": 9, + "qk_norm": True, + "cross_attn_norm": True, + "eps": 1e-6, + } + + +def _deterministic_inputs(dtype: torch.dtype) -> dict: + generator = torch.Generator(device="cpu").manual_seed(20260407) + chunk_size = 3 + latent_height = 4 + latent_width = 4 + frame_seqlen = latent_height * latent_width // 4 + token_count = chunk_size * frame_seqlen + + def randn(*shape: int) -> torch.Tensor: + return torch.randn(shape, dtype=dtype, generator=generator) + + return { + "x": [randn(16, chunk_size, latent_height, latent_width)], + "t": torch.tensor([679.0], dtype=dtype), + "context": [randn(8, 4096)], + "seq_len": token_count, + "y": [randn(20, chunk_size, latent_height, latent_width)], + # The released pipeline packs each six-channel Plucker ray over the + # VAE's 8x8 spatial stride, yielding 6 * 8 * 8 = 384 channels. + "dit_cond_dict": { + "c2ws_plucker_emb": [randn(1, 384, chunk_size, latent_height, latent_width)], + }, + "current_start": 0, + "max_attention_size": token_count, + "frame_seqlen": frame_seqlen, + "cross_attn_first_call": True, + } + + +def _to_device(value, device: torch.device): + if torch.is_tensor(value): + return value.to(device) + if isinstance(value, list): + return [_to_device(item, device) for item in value] + if isinstance(value, dict): + return {key: _to_device(item, device) for key, item in value.items()} + return value + + +def _new_caches(device: torch.device, dtype: torch.dtype, token_count: int) -> tuple[list[dict], list[dict]]: + self_cache = [{ + "k": torch.zeros((1, token_count, 40, 128), device=device, dtype=dtype), + "v": torch.zeros((1, token_count, 40, 128), device=device, dtype=dtype), + "global_end_index": torch.tensor([0], device=device, dtype=torch.long), + "local_end_index": torch.tensor([0], device=device, dtype=torch.long), + } for _ in range(40)] + cross_cache = [{ + "k": torch.zeros((1, 512, 40, 128), device=device, dtype=dtype), + "v": torch.zeros((1, 512, 40, 128), device=device, dtype=dtype), + "is_init": torch.tensor(0, device=device, dtype=torch.int32), + } for _ in range(40)] + return self_cache, cross_cache + + +def _run_official(inputs: dict, dtype: torch.dtype) -> torch.Tensor: + reference_dir = _reference_dir() + sys.path.insert(0, str(reference_dir)) + try: + model_module = importlib.import_module("wan.modules.model_fast") + # The official cross-attention calls its FlashAttention-only helper + # directly. Route it through the official module's own SDPA-capable + # dispatcher so both implementations execute the same public fallback. + model_module.flash_attention = model_module.attention + model = model_module.WanModelFast.from_pretrained( + _transformer_dir(), + torch_dtype=dtype, + low_cpu_mem_usage=True, + device_map={"": "cuda"}, + **_official_constructor_kwargs(), + ).eval() + device_inputs = _to_device(inputs, torch.device("cuda")) + self_cache, cross_cache = _new_caches(torch.device("cuda"), dtype, inputs["seq_len"]) + with ( + torch.inference_mode(), + torch.amp.autocast("cuda", dtype=dtype), + torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH), + ): + output = model(**device_inputs, kv_cache=self_cache, crossattn_cache=cross_cache)[0] + return output.detach().float().cpu() + finally: + sys.path.remove(str(reference_dir)) + if "model" in locals(): + del model + gc.collect() + torch.cuda.empty_cache() + + +def _run_fastvideo(inputs: dict, dtype: torch.dtype) -> torch.Tensor: + from fastvideo.configs.models.dits.lingbotworld_fast import LingBotWorldFastVideoConfig + from fastvideo.configs.pipelines.base import PipelineConfig + from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel + from fastvideo.fastvideo_args import FastVideoArgs + from fastvideo.models.dits.lingbotworld2.causal_fast import LingBotWorld2CausalFastTransformer3DModel + from fastvideo.models.loader.component_loader import TransformerLoader + + precision = "bf16" if dtype == torch.bfloat16 else "fp16" + maybe_init_distributed_environment_and_model_parallel(1, 1) + args = FastVideoArgs( + model_path=str(_model_dir()), + dit_cpu_offload=False, + dit_layerwise_offload=False, + use_fsdp_inference=False, + pipeline_config=PipelineConfig( + dit_config=LingBotWorldFastVideoConfig(), + dit_precision=precision, + ), + ) + args.device = torch.device("cuda") + model = TransformerLoader().load(str(_transformer_dir()), args).eval() + assert isinstance(model, LingBotWorld2CausalFastTransformer3DModel) + device_inputs = _to_device(inputs, torch.device("cuda")) + self_cache, cross_cache = _new_caches(torch.device("cuda"), dtype, inputs["seq_len"]) + try: + with ( + torch.inference_mode(), + torch.amp.autocast("cuda", dtype=dtype), + torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH), + ): + output = model(**device_inputs, kv_cache=self_cache, crossattn_cache=cross_cache)[0] + return output.detach().float().cpu() + finally: + del model + gc.collect() + cleanup_dist_env_and_memory() + + +def test_lingbotworld_fast_checkpoint_config_matches_released_variant() -> None: + config_path = _transformer_dir() / "config.json" + if not config_path.exists(): + pytest.skip(f"LingBot-World-Fast transformer config not found at {config_path}") + config = json.loads(config_path.read_text(encoding="utf-8")) + assert config["_class_name"] == "CausalLingBotWorldTransformer3DModel" + assert config["local_attn_size"] == -1 + assert config["num_frames_per_block"] == 3 + assert config["num_layers"] == 40 + assert config["num_attention_heads"] == 40 + assert config["attention_head_dim"] == 128 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Full LingBot-World-Fast DiT parity requires CUDA") +def test_lingbotworld_fast_full_transformer_matches_official() -> None: + if not _reference_dir().exists(): + pytest.skip( + "Official LingBot source is absent; set LINGBOTWORLD_FAST_REFERENCE_DIR " + "to a Robbyant/lingbot-world checkout.") + if not _has_real_weights(): + pytest.skip( + "Released transformer weights are absent; set LINGBOTWORLD_FAST_MODEL_DIR " + "to FastVideo/LingBot-World-Fast-Diffusers.") + + dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + inputs = _deterministic_inputs(dtype) + official = _run_official(inputs, dtype) + fastvideo = _run_fastvideo(inputs, dtype) + + difference = (fastvideo - official).abs() + relative_mean_drift = difference.mean() / official.abs().mean().clamp_min(1e-6) + print( + "LingBot-World-Fast full-DiT parity: " + f"max={difference.max():.6f}, mean={difference.mean():.6f}, " + f"relative_mean={relative_mean_drift:.6f}") + assert relative_mean_drift < 0.05 + assert_close(fastvideo, official, atol=0.1, rtol=0.1)