From 63053a28fc9c5a13cfb21aceb3f14d842aa01bfd Mon Sep 17 00:00:00 2001 From: Jonathan McCaffrey Date: Sat, 5 Sep 2026 11:11:54 -0700 Subject: [PATCH 1/2] Add LongSANA Runtime V2 text-to-video integration --- .../recipes/wan/autoencoder/vae.py | 2 +- integrations_v2/longsana/README.md | 92 +++ integrations_v2/longsana/VALIDATION.md | 100 +++ integrations_v2/longsana/__init__.py | 4 + integrations_v2/longsana/apps/__init__.py | 4 + integrations_v2/longsana/apps/t2v/__init__.py | 4 + integrations_v2/longsana/apps/t2v/adapter.py | 74 +++ integrations_v2/longsana/config.py | 71 +++ integrations_v2/longsana/impl/__init__.py | 4 + integrations_v2/longsana/impl/constants.py | 63 ++ integrations_v2/longsana/impl/model.py | 594 ++++++++++++++++++ integrations_v2/longsana/impl/pipeline.py | 142 +++++ integrations_v2/longsana/impl/scheduler.py | 101 +++ integrations_v2/longsana/impl/transformer.py | 379 +++++++++++ integrations_v2/longsana/pyproject.toml | 42 ++ .../longsana/resources/longsana_text.yaml | 20 + integrations_v2/longsana/scripts/benchmark.py | 296 +++++++++ .../longsana/scripts/operator_profile.py | 95 +++ integrations_v2/longsana/tests/test_smoke.py | 328 ++++++++++ .../longsana/tests/test_t2v_app.py | 86 +++ pyproject.toml | 2 + uv.lock | 27 + 22 files changed, 2529 insertions(+), 1 deletion(-) create mode 100644 integrations_v2/longsana/README.md create mode 100644 integrations_v2/longsana/VALIDATION.md create mode 100644 integrations_v2/longsana/__init__.py create mode 100644 integrations_v2/longsana/apps/__init__.py create mode 100644 integrations_v2/longsana/apps/t2v/__init__.py create mode 100644 integrations_v2/longsana/apps/t2v/adapter.py create mode 100644 integrations_v2/longsana/config.py create mode 100644 integrations_v2/longsana/impl/__init__.py create mode 100644 integrations_v2/longsana/impl/constants.py create mode 100644 integrations_v2/longsana/impl/model.py create mode 100644 integrations_v2/longsana/impl/pipeline.py create mode 100644 integrations_v2/longsana/impl/scheduler.py create mode 100644 integrations_v2/longsana/impl/transformer.py create mode 100644 integrations_v2/longsana/pyproject.toml create mode 100644 integrations_v2/longsana/resources/longsana_text.yaml create mode 100644 integrations_v2/longsana/scripts/benchmark.py create mode 100644 integrations_v2/longsana/scripts/operator_profile.py create mode 100644 integrations_v2/longsana/tests/test_smoke.py create mode 100644 integrations_v2/longsana/tests/test_t2v_app.py diff --git a/flashdreams/flashdreams/recipes/wan/autoencoder/vae.py b/flashdreams/flashdreams/recipes/wan/autoencoder/vae.py index 6783f3d23..ebf93f7de 100644 --- a/flashdreams/flashdreams/recipes/wan/autoencoder/vae.py +++ b/flashdreams/flashdreams/recipes/wan/autoencoder/vae.py @@ -1496,7 +1496,7 @@ def forward( *batch_shape, T, C, H, W = input.shape batch_size = math.prod(batch_shape) - z = input.reshape(batch_size, T, C, H, W) + z = input.reshape(batch_size, T, C, H, W).to(dtype=self.config.dtype) x = self.vae.decode(z.transpose(1, 2), cache=cache).transpose(1, 2) return x.reshape(*batch_shape, *x.shape[1:]) diff --git a/integrations_v2/longsana/README.md b/integrations_v2/longsana/README.md new file mode 100644 index 000000000..0ba6c157c --- /dev/null +++ b/integrations_v2/longsana/README.md @@ -0,0 +1,92 @@ + + +# LongSANA + +LongSANA 2B is a causal, text-to-video SANA-Video integration for FlashDreams +Runtime V2. It uses the public self-forcing 480p checkpoint and keeps a +constant-size recurrent attention state instead of retaining every prior token. + +The implementation reuses the existing SANA-WM Gemma/CHI prompt encoder, +normalization, timestep/text projection, and cross-attention components, plus +the shared Wan 2.1 VAE decoder and Runtime V2 diffusion/session machinery. + +## Model + +| application slug | resolution | rate | default rollout | +| --- | ---: | ---: | ---: | +| `t2v-longsana-2b-480p` | 832 x 480 | 16 FPS | 1,041 frames (~65 s) | + +The first autoregressive block emits 41 pixel frames. Every later block emits +40. The released sampler uses four self-forcing steps at raw timesteps +`1000, 960, 889, 727`, flow shift 7, CFG 1, and motion score 10. + +## Install and run + +```bash +uv sync --package flashdreams-longsana --extra dev --group test --inexact +uv run --no-sync flashdreams-run-v2 \ + t2v-longsana-2b-480p \ + --timeout 30 \ + --output-path artifacts/longsana.mp4 -- \ + --prompt "A red panda walks through a misty bamboo forest at sunrise." \ + --total-blocks 2 --seed 0 --no-compile +``` + +The generator, Gemma text encoder, and Wan VAE download from Hugging Face on +first use. The application accepts only the checkpoint's native 832 x 480 +resolution. + +`--timeout` bounds the shared interactive T2V presentation session. Use the +benchmark command below for exact generated-frame clips; a timeout-bounded +Runtime V2 MP4 may repeat its latest frame while the UI remains active. + +## Constant-memory cache + +Each of the 20 transformer blocks owns three recurrent tensors: + +- a cumulative rotated `V @ K^T` matrix; +- a cumulative positive-key sum used by the linear-attention denominator; +- the final frame needed by the causal temporal convolution. + +The state is updated in-place after a clean-timestep forward pass. Its size is +independent of generated duration (152.53 MiB for batch size one at the native +configuration), and absolute temporal RoPE positions advance across blocks. + +## Benchmark and validation + +See [VALIDATION.md](VALIDATION.md) for measured results and qualitative review. + +Run the checked-in diverse-prompt suite: + +```bash +uv run --no-sync python integrations_v2/longsana/scripts/benchmark.py \ + --output-dir artifacts/longsana_benchmark --blocks 2 +``` + +Use `--long-blocks 6` to extend the temporal-continuity case while keeping the +other cases short. The command writes MP4 clips, contact sheets, per-block +Runtime V2 stage timings, GPU allocation/peak measurements, cache-size history, +and a machine-readable `summary.json`. + +Capture a PyTorch operator trace for one warmed steady-state block with: + +```bash +uv run --no-sync python integrations_v2/longsana/scripts/operator_profile.py \ + --output-dir artifacts/longsana_profile +``` + +The first block includes lazy generator loading and is reported separately. +Steady-state throughput excludes prompt encoding but includes diffusion, Wan +decode, and the cache-finalization pass. Compare identical resolution, frame +rate, block count, seed, and prompt set when benchmarking another DiT. + +## Tests + +```bash +uv run --no-sync pytest integrations_v2/longsana/tests -m ci_cpu +uv run --no-sync ruff check integrations_v2/longsana \ + flashdreams/flashdreams/recipes/wan/autoencoder/vae.py +``` diff --git a/integrations_v2/longsana/VALIDATION.md b/integrations_v2/longsana/VALIDATION.md new file mode 100644 index 000000000..50e2e5894 --- /dev/null +++ b/integrations_v2/longsana/VALIDATION.md @@ -0,0 +1,100 @@ + + +# LongSANA validation baseline + +Measured September 4, 2026 from base commit +`d9b2516720ff07a4a171459905f41672cebc4d41`. + +## Configuration + +- GPU: NVIDIA RTX PRO 6000 Blackwell Workstation Edition, 94.97 GiB +- Python 3.12.3, PyTorch 2.12.1+cu130, CUDA 13.0 +- Generator: public `LongSANA_2B_480p_self_forcing` revision + `48283a1b034cecdfaf412a01be2ae202d2432a85` +- Wan 2.1 VAE: SANA-Video revision + `7dd4f2fcddc7db57597238d728e1f430129827ff` +- Output: 832 x 480 at 16 FPS +- Sampler: four self-forcing steps, shift 7, raw timesteps + `1000, 960, 889, 727`, motion score 10 +- Precision: BF16 DiT with FP32 linear-attention accumulation and the official + FP32 VAE + +The exact machine-readable results are in +`artifacts/longsana_benchmark/summary.json`; the operator trace and table are +in `artifacts/longsana_profile/`. + +## Correctness checks + +- The complete 2,057,553,344-parameter generator schema strictly loads all 418 + released checkpoint tensors with no missing, unexpected, or shape-mismatched + keys. +- Unit tests cover release configuration, prompt formatting, checkpoint + extraction, absolute temporal RoPE, upstream sampler precision/noise order, + recurrent-state lifecycle, fixed storage, block boundaries, application + defaults, resolution rejection, and entry-point registration. +- The installed `t2v-longsana-2b-480p` Runtime V2 application completed a + one-block run and emitted model-step stats. Its timeout-bounded interactive + MP4 repeats the latest presentation frame, so it is a runner smoke artifact, + not one of the exact-frame quality clips below. +- All 565 frames generated by the five-case suite were finite and decoded at + the advertised shape. +- The six-block case advanced from 11 initial latent frames through five + 10-frame blocks, producing 241 pixel frames without a cache-size increase. + +## Performance + +Steady-state values exclude prompt encoding and the first block's lazy +checkpoint load. + +| measurement | result | +| --- | ---: | +| End-to-end throughput (median across cases) | 9.93 FPS | +| DiT + clean cache commit throughput (median) | 19.86 FPS | +| End-to-end block latency (40 frames, median) | 4,029 ms | +| DiT diffusion latency (four forwards) | about 1,620 ms | +| Clean cache-commit latency (one forward) | about 404 ms | +| Wan decode latency | about 2,015 ms | +| Resident allocated GPU memory after a block | about 8.07 GiB | +| Peak allocated GPU memory | 37.30 GiB | +| Recurrent transformer state | 152.53 MiB | + +The six-block run's steady latencies ranged from 4,045 to 4,078 ms. Cache +storage stayed at exactly 152.526855 MiB after every block and the underlying +state buffers are updated in-place. At the native 15,600 steady-state tokens, +the normalized DiT baseline is about 405 ms per forward (five forwards per +block including cache commit). + +The warmed operator trace confirms that the FP32 Wan decoder is the largest +single stage and its large 3D convolutions drive peak memory. In the DiT, +GLUMB's pointwise/depthwise/temporal convolutions, dense projections, RoPE +copies, and elementwise kernels dominate. This makes the stage-separated +`diffuse_ms` plus `finalize_ms` the useful number for comparing the +LongSANA DiT against another decoder-independent architecture. + +## Qualitative review + +| case | observation | +| --- | --- | +| animal motion | Correct red-panda/bamboo composition and stable identity; walking displacement is modest. | +| dance + camera | Correct three-person neon-alley composition; poses evolve coherently, but the requested camera orbit is weak. | +| cooking interaction | Strong wok, flame, food, and hand composition; ordered flip/pour/plate semantics are not fully expressed in the short clip. | +| coastal aerial | Stable coast, surf, storm, and lighthouse; continuous scene geometry with restrained forward travel. | +| surreal long (six blocks) | Astronaut, horse, lunar landscape, lighting, and composition remain unusually stable across 15 seconds; gait evolves but camera motion is limited. | + +These are manual contact-sheet observations, not a VBench score. The results +validate execution, prompt alignment, temporal identity, and constant-memory +behavior. They also expose a quality characteristic worth measuring in an +apples-to-apples study: this checkpoint preserves scenes and subjects better +than it follows aggressive camera directions or multi-step action verbs. + +## Comparison guidance + +For a fair DiT comparison, report at least model parameters, latent token count, +network forwards per output block, milliseconds per network forward, +`diffuse_ms + finalize_ms`, recurrent-state bytes, resident/peak VRAM, and +output frames per second. Keep the VAE out of the DiT number; LongSANA's +official FP32 Wan decoder accounts for roughly half of this end-to-end block +latency and most of its peak allocation. diff --git a/integrations_v2/longsana/__init__.py b/integrations_v2/longsana/__init__.py new file mode 100644 index 000000000..d115a444c --- /dev/null +++ b/integrations_v2/longsana/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""LongSana constant-memory text-to-video integration.""" diff --git a/integrations_v2/longsana/apps/__init__.py b/integrations_v2/longsana/apps/__init__.py new file mode 100644 index 000000000..3cd705098 --- /dev/null +++ b/integrations_v2/longsana/apps/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""LongSana Runtime V2 applications.""" diff --git a/integrations_v2/longsana/apps/t2v/__init__.py b/integrations_v2/longsana/apps/t2v/__init__.py new file mode 100644 index 000000000..bca51f194 --- /dev/null +++ b/integrations_v2/longsana/apps/t2v/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""LongSana text-to-video application.""" diff --git a/integrations_v2/longsana/apps/t2v/adapter.py b/integrations_v2/longsana/apps/t2v/adapter.py new file mode 100644 index 000000000..51793c9f0 --- /dev/null +++ b/integrations_v2/longsana/apps/t2v/adapter.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LongSana text-to-video Runtime V2 application.""" + +from __future__ import annotations + +import dataclasses +from typing import Any + +from t2v import T2VApplication, T2VApplicationDefaults + +from flashdreams.api_v2.application import IApplication +from flashdreams.runtime_v2.session_desc import SessionDesc +from longsana.config import PIPELINE_LONGSANA_2B_480P +from longsana.impl.constants import ( + DEFAULT_VIDEO_FPS, + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, +) + +LONGSANA_T2V_DEFAULTS = T2VApplicationDefaults( + pipeline_config=PIPELINE_LONGSANA_2B_480P, + total_blocks=26, + pixel_width=DEFAULT_VIDEO_WIDTH, + pixel_height=DEFAULT_VIDEO_HEIGHT, + fps=DEFAULT_VIDEO_FPS, +) +"""A 26-block rollout emits 1,041 frames, about 65 seconds at 16 FPS.""" + + +class LongSanaT2VApplication(T2VApplication): + """LongSana 2B constant-memory text-to-video application.""" + + def __init__(self, pipeline_config: Any | None = None) -> None: + """ + Args: + pipeline_config: Optional stand-in used by tests. + """ + defaults = LONGSANA_T2V_DEFAULTS + if pipeline_config is not None: + defaults = dataclasses.replace( + defaults, + pipeline_config=pipeline_config, + ) + super().__init__(defaults=defaults) + + def _validate_frame_size(self, session_desc: SessionDesc, pipeline: Any) -> None: + """Require the native 832 by 480 release dimensions.""" + del pipeline + requested = (session_desc.video_width, session_desc.video_height) + expected = (DEFAULT_VIDEO_WIDTH, DEFAULT_VIDEO_HEIGHT) + if requested != expected: + raise ValueError( + f"LongSana 2B 480p requires {expected[0]}x{expected[1]} output, " + f"got {requested[0]}x{requested[1]}." + ) + + +def create_app() -> IApplication: + """Return a new LongSana text-to-video application.""" + return LongSanaT2VApplication() diff --git a/integrations_v2/longsana/config.py b/integrations_v2/longsana/config.py new file mode 100644 index 000000000..df21ee012 --- /dev/null +++ b/integrations_v2/longsana/config.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Public LongSana Runtime V2 pipeline configuration.""" + +from __future__ import annotations + +import torch + +from flashdreams.infra.diffusion.model import DiffusionModelConfig +from flashdreams.recipes.wan.autoencoder.vae import WanVAEDecoderConfig +from longsana.impl.constants import ( + DEFAULT_DENOISING_TIMESTEPS, + LONGSANA_TEXT_CONFIG_PATH, + LONGSANA_VAE_CHECKPOINT_PATH, +) +from longsana.impl.pipeline import LongSanaPipelineConfig +from longsana.impl.scheduler import LongSanaFlowMatchSchedulerConfig +from longsana.impl.transformer import LongSanaTransformerConfig +from sana_wm.impl.conditioning import SanaWMTextPromptEncoderConfig + +PIPELINE_LONGSANA_2B_480P = LongSanaPipelineConfig( + name="longsana-2b-480p", + enable_sync_and_profile=True, + encoder=None, + prompt_encoder=SanaWMTextPromptEncoderConfig( + config_path=LONGSANA_TEXT_CONFIG_PATH, + offload_text_encoder=True, + ), + decoder=WanVAEDecoderConfig( + checkpoint_path=LONGSANA_VAE_CHECKPOINT_PATH, + dtype=torch.float32, + use_cuda_graph=False, + use_compile=False, + ), + diffusion_model=DiffusionModelConfig( + seed=0, + context_noise=0, + transformer=LongSanaTransformerConfig(), + scheduler=LongSanaFlowMatchSchedulerConfig( + num_inference_steps=4, + shift=7.0, + denoising_timesteps=list(DEFAULT_DENOISING_TIMESTEPS), + warp_denoising_step=False, + num_train_timesteps=1000, + sigma_max=1.0, + sigma_min=0.0, + extra_one_step=True, + timestep_dtype=torch.float32, + enable_tqdm=True, + ), + ), +) +"""Official four-step LongSana 2B 480p pipeline.""" + +LONGSANA_CONFIGS: dict[str, LongSanaPipelineConfig] = { + PIPELINE_LONGSANA_2B_480P.name: PIPELINE_LONGSANA_2B_480P, +} +"""All public LongSana pipeline configurations.""" diff --git a/integrations_v2/longsana/impl/__init__.py b/integrations_v2/longsana/impl/__init__.py new file mode 100644 index 000000000..d485ecbe8 --- /dev/null +++ b/integrations_v2/longsana/impl/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""LongSana model and Runtime V2 adapters.""" diff --git a/integrations_v2/longsana/impl/constants.py b/integrations_v2/longsana/impl/constants.py new file mode 100644 index 000000000..2e9f65521 --- /dev/null +++ b/integrations_v2/longsana/impl/constants.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Released LongSana model constants and immutable artifact locations.""" + +from pathlib import Path + +DEFAULT_VIDEO_HEIGHT = 480 +"""Native LongSana release height in pixels.""" + +DEFAULT_VIDEO_WIDTH = 832 +"""Native LongSana release width in pixels.""" + +DEFAULT_VIDEO_FPS = 16 +"""Frame rate used by the public LongSana release.""" + +FIRST_LATENT_BLOCK_FRAMES = 11 +"""First block size for the release's 261-frame latent rollout.""" + +LATENT_BLOCK_FRAMES = 10 +"""Steady-state number of latent frames generated per AR block.""" + +MOTION_SCORE = 10 +"""Motion-score suffix used during LongSana self-forcing post-training.""" + +LONGSANA_REVISION = "48283a1b034cecdfaf412a01be2ae202d2432a85" +"""Immutable Hugging Face revision used for the public distilled checkpoint.""" + +LONGSANA_CHECKPOINT_PATH = ( + "https://huggingface.co/Efficient-Large-Model/" + "LongSANA_2B_480p_self_forcing/resolve/" + f"{LONGSANA_REVISION}/checkpoints/LongSANA_2B_480p_self_forcing.pt" +) +"""Public four-step LongSana generator checkpoint.""" + +SANA_VIDEO_REVISION = "7dd4f2fcddc7db57597238d728e1f430129827ff" +"""Immutable SANA-Video revision containing the Wan 2.1 VAE.""" + +LONGSANA_VAE_CHECKPOINT_PATH = ( + "https://huggingface.co/Efficient-Large-Model/SANA-Video_2B_480p/resolve/" + f"{SANA_VIDEO_REVISION}/vae/Wan2.1_VAE.pth" +) +"""Wan 2.1 VAE checkpoint paired with the LongSana release.""" + +LONGSANA_TEXT_CONFIG_PATH = str( + Path(__file__).resolve().parent.parent / "resources" / "longsana_text.yaml" +) +"""Packaged text-encoder settings copied from the upstream release config.""" + +DEFAULT_DENOISING_TIMESTEPS = [1000, 960, 889, 727] +"""Four raw self-forcing timesteps from the public LongSana checkpoint.""" diff --git a/integrations_v2/longsana/impl/model.py b/integrations_v2/longsana/impl/model.py new file mode 100644 index 000000000..0cfdfff85 --- /dev/null +++ b/integrations_v2/longsana/impl/model.py @@ -0,0 +1,594 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Checkpoint-compatible LongSana DiT with a constant-size recurrent state.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +from flashdreams.infra.config import InstantiateConfig +from sana_wm.impl.stage1_model import ( + RMSNorm, + SanaWMStage1Spec, + Stage1CrossAttention, + TextEmbedder, + TimestepEmbedder, +) + +LONGSANA_SPEC = SanaWMStage1Spec( + latent_channels=16, + hidden_size=2240, + text_dim=2304, + timestep_dim=256, + depth=20, + num_heads=20, + head_dim=112, + max_text_length=300, + latent_grid_size=(30, 52), + mlp_ratio=3, + temporal_kernel_size=3, +) +"""Architecture of the released LongSana 2B 480p generator.""" + + +@dataclass(kw_only=True) +class LongSanaBlockState: + """Constant-size recurrent state for one LongSana transformer block.""" + + value_key: Tensor | None = None + """Cumulative rotated value-key product, shaped [B, H, D, D].""" + + key_sum: Tensor | None = None + """Cumulative unrotated positive key sum, shaped [B, H, 1, D].""" + + conv_tail: Tensor | None = None + """Last spatial-FFN frame, shaped [B, C, 1, H*W].""" + + def num_bytes(self) -> int: + """Return bytes currently occupied by this block's recurrent tensors.""" + return sum( + tensor.numel() * tensor.element_size() + for tensor in (self.value_key, self.key_sum, self.conv_tail) + if tensor is not None + ) + + +@dataclass(kw_only=True) +class LongSanaNetworkConfig(InstantiateConfig): + """Config for the released LongSana generator architecture.""" + + _target: type["LongSanaModel"] = field(default_factory=lambda: LongSanaModel) + + spec: SanaWMStage1Spec = LONGSANA_SPEC + """Checkpoint-facing dimensions. Tests may substitute a small stand-in.""" + + patch_size: tuple[int, int, int] = (1, 2, 2) + """Temporal, height, and width patch factors.""" + + fp32_attention: bool = True + """Accumulate the rotated linear-attention numerator in float32.""" + + +class _PatchEmbed3D(nn.Module): + """3D convolutional patch embedder with upstream-compatible parameter names.""" + + def __init__( + self, + in_channels: int, + hidden_size: int, + patch_size: tuple[int, int, int], + ) -> None: + super().__init__() + self.patch_size = patch_size + self.kernel_size = patch_size + self.proj = nn.Conv3d( + in_channels, + hidden_size, + kernel_size=patch_size, + stride=patch_size, + bias=True, + ) + + def forward(self, x: Tensor) -> Tensor: + """Project and flatten BCTHW latents into BND tokens.""" + return self.proj(x).flatten(2).transpose(1, 2) + + +class LongSanaLinearAttention(nn.Module): + """ReLU-kernel linear attention with absolute RoPE and recurrent sums.""" + + def __init__(self, spec: SanaWMStage1Spec, *, fp32_attention: bool) -> None: + super().__init__() + self.heads = spec.num_heads + self.dim = spec.head_dim + self.eps = 1e-8 + self.fp32_attention = fp32_attention + self.qkv = nn.Linear(spec.hidden_size, 3 * spec.hidden_size, bias=False) + self.q_norm = RMSNorm(spec.hidden_size, eps=1e-5) + self.k_norm = RMSNorm(spec.hidden_size, eps=1e-5) + self.proj = nn.Linear(spec.hidden_size, spec.hidden_size) + + def forward( + self, + x: Tensor, + *, + rotary_emb: Tensor, + state: LongSanaBlockState, + update_state: bool, + ) -> Tensor: + """Apply attention using prior state and optionally commit this block.""" + batch, tokens, channels = x.shape + if channels != self.heads * self.dim: + raise ValueError( + f"channels={channels} != heads*head_dim={self.heads * self.dim}." + ) + + qkv = self.qkv(x).reshape(batch, tokens, 3, channels) + q, k, v = qkv.unbind(dim=2) + dtype = q.dtype + + q = self.q_norm(q).transpose(-1, -2) + k = self.k_norm(k).transpose(-1, -2) + v = v.transpose(-1, -2) + q = F.relu(q.reshape(batch, self.heads, self.dim, tokens)) + k = F.relu(k.reshape(batch, self.heads, self.dim, tokens)) + v = v.reshape(batch, self.heads, self.dim, tokens) + + q_rotated = _apply_causal_rope(q, rotary_emb) + k_rotated = _apply_causal_rope(k, rotary_emb) + if self.fp32_attention: + q_rotated = q_rotated.float() + k_rotated = k_rotated.float() + v = v.float() + + current_key_sum = k.sum(dim=-1, keepdim=True).transpose(-2, -1) + current_value_key = torch.matmul(v, k_rotated.transpose(-1, -2)) + total_key_sum = current_key_sum + total_value_key = current_value_key + if state.value_key is not None or state.key_sum is not None: + if state.value_key is None or state.key_sum is None: + raise RuntimeError( + "LongSana attention state is only partially initialized." + ) + total_value_key = current_value_key + state.value_key + total_key_sum = current_key_sum + state.key_sum + + denominator = 1.0 / (total_key_sum @ q + self.eps) + out = torch.matmul(total_value_key, q_rotated) + out = (out * denominator).to(dtype) + out = self.proj(out.reshape(batch, channels, tokens).permute(0, 2, 1)) + + if update_state: + _set_or_copy(state, "value_key", total_value_key) + _set_or_copy(state, "key_sum", total_key_sum) + return out + + +class LongSanaCausalGLUMBConvTemp(nn.Module): + """SANA GLUMB feed-forward layer with one-frame causal temporal state.""" + + def __init__(self, spec: SanaWMStage1Spec) -> None: + super().__init__() + inner = spec.mlp_inner_size + gated = spec.gated_mlp_size + self.inverted_conv = _Conv2dContainer(spec.hidden_size, inner, 1) + self.depth_conv = _Conv2dContainer( + inner, + inner, + 3, + groups=inner, + padding=1, + ) + self.point_conv = _Conv2dContainer( + gated, + spec.hidden_size, + 1, + bias=False, + ) + self.t_conv = nn.Conv2d( + spec.hidden_size, + spec.hidden_size, + kernel_size=(spec.temporal_kernel_size, 1), + padding=(spec.temporal_kernel_size // 2, 0), + bias=False, + ) + + def forward( + self, + x: Tensor, + *, + frames: int, + height: int, + width: int, + state: LongSanaBlockState, + update_state: bool, + ) -> Tensor: + """Run spatial GLUMB and a causal temporal convolution.""" + batch, tokens, channels = x.shape + if tokens != frames * height * width: + raise ValueError( + f"tokens={tokens} != frames*height*width={frames * height * width}." + ) + x_2d = x.reshape(batch * frames, height, width, channels).permute(0, 3, 1, 2) + x_2d = F.silu(self.inverted_conv(x_2d), inplace=True) + x_2d = self.depth_conv(x_2d) + value, gate = x_2d.chunk(2, dim=1) + x_2d = self.point_conv(value * F.silu(gate)) + + spatial = x_2d.view(batch, frames, channels, height * width).permute(0, 2, 1, 3) + padding = int(self.t_conv.kernel_size[0]) // 2 + conv_input = spatial + prefix = 0 + if state.conv_tail is not None: + conv_input = torch.cat((state.conv_tail[:, :, -padding:], spatial), dim=2) + prefix = conv_input.shape[2] - spatial.shape[2] + + temporal = self.t_conv(conv_input)[:, :, prefix:] + out = spatial + temporal + if update_state: + _set_or_copy(state, "conv_tail", spatial[:, :, -padding:]) + return out.permute(0, 2, 3, 1).reshape(batch, tokens, channels) + + +class _Conv2dContainer(nn.Module): + """Expose a convolution under the checkpoint-compatible conv name.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + *, + groups: int = 1, + bias: bool = True, + padding: int = 0, + ) -> None: + super().__init__() + self.conv = nn.Conv2d( + in_channels, + out_channels, + kernel_size, + groups=groups, + bias=bias, + padding=padding, + ) + + def forward(self, x: Tensor) -> Tensor: + """Apply the contained convolution.""" + return self.conv(x) + + +class LongSanaBlock(nn.Module): + """One checkpoint-compatible LongSana transformer block.""" + + def __init__(self, spec: SanaWMStage1Spec, *, fp32_attention: bool) -> None: + super().__init__() + self.norm1 = nn.LayerNorm( + spec.hidden_size, + elementwise_affine=False, + eps=1e-6, + ) + self.attn = LongSanaLinearAttention( + spec, + fp32_attention=fp32_attention, + ) + self.cross_attn = Stage1CrossAttention(spec) + self.norm2 = nn.LayerNorm( + spec.hidden_size, + elementwise_affine=False, + eps=1e-6, + ) + self.mlp = LongSanaCausalGLUMBConvTemp(spec) + self.scale_shift_table = nn.Parameter(torch.empty(6, spec.hidden_size)) + + def forward( + self, + x: Tensor, + y: Tensor, + timestep_modulation: Tensor, + *, + frames: int, + height: int, + width: int, + mask: Tensor | None, + rotary_emb: Tensor, + state: LongSanaBlockState, + update_state: bool, + ) -> Tensor: + """Run self-attention, text cross-attention, and causal GLUMB.""" + batch, _tokens, _channels = x.shape + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None] + timestep_modulation.reshape(batch, 6, -1) + ).chunk(6, dim=1) + + attn_input = _modulate(self.norm1(x), shift_msa, scale_msa) + attn_output = self.attn( + attn_input, + rotary_emb=rotary_emb, + state=state, + update_state=update_state, + ) + x = x + gate_msa * attn_output + x = x + self.cross_attn(x, y, mask=mask) + + mlp_input = _modulate(self.norm2(x), shift_mlp, scale_mlp) + mlp_output = self.mlp( + mlp_input, + frames=frames, + height=height, + width=width, + state=state, + update_state=update_state, + ) + return x + gate_mlp * mlp_output + + +class _LongSanaFinalLayer(nn.Module): + """Final AdaLN and spatial unpatch projection.""" + + def __init__( + self, + spec: SanaWMStage1Spec, + patch_size: tuple[int, int, int], + ) -> None: + super().__init__() + self.norm_final = nn.LayerNorm( + spec.hidden_size, + elementwise_affine=False, + eps=1e-6, + ) + self.linear = nn.Linear( + spec.hidden_size, + math.prod(patch_size) * spec.latent_channels, + ) + self.scale_shift_table = nn.Parameter(torch.empty(2, spec.hidden_size)) + + def forward(self, x: Tensor, timestep_embedding: Tensor) -> Tensor: + """Project hidden tokens into patched latent channels.""" + shift, scale = ( + self.scale_shift_table[None] + timestep_embedding[:, None] + ).chunk(2, dim=1) + return self.linear(_modulate(self.norm_final(x), shift, scale)) + + +class LongSanaModel(nn.Module): + """Released LongSana 2B model with Runtime V2-owned recurrent state.""" + + def __init__(self, config: LongSanaNetworkConfig) -> None: + super().__init__() + self.config = config + self.spec = config.spec + self.patch_size = config.patch_size + self.register_buffer( + "pos_embed", + torch.zeros(1, 1800, self.spec.hidden_size), + ) + if self.patch_size[0] != 1: + raise ValueError("LongSana requires temporal patch size 1.") + self.x_embedder = _PatchEmbed3D( + self.spec.latent_channels, + self.spec.hidden_size, + self.patch_size, + ) + self.t_embedder = TimestepEmbedder(self.spec) + self.t_block = nn.Sequential( + nn.SiLU(), + nn.Linear(self.spec.hidden_size, 6 * self.spec.hidden_size), + ) + self.y_embedder = TextEmbedder(self.spec) + self.attention_y_norm = RMSNorm(self.spec.hidden_size, eps=1e-5) + with torch.no_grad(): + self.attention_y_norm.weight.fill_(0.01) + self.blocks = nn.ModuleList( + [ + LongSanaBlock( + self.spec, + fp32_attention=config.fp32_attention, + ) + for _ in range(self.spec.depth) + ] + ) + self.final_layer = _LongSanaFinalLayer(self.spec, self.patch_size) + + def prepare_condition(self, condition: Tensor) -> Tensor: + """Project static Gemma features once for an entire rollout.""" + y = self.y_embedder(condition.to(dtype=self.dtype)) + return self.attention_y_norm(y) + + def forward( + self, + x: Tensor, + timestep: Tensor, + projected_condition: Tensor, + condition_mask: Tensor, + block_states: list[LongSanaBlockState], + *, + start_frame: int, + update_state: bool, + ) -> Tensor: + """Predict flow and optionally advance all per-block recurrent states.""" + if x.ndim != 5: + raise ValueError( + f"LongSana expects BCTHW input, got shape {tuple(x.shape)}." + ) + if len(block_states) != len(self.blocks): + raise ValueError( + f"Expected {len(self.blocks)} block states, got {len(block_states)}." + ) + batch, channels, frames, latent_height, latent_width = x.shape + if channels != self.spec.latent_channels: + raise ValueError( + f"Expected {self.spec.latent_channels} latent channels, got {channels}." + ) + if latent_height % self.patch_size[1] or latent_width % self.patch_size[2]: + raise ValueError( + "LongSana latent height and width must be divisible by spatial " + f"patch size {self.patch_size[1:]}, got {(latent_height, latent_width)}." + ) + + x = x.to(dtype=self.dtype) + height = latent_height // self.patch_size[1] + width = latent_width // self.patch_size[2] + tokens = self.x_embedder(x) + rope = causal_wan_rope( + head_dim=self.spec.head_dim, + start_frame=start_frame, + frames=frames, + height=height, + width=width, + device=x.device, + ) + + model_timestep = timestep.reshape(()).expand(batch).long().float() + timestep_embedding = self.t_embedder(model_timestep) + modulation = self.t_block(timestep_embedding) + y = projected_condition.to(dtype=self.dtype) + mask = condition_mask.to(device=x.device) + + for block, state in zip(self.blocks, block_states): + tokens = block( + tokens, + y, + modulation, + frames=frames, + height=height, + width=width, + mask=mask, + rotary_emb=rope, + state=state, + update_state=update_state, + ) + + output = self.final_layer(tokens, timestep_embedding) + return _unpatchify( + output, + frames=frames, + height=height, + width=width, + channels=self.spec.latent_channels, + patch_size=self.patch_size, + ) + + @property + def dtype(self) -> torch.dtype: + """Return the checkpoint parameter dtype.""" + return self.x_embedder.proj.weight.dtype + + +def causal_wan_rope( + *, + head_dim: int, + start_frame: int, + frames: int, + height: int, + width: int, + device: torch.device, + max_sequence_length: int = 1024, +) -> Tensor: + """Build upstream-compatible complex128 Wan RoPE at absolute frame positions.""" + end_frame = start_frame + frames + if start_frame < 0 or min(frames, height, width) <= 0: + raise ValueError( + "LongSana RoPE dimensions must be positive and start_frame non-negative." + ) + if max(end_frame, height, width) > max_sequence_length: + raise ValueError( + "LongSana RoPE position exceeds the released 1024-position table: " + f"end_frame={end_frame}, height={height}, width={width}." + ) + + temporal_complex = head_dim // 2 - 2 * (head_dim // 6) + height_complex = head_dim // 6 + width_complex = head_dim // 6 + temporal = _axis_rope( + max_sequence_length, + temporal_complex, + device, + )[start_frame:end_frame] + vertical = _axis_rope(max_sequence_length, height_complex, device)[:height] + horizontal = _axis_rope(max_sequence_length, width_complex, device)[:width] + + temporal = temporal[:, None, None].expand(frames, height, width, -1) + vertical = vertical[None, :, None].expand(frames, height, width, -1) + horizontal = horizontal[None, None, :].expand(frames, height, width, -1) + return torch.cat((temporal, vertical, horizontal), dim=-1).reshape( + 1, 1, frames * height * width, head_dim // 2 + ) + + +def _axis_rope(length: int, complex_dims: int, device: torch.device) -> Tensor: + if complex_dims == 0: + return torch.empty(length, 0, dtype=torch.complex128, device=device) + dim = complex_dims * 2 + positions = torch.arange(length, device=device) + frequency = 1.0 / ( + 10000.0 ** (torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim) + ) + angles = torch.outer(positions, frequency) + return torch.polar(torch.ones_like(angles), angles) + + +def _apply_causal_rope(hidden_states: Tensor, frequencies: Tensor) -> Tensor: + """Apply complex RoPE exactly as the released cached attention module.""" + complex_states = torch.view_as_complex( + hidden_states.permute(0, 1, 3, 2).to(torch.float64).unflatten(3, (-1, 2)) + ) + rotated = torch.view_as_real(complex_states * frequencies) + return rotated.flatten(3, 4).permute(0, 1, 3, 2).type_as(hidden_states) + + +def _unpatchify( + x: Tensor, + *, + frames: int, + height: int, + width: int, + channels: int, + patch_size: tuple[int, int, int], +) -> Tensor: + batch = x.shape[0] + pt, ph, pw = patch_size + expected_tokens = frames * height * width + if x.shape[1] != expected_tokens: + raise ValueError(f"Expected {expected_tokens} output tokens, got {x.shape[1]}.") + x = x.reshape(batch, frames, height, width, pt, ph, pw, channels) + x = x.permute(0, 7, 1, 4, 2, 5, 3, 6) + return x.reshape(batch, channels, frames * pt, height * ph, width * pw) + + +def _modulate(x: Tensor, shift: Tensor, scale: Tensor) -> Tensor: + return x * (1 + scale) + shift + + +def _set_or_copy(state: LongSanaBlockState, name: str, value: Tensor) -> None: + detached = value.detach() + current = getattr(state, name) + if current is None: + setattr(state, name, detached.clone()) + return + if current.shape != detached.shape or current.dtype != detached.dtype: + raise ValueError( + f"LongSana cache slot {name} changed shape or dtype: " + f"{tuple(current.shape)}/{current.dtype} -> " + f"{tuple(detached.shape)}/{detached.dtype}." + ) + current.copy_(detached) diff --git a/integrations_v2/longsana/impl/pipeline.py b/integrations_v2/longsana/impl/pipeline.py new file mode 100644 index 000000000..bfcb4da25 --- /dev/null +++ b/integrations_v2/longsana/impl/pipeline.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LongSana prompt-to-video orchestration on the shared Runtime V2 pipeline.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, cast + +import torch +from torch import Tensor + +from flashdreams.infra.pipeline import ( + StreamInferencePipeline, + StreamInferencePipelineCache, + StreamInferencePipelineConfig, +) +from longsana.impl.constants import ( + LONGSANA_TEXT_CONFIG_PATH, + MOTION_SCORE, +) +from longsana.impl.transformer import ( + LongSanaConditioning, + LongSanaTransformer, +) +from sana_wm.impl.conditioning import ( + SanaWMTextPromptEncoder, + SanaWMTextPromptEncoderConfig, + SanaWMTextPromptRequest, +) + + +@dataclass(kw_only=True) +class LongSanaPipelineConfig(StreamInferencePipelineConfig): + """Config for prompt encoding, LongSana diffusion, and Wan decoding.""" + + _target: type["LongSanaPipeline"] = field(default_factory=lambda: LongSanaPipeline) + + prompt_encoder: SanaWMTextPromptEncoderConfig = field( + default_factory=lambda: SanaWMTextPromptEncoderConfig( + config_path=LONGSANA_TEXT_CONFIG_PATH, + offload_text_encoder=True, + ) + ) + """Shared SANA-family Gemma/CHI prompt encoder.""" + + motion_score: int = MOTION_SCORE + """Training-time motion-score suffix added to every user prompt.""" + + +class LongSanaPipeline(StreamInferencePipeline): + """Runtime V2 pipeline with one prompt and a recurrent LongSana cache.""" + + config: LongSanaPipelineConfig + + def __init__(self, config: LongSanaPipelineConfig) -> None: + super().__init__(config) + self.config = config + self.prompt_encoder = cast( + SanaWMTextPromptEncoder, + config.prompt_encoder.setup(), + ) + + @torch.no_grad() + def initialize_cache( + self, + *, + text: list[str], + image: Any = None, + height: int | None = None, + width: int | None = None, + ) -> StreamInferencePipelineCache: + """Encode one prompt and construct a constant-memory rollout cache.""" + if image is not None: + raise ValueError("The released LongSana checkpoint is text-to-video only.") + if len(text) != 1 or not text[0].strip(): + raise ValueError("LongSana requires exactly one non-empty prompt.") + + transformer = self.diffusion_model.transformer + if not isinstance(transformer, LongSanaTransformer): + raise TypeError("LongSanaPipeline requires LongSanaTransformer.") + latent_height = transformer.config.latent_height if height is None else height + latent_width = transformer.config.latent_width if width is None else width + expected = ( + transformer.config.latent_height, + transformer.config.latent_width, + ) + if (latent_height, latent_width) != expected: + raise ValueError( + "This LongSana release is configured for latent size " + f"{expected}, got {(latent_height, latent_width)}." + ) + + prompt = f"{text[0].strip()} motion score: {self.config.motion_score}." + encoded = self.prompt_encoder( + SanaWMTextPromptRequest(prompt=prompt, negative_prompt="") + ) + conditioning = LongSanaConditioning( + condition=encoded.condition, + mask=encoded.condition_mask, + ) + return super().initialize_cache( + transformer_context={"conditioning": conditioning}, + ) + + @torch.no_grad() + def generate( + self, + autoregressive_index: int, + cache: StreamInferencePipelineCache, + input: Any = None, + ) -> Tensor: + """Select LongSana's first/steady block shape and generate one chunk.""" + transformer = self.diffusion_model.transformer + if not isinstance(transformer, LongSanaTransformer): + raise TypeError("LongSanaPipeline requires LongSanaTransformer.") + transformer.select_autoregressive_index(autoregressive_index) + return super().generate( + autoregressive_index=autoregressive_index, + cache=cache, + input=input, + ) + + def close(self) -> None: + """Release prompt and generator runtimes held by the resident pipeline.""" + self.prompt_encoder.release_runtime() + transformer = self.diffusion_model.transformer + if isinstance(transformer, LongSanaTransformer): + transformer.release_runtime() diff --git a/integrations_v2/longsana/impl/scheduler.py b/integrations_v2/longsana/impl/scheduler.py new file mode 100644 index 000000000..b06d7883b --- /dev/null +++ b/integrations_v2/longsana/impl/scheduler.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LongSana parity wrapper around the shared self-forcing flow scheduler.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch +from torch import Tensor +from tqdm import tqdm + +from flashdreams.infra.diffusion.scheduler.fm import ( + FlowMatchScheduler, + FlowMatchSchedulerConfig, +) +from flashdreams.infra.diffusion.scheduler import FlowPredictor + + +@dataclass(kw_only=True) +class LongSanaFlowMatchSchedulerConfig(FlowMatchSchedulerConfig): + """Flow scheduler preserving upstream LongSana precision and RNG layout.""" + + _target: type["LongSanaFlowMatchScheduler"] = field( + default_factory=lambda: LongSanaFlowMatchScheduler + ) + + +class LongSanaFlowMatchScheduler(FlowMatchScheduler): + """Four-step self-forcing scheduler with upstream B/T/C noise ordering.""" + + config: LongSanaFlowMatchSchedulerConfig + + def sample( + self, + initial_noise: Tensor, + predict_flow: FlowPredictor, + rng: torch.Generator | None = None, + ) -> Tensor: + """Denoise TCHW/BCTHW using upstream double-precision x0 conversion.""" + if initial_noise.ndim not in (4, 5): + raise ValueError( + "LongSana scheduler expects TCHW or BCTHW latents, " + f"got shape {tuple(initial_noise.shape)}." + ) + input_dtype = initial_noise.dtype + noisy = initial_noise + clean: Tensor | None = None + for index in tqdm( + range(self.denoising_step_list.shape[0]), + disable=not self.config.enable_tqdm, + desc="LongSanaFlowMatchScheduler", + ): + sigma = self.denoising_sigmas[index] + timestep = self.denoising_step_list[index].to(dtype=input_dtype) + if index > 0: + if clean is None: + raise RuntimeError("LongSana scheduler lost its previous x0.") + noise = _upstream_renoise_tensor(noisy, rng) + noisy = ((1.0 - sigma) * clean + sigma * noise).to(input_dtype) + flow = predict_flow(noisy, timestep) + clean = (noisy.double() - sigma.double() * flow.double()).to(input_dtype) + if clean is None: + raise RuntimeError("LongSana denoising timestep list is empty.") + return clean + + +def _upstream_renoise_tensor( + like: Tensor, + rng: torch.Generator | None, +) -> Tensor: + """Draw noise in upstream's flattened B/T/C/H/W element order.""" + unbatched = like.ndim == 4 + if unbatched: + frames, channels, height, width = like.shape + batch = 1 + else: + batch, channels, frames, height, width = like.shape + noise = torch.randn( + (batch * frames, channels, height, width), + device=like.device, + dtype=like.dtype, + generator=rng, + ) + if unbatched: + return noise + noise = noise.unflatten(0, (batch, frames)).permute(0, 2, 1, 3, 4).contiguous() + return noise diff --git a/integrations_v2/longsana/impl/transformer.py b/integrations_v2/longsana/impl/transformer.py new file mode 100644 index 000000000..f98d1c40f --- /dev/null +++ b/integrations_v2/longsana/impl/transformer.py @@ -0,0 +1,379 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Runtime V2 transformer adapter for the released LongSana generator.""" + +from __future__ import annotations + +import gc +import time +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any, cast + +import torch +import torch.nn as nn +from loguru import logger +from torch import Tensor + +from flashdreams.core.checkpoint.load import load_checkpoint +from flashdreams.infra.compile import compile_module +from flashdreams.infra.diffusion.transformer import ( + Transformer, + TransformerAutoregressiveCache, + TransformerConfig, +) +from longsana.impl.constants import ( + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, + FIRST_LATENT_BLOCK_FRAMES, + LATENT_BLOCK_FRAMES, + LONGSANA_CHECKPOINT_PATH, +) +from longsana.impl.model import ( + LongSanaBlockState, + LongSanaModel, + LongSanaNetworkConfig, +) + + +@dataclass(kw_only=True) +class LongSanaConditioning: + """Static positive-prompt inputs for one LongSana rollout.""" + + condition: Tensor + """Gemma prompt features shaped [B, 1, L, 2304].""" + + mask: Tensor + """Valid-token mask shaped [B, L].""" + + +@dataclass(kw_only=True) +class LongSanaTransformerCache(TransformerAutoregressiveCache): + """Prompt and constant-memory block state for one LongSana rollout.""" + + conditioning: LongSanaConditioning | None = None + """Static prompt conditioning.""" + + block_states: list[LongSanaBlockState] = field(default_factory=list) + """One three-tensor recurrent state per transformer block.""" + + first_block_frames: int = FIRST_LATENT_BLOCK_FRAMES + """Latent frames generated by AR step zero.""" + + block_frames: int = LATENT_BLOCK_FRAMES + """Latent frames generated by every later AR step.""" + + projected_condition: Tensor | None = None + """Prompt projection cached after its first model use.""" + + start_frame: int = 0 + """Absolute latent-frame position of the active block.""" + + next_index: int = 0 + """Next legal autoregressive index.""" + + active_index: int | None = None + """Currently-generating index, cleared by finalize.""" + + active_frames: int = 0 + """Latent-frame count of the active block.""" + + def start(self, autoregressive_index: int) -> None: + """Validate and mark the start of the next LongSana block.""" + if self.active_index is not None: + raise RuntimeError( + f"LongSana AR step {self.active_index} has not been finalized." + ) + if autoregressive_index != self.next_index: + raise ValueError( + f"Expected LongSana AR step {self.next_index}, " + f"got {autoregressive_index}." + ) + self.active_index = autoregressive_index + self.active_frames = ( + self.first_block_frames if autoregressive_index == 0 else self.block_frames + ) + + def finalize(self, autoregressive_index: int) -> None: + """Advance absolute positions after the clean-timestep cache update.""" + if self.active_index != autoregressive_index: + raise ValueError( + f"Active LongSana AR step is {self.active_index}, " + f"cannot finalize {autoregressive_index}." + ) + self.start_frame += self.active_frames + self.next_index += 1 + self.active_index = None + self.active_frames = 0 + + def state_bytes(self) -> int: + """Return recurrent block-state bytes, excluding static prompt tensors.""" + return sum(state.num_bytes() for state in self.block_states) + + +@dataclass(kw_only=True) +class LongSanaTransformerConfig(TransformerConfig): + """Config for the Runtime V2 LongSana transformer.""" + + _target: type["LongSanaTransformer"] = field( + default_factory=lambda: LongSanaTransformer + ) + + network: LongSanaNetworkConfig = field(default_factory=LongSanaNetworkConfig) + """Checkpoint-facing LongSana network architecture.""" + + checkpoint_path: str = LONGSANA_CHECKPOINT_PATH + """Public self-forcing checkpoint path or Hugging Face file URL.""" + + dtype: torch.dtype = torch.bfloat16 + """Generator parameter and activation dtype.""" + + latent_height: int = DEFAULT_VIDEO_HEIGHT // 8 + """Wan-latent height.""" + + latent_width: int = DEFAULT_VIDEO_WIDTH // 8 + """Wan-latent width.""" + + first_block_frames: int = FIRST_LATENT_BLOCK_FRAMES + """Latent frames in AR step zero.""" + + block_frames: int = LATENT_BLOCK_FRAMES + """Latent frames in each later AR step.""" + + compile_network: bool = False + """Compile the DiT after strict checkpoint loading.""" + + +class LongSanaTransformer(Transformer[LongSanaTransformerCache]): + """Lazy-loading Runtime V2 adapter for LongSana flow prediction.""" + + config: LongSanaTransformerConfig + + def __init__(self, config: LongSanaTransformerConfig) -> None: + super().__init__(config) + self.config = config + self._dummy = nn.Parameter(torch.empty(0, dtype=config.dtype)) + self._model_built = False + self._active_autoregressive_index = 0 + + @property + def latent_shape(self) -> tuple[int, ...]: + """Return the raw Wan latent shape for the selected AR block.""" + frames = ( + self.config.first_block_frames + if self._active_autoregressive_index == 0 + else self.config.block_frames + ) + return ( + frames, + self.config.network.spec.latent_channels, + self.config.latent_height, + self.config.latent_width, + ) + + def select_autoregressive_index(self, autoregressive_index: int) -> None: + """Select the block shape before Runtime V2 allocates initial noise.""" + if autoregressive_index < 0: + raise ValueError("LongSana autoregressive_index must be non-negative.") + self._active_autoregressive_index = autoregressive_index + + def initialize_autoregressive_cache( + self, + *, + conditioning: LongSanaConditioning | None = None, + **_: Any, + ) -> LongSanaTransformerCache: + """Create a prompt cache and empty constant-memory state.""" + return LongSanaTransformerCache( + conditioning=conditioning, + block_states=[ + LongSanaBlockState() for _ in range(self.config.network.spec.depth) + ], + first_block_frames=self.config.first_block_frames, + block_frames=self.config.block_frames, + ) + + def initial_noise( + self, + *, + latent_shape: tuple[int, ...], + rng: torch.Generator | None, + cache: LongSanaTransformerCache, + input: Any = None, + ) -> Tensor: + """Draw upstream B/C/T/H/W noise and expose it as Runtime V2 T/C/H/W.""" + del cache, input + frames, channels, height, width = latent_shape + noise = torch.randn( + (channels, frames, height, width), + device=self.device, + dtype=self.dtype, + generator=rng, + ) + return noise.permute(1, 0, 2, 3).contiguous() + + def predict_flow( + self, + noisy_latent: Tensor, + timestep: Tensor, + cache: LongSanaTransformerCache, + input: Any = None, + ) -> Tensor: + """Predict flow without mutating the recurrent LongSana state.""" + if input is not None: + raise ValueError("LongSana T2V does not accept a per-block encoder input.") + return self._predict( + noisy_latent=noisy_latent, + timestep=timestep, + cache=cache, + update_state=False, + ) + + def finalize_kv_cache( + self, + noisy_latent: Tensor, + timestep: Tensor, + cache: LongSanaTransformerCache, + input: Any = None, + ) -> None: + """Commit the clean block using LongSana's required zero-step forward.""" + if input is not None: + raise ValueError("LongSana T2V does not accept a per-block encoder input.") + _ = self._predict( + noisy_latent=noisy_latent, + timestep=timestep, + cache=cache, + update_state=True, + ) + + def patchify_and_maybe_split_cp(self, x: Any) -> Any: + """Keep Wan latents in raw BCTHW layout; the DiT owns patch embedding.""" + return x + + def unpatchify_and_maybe_gather_cp(self, x: Tensor) -> Tensor: + """Return raw BCTHW latents unchanged.""" + return x + + def release_runtime(self) -> None: + """Release loaded generator weights and CUDA caching allocations.""" + if hasattr(self, "model"): + del self.model + self._model_built = False + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def _predict( + self, + *, + noisy_latent: Tensor, + timestep: Tensor, + cache: LongSanaTransformerCache, + update_state: bool, + ) -> Tensor: + conditioning = cache.conditioning + if conditioning is None: + raise RuntimeError("LongSana cache was initialized without a text prompt.") + if cache.active_index is None: + raise RuntimeError( + "LongSana flow prediction ran outside an active AR step." + ) + self._ensure_model() + model = cast(LongSanaModel, self.model) + if cache.projected_condition is None: + cache.projected_condition = model.prepare_condition( + conditioning.condition.to(device=self.device) + ) + add_batch = noisy_latent.ndim == 4 + model_input = ( + noisy_latent.permute(1, 0, 2, 3).unsqueeze(0).contiguous() + if add_batch + else noisy_latent + ) + output = model( + model_input, + timestep, + cache.projected_condition, + conditioning.mask, + cache.block_states, + start_frame=cache.start_frame, + update_state=update_state, + ) + return ( + output.squeeze(0).permute(1, 0, 2, 3).contiguous() if add_batch else output + ) + + def _ensure_model(self) -> None: + if self._model_built: + return + started = time.perf_counter() + with torch.device(self.device): + model = self.config.network.setup() + model = model.to(dtype=self.config.dtype) + logger.info( + "[LongSana] Built {} ({:,} parameters)", + type(model).__name__, + sum(parameter.numel() for parameter in model.parameters()), + ) + + payload = load_checkpoint(self.config.checkpoint_path) + state_dict = longsana_state_dict(payload) + model.load_state_dict(state_dict, strict=True) + del payload, state_dict + gc.collect() + + model = model.eval() + if self.config.compile_network: + model = compile_module(model) + self.model = model + self._model_built = True + logger.info( + "[timing] LongSana generator build+load: {:.3f}s (dtype={})", + time.perf_counter() - started, + self.config.dtype, + ) + + +def longsana_state_dict(payload: Any) -> dict[str, Tensor]: + """Extract generator.model tensors from the public training checkpoint.""" + if not isinstance(payload, Mapping): + raise TypeError( + f"LongSana checkpoint must be a mapping, got {type(payload).__name__}." + ) + state: Any = payload.get("generator", payload) + if isinstance(state, Mapping): + state = state.get("state_dict", state) + if not isinstance(state, Mapping): + raise TypeError( + "LongSana checkpoint generator/state_dict entry must be a mapping." + ) + + model_prefix = "model." + prefixed = { + str(key)[len(model_prefix) :]: value + for key, value in state.items() + if str(key).startswith(model_prefix) and isinstance(value, Tensor) + } + if prefixed: + return prefixed + return { + str(key): value for key, value in state.items() if isinstance(value, Tensor) + } + + +def checkpoint_tensor_bytes(state_dict: Mapping[str, Tensor]) -> int: + """Return total tensor bytes in a normalized LongSana state dict.""" + return sum(tensor.numel() * tensor.element_size() for tensor in state_dict.values()) diff --git a/integrations_v2/longsana/pyproject.toml b/integrations_v2/longsana/pyproject.toml new file mode 100644 index 000000000..45ab452fc --- /dev/null +++ b/integrations_v2/longsana/pyproject.toml @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "flashdreams-longsana" +version = "0.1.0" +description = "LongSana 2B constant-memory streaming T2V for FlashDreams Runtime V2." +readme = "README.md" +requires-python = ">=3.10" +dependencies = [ + "flashdreams", + "flashdreams-sana-wm", + "flashdreams-t2v", + "mediapy>=1.1", +] + +[tool.uv.sources] +flashdreams = { workspace = true } +flashdreams-sana-wm = { workspace = true } +flashdreams-t2v = { workspace = true } + +[project.optional-dependencies] +dev = ["pytest>=8.0"] + +[project.entry-points."flashdreams.applications_v2"] +"t2v-longsana-2b-480p" = "longsana.apps.t2v.adapter:create_app" + +[tool.setuptools] +packages = ["longsana", "longsana.apps", "longsana.apps.t2v", "longsana.impl"] + +[tool.setuptools.package-dir] +longsana = "." + +[tool.setuptools.package-data] +longsana = ["resources/*.yaml"] + +[tool.uv] +managed = true diff --git a/integrations_v2/longsana/resources/longsana_text.yaml b/integrations_v2/longsana/resources/longsana_text.yaml new file mode 100644 index 000000000..275a10225 --- /dev/null +++ b/integrations_v2/longsana/resources/longsana_text.yaml @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +work_dir: "" +model: + mixed_precision: bf16 +text_encoder: + text_encoder_name: gemma-2-2b-it + y_norm: true + y_norm_scale_factor: 0.01 + model_max_length: 300 + chi_prompt: + - 'Given a user prompt, generate an "Enhanced prompt" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:' + - '- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.' + - '- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.' + - 'Here are examples of how to transform or refine prompts:' + - '- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.' + - '- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.' + - 'Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:' + - 'User Prompt: ' diff --git a/integrations_v2/longsana/scripts/benchmark.py b/integrations_v2/longsana/scripts/benchmark.py new file mode 100644 index 000000000..031e1b5f4 --- /dev/null +++ b/integrations_v2/longsana/scripts/benchmark.py @@ -0,0 +1,296 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark and visually validate LongSANA with a diverse prompt suite.""" + +from __future__ import annotations + +import argparse +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +import json +from pathlib import Path +import platform +import statistics +import subprocess +import time +from typing import Any + +import mediapy as media +import torch + +from longsana.config import PIPELINE_LONGSANA_2B_480P +from longsana.impl.constants import DEFAULT_VIDEO_FPS +from longsana.impl.transformer import LongSanaTransformerCache + + +@dataclass(frozen=True, kw_only=True) +class ValidationCase: + """One prompt and rollout length in the validation matrix.""" + + slug: str + prompt: str + blocks: int + seed: int + category: str + + +DEFAULT_PROMPTS = ( + ( + "animal_motion", + "A red panda walks briskly through a misty bamboo forest at sunrise, " + "tracking shot, realistic fur moving in the wind.", + "subject identity and articulated motion", + ), + ( + "dance_camera", + "Three street dancers perform fast synchronized choreography in a neon " + "alley at night while the camera smoothly circles them, cinematic.", + "multiple subjects, fast motion, and camera orbit", + ), + ( + "cooking_interaction", + "A chef flips vegetables in a flaming wok, pours sauce, and plates the " + "dish in a busy restaurant kitchen, close-up documentary camera.", + "object interaction and ordered actions", + ), + ( + "coastal_aerial", + "A continuous aerial shot flies over sea cliffs toward a lighthouse as " + "waves crash below and storm clouds roll across the coast, photorealistic.", + "large camera translation and scene continuity", + ), + ( + "surreal_long", + "An astronaut riding a white horse crosses a moonlit salt flat; reflections " + "ripple under the hooves as the camera follows from behind, cinematic.", + "long-horizon composition and identity", + ), +) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("artifacts/longsana_benchmark"), + ) + parser.add_argument("--blocks", type=int, default=2) + parser.add_argument("--long-blocks", type=int, default=6) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--device", default="cuda") + parser.add_argument( + "--case", + action="append", + dest="cases", + help="Run only this case slug; repeat to select more than one.", + ) + parser.add_argument( + "--no-contact-sheets", + action="store_true", + help="Skip ffmpeg contact-sheet creation.", + ) + args = parser.parse_args() + if args.blocks <= 0 or args.long_blocks <= 0: + parser.error("--blocks and --long-blocks must be positive") + return args + + +def _cases(args: argparse.Namespace) -> list[ValidationCase]: + selected = set(args.cases or ()) + known = {slug for slug, _prompt, _category in DEFAULT_PROMPTS} + unknown = selected - known + if unknown: + raise ValueError(f"Unknown validation case(s): {', '.join(sorted(unknown))}") + return [ + ValidationCase( + slug=slug, + prompt=prompt, + blocks=args.long_blocks if slug == "surreal_long" else args.blocks, + seed=args.seed + index, + category=category, + ) + for index, (slug, prompt, category) in enumerate(DEFAULT_PROMPTS) + if not selected or slug in selected + ] + + +def _gpu_info(device: str) -> dict[str, Any]: + info: dict[str, Any] = { + "device": device, + "cuda_available": torch.cuda.is_available(), + } + if torch.cuda.is_available() and torch.device(device).type == "cuda": + index = torch.cuda.current_device() + properties = torch.cuda.get_device_properties(index) + info.update( + { + "name": properties.name, + "total_memory_gib": properties.total_memory / 1024**3, + "compute_capability": [properties.major, properties.minor], + } + ) + return info + + +def _contact_sheet(video_path: Path, output_path: Path) -> None: + subprocess.run( + [ + "ffmpeg", + "-hide_banner", + "-loglevel", + "error", + "-y", + "-i", + str(video_path), + "-vf", + "fps=2,scale=416:240,tile=4x5", + "-frames:v", + "1", + str(output_path), + ], + check=True, + ) + + +def _reset_generator(pipeline: Any, seed: int) -> None: + pipeline.diffusion_model._rng = torch.Generator( # noqa: SLF001 + device=pipeline.device + ).manual_seed(seed) + + +def _run_case( + pipeline: Any, + case: ValidationCase, + output_dir: Path, + *, + contact_sheets: bool, +) -> dict[str, Any]: + _reset_generator(pipeline, case.seed) + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + + prompt_started = time.perf_counter() + cache = pipeline.initialize_cache(text=[case.prompt]) + prompt_encode_s = time.perf_counter() - prompt_started + if not isinstance(cache.transformer_cache, LongSanaTransformerCache): + raise TypeError("Benchmark requires a LongSanaTransformerCache.") + + chunks: list[torch.Tensor] = [] + block_metrics: list[dict[str, float]] = [] + cache_mib_history: list[float] = [] + for index in range(case.blocks): + frames = pipeline.generate(index, cache) + metrics = pipeline.finalize(index, cache) + if metrics is None: + raise RuntimeError("LongSANA profiling must be enabled for this benchmark.") + if not bool(torch.isfinite(frames).all()): + raise RuntimeError(f"Non-finite video output in {case.slug} block {index}.") + chunks.append(frames.detach().cpu()) + block_metrics.append(metrics) + cache_mib_history.append(cache.transformer_cache.state_bytes() / 1024**2) + + video = torch.cat(chunks).clamp(0, 1) + video_path = output_dir / f"{case.slug}_{case.blocks}blocks.mp4" + media.write_video( + video_path, + video.permute(0, 2, 3, 1).contiguous().numpy(), + fps=DEFAULT_VIDEO_FPS, + ) + if contact_sheets: + _contact_sheet( + video_path, + output_dir / f"{case.slug}_{case.blocks}blocks_contact.jpg", + ) + + steady = block_metrics[1:] or block_metrics + steady_total = [item["total_ms"] for item in steady] + steady_frames = sum(int(chunk.shape[0]) for chunk in chunks[1:] or chunks) + steady_elapsed_ms = sum(steady_total) + result = { + **asdict(case), + "video": str(video_path.resolve()), + "prompt_encode_s": prompt_encode_s, + "frames": int(video.shape[0]), + "shape": list(video.shape), + "finite": True, + "cache_mib_history": cache_mib_history, + "cache_constant": len({round(value, 6) for value in cache_mib_history}) == 1, + "block_metrics": block_metrics, + "steady_state": { + "blocks": len(steady), + "total_ms_median": statistics.median(steady_total), + "total_ms_mean": statistics.mean(steady_total), + "end_to_end_fps": steady_frames / (steady_elapsed_ms / 1000), + "diffusion_fps": steady_frames + / (sum(item["diffuse_ms"] + item["finalize_ms"] for item in steady) / 1000), + "decode_ms_median": statistics.median(item["decode_ms"] for item in steady), + "peak_memory_gib": max(item["mem_peak_gib"] for item in steady), + }, + } + case_path = output_dir / f"{case.slug}_{case.blocks}blocks.json" + case_path.write_text(json.dumps(result, indent=2) + "\n") + return result + + +def main() -> None: + """Run selected validation cases through the public LongSANA pipeline.""" + args = _parse_args() + cases = _cases(args) + args.output_dir.mkdir(parents=True, exist_ok=True) + if torch.device(args.device).type != "cuda" or not torch.cuda.is_available(): + raise RuntimeError("The released LongSANA benchmark requires a CUDA GPU.") + + setup_started = time.perf_counter() + pipeline = PIPELINE_LONGSANA_2B_480P.setup().to(args.device).eval() + setup_s = time.perf_counter() - setup_started + try: + results = [ + _run_case( + pipeline, + case, + args.output_dir, + contact_sheets=not args.no_contact_sheets, + ) + for case in cases + ] + finally: + pipeline.close() + + steady = [result["steady_state"] for result in results] + summary = { + "created_at": datetime.now(timezone.utc).isoformat(), + "pipeline": PIPELINE_LONGSANA_2B_480P.name, + "setup_s": setup_s, + "software": { + "python": platform.python_version(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + }, + "hardware": _gpu_info(args.device), + "cases": results, + "aggregate_steady_state": { + "end_to_end_fps_median": statistics.median( + item["end_to_end_fps"] for item in steady + ), + "diffusion_fps_median": statistics.median( + item["diffusion_fps"] for item in steady + ), + "total_ms_median": statistics.median( + item["total_ms_median"] for item in steady + ), + "peak_memory_gib_max": max(item["peak_memory_gib"] for item in steady), + "cache_mib": results[0]["cache_mib_history"][-1], + "all_outputs_finite": all(result["finite"] for result in results), + "all_caches_constant": all(result["cache_constant"] for result in results), + }, + } + summary_path = args.output_dir / "summary.json" + summary_path.write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary["aggregate_steady_state"], indent=2)) + print(f"Results: {summary_path.resolve()}") + + +if __name__ == "__main__": + main() diff --git a/integrations_v2/longsana/scripts/operator_profile.py b/integrations_v2/longsana/scripts/operator_profile.py new file mode 100644 index 000000000..1ffea4671 --- /dev/null +++ b/integrations_v2/longsana/scripts/operator_profile.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Capture an operator profile for one warmed LongSANA Runtime V2 block.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + +from longsana.config import PIPELINE_LONGSANA_2B_480P +from longsana.impl.transformer import LongSanaTransformerCache + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("artifacts/longsana_profile"), + ) + parser.add_argument( + "--prompt", + default="A red panda walks through a misty bamboo forest, tracking shot.", + ) + parser.add_argument("--device", default="cuda") + return parser.parse_args() + + +def main() -> None: + """Warm one block, then profile the complete next-block lifecycle.""" + args = _parse_args() + if torch.device(args.device).type != "cuda" or not torch.cuda.is_available(): + raise RuntimeError("The LongSANA operator profile requires a CUDA GPU.") + args.output_dir.mkdir(parents=True, exist_ok=True) + + pipeline = PIPELINE_LONGSANA_2B_480P.setup().to(args.device).eval() + try: + cache = pipeline.initialize_cache(text=[args.prompt]) + _ = pipeline.generate(0, cache) + _ = pipeline.finalize(0, cache) + torch.cuda.reset_peak_memory_stats() + + activities = [ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ] + with torch.profiler.profile( + activities=activities, + record_shapes=True, + profile_memory=True, + with_stack=False, + ) as profile: + frames = pipeline.generate(1, cache) + metrics = pipeline.finalize(1, cache) + + if metrics is None: + raise RuntimeError("LongSANA profiling must be enabled.") + transformer_cache = cache.transformer_cache + if not isinstance(transformer_cache, LongSanaTransformerCache): + raise TypeError("Profile requires a LongSanaTransformerCache.") + + trace_path = args.output_dir / "steady_state_trace.json" + table_path = args.output_dir / "steady_state_operators.txt" + summary_path = args.output_dir / "summary.json" + profile.export_chrome_trace(str(trace_path)) + table = profile.key_averages(group_by_input_shape=True).table( + sort_by="self_cuda_time_total", + row_limit=40, + ) + table_path.write_text(table + "\n") + summary = { + "pipeline": PIPELINE_LONGSANA_2B_480P.name, + "prompt": args.prompt, + "profiled_block": 1, + "frames": int(frames.shape[0]), + "shape": list(frames.shape), + "finite": bool(torch.isfinite(frames).all()), + "cache_mib": transformer_cache.state_bytes() / 1024**2, + "metrics": metrics, + "trace": str(trace_path.resolve()), + "operators": str(table_path.resolve()), + } + summary_path.write_text(json.dumps(summary, indent=2) + "\n") + print(table) + print(f"Results: {summary_path.resolve()}") + finally: + pipeline.close() + + +if __name__ == "__main__": + main() diff --git a/integrations_v2/longsana/tests/test_smoke.py b/integrations_v2/longsana/tests/test_smoke.py new file mode 100644 index 000000000..68aae2116 --- /dev/null +++ b/integrations_v2/longsana/tests/test_smoke.py @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU-safe contract and numerical tests for LongSana.""" + +from __future__ import annotations + +import torch +from torch import Tensor +import pytest + +from flashdreams.infra.diffusion.model import DiffusionModel +from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig +from flashdreams.recipes.wan.autoencoder.vae import WanVAEDecoderConfig +from longsana.config import LONGSANA_CONFIGS, PIPELINE_LONGSANA_2B_480P +from longsana.impl.constants import ( + DEFAULT_DENOISING_TIMESTEPS, + FIRST_LATENT_BLOCK_FRAMES, + LATENT_BLOCK_FRAMES, + LONGSANA_REVISION, + LONGSANA_TEXT_CONFIG_PATH, + SANA_VIDEO_REVISION, +) +from longsana.impl.model import ( + LongSanaBlockState, + LongSanaNetworkConfig, + causal_wan_rope, +) +from longsana.impl.pipeline import LongSanaPipelineConfig +from longsana.impl.scheduler import ( + LongSanaFlowMatchScheduler, + LongSanaFlowMatchSchedulerConfig, +) +from longsana.impl.transformer import ( + LongSanaConditioning, + LongSanaTransformer, + LongSanaTransformerCache, + LongSanaTransformerConfig, + longsana_state_dict, +) +from sana_wm.impl.stage1_model import SanaWMStage1Spec +from sana_wm.impl.transformer import _load_inference_config + +pytestmark = pytest.mark.ci_cpu + + +def _small_spec() -> SanaWMStage1Spec: + return SanaWMStage1Spec( + latent_channels=2, + hidden_size=16, + text_dim=12, + timestep_dim=8, + depth=2, + num_heads=2, + head_dim=8, + max_text_length=5, + latent_grid_size=(2, 2), + mlp_ratio=1, + temporal_kernel_size=3, + ) + + +def test_public_config_uses_runtime_v2_and_release_schedule() -> None: + """Keep the public pipeline on shared diffusion and Wan components.""" + config = PIPELINE_LONGSANA_2B_480P + + assert isinstance(config, LongSanaPipelineConfig) + assert config._target.__name__ == "LongSanaPipeline" + assert config.diffusion_model._target is DiffusionModel + assert isinstance(config.decoder, WanVAEDecoderConfig) + assert config.decoder.dtype is torch.float32 + assert isinstance( + config.diffusion_model.transformer, + LongSanaTransformerConfig, + ) + scheduler = config.diffusion_model.scheduler + assert isinstance(scheduler, LongSanaFlowMatchSchedulerConfig) + assert scheduler.denoising_timesteps == DEFAULT_DENOISING_TIMESTEPS + assert scheduler.shift == 7.0 + assert scheduler.warp_denoising_step is False + assert config.name in LONGSANA_CONFIGS + assert LONGSANA_REVISION in config.diffusion_model.transformer.checkpoint_path + assert SANA_VIDEO_REVISION in config.decoder.checkpoint_path + + +def test_packaged_text_config_matches_release() -> None: + """Use the released Gemma model, 300-token CHI prompt, and BF16 output.""" + config = _load_inference_config(LONGSANA_TEXT_CONFIG_PATH) + + assert config.model.mixed_precision == "bf16" + assert config.text_encoder.text_encoder_name == "gemma-2-2b-it" + assert config.text_encoder.model_max_length == 300 + assert config.text_encoder.y_norm_scale_factor == 0.01 + assert config.text_encoder.chi_prompt[-1] == "User Prompt: " + + +def test_full_model_schema_has_public_checkpoint_shape() -> None: + """Construct the 2B schema on meta without allocating model weights.""" + with torch.device("meta"): + model = LongSanaNetworkConfig().setup() + + assert len(model.state_dict()) == 418 + assert sum(parameter.numel() for parameter in model.parameters()) == 2_057_553_344 + assert model.state_dict()["x_embedder.proj.weight"].shape == ( + 2240, + 16, + 1, + 2, + 2, + ) + assert model.state_dict()["final_layer.linear.weight"].shape == (64, 2240) + + +def test_checkpoint_normalization_selects_generator_model() -> None: + """Discard critic/EMA entries and remove the wrapper's model prefix.""" + expected = torch.randn(2, 3) + payload = { + "generator": { + "model.layer.weight": expected, + "scheduler.buffer": torch.ones(1), + }, + "generator_ema": {"model.layer.weight": torch.zeros_like(expected)}, + "critic": {"weight": torch.ones(1)}, + } + + state = longsana_state_dict(payload) + + assert set(state) == {"layer.weight"} + assert state["layer.weight"] is expected + + +def test_recurrent_state_size_and_storage_are_constant() -> None: + """Accumulate a second block without retaining its token history.""" + torch.manual_seed(7) + spec = _small_spec() + model = LongSanaNetworkConfig(spec=spec).setup().eval() + condition = model.prepare_condition(torch.randn(1, 1, 5, spec.text_dim)) + mask = torch.ones(1, 5) + states = [LongSanaBlockState() for _ in range(spec.depth)] + latent = torch.randn(1, spec.latent_channels, 3, 4, 4) + + first = model( + latent, + torch.tensor(727.0), + condition, + mask, + states, + start_frame=0, + update_state=True, + ) + first_bytes = sum(state.num_bytes() for state in states) + pointers = [ + ( + state.value_key.data_ptr(), + state.key_sum.data_ptr(), + state.conv_tail.data_ptr(), + ) + for state in states + if state.value_key is not None + and state.key_sum is not None + and state.conv_tail is not None + ] + second = model( + latent, + torch.tensor(727.0), + condition, + mask, + states, + start_frame=3, + update_state=True, + ) + + assert first.shape == second.shape == latent.shape + assert sum(state.num_bytes() for state in states) == first_bytes + assert [ + ( + state.value_key.data_ptr(), + state.key_sum.data_ptr(), + state.conv_tail.data_ptr(), + ) + for state in states + if state.value_key is not None + and state.key_sum is not None + and state.conv_tail is not None + ] == pointers + + +def test_causal_rope_uses_absolute_frame_positions() -> None: + """A later block's temporal frequencies equal the full table slice.""" + complete = causal_wan_rope( + head_dim=8, + start_frame=0, + frames=5, + height=1, + width=1, + device=torch.device("cpu"), + ) + later = causal_wan_rope( + head_dim=8, + start_frame=3, + frames=2, + height=1, + width=1, + device=torch.device("cpu"), + ) + + assert later.dtype is torch.complex128 + torch.testing.assert_close(later, complete[:, :, 3:5]) + + +def test_scheduler_matches_upstream_precision_and_noise_order() -> None: + """Match LongSana's double x0 conversion and flattened B/T/C noise draw.""" + config = LongSanaFlowMatchSchedulerConfig( + num_inference_steps=4, + shift=7.0, + denoising_timesteps=list(DEFAULT_DENOISING_TIMESTEPS), + warp_denoising_step=False, + ) + scheduler = config.setup() + assert isinstance(scheduler, LongSanaFlowMatchScheduler) + initial = torch.linspace( + -1, + 1, + 1 * 2 * 3 * 2 * 2, + dtype=torch.bfloat16, + ).reshape(1, 2, 3, 2, 2) + + def predict_flow(noisy: Tensor, timestep: Tensor) -> Tensor: + return noisy * 0.125 + timestep / 4096 + + ours_rng = torch.Generator().manual_seed(123) + actual = scheduler.sample(initial, predict_flow, ours_rng) + + reference_rng = torch.Generator().manual_seed(123) + noisy = initial + expected = initial + for index, timestep in enumerate(scheduler.denoising_step_list): + sigma = scheduler.denoising_sigmas[index] + if index > 0: + batch, channels, frames, height, width = noisy.shape + noise = torch.randn( + (batch * frames, channels, height, width), + dtype=noisy.dtype, + generator=reference_rng, + ) + noise = noise.unflatten(0, (batch, frames)).permute(0, 2, 1, 3, 4) + noisy = ((1 - sigma) * expected + sigma * noise).to(initial.dtype) + flow = predict_flow(noisy, timestep.to(initial.dtype)) + expected = (noisy.double() - sigma.double() * flow.double()).to(initial.dtype) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_initial_noise_exposes_tchw_but_preserves_upstream_rng_order() -> None: + """Keep Runtime V2's public layout without changing LongSana's seeded noise.""" + spec = _small_spec() + transformer = LongSanaTransformerConfig( + network=LongSanaNetworkConfig(spec=spec), + dtype=torch.float32, + latent_height=4, + latent_width=4, + first_block_frames=3, + block_frames=2, + ).setup() + assert isinstance(transformer, LongSanaTransformer) + transformer.select_autoregressive_index(0) + cache = transformer.initialize_autoregressive_cache() + + actual_rng = torch.Generator().manual_seed(42) + actual = transformer.initial_noise( + latent_shape=transformer.latent_shape, + rng=actual_rng, + cache=cache, + ) + + expected_rng = torch.Generator().manual_seed(42) + expected = torch.randn( + (spec.latent_channels, 3, 4, 4), + dtype=torch.float32, + generator=expected_rng, + ).permute(1, 0, 2, 3) + + assert actual.shape == (3, spec.latent_channels, 4, 4) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_transformer_cache_tracks_release_block_boundaries() -> None: + """Advance 11 latent frames first and 10 thereafter.""" + cache = LongSanaTransformerCache( + conditioning=LongSanaConditioning( + condition=torch.empty(1, 1, 5, 12), + mask=torch.ones(1, 5), + ), + block_states=[LongSanaBlockState()], + ) + + cache.start(0) + assert cache.active_frames == FIRST_LATENT_BLOCK_FRAMES + cache.finalize(0) + assert cache.start_frame == FIRST_LATENT_BLOCK_FRAMES + + cache.start(1) + assert cache.active_frames == LATENT_BLOCK_FRAMES + cache.finalize(1) + assert cache.start_frame == FIRST_LATENT_BLOCK_FRAMES + LATENT_BLOCK_FRAMES + + +def test_generic_scheduler_would_not_preserve_longsana_rng_layout() -> None: + """Document why the integration uses its narrow scheduler subclass.""" + generic = FlowMatchSchedulerConfig( + num_inference_steps=4, + shift=7.0, + denoising_timesteps=list(DEFAULT_DENOISING_TIMESTEPS), + warp_denoising_step=False, + ).setup() + assert type(generic).__name__ == "FlowMatchScheduler" diff --git a/integrations_v2/longsana/tests/test_t2v_app.py b/integrations_v2/longsana/tests/test_t2v_app.py new file mode 100644 index 000000000..236d6ca67 --- /dev/null +++ b/integrations_v2/longsana/tests/test_t2v_app.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for the LongSana T2V application adapter.""" + +from __future__ import annotations + +import dataclasses +from pathlib import Path +import sys + +import pytest +from t2v.testing import FakeT2VPipelineConfig + +from longsana.apps.t2v.adapter import ( + LONGSANA_T2V_DEFAULTS, + LongSanaT2VApplication, +) +from longsana.impl.constants import ( + DEFAULT_VIDEO_FPS, + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, +) + +if sys.version_info >= (3, 11): + import tomllib +else: + import tomli as tomllib + +pytestmark = pytest.mark.ci_cpu + + +def test_application_advertises_native_one_minute_rollout() -> None: + """Expose release resolution, cadence, and 26-block long-video default.""" + app = LongSanaT2VApplication(pipeline_config=FakeT2VPipelineConfig()) + + description = app.session_desc() + + assert (description.video_width, description.video_height) == ( + DEFAULT_VIDEO_WIDTH, + DEFAULT_VIDEO_HEIGHT, + ) + assert description.frames_per_second_for_step == DEFAULT_VIDEO_FPS + assert LONGSANA_T2V_DEFAULTS.total_blocks == 26 + + +def test_application_rejects_non_native_resolution() -> None: + """Do not silently sample a resolution outside the validated release path.""" + app = LongSanaT2VApplication(pipeline_config=FakeT2VPipelineConfig()) + app.init( + [ + "--prompt", + "A red panda walks through a bamboo forest.", + "--device", + "cpu", + "--total-blocks", + "1", + ] + ) + requested = dataclasses.replace(app.session_desc(), video_width=640) + + with pytest.raises(ValueError, match="requires 832x480"): + app.create_session(requested) + + +def test_application_entry_point_is_registered() -> None: + """Expose the integration through flashdreams-run-v2 discovery.""" + path = Path(__file__).parents[1] / "pyproject.toml" + with path.open("rb") as handle: + project = tomllib.load(handle)["project"] + + assert project["entry-points"]["flashdreams.applications_v2"] == { + "t2v-longsana-2b-480p": "longsana.apps.t2v.adapter:create_app" + } diff --git a/pyproject.toml b/pyproject.toml index 7798b5be8..470b3b18a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,7 @@ extraPaths = [ "integrations_v2/fastvideo_causal_wan22", "integrations_v2/hy_worldplay", "integrations_v2/lingbot", + "integrations_v2/longsana", "integrations_v2/sana_wm", "integrations_v2/self_forcing", "integrations_v2/wan21", @@ -106,6 +107,7 @@ extra-paths = [ "integrations_v2/fastvideo_causal_wan22", "integrations_v2/hy_worldplay", "integrations_v2/lingbot", + "integrations_v2/longsana", "integrations_v2/sana_wm", "integrations_v2/self_forcing", "integrations_v2/wan21", diff --git a/uv.lock b/uv.lock index 85fe4549a..fa34b5273 100644 --- a/uv.lock +++ b/uv.lock @@ -32,6 +32,7 @@ members = [ "flashdreams-imgui-ui-demo", "flashdreams-interactive-drive-v2", "flashdreams-lingbot", + "flashdreams-longsana", "flashdreams-null-model", "flashdreams-omnidreams", "flashdreams-red-screen", @@ -1376,6 +1377,32 @@ requires-dist = [ ] provides-extras = ["dev"] +[[package]] +name = "flashdreams-longsana" +version = "0.1.0" +source = { editable = "integrations_v2/longsana" } +dependencies = [ + { name = "flashdreams" }, + { name = "flashdreams-sana-wm" }, + { name = "flashdreams-t2v" }, + { name = "mediapy" }, +] + +[package.optional-dependencies] +dev = [ + { name = "pytest" }, +] + +[package.metadata] +requires-dist = [ + { name = "flashdreams", editable = "flashdreams" }, + { name = "flashdreams-sana-wm", editable = "integrations_v2/sana_wm" }, + { name = "flashdreams-t2v", editable = "apps/t2v" }, + { name = "mediapy", specifier = ">=1.1" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, +] +provides-extras = ["dev"] + [[package]] name = "flashdreams-null-model" version = "0.1.0" From 8899b65676d1641516b11392284f326b518ee3fc Mon Sep 17 00:00:00 2001 From: Jonathan McCaffrey Date: Sun, 6 Sep 2026 12:55:51 -0700 Subject: [PATCH 2/2] Fix LongSANA session isolation and rollout limits --- integrations_v2/longsana/ACCELERATION.md | 220 ++++++++++++++++++ integrations_v2/longsana/README.md | 4 +- integrations_v2/longsana/apps/t2v/adapter.py | 10 + integrations_v2/longsana/config.py | 2 +- integrations_v2/longsana/impl/constants.py | 8 + integrations_v2/longsana/impl/model.py | 12 +- integrations_v2/longsana/impl/pipeline.py | 29 +-- integrations_v2/longsana/impl/scheduler.py | 2 +- integrations_v2/longsana/impl/transformer.py | 38 +-- integrations_v2/longsana/scripts/benchmark.py | 7 +- .../longsana/scripts/operator_profile.py | 1 - integrations_v2/longsana/tests/test_smoke.py | 64 ++++- .../longsana/tests/test_t2v_app.py | 39 +++- 13 files changed, 371 insertions(+), 65 deletions(-) create mode 100644 integrations_v2/longsana/ACCELERATION.md diff --git a/integrations_v2/longsana/ACCELERATION.md b/integrations_v2/longsana/ACCELERATION.md new file mode 100644 index 000000000..4865faa2e --- /dev/null +++ b/integrations_v2/longsana/ACCELERATION.md @@ -0,0 +1,220 @@ + + +# LongSANA acceleration scope + +This document scopes quality-preserving acceleration of the LongSANA Runtime V2 +pipeline, with emphasis on reuse and extension of `flashdreams.accelerated`. +The first target is steady-state latency at the released 832 x 480 resolution; +the recurrent cache must remain constant-memory and session-local. + +## Baseline and target + +The September 2026 RTX PRO 6000 Blackwell baseline for one steady 40-frame +output block is: + +| stage | latency | share | +| --- | ---: | ---: | +| Four denoising DiT forwards | about 1,620 ms | 40% | +| Required clean cache-commit forward | about 404 ms | 10% | +| FP32 Wan decode | about 2,015 ms | 50% | +| End to end | about 4,029 ms | 100% | + +This is 9.93 output FPS end to end and 19.86 FPS for the DiT plus commit. +Resident allocation is about 8.07 GiB, peak allocation is 37.30 GiB, and the +20-layer recurrent cache remains exactly 152.526855 MiB. + +A useful first milestone is a 1.5x steady-state speedup without a measurable +quality regression or cache growth. A 2x DiT speedup alone is bounded to about +1.33x end to end (roughly 3.03 seconds per block), as is a 2x decoder speedup. +Halving both stages would approach 2.02 seconds, or about 19.8 output FPS. + +## Compatibility map + +| LongSANA path | Existing accelerated component | Status | +| --- | --- | --- | +| Causal recurrent self-attention | None | Requires a new linear-attention primitive | +| Static text cross-attention | `OptimizedMultiHeadAttention` | Needs 112-wide heads, mask support, and an adapter | +| Dense and pointwise projections | `QuantizedNonPersistentLinear` | Opt-in candidate after checkpoint load | +| Wan VAE decoder | None | Use compile/graphs first; accelerated convolution is future work | + +LongSANA self-attention is not softmax attention. It applies a positive ReLU +kernel and updates cumulative `V @ K^T` and key-sum tensors. Replacing it with +`OptimizedMultiHeadAttention` would change the model, so the existing MHA path +is only applicable to cross-attention. + +LongSANA uses 20 heads with head dimension 112. The current +`OptimizedMultiHeadAttention` validates a power-of-two head dimension in +`[16, 256]`, so it rejects this model even though the cross-attention core is +conventional scaled-dot-product attention. The LongSANA path also applies a +prompt padding mask, while the accelerated `compute_kv` and `forward` +interfaces do not accept a mask, and its checkpoint exposes one fused +`kv_linear` rather than separate key and value accessors. + +Reuse therefore requires a checkpoint-preserving adapter, a compatible +cuDNN/SDPA path for 112, and valid-token handling. Compacting valid prompt +tokens before `compute_kv` may preserve the current batch-one semantics, but +must be covered by masked-reference parity tests. + +## Phase 0: lock the benchmark and parity gates + +Before optimizing, retain the current script outputs for first and steady +blocks: + +- end-to-end, diffusion, cache-commit, and decode latency; +- milliseconds per network forward and output FPS; +- resident and peak allocated VRAM; +- recurrent-cache bytes; +- backing-buffer addresses after extending the benchmark harness to record them; +- an operator trace after lazy initialization; +- seeded outputs for a two-block correctness case and a 24-block continuity + case. + +Run each candidate on the same prompt, seed, block count, precision, and GPU. +Report first-block compilation or graph-capture cost separately from warmed +steady-state performance. + +## Phase 1: exact, low-risk work + +### Cache static text K/V + +Each of the 20 cross-attention layers currently recomputes its text K/V +projection and K normalization on every denoising and commit forward even +though the 300-token prompt is fixed for the session. Extend the per-session +cache with one K/V pair per layer, compute it after prompt projection, and reuse +it for every block. + +This mirrors the static cross-attention cache owned by +`OptimizedMultiHeadAttention.compute_kv` and is the cleanest first reuse +point. It preserves checkpoint parameters and attention math. The added memory +is prompt-length dependent but duration independent and must be accounted for +separately from the recurrent self-attention state. + +### Compile and capture the two block shapes + +Measure `compile_network=True` for the DiT and the decoder's existing +`use_compile` and `use_cuda_graph` settings. LongSANA has two stable DiT +shapes: 11 latent frames for block zero and 10 for every steady block. Compile +both shapes. + +The recurrent state tensors are allocated on the first clean cache commit. +Capture the steady-state DiT graph only after those tensors exist, and keep +their addresses stable. Preserve the Wan path's intentional eager first decode +and capture only stable steady decoder calls. The clean commit forward is part +of the model algorithm and must be accelerated rather than skipped. + +### Cache or fuse RoPE construction + +`causal_wan_rope` currently rebuilds complex128 tables and +`_apply_causal_rope` casts Q/K through float64 in every transformer block. +First cache the immutable axis tables by device and slice them by absolute +frame position. Only consider a fused or lower-precision implementation after +numerical and video-quality A/B validation; the current precision matches the +released model. + +## Phase 2: extend `flashdreams.accelerated` + +### Recurrent causal linear attention + +Add a dedicated accelerated primitive rather than adapting softmax MHA. Its +interface must consume current Q/K/V plus the session's cumulative +`value_key` and `key_sum`, produce the normalized output, and optionally +update both states in-place during the clean commit. + +The kernel should fuse the highest-traffic operations where profitable: + +1. Q/K normalization, causal RoPE, and the positive feature map; +2. blockwise K and `V @ K^T` reductions plus the prior recurrent totals; +3. numerator and denominator contractions; +4. in-place final-state writes. + +It must support BF16 inputs, the reference mixed-precision state semantics +(FP32 `value_key` and reference-dtype `key_sum`), head dimension 112, both +11- and 10-frame blocks, and non-mutating denoising forwards. A Torch +implementation should remain selectable for parity and unsupported hardware. + +### Cross-attention compatibility + +Extend the accelerated attention policy so conventional cross-attention can +use `QKVFusionOption.FUSE_KV`, static `compute_kv`, and either +`SDPABackend.CUDNN` or a compatible fallback at head dimension 112. Add +mask or valid-token compaction semantics and a fused-`kv_linear` checkpoint +adapter. Do not pad silently unless tests prove padded normalization and +projection math are equivalent. + +Start with BF16 projections and SDPA. Evaluate projection FP8 separately using +`QuantizationOption`; simple FP8 SDPA is explicitly accuracy-sensitive and +should not be enabled by default. + +### Derived fused and quantized projections + +Preserve the 418-tensor checkpoint schema. `NonPersistentLinear` provides +derived, nonpersistent buffers but is still the same linear operation; replacing +an existing `nn.Linear` with it is not by itself an acceleration. Construct +nonpersistent packed weights only where a kernel consumes a genuinely fused +layout. + +After strict loading, evaluate `QuantizedNonPersistentLinear` as an opt-in for: + +- self-attention QKV (2240 to 6720) and output (2240 to 2240); +- cross-attention Q (2240 to 2240), KV (2240 to 4480), and output; +- timestep/modulation (2240 to 13440); +- GLUMB pointwise 1x1 convolutions. + +The first three groups are already single canonical `nn.Linear` operations, +so their opportunity is quantization or fusion into a larger kernel, not +projection packing alone. + +`linearize_stage1_ffn_for_quant` currently recognizes SANA-WM's +`GLUMBConvTemp`, not `LongSanaCausalGLUMBConvTemp`. Generalize that helper +or add a LongSANA-specific equivalent before applying +`QuantizedNonPersistentLinear`. Benchmark 1x1-convolution linearization +separately, then FP8 and INT8 variants independently so quality and speed +effects are attributable. + +## Phase 3: decoder and memory + +The official FP32 Wan decoder consumes about half the block latency and drives +the 37.30 GiB peak. Existing `flashdreams.accelerated` APIs do not provide a +drop-in VAE or convolution implementation, so first exhaust decoder compilation, +CUDA graphs, scheduling, and allocator reuse. + +Then evaluate reduced-precision decoder convolutions behind an opt-in setting. +If convolution remains the dominant bottleneck, scope accelerated 3D +convolution or a model-specific Wan decoder path as a separate project. Do not +combine decoder precision changes with DiT quantization in the same experiment. + +## Acceptance gates + +Every default-path optimization must retain: + +- strict loading of all 418 tensors and 2,057,553,344 parameters; +- finite outputs and exact first/steady output shapes; +- identical seeded tensors for exact BF16/FP32 modes within an agreed + operator-level tolerance; +- the same cache lifecycle, session isolation, and in-place backing storage; +- duration-independent recurrent state at exactly 152.526855 MiB; +- successful two-block Runtime V2 generation and a 24-block continuity run. + +Quantized or reduced-precision modes additionally require a diverse prompt set +covering people, animals, camera motion, text-like detail, fast motion, and long +scene continuity. Record latent error and stage-level drift, plus VBench, CLIP, +or FVD when available; otherwise retain blinded contact-sheet and full-video +review. Promote a mode to the default only after both quality and performance +gates pass. + +## Deliverables + +1. Exact-path baseline PR: static cross K/V, cached RoPE tables, compile/graph + measurements, and stage-separated benchmark results. +2. Accelerated linear-attention PR: reference/kernel parity tests, in-place + cache tests, and 112-dimension support. +3. Projection optimization PR: nonpersistent fusion followed by separately + gated FP8/INT8 configurations and quality reports. +4. Decoder PR or follow-up RFC: only if compile and graph capture leave the VAE + as the limiting stage. + +Each PR should report speedup against this baseline, not only isolated kernel +throughput, and should include an Amdahl projection for the next bottleneck. diff --git a/integrations_v2/longsana/README.md b/integrations_v2/longsana/README.md index 0ba6c157c..4ec5e2452 100644 --- a/integrations_v2/longsana/README.md +++ b/integrations_v2/longsana/README.md @@ -57,7 +57,9 @@ configuration), and absolute temporal RoPE positions advance across blocks. ## Benchmark and validation -See [VALIDATION.md](VALIDATION.md) for measured results and qualitative review. +See [VALIDATION.md](VALIDATION.md) for measured results and qualitative review, +and [ACCELERATION.md](ACCELERATION.md) for a staged `flashdreams.accelerated` +optimization scope. Run the checked-in diverse-prompt suite: diff --git a/integrations_v2/longsana/apps/t2v/adapter.py b/integrations_v2/longsana/apps/t2v/adapter.py index 51793c9f0..128d3df00 100644 --- a/integrations_v2/longsana/apps/t2v/adapter.py +++ b/integrations_v2/longsana/apps/t2v/adapter.py @@ -29,6 +29,7 @@ DEFAULT_VIDEO_FPS, DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH, + MAX_ROLLOUT_BLOCKS, ) LONGSANA_T2V_DEFAULTS = T2VApplicationDefaults( @@ -57,6 +58,15 @@ def __init__(self, pipeline_config: Any | None = None) -> None: ) super().__init__(defaults=defaults) + def _validate_total_blocks(self, total_blocks: int) -> None: + """Reject rollouts that exceed the released absolute RoPE table.""" + super()._validate_total_blocks(total_blocks) + if total_blocks > MAX_ROLLOUT_BLOCKS: + raise ValueError( + "LongSana supports at most " + f"{MAX_ROLLOUT_BLOCKS} blocks, got {total_blocks}." + ) + def _validate_frame_size(self, session_desc: SessionDesc, pipeline: Any) -> None: """Require the native 832 by 480 release dimensions.""" del pipeline diff --git a/integrations_v2/longsana/config.py b/integrations_v2/longsana/config.py index df21ee012..a3f2dad09 100644 --- a/integrations_v2/longsana/config.py +++ b/integrations_v2/longsana/config.py @@ -18,6 +18,7 @@ from __future__ import annotations import torch +from sana_wm.impl.conditioning import SanaWMTextPromptEncoderConfig from flashdreams.infra.diffusion.model import DiffusionModelConfig from flashdreams.recipes.wan.autoencoder.vae import WanVAEDecoderConfig @@ -29,7 +30,6 @@ from longsana.impl.pipeline import LongSanaPipelineConfig from longsana.impl.scheduler import LongSanaFlowMatchSchedulerConfig from longsana.impl.transformer import LongSanaTransformerConfig -from sana_wm.impl.conditioning import SanaWMTextPromptEncoderConfig PIPELINE_LONGSANA_2B_480P = LongSanaPipelineConfig( name="longsana-2b-480p", diff --git a/integrations_v2/longsana/impl/constants.py b/integrations_v2/longsana/impl/constants.py index 2e9f65521..1bde1cb8c 100644 --- a/integrations_v2/longsana/impl/constants.py +++ b/integrations_v2/longsana/impl/constants.py @@ -32,6 +32,14 @@ LATENT_BLOCK_FRAMES = 10 """Steady-state number of latent frames generated per AR block.""" +MAX_ROPE_POSITION = 1024 +"""Largest absolute latent position supported by the released RoPE table.""" + +MAX_ROLLOUT_BLOCKS = ( + 1 + (MAX_ROPE_POSITION - FIRST_LATENT_BLOCK_FRAMES) // LATENT_BLOCK_FRAMES +) +"""Maximum complete rollout that fits the released absolute RoPE table.""" + MOTION_SCORE = 10 """Motion-score suffix used during LongSana self-forcing post-training.""" diff --git a/integrations_v2/longsana/impl/model.py b/integrations_v2/longsana/impl/model.py index 0cfdfff85..08770eba4 100644 --- a/integrations_v2/longsana/impl/model.py +++ b/integrations_v2/longsana/impl/model.py @@ -23,9 +23,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch import Tensor - -from flashdreams.infra.config import InstantiateConfig from sana_wm.impl.stage1_model import ( RMSNorm, SanaWMStage1Spec, @@ -33,6 +30,10 @@ TextEmbedder, TimestepEmbedder, ) +from torch import Tensor + +from flashdreams.infra.config import InstantiateConfig +from longsana.impl.constants import MAX_ROPE_POSITION LONGSANA_SPEC = SanaWMStage1Spec( latent_channels=16, @@ -502,7 +503,7 @@ def causal_wan_rope( height: int, width: int, device: torch.device, - max_sequence_length: int = 1024, + max_sequence_length: int = MAX_ROPE_POSITION, ) -> Tensor: """Build upstream-compatible complex128 Wan RoPE at absolute frame positions.""" end_frame = start_frame + frames @@ -512,7 +513,8 @@ def causal_wan_rope( ) if max(end_frame, height, width) > max_sequence_length: raise ValueError( - "LongSana RoPE position exceeds the released 1024-position table: " + "LongSana RoPE position exceeds the released " + f"{MAX_ROPE_POSITION}-position table: " f"end_frame={end_frame}, height={height}, width={width}." ) diff --git a/integrations_v2/longsana/impl/pipeline.py b/integrations_v2/longsana/impl/pipeline.py index bfcb4da25..e28a57a2a 100644 --- a/integrations_v2/longsana/impl/pipeline.py +++ b/integrations_v2/longsana/impl/pipeline.py @@ -21,7 +21,11 @@ from typing import Any, cast import torch -from torch import Tensor +from sana_wm.impl.conditioning import ( + SanaWMTextPromptEncoder, + SanaWMTextPromptEncoderConfig, + SanaWMTextPromptRequest, +) from flashdreams.infra.pipeline import ( StreamInferencePipeline, @@ -36,11 +40,6 @@ LongSanaConditioning, LongSanaTransformer, ) -from sana_wm.impl.conditioning import ( - SanaWMTextPromptEncoder, - SanaWMTextPromptEncoderConfig, - SanaWMTextPromptRequest, -) @dataclass(kw_only=True) @@ -116,24 +115,6 @@ def initialize_cache( transformer_context={"conditioning": conditioning}, ) - @torch.no_grad() - def generate( - self, - autoregressive_index: int, - cache: StreamInferencePipelineCache, - input: Any = None, - ) -> Tensor: - """Select LongSana's first/steady block shape and generate one chunk.""" - transformer = self.diffusion_model.transformer - if not isinstance(transformer, LongSanaTransformer): - raise TypeError("LongSanaPipeline requires LongSanaTransformer.") - transformer.select_autoregressive_index(autoregressive_index) - return super().generate( - autoregressive_index=autoregressive_index, - cache=cache, - input=input, - ) - def close(self) -> None: """Release prompt and generator runtimes held by the resident pipeline.""" self.prompt_encoder.release_runtime() diff --git a/integrations_v2/longsana/impl/scheduler.py b/integrations_v2/longsana/impl/scheduler.py index b06d7883b..638bce387 100644 --- a/integrations_v2/longsana/impl/scheduler.py +++ b/integrations_v2/longsana/impl/scheduler.py @@ -23,11 +23,11 @@ from torch import Tensor from tqdm import tqdm +from flashdreams.infra.diffusion.scheduler import FlowPredictor from flashdreams.infra.diffusion.scheduler.fm import ( FlowMatchScheduler, FlowMatchSchedulerConfig, ) -from flashdreams.infra.diffusion.scheduler import FlowPredictor @dataclass(kw_only=True) diff --git a/integrations_v2/longsana/impl/transformer.py b/integrations_v2/longsana/impl/transformer.py index f98d1c40f..0003e1e80 100644 --- a/integrations_v2/longsana/impl/transformer.py +++ b/integrations_v2/longsana/impl/transformer.py @@ -41,6 +41,7 @@ FIRST_LATENT_BLOCK_FRAMES, LATENT_BLOCK_FRAMES, LONGSANA_CHECKPOINT_PATH, + MAX_ROPE_POSITION, ) from longsana.impl.model import ( LongSanaBlockState, @@ -102,10 +103,18 @@ def start(self, autoregressive_index: int) -> None: f"Expected LongSana AR step {self.next_index}, " f"got {autoregressive_index}." ) - self.active_index = autoregressive_index - self.active_frames = ( + active_frames = ( self.first_block_frames if autoregressive_index == 0 else self.block_frames ) + end_frame = self.start_frame + active_frames + if end_frame > MAX_ROPE_POSITION: + raise ValueError( + "LongSana rollout exceeds the released " + f"{MAX_ROPE_POSITION}-position RoPE table: " + f"block {autoregressive_index} ends at latent frame {end_frame}." + ) + self.active_index = autoregressive_index + self.active_frames = active_frames def finalize(self, autoregressive_index: int) -> None: """Advance absolute positions after the clean-timestep cache update.""" @@ -167,29 +176,17 @@ def __init__(self, config: LongSanaTransformerConfig) -> None: self.config = config self._dummy = nn.Parameter(torch.empty(0, dtype=config.dtype)) self._model_built = False - self._active_autoregressive_index = 0 @property def latent_shape(self) -> tuple[int, ...]: - """Return the raw Wan latent shape for the selected AR block.""" - frames = ( - self.config.first_block_frames - if self._active_autoregressive_index == 0 - else self.config.block_frames - ) + """Return the nominal first-block shape used by the diffusion interface.""" return ( - frames, + self.config.first_block_frames, self.config.network.spec.latent_channels, self.config.latent_height, self.config.latent_width, ) - def select_autoregressive_index(self, autoregressive_index: int) -> None: - """Select the block shape before Runtime V2 allocates initial noise.""" - if autoregressive_index < 0: - raise ValueError("LongSana autoregressive_index must be non-negative.") - self._active_autoregressive_index = autoregressive_index - def initialize_autoregressive_cache( self, *, @@ -215,8 +212,13 @@ def initial_noise( input: Any = None, ) -> Tensor: """Draw upstream B/C/T/H/W noise and expose it as Runtime V2 T/C/H/W.""" - del cache, input - frames, channels, height, width = latent_shape + del input + if cache.active_index is None or cache.active_frames <= 0: + raise RuntimeError( + "LongSana cache must start a block before drawing noise." + ) + _, channels, height, width = latent_shape + frames = cache.active_frames noise = torch.randn( (channels, frames, height, width), device=self.device, diff --git a/integrations_v2/longsana/scripts/benchmark.py b/integrations_v2/longsana/scripts/benchmark.py index 031e1b5f4..426ca5d3f 100644 --- a/integrations_v2/longsana/scripts/benchmark.py +++ b/integrations_v2/longsana/scripts/benchmark.py @@ -6,19 +6,18 @@ from __future__ import annotations import argparse -from dataclasses import asdict, dataclass -from datetime import datetime, timezone import json -from pathlib import Path import platform import statistics import subprocess import time +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path from typing import Any import mediapy as media import torch - from longsana.config import PIPELINE_LONGSANA_2B_480P from longsana.impl.constants import DEFAULT_VIDEO_FPS from longsana.impl.transformer import LongSanaTransformerCache diff --git a/integrations_v2/longsana/scripts/operator_profile.py b/integrations_v2/longsana/scripts/operator_profile.py index 1ffea4671..0a81faa6a 100644 --- a/integrations_v2/longsana/scripts/operator_profile.py +++ b/integrations_v2/longsana/scripts/operator_profile.py @@ -10,7 +10,6 @@ from pathlib import Path import torch - from longsana.config import PIPELINE_LONGSANA_2B_480P from longsana.impl.transformer import LongSanaTransformerCache diff --git a/integrations_v2/longsana/tests/test_smoke.py b/integrations_v2/longsana/tests/test_smoke.py index 68aae2116..6de6943e5 100644 --- a/integrations_v2/longsana/tests/test_smoke.py +++ b/integrations_v2/longsana/tests/test_smoke.py @@ -17,13 +17,8 @@ from __future__ import annotations -import torch -from torch import Tensor import pytest - -from flashdreams.infra.diffusion.model import DiffusionModel -from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig -from flashdreams.recipes.wan.autoencoder.vae import WanVAEDecoderConfig +import torch from longsana.config import LONGSANA_CONFIGS, PIPELINE_LONGSANA_2B_480P from longsana.impl.constants import ( DEFAULT_DENOISING_TIMESTEPS, @@ -31,6 +26,7 @@ LATENT_BLOCK_FRAMES, LONGSANA_REVISION, LONGSANA_TEXT_CONFIG_PATH, + MAX_ROPE_POSITION, SANA_VIDEO_REVISION, ) from longsana.impl.model import ( @@ -52,6 +48,11 @@ ) from sana_wm.impl.stage1_model import SanaWMStage1Spec from sana_wm.impl.transformer import _load_inference_config +from torch import Tensor + +from flashdreams.infra.diffusion.model import DiffusionModel +from flashdreams.infra.diffusion.scheduler.fm import FlowMatchSchedulerConfig +from flashdreams.recipes.wan.autoencoder.vae import WanVAEDecoderConfig pytestmark = pytest.mark.ci_cpu @@ -275,8 +276,8 @@ def test_initial_noise_exposes_tchw_but_preserves_upstream_rng_order() -> None: block_frames=2, ).setup() assert isinstance(transformer, LongSanaTransformer) - transformer.select_autoregressive_index(0) cache = transformer.initialize_autoregressive_cache() + cache.start(0) actual_rng = torch.Generator().manual_seed(42) actual = transformer.initial_noise( @@ -296,6 +297,41 @@ def test_initial_noise_exposes_tchw_but_preserves_upstream_rng_order() -> None: torch.testing.assert_close(actual, expected, rtol=0, atol=0) +def test_initial_noise_uses_each_sessions_active_block_shape() -> None: + """Keep first and steady block lengths isolated across interleaved sessions.""" + spec = _small_spec() + transformer = LongSanaTransformerConfig( + network=LongSanaNetworkConfig(spec=spec), + dtype=torch.float32, + latent_height=4, + latent_width=4, + first_block_frames=3, + block_frames=2, + ).setup() + assert isinstance(transformer, LongSanaTransformer) + first_session = transformer.initialize_autoregressive_cache() + steady_session = transformer.initialize_autoregressive_cache() + + steady_session.start(0) + steady_session.finalize(0) + steady_session.start(1) + first_session.start(0) + + steady = transformer.initial_noise( + latent_shape=transformer.latent_shape, + rng=torch.Generator().manual_seed(1), + cache=steady_session, + ) + first = transformer.initial_noise( + latent_shape=transformer.latent_shape, + rng=torch.Generator().manual_seed(2), + cache=first_session, + ) + + assert steady.shape == (2, spec.latent_channels, 4, 4) + assert first.shape == (3, spec.latent_channels, 4, 4) + + def test_transformer_cache_tracks_release_block_boundaries() -> None: """Advance 11 latent frames first and 10 thereafter.""" cache = LongSanaTransformerCache( @@ -317,6 +353,20 @@ def test_transformer_cache_tracks_release_block_boundaries() -> None: assert cache.start_frame == FIRST_LATENT_BLOCK_FRAMES + LATENT_BLOCK_FRAMES +def test_transformer_cache_rejects_rope_overflow_before_generation() -> None: + """Protect direct pipeline callers from exceeding absolute RoPE positions.""" + cache = LongSanaTransformerCache( + start_frame=MAX_ROPE_POSITION - LATENT_BLOCK_FRAMES + 1, + next_index=1, + ) + + with pytest.raises(ValueError, match="exceeds.*RoPE table"): + cache.start(1) + + assert cache.active_index is None + assert cache.active_frames == 0 + + def test_generic_scheduler_would_not_preserve_longsana_rng_layout() -> None: """Document why the integration uses its narrow scheduler subclass.""" generic = FlowMatchSchedulerConfig( diff --git a/integrations_v2/longsana/tests/test_t2v_app.py b/integrations_v2/longsana/tests/test_t2v_app.py index 236d6ca67..4dd71f468 100644 --- a/integrations_v2/longsana/tests/test_t2v_app.py +++ b/integrations_v2/longsana/tests/test_t2v_app.py @@ -18,12 +18,10 @@ from __future__ import annotations import dataclasses -from pathlib import Path import sys +from pathlib import Path import pytest -from t2v.testing import FakeT2VPipelineConfig - from longsana.apps.t2v.adapter import ( LONGSANA_T2V_DEFAULTS, LongSanaT2VApplication, @@ -32,7 +30,9 @@ DEFAULT_VIDEO_FPS, DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH, + MAX_ROLLOUT_BLOCKS, ) +from t2v.testing import FakeT2VPipelineConfig if sys.version_info >= (3, 11): import tomllib @@ -75,6 +75,39 @@ def test_application_rejects_non_native_resolution() -> None: app.create_session(requested) +def test_application_accepts_maximum_rope_bounded_rollout() -> None: + """Accept the last complete rollout that fits the absolute RoPE table.""" + app = LongSanaT2VApplication(pipeline_config=FakeT2VPipelineConfig()) + + app.init( + [ + "--prompt", + "A red panda walks through a bamboo forest.", + "--device", + "cpu", + "--total-blocks", + str(MAX_ROLLOUT_BLOCKS), + ] + ) + + +def test_application_rejects_rollout_beyond_rope_table() -> None: + """Fail before model setup rather than after a long partial generation.""" + app = LongSanaT2VApplication(pipeline_config=FakeT2VPipelineConfig()) + + with pytest.raises(ValueError, match="at most 102 blocks"): + app.init( + [ + "--prompt", + "A red panda walks through a bamboo forest.", + "--device", + "cpu", + "--total-blocks", + str(MAX_ROLLOUT_BLOCKS + 1), + ] + ) + + def test_application_entry_point_is_registered() -> None: """Expose the integration through flashdreams-run-v2 discovery.""" path = Path(__file__).parents[1] / "pyproject.toml"