diff --git a/docs/source/_static/performance/sana_wm/perf-bf16-0727.md b/docs/source/_static/performance/sana_wm/perf-bf16-0727.md new file mode 100644 index 000000000..64ba5dca3 --- /dev/null +++ b/docs/source/_static/performance/sana_wm/perf-bf16-0727.md @@ -0,0 +1,5 @@ +# SANA-WM BF16 Benchmark Data (ms) + +| device | official | flashdreams | +| --- | ---: | ---: | +| GB300 | 45823.92 | 42399.31 | diff --git a/docs/source/index.rst b/docs/source/index.rst index 16530b788..49971cc80 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -261,6 +261,14 @@ invocation, the checkpoint source, and the per-implementation knobs. Bidirectional Cosmos-Predict2 reference implementations (T2V / I2V, 2B). + .. grid-item-card:: SANA-WM (bidirectional) + :class-card: fd-feature + :link: models/sana_wm + :link-type: doc + + Bidirectional camera-controlled world model (Stage-1 DiT + LTX-2 + refiner, 2.6B). + .. Master toctree: one flat entry per top-level navbar item. Order here = order in the navbar. diff --git a/docs/source/models/index.rst b/docs/source/models/index.rst index 78d767002..3a68a014a 100644 --- a/docs/source/models/index.rst +++ b/docs/source/models/index.rst @@ -28,6 +28,7 @@ Models flashvsr hy_worldplay lingbot_world + sana_wm wan21 FlashDreams runs a growing family of world and video models (text-to-video, @@ -172,6 +173,14 @@ uses, and the settings you can tune. Bidirectional Cosmos-Predict2 reference implementations (T2V / I2V, 2B). + .. grid-item-card:: SANA-WM + :class-card: fd-feature + :link: /models/sana_wm + :link-type: doc + + Bidirectional camera-controlled world model (Stage-1 DiT + LTX-2 + refiner, 2.6B). + .. container:: fd-eyebrow Super-resolution diff --git a/docs/source/models/sana_wm.rst b/docs/source/models/sana_wm.rst new file mode 100644 index 000000000..a25299bda --- /dev/null +++ b/docs/source/models/sana_wm.rst @@ -0,0 +1,173 @@ +.. 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. + +SANA-WM +=================================== + +.. container:: fd-cta-row + + .. button-link:: https://nvlabs.github.io/Sana/ + :color: primary + + Project page + + .. button-link:: https://arxiv.org/abs/2410.10629 + :color: primary + + arXiv paper + + .. button-link:: https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional + :color: primary + + Model page + + .. button-link:: https://github.com/NVlabs/Sana + :color: primary + + Official code + +SANA-WM is a 2.6B bidirectional, camera-controlled world model from the +`NVlabs/Sana `_ family. Given a first frame, a +text prompt, and a camera trajectory it renders a video clip in a single +bidirectional pass. FlashDreams runs it through the ``sana-wm-bidirectional`` +runner, which pairs a native Stage-1 DiT (loading the public SANA-WM checkpoint +directly) with an LTX-2 refiner for the final decode. + +Requirements +------------ + +- **PyTorch**: >= 2.9. +- **Precision**: BF16 by default. FP8 Stage-1/refiner inference is available on + Hopper or newer GPUs (``sm_90+``) and FP4 on Blackwell (``sm_100+``); both + lower the memory footprint relative to the BF16 default. + +Installation +------------ + +.. code-block:: bash + + # from the repo root + uv sync --package flashdreams-sana-wm --extra dev + +Running the method +------------------ + +To run SANA-WM, launch the ``sana-wm-bidirectional`` runner with a first-frame +image, a prompt, and a camera trajectory: + +.. code-block:: bash + + uv run flashdreams-run sana-wm-bidirectional \ + --image-path ../Sana/asset/sana_wm/demo_0.png \ + --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ + --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ + --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ + --num-frames 161 \ + --output-dir outputs/sana_wm_bf16 + +The demo image, prompt, camera, and intrinsics files ship with the SANA-WM +release; any inputs with matching shapes work as well. + +Optional inputs and knobs +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +- ``--intrinsics-path`` is optional. When omitted, intrinsics are derived from + the first-frame size, assuming a centered principal point and a horizontal + field of view of ``--intrinsics-hfov-deg`` (default ``90``, matching the demo + intrinsics). +- ``--camera-path`` can be replaced by an ``--action`` DSL string, so a minimal + run needs only an image and a prompt: + + .. code-block:: bash + + uv run flashdreams-run sana-wm-bidirectional \ + --image-path my_frame.png \ + --prompt "a scene description; describe the world's own motion" \ + --action "w-100,dw-60,w-101" \ + --num-frames 161 \ + --output-dir outputs/mine + +- ``--no-refiner True`` runs Stage-1 only, for diagnostics. +- Quantized Stage-1/refiner inference is opt-in via ``--stage1-precision`` / + ``--refiner-precision`` (``fp8`` on ``sm_90+``, ``fp4`` on ``sm_100+``), with + ``--quant-backend`` (``torch-fp8`` / ``torch-fp4``) to force a backend. + +To inspect all supported CLI arguments and their default values, run: + +.. code-block:: bash + + uv run flashdreams-run sana-wm-bidirectional --help + +What to expect +-------------- + +- **Model checkpoint**: pulled from + ``huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional`` on first run and + cached under ``$HF_HOME``. The LTX-2 refiner and VAE weights are fetched from + their respective ``diffusers`` repositories on first use. +- **First launch**: a few minutes for the download and warmup; subsequent + launches reuse the caches. +- **Outputs**: ``outputs//sana_wm_generated.mp4``. + +See :doc:`/developer_guides/inference_pipeline_overview` for what one pass does +end-to-end. + +Profiling benchmark +------------------- + +Here is the BF16 profiling benchmark on post-load generation latency per +generated clip for FlashDreams SANA-WM compared to the +`official SANA-WM implementation `_ under +matched settings. On GB300, FlashDreams is about 7% faster than the official +implementation at median clip latency (full pipeline: Stage-1 DiT + LTX-2 +refiner, 10 measured runs). + +.. raw:: html + +
+
+
+

+ This chart shows post-load generation latency per generated clip in milliseconds for a 121-frame + full-pipeline BF16 run (Stage-1 DiT + LTX-2 refiner). The measured row used an NVIDIA GB300. + Model construction, checkpoint loading, video writing, and frame dumps are outside the timing boundary. + For the official SANA-WM implementation, see + this instruction. +

+
+
+ + +Citation +-------- + +If you use SANA-WM, please cite the original SANA work: + +.. code-block:: bibtex + + @misc{xie2024sana, + title={SANA: Efficient High-Resolution Image Synthesis with Linear Diffusion Transformers}, + author={Enze Xie and Junsong Chen and Junyu Chen and Han Cai and Haotian Tang and Yujun Lin and Zhekai Zhang and Muyang Li and Ligeng Zhu and Yao Lu and Song Han}, + year={2024}, + eprint={2410.10629}, + archivePrefix={arXiv}, + primaryClass={cs.CV} + } diff --git a/flashdreams/flashdreams/infra/diffusion/model/base.py b/flashdreams/flashdreams/infra/diffusion/model/base.py index 02dbb909c..57b4e5171 100644 --- a/flashdreams/flashdreams/infra/diffusion/model/base.py +++ b/flashdreams/flashdreams/infra/diffusion/model/base.py @@ -173,19 +173,19 @@ def generate( self.latent_shape, device=self.device, dtype=self.dtype ) dummy_latent = self.transformer.unpatchify_and_maybe_gather_cp(dummy_latent) - initial_noise = torch.randn( - dummy_latent.shape, - device=self.device, - dtype=self.dtype, - generator=self.rng, + initial_noise = self.transformer.initial_noise( + latent_shape=tuple(dummy_latent.shape), + rng=self.rng, + cache=cache, + input=input, ) else: - initial_noise = torch.randn( - self.latent_shape, - device=self.device, - dtype=self.dtype, - generator=self.rng, + initial_noise = self.transformer.initial_noise( + latent_shape=self.latent_shape, + rng=self.rng, + cache=cache, + input=input, ) def predict_flow(noisy_latent: Tensor, timestep: Tensor) -> Tensor: diff --git a/flashdreams/flashdreams/infra/diffusion/transformer/base.py b/flashdreams/flashdreams/infra/diffusion/transformer/base.py index c904e3d3a..1ff10749c 100644 --- a/flashdreams/flashdreams/infra/diffusion/transformer/base.py +++ b/flashdreams/flashdreams/infra/diffusion/transformer/base.py @@ -169,6 +169,37 @@ def initialize_autoregressive_cache(self, **context: Any) -> TransformerCacheT: """ return cast("TransformerCacheT", TransformerAutoregressiveCache()) + def initial_noise( + self, + *, + latent_shape: tuple[int, ...], + rng: torch.Generator | None, + cache: TransformerCacheT, + input: Any = None, + ) -> Tensor: + """Draw the starting latent for one denoising loop. + + Default is Gaussian noise in ``latent_shape``. Model integrations can + override this to seed from encoder outputs, pin I2V frames, or attach + per-step state to the cache while still using ``DiffusionModel``. + + Args: + latent_shape: Shape requested by the diffusion model. + rng: Per-model generator on ``self.device``, or ``None``. + cache: Per-rollout AR cache. + input: Same patchified encoder output passed to ``predict_flow``. + + Returns: + Initial latent tensor for scheduler sampling. + """ + del cache, input + return torch.randn( + latent_shape, + device=self.device, + dtype=self.dtype, + generator=rng, + ) + def postprocess_clean_latent( self, clean_latent: Tensor, diff --git a/integrations/sana/README.md b/integrations/sana/README.md new file mode 100644 index 000000000..2690a0e4e --- /dev/null +++ b/integrations/sana/README.md @@ -0,0 +1,171 @@ + + +# `sana_wm` + +FlashDreams SANA-WM integration for +[SANA-WM](https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional), +the 2.6B bidirectional camera-controlled world model released from NVlabs/Sana. + +The `sana-wm-bidirectional` runner uses FlashDreams config, runner, pipeline, +diffusion-model, scheduler, transformer, camera-conditioning, VAE decode, and +output-writing boundaries. The Stage-1 DiT module in this package loads the +public SANA-WM checkpoint directly. + +Current scope: + +| component | status | +| --- | --- | +| Stage-1 BF16 | FlashDreams DiT execution. | +| Stage-1 FP8 | PyTorch `_scaled_mm` backend. | +| Stage-1 FP4 | Triton quantization plus PyTorch `_scaled_mm`. | +| VAE decode | Direct `diffusers` LTX2 VAE use with upstream-matched tiling. | +| LTX-2 refiner | Direct `diffusers` LTX-2 transformer and Gemma connector use. | + +## Runner + +| slug | description | +| --- | --- | +| `sana-wm-bidirectional` | SANA-WM Stage-1 + LTX-2 refiner runner. | + +The FlashDreams package is named `sana_wm`. + +## Setup + +Install FlashDreams and the SANA integration into an environment with the +project's GPU runtime dependencies: + +```bash +uv sync --package flashdreams-sana-wm --extra dev +``` + +The examples below use the public demo image, prompt, camera, and intrinsics +files from the SANA-WM release. Equivalent inputs with the same shapes work as +well. + +## Run + +```bash +uv run flashdreams-run sana-wm-bidirectional \ + --image-path ../Sana/asset/sana_wm/demo_0.png \ + --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ + --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ + --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ + --num-frames 161 \ + --output-dir outputs/sana_wm_bf16 +``` + +Expected output: + +```text +outputs/sana_wm_bf16/sana-wm-bidirectional.mp4 +``` + +`--intrinsics-path` is optional. When omitted, intrinsics are derived from the +first-frame size assuming a centered principal point and a horizontal field of +view of `--intrinsics-hfov-deg` (default `90`, matching the demo intrinsics). +Likewise `--camera-path` may be replaced by an `--action` DSL string, which is +repeated or truncated to the requested snapped frame count. Explicit camera +paths are fitted to the same frame count instead of capping the video length. A +minimal run therefore needs only an image and a prompt: + +```bash +uv run flashdreams-run sana-wm-bidirectional \ + --image-path my_frame.png \ + --prompt "a scene description; describe the world's own motion" \ + --action "w-100,dw-60,w-101" \ + --num-frames 161 \ + --output-dir outputs/mine +``` + +For Stage-1-only diagnostics, add `--no-refiner True`. + +## BF16, FP8, and FP4 + +The runner defaults to BF16. Quantized Stage-1 and refiner inference are opt-in +and use Torch-based backends by default. + +| option | hardware | backend | +| --- | --- | --- | +| `--stage1-precision fp8` | Hopper or newer (`sm_90+`) | E4M3 `_scaled_mm` | +| `--stage1-precision fp4` | Blackwell (`sm_100+`) | Triton NVFP4 plus `_scaled_mm` | +| `--quant-backend torch-fp8` | Hopper or newer (`sm_90+`) | force FP8 | +| `--quant-backend torch-fp4` | Blackwell (`sm_100+`) | force FP4 | + +Stage-1 FP8 and FP4 follow the upstream precision CLI scope: self-attention, +cross-attention, and linearized FFN pointwise projections. This integration uses +Torch/Triton scaled-MM replacements instead of TransformerEngine. FP4 uses +NVFP4 W4A4 GEMMs with tiled 16-wide Hadamard rotation enabled, two-level +per-tensor-plus-block scaling, and 2D 16x16 block scaling for weights, matching +the upstream recipe behavior that keeps Stage-1 camera/action conditioning +coherent. + +FP8 smoke: + +```bash +uv run flashdreams-run sana-wm-bidirectional \ + --image-path ../Sana/asset/sana_wm/demo_0.png \ + --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ + --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ + --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ + --num-frames 161 \ + --output-dir outputs/sana_wm_fp8 \ + --stage1-precision fp8 \ + --refiner-precision fp8 +``` + +FP4 smoke: + +```bash +uv run flashdreams-run sana-wm-bidirectional \ + --image-path ../Sana/asset/sana_wm/demo_0.png \ + --prompt-path ../Sana/asset/sana_wm/demo_0.txt \ + --camera-path ../Sana/asset/sana_wm/demo_0_pose.npy \ + --intrinsics-path ../Sana/asset/sana_wm/demo_0_intrinsics.npy \ + --num-frames 161 \ + --output-dir outputs/sana_wm_fp4 \ + --stage1-precision fp4 \ + --refiner-precision fp4 +``` + +## Parity and benchmark + +The upstream comparison harness lives under `tests/parity_check/`, matching the +pattern used by the other benchmarked FlashDreams integrations. + +```bash +cd integrations/sana/tests/parity_check + +# Upstream + FlashDreams parity artifacts and frame diff. +bash run.sh + +# Matched upstream-vs-FlashDreams benchmark report. +DEVICE_LABEL="RTX PRO 6000 Blackwell" bash bench.sh +``` + +`bench.sh` writes `outputs/bench/bench.md` for review and `outputs/bench/perf.md` +for chart-ready data. Both use the same benchmark metric: post-load generation +latency per generated clip. SANA-WM renders each requested bidirectional clip in +one generation pass, rather than as independently timed frames. Stage-1 DiT, +conditioning/encode, VAE decode, optional refiner, memory, and frame-normalized +diagnostic rows are reported as breakdowns of that metric. + +For precision sweeps, copy the per-precision chart data from +`outputs/bench//perf.md` into the docs. Each model-card chart should +cover one precision and use GPU/device as the differentiating row, matching the +other FlashDreams model-card benchmark charts. + +## Tests + +CPU-safe tests cover import, config boundaries, action parsing, intrinsics, +camera conditioning, Stage-1 checkpoint schema, Stage-1 CPU forward shape, VAE +tiling, and low-precision backend selection: + +```bash +uv run --extra dev pytest integrations/sana/tests/test_smoke.py +``` + +GPU generation checks are heavyweight manual workflows and should stay out of +`ci_cpu`. diff --git a/integrations/sana/pyproject.toml b/integrations/sana/pyproject.toml new file mode 100644 index 000000000..2d63b6f77 --- /dev/null +++ b/integrations/sana/pyproject.toml @@ -0,0 +1,54 @@ +# 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. + +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "flashdreams-sana-wm" +version = "0.1.0" +description = "SANA-WM bidirectional world-model integration for flashdreams." +readme = "README.md" +requires-python = ">=3.10" +dependencies = [ + "accelerate>=1.0", + "diffusers>=0.36", + "flashdreams", + "imageio[ffmpeg]>=2.31", + "Pillow>=10", + "PyYAML>=6.0", + "safetensors>=0.5", + "torchvision>=0.26", + "transformers>=5.0,<6", +] + +[tool.uv.sources] +flashdreams = { workspace = true } + +[project.optional-dependencies] +dev = [ + "pytest>=8.0", +] + +[project.entry-points."flashdreams.runner_configs"] +"sana-wm-bidirectional" = "sana_wm.config:RUNNER_SANA_WM_BIDIRECTIONAL" + +[tool.setuptools.packages.find] +include = ["sana_wm*"] +exclude = ["tests"] + +[tool.uv] +managed = true diff --git a/integrations/sana/sana_wm/__init__.py b/integrations/sana/sana_wm/__init__.py new file mode 100644 index 000000000..c59615ca6 --- /dev/null +++ b/integrations/sana/sana_wm/__init__.py @@ -0,0 +1,23 @@ +# 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. + +"""SANA-WM runner for FlashDreams.""" + +from sana_wm.config import RUNNER_CONFIGS, RUNNER_SANA_WM_BIDIRECTIONAL + +__all__ = [ + "RUNNER_CONFIGS", + "RUNNER_SANA_WM_BIDIRECTIONAL", +] diff --git a/integrations/sana/sana_wm/_tools.py b/integrations/sana/sana_wm/_tools.py new file mode 100644 index 000000000..7562bfcba --- /dev/null +++ b/integrations/sana/sana_wm/_tools.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. + +"""Local file and checkpoint helpers for the SANA-WM integration.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +import torch + +HF_URI_SCHEME = "hf://" + + +def resolve_hf_path(path: str | Path) -> str: + """Resolve a local path or ``hf://owner/repo/subpath`` URI to a local path.""" + path_str = str(path) + if not path_str or Path(path_str).exists(): + return path_str + if not path_str.startswith(HF_URI_SCHEME): + return path_str + + from huggingface_hub import snapshot_download + + parts = path_str[len(HF_URI_SCHEME) :].split("/", 2) + if len(parts) < 2 or not parts[0] or not parts[1]: + raise ValueError( + f"Invalid HF path {path_str!r}; expected hf:///[/]." + ) + repo_id = f"{parts[0]}/{parts[1]}" + subpath = parts[2] if len(parts) > 2 else "" + allow_patterns = None + if subpath: + allow_patterns = [subpath, f"{subpath}/*", f"{subpath}/**"] + local_root = snapshot_download(repo_id=repo_id, allow_patterns=allow_patterns) + return os.path.join(local_root, subpath) if subpath else local_root + + +def find_model(model_name: str) -> dict[str, Any]: + """Load a SANA checkpoint from a local path or resolved HF artefact.""" + resolved = resolve_hf_path(model_name) + if not os.path.isfile(resolved): + raise FileNotFoundError(f"Could not find SANA checkpoint at {resolved}") + + if resolved.endswith(".safetensors"): + import safetensors.torch + + return {"state_dict": safetensors.torch.load_file(resolved, device="cpu")} + if resolved.endswith(".safetensors.index.json"): + import safetensors.torch + + with open(resolved, encoding="utf-8") as handle: + index = json.load(handle)["weight_map"] + state_dict = {} + for shard in sorted(set(index.values())): + shard_path = os.path.join(os.path.dirname(resolved), shard) + state_dict.update(safetensors.torch.load_file(shard_path, device="cpu")) + return {"state_dict": state_dict} + return torch.load(resolved, map_location=lambda storage, _loc: storage) + + +def hf_download_or_fpath(path: str | Path) -> str: + """Compatibility alias retained for local helper call sites.""" + return resolve_hf_path(path) + + +__all__ = [ + "find_model", + "hf_download_or_fpath", + "resolve_hf_path", +] diff --git a/integrations/sana/sana_wm/camera.py b/integrations/sana/sana_wm/camera.py new file mode 100644 index 000000000..063345910 --- /dev/null +++ b/integrations/sana/sana_wm/camera.py @@ -0,0 +1,610 @@ +# 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. + +"""SANA-WM camera-pose DSL, intrinsics, and Plucker conditioning helpers.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import torch + +from sana_wm.constants import ( + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, + SANA_WM_VAE_SPATIAL_COMPRESSION, + SANA_WM_VAE_TEMPORAL_COMPRESSION, +) + +FPS = 16 +"""Camera-control integration frame rate.""" + +DEFAULT_TRANSLATION_SPEED = 0.025 +"""Default per-frame translation magnitude used by SANA-WM.""" + +DEFAULT_ROTATION_SPEED_DEG = 0.6 +"""Default per-frame rotation magnitude in degrees.""" + +DEFAULT_PITCH_LIMIT_DEG = 60.0 +"""Maximum absolute camera pitch in degrees.""" + +TAU_PRESS = 0.45 +"""Velocity smoothing time constant when a key is pressed.""" + +TAU_COAST = 1.0 +"""Velocity smoothing time constant when controls are released.""" + +DSL_KEY_TO_CONTROL: dict[str, str] = { + "w": "forward", + "s": "back", + "a": "yaw_left", + "d": "yaw_right", + "i": "pitch_up", + "k": "pitch_down", + "j": "strafe_left", + "l": "strafe_right", +} +"""Mapping from SANA-WM action DSL letters to canonical controls.""" + +ALLOWED_ACTION_KEYS = frozenset(DSL_KEY_TO_CONTROL) +"""Action DSL keys accepted in ``-`` segments.""" + + +@dataclass +class VelocityState: + """Per-frame camera velocity in OpenCV camera coordinates.""" + + tx: float = 0.0 + """Forward translation velocity.""" + + sx: float = 0.0 + """Rightward strafe velocity.""" + + yaw: float = 0.0 + """Positive yaw-right angular velocity in radians.""" + + pitch: float = 0.0 + """Positive pitch-up angular velocity in radians.""" + + def snap_to(self, target: "VelocityState") -> None: + """Copy all velocity components from ``target``.""" + self.tx, self.sx = target.tx, target.sx + self.yaw, self.pitch = target.yaw, target.pitch + + def step_toward(self, target: "VelocityState", dt: float) -> None: + """Ease velocity components toward ``target`` over one timestep.""" + for attr in ("tx", "sx", "yaw", "pitch"): + cur = getattr(self, attr) + tgt = getattr(target, attr) + tau = TAU_PRESS if abs(tgt) > 1e-12 else TAU_COAST + alpha = 1.0 - math.exp(-dt / tau) + setattr(self, attr, cur + alpha * (tgt - cur)) + + +class CameraPoseIntegrator: + """Integrate camera velocity into camera-to-world poses.""" + + def __init__( + self, pitch_limit_rad: float = math.radians(DEFAULT_PITCH_LIMIT_DEG) + ) -> None: + self.pose = np.eye(4, dtype=np.float64) + self.pitch = 0.0 + self.pitch_limit = float(pitch_limit_rad) + + def step(self, velocity: VelocityState) -> np.ndarray: + """Advance one frame and return the current camera-to-world pose.""" + new_pitch = max( + -self.pitch_limit, min(self.pitch_limit, self.pitch + velocity.pitch) + ) + pitch_step = new_pitch - self.pitch + self.pitch = new_pitch + + rotation = self.pose[:3, :3] + rotation_new = _rot_y(velocity.yaw) @ rotation @ _rot_x(pitch_step) + + forward = rotation_new[:, 2].copy() + forward[1] = 0.0 + right = rotation_new[:, 0].copy() + right[1] = 0.0 + forward_norm = float(np.linalg.norm(forward)) + right_norm = float(np.linalg.norm(right)) + if forward_norm > 0: + forward /= forward_norm + 1e-6 + if right_norm > 0: + right /= right_norm + 1e-6 + + translation = self.pose[:3, 3] + forward * velocity.tx + right * velocity.sx + self.pose = np.eye(4, dtype=np.float64) + self.pose[:3, :3] = rotation_new + self.pose[:3, 3] = translation + return self.pose.copy() + + +def _rot_x(angle_rad: float) -> np.ndarray: + c, s = math.cos(angle_rad), math.sin(angle_rad) + return np.array([[1.0, 0.0, 0.0], [0.0, c, -s], [0.0, s, c]], dtype=np.float64) + + +def _rot_y(angle_rad: float) -> np.ndarray: + c, s = math.cos(angle_rad), math.sin(angle_rad) + return np.array([[c, 0.0, s], [0.0, 1.0, 0.0], [-s, 0.0, c]], dtype=np.float64) + + +def controls_to_target_velocity( + controls: set[str], + *, + translation_speed: float = DEFAULT_TRANSLATION_SPEED, + rotation_speed_rad: float | None = None, +) -> VelocityState: + """Map canonical control tokens to a target velocity.""" + if rotation_speed_rad is None: + rotation_speed_rad = math.radians(DEFAULT_ROTATION_SPEED_DEG) + forward = (1.0 if "forward" in controls else 0.0) - ( + 1.0 if "back" in controls else 0.0 + ) + strafe = (1.0 if "strafe_right" in controls else 0.0) - ( + 1.0 if "strafe_left" in controls else 0.0 + ) + yaw = (1.0 if "yaw_right" in controls else 0.0) - ( + 1.0 if "yaw_left" in controls else 0.0 + ) + pitch = (1.0 if "pitch_up" in controls else 0.0) - ( + 1.0 if "pitch_down" in controls else 0.0 + ) + return VelocityState( + tx=forward * translation_speed, + sx=strafe * translation_speed, + yaw=yaw * rotation_speed_rad, + pitch=pitch * rotation_speed_rad, + ) + + +def parse_action_string(action: str) -> list[list[str]]: + """Expand a SANA-WM action string into per-frame held keys. + + Args: + action: Comma-separated ``-`` segments. ``none-N`` + holds the pose for ``N`` frames. + + Returns: + Per-frame lists of DSL keys. + + Raises: + ValueError: The action string is empty or contains an invalid segment. + """ + cleaned = "".join(action.replace("，", ",").split()) + if not cleaned: + raise ValueError("action string is empty") + per_frame: list[list[str]] = [] + for segment in cleaned.split(","): + if not segment or "-" not in segment: + raise ValueError( + f"Invalid action segment {segment!r}: expected '-'." + ) + keys_part, duration_part = segment.rsplit("-", 1) + if not duration_part.isdigit() or int(duration_part) <= 0: + raise ValueError(f"Action segment {segment!r} has a non-positive duration.") + keys_lower = keys_part.lower() + if keys_lower == "none": + keys: list[str] = [] + else: + bad = sorted({key for key in keys_lower if key not in ALLOWED_ACTION_KEYS}) + if bad: + raise ValueError( + f"Action segment {segment!r} contains unknown keys {bad}; " + f"allowed: {''.join(sorted(ALLOWED_ACTION_KEYS))}." + ) + keys = sorted(set(keys_lower)) + per_frame.extend([keys] * int(duration_part)) + return per_frame + + +def action_string_to_c2w( + action: str, + *, + translation_speed: float = DEFAULT_TRANSLATION_SPEED, + rotation_speed_deg: float = DEFAULT_ROTATION_SPEED_DEG, + pitch_limit_deg: float = DEFAULT_PITCH_LIMIT_DEG, + smooth: bool = True, + num_frames: int | None = None, +) -> np.ndarray: + """Roll out a camera-to-world trajectory from an action string. + + Args: + action: Comma-separated ``-`` segments. + translation_speed: Per-frame translation magnitude. + rotation_speed_deg: Per-frame angular magnitude in degrees. + pitch_limit_deg: Maximum absolute pitch in degrees. + smooth: Whether to use the SANA-WM velocity smoothing model. + num_frames: Optional output frame count. When provided, the action + prompt is repeated or truncated to produce exactly this many + camera poses. + + Returns: + ``[F + 1, 4, 4]`` camera-to-world matrices in OpenCV convention. + """ + per_frame = parse_action_string(action) + if num_frames is not None: + if num_frames < 1: + raise ValueError(f"num_frames must be positive, got {num_frames}.") + target_steps = num_frames - 1 + if target_steps == 0: + per_frame = [] + elif len(per_frame) < target_steps: + repeats = math.ceil(target_steps / len(per_frame)) + per_frame = (per_frame * repeats)[:target_steps] + else: + per_frame = per_frame[:target_steps] + integrator = CameraPoseIntegrator(math.radians(pitch_limit_deg)) + velocity = VelocityState() + poses = [integrator.pose.copy()] + last_controls: set[str] = set() + dt = 1.0 / FPS + rotation_speed_rad = math.radians(rotation_speed_deg) + + for keys in per_frame: + controls = {DSL_KEY_TO_CONTROL[key] for key in keys} + target = controls_to_target_velocity( + controls, + translation_speed=translation_speed, + rotation_speed_rad=rotation_speed_rad, + ) + if smooth: + if controls - last_controls: + velocity.snap_to(target) + else: + velocity.step_toward(target, dt) + last_controls = controls + else: + velocity = target + poses.append(integrator.step(velocity)) + + return np.stack(poses, axis=0).astype(np.float32) + + +def fit_camera_trajectory(c2w: np.ndarray, num_frames: int) -> np.ndarray: + """Return a camera-to-world trajectory fitted to ``num_frames`` poses. + + Translations are interpolated linearly. Rotations are interpolated in + matrix space and projected back to the nearest proper rotation, avoiding a + new SciPy dependency for runner-time camera fitting. + """ + c2w = np.asarray(c2w, dtype=np.float32) + if c2w.ndim != 3 or c2w.shape[1:] != (4, 4): + raise ValueError(f"camera trajectory must be [F, 4, 4], got {c2w.shape}.") + if num_frames < 1: + raise ValueError(f"num_frames must be positive, got {num_frames}.") + if c2w.shape[0] == num_frames: + return c2w.copy() + if c2w.shape[0] == 1: + return np.broadcast_to(c2w[:1], (num_frames, 4, 4)).copy() + + old_t = np.linspace(0.0, 1.0, c2w.shape[0], dtype=np.float32) + new_t = np.linspace(0.0, 1.0, num_frames, dtype=np.float32) + fitted = np.broadcast_to(np.eye(4, dtype=np.float32), (num_frames, 4, 4)).copy() + + for axis in range(3): + fitted[:, axis, 3] = np.interp( + new_t, + old_t, + c2w[:, axis, 3], + ).astype(np.float32) + + rotations = c2w[:, :3, :3].reshape(c2w.shape[0], 9) + interp_rotations = np.empty((num_frames, 9), dtype=np.float32) + for idx in range(9): + interp_rotations[:, idx] = np.interp( + new_t, + old_t, + rotations[:, idx], + ).astype(np.float32) + for frame_idx, rotation in enumerate(interp_rotations.reshape(num_frames, 3, 3)): + u, _, vh = np.linalg.svd(rotation.astype(np.float64), full_matrices=False) + projected = u @ vh + if np.linalg.det(projected) < 0.0: + u[:, -1] *= -1.0 + projected = u @ vh + fitted[frame_idx, :3, :3] = projected.astype(np.float32) + return fitted + + +def fit_intrinsics_sequence(arr: np.ndarray, num_frames: int) -> np.ndarray: + """Return ``arr`` fitted to ``num_frames`` along axis 0.""" + arr = np.asarray(arr, dtype=np.float32) + if arr.shape[0] == num_frames: + return arr.copy() + if arr.shape[0] > num_frames: + return arr[:num_frames].copy() + if arr.shape[0] == 1: + return np.broadcast_to(arr[:1], (num_frames, *arr.shape[1:])).copy() + + old_t = np.linspace(0.0, 1.0, arr.shape[0], dtype=np.float32) + new_t = np.linspace(0.0, 1.0, num_frames, dtype=np.float32) + flat = arr.reshape(arr.shape[0], -1) + fitted = np.empty((num_frames, flat.shape[1]), dtype=np.float32) + for idx in range(flat.shape[1]): + fitted[:, idx] = np.interp(new_t, old_t, flat[:, idx]).astype(np.float32) + return fitted.reshape((num_frames, *arr.shape[1:])) + + +def load_intrinsics(path: Path, num_frames: int) -> np.ndarray: + """Return ``[num_frames, 4]`` intrinsics as ``[fx, fy, cx, cy]``. + + Args: + path: ``.npy`` file shaped ``[3, 3]``, ``[F, 3, 3]``, ``[4]``, or + ``[F, 4]``. + num_frames: Number of frames required by the rollout. + + Returns: + Per-frame vector intrinsics. + + Raises: + ValueError: The file shape is unsupported. + """ + arr = np.load(path).astype(np.float32) + if arr.shape == (4,): + return np.broadcast_to(arr, (num_frames, 4)).copy() + if arr.shape == (3, 3): + vector = np.array( + [arr[0, 0], arr[1, 1], arr[0, 2], arr[1, 2]], dtype=np.float32 + ) + return np.broadcast_to(vector, (num_frames, 4)).copy() + if arr.ndim == 2 and arr.shape[1] == 4: + return fit_intrinsics_sequence(arr, num_frames) + if arr.ndim == 3 and arr.shape[1:] == (3, 3): + matrix = fit_intrinsics_sequence(arr, num_frames) + return np.stack( + [matrix[:, 0, 0], matrix[:, 1, 1], matrix[:, 0, 2], matrix[:, 1, 2]], + axis=1, + ) + raise ValueError( + f"Unsupported intrinsics shape {arr.shape} for num_frames={num_frames}; " + "expected (3,3), (F,3,3), (4,), or (F,4)." + ) + + +def default_intrinsics_vec4( + src_size: tuple[int, int], + num_frames: int, + *, + hfov_deg: float = 90.0, +) -> np.ndarray: + """Derive ``[fx, fy, cx, cy]`` intrinsics from an image size. + + Used when the caller does not supply an intrinsics file. Assumes square + pixels and a principal point at the image center, with focal length set by + a horizontal field of view. The public SANA-WM demo intrinsics correspond + to roughly ``90`` degrees; lower values are narrower/more zoomed-in. The + returned intrinsics are in the source image's pixel coordinates, matching + :func:`load_intrinsics`, so they flow through + :func:`transform_intrinsics_for_crop` identically. + + Args: + src_size: Source image size as ``(width, height)``. + num_frames: Number of frames required by the rollout. + hfov_deg: Horizontal field of view in degrees. + + Returns: + ``[num_frames, 4]`` intrinsics as ``[fx, fy, cx, cy]``. + """ + src_w, src_h = src_size + if src_w <= 0 or src_h <= 0: + raise ValueError(f"src_size must be positive, got {src_size}.") + if not 0.0 < hfov_deg < 180.0: + raise ValueError(f"hfov_deg must be in (0, 180), got {hfov_deg}.") + focal = 0.5 * src_w / math.tan(math.radians(hfov_deg) / 2.0) + vector = np.array( + [focal, focal, src_w / 2.0, src_h / 2.0], dtype=np.float32 + ) + return np.broadcast_to(vector, (num_frames, 4)).copy() + + +def snap_num_frames( + num_frames: int, + *, + stride: int = SANA_WM_VAE_TEMPORAL_COMPRESSION, + upper_bound: int | None = None, +) -> int: + """Snap frame count to the nearest ``stride * k + 1`` value.""" + if num_frames < 1: + return 1 + if (num_frames - 1) % stride == 0: + return num_frames + floor_cand = num_frames - ((num_frames - 1) % stride) + ceil_cand = floor_cand + stride + snapped = ( + floor_cand + if (num_frames - floor_cand) < (ceil_cand - num_frames) + else ceil_cand + ) + if upper_bound is not None and snapped > upper_bound: + snapped = floor_cand + return max(snapped, 1) + + +def transform_intrinsics_for_crop( + intrinsics_vec4: np.ndarray, + src_size: tuple[int, int], + resized_size: tuple[int, int], + crop_offset: tuple[int, int], +) -> np.ndarray: + """Adjust ``[fx, fy, cx, cy]`` to match resize-then-center-crop pixels.""" + src_w, src_h = src_size + resized_w, resized_h = resized_size + crop_left, crop_top = crop_offset + sx, sy = resized_w / src_w, resized_h / src_h + out = intrinsics_vec4.copy() + out[..., 0] *= sx + out[..., 2] = out[..., 2] * sx - crop_left + out[..., 1] *= sy + out[..., 3] = out[..., 3] * sy - crop_top + return out + + +def resize_center_crop_geometry( + src_size: tuple[int, int], + *, + target_h: int = DEFAULT_VIDEO_HEIGHT, + target_w: int = DEFAULT_VIDEO_WIDTH, +) -> tuple[tuple[int, int], tuple[int, int]]: + """Return resize dimensions and center-crop offset for an input image size. + + Args: + src_size: Source image size as ``(width, height)``. + target_h: Target crop height. + target_w: Target crop width. + + Returns: + ``(resized_size, crop_offset)`` where sizes are ``(width, height)`` + and offsets are ``(left, top)``. + """ + src_w, src_h = src_size + scale = max(target_h / src_h, target_w / src_w) + resized_w = max(target_w, int(round(src_w * scale))) + resized_h = max(target_h, int(round(src_h * scale))) + left = (resized_w - target_w) // 2 + top = (resized_h - target_h) // 2 + return (resized_w, resized_h), (left, top) + + +def get_pose_inverse(transform: torch.Tensor) -> torch.Tensor: + """Return inverse homogeneous transforms using rotation transpose.""" + rotation = transform[..., :3, :3] + translation = transform[..., :3, 3] + rotation_inv = rotation.transpose(-1, -2) + translation_inv = -torch.matmul(rotation_inv, translation.unsqueeze(-1)).squeeze(-1) + out = torch.eye(4, dtype=transform.dtype, device=transform.device).repeat( + transform.shape[:-2] + (1, 1) + ) + out[..., :3, :3] = rotation_inv + out[..., :3, 3] = translation_inv + return out + + +def compute_raymap( + intrinsics: torch.Tensor, + poses: torch.Tensor, + height: int, + width: int, + *, + use_plucker: bool = True, +) -> torch.Tensor: + """Compute world-space ray directions and moments for camera poses. + + Args: + intrinsics: ``[T, 4]`` tensor of ``[fx, fy, cx, cy]``. + poses: ``[T, 4, 4]`` camera-to-world matrices. + height: Raymap height. + width: Raymap width. + use_plucker: ``True`` returns ``[direction, moment]``; otherwise + returns ``[origin, direction]``. + + Returns: + ``[T, height, width, 6]`` raymap tensor. + """ + num_frames = intrinsics.shape[0] + device, dtype = intrinsics.device, intrinsics.dtype + y_grid, x_grid = torch.meshgrid( + torch.arange(height, device=device, dtype=dtype), + torch.arange(width, device=device, dtype=dtype), + indexing="ij", + ) + x_grid = x_grid[None].expand(num_frames, -1, -1) + y_grid = y_grid[None].expand(num_frames, -1, -1) + + fx = intrinsics[:, 0].view(num_frames, 1, 1) + fy = intrinsics[:, 1].view(num_frames, 1, 1) + cx = intrinsics[:, 2].view(num_frames, 1, 1) + cy = intrinsics[:, 3].view(num_frames, 1, 1) + + dirs_cam = torch.stack( + [(x_grid - cx) / fx, (y_grid - cy) / fy, torch.ones_like(x_grid)], + dim=-1, + ) + rotation = poses[:, :3, :3] + translation = poses[:, :3, 3] + dirs_world = torch.einsum("tij,thwj->thwi", rotation, dirs_cam) + dirs_world = dirs_world / torch.norm(dirs_world, dim=-1, keepdim=True) + origins = translation.view(num_frames, 1, 1, 3).expand_as(dirs_world) + if use_plucker: + return torch.cat([dirs_world, torch.cross(origins, dirs_world, dim=-1)], dim=-1) + return torch.cat([origins, dirs_world], dim=-1) + + +def prepare_camera( + poses_c2w: np.ndarray, + intrinsics_vec4: np.ndarray, + *, + target_size: tuple[int, int] = (DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH), + vae_stride: tuple[int, int, int] = ( + SANA_WM_VAE_TEMPORAL_COMPRESSION, + SANA_WM_VAE_SPATIAL_COMPRESSION, + SANA_WM_VAE_SPATIAL_COMPRESSION, + ), +) -> dict[str, torch.Tensor]: + """Build SANA-WM raymap and chunk-Plucker conditioning tensors.""" + num_frames = poses_c2w.shape[0] + vae_time_stride, vae_spatial_stride = vae_stride[0], vae_stride[-1] + pixel_h, pixel_w = target_size + latent_h = pixel_h // vae_spatial_stride + latent_w = pixel_w // vae_spatial_stride + latent_frames = (num_frames - 1) // vae_time_stride + 1 + + poses = torch.from_numpy(poses_c2w).float() + first_inv = get_pose_inverse(poses[0:1]).squeeze(0) + poses_rel = torch.matmul(first_inv, poses[1:]) + poses = torch.cat([torch.eye(4).unsqueeze(0), poses_rel], dim=0) + + intrinsics = torch.from_numpy(intrinsics_vec4).float() + intrinsics_latent = intrinsics.clone() + intrinsics_latent[:, [0, 2]] *= latent_w / float(pixel_w) + intrinsics_latent[:, [1, 3]] *= latent_h / float(pixel_h) + + time_indices = torch.arange(0, num_frames, vae_time_stride) + if len(time_indices) > latent_frames: + time_indices = time_indices[:latent_frames] + raymap = torch.cat( + [ + poses[time_indices].reshape(len(time_indices), -1), + intrinsics_latent[time_indices], + ], + dim=-1, + ) + + chunks = [] + for start in time_indices - (vae_time_stride - 1): + chunk_start = max(0, int(start)) + chunk_end = chunk_start + vae_time_stride + chunk_poses = poses[chunk_start:chunk_end] + chunk_intrinsics = intrinsics_latent[chunk_start:chunk_end] + if chunk_poses.shape[0] < vae_time_stride: + pad = vae_time_stride - chunk_poses.shape[0] + chunk_poses = torch.cat([chunk_poses, chunk_poses[-1:].repeat(pad, 1, 1)]) + chunk_intrinsics = torch.cat( + [chunk_intrinsics, chunk_intrinsics[-1:].repeat(pad, 1)] + ) + plucker = compute_raymap( + chunk_intrinsics, + chunk_poses, + latent_h, + latent_w, + use_plucker=True, + ) + chunks.append(plucker.permute(0, 3, 1, 2).reshape(-1, latent_h, latent_w)) + chunk_plucker = torch.stack(chunks).permute(1, 0, 2, 3) + return {"raymap": raymap, "chunk_plucker": chunk_plucker} diff --git a/integrations/sana/sana_wm/conditioning.py b/integrations/sana/sana_wm/conditioning.py new file mode 100644 index 000000000..6be23039b --- /dev/null +++ b/integrations/sana/sana_wm/conditioning.py @@ -0,0 +1,552 @@ +# 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. + +"""SANA-WM prompt, first-frame, and camera conditioning components.""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +from torch import Tensor + +from flashdreams.infra.encoder import ( + Encoder, + EncoderConfig, + StreamingEncoder, + StreamingEncoderCache, +) +from sana_wm.camera import prepare_camera +from sana_wm.constants import ( + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, + SANA_WM_CONFIG_PATH, +) +from sana_wm.transformer import ( + QuantBackend, + SanaWMStage1Conditioning, + _chunk_index_from_config, + _get_tokenizer_and_text_encoder, + _get_vae, + _get_weight_dtype, + _load_inference_config, + _vae_encode_ltx2, +) + + +@dataclass(kw_only=True) +class SanaWMTextPromptRequest: + """Raw text prompt inputs for SANA-WM Stage 1.""" + + prompt: str + negative_prompt: str = "" + + +@dataclass(kw_only=True) +class SanaWMTextConditioning: + """Encoded positive and negative prompt tensors.""" + + condition: Tensor + condition_mask: Tensor + negative: Tensor + negative_mask: Tensor + + +@dataclass(kw_only=True) +class SanaWMCameraRequest: + """Raw camera trajectory inputs for SANA-WM conditioning.""" + + poses_c2w: np.ndarray + intrinsics_vec4: np.ndarray + + +@dataclass(kw_only=True) +class SanaWMI2VConditioning: + """Encoded first-frame and camera tensors for Stage-1 diffusion.""" + + first_latent: Tensor + camera: dict[str, Tensor] + num_frames: int + + +@dataclass(kw_only=True) +class SanaWMI2VConditioningRequest: + """Raw one-shot SANA-WM I2V rollout inputs.""" + + image: Any + prompt: str + poses_c2w: np.ndarray + intrinsics_vec4: np.ndarray + num_frames: int + fps: int + steps: int + cfg_scale: float + flow_shift: float | None + seed: int + negative_prompt: str = "" + + +@dataclass(kw_only=True) +class SanaWMTextPromptEncoderConfig(EncoderConfig): + """Config for the Stage-1 prompt encoder component.""" + + _target: type["SanaWMTextPromptEncoder"] = field( + default_factory=lambda: SanaWMTextPromptEncoder + ) + + config_path: str = SANA_WM_CONFIG_PATH + """SANA-WM inference YAML path or ``hf://`` URI.""" + + stage1_precision: str = "bf16" + """Stage-1 precision; quantized paths pad text tokens for scaled-MM.""" + + quant_backend: QuantBackend = "auto" + """Low-precision backend selector, included in the prompt cache key.""" + + offload_text_encoder: bool = False + """Move the text encoder back to CPU after prompt encoding.""" + + +class SanaWMTextPromptEncoder(Encoder): + """Encode SANA-WM Stage-1 positive and negative prompts.""" + + config: SanaWMTextPromptEncoderConfig + + def __init__(self, config: SanaWMTextPromptEncoderConfig) -> None: + super().__init__(config) + self.config = config + self._dummy = nn.Parameter(torch.empty(0)) + self._runtime_config: Any | None = None + self._text_encoder_built = False + self._prompt_cache: dict[ + tuple[object, ...], tuple[Tensor, Tensor, Tensor, Tensor] + ] = {} + + @property + def device(self) -> torch.device: + return self._dummy.device + + def forward(self, input: SanaWMTextPromptRequest) -> SanaWMTextConditioning: + """Encode prompt strings into Stage-1 text embeddings and masks.""" + cond, cond_mask, neg, neg_mask = self._encode_prompts( + input.prompt, + input.negative_prompt, + ) + cond, cond_mask, neg, neg_mask = self._pad_text_for_quant( + cond, + cond_mask, + neg, + neg_mask, + ) + return SanaWMTextConditioning( + condition=cond, + condition_mask=cond_mask, + negative=neg, + negative_mask=neg_mask, + ) + + def release_runtime(self) -> None: + """Release prompt encoder tensors.""" + self._prompt_cache.clear() + for attr in ("text_encoder", "tokenizer"): + if hasattr(self, attr): + setattr(self, attr, None) + self._text_encoder_built = False + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def _ensure_runtime_config(self) -> Any: + if self._runtime_config is None: + self._runtime_config = _load_inference_config(self.config.config_path) + return self._runtime_config + + def _ensure_weight_dtype(self) -> torch.dtype: + return _get_weight_dtype(self._ensure_runtime_config().model.mixed_precision) + + def _ensure_text_encoder(self) -> None: + if self._text_encoder_built: + return + cfg = self._ensure_runtime_config() + self.tokenizer, self.text_encoder = _get_tokenizer_and_text_encoder( + name=cfg.text_encoder.text_encoder_name, + device=self.device, + ) + if self.config.offload_text_encoder: + self.text_encoder.to("cpu") + self._text_encoder_built = True + + def _encode_prompts( + self, + prompt: str, + negative_prompt: str, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + cfg = self._ensure_runtime_config() + self._ensure_text_encoder() + max_length = cfg.text_encoder.model_max_length + chi_prompt = "\n".join(cfg.text_encoder.chi_prompt or []) + if chi_prompt: + prompt = chi_prompt + prompt + max_length_all = len(self.tokenizer.encode(chi_prompt)) + max_length - 2 + else: + max_length_all = max_length + + key = ( + prompt, + negative_prompt, + str(self.device), + str(self._ensure_weight_dtype()), + self.config.stage1_precision, + self.config.quant_backend, + ) + if key in self._prompt_cache: + return self._prompt_cache[key] + + move_text_encoder = self.config.offload_text_encoder or ( + _module_device(self.text_encoder) != self.device + ) + if move_text_encoder: + self.text_encoder.to(self.device) + + def encode(text: str, length: int) -> tuple[Tensor, Tensor]: + tokens = self.tokenizer( + [text], + max_length=length, + padding="max_length", + truncation=True, + return_tensors="pt", + ).to(self.device) + return self.text_encoder(tokens.input_ids, tokens.attention_mask)[0], ( + tokens.attention_mask + ) + + try: + cond, cond_mask = encode(prompt, max_length_all) + select = [0] + list(range(-max_length + 1, 0)) + cond = cond[:, None][:, :, select] + cond_mask = cond_mask[:, select] + neg, neg_mask = encode(negative_prompt, max_length) + result = (cond, cond_mask, neg[:, None], neg_mask) + self._prompt_cache.clear() + self._prompt_cache[key] = result + return result + finally: + if move_text_encoder: + self.text_encoder.to("cpu") + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def _pad_text_for_quant( + self, + cond: Tensor, + cond_mask: Tensor, + neg: Tensor, + neg_mask: Tensor, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + if self.config.stage1_precision == "bf16": + return cond, cond_mask, neg, neg_mask + multiple = int(os.environ.get("SANA_WM_STAGE1_NVFP4_TEXT_PAD_MULTIPLE", "8")) + if multiple <= 1: + return cond, cond_mask, neg, neg_mask + + def pad_pair(text: Tensor, mask: Tensor) -> tuple[Tensor, Tensor]: + pad = (-text.shape[-2]) % multiple + if pad == 0: + return text, mask + text_shape = list(text.shape) + text_shape[-2] = pad + mask_shape = list(mask.shape) + mask_shape[-1] = pad + return ( + torch.cat([text, text.new_zeros(text_shape)], dim=-2), + torch.cat([mask, mask.new_zeros(mask_shape)], dim=-1), + ) + + cond, cond_mask = pad_pair(cond, cond_mask) + neg, neg_mask = pad_pair(neg, neg_mask) + return cond, cond_mask, neg, neg_mask + + +@dataclass(kw_only=True) +class SanaWMFirstFrameEncoderConfig(EncoderConfig): + """Config for the first-frame VAE encoder component.""" + + _target: type["SanaWMFirstFrameEncoder"] = field( + default_factory=lambda: SanaWMFirstFrameEncoder + ) + + config_path: str = SANA_WM_CONFIG_PATH + """SANA-WM inference YAML path or ``hf://`` URI.""" + + offload_vae: bool = False + """Move the VAE to CPU after first-frame encoding.""" + + +class SanaWMFirstFrameEncoder(Encoder): + """Encode the input first frame with the LTX-2 VAE.""" + + config: SanaWMFirstFrameEncoderConfig + + def __init__(self, config: SanaWMFirstFrameEncoderConfig) -> None: + super().__init__(config) + self.config = config + self._dummy = nn.Parameter(torch.empty(0)) + self._runtime_config: Any | None = None + self._vae_built = False + + @property + def device(self) -> torch.device: + return self._dummy.device + + def forward(self, input: Any) -> Tensor: + """Encode a PIL-like RGB image into a single latent frame.""" + from torchvision import transforms as T + + cfg = self._ensure_runtime_config() + weight_dtype = _get_weight_dtype(cfg.model.mixed_precision) + self._ensure_vae() + if self.config.offload_vae: + self.vae.to(self.device) + + image = (T.ToTensor()(input) * 2.0 - 1.0).unsqueeze(0).unsqueeze(2) + latent = _vae_encode_ltx2( + cfg.vae.vae_type, + self.vae, + image.to(self.device, dtype=self.vae_dtype), + device=self.device, + ).to(weight_dtype) + if self.config.offload_vae: + self.vae.to("cpu") + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return latent + + def _ensure_runtime_config(self) -> Any: + if self._runtime_config is None: + self._runtime_config = _load_inference_config(self.config.config_path) + return self._runtime_config + + def _ensure_vae(self) -> None: + if self._vae_built: + return + cfg = self._ensure_runtime_config() + self.vae_dtype = _get_weight_dtype(cfg.vae.weight_dtype) + from sana_wm._tools import resolve_hf_path + + cfg.vae.vae_pretrained = resolve_hf_path(cfg.vae.vae_pretrained) + self.vae = _get_vae( + cfg.vae.vae_type, + cfg.vae.vae_pretrained, + device=self.device, + dtype=self.vae_dtype, + config=cfg.vae, + ) + self._vae_built = True + + +@dataclass(kw_only=True) +class SanaWMCameraConditioningEncoderConfig(EncoderConfig): + """Config for the camera/raymap conditioning component.""" + + _target: type["SanaWMCameraConditioningEncoder"] = field( + default_factory=lambda: SanaWMCameraConditioningEncoder + ) + + height: int = DEFAULT_VIDEO_HEIGHT + width: int = DEFAULT_VIDEO_WIDTH + + +class SanaWMCameraConditioningEncoder(Encoder): + """Build SANA-WM raymap and chunk-Plucker camera tensors.""" + + config: SanaWMCameraConditioningEncoderConfig + + def __init__(self, config: SanaWMCameraConditioningEncoderConfig) -> None: + super().__init__(config) + self.config = config + + def forward(self, input: SanaWMCameraRequest) -> dict[str, Tensor]: + """Encode raw poses and intrinsics into Stage-1 camera tensors.""" + return prepare_camera( + input.poses_c2w, + input.intrinsics_vec4, + target_size=(self.config.height, self.config.width), + ) + + +@dataclass(kw_only=True) +class SanaWMConditioningEncoderConfig(EncoderConfig): + """Config for the pipeline-level SANA-WM conditioning encoder.""" + + _target: type["SanaWMConditioningEncoder"] = field( + default_factory=lambda: SanaWMConditioningEncoder + ) + + config_path: str = SANA_WM_CONFIG_PATH + """SANA-WM inference YAML path or ``hf://`` URI.""" + + text_encoder: SanaWMTextPromptEncoderConfig = field( + default_factory=SanaWMTextPromptEncoderConfig + ) + first_frame_encoder: SanaWMFirstFrameEncoderConfig = field( + default_factory=SanaWMFirstFrameEncoderConfig + ) + camera_encoder: SanaWMCameraConditioningEncoderConfig = field( + default_factory=SanaWMCameraConditioningEncoderConfig + ) + height: int = DEFAULT_VIDEO_HEIGHT + width: int = DEFAULT_VIDEO_WIDTH + + +class SanaWMConditioningEncoder(StreamingEncoder[StreamingEncoderCache]): + """Prepare all per-rollout inputs for the Sana Stage-1 sampler.""" + + config: SanaWMConditioningEncoderConfig + + def __init__(self, config: SanaWMConditioningEncoderConfig) -> None: + super().__init__(config) + self.config = config + self.text_encoder = config.text_encoder.setup() + self.first_frame_encoder = config.first_frame_encoder.setup() + self.camera_encoder = config.camera_encoder.setup() + self._runtime_config: Any | None = None + + def initialize_autoregressive_cache(self, **_context: Any) -> StreamingEncoderCache: + """Return an empty cache for the one-shot conditioning encoder.""" + return StreamingEncoderCache() + + @torch.inference_mode() + def forward( + self, + input: SanaWMI2VConditioningRequest, + autoregressive_index: int = 0, + cache: StreamingEncoderCache | None = None, + ) -> SanaWMStage1Conditioning: + """Encode raw rollout inputs into Stage-1 conditioning.""" + del cache + if autoregressive_index != 0: + raise ValueError("SANA-WM bidirectional inference has one AR step.") + + cfg = self._ensure_runtime_config() + text = self.text_encoder( + SanaWMTextPromptRequest( + prompt=input.prompt, + negative_prompt=input.negative_prompt, + ) + ) + first_latent = self.first_frame_encoder(input.image) + weight_dtype = _get_weight_dtype(cfg.model.mixed_precision) + camera = self.camera_encoder( + SanaWMCameraRequest( + poses_c2w=input.poses_c2w, + intrinsics_vec4=input.intrinsics_vec4, + ) + ) + raymap = camera["raymap"].unsqueeze(0).to( + first_latent.device, + dtype=weight_dtype, + ) + chunk_plucker = camera["chunk_plucker"].unsqueeze(0).to( + first_latent.device, + dtype=weight_dtype, + ) + + model_kwargs_extra: dict[str, object] = {} + if input.cfg_scale > 1.0: + model_kwargs_extra["negative_mask"] = text.negative_mask + uncondition = text.negative + else: + uncondition = None + + vae_stride = cfg.vae.vae_stride + latent_t = (input.num_frames - 1) // int(vae_stride[0]) + 1 + latent_h = self.config.height // int(vae_stride[-1]) + latent_w = self.config.width // int(vae_stride[-1]) + chunk_index = _chunk_index_from_config(cfg, num_frames=latent_t) + model_kwargs: dict[str, object] = { + "data_info": { + "img_hw": torch.tensor( + [[self.config.height, self.config.width]], + dtype=torch.float, + device=first_latent.device, + ), + "condition_frame_info": {0: 0.0}, + }, + "mask": text.condition_mask, + "camera_conditions": raymap, + "chunk_plucker": chunk_plucker, + **model_kwargs_extra, + } + if chunk_index is not None: + model_kwargs["chunk_index"] = chunk_index + + return SanaWMStage1Conditioning( + condition=text.condition, + uncondition=uncondition, + model_kwargs=model_kwargs, + first_latent=first_latent, + latent_shape=( + 1, + int(first_latent.shape[1]), + latent_t, + latent_h, + latent_w, + ), + cfg_scale=float(input.cfg_scale), + flow_shift=self._resolve_flow_shift(input.flow_shift), + steps=int(input.steps), + seed=int(input.seed), + ) + + def _ensure_runtime_config(self) -> Any: + if self._runtime_config is None: + self._runtime_config = _load_inference_config(self.config.config_path) + return self._runtime_config + + def _resolve_flow_shift(self, override: float | None) -> float: + cfg = self._ensure_runtime_config() + if override is not None: + return float(override) + if cfg.scheduler.inference_flow_shift is not None: + return float(cfg.scheduler.inference_flow_shift) + return float(cfg.scheduler.flow_shift) + + +def _module_device(module: nn.Module) -> torch.device: + try: + return next(module.parameters()).device + except StopIteration: + return torch.device("cpu") + + +__all__ = [ + "SanaWMCameraConditioningEncoder", + "SanaWMCameraConditioningEncoderConfig", + "SanaWMCameraRequest", + "SanaWMConditioningEncoder", + "SanaWMConditioningEncoderConfig", + "SanaWMFirstFrameEncoder", + "SanaWMFirstFrameEncoderConfig", + "SanaWMI2VConditioning", + "SanaWMI2VConditioningRequest", + "SanaWMTextConditioning", + "SanaWMTextPromptEncoder", + "SanaWMTextPromptEncoderConfig", + "SanaWMTextPromptRequest", +] diff --git a/integrations/sana/sana_wm/config.py b/integrations/sana/sana_wm/config.py new file mode 100644 index 000000000..e08b2c679 --- /dev/null +++ b/integrations/sana/sana_wm/config.py @@ -0,0 +1,58 @@ +# 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. + +"""Static configs for the SANA-WM runner.""" + +from __future__ import annotations + +from flashdreams.infra.pipeline import StreamInferencePipelineConfig +from flashdreams.infra.runner import RunnerConfig +from sana_wm.conditioning import SanaWMConditioningEncoderConfig +from sana_wm.decoder import SanaWMVideoDecoderConfig +from sana_wm.diffusion import SanaWMDiffusionModelConfig +from sana_wm.runner import SanaWMRunnerConfig +from sana_wm.scheduler import SanaWMLTXEulerSchedulerConfig +from sana_wm.transformer import SanaWMTransformerConfig + + +PIPELINE_SANA_WM_BIDIRECTIONAL = StreamInferencePipelineConfig( + name="sana-wm-bidirectional", + encoder=SanaWMConditioningEncoderConfig(), + diffusion_model=SanaWMDiffusionModelConfig( + transformer=SanaWMTransformerConfig(), + scheduler=SanaWMLTXEulerSchedulerConfig(), + seed=42, + ), + decoder=SanaWMVideoDecoderConfig(), +) +"""FlashDreams SANA-WM pipeline.""" + +RUNNER_SANA_WM_BIDIRECTIONAL = SanaWMRunnerConfig( + runner_name=PIPELINE_SANA_WM_BIDIRECTIONAL.name, + description="SANA-WM bidirectional I2V runner (Stage-1 DiT + LTX-2 refiner).", + pipeline=PIPELINE_SANA_WM_BIDIRECTIONAL, +) +"""SANA-WM runner config.""" + +RUNNER_CONFIGS: dict[str, RunnerConfig] = { + cfg.runner_name: cfg for cfg in (RUNNER_SANA_WM_BIDIRECTIONAL,) +} +"""SANA-WM runner configs keyed by ``runner_name``.""" + +__all__ = [ + "PIPELINE_SANA_WM_BIDIRECTIONAL", + "RUNNER_CONFIGS", + "RUNNER_SANA_WM_BIDIRECTIONAL", +] diff --git a/integrations/sana/sana_wm/constants.py b/integrations/sana/sana_wm/constants.py new file mode 100644 index 000000000..d45f0c974 --- /dev/null +++ b/integrations/sana/sana_wm/constants.py @@ -0,0 +1,49 @@ +# 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. + +"""Constants for the public SANA-WM bidirectional release.""" + +SANA_WM_HF_REPO = "Efficient-Large-Model/SANA-WM_bidirectional" +"""Hugging Face repository containing SANA-WM bidirectional artefacts.""" + +SANA_WM_MODEL_PATH = f"hf://{SANA_WM_HF_REPO}/dit/sana_wm_1600m_720p.safetensors" +"""Default Stage-1 SANA-WM DiT checkpoint.""" + +SANA_WM_CONFIG_PATH = f"hf://{SANA_WM_HF_REPO}/config.yaml" +"""Default inference YAML.""" + +SANA_WM_REFINER_ROOT = f"hf://{SANA_WM_HF_REPO}/refiner" +"""Default LTX-2 refiner root.""" + +SANA_WM_REFINER_GEMMA_ROOT = f"hf://{SANA_WM_HF_REPO}/refiner/text_encoder" +"""Default Gemma text-encoder root used by the refiner.""" + +DEFAULT_VIDEO_HEIGHT = 704 +"""SANA-WM bidirectional output height in pixels.""" + +DEFAULT_VIDEO_WIDTH = 1280 +"""SANA-WM bidirectional output width in pixels.""" + +SANA_WM_VAE_TEMPORAL_COMPRESSION = 8 +"""Temporal compression ratio of the LTX2 VAE used by SANA-WM.""" + +SANA_WM_VAE_SPATIAL_COMPRESSION = 32 +"""Spatial compression ratio of the LTX2 VAE used by SANA-WM.""" + +DEFAULT_FPS = 16 +"""Frame rate used by the public SANA-WM examples.""" + +DEFAULT_ACTION = "w-100,dw-60,w-100,aw-60" +"""Default SANA-WM demo action string.""" diff --git a/integrations/sana/sana_wm/decoder.py b/integrations/sana/sana_wm/decoder.py new file mode 100644 index 000000000..38fcf5794 --- /dev/null +++ b/integrations/sana/sana_wm/decoder.py @@ -0,0 +1,526 @@ +# 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. + +"""SANA-WM latent refiner and VAE decoder components.""" + +from __future__ import annotations + +import gc +import time +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +from loguru import logger +from torch import Tensor + +from flashdreams.infra.config import InstantiateConfig +from flashdreams.infra.decoder import ( + DecoderConfig, + StreamingDecoderCache, + StreamingVideoDecoder, +) +from sana_wm._tools import resolve_hf_path +from sana_wm.constants import ( + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, + SANA_WM_CONFIG_PATH, + SANA_WM_REFINER_GEMMA_ROOT, + SANA_WM_REFINER_ROOT, + SANA_WM_VAE_SPATIAL_COMPRESSION, + SANA_WM_VAE_TEMPORAL_COMPRESSION, +) +from sana_wm.transformer import ( + Precision, + QuantBackend, + _avoid_degenerate_tile_tail, + _get_vae, + _get_weight_dtype, + _load_inference_config, + _vae_decode_ltx2, +) + + +@dataclass(kw_only=True) +class SanaWMDecodedVideo: + """Decoded SANA-WM video outputs.""" + + video_hwc: np.ndarray + stage1_video_hwc: np.ndarray | None = None + + +@dataclass(kw_only=True) +class SanaWMVideoDecoderCache(StreamingDecoderCache): + """Per-rollout settings for the SANA-WM decoder/refiner stage.""" + + prompt: str = "" + fps: int = 16 + save_stage1: bool = False + refiner_seed: int = 42 + sink_size: int = 1 + + +@dataclass(kw_only=True) +class SanaWMLTX2VAEDecoderConfig(InstantiateConfig): + """Config for the LTX-2 VAE latent-to-video decoder.""" + + _target: type["SanaWMLTX2VAEDecoder"] = field( + default_factory=lambda: SanaWMLTX2VAEDecoder + ) + + config_path: str = SANA_WM_CONFIG_PATH + """SANA-WM inference YAML path or ``hf://`` URI.""" + + height: int = DEFAULT_VIDEO_HEIGHT + width: int = DEFAULT_VIDEO_WIDTH + + offload_vae: bool = False + """Move the VAE to CPU between decode calls.""" + + vae_tile_sample_min_width: int = 512 + vae_tile_sample_stride_width: int = 448 + vae_tile_sample_min_height: int = 512 + vae_tile_sample_stride_height: int = 448 + vae_tile_sample_min_num_frames: int = 96 + vae_tile_sample_stride_num_frames: int = 64 + + vae_oom_retry_tile_sample_min_width: int = 128 + vae_oom_retry_tile_sample_stride_width: int = 64 + vae_oom_retry_tile_sample_min_height: int = 128 + vae_oom_retry_tile_sample_stride_height: int = 64 + vae_oom_retry_tile_sample_min_num_frames: int = 16 + vae_oom_retry_tile_sample_stride_num_frames: int = 8 + + +class SanaWMLTX2VAEDecoder(nn.Module): + """Decode SANA-WM LTX-2 VAE latents into HWC uint8 video.""" + + config: SanaWMLTX2VAEDecoderConfig + + def __init__(self, config: SanaWMLTX2VAEDecoderConfig) -> None: + super().__init__() + self.config = config + self._dummy = nn.Parameter(torch.empty(0)) + self._runtime_config: Any | None = None + self._vae_built = False + + @property + def device(self) -> torch.device: + return self._dummy.device + + @torch.inference_mode() + def decode_latents(self, latents: Tensor) -> np.ndarray: + """Decode VAE latents to ``uint8`` HWC video.""" + self._ensure_vae() + if self.config.offload_vae: + self.vae.to(self.device) + samples = latents.to(device=self.device, dtype=self.vae_dtype) + if torch.cuda.is_available(): + torch.cuda.synchronize() + t0 = time.perf_counter() + retry_decode = False + try: + decoded = self._vae_decode(samples) + except torch.OutOfMemoryError: + retry_decode = True + + if retry_decode: + if torch.cuda.is_available(): + torch.cuda.synchronize() + torch.cuda.empty_cache() + gc.collect() + logger.warning( + "[sana-vae] decode OOM; retrying with smaller tiles " + "width={} stride_width={} height={} stride_height={} " + "frames={} stride_frames={}", + self.config.vae_oom_retry_tile_sample_min_width, + self.config.vae_oom_retry_tile_sample_stride_width, + self.config.vae_oom_retry_tile_sample_min_height, + self.config.vae_oom_retry_tile_sample_stride_height, + self.config.vae_oom_retry_tile_sample_min_num_frames, + self.config.vae_oom_retry_tile_sample_stride_num_frames, + ) + self._configure_vae_tiling( + tile_sample_min_width=self.config.vae_oom_retry_tile_sample_min_width, + tile_sample_stride_width=( + self.config.vae_oom_retry_tile_sample_stride_width + ), + tile_sample_min_height=self.config.vae_oom_retry_tile_sample_min_height, + tile_sample_stride_height=( + self.config.vae_oom_retry_tile_sample_stride_height + ), + tile_sample_min_num_frames=( + self.config.vae_oom_retry_tile_sample_min_num_frames + ), + tile_sample_stride_num_frames=( + self.config.vae_oom_retry_tile_sample_stride_num_frames + ), + ) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + decoded = self._vae_decode(samples) + if torch.cuda.is_available(): + torch.cuda.synchronize() + logger.info( + "[timing] vae decode: {:.3f}s (latent T={} -> pixels {})", + time.perf_counter() - t0, + latents.shape[2], + tuple(decoded.shape) if isinstance(decoded, Tensor) else "list", + ) + if isinstance(decoded, list): + decoded = torch.stack(decoded, dim=0) + video = ( + torch.clamp(127.5 * decoded + 127.5, 0, 255) + .to(torch.uint8) + .permute(0, 2, 3, 4, 1) + .contiguous() + .cpu() + .numpy()[0] + ) + if self.config.offload_vae: + self.vae.to("cpu") + del samples, decoded + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return video + + def _ensure_runtime_config(self) -> Any: + if self._runtime_config is None: + self._runtime_config = _load_inference_config(self.config.config_path) + return self._runtime_config + + def _ensure_vae(self) -> None: + if self._vae_built: + return + cfg = self._ensure_runtime_config() + self.vae_dtype = _get_weight_dtype(cfg.vae.weight_dtype) + cfg.vae.vae_pretrained = resolve_hf_path(cfg.vae.vae_pretrained) + self.vae = _get_vae( + cfg.vae.vae_type, + cfg.vae.vae_pretrained, + device=self.device, + dtype=self.vae_dtype, + config=cfg.vae, + ) + if hasattr(self.vae, "enable_tiling"): + self._configure_vae_tiling() + self._vae_built = True + + def _configure_vae_tiling( + self, + *, + tile_sample_min_width: int | None = None, + tile_sample_stride_width: int | None = None, + tile_sample_min_height: int | None = None, + tile_sample_stride_height: int | None = None, + tile_sample_min_num_frames: int | None = None, + tile_sample_stride_num_frames: int | None = None, + ) -> None: + vae = self.vae + min_width = int(tile_sample_min_width or self.config.vae_tile_sample_min_width) + stride_width = int( + tile_sample_stride_width or self.config.vae_tile_sample_stride_width + ) + min_height = int(tile_sample_min_height or self.config.vae_tile_sample_min_height) + stride_height = int( + tile_sample_stride_height or self.config.vae_tile_sample_stride_height + ) + spatial_ratio = int(getattr(vae, "spatial_compression_ratio", 1)) + stride_width = _avoid_degenerate_tile_tail( + sample_extent=self.config.width, + sample_tile_min=min_width, + sample_stride=stride_width, + compression_ratio=spatial_ratio, + ) + stride_height = _avoid_degenerate_tile_tail( + sample_extent=self.config.height, + sample_tile_min=min_height, + sample_stride=stride_height, + compression_ratio=spatial_ratio, + ) + min_frames = int( + tile_sample_min_num_frames + or self.config.vae_tile_sample_min_num_frames + ) + stride_frames = int( + tile_sample_stride_num_frames + or self.config.vae_tile_sample_stride_num_frames + ) + temporal_ratio = int(getattr(vae, "temporal_compression_ratio", 1)) + if temporal_ratio > 1: + min_frames = max(min_frames, temporal_ratio) + stride_frames = max(stride_frames, temporal_ratio) + kwargs = { + "tile_sample_min_height": min_height, + "tile_sample_stride_height": stride_height, + "tile_sample_min_width": min_width, + "tile_sample_stride_width": stride_width, + "tile_sample_min_num_frames": min_frames, + "tile_sample_stride_num_frames": stride_frames, + } + if hasattr(vae, "enable_tiling"): + try: + vae.enable_tiling(**kwargs) + except TypeError: + vae.enable_tiling() + for name, value in kwargs.items(): + if hasattr(vae, name): + setattr(vae, name, value) + if hasattr(vae, "use_framewise_encoding"): + vae.use_framewise_encoding = True + if hasattr(vae, "use_framewise_decoding"): + vae.use_framewise_decoding = True + logger.info( + "[sana-vae] tiling width={} stride_width={} height={} " + "stride_height={} frames={} stride_frames={}", + min_width, + stride_width, + min_height, + stride_height, + min_frames, + stride_frames, + ) + + def _vae_decode(self, latents: Tensor) -> Tensor: + return _vae_decode_ltx2( + self._ensure_runtime_config().vae.vae_type, + self.vae, + latents, + ) + + +@dataclass(kw_only=True) +class SanaWMLTX2LatentRefinerConfig(InstantiateConfig): + """Config for the optional LTX-2 latent refiner component.""" + + _target: type["SanaWMLTX2LatentRefiner"] = field( + default_factory=lambda: SanaWMLTX2LatentRefiner + ) + + refiner_root: str = SANA_WM_REFINER_ROOT + refiner_gemma_root: str = SANA_WM_REFINER_GEMMA_ROOT + refiner_precision: Precision = "bf16" + quant_backend: QuantBackend = "torch" + offload_refiner: bool = False + cache_text_encoder: bool = False + """Keep Gemma cached on CPU after prompt encoding for repeated pipeline use.""" + + +class SanaWMLTX2LatentRefiner(nn.Module): + """Run the optional LTX-2 refinement stage over Stage-1 latents.""" + + config: SanaWMLTX2LatentRefinerConfig + + def __init__(self, config: SanaWMLTX2LatentRefinerConfig) -> None: + super().__init__() + self.config = config + self._dummy = nn.Parameter(torch.empty(0)) + self._refiner_built = False + + @property + def device(self) -> torch.device: + return self._dummy.device + + @torch.inference_mode() + def refine_latents( + self, + *, + latents: Tensor, + prompt: str, + fps: int, + sink_size: int, + seed: int, + ) -> Tensor: + """Run the LTX-2 refiner.""" + self._ensure_refiner() + refined = self.refiner.refine_latents( + latents, + prompt, + fps=float(fps), + sink_size=int(sink_size), + seed=int(seed), + progress=True, + ) + if self.config.offload_refiner: + self.release_runtime() + return refined + + def release_runtime(self) -> None: + """Release refiner tensors.""" + if not self._refiner_built: + return + del self.refiner + self._refiner_built = False + if torch.cuda.is_available(): + torch.cuda.empty_cache() + gc.collect() + + def _ensure_refiner(self) -> None: + if self._refiner_built: + return + from sana_wm.refiner import SanaWMLTX2Refiner + + t0 = time.perf_counter() + compute_dtype = ( + torch.bfloat16 + if self.config.refiner_precision in {"fp8", "fp4"} + else _get_weight_dtype(self.config.refiner_precision) + ) + self.refiner = SanaWMLTX2Refiner( + refiner_root=resolve_hf_path(self.config.refiner_root), + gemma_root=resolve_hf_path(self.config.refiner_gemma_root), + dtype=compute_dtype, + device=self.device, + precision=self.config.refiner_precision, + quant_backend=self.config.quant_backend, + cache_text_encoder=self.config.cache_text_encoder, + ) + self._refiner_built = True + logger.info( + "[timing] refiner build: {:.3f}s (precision={})", + time.perf_counter() - t0, + self.config.refiner_precision, + ) + + +@dataclass(kw_only=True) +class SanaWMVideoDecoderConfig(DecoderConfig): + """Config for the SANA-WM latent refiner plus VAE decode boundary.""" + + _target: type["SanaWMVideoDecoder"] = field( + default_factory=lambda: SanaWMVideoDecoder + ) + + vae_decoder: SanaWMLTX2VAEDecoderConfig = field( + default_factory=SanaWMLTX2VAEDecoderConfig + ) + refiner: SanaWMLTX2LatentRefinerConfig | None = field( + default_factory=SanaWMLTX2LatentRefinerConfig + ) + + +class SanaWMVideoDecoder(StreamingVideoDecoder[SanaWMVideoDecoderCache]): + """Decode Stage-1 latents, optionally through the LTX-2 refiner.""" + + config: SanaWMVideoDecoderConfig + + def __init__(self, config: SanaWMVideoDecoderConfig) -> None: + super().__init__(config) + self.config = config + self.vae_decoder = config.vae_decoder.setup() + self.refiner = config.refiner.setup() if config.refiner is not None else None + + def initialize_autoregressive_cache( + self, + **context: Any, + ) -> SanaWMVideoDecoderCache: + """Build per-rollout decode/refiner settings.""" + return SanaWMVideoDecoderCache(**context) + + @torch.inference_mode() + def forward( + self, + input: Tensor, + autoregressive_index: int = 0, + cache: SanaWMVideoDecoderCache | None = None, + ) -> SanaWMDecodedVideo: + """Refine/decode one SANA-WM latent rollout.""" + if autoregressive_index != 0: + raise ValueError("SANA-WM bidirectional inference has one AR step.") + cache = cache or SanaWMVideoDecoderCache() + stage1_latent = input + output_latent = input + if self.refiner is not None: + output_latent = self.refiner.refine_latents( + latents=stage1_latent, + prompt=cache.prompt, + fps=cache.fps, + sink_size=cache.sink_size, + seed=cache.refiner_seed, + ) + # refine_latents() already releases the refiner when offload_refiner + # is set; keep it resident otherwise so a reused pipeline avoids a + # full refiner rebuild + re-quantization on the next call. + elif cache.save_stage1: + logger.info( + "SANA-WM is already running without the refiner; " + "--save-stage1 does not create an extra output." + ) + + video_hwc = self.vae_decoder.decode_latents(output_latent) + if self.refiner is not None: + video_hwc = video_hwc[1:] + stage1_video_hwc = None + if cache.save_stage1 and self.refiner is not None: + stage1_video_hwc = self.vae_decoder.decode_latents(stage1_latent) + return SanaWMDecodedVideo( + video_hwc=video_hwc, + stage1_video_hwc=stage1_video_hwc, + ) + + @property + def spatial_compression_ratio(self) -> int: + """Pixel side divided by latent side.""" + return SANA_WM_VAE_SPATIAL_COMPRESSION + + @property + def temporal_compression_ratio(self) -> int: + """Pixel frame compression ratio after the first latent.""" + return SANA_WM_VAE_TEMPORAL_COMPRESSION + + def get_output_temporal_size( + self, + autoregressive_index: int, + input_temporal_size: int, + ) -> int: + """Return decoded pixel frames for a SANA-WM latent sequence.""" + if autoregressive_index != 0: + raise ValueError("SANA-WM bidirectional inference has one AR step.") + if input_temporal_size <= 0: + raise ValueError( + f"input_temporal_size must be positive, got {input_temporal_size}." + ) + return 1 + (input_temporal_size - 1) * self.temporal_compression_ratio + + def get_input_temporal_size( + self, + autoregressive_index: int, + output_temporal_size: int, + ) -> int: + """Return latent frames required to decode ``output_temporal_size`` pixels.""" + if autoregressive_index != 0: + raise ValueError("SANA-WM bidirectional inference has one AR step.") + ratio = self.temporal_compression_ratio + remainder = (output_temporal_size - 1) % ratio + if output_temporal_size <= 0 or remainder != 0: + raise ValueError( + "SANA-WM output frame count must be positive and equal to 8k+1; " + f"got {output_temporal_size}." + ) + return ((output_temporal_size - 1) // ratio) + 1 + + +__all__ = [ + "SanaWMDecodedVideo", + "SanaWMLTX2LatentRefiner", + "SanaWMLTX2LatentRefinerConfig", + "SanaWMLTX2VAEDecoder", + "SanaWMLTX2VAEDecoderConfig", + "SanaWMVideoDecoder", + "SanaWMVideoDecoderCache", + "SanaWMVideoDecoderConfig", +] diff --git a/integrations/sana/sana_wm/diffusion.py b/integrations/sana/sana_wm/diffusion.py new file mode 100644 index 000000000..6b67f01b5 --- /dev/null +++ b/integrations/sana/sana_wm/diffusion.py @@ -0,0 +1,37 @@ +# 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. + +"""SANA-WM diffusion config bound to the shared FlashDreams diffusion model.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from flashdreams.infra.diffusion.model import DiffusionModel, DiffusionModelConfig + + +@dataclass(kw_only=True) +class SanaWMDiffusionModelConfig(DiffusionModelConfig): + """Diffusion model config for SANA-WM's Stage-1 sampler. + + Sana-specific behavior lives in the configured transformer and scheduler; + this config intentionally instantiates the common FlashDreams + :class:`DiffusionModel` instead of a custom ``generate`` implementation. + """ + + _target: type[DiffusionModel] = field(default_factory=lambda: DiffusionModel) + + +__all__ = ["SanaWMDiffusionModelConfig"] diff --git a/integrations/sana/sana_wm/ops/__init__.py b/integrations/sana/sana_wm/ops/__init__.py new file mode 100644 index 000000000..07bb2847d --- /dev/null +++ b/integrations/sana/sana_wm/ops/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""SANA-WM integration-local Triton kernels.""" diff --git a/integrations/sana/sana_wm/quant.py b/integrations/sana/sana_wm/quant.py new file mode 100644 index 000000000..d285e73f1 --- /dev/null +++ b/integrations/sana/sana_wm/quant.py @@ -0,0 +1,689 @@ +# 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. + +"""Low-precision linear helpers for SANA-WM.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Any, Literal + +import torch +import triton +import triton.language as tl +from torch import nn + +FP8_MAX_E4M3 = 448.0 +FP8_SCALE_EPS = 1.0e-12 +FP4_MAX_E2M1 = 6.0 +NVFP4_BLOCK_SIZE = 16 +NVFP4_E4M3_SCALE_EPS = 0.015625 +NVFP4_GLOBAL_SCALE_EPS = 1.0e-12 +_RHT16_SCALE = 0.25 +_RHT16_SIGNS = ( + 1.0, + -1.0, + -1.0, + 1.0, + -1.0, + 1.0, + -1.0, + 1.0, + 1.0, + 1.0, + -1.0, + -1.0, + -1.0, + -1.0, + 1.0, + 1.0, +) + + +@dataclass(frozen=True) +class TorchScaledMMFP8Recipe: + """FP8 recipe marker for SANA-WM Linear replacement.""" + + name: str = "torch_scaled_mm_fp8" + precision: Literal["fp8"] = "fp8" + + +@dataclass(frozen=True) +class TorchScaledMMFP4Recipe: + """FP4 recipe marker for SANA-WM Linear replacement.""" + + name: str = "torch_scaled_mm_fp4" + precision: Literal["fp4"] = "fp4" + use_rht: bool = True + use_global_scale: bool = True + weight_scale_2d: bool = True + + +QuantRecipe = TorchScaledMMFP8Recipe | TorchScaledMMFP4Recipe + + +@triton.jit +def _pack_fp32_to_fp4_pairs(values): + packed = tl.inline_asm_elementwise( + asm=""" + { + .reg .b8 byte0, byte1, byte2, byte3; + cvt.rn.satfinite.e2m1x2.f32 byte0, $5, $1; + cvt.rn.satfinite.e2m1x2.f32 byte1, $6, $2; + cvt.rn.satfinite.e2m1x2.f32 byte2, $7, $3; + cvt.rn.satfinite.e2m1x2.f32 byte3, $8, $4; + mov.b32 $0, {byte0, byte1, byte2, byte3}; + } + """, + constraints=("=r,r,r,r,r,r,r,r,r"), + args=values, + dtype=tl.uint8, + is_pure=True, + pack=4, + ) + return packed + + +@triton.jit +def _quantize_nvfp4_kernel( + input_ptr, + qdata_ptr, + scale_ptr, + global_scale_ptr, + stride_m, + stride_k, + rows: tl.constexpr, + cols: tl.constexpr, + mask_scales: tl.constexpr, + scale_2d: tl.constexpr, + has_global_scale: tl.constexpr, +): + fp4_max = 6.0 + fp8_max = 448.0 + scale_eps = 0.015625 + global_scale = 1.0 + if has_global_scale: + global_scale = tl.load(global_scale_ptr).to(tl.float32) + + pid_k = tl.program_id(0) + pid_m = tl.program_id(1) + + offsets_m = pid_m * 128 + tl.arange(0, 128)[:, None] + offsets_k = pid_k * 64 + tl.arange(0, 64)[None, :] + mask = (offsets_m < rows) & (offsets_k < cols) + values = tl.load( + input_ptr + offsets_m * stride_m + offsets_k * stride_k, + mask=mask, + other=0.0, + ) + values = values.to(tl.float32) + if scale_2d: + values_4d = values.reshape(8, 16, 4, 16) + block_amax = tl.max(tl.max(tl.abs(values_4d), axis=3), axis=1) + scales_f32 = block_amax / fp4_max + if has_global_scale: + scales_f32 = scales_f32 / global_scale + scales_f32 = tl.clamp(scales_f32, scale_eps, fp8_max) + scales_tile = scales_f32.to(tl.float8e4nv) + scales = tl.broadcast_to(tl.expand_dims(scales_tile, 1), (8, 16, 4)) + scales = scales.reshape(128, 4) + else: + values_3d = values.reshape(128, 4, 16) + block_amax = tl.max(tl.abs(values_3d), axis=2) + scales_f32 = block_amax / fp4_max + if has_global_scale: + scales_f32 = scales_f32 / global_scale + scales_f32 = tl.clamp(scales_f32, scale_eps, fp8_max) + scales = scales_f32.to(tl.float8e4nv) + + values = values.reshape(128, 4, 16) + scale_product = scales.to(tl.float32) + if has_global_scale: + scale_product = scale_product * global_scale + quantized = tl.div_rn(values, scale_product[:, :, None]) + + if mask_scales: + scale_offsets_k = pid_k * 4 + tl.arange(0, 4)[None, :] + scale_mask = (offsets_m < rows) & (scale_offsets_k < tl.cdiv(cols, 16)) + scales = tl.where(scale_mask, scales, 0.0) + + packed_scales = scales.reshape(4, 32, 4).permute(1, 0, 2).reshape(32, 16) + scale_m = tl.arange(0, 32)[:, None] + scale_k = tl.arange(0, 16)[None, :] + tl.store( + scale_ptr + (pid_m * tl.num_programs(0) + pid_k) * (32 * 16) + scale_m * 16 + scale_k, + packed_scales, + ) + + packed = _pack_fp32_to_fp4_pairs(quantized.reshape(128, 32, 2).split()) + q_offsets_m = pid_m * 128 + tl.arange(0, 128)[:, None] + q_offsets_k = pid_k * 32 + tl.arange(0, 32)[None, :] + q_mask = (q_offsets_m < rows) & (q_offsets_k < cols // 2) + tl.store(qdata_ptr + q_offsets_m * (cols // 2) + q_offsets_k, packed, mask=q_mask) + + +def _ceil_div(value: int, divisor: int) -> int: + return (value + divisor - 1) // divisor + + +def _swizzled_nvfp4_scale_shape(rows: int, cols: int) -> tuple[int, int]: + return _ceil_div(rows, 128) * 32, _ceil_div(cols, 64) * 16 + + +def quantize_nvfp4_swizzled( + input: torch.Tensor, + *, + global_scale: torch.Tensor | None = None, + scale_2d: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a 2D CUDA tensor to NVFP4 qdata plus swizzled E4M3 scales.""" + _require_fp4_dtype() + if input.dim() != 2: + raise ValueError(f"NVFP4 quantization requires a 2D tensor, got {input.dim()}D.") + if not input.is_cuda: + raise RuntimeError("NVFP4 quantization requires CUDA input.") + rows, cols = input.shape + if cols % NVFP4_BLOCK_SIZE != 0: + raise ValueError( + f"NVFP4 quantization requires last dim divisible by {NVFP4_BLOCK_SIZE}, got {cols}." + ) + input_2d = input.contiguous() + scale_rows, scale_cols = _swizzled_nvfp4_scale_shape(rows, cols) + qdata = torch.empty((rows, cols // 2), device=input.device, dtype=torch.uint8) + scales = torch.empty( + (scale_rows, scale_cols), + device=input.device, + dtype=torch.float8_e4m3fn, + ) + grid = (_ceil_div(cols, 64), _ceil_div(rows, 128)) + _quantize_nvfp4_kernel[grid]( + input_2d, + qdata, + scales, + global_scale if global_scale is not None else input_2d, + input_2d.stride(0), + input_2d.stride(1), + rows, + cols, + mask_scales=(rows % 128 != 0 or cols % 64 != 0), + scale_2d=scale_2d, + has_global_scale=global_scale is not None, + ) + return qdata, scales + + +def nvfp4_global_scale(input: torch.Tensor) -> torch.Tensor: + """Return the per-tensor FP32 global scale for hierarchical NVFP4.""" + amax = input.detach().abs().amax().to(torch.float32) + scale = amax / (FP8_MAX_E4M3 * FP4_MAX_E2M1) + return scale.clamp_min(NVFP4_GLOBAL_SCALE_EPS).reshape(()) + + +def apply_rht16(input: torch.Tensor) -> torch.Tensor: + """Apply the tiled 16-wide random Hadamard transform used by NVFP4.""" + if input.shape[-1] % 16 != 0: + raise ValueError( + f"RHT16 requires the last dimension to be divisible by 16, got {input.shape[-1]}." + ) + original_shape = input.shape + work = input.to(torch.float32).reshape(-1, 16) + signs = torch.tensor(_RHT16_SIGNS, device=input.device, dtype=work.dtype) + work = work * signs + width = 1 + while width < 16: + work = work.reshape(-1, 16 // (2 * width), 2, width) + left = work[:, :, 0, :] + right = work[:, :, 1, :] + work = torch.stack((left + right, left - right), dim=2).reshape(-1, 16) + width *= 2 + return (work * _RHT16_SCALE).reshape(original_shape).contiguous() + + +class TorchScaledMMFP8Linear(nn.Module): + """Inference-only FP8 Linear using ``torch._scaled_mm``. + + The module stores the source weight as E4M3 FP8 with per-output-channel + scales. At runtime it quantizes the flattened activation rows to E4M3 with + per-row scales, then calls ``torch._scaled_mm`` and returns BF16 output. + This intentionally mirrors the shape contract of ``nn.Linear`` so it can + replace eligible SANA-WM linear layers without changing call sites. + """ + + in_features: int + out_features: int + + def __init__( + self, + *, + weight: torch.Tensor, + weight_fp8: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, + ) -> None: + super().__init__() + if weight.shape != weight_fp8.shape: + raise ValueError( + "weight and weight_fp8 must have matching shape, got " + f"{tuple(weight.shape)} and {tuple(weight_fp8.shape)}." + ) + if weight_fp8.dim() != 2: + raise ValueError(f"weight_fp8 must be 2D, got {tuple(weight_fp8.shape)}.") + if weight_scale.shape != (weight_fp8.shape[0], 1): + raise ValueError( + "weight_scale must have shape " + f"({weight_fp8.shape[0]}, 1), got {tuple(weight_scale.shape)}." + ) + if bias is not None and bias.shape != (weight_fp8.shape[0],): + raise ValueError( + f"bias must have shape ({weight_fp8.shape[0]},), got {tuple(bias.shape)}." + ) + self.out_features = int(weight_fp8.shape[0]) + self.in_features = int(weight_fp8.shape[1]) + self.out_dtype = out_dtype + self.register_buffer("weight", weight.detach().contiguous()) + self.register_buffer("weight_fp8", weight_fp8.contiguous()) + self.register_buffer("weight_scale", weight_scale.to(torch.float32).contiguous()) + if bias is None: + self.register_buffer("bias", None) + else: + self.register_buffer("bias", bias.detach().to(out_dtype).contiguous()) + + @classmethod + def from_linear( + cls, + source: nn.Linear, + *, + out_dtype: torch.dtype, + ) -> "TorchScaledMMFP8Linear": + """Create an FP8 replacement from a source ``nn.Linear``.""" + _require_fp8_dtype() + weight_f32 = source.weight.detach().to(torch.float32) + weight_scale = ( + weight_f32.abs().amax(dim=1, keepdim=True).clamp_min(FP8_SCALE_EPS) + / FP8_MAX_E4M3 + ) + weight_fp8 = torch.clamp( + weight_f32 / weight_scale, + -FP8_MAX_E4M3, + FP8_MAX_E4M3, + ).to(torch.float8_e4m3fn) + bias = source.bias.detach() if source.bias is not None else None + replacement = cls( + weight=source.weight.detach().to(device=source.weight.device), + weight_fp8=weight_fp8.to(device=source.weight.device), + weight_scale=weight_scale.to(device=source.weight.device), + bias=bias.to(device=source.weight.device) if bias is not None else None, + out_dtype=out_dtype, + ) + replacement.train(source.training) + return replacement + + def forward(self, input: torch.Tensor) -> torch.Tensor: + """Apply the quantized linear projection.""" + if input.shape[-1] != self.in_features: + raise ValueError( + f"expected input last dim {self.in_features}, got {input.shape[-1]}." + ) + if not input.is_cuda: + raise RuntimeError("TorchScaledMMFP8Linear requires CUDA input.") + _require_scaled_mm() + _require_fp8_dtype() + + leading_shape = input.shape[:-1] + input_2d = input.reshape(-1, self.in_features) + if input_2d.numel() == 0: + return input.new_empty( + (*leading_shape, self.out_features), + dtype=self.out_dtype, + ) + + input_f32 = input_2d.to(torch.float32) + input_scale = ( + input_f32.abs().amax(dim=1, keepdim=True).clamp_min(FP8_SCALE_EPS) + / FP8_MAX_E4M3 + ) + input_fp8 = torch.clamp( + input_f32 / input_scale, + -FP8_MAX_E4M3, + FP8_MAX_E4M3, + ).to(torch.float8_e4m3fn) + input_fp8 = input_fp8.contiguous() + + # ``_scaled_mm`` expects the B operand as a column-major transpose view, + # so keep the transpose stride instead of forcing it contiguous. + output = torch._scaled_mm( + input_fp8, + self.weight_fp8.t(), + scale_a=input_scale.contiguous(), + scale_b=self.weight_scale.t().contiguous(), + out_dtype=self.out_dtype, + ) + if self.bias is not None: + output = output + self.bias.to(device=output.device, dtype=output.dtype) + return output.reshape(*leading_shape, self.out_features) + + +class TorchScaledMMFP4Linear(nn.Module): + """Inference-only NVFP4 Linear using Triton quantization and ``_scaled_mm``.""" + + in_features: int + out_features: int + + def __init__( + self, + *, + weight: torch.Tensor, + weight_qdata: torch.Tensor, + weight_scale: torch.Tensor, + weight_global_scale: torch.Tensor | None, + bias: torch.Tensor | None, + out_dtype: torch.dtype, + use_rht: bool, + ) -> None: + super().__init__() + expected_weight_shape = (weight_qdata.shape[0], weight_qdata.shape[1] * 2) + if weight.shape != expected_weight_shape: + raise ValueError( + "weight must match unpacked weight_qdata shape " + f"{expected_weight_shape}, got {tuple(weight.shape)}." + ) + if weight_qdata.dim() != 2: + raise ValueError(f"weight_qdata must be 2D, got {tuple(weight_qdata.shape)}.") + if weight_qdata.shape[1] * 2 % NVFP4_BLOCK_SIZE != 0: + raise ValueError( + "weight_qdata must represent a K dimension divisible by " + f"{NVFP4_BLOCK_SIZE}, got packed shape {tuple(weight_qdata.shape)}." + ) + if bias is not None and bias.shape != (weight_qdata.shape[0],): + raise ValueError( + f"bias must have shape ({weight_qdata.shape[0]},), got {tuple(bias.shape)}." + ) + self.out_features = int(weight_qdata.shape[0]) + self.in_features = int(weight_qdata.shape[1] * 2) + self.out_dtype = out_dtype + self.use_rht = use_rht + self.register_buffer("weight", weight.detach().contiguous()) + self.register_buffer("weight_qdata", weight_qdata.contiguous()) + self.register_buffer("weight_scale", weight_scale.contiguous()) + if weight_global_scale is None: + self.register_buffer("weight_global_scale", None) + else: + self.register_buffer( + "weight_global_scale", + weight_global_scale.detach().to(torch.float32).reshape(()), + ) + if bias is None: + self.register_buffer("bias", None) + else: + self.register_buffer("bias", bias.detach().to(out_dtype).contiguous()) + + @classmethod + def from_linear( + cls, + source: nn.Linear, + *, + out_dtype: torch.dtype, + use_rht: bool = True, + use_global_scale: bool = True, + weight_scale_2d: bool = True, + ) -> "TorchScaledMMFP4Linear": + """Create an NVFP4 replacement from a source ``nn.Linear``.""" + _require_fp4_dtype() + weight_for_quant = source.weight.detach() + if use_rht: + weight_for_quant = apply_rht16(weight_for_quant) + weight_global_scale = ( + nvfp4_global_scale(weight_for_quant) if use_global_scale else None + ) + weight_qdata, weight_scale = quantize_nvfp4_swizzled( + weight_for_quant, + global_scale=weight_global_scale, + scale_2d=weight_scale_2d, + ) + bias = source.bias.detach() if source.bias is not None else None + replacement = cls( + weight=source.weight.detach().to(device=source.weight.device), + weight_qdata=weight_qdata, + weight_scale=weight_scale, + weight_global_scale=weight_global_scale, + bias=bias.to(device=source.weight.device) if bias is not None else None, + out_dtype=out_dtype, + use_rht=use_rht, + ) + replacement.train(source.training) + return replacement + + def forward(self, input: torch.Tensor) -> torch.Tensor: + """Apply the quantized linear projection.""" + if input.shape[-1] != self.in_features: + raise ValueError( + f"expected input last dim {self.in_features}, got {input.shape[-1]}." + ) + if not input.is_cuda: + raise RuntimeError("TorchScaledMMFP4Linear requires CUDA input.") + _require_scaled_mm() + _require_fp4_dtype() + + leading_shape = input.shape[:-1] + input_2d = input.reshape(-1, self.in_features) + if input_2d.numel() == 0: + return input.new_empty( + (*leading_shape, self.out_features), + dtype=self.out_dtype, + ) + + input_for_quant = apply_rht16(input_2d) if self.use_rht else input_2d + input_global_scale = ( + nvfp4_global_scale(input_for_quant) + if self.weight_global_scale is not None + else None + ) + input_qdata, input_scale = quantize_nvfp4_swizzled( + input_for_quant, + global_scale=input_global_scale, + ) + add_bias_after_scale = ( + self.bias is not None + and self.weight_global_scale is not None + and input_global_scale is not None + ) + output = torch._scaled_mm( + input_qdata.view(torch.float4_e2m1fn_x2), + self.weight_qdata.t().view(torch.float4_e2m1fn_x2), + input_scale.view(torch.float8_e4m3fn), + self.weight_scale.view(torch.float8_e4m3fn), + bias=None + if add_bias_after_scale + else ( + self.bias.to(device=input.device, dtype=self.out_dtype) + if self.bias is not None + else None + ), + out_dtype=self.out_dtype, + ) + if input_global_scale is not None and self.weight_global_scale is not None: + output = output * ( + input_global_scale.to(device=output.device, dtype=output.dtype) + * self.weight_global_scale.to(device=output.device, dtype=output.dtype) + ) + if add_bias_after_scale: + output = output + self.bias.to(device=output.device, dtype=output.dtype) + return output.reshape(*leading_shape, self.out_features) + + +def replace_linear_with_torch_fp8( + module: nn.Module, + *, + recipe: Any, + params_dtype: torch.dtype, + skip_patterns: tuple[str, ...], + include_patterns: tuple[str, ...] | None = None, + prefix: str = "", +) -> tuple[int, int]: + """Replace eligible ``nn.Linear`` modules with ``TorchScaledMMFP8Linear``.""" + return _replace_linear_with_quant( + module, + recipe=TorchScaledMMFP8Recipe(), + params_dtype=params_dtype, + skip_patterns=skip_patterns, + include_patterns=include_patterns, + prefix=prefix, + ) + + +def replace_linear_with_torch_fp4( + module: nn.Module, + *, + recipe: Any, + params_dtype: torch.dtype, + skip_patterns: tuple[str, ...], + include_patterns: tuple[str, ...] | None = None, + prefix: str = "", +) -> tuple[int, int]: + """Replace eligible ``nn.Linear`` modules with ``TorchScaledMMFP4Linear``.""" + return _replace_linear_with_quant( + module, + recipe=TorchScaledMMFP4Recipe(), + params_dtype=params_dtype, + skip_patterns=skip_patterns, + include_patterns=include_patterns, + prefix=prefix, + ) + + +def replace_linear_with_quant( + module: nn.Module, + *, + recipe: QuantRecipe, + params_dtype: torch.dtype, + skip_patterns: tuple[str, ...], + include_patterns: tuple[str, ...] | None = None, + prefix: str = "", +) -> tuple[int, int]: + """Replace eligible ``nn.Linear`` modules with the requested backend.""" + return _replace_linear_with_quant( + module, + recipe=recipe, + params_dtype=params_dtype, + skip_patterns=skip_patterns, + include_patterns=include_patterns, + prefix=prefix, + ) + + +def _replace_linear_with_quant( + module: nn.Module, + *, + recipe: QuantRecipe, + params_dtype: torch.dtype, + skip_patterns: tuple[str, ...], + include_patterns: tuple[str, ...] | None, + prefix: str, +) -> tuple[int, int]: + converted = 0 + skipped = 0 + for name, child in list(module.named_children()): + child_prefix = f"{prefix}.{name}" if prefix else name + if _name_matches(skip_patterns, child_prefix): + skipped += 1 + continue + if isinstance(child, nn.Linear): + if include_patterns is not None and not _name_matches( + include_patterns, + child_prefix, + ): + skipped += 1 + continue + if child.in_features % 16 != 0 or child.out_features % 16 != 0: + skipped += 1 + continue + if recipe.precision == "fp4" and child.in_features % 32 != 0: + skipped += 1 + continue + out_dtype = ( + params_dtype + if params_dtype in {torch.bfloat16, torch.float16} + else torch.bfloat16 + ) + if recipe.precision == "fp8": + replacement = TorchScaledMMFP8Linear.from_linear( + child, + out_dtype=out_dtype, + ) + elif recipe.precision == "fp4": + replacement = TorchScaledMMFP4Linear.from_linear( + child, + out_dtype=out_dtype, + use_rht=recipe.use_rht, + use_global_scale=recipe.use_global_scale, + weight_scale_2d=recipe.weight_scale_2d, + ) + else: + raise ValueError(f"Unsupported SANA-WM quant recipe: {recipe!r}.") + setattr(module, name, replacement) + converted += 1 + continue + child_converted, child_skipped = _replace_linear_with_quant( + child, + recipe=recipe, + params_dtype=params_dtype, + skip_patterns=skip_patterns, + include_patterns=include_patterns, + prefix=child_prefix, + ) + converted += child_converted + skipped += child_skipped + return converted, skipped + + +def _name_matches(patterns: tuple[str, ...], name: str) -> bool: + return any(re.search(pattern, name) for pattern in patterns) + + +def _require_scaled_mm() -> None: + if not hasattr(torch, "_scaled_mm"): + raise RuntimeError("torch._scaled_mm is required for SANA-WM quantization.") + + +def _require_fp8_dtype() -> None: + if not hasattr(torch, "float8_e4m3fn"): + raise RuntimeError("torch.float8_e4m3fn is required for the SANA-WM FP8 backend.") + + +def _require_fp4_dtype() -> None: + if not hasattr(torch, "float4_e2m1fn_x2"): + raise RuntimeError( + "torch.float4_e2m1fn_x2 is required for the SANA-WM FP4 backend." + ) + if not hasattr(torch, "float8_e4m3fn"): + raise RuntimeError( + "torch.float8_e4m3fn scales are required for the SANA-WM FP4 backend." + ) + + +__all__ = [ + "TorchScaledMMFP4Linear", + "TorchScaledMMFP4Recipe", + "TorchScaledMMFP8Linear", + "TorchScaledMMFP8Recipe", + "apply_rht16", + "nvfp4_global_scale", + "quantize_nvfp4_swizzled", + "replace_linear_with_quant", + "replace_linear_with_torch_fp4", + "replace_linear_with_torch_fp8", +] diff --git a/integrations/sana/sana_wm/refiner.py b/integrations/sana/sana_wm/refiner.py new file mode 100644 index 000000000..8a1246281 --- /dev/null +++ b/integrations/sana/sana_wm/refiner.py @@ -0,0 +1,700 @@ +# 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. + +"""LTX-2 refiner used by the SANA-WM integration.""" + +from __future__ import annotations + +import time +from pathlib import Path +from typing import Literal + +import torch +import torch.nn as nn +import torch.nn.functional as F +from loguru import logger +from torch import Tensor + +from sana_wm.quant import ( + TorchScaledMMFP4Recipe, + TorchScaledMMFP8Recipe, + replace_linear_with_quant, +) + +Precision = Literal["bf16", "fp8", "fp4"] +QuantBackend = Literal["auto", "torch", "torch-fp8", "torch-fp4"] + +_REFINER_QUANT_SKIP_DEFAULTS = ( + r"^proj_in$", + r"^proj_out$", + r"(^|\.)audio_", + r"audio_to_video", + r"video_to_audio", + r"av_cross_attn", + r"caption_projection", + r"time_embed", +) + + +def _move_tensor_attr(module: nn.Module, name: str, device: torch.device | str) -> None: + tensor = getattr(module, name, None) + if isinstance(tensor, nn.Parameter): + setattr(module, name, nn.Parameter(tensor.to(device), requires_grad=tensor.requires_grad)) + elif isinstance(tensor, Tensor): + setattr(module, name, tensor.to(device)) + + +def _offload_video_unused_audio_modules( + transformer: nn.Module, + device: torch.device | str = "cpu", +) -> None: + for name in ( + "audio_proj_in", + "audio_caption_projection", + "audio_time_embed", + "av_cross_attn_video_scale_shift", + "av_cross_attn_audio_scale_shift", + "av_cross_attn_video_a2v_gate", + "av_cross_attn_audio_v2a_gate", + "audio_rope", + "cross_attn_rope", + "cross_attn_audio_rope", + "audio_norm_out", + "audio_proj_out", + ): + child = getattr(transformer, name, None) + if isinstance(child, nn.Module): + child.to(device) + for block in getattr(transformer, "transformer_blocks", ()): + for name in ( + "audio_norm1", + "audio_attn1", + "audio_norm2", + "audio_attn2", + "audio_to_video_norm", + "audio_to_video_attn", + "video_to_audio_norm", + "video_to_audio_attn", + "audio_norm3", + "audio_ff", + ): + child = getattr(block, name, None) + if isinstance(child, nn.Module): + child.to(device) + + +def _move_ltx2_video_modules_to( + transformer: nn.Module, + device: torch.device | str, +) -> None: + for name in ("proj_in", "caption_projection", "time_embed", "rope", "norm_out", "proj_out"): + child = getattr(transformer, name, None) + if isinstance(child, nn.Module): + child.to(device) + _move_tensor_attr(transformer, "scale_shift_table", device) + for block in getattr(transformer, "transformer_blocks", ()): + _move_tensor_attr(block, "scale_shift_table", device) + for name in ("norm1", "attn1", "norm2", "attn2", "norm3", "ff"): + child = getattr(block, name, None) + if isinstance(child, nn.Module): + child.to(device) + + +class SanaWMLTX2Refiner(nn.Module): + """Run the SANA-WM LTX-2 latent refiner without importing a Sana checkout.""" + + def __init__( + self, + *, + refiner_root: str | Path, + gemma_root: str | Path, + dtype: torch.dtype, + device: torch.device | str, + precision: Precision = "bf16", + quant_backend: QuantBackend = "torch", + text_max_sequence_length: int = 1024, + cache_text_encoder: bool = False, + ) -> None: + super().__init__() + self.refiner_root = Path(refiner_root) + self.gemma_root = Path(gemma_root) + self.dtype = dtype + self.device = torch.device(device) + self.precision = precision + self.quant_backend = quant_backend + self.text_max_sequence_length = int(text_max_sequence_length) + self.cache_text_encoder = bool(cache_text_encoder) + self._quantized = False + self._text_encoder_built = False + self.transformer, self.connectors = self._load_diffusers_components() + + @torch.inference_mode() + def refine_latents( + self, + sana_latent: Tensor, + prompt: str, + *, + fps: float, + sink_size: int = 1, + seed: int = 42, + progress: bool = True, + sigmas: tuple[float, ...] = (0.909375, 0.725, 0.421875, 0.0), + ) -> Tensor: + """Refine Stage-1 VAE latents with the sink-bidirectional LTX-2 path.""" + if sana_latent.shape[2] <= sink_size: + raise ValueError( + f"Stage-1 latent has {sana_latent.shape[2]} frames but " + f"sink_size={sink_size}." + ) + + prompt_embeds, prompt_attention_mask = self._encode_prompt(prompt) + + _move_ltx2_video_modules_to(self.transformer, self.device) + _offload_video_unused_audio_modules(self.transformer, "cpu") + self.transformer.eval() + self._prepare_quantization() + + z = sana_latent.to(device=self.device, dtype=self.dtype) + sigmas_t = torch.tensor(sigmas, dtype=torch.float32, device=self.device) + start_sigma = float(sigmas_t[0]) + sink = z[:, :, :sink_size].contiguous() + current = z[:, :, sink_size:].contiguous() + generator = torch.Generator(device=self.device).manual_seed(int(seed)) + eps = torch.randn( + current.shape, + generator=generator, + device=self.device, + dtype=self.dtype, + ) + noisy = (1.0 - start_sigma) * current + start_sigma * eps + + iterator = range(len(sigmas_t) - 1) + if progress: + from tqdm.auto import tqdm + + iterator = tqdm(iterator, desc="refiner", unit="step") + + for step_index in iterator: + sigma = sigmas_t[step_index] + denoised = self._predict_current_x0( + sink=sink, + noisy_current=noisy, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + sigma=sigma, + fps=fps, + ) + noisy_tokens = _pack_latents( + noisy, + patch_size=self.transformer.config.patch_size, + patch_size_t=self.transformer.config.patch_size_t, + ) + velocity = (noisy_tokens.float() - denoised.float()) / sigma.float() + next_tokens = ( + noisy_tokens.float() + + velocity * (sigmas_t[step_index + 1] - sigma).float() + ) + noisy = _unpack_latents( + next_tokens.to(self.dtype), + num_frames=noisy.shape[2], + height=noisy.shape[3], + width=noisy.shape[4], + patch_size=self.transformer.config.patch_size, + patch_size_t=self.transformer.config.patch_size_t, + ) + + return torch.cat([sink, noisy], dim=2) + + def _load_diffusers_components(self) -> tuple[nn.Module, nn.Module]: + from diffusers.models.transformers.transformer_ltx2 import ( + LTX2VideoTransformer3DModel, + ) + from diffusers.pipelines.ltx2 import LTX2TextConnectors + + transformer = LTX2VideoTransformer3DModel.from_pretrained( + self.refiner_root, + subfolder="transformer", + torch_dtype=self.dtype, + ).eval() + connectors = LTX2TextConnectors.from_pretrained( + self.refiner_root, + subfolder="connectors", + torch_dtype=self.dtype, + ).eval() + return transformer, connectors + + def _prepare_quantization(self) -> None: + if self._quantized or self.precision == "bf16": + return + recipe = ( + TorchScaledMMFP8Recipe() + if self.precision == "fp8" + else TorchScaledMMFP4Recipe() + ) + converted, skipped = replace_linear_with_quant( + self.transformer, + recipe=recipe, + params_dtype=self.dtype, + skip_patterns=_REFINER_QUANT_SKIP_DEFAULTS, + ) + if converted <= 0: + raise RuntimeError( + f"SANA-WM refiner {self.precision} converted no Linear " + f"layers; skipped={skipped}." + ) + self._quantized = True + recipe_detail = "" + if isinstance(recipe, TorchScaledMMFP4Recipe): + recipe_detail = ( + f" rht={recipe.use_rht}" + f" global_scale={recipe.use_global_scale}" + f" weight_scale_2d={recipe.weight_scale_2d}" + ) + logger.info( + "[refiner-quant] precision={}{} converted {} Linear layers (skipped {})", + self.precision, + recipe_detail, + converted, + skipped, + ) + + def _ensure_text_encoder(self) -> None: + """Load the Gemma tokenizer + encoder once and cache them in CPU RAM. + + The encoder is ~20 GB. By default it is released after the one-shot + prompt encode to match upstream's single-generation path and avoid a + large GPU-to-CPU copy. Set ``cache_text_encoder=True`` for repeated + pipeline reuse, where paying the copy once can beat reloading Gemma on + every call. + """ + if self._text_encoder_built: + return + from transformers import AutoTokenizer, Gemma3ForConditionalGeneration + + t0 = time.perf_counter() + tokenizer = AutoTokenizer.from_pretrained(self.gemma_root) + tokenizer.padding_side = "left" + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + self.tokenizer = tokenizer + # Built on CPU; moved to the GPU on demand for each encode call only. + self.text_encoder = Gemma3ForConditionalGeneration.from_pretrained( + self.gemma_root, + dtype=self.dtype, + low_cpu_mem_usage=True, + ).eval() + self._text_encoder_built = True + logger.info( + "[timing] refiner text-encoder build: {:.3f}s", + time.perf_counter() - t0, + ) + + @torch.inference_mode() + def _encode_prompt(self, prompt: str) -> tuple[Tensor, Tensor]: + self._ensure_text_encoder() + tokenizer = self.tokenizer + + text_inputs = tokenizer( + [prompt.strip()], + padding="max_length", + max_length=self.text_max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + input_ids = text_inputs.input_ids.to(self.device) + attention_mask = text_inputs.attention_mask.to(self.device) + + # Pull the cached encoder onto the GPU only for the forward, then park + # it back on the CPU so it is absent during denoise/decode. + self.text_encoder.to(self.device) + text_backbone: nn.Module | None = None + try: + text_backbone = getattr(self.text_encoder, "model", self.text_encoder) + outputs = text_backbone( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + hidden_states = torch.stack(outputs.hidden_states, dim=-1) + sequence_lengths = attention_mask.sum(dim=-1) + del outputs + finally: + release_text_encoder = not self.cache_text_encoder + if release_text_encoder: + del text_backbone + del self.text_encoder + self._text_encoder_built = False + if torch.cuda.is_available(): + torch.cuda.empty_cache() + else: + self.text_encoder.to("cpu") + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + prompt_embeds = _pack_text_embeds( + hidden_states, + sequence_lengths, + device=self.device, + padding_side=tokenizer.padding_side, + ).to(dtype=self.dtype) + del hidden_states + + self.connectors.to(self.device) + connector_prompt_embeds, _, connector_attention_mask = self.connectors( + prompt_embeds, + attention_mask, + ) + self.connectors.to("cpu") + del prompt_embeds, attention_mask + return ( + connector_prompt_embeds.to(device=self.device, dtype=self.dtype), + connector_attention_mask.to(device=self.device), + ) + + def _predict_current_x0( + self, + *, + sink: Tensor, + noisy_current: Tensor, + prompt_embeds: Tensor, + prompt_attention_mask: Tensor, + sigma: Tensor, + fps: float, + ) -> Tensor: + full_latent = torch.cat([sink, noisy_current], dim=2) + batch_size, _, num_frames, height, width = full_latent.shape + latent_tokens = _pack_latents( + full_latent, + patch_size=self.transformer.config.patch_size, + patch_size_t=self.transformer.config.patch_size_t, + ) + n_context_tokens = _pack_latents( + sink, + patch_size=self.transformer.config.patch_size, + patch_size_t=self.transformer.config.patch_size_t, + ).shape[1] + + raw_timestep = torch.zeros( + batch_size, + latent_tokens.shape[1], + 1, + dtype=torch.float32, + device=self.device, + ) + raw_timestep[:, n_context_tokens:, 0] = sigma.float() + model_timestep = ( + raw_timestep.squeeze(-1) + * float(self.transformer.config.timestep_scale_multiplier) + ) + + velocity = self._forward_video_only( + hidden_states=latent_tokens, + encoder_hidden_states=prompt_embeds, + timestep=model_timestep, + encoder_attention_mask=prompt_attention_mask, + num_frames=num_frames, + height=height, + width=width, + fps=fps, + n_context_tokens=n_context_tokens, + ) + denoised = latent_tokens.float() - velocity.float() * raw_timestep + return denoised[:, n_context_tokens:, :].to(self.dtype) + + def _forward_video_only( + self, + *, + hidden_states: Tensor, + encoder_hidden_states: Tensor, + timestep: Tensor, + encoder_attention_mask: Tensor | None, + num_frames: int, + height: int, + width: int, + fps: float, + n_context_tokens: int, + ) -> Tensor: + transformer = self.transformer + batch_size = hidden_states.size(0) + encoder_attention_mask = _prepare_encoder_attention_mask( + encoder_attention_mask, + hidden_states.dtype, + ) + video_coords = transformer.rope.prepare_video_coords( + batch_size, + num_frames, + height, + width, + hidden_states.device, + fps=fps, + ) + video_rotary_emb = transformer.rope(video_coords, device=hidden_states.device) + + hidden_states = transformer.proj_in(hidden_states) + temb, embedded_timestep = transformer.time_embed( + timestep.flatten(), + batch_size=batch_size, + hidden_dtype=hidden_states.dtype, + ) + temb = temb.view(batch_size, -1, temb.size(-1)) + embedded_timestep = embedded_timestep.view( + batch_size, + -1, + embedded_timestep.size(-1), + ) + + encoder_hidden_states = transformer.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view( + batch_size, + -1, + hidden_states.size(-1), + ) + + for block in transformer.transformer_blocks: + hidden_states = _forward_video_block( + block=block, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + video_rotary_emb=video_rotary_emb, + encoder_attention_mask=encoder_attention_mask, + n_context_tokens=n_context_tokens, + ) + + scale_shift_values = ( + transformer.scale_shift_table[None, None] + + embedded_timestep[:, :, None] + ) + shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] + hidden_states = transformer.norm_out(hidden_states) + hidden_states = hidden_states * (1 + scale) + shift + return transformer.proj_out(hidden_states) + + +def _forward_video_block( + *, + block: nn.Module, + hidden_states: Tensor, + encoder_hidden_states: Tensor, + temb: Tensor, + video_rotary_emb: tuple[Tensor, Tensor], + encoder_attention_mask: Tensor | None, + n_context_tokens: int, +) -> Tensor: + batch_size = hidden_states.size(0) + norm_hidden_states = block.norm1(hidden_states) + num_ada_params = block.scale_shift_table.shape[0] + ada_values = block.scale_shift_table[None, None].to(temb.device) + temb.reshape( + batch_size, + temb.size(1), + num_ada_params, + -1, + ) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + ada_values.unbind(dim=2) + ) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + + attn_hidden_states = _streaming_self_attention( + attn=block.attn1, + hidden_states=norm_hidden_states, + query_rotary_emb=video_rotary_emb, + n_context_tokens=n_context_tokens, + ) + hidden_states = hidden_states + attn_hidden_states * gate_msa + + norm_hidden_states = block.norm2(hidden_states) + attn_hidden_states = block.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + query_rotary_emb=None, + attention_mask=encoder_attention_mask, + ) + hidden_states = hidden_states + attn_hidden_states + + norm_hidden_states = block.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp + return hidden_states + block.ff(norm_hidden_states) * gate_mlp + + +def _streaming_self_attention( + *, + attn: nn.Module, + hidden_states: Tensor, + query_rotary_emb: tuple[Tensor, Tensor], + n_context_tokens: int, +) -> Tensor: + from diffusers.models.transformers.transformer_ltx2 import ( + apply_interleaved_rotary_emb, + apply_split_rotary_emb, + ) + + gate_logits = ( + attn.to_gate_logits(hidden_states) + if attn.to_gate_logits is not None + else None + ) + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + query = attn.norm_q(query) + key = attn.norm_k(key) + + if attn.rope_type == "interleaved": + query = apply_interleaved_rotary_emb(query, query_rotary_emb) + key = apply_interleaved_rotary_emb(key, query_rotary_emb) + elif attn.rope_type == "split": + query = apply_split_rotary_emb(query, query_rotary_emb) + key = apply_split_rotary_emb(key, query_rotary_emb) + else: + raise ValueError(f"Unsupported LTX-2 RoPE type: {attn.rope_type}") + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + if n_context_tokens <= 0 or n_context_tokens >= query.shape[1]: + hidden_states = _refiner_attention(query, key, value) + else: + context_hidden_states = _refiner_attention( + query[:, :n_context_tokens], + key[:, :n_context_tokens], + value[:, :n_context_tokens], + ) + current_hidden_states = _refiner_attention( + query[:, n_context_tokens:], + key, + value, + ) + hidden_states = torch.cat( + [context_hidden_states, current_hidden_states], + dim=1, + ) + + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + if gate_logits is not None: + hidden_states = hidden_states.unflatten(2, (attn.heads, -1)) + gates = (2.0 * torch.sigmoid(gate_logits)).unsqueeze(-1) + hidden_states = hidden_states * gates + hidden_states = hidden_states.flatten(2, 3) + hidden_states = attn.to_out[0](hidden_states) + return attn.to_out[1](hidden_states) + + +def _refiner_attention(query: Tensor, key: Tensor, value: Tensor) -> Tensor: + hidden_states = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + return hidden_states.transpose(1, 2) + + +def _pack_text_embeds( + text_hidden_states: Tensor, + sequence_lengths: Tensor, + *, + device: torch.device | str, + padding_side: str = "left", + scale_factor: int = 8, + eps: float = 1e-6, +) -> Tensor: + batch_size, seq_len, hidden_dim, _ = text_hidden_states.shape + original_dtype = text_hidden_states.dtype + token_indices = torch.arange(seq_len, device=device).unsqueeze(0) + if padding_side == "right": + mask = token_indices < sequence_lengths[:, None] + elif padding_side == "left": + start_indices = seq_len - sequence_lengths[:, None] + mask = token_indices >= start_indices + else: + raise ValueError( + f"padding_side must be 'left' or 'right', got {padding_side!r}." + ) + mask = mask[:, :, None, None] + + masked = text_hidden_states.masked_fill(~mask, 0.0) + denom = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1) + masked_mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps) + x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin( + dim=(1, 2), + keepdim=True, + ) + x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax( + dim=(1, 2), + keepdim=True, + ) + normalized = (text_hidden_states - masked_mean) / (x_max - x_min + eps) + normalized = normalized * scale_factor + normalized = normalized.flatten(2) + mask_flat = mask.squeeze(-1).expand(-1, -1, normalized.shape[-1]) + normalized = normalized.masked_fill(~mask_flat, 0.0) + return normalized.to(dtype=original_dtype) + + +def _pack_latents(latents: Tensor, patch_size: int = 1, patch_size_t: int = 1) -> Tensor: + batch_size, _, num_frames, height, width = latents.shape + post_patch_num_frames = num_frames // patch_size_t + post_patch_height = height // patch_size + post_patch_width = width // patch_size + latents = latents.reshape( + batch_size, + -1, + post_patch_num_frames, + patch_size_t, + post_patch_height, + patch_size, + post_patch_width, + patch_size, + ) + latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7) + return latents.flatten(4, 7).flatten(1, 3) + + +def _unpack_latents( + latents: Tensor, + *, + num_frames: int, + height: int, + width: int, + patch_size: int = 1, + patch_size_t: int = 1, +) -> Tensor: + batch_size = latents.size(0) + latents = latents.reshape( + batch_size, + num_frames, + height, + width, + -1, + patch_size_t, + patch_size, + patch_size, + ) + latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7) + return latents.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + +def _prepare_encoder_attention_mask(mask: Tensor | None, dtype: torch.dtype) -> Tensor | None: + if mask is None: + return None + if mask.ndim != 2: + return mask + if bool(torch.all(mask)): + return None + return ((1 - mask.to(dtype)) * -10000.0).unsqueeze(1) diff --git a/integrations/sana/sana_wm/runner.py b/integrations/sana/sana_wm/runner.py new file mode 100644 index 000000000..ec7c73161 --- /dev/null +++ b/integrations/sana/sana_wm/runner.py @@ -0,0 +1,602 @@ +# 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. + +"""SANA-WM bidirectional runner.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Literal + +import numpy as np +import torch +from loguru import logger + +from flashdreams.infra.config import derive_config +from flashdreams.infra.pipeline import StreamInferencePipelineConfig +from flashdreams.infra.runner import Runner, RunnerConfig +from flashdreams.infra.runner_io import ensure_output_dir, resolve_prompt_value, runner_artifact_path +from sana_wm.camera import ( + action_string_to_c2w, + default_intrinsics_vec4, + fit_camera_trajectory, + load_intrinsics, + resize_center_crop_geometry, + snap_num_frames, + transform_intrinsics_for_crop, +) +from sana_wm.conditioning import SanaWMI2VConditioningRequest +from sana_wm.constants import ( + DEFAULT_ACTION, + DEFAULT_FPS, + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, + SANA_WM_CONFIG_PATH, + SANA_WM_MODEL_PATH, + SANA_WM_REFINER_GEMMA_ROOT, + SANA_WM_REFINER_ROOT, + SANA_WM_VAE_TEMPORAL_COMPRESSION, +) +from sana_wm.decoder import SanaWMDecodedVideo + +SamplingAlgo = Literal["auto", "flow_euler_ltx"] +"""Sampling algorithms exposed by the SANA-WM runner.""" + +Precision = Literal["bf16", "fp8", "fp4"] +"""SANA-WM Stage-1/refiner precision modes.""" + +QuantBackend = Literal["auto", "torch", "torch-fp8", "torch-fp4"] +"""Low-precision linear backend for FP8/FP4 SANA-WM paths.""" + +ResolvedQuantBackend = Literal["torch", "torch-fp8", "torch-fp4"] +"""Concrete low-precision linear backend selected after resolving ``auto``.""" + + +@dataclass(kw_only=True) +class SanaWMRunnerConfig(RunnerConfig): + """Runner config for the SANA-WM bidirectional release.""" + + _target: type["SanaWMRunner"] = field(default_factory=lambda: SanaWMRunner) + + device: str = "auto" + """Torch device for SANA-WM. ``"auto"`` picks CUDA when available.""" + + image_path: Path | None = None + """Path to the first-frame RGB image. Required at ``run()`` time.""" + + prompt: str = "" + """Inline text prompt. A non-empty value wins over ``prompt_path``.""" + + prompt_path: Path | None = None + """Fallback prompt file read when ``prompt`` is empty.""" + + camera_path: Path | None = None + """Optional ``.npy`` camera-to-world trajectory shaped ``[F, 4, 4]``.""" + + intrinsics_path: Path | None = None + """Optional ``.npy`` intrinsics shaped ``[3, 3]``, ``[F, 3, 3]``, + ``[4]``, or ``[F, 4]``. When omitted, intrinsics are derived from the + first-frame size using ``intrinsics_hfov_deg`` with a centered principal + point.""" + + intrinsics_hfov_deg: float = 90.0 + """Horizontal field of view in degrees used to derive intrinsics when + ``intrinsics_path`` is not provided. The public demo intrinsics correspond + to ~90 degrees; lower values are narrower/more zoomed-in.""" + + action: str | None = DEFAULT_ACTION + """Action DSL used to derive camera motion when ``camera_path`` is not + provided. The action is repeated or truncated to match ``num_frames``.""" + + translation_speed: float = 0.025 + """Per-frame action translation speed.""" + + rotation_speed_deg: float = 0.6 + """Per-frame action rotation speed in degrees.""" + + num_frames: int = 161 + """Requested output frames before LTX2-VAE stride snapping.""" + + fps: int = DEFAULT_FPS + """Output video frame rate.""" + + step: int = 60 + """Stage-1 DiT sampling steps.""" + + cfg_scale: float = 5.0 + """Classifier-free guidance scale for Stage 1.""" + + flow_shift: float | None = None + """Optional scheduler flow-shift override.""" + + sampling_algo: SamplingAlgo = "auto" + """Stage-1 sampler. ``"auto"`` uses ``"flow_euler_ltx"``.""" + + save_stage1: bool = False + """Also decode the unrefined Stage-1 latent when the refiner is enabled.""" + + negative_prompt: str = "" + """Negative prompt used when ``cfg_scale > 1``.""" + + seed: int = 42 + """Stage-1 random seed.""" + + config_path: str = SANA_WM_CONFIG_PATH + """SANA-WM inference YAML path or ``hf://`` URI.""" + + model_path: str = SANA_WM_MODEL_PATH + """Stage-1 checkpoint path or ``hf://`` URI.""" + + stage1_precision: Precision = "bf16" + """Stage-1 DiT compute precision. ``"bf16"`` is the default; ``"fp8"`` + requires Hopper or newer; ``"fp4"`` requires Blackwell.""" + + no_refiner: bool = False + """Skip the LTX-2 refiner and decode Stage-1 latents directly.""" + + refiner_precision: Precision = "bf16" + """LTX-2 refiner compute precision. ``"bf16"`` is the default; + ``"fp8"`` requires Hopper or newer; ``"fp4"`` requires Blackwell. + Ignored when ``no_refiner`` is ``True``.""" + + quant_backend: QuantBackend = "auto" + """Backend for quantized linear layers. ``"auto"`` and ``"torch"`` allow + both FP8 and FP4 Torch replacements. ``"torch-fp8"`` and ``"torch-fp4"`` + select one ``torch._scaled_mm`` replacement explicitly.""" + + refiner_root: str = SANA_WM_REFINER_ROOT + """LTX-2 refiner root path or ``hf://`` URI.""" + + refiner_gemma_root: str = SANA_WM_REFINER_GEMMA_ROOT + """Gemma text-encoder root for the LTX-2 refiner.""" + + refiner_seed: int = 42 + """Refiner random seed.""" + + sink_size: int = 1 + """Number of sink latent frames used by the refiner.""" + + offload_vae: bool = False + """Move the VAE to CPU between encode/decode phases.""" + + offload_stage1: bool = False + """Tear down the Stage-1 DiT after sampling to free memory for decode/ + refine. Default keeps it resident (fastest); enable on memory-constrained + GPUs.""" + + offload_refiner: bool = False + """Build and release the refiner only around refinement.""" + + offload_text_encoder: bool = False + """Move the Stage-1 text encoder to CPU between prompt encodes.""" + + +class SanaWMRunner(Runner[SanaWMRunnerConfig, object]): + """CLI driver for SANA-WM configs.""" + + config: SanaWMRunnerConfig + + def __init__(self, config: SanaWMRunnerConfig) -> None: + self.config = config + self.local_rank = int(os.environ.get("LOCAL_RANK", "0")) + if torch.distributed.is_initialized(): + self.world_size = torch.distributed.get_world_size() + self.global_rank = torch.distributed.get_rank() + else: + self.world_size = int(os.environ.get("WORLD_SIZE", "1")) + self.global_rank = int(os.environ.get("RANK", "0")) + self.is_rank_zero = self.global_rank == 0 + + def _resolve_prompt(self) -> str: + """Resolve the prompt from inline text or ``prompt_path``.""" + cfg = self.config + if cfg.prompt: + return resolve_prompt_value(cfg.prompt) + if cfg.prompt_path is None: + raise ValueError("SanaWMRunner requires --prompt or --prompt-path.") + return resolve_prompt_value(cfg.prompt_path) + + def _resolve_device(self) -> torch.device: + """Return the device used by SANA-WM.""" + if self.config.device == "auto": + if torch.cuda.is_available(): + return torch.device(f"cuda:{self.local_rank}") + return torch.device("cpu") + if self.config.device == "cuda" and torch.cuda.is_available(): + return torch.device(f"cuda:{self.local_rank}") + return torch.device(self.config.device) + + def _resolve_trajectory(self, *, num_frames: int) -> np.ndarray: + """Load, fit, or roll out the camera-to-world trajectory.""" + if self.config.camera_path is not None: + c2w = np.load(self.config.camera_path).astype(np.float32) + if c2w.ndim != 3 or c2w.shape[1:] != (4, 4): + raise ValueError( + f"--camera-path must be a [F, 4, 4] .npy; got {c2w.shape}." + ) + if c2w.shape[0] != num_frames and self.is_rank_zero: + logger.info( + "Fitting --camera-path trajectory from {} to {} frames.", + c2w.shape[0], + num_frames, + ) + return fit_camera_trajectory(c2w, num_frames) + if not self.config.action: + raise ValueError("SanaWMRunner requires --camera-path or --action.") + if self.is_rank_zero: + logger.info( + "No --camera-path provided; deriving a {}-frame trajectory " + "from --action.", + num_frames, + ) + return action_string_to_c2w( + self.config.action, + translation_speed=self.config.translation_speed, + rotation_speed_deg=self.config.rotation_speed_deg, + num_frames=num_frames, + ) + + def run(self) -> None: + """Run SANA-WM bidirectional inference and write outputs.""" + cfg = self.config + if cfg.image_path is None: + raise ValueError("SanaWMRunner requires --image-path.") + + device = self._resolve_device() + quant_backend = _resolve_quant_backend( + cfg.quant_backend, + _active_quantized_precisions( + stage1_precision=cfg.stage1_precision, + refiner_precision=cfg.refiner_precision, + refiner_enabled=not cfg.no_refiner, + ), + ) + _validate_precision_request( + device=device, + stage1_precision=cfg.stage1_precision, + refiner_precision=cfg.refiner_precision, + refiner_enabled=not cfg.no_refiner, + quant_backend=cfg.quant_backend, + ) + prompt = self._resolve_prompt() + image, c2w, intrinsics_vec4, num_frames = self._prepare_inputs() + pipeline_cfg = _pipeline_config( + cfg, + quant_backend=quant_backend, + ) + pipeline = pipeline_cfg.setup().to(device).eval() + sampling_algo = self._sampling_algo() + if sampling_algo != "flow_euler_ltx": + raise ValueError( + "SANA-WM requires flow_euler_ltx for the " + f"bidirectional runner; got {sampling_algo!r}." + ) + cache = pipeline.initialize_cache( + decoder_context={ + "prompt": prompt, + "fps": cfg.fps, + "save_stage1": cfg.save_stage1, + "refiner_seed": cfg.refiner_seed, + "sink_size": cfg.sink_size, + } + ) + with torch.inference_mode(): + decoded = pipeline.generate( + 0, + cache, + input=SanaWMI2VConditioningRequest( + image=image, + prompt=prompt, + poses_c2w=c2w, + intrinsics_vec4=intrinsics_vec4, + num_frames=num_frames, + fps=cfg.fps, + steps=cfg.step, + cfg_scale=cfg.cfg_scale, + flow_shift=cfg.flow_shift, + seed=cfg.seed, + negative_prompt=cfg.negative_prompt, + ), + ) + pipeline.finalize(0, cache) + if not isinstance(decoded, SanaWMDecodedVideo): + raise TypeError( + "SANA-WM pipeline decoder returned " + f"{type(decoded).__name__}, expected SanaWMDecodedVideo." + ) + if not self.is_rank_zero: + return + ensure_output_dir(cfg.output_dir) + _write_video( + runner_artifact_path(cfg.output_dir, cfg.runner_name, "mp4"), + decoded.video_hwc, + cfg.fps, + ) + if decoded.stage1_video_hwc is not None: + _write_video( + runner_artifact_path(cfg.output_dir, f"{cfg.runner_name}_stage1", "mp4"), + decoded.stage1_video_hwc, + cfg.fps, + ) + + def _prepare_inputs(self) -> tuple[object, np.ndarray, np.ndarray, int]: + """Load/crop the input image and prepare c2w/intrinsics.""" + from PIL import Image + + cfg = self.config + assert cfg.image_path is not None + image = Image.open(cfg.image_path).convert("RGB") + snapped = snap_num_frames( + cfg.num_frames, + stride=SANA_WM_VAE_TEMPORAL_COMPRESSION, + ) + if snapped != cfg.num_frames and self.is_rank_zero: + logger.warning( + "SANA-WM requires num_frames = 8k+1; requested {} snapped to {}.", + cfg.num_frames, + snapped, + ) + num_frames = snapped + c2w = self._resolve_trajectory(num_frames=num_frames) + + resized_size, crop_offset = resize_center_crop_geometry( + image.size, + target_h=DEFAULT_VIDEO_HEIGHT, + target_w=DEFAULT_VIDEO_WIDTH, + ) + resized = image.resize(resized_size, Image.LANCZOS) + left, top = crop_offset + cropped = resized.crop( + ( + left, + top, + left + DEFAULT_VIDEO_WIDTH, + top + DEFAULT_VIDEO_HEIGHT, + ) + ) + if cfg.intrinsics_path is None: + intrinsics_src = default_intrinsics_vec4( + image.size, num_frames, hfov_deg=cfg.intrinsics_hfov_deg + ) + if self.is_rank_zero: + logger.info( + "No --intrinsics-path provided; deriving intrinsics from " + "image size {} at hfov={} deg (principal point centered).", + image.size, + cfg.intrinsics_hfov_deg, + ) + else: + intrinsics_src = load_intrinsics(cfg.intrinsics_path, num_frames) + intrinsics_vec4 = transform_intrinsics_for_crop( + intrinsics_src, + image.size, + resized_size, + crop_offset, + ) + return cropped, c2w, intrinsics_vec4, num_frames + + def _sampling_algo(self) -> SamplingAlgo: + """Resolve the sampler selection.""" + if self.config.sampling_algo == "auto": + return "flow_euler_ltx" + return self.config.sampling_algo + + +def _pipeline_config( + cfg: SanaWMRunnerConfig, + *, + quant_backend: ResolvedQuantBackend, +) -> StreamInferencePipelineConfig: + """Apply CLI runtime fields to the SANA-WM pipeline literal.""" + scheduler_updates: dict[str, object] = {"num_inference_steps": cfg.step} + if cfg.flow_shift is not None: + scheduler_updates["shift"] = cfg.flow_shift + return derive_config( + cfg.pipeline, + diffusion_model=dict( + seed=cfg.seed, + scheduler=scheduler_updates, + transformer=dict( + config_path=cfg.config_path, + checkpoint_path=cfg.model_path, + stage1_precision=cfg.stage1_precision, + quant_backend=quant_backend, + offload_stage1=cfg.offload_stage1, + ), + ), + encoder=dict( + config_path=cfg.config_path, + text_encoder=dict( + config_path=cfg.config_path, + stage1_precision=cfg.stage1_precision, + quant_backend=quant_backend, + offload_text_encoder=cfg.offload_text_encoder, + ), + first_frame_encoder=dict( + config_path=cfg.config_path, + offload_vae=cfg.offload_vae, + ), + camera_encoder=dict( + height=DEFAULT_VIDEO_HEIGHT, + width=DEFAULT_VIDEO_WIDTH, + ), + height=DEFAULT_VIDEO_HEIGHT, + width=DEFAULT_VIDEO_WIDTH, + ), + decoder=dict( + vae_decoder=dict( + config_path=cfg.config_path, + offload_vae=cfg.offload_vae, + ), + refiner=( + None + if cfg.no_refiner + else dict( + refiner_root=cfg.refiner_root, + refiner_gemma_root=cfg.refiner_gemma_root, + refiner_precision=cfg.refiner_precision, + quant_backend=quant_backend, + offload_refiner=cfg.offload_refiner, + ) + ), + ), + ) + + +def _write_video(path: Path, video_hwc: np.ndarray, fps: int) -> Path: + """Write an HWC uint8 video to ``path``.""" + import imageio.v3 as iio + + iio.imwrite(path, video_hwc, fps=fps) + logger.info("Saved {}", path) + return path + + +def _validate_precision_request( + *, + device: torch.device, + stage1_precision: Precision, + refiner_precision: Precision, + refiner_enabled: bool, + quant_backend: QuantBackend, +) -> None: + """Fail early when requested quantized precision cannot run.""" + quantized = _active_quantized_precisions( + stage1_precision=stage1_precision, + refiner_precision=refiner_precision, + refiner_enabled=refiner_enabled, + ) + if not quantized: + return + + if device.type != "cuda" or not torch.cuda.is_available(): + raise ValueError( + "SANA-WM fp8/fp4 precision requires a CUDA device; " + f"resolved device is {device}." + ) + + capability = torch.cuda.get_device_capability(device) + major, minor = capability + sm_name = f"sm_{major}{minor}" + if "fp8" in quantized and major < 9: + raise ValueError( + "SANA-WM fp8 precision requires a Hopper or newer GPU " + f"(sm_90+); detected {sm_name}." + ) + if "fp4" in quantized and major < 10: + raise ValueError( + "SANA-WM fp4/NVFP4 precision requires a Blackwell GPU " + f"(sm_100+); detected {sm_name}. Use bf16 or fp8 on this GPU." + ) + + resolved_backend = _resolve_quant_backend(quant_backend, quantized) + if resolved_backend == "torch-fp8": + _validate_torch_fp8_backend(quantized) + elif resolved_backend == "torch-fp4": + _validate_torch_fp4_backend(quantized) + elif resolved_backend == "torch": + _validate_quant_backend(quantized) + else: + raise ValueError(f"Unsupported SANA-WM quant backend: {quant_backend!r}.") + + +def _active_quantized_precisions( + *, + stage1_precision: Precision, + refiner_precision: Precision, + refiner_enabled: bool, +) -> list[Precision]: + """Return active non-BF16 precision requests in execution order.""" + active_precisions = [stage1_precision] + if refiner_enabled: + active_precisions.append(refiner_precision) + return [precision for precision in active_precisions if precision != "bf16"] + + +def _resolve_quant_backend( + quant_backend: QuantBackend, + quantized_precisions: list[Precision], +) -> ResolvedQuantBackend: + """Resolve ``auto`` to the Torch low-precision backend.""" + del quantized_precisions + if quant_backend == "auto": + return "torch" + if quant_backend in {"torch", "torch-fp8", "torch-fp4"}: + return quant_backend + raise ValueError(f"Unsupported SANA-WM quant backend: {quant_backend!r}.") + + +def _validate_torch_fp8_backend(precisions: list[Precision]) -> None: + """Validate the PyTorch scaled-MM backend request.""" + unsupported = sorted({precision for precision in precisions if precision != "fp8"}) + if unsupported: + raise ValueError( + "SANA-WM --quant-backend torch-fp8 accepts fp8 only; " + f"unsupported requested precision(s): {', '.join(unsupported)}. " + "Use --quant-backend torch-fp4 or --quant-backend torch for fp4." + ) + _validate_torch_fp8_primitives() + + +def _validate_torch_fp4_backend(precisions: list[Precision]) -> None: + """Validate the PyTorch NVFP4 backend request.""" + unsupported = sorted({precision for precision in precisions if precision != "fp4"}) + if unsupported: + raise ValueError( + "SANA-WM --quant-backend torch-fp4 accepts fp4 only; " + f"unsupported requested precision(s): {', '.join(unsupported)}. " + "Use --quant-backend torch-fp8 or --quant-backend torch for fp8." + ) + _validate_torch_fp4_primitives() + + +def _validate_quant_backend(precisions: list[Precision]) -> None: + """Validate PyTorch low-precision primitives.""" + if "fp8" in precisions: + _validate_torch_fp8_primitives() + if "fp4" in precisions: + _validate_torch_fp4_primitives() + + +def _validate_torch_fp8_primitives() -> None: + if not hasattr(torch, "_scaled_mm") or not hasattr(torch, "float8_e4m3fn"): + raise RuntimeError( + "SANA-WM --quant-backend torch-fp8 requires PyTorch with " + "torch._scaled_mm and torch.float8_e4m3fn support." + ) + + +def _validate_torch_fp4_primitives() -> None: + missing = [ + name + for name in ("_scaled_mm", "float4_e2m1fn_x2", "float8_e4m3fn") + if not hasattr(torch, name) + ] + if missing: + raise RuntimeError( + "SANA-WM fp4 requires PyTorch with " + f"{', '.join(f'torch.{name}' for name in missing)} support." + ) + + +__all__ = [ + "Precision", + "QuantBackend", + "SamplingAlgo", + "SanaWMRunner", + "SanaWMRunnerConfig", +] diff --git a/integrations/sana/sana_wm/scheduler.py b/integrations/sana/sana_wm/scheduler.py new file mode 100644 index 000000000..168c377db --- /dev/null +++ b/integrations/sana/sana_wm/scheduler.py @@ -0,0 +1,196 @@ +# 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. + +"""SANA-WM LTX-style Euler scheduler boundary.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np +import torch +from torch import Tensor + +from flashdreams.infra.diffusion.scheduler import ( + FlowPredictor, + Scheduler, + SchedulerConfig, +) + + +@dataclass(kw_only=True) +class SanaWMLTXEulerSchedulerConfig(SchedulerConfig): + """Config for SANA-WM's LTX-style flow-matching Euler scheduler.""" + + _target: type["SanaWMLTXEulerScheduler"] = field( + default_factory=lambda: SanaWMLTXEulerScheduler + ) + + num_inference_steps: int = 60 + """Default number of Euler steps when using the generic scheduler API.""" + + shift: float = 9.8 + """Default flow-match schedule shift when using the generic scheduler API.""" + + num_train_timesteps: int = 1000 + """Training timestep scale used by the public SANA-WM release.""" + + +class SanaWMLTXEulerScheduler(Scheduler): + """Euler scheduler with SANA-WM per-token timestep support.""" + + config: SanaWMLTXEulerSchedulerConfig + + def __init__(self, config: SanaWMLTXEulerSchedulerConfig) -> None: + super().__init__(config) + self.config = config + + def timesteps( + self, + *, + num_inference_steps: int, + shift: float, + device: torch.device | str, + ) -> Tensor: + """Return diffusers-compatible FlowMatch Euler timesteps.""" + steps = int(num_inference_steps) + if steps <= 0: + raise ValueError(f"num_inference_steps must be > 0, got {steps}.") + num_train_timesteps = int(self.config.num_train_timesteps) + init_timesteps = np.linspace( + 1, + num_train_timesteps, + num_train_timesteps, + dtype=np.float32, + )[::-1].copy() + init_sigmas = init_timesteps / num_train_timesteps + init_sigmas = shift * init_sigmas / (1.0 + (shift - 1.0) * init_sigmas) + sigma_min = float(init_sigmas[-1]) + sigma_max = float(init_sigmas[0]) + timesteps = np.linspace( + sigma_max * num_train_timesteps, + sigma_min * num_train_timesteps, + steps, + ) + sigmas = timesteps / num_train_timesteps + sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas) + sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) + return torch.from_numpy(sigmas * num_train_timesteps).to(device=device) + + def step_ltx( + self, + *, + model_output: Tensor, + timestep: Tensor, + next_timestep: Tensor, + sample: Tensor, + per_token_timesteps: Tensor | None = None, + schedule_timesteps: Tensor | None = None, + ) -> Tensor: + """Apply one SANA-WM LTX Euler step in token layout. + + Args: + model_output: Flow tensor in ``[B, N, C]`` token layout. + timestep: Current scalar training-scale timestep. + next_timestep: Next scalar training-scale timestep. + sample: Current latent in ``[B, N, C]`` token layout. + per_token_timesteps: Optional ``[B, N]`` timestep table. Tokens + with timestep ``0`` stay at sigma zero, which matches the + first-frame-pinning branch used by SANA-WM. + schedule_timesteps: Optional full scheduler timestep table used + to find the next lower sigma for each per-token timestep. + + Returns: + Updated sample in ``[B, N, C]`` layout. + """ + num_train_timesteps = float(self.config.num_train_timesteps) + sample = sample.to(torch.float32) + if per_token_timesteps is None: + sigma = timestep.to(device=sample.device, dtype=sample.dtype) + sigma_next = next_timestep.to(device=sample.device, dtype=sample.dtype) + dt = (sigma_next - sigma) / num_train_timesteps + return (sample + dt * model_output).to(model_output.dtype) + + per_token_sigmas = ( + per_token_timesteps.to(device=sample.device, dtype=sample.dtype) + / num_train_timesteps + ) + if schedule_timesteps is not None: + sigmas = ( + schedule_timesteps.to(device=sample.device, dtype=sample.dtype) + / num_train_timesteps + ) + sigmas = sigmas[:, None, None] + lower_mask = sigmas < per_token_sigmas[None] - 1e-6 + lower_sigmas = (lower_mask * sigmas).max(dim=0).values + else: + next_sigma_scalar = ( + next_timestep.to(device=sample.device, dtype=sample.dtype) + / num_train_timesteps + ) + lower_sigmas = torch.where( + per_token_sigmas > next_sigma_scalar + 1e-6, + next_sigma_scalar.expand_as(per_token_sigmas), + torch.zeros_like(per_token_sigmas), + ) + dt = per_token_sigmas - lower_sigmas + return sample + dt.unsqueeze(-1) * model_output + + def sample( + self, + initial_noise: Tensor, + predict_flow: FlowPredictor, + rng: torch.Generator | None = None, + ) -> Tensor: + """Run the generic scalar-timestep Euler loop.""" + del rng + timesteps = self.timesteps( + num_inference_steps=self.config.num_inference_steps, + shift=self.config.shift, + device=initial_noise.device, + ) + noisy = initial_noise + for index, timestep in enumerate(timesteps[:-1]): + flow = predict_flow(noisy, timestep.to(dtype=initial_noise.dtype)) + noisy = self.step_ltx( + model_output=flow, + timestep=timestep, + next_timestep=timesteps[index + 1], + sample=noisy, + ) + return noisy.to(initial_noise.dtype) + + def add_noise( + self, + clean_input: Tensor, + timestep: Tensor, + rng: torch.Generator | None = None, + ) -> Tensor: + """Apply forward flow-match corruption at ``timestep``.""" + sigma = timestep.to(device=clean_input.device, dtype=clean_input.dtype) + sigma = sigma / float(self.config.num_train_timesteps) + noise = torch.randn( + clean_input.shape, + generator=rng, + device=clean_input.device, + dtype=clean_input.dtype, + ) + return ((1.0 - sigma) * clean_input + sigma * noise).to(clean_input.dtype) + + +__all__ = [ + "SanaWMLTXEulerScheduler", + "SanaWMLTXEulerSchedulerConfig", +] diff --git a/integrations/sana/sana_wm/stage1_model.py b/integrations/sana/sana_wm/stage1_model.py new file mode 100644 index 000000000..2f93cecce --- /dev/null +++ b/integrations/sana/sana_wm/stage1_model.py @@ -0,0 +1,2498 @@ +# 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. + +"""SANA-WM Stage-1 DiT module definitions. + +The module names and tensor shapes are checkpoint-facing API: keep them stable +unless the public checkpoint contract changes. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +import math +import os + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor + +try: + import triton + import triton.language as tl +except ImportError: # pragma: no cover - exercised in minimal CPU environments. + triton = None + tl = None + + +@dataclass(frozen=True) +class SanaWMStage1Spec: + """Static architecture values for the public SANA-WM bidirectional DiT.""" + + latent_channels: int = 128 + hidden_size: int = 2240 + text_dim: int = 2304 + timestep_dim: int = 256 + depth: int = 20 + num_heads: int = 20 + head_dim: int = 112 + max_text_length: int = 300 + latent_grid_size: tuple[int, int] = (22, 40) + mlp_ratio: int = 3 + conv_kernel_size: int = 4 + temporal_kernel_size: int = 3 + plucker_channels: int = 48 + raymap_channels: int = 3 + softmax_every_n: int = 4 + chunk_size: int | None = None + chunk_split_strategy: str = "first_chunk_plus_one" + + @property + def mlp_inner_size(self) -> int: + """Return the GLUMBConv hidden expansion width.""" + return self.hidden_size * self.mlp_ratio * 2 + + @property + def gated_mlp_size(self) -> int: + """Return the width consumed by the pointwise output projection.""" + return self.hidden_size * self.mlp_ratio + + def block_uses_gdn(self, index: int) -> bool: + """Return whether a block has GDN convolution checkpoint tensors.""" + return (index + 1) % self.softmax_every_n != 0 + + +SANA_WM_STAGE1_SPEC = SanaWMStage1Spec() +"""Architecture spec for the public SANA-WM bidirectional Stage-1 checkpoint.""" + + +def _env_flag_enabled(name: str) -> bool: + return os.environ.get(name, "").lower() in {"1", "true", "yes", "on"} + + +_DISABLE_UCPE_FAST = _env_flag_enabled("SANA_WM_STAGE1_DISABLE_UCPE_FAST") +_DISABLE_GLU_FAST = _env_flag_enabled("SANA_WM_STAGE1_DISABLE_GLU_FAST") +_DISABLE_RMS_RELU_FAST = _env_flag_enabled("SANA_WM_STAGE1_DISABLE_RMS_RELU_FAST") + + +@dataclass(frozen=True) +class _CameraProjectionCache: + raymats: Tensor + proj: Tensor + proj_q: Tensor + proj_kv: Tensor + rope_cam: Tensor | None + + +class RMSNorm(nn.Module): + """RMSNorm parameter container matching the SANA checkpoint schema.""" + + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.empty(hidden_size)) + self.eps = eps + nn.init.ones_(self.weight) + + def forward(self, x: Tensor) -> Tensor: + """Apply RMS normalization using the stored scale parameter.""" + dtype = x.dtype + normed = x.float() * torch.rsqrt( + x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps + ) + return (normed * self.weight).to(dtype=dtype) + + +class Conv3dProjector(nn.Module): + """Named 1x1x1 projection used by latent, ray, and Plucker embedders.""" + + def __init__(self, in_channels: int, hidden_size: int) -> None: + super().__init__() + self.proj = nn.Conv3d(in_channels, hidden_size, kernel_size=1) + + def forward(self, x: Tensor) -> Tensor: + """Project a 5D tensor to hidden channels.""" + return self.proj(x) + + +class TimestepEmbedder(nn.Module): + """Two-layer timestep embedder with checkpoint-compatible names.""" + + def __init__(self, spec: SanaWMStage1Spec) -> None: + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(spec.timestep_dim, spec.hidden_size), + nn.SiLU(), + nn.Linear(spec.hidden_size, spec.hidden_size), + ) + + def forward(self, t: Tensor) -> Tensor: + """Project timestep features into the DiT hidden width.""" + t_freq = _timestep_embedding(t.flatten(), self.mlp[0].in_features) + return self.mlp(t_freq.to(device=t.device, dtype=self.mlp[0].weight.dtype)) + + +class TextProjection(nn.Module): + """Two-layer text projection matching ``y_embedder.y_proj`` keys.""" + + def __init__(self, spec: SanaWMStage1Spec) -> None: + super().__init__() + self.fc1 = nn.Linear(spec.text_dim, spec.hidden_size) + self.act = nn.SiLU() + self.fc2 = nn.Linear(spec.hidden_size, spec.hidden_size) + + def forward(self, y: Tensor) -> Tensor: + """Project text encoder activations into the DiT hidden width.""" + return self.fc2(self.act(self.fc1(y))) + + +class TextEmbedder(nn.Module): + """Text embedding/projection container matching SANA checkpoint keys.""" + + def __init__(self, spec: SanaWMStage1Spec) -> None: + super().__init__() + self.y_embedding = nn.Parameter(torch.empty(spec.max_text_length, spec.text_dim)) + self.y_proj = TextProjection(spec) + + def forward(self, y: Tensor | None) -> Tensor: + """Project text embeddings, using the learned null embedding when absent.""" + if y is None: + y = self.y_embedding.unsqueeze(0) + if y.ndim == 4 and y.shape[1] == 1: + y = y.squeeze(1) + return self.y_proj(y) + + +class FinalLayer(nn.Module): + """Final AdaLN and latent-channel projection container.""" + + def __init__(self, spec: SanaWMStage1Spec) -> None: + super().__init__() + self.norm_final = nn.LayerNorm( + spec.hidden_size, + elementwise_affine=False, + eps=1e-6, + ) + self.scale_shift_table = nn.Parameter(torch.empty(2, spec.hidden_size)) + self.linear = nn.Linear(spec.hidden_size, spec.latent_channels) + + def forward(self, x: Tensor, t: Tensor, *, frames: int) -> Tensor: + """Project hidden tokens back to latent channels.""" + if t.ndim > 2: + batch, tokens, channels = x.shape + shift, scale = ( + self.scale_shift_table[None, None, :, :] + t.transpose(1, 2) + ).chunk(2, dim=-2) + x = _modulate( + self.norm_final(x).reshape(batch, frames, -1, channels), + shift, + scale, + ).reshape(batch, tokens, channels) + else: + shift, scale = (self.scale_shift_table[None] + t[:, None]).chunk( + 2, + dim=1, + ) + x = _modulate(self.norm_final(x), shift, scale) + return self.linear(x) + + +class Conv2dContainer(nn.Module): + """Expose a convolution under a stable ``.conv`` attribute.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int], + *, + groups: int = 1, + bias: bool = True, + padding: int | tuple[int, int] = 0, + ) -> None: + super().__init__() + self.conv = nn.Conv2d( + in_channels, + out_channels, + kernel_size=kernel_size, + groups=groups, + bias=bias, + padding=padding, + ) + + def forward(self, x: Tensor) -> Tensor: + """Apply the contained convolution.""" + return self.conv(x) + + +class GLUMBConvTemp(nn.Module): + """Checkpoint-compatible GLUMBConvTemp feed-forward container.""" + + def __init__(self, spec: SanaWMStage1Spec) -> None: + super().__init__() + self.inverted_conv = Conv2dContainer(spec.hidden_size, spec.mlp_inner_size, 1) + self.depth_conv = Conv2dContainer( + spec.mlp_inner_size, + spec.mlp_inner_size, + 3, + groups=spec.mlp_inner_size, + padding=1, + ) + self.point_conv = Conv2dContainer( + spec.gated_mlp_size, + 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, + ) + nn.init.zeros_(self.t_conv.weight) + + def forward(self, x: Tensor, *, frames: int, height: int, width: int) -> Tensor: + """Run spatial GLU plus temporal aggregation.""" + batch, tokens, channels = x.shape + 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( + _silu_multiply(value, gate, inplace=not torch.is_grad_enabled()) + ) + + x_time = x_2d.view(batch, frames, channels, height * width).permute(0, 2, 1, 3) + x_time = x_time + self.t_conv(x_time) + return x_time.permute(0, 2, 3, 1).reshape(batch, tokens, channels) + + +class _LinearizedPointwiseConv2d(nn.Module): + """Run an existing 1x1 ``Conv2dContainer`` through a Linear module.""" + + def __init__(self, conv_layer: nn.Module) -> None: + super().__init__() + conv = getattr(conv_layer, "conv", None) + if not isinstance(conv, nn.Conv2d): + raise ValueError("expected Conv2dContainer.conv to be nn.Conv2d.") + if conv.kernel_size != (1, 1) or conv.stride != (1, 1) or conv.padding != (0, 0): + raise ValueError("only exact 1x1 pointwise Conv2d can be linearized.") + if conv.dilation != (1, 1) or conv.groups != 1: + raise ValueError("grouped or dilated pointwise Conv2d cannot be linearized.") + self.linear = nn.Linear( + conv.in_channels, + conv.out_channels, + bias=conv.bias is not None, + device=conv.weight.device, + dtype=conv.weight.dtype, + ) + with torch.no_grad(): + self.linear.weight.copy_(conv.weight.flatten(1)) + if conv.bias is not None: + self.linear.bias.copy_(conv.bias) + + def forward(self, x: Tensor) -> Tensor: + """Apply the pointwise projection while preserving NCHW layout.""" + if x.dim() != 4: + raise ValueError(f"expected NCHW input, got shape {tuple(x.shape)}.") + batch, _channels, height, width = x.shape + x_nhwc = x.permute(0, 2, 3, 1).reshape(batch * height * width, -1) + y = self.linear(x_nhwc) + return y.reshape(batch, height, width, -1).permute(0, 3, 1, 2).contiguous() + + +def linearize_stage1_ffn_for_quant(module: nn.Module) -> tuple[int, int]: + """Expose Stage-1 pointwise FFN convolutions as Linears for quantization.""" + converted = 0 + skipped = 0 + for child in module.modules(): + if not isinstance(child, GLUMBConvTemp): + continue + for attr in ("inverted_conv", "point_conv"): + pointwise = getattr(child, attr) + if isinstance(pointwise, _LinearizedPointwiseConv2d): + continue + try: + replacement = _LinearizedPointwiseConv2d(pointwise) + except ValueError: + skipped += 1 + continue + replacement.train(pointwise.training) + setattr(child, attr, replacement) + converted += 1 + return converted, skipped + + +class Stage1SelfAttention(nn.Module): + """Self/camera attention parameter container for one Stage-1 block.""" + + def __init__(self, spec: SanaWMStage1Spec, *, use_gdn_convs: bool) -> None: + super().__init__() + self.heads = spec.num_heads + self.dim = spec.head_dim + self.eps = 1e-6 + self.use_gdn_convs = use_gdn_convs + self.patch_size = (1, 1, 1) + self.A_log = nn.Parameter(torch.empty(spec.num_heads)) + self.beta_proj = nn.Linear(spec.hidden_size, spec.num_heads) + self.dt_bias = nn.Parameter(torch.empty(spec.num_heads)) + self.gate_proj = nn.Linear(spec.hidden_size, spec.num_heads) + self.k_norm = RMSNorm(spec.hidden_size) + self.k_norm_cam = RMSNorm(spec.hidden_size) + self.k_proj_cam = nn.Linear(spec.hidden_size, spec.hidden_size) + self.out_proj_cam = nn.Linear(spec.hidden_size, spec.hidden_size) + self.output_gate = nn.Linear(spec.hidden_size, spec.hidden_size) + self.proj = nn.Linear(spec.hidden_size, spec.hidden_size) + self.q_norm = RMSNorm(spec.hidden_size) + self.q_norm_cam = RMSNorm(spec.hidden_size) + self.q_proj_cam = nn.Linear(spec.hidden_size, spec.hidden_size) + self.qkv = nn.Linear(spec.hidden_size, 3 * spec.hidden_size, bias=False) + self.recall_gate = nn.Parameter(torch.empty(1)) + self.v_proj_cam = nn.Linear(spec.hidden_size, spec.hidden_size) + if use_gdn_convs: + self.conv_k = nn.Conv1d( + spec.hidden_size, + spec.hidden_size, + kernel_size=spec.conv_kernel_size, + groups=spec.hidden_size, + bias=False, + ) + self.conv_k_cam = nn.Conv1d( + spec.hidden_size, + spec.hidden_size, + kernel_size=spec.conv_kernel_size, + groups=spec.hidden_size, + bias=False, + ) + nn.init.zeros_(self.A_log) + nn.init.constant_(self.dt_bias, -5.0) + nn.init.zeros_(self.recall_gate) + + def forward( + self, + x: Tensor, + *_args: object, + HW: tuple[int, int, int] | None = None, + rotary_emb: Tensor | None = None, + camera_conditions: Tensor | None = None, + camera_cache: _CameraProjectionCache | None = None, + apply_output_gate: bool = True, + **kwargs: object, + ) -> Tensor: + """Run SANA-WM self/camera attention.""" + if HW is None: + raise ValueError("SANA-WM Stage-1 attention requires HW=(T, H, W).") + batch, tokens, channels = x.shape + if channels != self.heads * self.dim: + raise ValueError( + f"channels={channels} != heads*dim={self.heads * self.dim}" + ) + + precomputed_gates = self._compute_frame_gates(x, HW) + main_raw = self._forward_gdn_main( + x, + HW=HW, + rotary_emb=rotary_emb, + precomputed_gates=precomputed_gates, + ) + + cam_contrib: Tensor | int = 0 + if camera_conditions is not None: + if self.use_gdn_convs: + cam_raw = self._forward_gdn_camera( + x, + HW=HW, + rotary_emb=rotary_emb, + camera_conditions=camera_conditions, + camera_cache=camera_cache, + precomputed_gates=precomputed_gates, + ) + else: + cam_raw = self._forward_softmax_camera( + x, + HW=HW, + rotary_emb=rotary_emb, + camera_conditions=camera_conditions, + camera_cache=camera_cache, + chunk_size=kwargs.get("chunk_size"), + chunk_split_strategy=str(kwargs.get("chunk_split_strategy", "uniform")), + chunk_index=kwargs.get("chunk_index"), + ) + cam_contrib = self.out_proj_cam(cam_raw) + + combined = main_raw + cam_contrib + if apply_output_gate: + combined = _apply_output_gate( + combined, + x, + self.output_gate.weight, + self.output_gate.bias, + ) + combined = self.proj(combined.to(dtype=self.proj.weight.dtype)) + del kwargs + return combined + + def _forward_gdn_main( + self, + x: Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + precomputed_gates: tuple[Tensor, Tensor], + ) -> Tensor: + batch, tokens, channels = x.shape + qkv = self.qkv(x).reshape(batch, tokens, 3, self.heads, self.dim) + if hasattr(self, "conv_k"): + k_raw = qkv[:, :, 1].reshape(batch, tokens, channels) + k_conv = _apply_bidirectional_temporal_conv(k_raw, self.conv_k, HW) + qkv[:, :, 1] = k_conv.reshape(batch, tokens, self.heads, self.dim) + + q = _rmsnorm_relu_heads( + qkv[:, :, 0], + self.q_norm.weight, + self.q_norm.eps, + ) + k = _rmsnorm_relu_heads( + qkv[:, :, 1], + self.k_norm.weight, + self.k_norm.eps, + scale=_gdn_key_scale(self.dim, HW), + ) + v = qkv[:, :, 2] + v = v.float() + q_rot = _apply_complex_rope(q, rotary_emb) + k_rot = _apply_complex_rope(k, rotary_emb) + beta, decay = precomputed_gates + out = _bidirectional_gdn_scan( + q=q, + k=k, + q_rot=q_rot, + k_rot=k_rot, + v=v, + beta=beta, + decay=decay, + HW=HW, + eps=self.eps, + ) + return out.reshape(batch, tokens, channels).to(dtype=x.dtype) + + def _forward_softmax_main( + self, + x: Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + chunk_size: int | None, + chunk_split_strategy: str, + chunk_index: list[int] | None, + ) -> Tensor: + batch, tokens, channels = x.shape + qkv = self.qkv(x).reshape(batch, tokens, 3, self.heads, self.dim) + q, k, v = qkv.unbind(dim=2) + q = self.q_norm(q.reshape(batch, tokens, channels)).reshape( + batch, + tokens, + self.heads, + self.dim, + ) + k = self.k_norm(k.reshape(batch, tokens, channels)).reshape( + batch, + tokens, + self.heads, + self.dim, + ) + q = _apply_complex_rope(q, rotary_emb) + k = _apply_complex_rope(k, rotary_emb) + out = _scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + HW=HW, + chunk_size=chunk_size, + chunk_split_strategy=chunk_split_strategy, + chunk_index=chunk_index, + ) + return out.transpose(1, 2).reshape(batch, tokens, channels).to(dtype=x.dtype) + + def _forward_gdn_camera( + self, + x: Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + camera_conditions: Tensor, + camera_cache: _CameraProjectionCache | None, + precomputed_gates: tuple[Tensor, Tensor], + ) -> Tensor: + q_cam, k_cam, v_cam = _camera_qkv(self, x, HW) + q_trans, k_trans, v_trans, inflation_sq, output_projector = _prepare_ucpe_qkv( + q_cam, + k_cam, + v_cam, + camera_conditions=camera_conditions, + HW=HW, + rotary_emb=rotary_emb, + camera_cache=camera_cache, + q_norm_weight=self.q_norm_cam.weight, + k_norm_weight=self.k_norm_cam.weight, + norm_eps=self.q_norm_cam.eps, + ) + beta, decay = precomputed_gates + frame_inflation = inflation_sq.reshape( + x.shape[0], + self.heads, + HW[0], + HW[1] * HW[2], + ).mean(dim=-1) + beta = beta / frame_inflation.unsqueeze(-1).clamp_min(1.0) + out = _bidirectional_numerator_scan( + q=q_trans, + k=k_trans, + v=v_trans, + beta=beta, + decay=decay, + HW=HW, + ) + out = output_projector(out) + return out.reshape(x.shape[0], x.shape[1], -1).to(dtype=x.dtype) + + def _forward_softmax_camera( + self, + x: Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + camera_conditions: Tensor, + camera_cache: _CameraProjectionCache | None, + chunk_size: int | None, + chunk_split_strategy: str, + chunk_index: list[int] | None, + ) -> Tensor: + q_cam, k_cam, v_cam = _camera_qkv(self, x, HW) + q_trans, k_trans, v_trans, output_projector = _prepare_ucpe_qkv_softmax( + q_cam, + k_cam, + v_cam, + camera_conditions=camera_conditions, + HW=HW, + rotary_emb=rotary_emb, + camera_cache=camera_cache, + q_norm_weight=self.q_norm_cam.weight, + k_norm_weight=self.k_norm_cam.weight, + norm_eps=self.q_norm_cam.eps, + ) + out = _scaled_dot_product_attention( + q_trans.transpose(1, 2), + k_trans.transpose(1, 2), + v_trans.transpose(1, 2), + HW=HW, + chunk_size=chunk_size, + chunk_split_strategy=chunk_split_strategy, + chunk_index=chunk_index, + ) + return output_projector(out.transpose(1, 2)).reshape( + x.shape[0], + x.shape[1], + -1, + ).to(dtype=x.dtype) + + def _compute_frame_gates( + self, + x: Tensor, + HW: tuple[int, int, int], + ) -> tuple[Tensor, Tensor]: + batch, tokens, channels = x.shape + frames, height, width = HW + spatial = height * width + if tokens != frames * spatial: + raise ValueError(f"tokens={tokens} != T*H*W={frames * spatial}") + beta, decay = _compute_frame_gates( + x, + frames, + spatial, + self.heads, + self.beta_proj.weight, + self.beta_proj.bias, + self.gate_proj.weight, + self.gate_proj.bias, + self.dt_bias, + self.A_log, + ) + return beta.float(), decay.float() + + +class Stage1CrossAttention(nn.Module): + """Cross-attention parameter container for one Stage-1 block.""" + + def __init__(self, spec: SanaWMStage1Spec) -> None: + super().__init__() + self.num_heads = spec.num_heads + self.head_dim = spec.hidden_size // spec.num_heads + self.k_norm = RMSNorm(spec.hidden_size) + self.kv_linear = nn.Linear(spec.hidden_size, 2 * spec.hidden_size) + self.proj = nn.Linear(spec.hidden_size, spec.hidden_size) + self.q_linear = nn.Linear(spec.hidden_size, spec.hidden_size) + self.q_norm = RMSNorm(spec.hidden_size) + + def forward( + self, + x: Tensor, + y: Tensor, + *, + mask: Tensor | None = None, + **_kwargs: object, + ) -> Tensor: + """Run text cross-attention with checkpoint-compatible Q/K norms.""" + batch, tokens, channels = x.shape + q = self.q_norm(self.q_linear(x)).view( + batch, + tokens, + self.num_heads, + self.head_dim, + ) + kv = self.kv_linear(y).view(batch, -1, 2, channels) + k, v = kv.unbind(dim=2) + k = self.k_norm(k).view(batch, -1, self.num_heads, self.head_dim) + v = v.view(batch, -1, self.num_heads, self.head_dim) + + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + attn_mask = None + if mask is not None: + attn_mask = (1 - mask.to(dtype=q.dtype)) * -10000.0 + if attn_mask.ndim == 2: + attn_mask = attn_mask[:, None, None].repeat(1, self.num_heads, 1, 1) + out = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + dropout_p=0.0, + is_causal=False, + ) + out = out.transpose(1, 2).reshape(batch, tokens, channels) + return self.proj(out.to(dtype=self.proj.weight.dtype)) + + +class SanaWMStage1Block(nn.Module): + """One checkpoint-compatible SANA-WM Stage-1 transformer block.""" + + def __init__(self, spec: SanaWMStage1Spec, *, index: int) -> None: + super().__init__() + self.norm1 = nn.LayerNorm( + spec.hidden_size, + elementwise_affine=False, + eps=1e-6, + ) + self.norm2 = nn.LayerNorm( + spec.hidden_size, + elementwise_affine=False, + eps=1e-6, + ) + self.scale_shift_table = nn.Parameter(torch.empty(6, spec.hidden_size)) + self.attn = Stage1SelfAttention( + spec, + use_gdn_convs=spec.block_uses_gdn(index), + ) + self.cross_attn = Stage1CrossAttention(spec) + self.mlp = GLUMBConvTemp(spec) + self.plucker_proj = nn.Linear(spec.hidden_size, spec.hidden_size) + + def forward( + self, + x: Tensor, + y: Tensor, + t: Tensor, + *, + frames: int, + height: int, + width: int, + mask: Tensor | None = None, + plucker_emb: Tensor | None = None, + **kwargs: object, + ) -> Tensor: + """Run one Stage-1 transformer block.""" + batch, tokens, channels = x.shape + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None, None, :, :] + t.reshape(batch, frames, 6, -1) + ).chunk(6, dim=-2) + + x_norm = self.norm1(x).reshape(batch, frames, -1, channels) + attn_in = _modulate(x_norm, shift_msa, scale_msa).reshape(batch, tokens, channels) + attn_out = self.attn(attn_in, **kwargs).reshape(batch, frames, -1, channels) + x = x + (gate_msa * attn_out).reshape(batch, tokens, channels) + + if plucker_emb is not None: + x = x + self.plucker_proj(plucker_emb) + + x = x + self.cross_attn(x, y, mask=mask) + + x_norm = self.norm2(x).reshape(batch, frames, -1, channels) + mlp_in = _modulate(x_norm, shift_mlp, scale_mlp).reshape(batch, tokens, channels) + mlp_out = self.mlp(mlp_in, frames=frames, height=height, width=width) + mlp_out = mlp_out.reshape(batch, frames, -1, channels) + return x + (gate_mlp * mlp_out).reshape(batch, tokens, channels) + + +class SanaWMStage1Model(nn.Module): + """Checkpoint-compatible Stage-1 SANA-WM DiT shell.""" + + def __init__(self, spec: SanaWMStage1Spec = SANA_WM_STAGE1_SPEC) -> None: + super().__init__() + self.spec = spec + self.x_embedder = Conv3dProjector(spec.latent_channels, spec.hidden_size) + self.t_embedder = TimestepEmbedder(spec) + self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(spec.hidden_size, 6 * spec.hidden_size)) + self.y_embedder = TextEmbedder(spec) + self.attention_y_norm = RMSNorm(spec.hidden_size) + self.raymap_embedder = Conv3dProjector(spec.raymap_channels, spec.hidden_size) + self.plucker_embedder = Conv3dProjector(spec.plucker_channels, spec.hidden_size) + self.blocks = nn.ModuleList( + SanaWMStage1Block(spec, index=index) for index in range(spec.depth) + ) + num_pos_tokens = spec.latent_grid_size[0] * spec.latent_grid_size[1] // 2 + 44 + self.pos_embed = nn.Parameter(torch.empty(1, num_pos_tokens, spec.hidden_size)) + self.final_layer = FinalLayer(spec) + + def forward( + self, + x: Tensor, + timestep: Tensor, + y: Tensor, + *, + mask: Tensor | None = None, + chunk_plucker: Tensor | None = None, + **kwargs: object, + ) -> Tensor: + """Run the SANA-WM Stage-1 DiT.""" + batch, _channels, frames, height, width = x.shape + x = self.x_embedder(x) + x = x.permute(0, 2, 3, 4, 1).reshape(batch, frames * height * width, -1) + + plucker_emb = kwargs.get("chunk_plucker_emb") + if isinstance(plucker_emb, Tensor): + if ( + plucker_emb.ndim != 3 + or plucker_emb.shape[1:] != x.shape[1:] + or plucker_emb.shape[0] not in (1, batch) + ): + raise ValueError( + "chunk_plucker_emb must have shape [1|B, T*H*W, hidden], " + f"got {tuple(plucker_emb.shape)} for x={tuple(x.shape)}." + ) + plucker_emb = plucker_emb.to(device=x.device, dtype=x.dtype) + elif chunk_plucker is not None: + plucker_emb = self.prepare_plucker_embedding(chunk_plucker) + + rotary_emb = kwargs.get("rotary_emb") + if isinstance(rotary_emb, Tensor): + rotary_emb = rotary_emb.to(device=x.device) + else: + rotary_emb = _wan_rope_complex( + self.spec.head_dim, + frames, + height, + width, + x.device, + ) + raw_camera_conditions = kwargs.get("camera_conditions") + camera_conditions = ( + raw_camera_conditions.to(device=x.device, dtype=x.dtype) + if isinstance(raw_camera_conditions, Tensor) + else None + ) + camera_cache_value = kwargs.get("camera_cache") + camera_cache = ( + _camera_projection_cache_to( + camera_cache_value, + device=x.device, + dtype=x.dtype, + ) + if isinstance(camera_cache_value, _CameraProjectionCache) + else ( + _prepare_camera_projection_cache( + camera_conditions, + HW=(frames, height, width), + rotary_emb=rotary_emb, + head_dim=self.spec.head_dim, + ) + if camera_conditions is not None + else None + ) + ) + camera_signal = ( + camera_cache.raymats if camera_cache is not None else camera_conditions + ) + block_kwargs = { + key: value + for key, value in kwargs.items() + if key + not in { + "camera_conditions", + "chunk_plucker_emb", + "camera_cache", + "rotary_emb", + } + } + + y = self.y_embedder(y) + y = self.attention_y_norm(y) + if mask is not None and mask.ndim > 2: + mask = mask.squeeze(1).squeeze(1) + + timestep_embed = self.t_embedder(timestep.flatten()) + timestep_embed = timestep_embed.unflatten(0, timestep.shape) + block_t = self.t_block(timestep_embed).reshape( + batch, + 1, + frames, + 6 * self.spec.hidden_size, + ) + + for block in self.blocks: + x = block( + x, + y, + block_t, + frames=frames, + height=height, + width=width, + mask=mask, + plucker_emb=plucker_emb, + HW=(frames, height, width), + rotary_emb=rotary_emb, + camera_conditions=camera_signal, + camera_cache=camera_cache, + chunk_size=self.spec.chunk_size, + chunk_split_strategy=self.spec.chunk_split_strategy, + **block_kwargs, + ) + + x = self.final_layer(x, timestep_embed, frames=frames) + return x.reshape(batch, frames, height, width, -1).permute(0, 4, 1, 2, 3) + + def prepare_plucker_embedding(self, chunk_plucker: Tensor) -> Tensor: + """Project static Plucker conditioning into Stage-1 token space.""" + batch = chunk_plucker.shape[0] + plucker_emb = self.plucker_embedder(chunk_plucker) + return plucker_emb.permute(0, 2, 3, 4, 1).reshape( + batch, + -1, + self.spec.hidden_size, + ) + + def prepare_camera_projection_cache( + self, + camera_conditions: Tensor, + *, + frames: int, + height: int, + width: int, + ) -> tuple[Tensor, _CameraProjectionCache]: + """Precompute static RoPE and camera projection tensors for a rollout.""" + param = next(self.parameters()) + camera_conditions = camera_conditions.to(device=param.device, dtype=param.dtype) + rotary_emb = _wan_rope_complex( + self.spec.head_dim, + frames, + height, + width, + param.device, + ) + camera_cache = _prepare_camera_projection_cache( + camera_conditions, + HW=(frames, height, width), + rotary_emb=rotary_emb, + head_dim=self.spec.head_dim, + ) + return rotary_emb, camera_cache + + +def _modulate(x: Tensor, shift: Tensor, scale: Tensor) -> Tensor: + return x * (1 + scale) + shift + + +def _timestep_embedding(t: Tensor, dim: int, max_period: int = 10000) -> Tensor: + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) + / half + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + +def _gdn_key_scale(head_dim: int, HW: tuple[int, int, int]) -> float: + return (head_dim**-0.5) * ((HW[1] * HW[2]) ** -0.5) + + +def _sdpa_needs_head_pad(head_dim: int) -> bool: + return head_dim not in (32, 64, 128, 256) and head_dim < 256 + + +def _scaled_dot_product_attention( + xq: Tensor, + xk: Tensor, + xv: Tensor, + *, + HW: tuple[int, int, int] | None = None, + chunk_size: int | None = None, + chunk_split_strategy: str = "uniform", + chunk_index: list[int] | None = None, +) -> Tensor: + if HW is not None and chunk_size is not None and chunk_size < HW[0]: + frames, height, width = HW + spatial = height * width + boundaries = _normalize_chunk_index( + chunk_index, + frames, + chunk_size, + chunk_split_strategy, + ) + out_chunks = [] + for start, end in zip(boundaries[:-1], boundaries[1:]): + q_chunk = xq[:, :, start * spatial : end * spatial] + out_chunks.append( + _scaled_dot_product_attention_full( + q_chunk, + xk[:, :, : end * spatial], + xv[:, :, : end * spatial], + ), + ) + return torch.cat(out_chunks, dim=2) + return _scaled_dot_product_attention_full(xq, xk, xv) + + +def _scaled_dot_product_attention_full(xq: Tensor, xk: Tensor, xv: Tensor) -> Tensor: + head_dim = xq.shape[-1] + if not _sdpa_needs_head_pad(head_dim): + return F.scaled_dot_product_attention( + xq, + xk, + xv, + dropout_p=0.0, + is_causal=False, + ) + pad_to = 128 if head_dim <= 128 else 256 + pad_size = pad_to - head_dim + xq_padded = F.pad(xq, (0, pad_size)) + xk_padded = F.pad(xk, (0, pad_size)) + xv_padded = F.pad(xv, (0, pad_size)) + out = F.scaled_dot_product_attention( + xq_padded, + xk_padded, + xv_padded, + dropout_p=0.0, + is_causal=False, + scale=head_dim**-0.5, + ) + return out[..., :head_dim] + + +def _normalize_chunk_index( + chunk_index: list[int] | None, + frames: int, + chunk_size: int, + chunk_split_strategy: str, +) -> list[int]: + if chunk_index is None: + chunk_index = _chunk_index_from_chunk_size( + frames, + chunk_size, + chunk_split_strategy, + ) + else: + chunk_index = list(chunk_index) + if not chunk_index or chunk_index[0] != 0: + chunk_index = [0] + [idx for idx in chunk_index if idx > 0] + chunk_index = [idx for idx in chunk_index if idx < frames] + if not chunk_index: + chunk_index = [0] + if chunk_index[-1] != frames: + chunk_index.append(frames) + return chunk_index + + +def _chunk_index_from_chunk_size( + frames: int, + chunk_size: int, + chunk_split_strategy: str, +) -> list[int]: + if chunk_size <= 0: + raise ValueError(f"chunk_size must be > 0, got {chunk_size}.") + if frames <= 0: + raise ValueError(f"frames must be > 0, got {frames}.") + strategy = "uniform" if chunk_split_strategy is None else str(chunk_split_strategy) + strategy = strategy.lower() + if strategy in ("uniform", "default"): + indices = list(range(0, frames, chunk_size)) + if len(indices) > 1 and (frames - indices[-1]) < chunk_size: + indices.pop() + return indices + if strategy in ("first_frame", "first_frame_alone", "first_frame_only"): + if frames <= 1: + return [0] + indices = [0] + list(range(1, frames, chunk_size)) + if len(indices) > 2 and (frames - indices[-1]) < chunk_size: + indices.pop() + return indices + if strategy in ("first_plus_one", "first_chunk_plus_one"): + if frames <= chunk_size + 1: + return [0] + indices = [0] + list(range(chunk_size + 1, frames, chunk_size)) + if len(indices) > 1 and (frames - indices[-1]) < chunk_size: + indices.pop() + return indices + raise ValueError( + "Unknown chunk_split_strategy " + f"{chunk_split_strategy!r}; expected uniform, first_frame, or first_plus_one." + ) + + +def _prepare_camera_projection_cache( + camera_conditions: Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + head_dim: int, +) -> _CameraProjectionCache: + batch, frames = camera_conditions.shape[:2] + latent_frames, height, width = HW + if frames != latent_frames: + raise ValueError( + f"camera_conditions has {frames} frames but latent grid has {latent_frames}." + ) + tokens = frames * height * width + raymats = _camera_ray_mats(camera_conditions.float(), HW).to( + dtype=camera_conditions.dtype, + ) + proj = raymats.reshape(batch, tokens, 4, 4) + proj_q = proj.transpose(-1, -2).contiguous() + proj_kv = _invert_se3(proj).contiguous() + rope_cam = _slice_rope_for_camera(rotary_emb, head_dim) + return _CameraProjectionCache( + raymats=raymats, + proj=proj, + proj_q=proj_q, + proj_kv=proj_kv, + rope_cam=rope_cam, + ) + + +def _camera_projection_cache_to( + cache: _CameraProjectionCache, + *, + device: torch.device, + dtype: torch.dtype, +) -> _CameraProjectionCache: + if ( + cache.raymats.device == device + and cache.proj.device == device + and cache.proj_q.device == device + and cache.proj_kv.device == device + and cache.raymats.dtype == dtype + and cache.proj.dtype == dtype + and cache.proj_q.dtype == dtype + and cache.proj_kv.dtype == dtype + and (cache.rope_cam is None or cache.rope_cam.device == device) + ): + return cache + return _CameraProjectionCache( + raymats=cache.raymats.to(device=device, dtype=dtype), + proj=cache.proj.to(device=device, dtype=dtype), + proj_q=cache.proj_q.to(device=device, dtype=dtype), + proj_kv=cache.proj_kv.to(device=device, dtype=dtype), + rope_cam=( + cache.rope_cam.to(device=device) if cache.rope_cam is not None else None + ), + ) + + +@torch.compile +def _apply_output_gate( + out: Tensor, + gate_x: Tensor, + gate_weight: Tensor, + gate_bias: Tensor, +) -> Tensor: + gate_values = F.silu(F.linear(gate_x, gate_weight, gate_bias).float()) + return out * gate_values.to(dtype=out.dtype) + + +@torch.compile +def _compute_frame_gates( + x: Tensor, + frames: int, + spatial: int, + heads: int, + beta_weight: Tensor, + beta_bias: Tensor, + gate_weight: Tensor, + gate_bias: Tensor, + dt_bias: Tensor, + A_log: Tensor, +) -> tuple[Tensor, Tensor]: + batch, tokens, channels = x.shape + beta = F.linear(x, beta_weight, beta_bias).sigmoid() + beta = beta.reshape(batch, frames, spatial, heads).permute(0, 3, 1, 2) + x_frame = x.reshape(batch, frames, spatial, channels).mean(dim=2) + gate = F.linear(x_frame, gate_weight, gate_bias).float() + dt = dt_bias.float().view(1, 1, -1) + a_val = A_log.float().exp().view(1, 1, -1) + decay = (-a_val * F.softplus(gate + dt)).exp().transpose(1, 2) + return beta, decay + + +if triton is not None and tl is not None: + + @triton.jit + def _bidirectional_temporal_conv_kernel( + x_ptr, + weight_ptr, + out_ptr, + x_stride_b, + x_stride_t, + x_stride_c, + weight_stride_c, + weight_stride_k, + out_stride_b, + out_stride_t, + out_stride_c, + frames: tl.constexpr, + spatial: tl.constexpr, + channels, + kernel_size: tl.constexpr, + block_f: tl.constexpr, + block_s: tl.constexpr, + block_c: tl.constexpr, + ) -> None: + batch_idx = tl.program_id(0) + spatial_offsets = tl.program_id(1) * block_s + tl.arange(0, block_s) + channel_offsets = tl.program_id(2) * block_c + tl.arange(0, block_c) + frame_offsets = tl.arange(0, block_f) + + frame_mask = frame_offsets < frames + spatial_mask = spatial_offsets < spatial + channel_mask = channel_offsets < channels + element_mask = ( + frame_mask[:, None, None] + & spatial_mask[None, :, None] + & channel_mask[None, None, :] + ) + + token_offsets = frame_offsets[:, None, None] * spatial + spatial_offsets[ + None, + :, + None, + ] + x_offsets = ( + batch_idx * x_stride_b + + token_offsets * x_stride_t + + channel_offsets[None, None, :] * x_stride_c + ) + out_offsets = ( + batch_idx * out_stride_b + + token_offsets * out_stride_t + + channel_offsets[None, None, :] * out_stride_c + ) + + center = tl.load( + weight_ptr + + channel_offsets * weight_stride_c + + (kernel_size - 1) * weight_stride_k, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + value = tl.load(x_ptr + x_offsets, mask=element_mask, other=0.0).to( + tl.float32 + ) + acc = value * center[None, None, :] + + for offset in range(1, kernel_size): + coeff = tl.load( + weight_ptr + + channel_offsets * weight_stride_c + + (kernel_size - 1 - offset) * weight_stride_k, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + + prev_frame_offsets = frame_offsets - offset + prev_token_offsets = prev_frame_offsets[:, None, None] * spatial + ( + spatial_offsets[None, :, None] + ) + prev_mask = element_mask & (frame_offsets[:, None, None] >= offset) + prev_offsets = ( + batch_idx * x_stride_b + + prev_token_offsets * x_stride_t + + channel_offsets[None, None, :] * x_stride_c + ) + prev = tl.load(x_ptr + prev_offsets, mask=prev_mask, other=0.0).to( + tl.float32 + ) + + next_frame_offsets = frame_offsets + offset + next_token_offsets = next_frame_offsets[:, None, None] * spatial + ( + spatial_offsets[None, :, None] + ) + next_mask = element_mask & (next_frame_offsets[:, None, None] < frames) + next_offsets = ( + batch_idx * x_stride_b + + next_token_offsets * x_stride_t + + channel_offsets[None, None, :] * x_stride_c + ) + next_value = tl.load(x_ptr + next_offsets, mask=next_mask, other=0.0).to( + tl.float32 + ) + acc += (prev + next_value) * coeff[None, None, :] + + tl.store(out_ptr + out_offsets, acc, mask=element_mask) + + @triton.jit + def _ucpe_first_half_kernel( + x_ptr, + matrix_ptr, + out_ptr, + x_stride_b, + x_stride_n, + x_stride_h, + x_stride_d, + matrix_stride_b, + matrix_stride_n, + matrix_stride_i, + matrix_stride_j, + out_stride_b, + out_stride_n, + out_stride_h, + out_stride_d, + tokens: tl.constexpr, + groups4: tl.constexpr, + matrix_batch: tl.constexpr, + block_n: tl.constexpr, + ) -> None: + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + packed_idx = tl.program_id(2) + group_idx = packed_idx % groups4 + token_block = packed_idx // groups4 + token_offsets = token_block * block_n + tl.arange(0, block_n) + token_mask = token_offsets < tokens + + matrix_batch_idx = 0 if matrix_batch == 1 else batch_idx + col0 = group_idx * 4 + col1 = col0 + 1 + col2 = col0 + 2 + col3 = col0 + 3 + + x_base = ( + batch_idx * x_stride_b + + token_offsets * x_stride_n + + head_idx * x_stride_h + ) + x0 = tl.load( + x_ptr + x_base + col0 * x_stride_d, + mask=token_mask, + other=0.0, + ).to( + tl.float32 + ) + x1 = tl.load( + x_ptr + x_base + col1 * x_stride_d, + mask=token_mask, + other=0.0, + ).to( + tl.float32 + ) + x2 = tl.load( + x_ptr + x_base + col2 * x_stride_d, + mask=token_mask, + other=0.0, + ).to( + tl.float32 + ) + x3 = tl.load( + x_ptr + x_base + col3 * x_stride_d, + mask=token_mask, + other=0.0, + ).to( + tl.float32 + ) + + matrix_base = ( + matrix_batch_idx * matrix_stride_b + token_offsets * matrix_stride_n + ) + out_base = ( + batch_idx * out_stride_b + + token_offsets * out_stride_n + + head_idx * out_stride_h + ) + for row in range(4): + row_base = matrix_base + row * matrix_stride_i + m0 = tl.load( + matrix_ptr + row_base + 0 * matrix_stride_j, + mask=token_mask, + other=0.0, + ).to(tl.float32) + m1 = tl.load( + matrix_ptr + row_base + 1 * matrix_stride_j, + mask=token_mask, + other=0.0, + ).to(tl.float32) + m2 = tl.load( + matrix_ptr + row_base + 2 * matrix_stride_j, + mask=token_mask, + other=0.0, + ).to(tl.float32) + m3 = tl.load( + matrix_ptr + row_base + 3 * matrix_stride_j, + mask=token_mask, + other=0.0, + ).to(tl.float32) + value = (x0 * m0 + x1 * m1) + (x2 * m2 + x3 * m3) + tl.store( + out_ptr + out_base + (col0 + row) * out_stride_d, + value, + mask=token_mask, + ) + + @triton.jit + def _ucpe_rope_second_half_kernel( + x_ptr, + rope_ptr, + out_ptr, + x_stride_b, + x_stride_n, + x_stride_h, + x_stride_d, + rope_stride_n, + rope_stride_p, + rope_stride_ri, + out_stride_b, + out_stride_n, + out_stride_h, + out_stride_d, + tokens: tl.constexpr, + half: tl.constexpr, + rope_pairs: tl.constexpr, + inverse_rope: tl.constexpr, + block_n: tl.constexpr, + ) -> None: + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + packed_idx = tl.program_id(2) + pair_idx = packed_idx % rope_pairs + token_block = packed_idx // rope_pairs + token_offsets = token_block * block_n + tl.arange(0, block_n) + token_mask = token_offsets < tokens + + real_d = half + pair_idx * 2 + imag_d = real_d + 1 + x_base = ( + batch_idx * x_stride_b + + token_offsets * x_stride_n + + head_idx * x_stride_h + ) + real = tl.load( + x_ptr + x_base + real_d * x_stride_d, + mask=token_mask, + other=0.0, + ).to( + tl.float32 + ) + imag = tl.load( + x_ptr + x_base + imag_d * x_stride_d, + mask=token_mask, + other=0.0, + ).to( + tl.float32 + ) + rope_base = token_offsets * rope_stride_n + pair_idx * rope_stride_p + rope_real = tl.load( + rope_ptr + rope_base + 0 * rope_stride_ri, + mask=token_mask, + other=0.0, + ).to(tl.float32) + rope_imag = tl.load( + rope_ptr + rope_base + 1 * rope_stride_ri, + mask=token_mask, + other=0.0, + ).to(tl.float32) + if inverse_rope: + rope_imag = -rope_imag + + out_base = ( + batch_idx * out_stride_b + + token_offsets * out_stride_n + + head_idx * out_stride_h + ) + tl.store( + out_ptr + out_base + real_d * out_stride_d, + real * rope_real - imag * rope_imag, + mask=token_mask, + ) + tl.store( + out_ptr + out_base + imag_d * out_stride_d, + real * rope_imag + imag * rope_real, + mask=token_mask, + ) + + @triton.jit + def _ucpe_norms_kernel( + x_ptr, + out_ptr, + pre_ptr, + post_ptr, + x_stride_b, + x_stride_n, + x_stride_h, + x_stride_d, + out_stride_b, + out_stride_n, + out_stride_h, + out_stride_d, + pre_stride_b, + pre_stride_h, + pre_stride_n, + post_stride_b, + post_stride_h, + post_stride_n, + tokens: tl.constexpr, + dim: tl.constexpr, + block_n: tl.constexpr, + block_d: tl.constexpr, + ) -> None: + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + token_block = tl.program_id(2) + token_offsets = token_block * block_n + tl.arange(0, block_n) + dim_offsets = tl.arange(0, block_d) + mask = (token_offsets[:, None] < tokens) & (dim_offsets[None, :] < dim) + + x_offsets = ( + batch_idx * x_stride_b + + token_offsets[:, None] * x_stride_n + + head_idx * x_stride_h + + dim_offsets[None, :] * x_stride_d + ) + out_offsets = ( + batch_idx * out_stride_b + + token_offsets[:, None] * out_stride_n + + head_idx * out_stride_h + + dim_offsets[None, :] * out_stride_d + ) + x_values = tl.load(x_ptr + x_offsets, mask=mask, other=0.0).to(tl.float32) + out_values = tl.load(out_ptr + out_offsets, mask=mask, other=0.0).to( + tl.float32 + ) + pre = tl.sum(x_values * x_values, axis=1) + post = tl.sum(out_values * out_values, axis=1) + token_mask = token_offsets < tokens + pre_offsets = ( + batch_idx * pre_stride_b + + head_idx * pre_stride_h + + token_offsets * pre_stride_n + ) + post_offsets = ( + batch_idx * post_stride_b + + head_idx * post_stride_h + + token_offsets * post_stride_n + ) + tl.store(pre_ptr + pre_offsets, pre, mask=token_mask) + tl.store(post_ptr + post_offsets, post, mask=token_mask) + + @triton.jit + def _silu_multiply_kernel( + value_ptr, + gate_ptr, + out_ptr, + value_stride_n, + value_stride_c, + value_stride_h, + value_stride_w, + gate_stride_n, + gate_stride_c, + gate_stride_h, + gate_stride_w, + channels: tl.constexpr, + height: tl.constexpr, + width: tl.constexpr, + total: tl.constexpr, + block: tl.constexpr, + ) -> None: + offsets = tl.program_id(0) * block + tl.arange(0, block) + mask = offsets < total + channel_offsets = offsets % channels + spatial_offsets = offsets // channels + width_offsets = spatial_offsets % width + height_offsets = (spatial_offsets // width) % height + batch_offsets = spatial_offsets // (height * width) + + value_offsets = ( + batch_offsets * value_stride_n + + channel_offsets * value_stride_c + + height_offsets * value_stride_h + + width_offsets * value_stride_w + ) + gate_offsets = ( + batch_offsets * gate_stride_n + + channel_offsets * gate_stride_c + + height_offsets * gate_stride_h + + width_offsets * gate_stride_w + ) + value = tl.load(value_ptr + value_offsets, mask=mask, other=0.0).to(tl.float32) + gate = tl.load(gate_ptr + gate_offsets, mask=mask, other=0.0).to(tl.float32) + silu_gate = gate / (1.0 + tl.exp(-gate)) + tl.store(out_ptr + offsets, value * silu_gate, mask=mask) + + @triton.jit + def _rmsnorm_relu_heads_kernel( + x_ptr, + weight_ptr, + out_ptr, + x_stride_b, + x_stride_n, + x_stride_h, + x_stride_d, + out_stride_b, + out_stride_n, + out_stride_h, + out_stride_d, + tokens: tl.constexpr, + heads: tl.constexpr, + dim: tl.constexpr, + channels: tl.constexpr, + eps: tl.constexpr, + scale: tl.constexpr, + block_c: tl.constexpr, + ) -> None: + row = tl.program_id(0) + batch_idx = row // tokens + token_idx = row - batch_idx * tokens + channel_offsets = tl.arange(0, block_c) + mask = channel_offsets < channels + head_offsets = channel_offsets // dim + dim_offsets = channel_offsets - head_offsets * dim + + x_offsets = ( + batch_idx * x_stride_b + + token_idx * x_stride_n + + head_offsets * x_stride_h + + dim_offsets * x_stride_d + ) + values = tl.load(x_ptr + x_offsets, mask=mask, other=0.0).to(tl.float32) + mean_sq = tl.sum(values * values, axis=0) / channels + inv_rms = tl.rsqrt(mean_sq + eps) + weights = tl.load(weight_ptr + channel_offsets, mask=mask, other=0.0).to( + tl.float32 + ) + out_values = tl.maximum(values * inv_rms * weights * scale, 0.0) + out_offsets = ( + batch_idx * out_stride_b + + token_idx * out_stride_n + + head_offsets * out_stride_h + + dim_offsets * out_stride_d + ) + tl.store(out_ptr + out_offsets, out_values, mask=mask) + +else: + _bidirectional_temporal_conv_kernel = None + _ucpe_first_half_kernel = None + _ucpe_rope_second_half_kernel = None + _ucpe_norms_kernel = None + _silu_multiply_kernel = None + _rmsnorm_relu_heads_kernel = None + + +def _apply_bidirectional_temporal_conv_fast( + x: Tensor, + conv: nn.Conv1d, + HW: tuple[int, int, int], +) -> Tensor: + batch, tokens, channels = x.shape + frames, height, width = HW + spatial = height * width + out = torch.empty_like(x) + assert triton is not None + assert _bidirectional_temporal_conv_kernel is not None + block_f = triton.next_power_of_2(frames) + block_s = 4 + block_c = 128 + grid = (batch, triton.cdiv(spatial, block_s), triton.cdiv(channels, block_c)) + weight = conv.weight[:, 0, :] + _bidirectional_temporal_conv_kernel[grid]( + x, + weight, + out, + x.stride(0), + x.stride(1), + x.stride(2), + weight.stride(0), + weight.stride(1), + out.stride(0), + out.stride(1), + out.stride(2), + frames, + spatial, + channels, + int(conv.kernel_size[0]), + block_f, + block_s, + block_c, + ) + return out + + +def _apply_bidirectional_temporal_conv_eager( + x: Tensor, + conv: nn.Conv1d, + HW: tuple[int, int, int], +) -> Tensor: + batch, tokens, channels = x.shape + frames, height, width = HW + spatial = height * width + if tokens != frames * spatial: + raise ValueError(f"tokens={tokens} != T*H*W={frames * spatial}") + temporal = ( + x.reshape(batch, frames, spatial, channels) + .permute(0, 2, 3, 1) + .reshape(batch * spatial, channels, frames) + ) + y_fwd = _causal_depthwise_conv(temporal, conv) + y_bwd = _causal_depthwise_conv(temporal.flip(-1), conv).flip(-1) + center = conv.weight[:, 0, -1].view(1, channels, 1) + out = y_fwd + y_bwd - temporal * center + return ( + out.reshape(batch, spatial, channels, frames) + .permute(0, 3, 1, 2) + .reshape(batch, tokens, channels) + .to(dtype=x.dtype) + ) + + +def _apply_bidirectional_temporal_conv( + x: Tensor, + conv: nn.Conv1d, + HW: tuple[int, int, int], +) -> Tensor: + if ( + x.is_cuda + and conv.weight.is_cuda + and conv.weight.dtype == x.dtype + and conv.bias is None + and conv.in_channels == x.shape[-1] + and conv.out_channels == x.shape[-1] + and conv.groups == x.shape[-1] + and conv.stride == (1,) + and conv.padding == (0,) + and conv.dilation == (1,) + and not torch.is_grad_enabled() + and triton is not None + and _bidirectional_temporal_conv_kernel is not None + ): + batch, tokens, _channels = x.shape + frames, height, width = HW + if tokens == frames * height * width: + return _apply_bidirectional_temporal_conv_fast(x, conv, HW) + return _apply_bidirectional_temporal_conv_eager(x, conv, HW) + + +def _causal_depthwise_conv(x: Tensor, conv: nn.Conv1d) -> Tensor: + kernel = int(conv.kernel_size[0]) + padded = F.pad(x, (kernel - 1, 0)) + return F.conv1d( + padded, + conv.weight.to(dtype=x.dtype), + bias=None, + stride=1, + padding=0, + dilation=1, + groups=conv.groups, + ) + + +def _wan_rope_complex( + head_dim: int, + frames: int, + height: int, + width: int, + device: torch.device, +) -> Tensor: + t_size = head_dim // 2 - 2 * (head_dim // 6) + h_size = head_dim // 6 + w_size = head_dim // 6 + freqs_t = _axis_rope_complex(frames, t_size, device) + freqs_h = _axis_rope_complex(height, h_size, device) + freqs_w = _axis_rope_complex(width, w_size, device) + expanded_t = freqs_t[:, None, None, :].expand(frames, height, width, t_size) + expanded_h = freqs_h[None, :, None, :].expand(frames, height, width, h_size) + expanded_w = freqs_w[None, None, :, :].expand(frames, height, width, w_size) + freqs = torch.cat([expanded_t, expanded_h, expanded_w], dim=-1) + return freqs.reshape(1, 1, frames * height * width, -1) + + +def _axis_rope_complex(length: int, complex_dims: int, device: torch.device) -> Tensor: + if complex_dims == 0: + return torch.empty(length, 0, dtype=torch.complex64, device=device) + dim = complex_dims * 2 + exponent = torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim + freqs = 1.0 / (10000.0**exponent) + positions = torch.arange(length, dtype=torch.float32, device=device) + angles = positions[:, None] * freqs[None] + return torch.polar(torch.ones_like(angles), angles) + + +def _apply_complex_rope(x: Tensor, rotary_emb: Tensor | None) -> Tensor: + if rotary_emb is None: + return x + batch, tokens, heads, dim = x.shape + freqs = rotary_emb.squeeze(0).squeeze(0) + if freqs.shape[0] != tokens or freqs.shape[1] != dim // 2: + raise ValueError( + f"RoPE shape {tuple(freqs.shape)} is incompatible with {(tokens, dim)}." + ) + x_float = x.float() + x_complex = torch.view_as_complex( + x_float.reshape(batch, tokens, heads, dim // 2, 2) + ) + rotated = torch.view_as_real(x_complex * freqs[None, :, None, :]).flatten(-2) + return rotated.to(dtype=x.dtype) + + +def _slice_rope_for_camera(rotary_emb: Tensor | None, head_dim: int) -> Tensor | None: + if rotary_emb is None: + return None + orig_t_size = head_dim // 2 - 2 * (head_dim // 6) + orig_h_size = head_dim // 6 + new_head_dim = head_dim // 2 + new_t_size = new_head_dim // 2 - 2 * (new_head_dim // 6) + new_h_size = new_head_dim // 6 + new_w_size = new_head_dim // 6 + t_part = rotary_emb[..., :new_t_size] + h_part = rotary_emb[..., orig_t_size : orig_t_size + new_h_size] + w_part = rotary_emb[ + ..., + orig_t_size + + orig_h_size : orig_t_size + + orig_h_size + + new_w_size, + ] + return torch.cat([t_part, h_part, w_part], dim=-1) + + +def _bidirectional_gdn_scan( + *, + q: Tensor, + k: Tensor, + q_rot: Tensor, + k_rot: Tensor, + v: Tensor, + beta: Tensor, + decay: Tensor, + HW: tuple[int, int, int], + eps: float, +) -> Tensor: + m_hist, z_hist = _gdn_histories( + k=k, + k_rot=k_rot, + v=v, + beta=beta, + decay=decay, + HW=HW, + include_denominator=True, + ) + batch, tokens, heads, dim = q.shape + frames, height, width = HW + spatial = height * width + q = q.reshape(batch, frames, spatial, heads, dim).permute(0, 3, 1, 2, 4) + q_rot = q_rot.reshape(batch, frames, spatial, heads, dim).permute( + 0, + 3, + 1, + 2, + 4, + ) + num = torch.einsum("bhfsd,bhfde->bhfse", q_rot.float(), m_hist) + den = torch.einsum("bhfsd,bhfd->bhfs", q.float(), z_hist) + out = num / (den[..., None] + eps) + return out.permute(0, 2, 3, 1, 4).reshape(batch, tokens, heads, dim) + + +def _bidirectional_numerator_scan( + *, + q: Tensor, + k: Tensor, + v: Tensor, + beta: Tensor, + decay: Tensor, + HW: tuple[int, int, int], +) -> Tensor: + m_hist, _z_hist = _gdn_histories( + k=k, + k_rot=k, + v=v, + beta=beta, + decay=decay, + HW=HW, + include_denominator=False, + ) + batch, tokens, heads, dim = q.shape + frames, height, width = HW + spatial = height * width + q = q.reshape(batch, frames, spatial, heads, dim).permute(0, 3, 1, 2, 4) + out = torch.einsum("bhfsd,bhfde->bhfse", q.float(), m_hist) + return out.permute(0, 2, 3, 1, 4).reshape(batch, tokens, heads, dim) + + +def _gdn_histories( + *, + k: Tensor, + k_rot: Tensor, + v: Tensor, + beta: Tensor, + decay: Tensor, + HW: tuple[int, int, int], + include_denominator: bool, +) -> tuple[Tensor, Tensor | None]: + batch, tokens, heads, dim = k.shape + frames, height, width = HW + spatial = height * width + k = k.reshape(batch, frames, spatial, heads, dim).permute(0, 3, 1, 2, 4).float() + k_rot = ( + k_rot.reshape(batch, frames, spatial, heads, dim) + .permute(0, 3, 1, 2, 4) + .float() + ) + v = v.reshape(batch, frames, spatial, heads, dim).permute(0, 3, 1, 2, 4).float() + eye = torch.eye(dim, dtype=torch.float32, device=k.device).view(1, 1, 1, dim, dim) + + # The per-frame GDN transition is linear in the running state, so the + # spatial reductions that build each frame's transition are batched over all + # frames in one einsum, leaving only the (light) sequential state scan as a + # loop — run forward and reverse to form the bidirectional history. + p_kv = torch.einsum("bhfsd,bhfse->bhfde", k_rot, beta[..., None] * k_rot) + a_kv = torch.einsum("bhfsd,bhfse->bhfde", k_rot, beta[..., None] * v) + trans_m = decay[..., None, None] * (eye - p_kv) + m_state = torch.zeros(batch, heads, dim, dim, dtype=torch.float32, device=k.device) + m_forward: list[Tensor] = [] + for frame in range(frames): + m_state = torch.einsum("bhde,bhef->bhdf", trans_m[:, :, frame], m_state) + m_state = m_state + a_kv[:, :, frame] + m_forward.append(m_state) + m_hist = torch.stack(m_forward, dim=2) + m_state = torch.zeros(batch, heads, dim, dim, dtype=torch.float32, device=k.device) + for src in range(frames - 1, 0, -1): + m_state = torch.einsum("bhde,bhef->bhdf", trans_m[:, :, src], m_state) + m_state = m_state + a_kv[:, :, src] + m_hist[:, :, src - 1] = m_hist[:, :, src - 1] + m_state + + z_hist = None + if include_denominator: + beta_k = beta[..., None] * k + p_z = torch.einsum("bhfsd,bhfse->bhfde", k, beta_k) + b_z = beta_k.sum(dim=3) + trans_z = decay[..., None, None] * (eye - p_z) + z_state = torch.zeros(batch, heads, dim, dtype=torch.float32, device=k.device) + z_forward: list[Tensor] = [] + for frame in range(frames): + z_state = torch.einsum("bhde,bhe->bhd", trans_z[:, :, frame], z_state) + z_state = z_state + b_z[:, :, frame] + z_forward.append(z_state) + z_hist = torch.stack(z_forward, dim=2) + z_state = torch.zeros(batch, heads, dim, dtype=torch.float32, device=k.device) + for src in range(frames - 1, 0, -1): + z_state = torch.einsum("bhde,bhe->bhd", trans_z[:, :, src], z_state) + z_state = z_state + b_z[:, :, src] + z_hist[:, :, src - 1] = z_hist[:, :, src - 1] + z_state + + return m_hist, z_hist + + +def _camera_qkv( + module: Stage1SelfAttention, + x: Tensor, + HW: tuple[int, int, int], +) -> tuple[Tensor, Tensor, Tensor]: + batch, tokens, _channels = x.shape + qkv_weight = torch.cat( + [ + module.q_proj_cam.weight, + module.k_proj_cam.weight, + module.v_proj_cam.weight, + ], + ) + qkv_bias = torch.cat( + [ + module.q_proj_cam.bias, + module.k_proj_cam.bias, + module.v_proj_cam.bias, + ], + ) + q_raw_flat, k_raw_flat, v_raw_flat = F.linear(x, qkv_weight, qkv_bias).chunk( + 3, + dim=-1, + ) + if hasattr(module, "conv_k_cam"): + k_raw_flat = _apply_bidirectional_temporal_conv(k_raw_flat, module.conv_k_cam, HW) + q_raw = q_raw_flat.reshape(batch, tokens, module.heads, module.dim) + k_raw = k_raw_flat.reshape(batch, tokens, module.heads, module.dim) + v_raw = v_raw_flat.reshape(batch, tokens, module.heads, module.dim) + return q_raw, k_raw, v_raw + + +def _prepare_ucpe_qkv( + q_raw: Tensor, + k_raw: Tensor, + v_raw: Tensor, + *, + camera_conditions: Tensor, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + camera_cache: _CameraProjectionCache | None, + q_norm_weight: Tensor, + k_norm_weight: Tensor, + norm_eps: float, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Callable[[Tensor], Tensor]]: + batch, tokens, heads, dim = q_raw.shape + q_norm = _rmsnorm_relu_heads(q_raw, q_norm_weight, norm_eps) + k_norm = _rmsnorm_relu_heads( + k_raw, + k_norm_weight, + norm_eps, + scale=_gdn_key_scale(dim, HW), + ) + v_float = v_raw.float() + if camera_cache is None: + camera_cache = _prepare_camera_projection_cache( + camera_conditions, + HW=HW, + rotary_emb=rotary_emb, + head_dim=dim, + ) + + q_trans = _ucpe_transform_apply( + q_norm, + camera_cache.proj_q, + camera_cache.rope_cam, + inverse_rope=False, + ) + k_trans, k_pre_sq, k_post_sq = _ucpe_transform( + k_norm, + camera_cache.proj_kv, + camera_cache.rope_cam, + inverse_rope=False, + ) + v_trans = _ucpe_transform_apply( + v_float, + camera_cache.proj_kv, + camera_cache.rope_cam, + inverse_rope=False, + ) + inflation_sq = k_post_sq.clamp_min(1e-12) / k_pre_sq.clamp_min(1e-12) + + def output_projector(out: Tensor) -> Tensor: + projected = _ucpe_transform_apply( + out.float(), + camera_cache.proj, + camera_cache.rope_cam, + inverse_rope=True, + ) + return projected.to(dtype=out.dtype) + + return q_trans, k_trans, v_trans, inflation_sq, output_projector + + +def _prepare_ucpe_qkv_softmax( + q_raw: Tensor, + k_raw: Tensor, + v_raw: Tensor, + *, + camera_conditions: Tensor, + HW: tuple[int, int, int], + rotary_emb: Tensor | None, + camera_cache: _CameraProjectionCache | None, + q_norm_weight: Tensor, + k_norm_weight: Tensor, + norm_eps: float, +) -> tuple[Tensor, Tensor, Tensor, Callable[[Tensor], Tensor]]: + batch, tokens, heads, dim = q_raw.shape + q_inv = _inv_rms(q_raw, norm_eps) + k_inv = _inv_rms(k_raw, norm_eps) + q_weight = q_norm_weight.float().view(heads, dim) + k_weight = k_norm_weight.float().view(heads, dim) + q_norm = ( + q_raw.float() * q_inv[:, :, None, None] * q_weight[None, None] + ).to(dtype=q_raw.dtype) + k_norm = ( + k_raw.float() * k_inv[:, :, None, None] * k_weight[None, None] + ).to(dtype=k_raw.dtype) + if camera_cache is None: + camera_cache = _prepare_camera_projection_cache( + camera_conditions, + HW=HW, + rotary_emb=rotary_emb, + head_dim=dim, + ) + q_trans = _ucpe_transform_apply( + q_norm, + camera_cache.proj_q, + camera_cache.rope_cam, + inverse_rope=False, + ).to(dtype=q_raw.dtype) + kv_trans = _ucpe_transform_apply( + torch.cat([k_norm, v_raw], dim=2), + camera_cache.proj_kv, + camera_cache.rope_cam, + inverse_rope=False, + ).to(dtype=q_raw.dtype) + k_trans, v_trans = kv_trans.chunk(2, dim=2) + + def output_projector(out: Tensor) -> Tensor: + projected = _ucpe_transform_apply( + out.to(dtype=q_raw.dtype), + camera_cache.proj, + camera_cache.rope_cam, + inverse_rope=True, + ) + return projected.to(dtype=out.dtype) + + return q_trans, k_trans, v_trans, output_projector + + +def _inv_rms(x: Tensor, eps: float) -> Tensor: + batch, tokens, heads, dim = x.shape + channels = heads * dim + return torch.rsqrt(x.float().pow(2).sum(dim=(-1, -2)) / channels + eps) + + +def _rmsnorm_relu_heads( + x: Tensor, + weight: Tensor, + eps: float, + *, + scale: float = 1.0, +) -> Tensor: + fast = _rmsnorm_relu_heads_fast(x, weight, eps, scale=scale) + if fast is not None: + return fast + batch, tokens, heads, dim = x.shape + inv_rms = _inv_rms(x, eps) + norm_weight = weight.float().view(heads, dim) + out = F.relu(x.float() * inv_rms[:, :, None, None] * norm_weight[None, None]) + if scale != 1.0: + out = out * scale + return out + + +def _rmsnorm_relu_heads_fast( + x: Tensor, + weight: Tensor, + eps: float, + *, + scale: float, +) -> Tensor | None: + batch, tokens, heads, dim = x.shape + channels = heads * dim + block_c = 1 << (channels - 1).bit_length() + if ( + _DISABLE_RMS_RELU_FAST + or not x.is_cuda + or not weight.is_cuda + or torch.is_grad_enabled() + or triton is None + or _rmsnorm_relu_heads_kernel is None + or weight.numel() != channels + or block_c > 4096 + ): + return None + out = torch.empty( + (batch, tokens, heads, dim), + device=x.device, + dtype=torch.float32, + ) + grid = (batch * tokens,) + _rmsnorm_relu_heads_kernel[grid]( + x, + weight, + out, + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + tokens, + heads, + dim, + channels, + float(eps), + float(scale), + block_c, + ) + return out + + +def _ucpe_transform( + x: Tensor, + matrix: Tensor, + rotary_emb: Tensor | None, + *, + inverse_rope: bool, +) -> tuple[Tensor, Tensor, Tensor]: + out = _ucpe_transform_apply(x, matrix, rotary_emb, inverse_rope=inverse_rope) + fast_norms = _ucpe_transform_norms_fast(x, out) + if fast_norms is not None: + pre_sq, post_sq = fast_norms + return out, pre_sq, post_sq + pre_sq = _ucpe_transform_norms_eager(x) + post_sq = _ucpe_transform_norms_eager(out) + return out, pre_sq, post_sq + + +def _ucpe_transform_norms_eager(x: Tensor) -> Tensor: + return x.float().pow(2).sum(dim=-1).transpose(1, 2).contiguous() + + +def _ucpe_transform_apply( + x: Tensor, + matrix: Tensor, + rotary_emb: Tensor | None, + *, + inverse_rope: bool, +) -> Tensor: + fast = _ucpe_transform_apply_fast( + x, + matrix, + rotary_emb, + inverse_rope=inverse_rope, + ) + if fast is not None: + return fast + return _ucpe_transform_apply_eager( + x, + matrix, + rotary_emb, + inverse_rope=inverse_rope, + ) + + +def _ucpe_transform_apply_eager( + x: Tensor, + matrix: Tensor, + rotary_emb: Tensor | None, + *, + inverse_rope: bool, +) -> Tensor: + batch, tokens, heads, dim = x.shape + half = dim // 2 + if half % 4 != 0: + raise ValueError(f"UCPE requires head_dim/2 divisible by 4, got {half}.") + first = x[..., :half].reshape(batch, tokens, heads, half // 4, 4) + first_out = torch.einsum("bnij,bnhgj->bnhgi", matrix.float(), first.float()) + first_out = first_out.reshape(batch, tokens, heads, half) + second = x[..., half:] + if inverse_rope and rotary_emb is not None: + second_out = _apply_complex_rope(second, rotary_emb.conj()) + else: + second_out = _apply_complex_rope(second, rotary_emb) + out = torch.cat([first_out, second_out], dim=-1) + return out + + +def _ucpe_transform_apply_fast( + x: Tensor, + matrix: Tensor, + rotary_emb: Tensor | None, + *, + inverse_rope: bool, +) -> Tensor | None: + batch, tokens, heads, dim = x.shape + half = dim // 2 + if half % 4 != 0: + raise ValueError(f"UCPE requires head_dim/2 divisible by 4, got {half}.") + if ( + _DISABLE_UCPE_FAST + or not x.is_cuda + or not matrix.is_cuda + or rotary_emb is None + or not rotary_emb.is_cuda + or torch.is_grad_enabled() + or triton is None + or _ucpe_first_half_kernel is None + or _ucpe_rope_second_half_kernel is None + or matrix.shape[-2:] != (4, 4) + or matrix.shape[1] != tokens + or matrix.shape[0] not in (1, batch) + ): + return None + rope = torch.view_as_real(rotary_emb.squeeze(0).squeeze(0)) + if rope.shape != (tokens, half // 2, 2): + return None + + out = torch.empty_like(x, dtype=torch.float32) + groups4 = half // 4 + rope_pairs = half // 2 + block_n = 128 + first_grid = ( + batch, + heads, + groups4 * triton.cdiv(tokens, block_n), + ) + _ucpe_first_half_kernel[first_grid]( + x, + matrix, + out, + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + matrix.stride(0), + matrix.stride(1), + matrix.stride(2), + matrix.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + tokens, + groups4, + matrix.shape[0], + block_n, + ) + rope_grid = ( + batch, + heads, + rope_pairs * triton.cdiv(tokens, block_n), + ) + _ucpe_rope_second_half_kernel[rope_grid]( + x, + rope, + out, + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + rope.stride(0), + rope.stride(1), + rope.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + tokens, + half, + rope_pairs, + inverse_rope, + block_n, + ) + return out + + +def _ucpe_transform_norms_fast(x: Tensor, out: Tensor) -> tuple[Tensor, Tensor] | None: + batch, tokens, heads, dim = x.shape + block_d = 1 << (dim - 1).bit_length() + if ( + _DISABLE_UCPE_FAST + or not x.is_cuda + or not out.is_cuda + or torch.is_grad_enabled() + or triton is None + or _ucpe_norms_kernel is None + or out.shape != x.shape + or block_d > 256 + ): + return None + + pre_sq = torch.empty((batch, heads, tokens), device=x.device, dtype=torch.float32) + post_sq = torch.empty_like(pre_sq) + block_n = 16 + grid = ( + batch, + heads, + triton.cdiv(tokens, block_n), + ) + _ucpe_norms_kernel[grid]( + x, + out, + pre_sq, + post_sq, + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + pre_sq.stride(0), + pre_sq.stride(1), + pre_sq.stride(2), + post_sq.stride(0), + post_sq.stride(1), + post_sq.stride(2), + tokens, + dim, + block_n, + block_d, + ) + return pre_sq, post_sq + + +def _silu_multiply(value: Tensor, gate: Tensor, *, inplace: bool) -> Tensor: + fast = _silu_multiply_fast(value, gate) + if fast is not None: + return fast + return value * F.silu(gate, inplace=inplace) + + +def _silu_multiply_fast(value: Tensor, gate: Tensor) -> Tensor | None: + if ( + _DISABLE_GLU_FAST + or value.shape != gate.shape + or value.dim() != 4 + or not value.is_cuda + or not gate.is_cuda + or torch.is_grad_enabled() + or triton is None + or _silu_multiply_kernel is None + ): + return None + batch, channels, height, width = value.shape + total = batch * channels * height * width + if total <= 0: + return torch.empty_like(value, memory_format=torch.channels_last) + out = torch.empty( + (batch, channels, height, width), + device=value.device, + dtype=value.dtype, + memory_format=torch.channels_last, + ) + block = 128 + grid = (triton.cdiv(total, block),) + _silu_multiply_kernel[grid]( + value, + gate, + out, + value.stride(0), + value.stride(1), + value.stride(2), + value.stride(3), + gate.stride(0), + gate.stride(1), + gate.stride(2), + gate.stride(3), + channels, + height, + width, + total, + block, + ) + return out + + +def _camera_ray_mats( + camera_conditions: Tensor, + HW: tuple[int, int, int], +) -> Tensor: + batch, frames = camera_conditions.shape[:2] + latent_frames, height, width = HW + if frames != latent_frames: + raise ValueError( + f"camera_conditions has {frames} frames but latent grid has {latent_frames}." + ) + c2w = camera_conditions[..., :16].reshape(batch, frames, 4, 4) + fx = camera_conditions[..., 16] + fy = camera_conditions[..., 17] + cx = camera_conditions[..., 18] + cy = camera_conditions[..., 19] + y_grid, x_grid = torch.meshgrid( + torch.arange(height, dtype=torch.float32, device=camera_conditions.device), + torch.arange(width, dtype=torch.float32, device=camera_conditions.device), + indexing="ij", + ) + x = (x_grid.view(1, 1, height, width) - cx[:, :, None, None]) / fx[ + :, + :, + None, + None, + ].clamp_min(1e-6) + y = (y_grid.view(1, 1, height, width) - cy[:, :, None, None]) / fy[ + :, + :, + None, + None, + ].clamp_min(1e-6) + dirs_cam = torch.stack([x, y, torch.ones_like(x)], dim=-1) + dirs_cam = F.normalize(dirs_cam, dim=-1, eps=1e-6) + rotation = c2w[..., :3, :3] + translation = c2w[..., :3, 3] + dirs_world = torch.einsum("btij,bthwj->bthwi", rotation, dirs_cam) + cam_y = rotation[..., :, 1].view(batch, frames, 1, 1, 3).expand_as(dirs_world) + z_ray = F.normalize(dirs_world, dim=-1, eps=1e-6) + x_ray = F.normalize(torch.cross(cam_y, z_ray, dim=-1), dim=-1, eps=1e-6) + y_ray = F.normalize(torch.cross(z_ray, x_ray, dim=-1), dim=-1, eps=1e-6) + ray_to_world = torch.stack([x_ray, y_ray, z_ray], dim=-1) + world_to_ray = ray_to_world.transpose(-1, -2) + origin = translation.view(batch, frames, 1, 1, 3).expand_as(dirs_world) + trans = -torch.einsum("bthwij,bthwj->bthwi", world_to_ray, origin) + mats = torch.zeros( + batch, + frames, + height, + width, + 4, + 4, + dtype=torch.float32, + device=camera_conditions.device, + ) + mats[..., :3, :3] = world_to_ray + mats[..., :3, 3] = trans + mats[..., 3, 3] = 1.0 + return mats + + +def _invert_se3(transforms: Tensor) -> Tensor: + rotation_inv = transforms[..., :3, :3].transpose(-1, -2) + out = torch.zeros_like(transforms) + out[..., :3, :3] = rotation_inv + out[..., :3, 3] = -torch.einsum( + "...ij,...j->...i", + rotation_inv, + transforms[..., :3, 3], + ) + out[..., 3, 3] = 1.0 + return out + + +__all__ = [ + "SANA_WM_STAGE1_SPEC", + "GLUMBConvTemp", + "RMSNorm", + "SanaWMStage1Block", + "SanaWMStage1Model", + "SanaWMStage1Spec", + "Stage1CrossAttention", + "Stage1SelfAttention", + "linearize_stage1_ffn_for_quant", +] diff --git a/integrations/sana/sana_wm/transformer.py b/integrations/sana/sana_wm/transformer.py new file mode 100644 index 000000000..c05a71c59 --- /dev/null +++ b/integrations/sana/sana_wm/transformer.py @@ -0,0 +1,854 @@ +# 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. + +"""FlashDreams adapter for the SANA-WM Stage-1 DiT.""" + +from __future__ import annotations + +import gc +import time +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any, Literal, cast + +import torch +import torch.nn as nn +import yaml +from loguru import logger +from torch import Tensor + +from flashdreams.infra.diffusion.transformer import ( + Transformer, + TransformerAutoregressiveCache, + TransformerConfig, +) +from sana_wm._tools import find_model, resolve_hf_path +from sana_wm.constants import ( + DEFAULT_VIDEO_HEIGHT, + DEFAULT_VIDEO_WIDTH, + SANA_WM_CONFIG_PATH, + SANA_WM_MODEL_PATH, +) +from sana_wm.quant import ( + TorchScaledMMFP4Recipe, + TorchScaledMMFP8Recipe, + replace_linear_with_quant, +) +from sana_wm.stage1_model import SanaWMStage1Model, linearize_stage1_ffn_for_quant + +Precision = Literal["bf16", "fp8", "fp4"] +QuantBackend = Literal["auto", "torch", "torch-fp8", "torch-fp4"] + +_STAGE1_QUANT_SKIP_DEFAULTS = ( + "^x_embedder", + "^raymap_embedder", + "^plucker_embedder", + "^t_embedder", + "^t_block", + "^y_embedder", + "^final_layer", +) +_STAGE1_QUANT_INCLUDE_DEFAULTS = ( + r"^blocks\.\d+\.attn\.qkv$", + r"^blocks\.\d+\.attn\.proj$", + r"^blocks\.\d+\.attn\.beta_proj$", + r"^blocks\.\d+\.attn\.gate_proj$", + r"^blocks\.\d+\.attn\.output_gate$", + r"^blocks\.\d+\.cross_attn\.", + r"^blocks\.\d+\.mlp\.inverted_conv\.linear$", + r"^blocks\.\d+\.mlp\.point_conv\.linear$", +) + + +@dataclass(kw_only=True) +class SanaWMStage1Conditioning: + """Per-rollout inputs needed by the SANA-WM Stage-1 sampler.""" + + condition: Tensor + uncondition: Tensor | None + model_kwargs: dict[str, object] + first_latent: Tensor + latent_shape: tuple[int, int, int, int, int] + cfg_scale: float + flow_shift: float + steps: int + seed: int + + +@dataclass(kw_only=True) +class SanaWMTransformerCache(TransformerAutoregressiveCache): + """AR cache for the one-shot SANA-WM Stage-1 rollout.""" + + conditioning: SanaWMStage1Conditioning | None = None + + +@dataclass(kw_only=True) +class SanaWMTransformerConfig(TransformerConfig): + """Config for the SANA-WM Stage-1 transformer adapter.""" + + _target: type["SanaWMTransformer"] = field( + default_factory=lambda: SanaWMTransformer + ) + + config_path: str = SANA_WM_CONFIG_PATH + """SANA-WM inference YAML path or ``hf://`` URI.""" + + checkpoint_path: str = SANA_WM_MODEL_PATH + """SANA-WM Stage-1 checkpoint path or ``hf://`` URI.""" + + stage1_precision: Precision = "bf16" + """Stage-1 precision requested by the runner.""" + + quant_backend: QuantBackend = "auto" + """Low-precision linear replacement backend for FP8/FP4 modes.""" + + height: int = DEFAULT_VIDEO_HEIGHT + """Pixel height used by the public SANA-WM bidirectional release.""" + + width: int = DEFAULT_VIDEO_WIDTH + """Pixel width used by the public SANA-WM bidirectional release.""" + + offload_stage1: bool = False + """Tear down the Stage-1 DiT after sampling, before decode/refine. Default + keeps it resident (fastest when memory is available); enable only to free + Stage-1 memory for the VAE/refiner on memory-constrained GPUs.""" + + +class SanaWMTransformer(Transformer[SanaWMTransformerCache]): + """FlashDreams adapter for the SANA-WM Stage-1 model call.""" + + def __init__(self, config: SanaWMTransformerConfig) -> None: + super().__init__(config) + self.config: SanaWMTransformerConfig = config + self._dummy = nn.Parameter(torch.empty(0)) + self._runtime_config: Any | None = None + self._model_path: str | None = None + self.weight_dtype: torch.dtype | None = None + self._model_built = False + self._stage1_quantized = False + + @property + def latent_shape(self) -> tuple[int, ...]: + """Return the current rollout latent shape or the public default.""" + if self._runtime_config is None: + return (1, 128, 21, 22, 40) + cfg = self._ensure_runtime_config() + latent_t = (161 - 1) // int(cfg.vae.vae_stride[0]) + 1 + return ( + 1, + int(cfg.vae.vae_latent_dim), + latent_t, + self.config.height // int(cfg.vae.vae_stride[-1]), + self.config.width // int(cfg.vae.vae_stride[-1]), + ) + + def initialize_autoregressive_cache( + self, + *, + conditioning: SanaWMStage1Conditioning | None = None, + **_: Any, + ) -> SanaWMTransformerCache: + """Build a cache containing per-rollout SANA-WM conditioning.""" + return SanaWMTransformerCache(conditioning=conditioning) + + def predict_flow( + self, + noisy_latent: Tensor, + timestep: Tensor, + cache: SanaWMTransformerCache, + input: Any = None, + model_kwargs: dict[str, object] | None = None, + ) -> Tensor: + """Execute one SANA-WM DiT flow prediction.""" + if isinstance(input, SanaWMStage1Conditioning): + cache.conditioning = input + return self._predict_conditioned( + noisy_latent=noisy_latent, + timestep=timestep, + conditioning=input, + ) + + conditioning = _require_conditioning(cache) + prompt_embeds = ( + cast(Tensor, input) if input is not None else conditioning.condition + ) + kwargs = conditioning.model_kwargs if model_kwargs is None else model_kwargs + return self._predict_with_prompt( + noisy_latent, + timestep, + prompt_embeds, + kwargs, + ) + + def finalize_kv_cache( + self, + noisy_latent: Tensor, + timestep: Tensor, + cache: SanaWMTransformerCache, + input: Any = None, + ) -> None: + """SANA-WM bidirectional inference has no streaming KV cache to advance.""" + del noisy_latent, timestep, cache, input + + def patchify_and_maybe_split_cp(self, x: Any) -> Any: + """SANA-WM Stage-1 latents are already in model layout.""" + return x + + def unpatchify_and_maybe_gather_cp(self, x: Tensor) -> Tensor: + """SANA-WM Stage-1 latents are already in model layout.""" + return x + + def initial_noise( + self, + *, + latent_shape: tuple[int, ...], + rng: torch.Generator | None, + cache: SanaWMTransformerCache, + input: Any = None, + ) -> Tensor: + """Draw SANA-WM's first-frame-pinned Stage-1 initial latent.""" + del latent_shape, rng + if isinstance(input, SanaWMStage1Conditioning): + cache.conditioning = input + conditioning = input + else: + conditioning = _require_conditioning(cache) + return self.initial_latents(conditioning) + + def postprocess_clean_latent( + self, + clean_latent: Tensor, + cache: SanaWMTransformerCache, + input: Any = None, + ) -> Tensor: + """Release Stage-1 runtime before decode/refine only when offloading. + + By default the Stage-1 DiT stays resident so a reused pipeline never + reloads and re-quantizes it; ``offload_stage1`` restores the old + free-before-decode behavior for memory-constrained GPUs. + """ + del input + if self.config.offload_stage1: + self.release_stage1_runtime(cache) + elif cache is not None: + cache.conditioning = None + return clean_latent + + def initial_latents(self, conditioning: SanaWMStage1Conditioning) -> Tensor: + """Draw Stage-1 initial noise and pin the encoded first frame.""" + generator = torch.Generator(device=self.device).manual_seed(conditioning.seed) + latents = torch.randn( + conditioning.latent_shape, + dtype=self._ensure_weight_dtype(), + device=self.device, + generator=generator, + ) + latents[:, :, :1] = conditioning.first_latent + return latents + + def release_stage1_runtime( + self, + cache: SanaWMTransformerCache | None = None, + ) -> None: + """Release Stage-1-only tensors before VAE/refiner work.""" + free_before_gib: float | None = None + if torch.cuda.is_available(): + try: + free_before, _total = torch.cuda.mem_get_info() + free_before_gib = free_before / (1024**3) + except RuntimeError: + free_before_gib = None + if cache is not None: + cache.conditioning = None + + for attr in ("model",): + module = getattr(self, attr, None) + if module is None: + continue + try: + module.to("meta") + except Exception: + try: + module.to("cpu") + except Exception: + pass + setattr(self, attr, None) + + self._model_built = False + self._stage1_quantized = False + if torch.cuda.is_available(): + torch.cuda.empty_cache() + try: + free_after, _total = torch.cuda.mem_get_info() + if free_before_gib is not None: + logger.info( + "[stage1] released Stage-1 runtime before decode/refine " + "(free CUDA memory: {:.2f} -> {:.2f} GiB)", + free_before_gib, + free_after / (1024**3), + ) + else: + logger.info( + "[stage1] released Stage-1 runtime before decode/refine" + ) + except RuntimeError: + logger.info("[stage1] released Stage-1 runtime before decode/refine") + gc.collect() + + def _predict_conditioned( + self, + *, + noisy_latent: Tensor, + timestep: Tensor, + conditioning: SanaWMStage1Conditioning, + ) -> Tensor: + self._prepare_static_model_kwargs(conditioning) + model_timestep = _conditioned_frame_timestep( + noisy_latent=noisy_latent, + timestep=timestep, + conditioning=conditioning, + num_train_timesteps=1000, + ) + if conditioning.cfg_scale <= 1.0: + flow = self._predict_with_prompt( + noisy_latent, + model_timestep, + conditioning.condition, + _condition_model_kwargs(conditioning.model_kwargs), + ) + return _zero_conditioned_frame_flow(flow, conditioning) + if conditioning.uncondition is None: + raise RuntimeError("CFG was requested without negative prompt embeds.") + + noise_pred = self._predict_with_prompt( + torch.cat([noisy_latent, noisy_latent], dim=0), + torch.cat([model_timestep, model_timestep], dim=0), + torch.cat([conditioning.uncondition, conditioning.condition], dim=0), + _batched_cfg_model_kwargs(conditioning.model_kwargs), + ) + flow = _cfg_guidance(noise_pred, conditioning.cfg_scale) + return _zero_conditioned_frame_flow(flow, conditioning) + + def _prepare_static_model_kwargs( + self, + conditioning: SanaWMStage1Conditioning, + ) -> None: + """Cache model projections for conditioning that is constant per rollout.""" + model_kwargs = conditioning.model_kwargs + self._ensure_model() + chunk_plucker = model_kwargs.get("chunk_plucker") + prepare_plucker = getattr(self.model, "prepare_plucker_embedding", None) + if ( + "chunk_plucker_emb" not in model_kwargs + and isinstance(chunk_plucker, Tensor) + and callable(prepare_plucker) + ): + with torch.inference_mode(): + model_kwargs["chunk_plucker_emb"] = prepare_plucker(chunk_plucker) + del model_kwargs["chunk_plucker"] + + camera_conditions = model_kwargs.get("camera_conditions") + prepare_camera = getattr(self.model, "prepare_camera_projection_cache", None) + if ( + "camera_cache" not in model_kwargs + and isinstance(camera_conditions, Tensor) + and camera_conditions.shape[0] == 1 + and callable(prepare_camera) + ): + _batch, _channels, frames, height, width = conditioning.latent_shape + with torch.inference_mode(): + rotary_emb, camera_cache = prepare_camera( + camera_conditions, + frames=frames, + height=height, + width=width, + ) + model_kwargs["rotary_emb"] = rotary_emb + model_kwargs["camera_cache"] = camera_cache + + def _predict_with_prompt( + self, + noisy_latent: Tensor, + timestep: Tensor, + prompt_embeds: Tensor, + model_kwargs: dict[str, object], + ) -> Tensor: + self._ensure_model() + return self.model( + noisy_latent, + timestep, + prompt_embeds, + **model_kwargs, + ) + + def _ensure_runtime_config(self) -> Any: + if self._runtime_config is not None: + return self._runtime_config + self._runtime_config = _load_inference_config(self.config.config_path) + return self._runtime_config + + def _ensure_weight_dtype(self) -> torch.dtype: + if self.weight_dtype is None: + self.weight_dtype = _get_weight_dtype( + self._ensure_runtime_config().model.mixed_precision + ) + return self.weight_dtype + + def _ensure_model(self) -> None: + if self._model_built: + self._prepare_stage1_quant() + return + t0 = time.perf_counter() + cfg = self._ensure_runtime_config() + weight_dtype = self._ensure_weight_dtype() + model = SanaWMStage1Model().to(self.device) + logger.info( + "[Sana] Loaded {} ({:,} params)", + cfg.model.model, + sum(p.numel() for p in model.parameters()), + ) + self._model_path = resolve_hf_path(self.config.checkpoint_path) + state = find_model(self._model_path) + if "generator" in state: + state = state["generator"] + if "state_dict" not in state: + state = { + "state_dict": { + (k[len("model.") :] if k.startswith("model.") else k): v + for k, v in state.items() + } + } + missing, unexpected = model.load_state_dict(state["state_dict"], strict=True) + if missing: + logger.warning("[Sana] Missing keys: {}", missing) + if unexpected: + logger.warning("[Sana] Unexpected keys: {}", unexpected) + self.model = model.eval().to(weight_dtype) + self._model_built = True + self._prepare_stage1_quant() + logger.info( + "[timing] stage1 build+load+quant: {:.3f}s (precision={})", + time.perf_counter() - t0, + self.config.stage1_precision, + ) + + def _prepare_stage1_quant(self) -> None: + if self._stage1_quantized or self.config.stage1_precision == "bf16": + return + if self.config.stage1_precision == "fp8": + recipe = TorchScaledMMFP8Recipe() + else: + recipe = TorchScaledMMFP4Recipe() + linearized, linearize_skipped = linearize_stage1_ffn_for_quant(self.model) + if linearized > 0 or linearize_skipped > 0: + logger.info( + "[stage1-quant] linearized {} FFN pointwise convs (skipped {})", + linearized, + linearize_skipped, + ) + converted, skipped = replace_linear_with_quant( + self.model, + recipe=recipe, + params_dtype=self._ensure_weight_dtype(), + skip_patterns=_STAGE1_QUANT_SKIP_DEFAULTS, + include_patterns=_stage1_quant_include_patterns(), + ) + if converted <= 0: + raise RuntimeError( + f"SANA-WM {self.config.stage1_precision} converted no " + f"Stage-1 Linear layers; skipped={skipped}." + ) + self._stage1_quantized = True + recipe_detail = "" + if isinstance(recipe, TorchScaledMMFP4Recipe): + recipe_detail = ( + f" rht={recipe.use_rht}" + f" global_scale={recipe.use_global_scale}" + f" weight_scale_2d={recipe.weight_scale_2d}" + ) + logger.info( + "[stage1-quant] precision={}{} converted {} Linear layers (skipped {})", + self.config.stage1_precision, + recipe_detail, + converted, + skipped, + ) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def _stage1_quant_include_patterns() -> tuple[str, ...]: + """Return Stage-1 Linear names eligible for FP8/FP4 quantization.""" + return _STAGE1_QUANT_INCLUDE_DEFAULTS + + +def _require_conditioning( + cache: SanaWMTransformerCache, +) -> SanaWMStage1Conditioning: + if cache.conditioning is None: + raise RuntimeError("SANA-WM cache was initialized without conditioning.") + return cache.conditioning + + +def _condition_model_kwargs(model_kwargs: dict[str, object]) -> dict[str, object]: + """Return model kwargs for the positive prompt branch.""" + return { + key: value + for key, value in model_kwargs.items() + if key != "negative_mask" and not key.startswith("_") + } + + +def _batched_cfg_model_kwargs(model_kwargs: dict[str, object]) -> dict[str, object]: + """Return model kwargs for a single batched negative/positive CFG forward.""" + kwargs = _condition_model_kwargs(model_kwargs) + mask = kwargs.get("mask") + negative_mask = model_kwargs.get("negative_mask") + if isinstance(mask, Tensor): + mask_uncond = negative_mask if isinstance(negative_mask, Tensor) else mask + kwargs["mask"] = torch.cat([mask_uncond, mask], dim=0) + camera_conditions = kwargs.get("camera_conditions") + if isinstance(camera_conditions, Tensor) and camera_conditions.shape[0] != 1: + kwargs["camera_conditions"] = torch.cat( + [camera_conditions, camera_conditions], + dim=0, + ) + chunk_plucker = kwargs.get("chunk_plucker") + if isinstance(chunk_plucker, Tensor): + kwargs["chunk_plucker"] = torch.cat([chunk_plucker, chunk_plucker], dim=0) + return kwargs + + +def _cfg_guidance(noise_pred: Tensor, cfg_scale: float) -> Tensor: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2, dim=0) + return noise_pred_uncond + cfg_scale * (noise_pred_text - noise_pred_uncond) + + +def _conditioned_frame_timestep( + *, + noisy_latent: Tensor, + timestep: Tensor, + conditioning: SanaWMStage1Conditioning, + num_train_timesteps: int, +) -> Tensor: + """Expand scheduler timesteps and pin conditioned frames to timestep zero.""" + batch, _channels, frames, _height, _width = noisy_latent.shape + if timestep.ndim == 0: + frame_timestep = timestep.reshape(1, 1, 1).expand(batch, 1, frames) + elif timestep.ndim == 1: + frame_timestep = timestep.reshape(batch, 1, 1).expand(batch, 1, frames) + elif timestep.ndim == 3: + frame_timestep = timestep + elif timestep.ndim == noisy_latent.ndim: + frame_timestep = timestep[:, :1, :, 0, 0] + else: + raise ValueError( + "SANA-WM timestep must be scalar, [B], [B, 1, T], or latent-shaped; " + f"got shape={tuple(timestep.shape)}." + ) + if frame_timestep.shape != (batch, 1, frames): + raise ValueError( + "SANA-WM timestep shape is incompatible with the latent: " + f"got {tuple(frame_timestep.shape)}, expected {(batch, 1, frames)}." + ) + + frame_timestep = frame_timestep.to(device=noisy_latent.device, dtype=torch.float32) + condition_frame_mask = _condition_frame_mask( + conditioning, + batch=batch, + frames=frames, + device=noisy_latent.device, + ) + max_timestep = (1.0 - condition_frame_mask) * float(num_train_timesteps) + return torch.minimum(frame_timestep, max_timestep) + + +def _zero_conditioned_frame_flow( + flow: Tensor, + conditioning: SanaWMStage1Conditioning, +) -> Tensor: + """Prevent scalar scheduler updates from moving pinned condition frames.""" + batch, _channels, frames, _height, _width = flow.shape + condition_frame_mask = _condition_frame_mask( + conditioning, + batch=batch, + frames=frames, + device=flow.device, + ).to(dtype=torch.bool) + return torch.where(condition_frame_mask[:, :, :, None, None], 0, flow) + + +def _condition_frame_mask( + conditioning: SanaWMStage1Conditioning, + *, + batch: int, + frames: int, + device: torch.device, +) -> Tensor: + data_info = conditioning.model_kwargs.get("data_info", {}) + condition_frame_info = ( + data_info.get("condition_frame_info", {}) + if isinstance(data_info, dict) + else {} + ) + mask = torch.zeros(batch, 1, frames, dtype=torch.float32, device=device) + if not isinstance(condition_frame_info, dict): + return mask + for frame_idx in condition_frame_info: + index = int(frame_idx) + if 0 <= index < frames: + mask[:, :, index] = 1.0 + return mask + + +def _avoid_degenerate_tile_tail( + *, + sample_extent: int, + sample_tile_min: int, + sample_stride: int, + compression_ratio: int, +) -> int: + """Choose a tile stride that avoids final latent tiles of size one.""" + if compression_ratio <= 1: + return sample_stride + latent_extent = max(1, sample_extent // compression_ratio) + latent_tile_min = max(1, sample_tile_min // compression_ratio) + requested_latent_stride = max(1, sample_stride // compression_ratio) + + def tail_size(latent_stride: int) -> int: + last_start = ((latent_extent - 1) // latent_stride) * latent_stride + return latent_extent - last_start + + candidates = range( + min(requested_latent_stride, latent_tile_min), + 1, + -1, + ) + for latent_stride in candidates: + if tail_size(latent_stride) > 1: + return latent_stride * compression_ratio + + for latent_stride in range( + requested_latent_stride + 1, + latent_tile_min + 1, + ): + if tail_size(latent_stride) > 1: + return latent_stride * compression_ratio + + return sample_stride + + +def _load_inference_config(config_path: str) -> Any: + with open(resolve_hf_path(config_path), encoding="utf-8") as handle: + raw = yaml.safe_load(handle) or {} + if not isinstance(raw, dict): + raise TypeError(f"SANA-WM config must be a mapping, got {type(raw).__name__}.") + raw.setdefault("work_dir", "") + return _to_namespace(raw) + + +def _get_vae(*args: Any, **kwargs: Any) -> nn.Module: + name, model_path = args[:2] + device = kwargs["device"] + dtype = kwargs["dtype"] + if "LTX2VAE_diffusers" not in str(name): + raise ValueError(f"Unsupported SANA-WM VAE type: {name!r}") + from diffusers import AutoencoderKLLTX2Video + + return ( + AutoencoderKLLTX2Video.from_pretrained( + model_path, + subfolder="vae", + torch_dtype=dtype, + ) + .to(device) + .eval() + ) + + +def _get_tokenizer_and_text_encoder(*args: Any, **kwargs: Any) -> tuple[Any, nn.Module]: + name = kwargs.get("name", args[0] if args else "T5") + device = kwargs.get("device", "cuda") + model_id = _TEXT_ENCODER_MODEL_IDS.get(str(name)) + if model_id is None: + raise ValueError(f"Unsupported SANA-WM text encoder: {name!r}") + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + T5EncoderModel, + T5Tokenizer, + ) + + if "T5" in str(name): + tokenizer = T5Tokenizer.from_pretrained(model_id) + text_encoder = T5EncoderModel.from_pretrained( + model_id, + torch_dtype=torch.float16, + ).to(device) + return tokenizer, text_encoder.eval() + + tokenizer = AutoTokenizer.from_pretrained(model_id) + tokenizer.padding_side = "right" + text_encoder = ( + AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16) + .get_decoder() + .to(device) + .eval() + ) + return tokenizer, text_encoder + + +def _get_weight_dtype(value: str) -> torch.dtype: + normalized = str(value).lower() + if normalized in {"bf16", "bfloat16"}: + return torch.bfloat16 + if normalized in {"fp16", "float16", "half"}: + return torch.float16 + if normalized in {"fp32", "float32", "float"}: + return torch.float32 + raise ValueError(f"Unsupported dtype value: {value!r}") + + +class _ConfigNamespace(SimpleNamespace): + """Simple YAML-backed config object with dict-like ``get`` support.""" + + def get(self, name: str, default: Any = None) -> Any: + return getattr(self, name, default) + + +_TEXT_ENCODER_MODEL_IDS = { + "T5": "DeepFloyd/t5-v1_1-xxl", + "T5-small": "google/t5-v1_1-small", + "T5-base": "google/t5-v1_1-base", + "T5-large": "google/t5-v1_1-large", + "T5-xl": "google/t5-v1_1-xl", + "T5-xxl": "google/t5-v1_1-xxl", + "gemma-2b": "google/gemma-2b", + "gemma-2b-it": "google/gemma-2b-it", + "gemma-2-2b": "google/gemma-2-2b", + "gemma-2-2b-it": "Efficient-Large-Model/gemma-2-2b-it", + "gemma-2-9b": "google/gemma-2-9b", + "gemma-2-9b-it": "google/gemma-2-9b-it", +} + + +def _vae_encode_ltx2(name: str, vae: nn.Module, images: Tensor, *, device: torch.device) -> Tensor: + if "LTX2VAE_diffusers" not in name: + raise ValueError(f"Unsupported SANA-WM VAE encode type: {name!r}") + dtype = images.dtype + vae_device = next(vae.parameters()).device + vae_dtype = next(vae.parameters()).dtype + posterior = vae.encode(images.to(device=vae_device, dtype=vae_dtype)).latent_dist + z = posterior.mode() + latents_mean = vae.latents_mean.view(1, -1, 1, 1, 1).to(z.device, z.dtype) + latents_std = vae.latents_std.view(1, -1, 1, 1, 1).to(z.device, z.dtype) + z = (z - latents_mean) * vae.config.scaling_factor / latents_std + return z.to(device=device, dtype=dtype) + + +def _vae_decode_ltx2(name: str, vae: nn.Module, latents: Tensor) -> Tensor: + if "LTX2VAE_diffusers" not in name: + raise ValueError(f"Unsupported SANA-WM VAE decode type: {name!r}") + vae_device = next(vae.parameters()).device + vae_dtype = next(vae.parameters()).dtype + latents_mean = vae.latents_mean.view(1, -1, 1, 1, 1).to( + latents.device, + latents.dtype, + ) + latents_std = vae.latents_std.view(1, -1, 1, 1, 1).to( + latents.device, + latents.dtype, + ) + scaled = latents * latents_std / vae.config.scaling_factor + latents_mean + return vae.decode( + scaled.to(device=vae_device, dtype=vae_dtype), + temb=None, + return_dict=False, + )[0] + + +def _to_namespace(value: Any) -> Any: + if isinstance(value, dict): + return _ConfigNamespace( + **{str(key): _to_namespace(child) for key, child in value.items()} + ) + if isinstance(value, list): + return [_to_namespace(child) for child in value] + return value + + +def _chunk_index_from_config(config: Any, *, num_frames: int) -> list[int] | None: + model = getattr(config, "model", None) + if model is None: + return None + chunk_index = model.get("chunk_index", None) + chunk_size = model.get("chunk_size", None) + strategy = model.get("chunk_split_strategy", "uniform") + if chunk_index is not None: + if not isinstance(chunk_index, (list, tuple)): + raise TypeError( + f"chunk_index must be a list, got {type(chunk_index).__name__}" + ) + if len(chunk_index) == 0: + raise ValueError("chunk_index cannot be empty.") + return [int(index) for index in chunk_index] + if chunk_size is None: + return None + return _chunk_index_from_chunk_size( + num_frames, + int(chunk_size), + strategy=str(strategy), + ) + + +def _chunk_index_from_chunk_size( + num_frames: int, + chunk_size: int, + *, + strategy: str = "uniform", +) -> list[int]: + if num_frames <= 0: + raise ValueError(f"num_frames must be > 0, got {num_frames}.") + if chunk_size <= 0: + raise ValueError(f"chunk_size must be > 0, got {chunk_size}.") + normalized = strategy.lower() + if normalized in {"uniform", "default"}: + indices = list(range(0, num_frames, chunk_size)) + if len(indices) > 1 and (num_frames - indices[-1]) < chunk_size: + indices.pop() + return indices + if normalized in {"first_frame", "first_frame_alone", "first_frame_only"}: + if num_frames <= 1: + return [0] + indices = [0] + list(range(1, num_frames, chunk_size)) + if len(indices) > 2 and (num_frames - indices[-1]) < chunk_size: + indices.pop() + return indices + if normalized in {"first_plus_one", "first_chunk_plus_one"}: + if num_frames <= chunk_size + 1: + return [0] + indices = [0] + list(range(chunk_size + 1, num_frames, chunk_size)) + if len(indices) > 1 and (num_frames - indices[-1]) < chunk_size: + indices.pop() + return indices + raise ValueError(f"Unknown chunk_split_strategy {strategy!r}.") + + +__all__ = [ + "SanaWMTransformer", + "SanaWMTransformerCache", + "SanaWMTransformerConfig", + "SanaWMStage1Conditioning", +] diff --git a/integrations/sana/tests/parity_check/.gitignore b/integrations/sana/tests/parity_check/.gitignore new file mode 100644 index 000000000..2df75f55e --- /dev/null +++ b/integrations/sana/tests/parity_check/.gitignore @@ -0,0 +1,9 @@ +/Sana/ +/.venv/ +/outputs/ +__pycache__/ +*.pyc +*.npy +*.mp4 +*.json +*.log diff --git a/integrations/sana/tests/parity_check/README.md b/integrations/sana/tests/parity_check/README.md new file mode 100644 index 000000000..5171f7376 --- /dev/null +++ b/integrations/sana/tests/parity_check/README.md @@ -0,0 +1,112 @@ + + +# SANA-WM parity and benchmark harness + +This harness compares upstream +[`NVlabs/Sana`](https://github.com/NVlabs/Sana) SANA-WM against the in-tree +FlashDreams `sana-wm-bidirectional` integration on the same demo image, prompt, +camera trajectory, intrinsics, seed, resolution, and precision settings. + +All dependencies live in this directory's isolated `./.venv`. `run.sh` reuses or +clones the upstream checkout, pins it to `6298508`, and applies `changes.patch` +idempotently. The patch is instrumentation-only: it adds upstream CLI flags for +frame dumps, stats JSON, precision selection, cuDNN SDPA selection, and Stage-1 +`torch.compile` without changing the generation algorithm. A small +`compat/mmcv` shim covers the registry/logging imports needed by the inference +path. + +## Run parity + +```bash +cd integrations/sana/tests/parity_check +bash run.sh +``` + +Defaults: + +- `SANA_REPO=$HOME/dev/Sana` +- upstream pin `6298508` +- demo input `asset/sana_wm/demo_0.*` +- `NUM_FRAMES=121` +- `SEED=42` +- `NO_REFINER=1` to isolate Stage-1 plus SANA VAE decode +- `STAGE1_PRECISION=bf16` +- `REFINER_PRECISION=$STAGE1_PRECISION` +- `QUANT_BACKEND=auto` +- `FORCE_CUDNN_SDPA=0` + +Outputs are written under `outputs/parity/`: + +- `upstream/frames.npy` +- `upstream/stats.json` +- `flashdreams/frames.npy` +- `flashdreams/stats.json` +- `parity.json` + +Set `NO_REFINER=0` to compare the full Stage-1 + LTX-2 refiner path. Set +`COMPILE_STAGE1=1` to wrap each Stage-1 DiT with `torch.compile`; this is opt-in +because the pinned upstream SANA-WM stack can fail during TorchInductor Triton +compilation on current PyTorch/Triton builds. + +## Run benchmark + +```bash +cd integrations/sana/tests/parity_check +bash bench.sh +``` + +Benchmark defaults discard one warmup run and measure three additional runs: +`WARMUP_RUNS=1 MEASURED_RUNS=3 COMPILE_STAGE1=0 NO_REFINER=1 FORCE_CUDNN_SDPA=0`. +Outputs are under `outputs/bench/`: + +- `bench.json` - machine-readable inputs, medians, p90s, memory, and stage + timings. +- `bench.md` - human-readable report. +- `perf.md` - chart-ready data using the same benchmark metric as `bench.md`. + +The benchmark metric is post-load generation latency per generated clip. Model +construction, checkpoint loading, video writing, and frame dumps are outside the +timing boundary. SANA-WM renders each requested bidirectional clip in one +generation pass, rather than as independently timed frames. With the default +`NO_REFINER=1`, the timed work covers conditioning, Stage-1 DiT, and SANA VAE +decode. Set `NO_REFINER=0` to benchmark the full Stage-1 + LTX-2 refiner path +with the same timing boundary. + +`bench.md` also reports Stage-1 DiT, conditioning/encode, VAE decode, optional +refiner, memory, and frame-normalized diagnostic breakdowns. Those rows explain +the benchmark result; they are not a second benchmark metric. + +The model card should use one chart per precision. Each chart file should have +GPU/device as the first column and implementation as the series columns, for +example: + +```markdown +| device | official | flashdreams | +| --- | ---: | ---: | +| GB202 | 77826.74 | 76938.72 | +``` + +In precision-sweep mode, the chart-ready files are +`outputs/bench//perf.md`. The top-level `outputs/bench/perf.md` is a +precision summary, not the model-card chart data. + +Set `FORCE_CUDNN_SDPA=1` only for backend-isolation probes. It is not the +default SANA-WM benchmark setting. + +Set `DEVICE_LABEL` when generating chart data for docs: + +```bash +DEVICE_LABEL="RTX PRO 6000 Blackwell" bash bench.sh +``` + +Run the precision sweep used by the SANA-WM model card with: + +```bash +DEVICE_LABEL="RTX PRO 6000 Blackwell" BENCH_PRECISIONS=bf16,fp8,fp4 bash bench.sh +``` + +The scripts do not set allocator overrides or GPU wait loops. If GPU contention +matters in your environment, handle it outside the harness. diff --git a/integrations/sana/tests/parity_check/bench.sh b/integrations/sana/tests/parity_check/bench.sh new file mode 100644 index 000000000..c6f1a105f --- /dev/null +++ b/integrations/sana/tests/parity_check/bench.sh @@ -0,0 +1,269 @@ +#!/usr/bin/env bash +# 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. + +# Stack-matched SANA-WM upstream vs FlashDreams benchmark. The chart metric is +# post-load generation latency per generated clip. Defaults to Stage-1-only so +# conditioning, Stage-1 DiT, and SANA VAE decode are measured without the LTX-2 +# refiner. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +_abspath() { + case "$1" in + /*) printf '%s\n' "$1" ;; + *) printf '%s/%s\n' "${PWD}" "$1" ;; + esac +} + +SANA_REPO="$(_abspath "${SANA_REPO:-${SCRIPT_DIR}/Sana}")" +PATCH_FILE="${SCRIPT_DIR}/changes.patch" +REPO_URL="https://github.com/NVlabs/Sana.git" +PIN_COMMIT="6298508" + +OUTPUT_DIR="$(_abspath "${OUTPUT_DIR:-${SCRIPT_DIR}/outputs/bench}")" +UPSTREAM_ROOT="${OUTPUT_DIR}/upstream" +NATIVE_ROOT="${OUTPUT_DIR}/flashdreams" +IMAGE_PATH="$(_abspath "${IMAGE_PATH:-${SANA_REPO}/asset/sana_wm/demo_0.png}")" +PROMPT_PATH="$(_abspath "${PROMPT_PATH:-${SANA_REPO}/asset/sana_wm/demo_0.txt}")" +CAMERA_PATH="$(_abspath "${CAMERA_PATH:-${SANA_REPO}/asset/sana_wm/demo_0_pose.npy}")" +INTRINSICS_PATH="$(_abspath "${INTRINSICS_PATH:-${SANA_REPO}/asset/sana_wm/demo_0_intrinsics.npy}")" +NUM_FRAMES="${NUM_FRAMES:-121}" +FPS="${FPS:-16}" +STEP="${STEP:-60}" +CFG_SCALE="${CFG_SCALE:-5.0}" +SEED="${SEED:-42}" +NO_REFINER="${NO_REFINER:-1}" +FORCE_CUDNN_SDPA="${FORCE_CUDNN_SDPA:-0}" +COMPILE_STAGE1="${COMPILE_STAGE1:-0}" +WARMUP_RUNS="${WARMUP_RUNS:-1}" +MEASURED_RUNS="${MEASURED_RUNS:-3}" +DEVICE_LABEL="${DEVICE_LABEL:-GPU}" +STAGE1_PRECISION="${STAGE1_PRECISION:-bf16}" +REFINER_PRECISION="${REFINER_PRECISION:-${STAGE1_PRECISION}}" +QUANT_BACKEND="${QUANT_BACKEND:-auto}" +CHART_LABEL="${CHART_LABEL:-${DEVICE_LABEL}}" +BENCH_PRECISIONS="${BENCH_PRECISIONS:-}" + +_is_true() { + case "${1,,}" in + 1|true|yes|on) return 0 ;; + *) return 1 ;; + esac +} + +if [[ -n "${BENCH_PRECISIONS}" ]]; then + mkdir -p "${OUTPUT_DIR}" + SWEEP_ITEMS=() + IFS=',' read -r -a PRECISION_LIST <<< "${BENCH_PRECISIONS}" + for PRECISION_RAW in "${PRECISION_LIST[@]}"; do + PRECISION="${PRECISION_RAW//[[:space:]]/}" + if [[ -z "${PRECISION}" ]]; then + continue + fi + case "${PRECISION}" in + bf16|fp8|fp4) ;; + *) + echo "[bench] ERROR: unsupported BENCH_PRECISIONS entry: ${PRECISION}" >&2 + exit 1 + ;; + esac + PRECISION_LABEL="${PRECISION^^}" + PRECISION_OUTPUT_DIR="${OUTPUT_DIR}/${PRECISION}" + echo "[bench] precision sweep row ${PRECISION_LABEL} -> ${PRECISION_OUTPUT_DIR}" + BENCH_PRECISIONS="" \ + OUTPUT_DIR="${PRECISION_OUTPUT_DIR}" \ + STAGE1_PRECISION="${PRECISION}" \ + REFINER_PRECISION="${PRECISION}" \ + QUANT_BACKEND="${QUANT_BACKEND}" \ + CHART_LABEL="${CHART_LABEL}" \ + bash "${SCRIPT_DIR}/bench.sh" + SWEEP_ITEMS+=(--item "${PRECISION_LABEL}:${PRECISION_OUTPUT_DIR}/bench.json") + done + if [[ "${#SWEEP_ITEMS[@]}" -eq 0 ]]; then + echo "[bench] ERROR: BENCH_PRECISIONS produced no precision rows." >&2 + exit 1 + fi + echo "[bench] aggregating precision sweep -> ${OUTPUT_DIR}/bench.md" + ( cd "${SCRIPT_DIR}" && \ + uv run python "${SCRIPT_DIR}/bench_sweep_summary.py" \ + "${SWEEP_ITEMS[@]}" \ + --output-json "${OUTPUT_DIR}/bench.json" \ + --output-md "${OUTPUT_DIR}/bench.md" \ + --output-chart-md "${OUTPUT_DIR}/perf.md" ) + echo "[bench] done." + echo " precision summary: ${OUTPUT_DIR}/bench.md" + echo " precision data: ${OUTPUT_DIR}/perf.md" + echo " chart data: ${OUTPUT_DIR}/{bf16,fp8,fp4}/perf.md" + exit 0 +fi + +if [[ ! -d "${SANA_REPO}/.git" ]]; then + echo "[setup] cloning ${REPO_URL} -> ${SANA_REPO}" + git clone "${REPO_URL}" "${SANA_REPO}" +else + echo "[setup] repo already present at ${SANA_REPO}, skipping clone" +fi + +cd "${SANA_REPO}" +CURRENT_COMMIT="$(git rev-parse --short HEAD)" +if [[ "${CURRENT_COMMIT}" != "${PIN_COMMIT}" ]]; then + echo "[setup] checking out pinned commit ${PIN_COMMIT}" + git checkout "${PIN_COMMIT}" +else + echo "[setup] already at pinned commit ${PIN_COMMIT}, skipping checkout" +fi + +if git apply --reverse --check "${PATCH_FILE}" >/dev/null 2>&1; then + echo "[setup] patch already applied, skipping" +elif git apply --check "${PATCH_FILE}" >/dev/null 2>&1; then + echo "[setup] applying ${PATCH_FILE}" + git apply "${PATCH_FILE}" +else + echo "[setup] ERROR: ${PATCH_FILE} neither cleanly applies nor is already applied." >&2 + exit 1 +fi + +echo "[setup] ensuring Python deps via uv sync (isolated venv)" +( cd "${SCRIPT_DIR}" && uv sync ) +UPSTREAM_PYTHONPATH="${SCRIPT_DIR}/compat:${SANA_REPO}" +if [[ -n "${PYTHONPATH:-}" ]]; then + UPSTREAM_PYTHONPATH="${UPSTREAM_PYTHONPATH}:${PYTHONPATH}" +fi + +mkdir -p "${UPSTREAM_ROOT}" "${NATIVE_ROOT}" + +# Remove stale run_* dirs from prior (possibly longer) benchmarks. bench_summary +# reads every run_* dir, so leftovers from an earlier run with more MEASURED_RUNS +# would silently pollute this invocation's aggregate. +rm -rf "${UPSTREAM_ROOT:?}"/run_* "${NATIVE_ROOT:?}"/run_* + +UPSTREAM_REFINER_ARGS=() +NATIVE_REFINER_ARGS=() +if _is_true "${NO_REFINER}"; then + UPSTREAM_REFINER_ARGS+=(--no_refiner) + NATIVE_REFINER_ARGS+=(--no-refiner) +fi +UPSTREAM_BACKEND_ARGS=() +NATIVE_BACKEND_ARGS=() +if _is_true "${FORCE_CUDNN_SDPA}"; then + UPSTREAM_BACKEND_ARGS+=(--force_cudnn_sdpa) + NATIVE_BACKEND_ARGS+=(--force-cudnn-sdpa) +fi +UPSTREAM_COMPILE_ARGS=() +NATIVE_COMPILE_ARGS=() +if _is_true "${COMPILE_STAGE1}"; then + UPSTREAM_COMPILE_ARGS+=(--compile_stage1) + NATIVE_COMPILE_ARGS+=(--compile-stage1) +fi +UPSTREAM_PRECISION_ARGS=( + --stage1_precision "${STAGE1_PRECISION}" + --refiner_precision "${REFINER_PRECISION}" +) +NATIVE_PRECISION_ARGS=( + --stage1-precision "${STAGE1_PRECISION}" + --refiner-precision "${REFINER_PRECISION}" + --quant-backend "${QUANT_BACKEND}" +) + +TOTAL_RUNS=$(( WARMUP_RUNS + MEASURED_RUNS )) +for ((i = 0; i < TOTAL_RUNS; i++)); do + UPSTREAM_OUT="${UPSTREAM_ROOT}/run_${i}" + NATIVE_OUT="${NATIVE_ROOT}/run_${i}" + mkdir -p "${UPSTREAM_OUT}" "${NATIVE_OUT}" + + echo "[bench] upstream run ${i}/${TOTAL_RUNS}" + ( cd "${SCRIPT_DIR}" && \ + PYTHONPATH="${UPSTREAM_PYTHONPATH}" \ + uv run python "${SANA_REPO}/inference_video_scripts/wm/inference_sana_wm.py" \ + --image "${IMAGE_PATH}" \ + --prompt "${PROMPT_PATH}" \ + --camera "${CAMERA_PATH}" \ + --intrinsics "${INTRINSICS_PATH}" \ + --output_dir "${UPSTREAM_OUT}" \ + --name upstream \ + --num_frames "${NUM_FRAMES}" \ + --fps "${FPS}" \ + --step "${STEP}" \ + --cfg_scale "${CFG_SCALE}" \ + --seed "${SEED}" \ + --no_action_overlay \ + --stats_json "${UPSTREAM_OUT}/stats.json" \ + "${UPSTREAM_PRECISION_ARGS[@]}" \ + "${UPSTREAM_BACKEND_ARGS[@]}" \ + "${UPSTREAM_COMPILE_ARGS[@]}" \ + "${UPSTREAM_REFINER_ARGS[@]}" ) + + echo "[bench] FlashDreams run ${i}/${TOTAL_RUNS}" + ( cd "${SCRIPT_DIR}" && \ + uv run python "${SCRIPT_DIR}/run_native.py" \ + --image-path "${IMAGE_PATH}" \ + --prompt-path "${PROMPT_PATH}" \ + --camera-path "${CAMERA_PATH}" \ + --intrinsics-path "${INTRINSICS_PATH}" \ + --output-dir "${NATIVE_OUT}" \ + --name flashdreams \ + --num-frames "${NUM_FRAMES}" \ + --fps "${FPS}" \ + --step "${STEP}" \ + --cfg-scale "${CFG_SCALE}" \ + --seed "${SEED}" \ + --stats-json "${NATIVE_OUT}/stats.json" \ + "${NATIVE_PRECISION_ARGS[@]}" \ + "${NATIVE_BACKEND_ARGS[@]}" \ + "${NATIVE_COMPILE_ARGS[@]}" \ + "${NATIVE_REFINER_ARGS[@]}" ) +done + +SUMMARY_JSON="${OUTPUT_DIR}/bench.json" +SUMMARY_MD="${OUTPUT_DIR}/bench.md" +SUMMARY_CHART_MD="${OUTPUT_DIR}/perf.md" +SUMMARY_FLAGS=() +if _is_true "${NO_REFINER}"; then + SUMMARY_FLAGS+=(--no-refiner) +fi +if _is_true "${COMPILE_STAGE1}"; then + SUMMARY_FLAGS+=(--compile-stage1) +fi +if _is_true "${FORCE_CUDNN_SDPA}"; then + SUMMARY_FLAGS+=(--force-cudnn-sdpa) +fi +echo "[bench] summarising -> ${SUMMARY_MD}" +( cd "${SCRIPT_DIR}" && \ + uv run python "${SCRIPT_DIR}/bench_summary.py" \ + --upstream-dir "${UPSTREAM_ROOT}" \ + --flashdreams-dir "${NATIVE_ROOT}" \ + --warmup-runs "${WARMUP_RUNS}" \ + --image-path "${IMAGE_PATH}" \ + --prompt-path "${PROMPT_PATH}" \ + --camera-path "${CAMERA_PATH}" \ + --intrinsics-path "${INTRINSICS_PATH}" \ + --num-frames "${NUM_FRAMES}" \ + --seed "${SEED}" \ + --device-label "${DEVICE_LABEL}" \ + --chart-label "${CHART_LABEL}" \ + --stage1-precision "${STAGE1_PRECISION}" \ + --refiner-precision "${REFINER_PRECISION}" \ + --quant-backend "${QUANT_BACKEND}" \ + "${SUMMARY_FLAGS[@]}" \ + --output-json "${SUMMARY_JSON}" \ + --output-md "${SUMMARY_MD}" \ + --output-chart-md "${SUMMARY_CHART_MD}" ) + +echo "[bench] done." +echo " summary: ${SUMMARY_MD}" +echo " chart data: ${SUMMARY_CHART_MD}" diff --git a/integrations/sana/tests/parity_check/bench_summary.py b/integrations/sana/tests/parity_check/bench_summary.py new file mode 100644 index 000000000..fa76ed17f --- /dev/null +++ b/integrations/sana/tests/parity_check/bench_summary.py @@ -0,0 +1,355 @@ +# 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. + +"""Summarize SANA-WM upstream and FlashDreams benchmark stats.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path +from typing import Any + + +def _load_stats(root: Path) -> list[dict[str, Any]]: + paths = sorted(root.glob("run_*/stats.json")) + if not paths and (root / "stats.json").exists(): + paths = [root / "stats.json"] + return [json.loads(path.read_text(encoding="utf-8")) | {"_path": str(path)} for path in paths] + + +def _median(values: list[float]) -> float | None: + return statistics.median(values) if values else None + + +def _p90(values: list[float]) -> float | None: + if not values: + return None + ordered = sorted(values) + index = max(0, min(len(ordered) - 1, round(0.9 * (len(ordered) - 1)))) + return ordered[index] + + +def _upstream_stage_ms(item: dict[str, Any], key: str) -> float | None: + timings = item.get("timings_s") + if not isinstance(timings, dict): + return None + for candidate in key.split("|"): + value = timings.get(candidate) + if isinstance(value, (int, float)): + return float(value) * 1000.0 + return None + + +def _native_stage_ms(item: dict[str, Any], key: str) -> float | None: + stats = item.get("stats_ms") + if not isinstance(stats, dict): + return None + value = stats.get(key) + return float(value) if isinstance(value, (int, float)) else None + + +def _upstream_mem_gib(item: dict[str, Any]) -> float | None: + value = item.get("mem_peak_gib") + return float(value) if isinstance(value, (int, float)) else None + + +def _native_mem_gib(item: dict[str, Any]) -> float | None: + stats = item.get("stats_ms") + if not isinstance(stats, dict): + return None + value = stats.get("mem_peak_gib") + return float(value) if isinstance(value, (int, float)) else None + + +def _collect( + items: list[dict[str, Any]], + warmup_runs: int, + stage_reader, + mem_reader, + stages: dict[str, str], +) -> dict[str, Any]: + kept = items[warmup_runs:] + rows: dict[str, Any] = { + "runs_total": len(items), + "warmup_runs": warmup_runs, + "runs_measured": len(kept), + "paths": [item.get("_path") for item in kept], + } + wall = [ + float(item["wall_s"]) + for item in kept + if isinstance(item.get("wall_s"), (int, float)) + ] + rows["wall_median_s"] = _median(wall) + rows["wall_p90_s"] = _p90(wall) + memory = [ + value + for item in kept + if (value := mem_reader(item)) is not None + ] + rows["mem_peak_median_gib"] = _median(memory) + rows["mem_peak_p90_gib"] = _p90(memory) + for label, key in stages.items(): + values = [ + value + for item in kept + if (value := stage_reader(item, key)) is not None + ] + rows[f"{label}_median_ms"] = _median(values) + rows[f"{label}_p90_ms"] = _p90(values) + return rows + + +def _fmt(value: Any, suffix: str = "") -> str: + if value is None: + return "n/a" + if isinstance(value, float): + return f"{value:.2f}{suffix}" + return f"{value}{suffix}" + + +def _ms_per_frame(wall_s: float | None, num_frames: int) -> float | None: + if wall_s is None or num_frames <= 0: + return None + return wall_s * 1000.0 / num_frames + + +def _ms_per_clip(wall_s: float | None) -> float | None: + if wall_s is None: + return None + return wall_s * 1000.0 + + +def _metric_value(summary: dict[str, Any], side: str, key: str) -> float | None: + value = summary[side].get(key) + return float(value) if isinstance(value, (int, float)) else None + + +def _generation_ms_per_clip( + summary: dict[str, Any], + side: str, + *, + percentile: str = "median", +) -> float | None: + return _ms_per_clip(_metric_value(summary, side, f"wall_{percentile}_s")) + + +def _generation_ms_per_frame( + summary: dict[str, Any], + side: str, + *, + percentile: str = "median", +) -> float | None: + num_frames = int(summary["inputs"]["num_frames"]) + return _ms_per_frame( + _metric_value(summary, side, f"wall_{percentile}_s"), + num_frames, + ) + + +def _sum_optional(*values: float | None) -> float | None: + if any(value is None for value in values): + return None + return sum(float(value) for value in values) + + +def _render_markdown(summary: dict[str, Any]) -> str: + upstream = summary["upstream"] + native = summary["flashdreams"] + rows = [ + "# SANA-WM parity harness benchmark", + "", + "## Inputs", + "", + f"- image: `{summary['inputs']['image_path']}`", + f"- prompt: `{summary['inputs']['prompt_path']}`", + f"- camera: `{summary['inputs']['camera_path']}`", + f"- intrinsics: `{summary['inputs']['intrinsics_path']}`", + f"- num_frames: `{summary['inputs']['num_frames']}`", + f"- seed: `{summary['inputs']['seed']}`", + f"- no_refiner: `{summary['inputs']['no_refiner']}`", + f"- stage1_precision: `{summary['inputs']['stage1_precision']}`", + f"- refiner_precision: `{summary['inputs']['refiner_precision']}`", + f"- quant_backend: `{summary['inputs']['quant_backend']}`", + f"- compile_stage1: `{summary['inputs']['compile_stage1']}`", + f"- force_cudnn_sdpa: `{summary['inputs']['force_cudnn_sdpa']}`", + f"- warmup runs discarded: `{summary['inputs']['warmup_runs']}`", + "", + "## Benchmark metric", + "", + "The chart metric is post-load generation latency per generated clip.", + "Model construction, checkpoint loading, video writing, and frame dumps are outside this timing boundary.", + "SANA-WM renders each requested bidirectional clip in one generation pass, not as independently timed frames.", + "With the default `NO_REFINER=1`, the timed work covers conditioning, Stage-1 DiT, and SANA VAE decode.", + "", + "| metric | upstream | FlashDreams |", + "| --- | ---: | ---: |", + f"| measured runs | {upstream['runs_measured']} | {native['runs_measured']} |", + f"| generation median / clip | {_fmt(_generation_ms_per_clip(summary, 'upstream'), ' ms')} | {_fmt(_generation_ms_per_clip(summary, 'flashdreams'), ' ms')} |", + f"| generation p90 / clip | {_fmt(_generation_ms_per_clip(summary, 'upstream', percentile='p90'), ' ms')} | {_fmt(_generation_ms_per_clip(summary, 'flashdreams', percentile='p90'), ' ms')} |", + f"| generation median / frame, diagnostic | {_fmt(_generation_ms_per_frame(summary, 'upstream'), ' ms')} | {_fmt(_generation_ms_per_frame(summary, 'flashdreams'), ' ms')} |", + f"| generation p90 / frame, diagnostic | {_fmt(_generation_ms_per_frame(summary, 'upstream', percentile='p90'), ' ms')} | {_fmt(_generation_ms_per_frame(summary, 'flashdreams', percentile='p90'), ' ms')} |", + f"| wall median | {_fmt(upstream['wall_median_s'], ' s')} | {_fmt(native['wall_median_s'], ' s')} |", + f"| wall p90 | {_fmt(upstream['wall_p90_s'], ' s')} | {_fmt(native['wall_p90_s'], ' s')} |", + "", + "## Timing breakdown", + "", + "| stage | upstream median | FlashDreams median |", + "| --- | ---: | ---: |", + f"| Stage-1 incl. conditioning median | {_fmt(upstream['stage1_total_median_ms'], ' ms')} | {_fmt(_sum_optional(native['encode_median_ms'], native['dit_median_ms']), ' ms')} |", + f"| Stage-1 DiT median | {_fmt(upstream['dit_median_ms'], ' ms')} | {_fmt(native['dit_median_ms'], ' ms')} |", + f"| conditioning/encode median | n/a | {_fmt(native['encode_median_ms'], ' ms')} |", + ] + if summary["inputs"]["no_refiner"]: + # No refiner: FlashDreams `decode_ms` and upstream `vae_decode_s` are both + # the pure SANA VAE decode, so they compare directly. + rows.append( + f"| VAE decode median | {_fmt(upstream['vae_decode_median_ms'], ' ms')} | {_fmt(native['vae_decode_median_ms'], ' ms')} |" + ) + else: + # Refiner enabled: both sides bundle the refiner denoise together with its + # VAE decode into a single measurement (upstream `refiner_s`, FlashDreams + # `decode_ms`). They are only apples-to-apples as one combined row; the + # standalone upstream `vae_decode_s` is the Stage-1 decode and is NOT + # comparable to FlashDreams `decode_ms`. + rows.append( + f"| refiner + VAE decode median | {_fmt(upstream.get('refiner_median_ms'), ' ms')} | {_fmt(native['vae_decode_median_ms'], ' ms')} |" + ) + rows.append( + f"| peak memory median | {_fmt(upstream['mem_peak_median_gib'], ' GiB')} | {_fmt(native['mem_peak_median_gib'], ' GiB')} |" + ) + rows.append("") + return "\n".join(rows) + + +def _render_chart_markdown(summary: dict[str, Any], device_label: str) -> str: + official = _generation_ms_per_clip(summary, "upstream") + flashdreams = _generation_ms_per_clip(summary, "flashdreams") + if official is None or flashdreams is None: + raise ValueError("cannot render chart data without wall-clock stats") + return "\n".join( + [ + "# SANA-WM Benchmark Data (ms)", + "", + "| device | official | flashdreams |", + "| --- | ---: | ---: |", + f"| {device_label} | {official:.2f} | {flashdreams:.2f} |", + "", + ] + ) + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--upstream-dir", type=Path, required=True) + parser.add_argument("--flashdreams-dir", type=Path, required=True) + parser.add_argument("--warmup-runs", type=int, default=1) + parser.add_argument("--image-path", type=Path, required=True) + parser.add_argument("--prompt-path", type=Path, required=True) + parser.add_argument("--camera-path", type=Path, required=True) + parser.add_argument("--intrinsics-path", type=Path, required=True) + parser.add_argument("--num-frames", type=int, required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--no-refiner", action="store_true") + parser.add_argument("--stage1-precision", choices=["bf16", "fp8", "fp4"], default="bf16") + parser.add_argument("--refiner-precision", choices=["bf16", "fp8", "fp4"], default="bf16") + parser.add_argument( + "--quant-backend", + choices=["auto", "torch", "torch-fp8", "torch-fp4"], + default="auto", + ) + parser.add_argument("--compile-stage1", action="store_true") + parser.add_argument("--force-cudnn-sdpa", action="store_true") + parser.add_argument("--device-label", default="GPU") + parser.add_argument("--chart-label", default=None) + parser.add_argument("--output-json", type=Path, required=True) + parser.add_argument("--output-md", type=Path, required=True) + parser.add_argument("--output-chart-md", type=Path, default=None) + args = parser.parse_args(argv) + + upstream_items = _load_stats(args.upstream_dir) + native_items = _load_stats(args.flashdreams_dir) + summary = { + "inputs": { + "image_path": str(args.image_path), + "prompt_path": str(args.prompt_path), + "camera_path": str(args.camera_path), + "intrinsics_path": str(args.intrinsics_path), + "num_frames": args.num_frames, + "seed": args.seed, + "no_refiner": args.no_refiner, + "stage1_precision": args.stage1_precision, + "refiner_precision": args.refiner_precision, + "quant_backend": args.quant_backend, + "compile_stage1": args.compile_stage1, + "force_cudnn_sdpa": args.force_cudnn_sdpa, + "warmup_runs": args.warmup_runs, + "device_label": args.device_label, + "chart_label": args.chart_label or args.device_label, + }, + "upstream": _collect( + upstream_items, + args.warmup_runs, + _upstream_stage_ms, + _upstream_mem_gib, + { + "dit": "stage1_dit_s|stage1_sample_s", + "stage1_total": "stage1_sample_s", + "refiner": "refiner_s", + "vae_decode": "vae_decode_s", + }, + ), + "flashdreams": _collect( + native_items, + args.warmup_runs, + _native_stage_ms, + _native_mem_gib, + { + "encode": "encode_ms", + "dit": "diffuse_ms", + "vae_decode": "decode_ms", + }, + ), + } + summary["benchmark"] = { + "metric": "generation_ms_per_clip", + "unit": "ms", + "timing_boundary": ( + "pipeline.generate after model setup; excludes model construction, " + "checkpoint loading, video writing, and frame dumps" + ), + "device_label": args.device_label, + "chart_label": args.chart_label or args.device_label, + "official": _generation_ms_per_clip(summary, "upstream"), + "flashdreams": _generation_ms_per_clip(summary, "flashdreams"), + } + args.output_json.parent.mkdir(parents=True, exist_ok=True) + args.output_json.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") + args.output_md.parent.mkdir(parents=True, exist_ok=True) + report = _render_markdown(summary) + args.output_md.write_text(report, encoding="utf-8") + if args.output_chart_md is not None: + args.output_chart_md.parent.mkdir(parents=True, exist_ok=True) + args.output_chart_md.write_text( + _render_chart_markdown(summary, args.chart_label or args.device_label), + encoding="utf-8", + ) + print(report) + + +if __name__ == "__main__": + main() diff --git a/integrations/sana/tests/parity_check/bench_sweep_summary.py b/integrations/sana/tests/parity_check/bench_sweep_summary.py new file mode 100644 index 000000000..258b752ff --- /dev/null +++ b/integrations/sana/tests/parity_check/bench_sweep_summary.py @@ -0,0 +1,129 @@ +# 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. + +"""Aggregate SANA-WM precision benchmark summaries.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + + +def _load_item(raw: str) -> dict[str, Any]: + if ":" not in raw: + raise argparse.ArgumentTypeError( + f"expected LABEL:/path/to/bench.json, got {raw!r}" + ) + label, path_raw = raw.split(":", 1) + label = label.strip() + path = Path(path_raw) + if not label: + raise argparse.ArgumentTypeError(f"empty label in {raw!r}") + if not path.exists(): + raise argparse.ArgumentTypeError(f"bench summary does not exist: {path}") + summary = json.loads(path.read_text(encoding="utf-8")) + benchmark = summary.get("benchmark") + if not isinstance(benchmark, dict): + raise argparse.ArgumentTypeError(f"{path} has no benchmark object") + official = benchmark.get("official") + flashdreams = benchmark.get("flashdreams") + if not isinstance(official, (int, float)) or not isinstance( + flashdreams, (int, float) + ): + raise argparse.ArgumentTypeError( + f"{path} benchmark must contain numeric official and flashdreams values" + ) + return { + "label": label, + "path": str(path), + "official": float(official), + "flashdreams": float(flashdreams), + "inputs": summary.get("inputs", {}), + } + + +def _render_chart(rows: list[dict[str, Any]]) -> str: + lines = [ + "# SANA-WM Precision Sweep Summary (ms)", + "", + "| precision | official | flashdreams |", + "| --- | ---: | ---: |", + ] + for row in rows: + lines.append( + f"| {row['label']} | {row['official']:.2f} | {row['flashdreams']:.2f} |" + ) + lines.append("") + return "\n".join(lines) + + +def _render_report(rows: list[dict[str, Any]]) -> str: + lines = [ + "# SANA-WM precision benchmark sweep", + "", + "The chart metric is post-load generation latency per generated clip.", + "Model-card chart data is grouped by GPU/device in each precision subdirectory's `perf.md`.", + "", + "| precision | official | FlashDreams | source |", + "| --- | ---: | ---: | --- |", + ] + for row in rows: + lines.append( + f"| {row['label']} | {row['official']:.2f} ms | " + f"{row['flashdreams']:.2f} ms | `{row['path']}` |" + ) + lines.append("") + return "\n".join(lines) + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--item", + action="append", + type=_load_item, + required=True, + help="Precision row as LABEL:/path/to/bench.json. Repeat once per row.", + ) + parser.add_argument("--output-json", type=Path, required=True) + parser.add_argument("--output-md", type=Path, required=True) + parser.add_argument("--output-chart-md", type=Path, required=True) + args = parser.parse_args(argv) + + payload = { + "benchmark": { + "metric": "generation_ms_per_clip", + "unit": "ms", + "timing_boundary": ( + "pipeline.generate after model setup; excludes model construction, " + "checkpoint loading, video writing, and frame dumps" + ), + }, + "rows": args.item, + } + args.output_json.parent.mkdir(parents=True, exist_ok=True) + args.output_json.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + args.output_md.parent.mkdir(parents=True, exist_ok=True) + report = _render_report(args.item) + args.output_md.write_text(report, encoding="utf-8") + args.output_chart_md.parent.mkdir(parents=True, exist_ok=True) + args.output_chart_md.write_text(_render_chart(args.item), encoding="utf-8") + print(report) + + +if __name__ == "__main__": + main() diff --git a/integrations/sana/tests/parity_check/changes.patch b/integrations/sana/tests/parity_check/changes.patch new file mode 100644 index 000000000..10f349d4c --- /dev/null +++ b/integrations/sana/tests/parity_check/changes.patch @@ -0,0 +1,215 @@ +diff --git a/inference_video_scripts/wm/inference_sana_wm.py b/inference_video_scripts/wm/inference_sana_wm.py +index 5e07197..17db817 100644 +--- a/inference_video_scripts/wm/inference_sana_wm.py ++++ b/inference_video_scripts/wm/inference_sana_wm.py +@@ -937,6 +937,101 @@ def write_video(output_dir: Path, name: str, video_hwc: np.ndarray, fps: int, lo + return video_path + + ++def write_json(path: Path, payload: dict[str, object]) -> None: ++ path.parent.mkdir(parents=True, exist_ok=True) ++ path.write_text(json.dumps(payload, indent=2, default=str) + "\n", encoding="utf-8") ++ ++ ++def apply_backend_defaults() -> None: ++ torch.backends.cudnn.benchmark = True ++ torch.set_float32_matmul_precision("high") ++ if torch.cuda.is_available(): ++ torch.backends.cuda.enable_flash_sdp(False) ++ torch.backends.cuda.enable_math_sdp(True) ++ torch.backends.cuda.enable_mem_efficient_sdp(False) ++ torch.backends.cuda.enable_cudnn_sdp(True) ++ try: ++ import torch._inductor.config as inductor_config ++ ++ inductor_config.coordinate_descent_tuning = True ++ inductor_config.epilogue_fusion = True ++ except Exception: ++ pass ++ ++ ++def apply_precision_args(args: argparse.Namespace, logger: logging.Logger) -> None: ++ stage1_precision = args.stage1_precision ++ refiner_precision = "bf16" if args.no_refiner else args.refiner_precision ++ ++ if stage1_precision == "bf16": ++ os.environ.pop("SANA_WM_STAGE1_NVFP4", None) ++ os.environ.pop("SANA_WM_STAGE1_NVFP4_MODE", None) ++ os.environ.pop("SANA_WM_STAGE1_NVFP4_SCOPE", None) ++ os.environ.pop("SANA_WM_STAGE1_LINEARIZE_FFN", None) ++ os.environ.pop("SANA_WM_STAGE1_QUANT", None) ++ else: ++ os.environ["SANA_WM_STAGE1_NVFP4"] = "1" ++ os.environ["SANA_WM_STAGE1_NVFP4_MODE"] = "self_attn+cross+ffn" ++ os.environ["SANA_WM_STAGE1_NVFP4_SCOPE"] = "block" ++ os.environ["SANA_WM_STAGE1_LINEARIZE_FFN"] = "1" ++ os.environ["SANA_WM_STAGE1_QUANT"] = ( ++ "fp8block" if stage1_precision == "fp8" else "nvfp4" ++ ) ++ ++ if refiner_precision == "bf16": ++ os.environ.pop("SANA_WM_REFINER_NVFP4", None) ++ os.environ.pop("SANA_WM_REFINER_QUANT", None) ++ else: ++ os.environ["SANA_WM_REFINER_NVFP4"] = "1" ++ os.environ["SANA_WM_REFINER_QUANT"] = ( ++ "fp8block" if refiner_precision == "fp8" else "nvfp4" ++ ) ++ ++ if stage1_precision != "bf16" or refiner_precision != "bf16": ++ need = ( ++ "NVFP4BlockScaling" ++ if "fp4" in (stage1_precision, refiner_precision) ++ else "Float8BlockScaling" ++ ) ++ try: ++ import transformer_engine.common.recipe as te_recipe ++ import transformer_engine.pytorch # noqa: F401 ++ ++ if not hasattr(te_recipe, need): ++ raise ImportError(f"installed Transformer Engine lacks {need}") ++ except Exception as exc: ++ raise SystemExit( ++ "--stage1_precision/--refiner_precision fp8/fp4 require NVIDIA " ++ f"Transformer Engine >= 2.x for {need} ({exc})." ++ ) ++ ++ if "fp4" in (stage1_precision, refiner_precision) and torch.cuda.is_available(): ++ major, minor = torch.cuda.get_device_capability() ++ if major < 10: ++ logger.warning( ++ "fp4 (NVFP4) requires a Blackwell GPU (sm_100+); detected sm_%d%d.", ++ major, ++ minor, ++ ) ++ logger.info("[precision] stage1=%s refiner=%s", stage1_precision, refiner_precision) ++ ++ ++def time_bound_method(obj: object, name: str, timing_key: str, timings: dict[str, float]) -> None: ++ original = getattr(obj, name) ++ ++ def wrapped(*args, **kwargs): ++ if torch.cuda.is_available(): ++ torch.cuda.synchronize() ++ t0 = time.perf_counter() ++ result = original(*args, **kwargs) ++ if torch.cuda.is_available(): ++ torch.cuda.synchronize() ++ timings[timing_key] = time.perf_counter() - t0 ++ return result ++ ++ setattr(obj, name, wrapped) ++ ++ + # ============================================================================ + # Pipeline + # ============================================================================ +@@ -2132,6 +2227,30 @@ def _build_parser() -> argparse.ArgumentParser: + action="store_true", + help="Skip rendering the WASD + joystick overlay on the output video.", + ) ++ p.add_argument("--dump_frames", type=Path, default=None, help="Optional .npy path for decoded uint8 frames.") ++ p.add_argument("--stats_json", type=Path, default=None, help="Optional JSON path for timing and memory stats.") ++ p.add_argument( ++ "--force_cudnn_sdpa", ++ action="store_true", ++ help="Force PyTorch scaled_dot_product_attention to use the cuDNN backend.", ++ ) ++ p.add_argument( ++ "--compile_stage1", ++ action="store_true", ++ help="Wrap the Stage-1 DiT with torch.compile(mode='max-autotune-no-cudagraphs').", ++ ) ++ p.add_argument( ++ "--stage1_precision", ++ choices=["bf16", "fp8", "fp4"], ++ default="bf16", ++ help="Stage-1 DiT precision for benchmark parity.", ++ ) ++ p.add_argument( ++ "--refiner_precision", ++ choices=["bf16", "fp8", "fp4"], ++ default="bf16", ++ help="LTX-2 refiner precision for benchmark parity.", ++ ) + + # Weights and config. + p.add_argument( +@@ -2222,7 +2341,11 @@ def _snap_num_frames(n: int, stride: int = 8, *, upper_bound: int | None = None) + def main() -> None: + args = _build_parser().parse_args() + ++ if args.force_cudnn_sdpa: ++ apply_backend_defaults() ++ + logger = get_root_logger() ++ apply_precision_args(args, logger) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + image = Image.open(args.image).convert("RGB") +@@ -2276,6 +2399,15 @@ def main() -> None: + offload_refiner=args.offload_refiner, + logger=logger, + ) ++ if args.compile_stage1: ++ pipeline.model = torch.compile(pipeline.model, mode="max-autotune-no-cudagraphs") ++ ++ timings: dict[str, float] = {} ++ time_bound_method(pipeline, "_sample_stage1", "stage1_sample_s", timings) ++ time_bound_method(pipeline, "_dispatch_solver", "stage1_dit_s", timings) ++ time_bound_method(pipeline, "_decode_with_sana_vae", "vae_decode_s", timings) ++ if refiner is not None: ++ time_bound_method(pipeline, "_refine", "refiner_s", timings) + + denoising_step_list: list[int] | None = None + if args.denoising_step_list: +@@ -2308,14 +2440,48 @@ def main() -> None: + save_stage1=args.save_stage1, + ) + ++ if torch.cuda.is_available(): ++ torch.cuda.reset_peak_memory_stats(device) ++ torch.cuda.synchronize(device) ++ wall_start = time.perf_counter() + out = pipeline.generate(cropped, prompt, c2w, intrinsics_vec4, params) ++ if torch.cuda.is_available(): ++ torch.cuda.synchronize(device) ++ wall_s = time.perf_counter() - wall_start + video_hwc = out["video"] + + if not args.no_action_overlay: + logger.info("Compositing action overlay onto the output video.") + video_hwc = apply_overlay(video_hwc, out["c2w"]) + +- write_video(args.output_dir, args.name, video_hwc, params.fps, logger) ++ video_path = write_video(args.output_dir, args.name, video_hwc, params.fps, logger) ++ ++ if args.dump_frames is not None: ++ args.dump_frames.parent.mkdir(parents=True, exist_ok=True) ++ np.save(args.dump_frames, np.asarray(video_hwc, dtype=np.uint8)) ++ ++ if args.stats_json is not None: ++ stats: dict[str, object] = { ++ "backend": "upstream", ++ "entrypoint": "inference_video_scripts/wm/inference_sana_wm.py", ++ "video_path": str(video_path), ++ "video_shape": list(np.asarray(video_hwc).shape), ++ "timings_s": timings, ++ "num_frames": num_frames, ++ "seed": args.seed, ++ "fps": params.fps, ++ "step": params.step, ++ "cfg_scale": params.cfg_scale, ++ "no_refiner": args.no_refiner, ++ "stage1_precision": args.stage1_precision, ++ "refiner_precision": args.refiner_precision, ++ "compile_stage1": args.compile_stage1, ++ "force_cudnn_sdpa": args.force_cudnn_sdpa, ++ "wall_s": wall_s, ++ } ++ if torch.cuda.is_available(): ++ stats["mem_peak_gib"] = torch.cuda.max_memory_allocated(device) / (1024**3) ++ write_json(args.stats_json, stats) + + stage1_video = out.get("stage1_video") + if stage1_video is not None: diff --git a/integrations/sana/tests/parity_check/compat/mmcv/__init__.py b/integrations/sana/tests/parity_check/compat/mmcv/__init__.py new file mode 100644 index 000000000..9b2aebaae --- /dev/null +++ b/integrations/sana/tests/parity_check/compat/mmcv/__init__.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Minimal MMCV compatibility used by the SANA-WM parity harness. + +The upstream inference entrypoint only needs registry construction plus a few +utility functions at import time. This shim avoids installing old mmcv builds in +the isolated parity venv. +""" + +from __future__ import annotations + +import os +import pickle +from pathlib import Path +from typing import Any, Callable + + +class Registry: + """Small subset of ``mmcv.Registry`` used by SANA model registration.""" + + def __init__(self, name: str) -> None: + self.name = name + self.module_dict: dict[str, Any] = {} + + def __contains__(self, key: str) -> bool: + return key in self.module_dict + + def get(self, key: str) -> Any: + return self.module_dict.get(key) + + def register_module( + self, + module: Any | None = None, + *, + name: str | None = None, + force: bool = False, + ) -> Callable[[Any], Any] | Any: + """Register ``module`` or return a class decorator.""" + + def _register(obj: Any) -> Any: + key = name or obj.__name__ + if key in self.module_dict and not force: + raise KeyError(f"{key!r} is already registered in {self.name}.") + self.module_dict[key] = obj + return obj + + if module is not None: + return _register(module) + return _register + + def build( + self, + cfg: dict[str, Any], + *, + default_args: dict[str, Any] | None = None, + ) -> Any: + return build_from_cfg(cfg, self, default_args=default_args) + + +class Config(dict): + """Tiny dict-backed stand-in for import-time ``mmcv.Config`` references.""" + + @classmethod + def fromfile(cls, filename: str | os.PathLike[str]) -> "Config": + raise NotImplementedError( + "mmcv.Config.fromfile is not implemented in the parity harness shim." + ) + + +def build_from_cfg( + cfg: dict[str, Any], + registry: Registry, + *, + default_args: dict[str, Any] | None = None, +) -> Any: + """Instantiate a registered object from a config dict.""" + + if not isinstance(cfg, dict): + raise TypeError(f"cfg must be a dict, got {type(cfg).__name__}.") + if "type" not in cfg: + raise KeyError("cfg must contain key 'type'.") + args = dict(default_args or {}) + args.update({k: v for k, v in cfg.items() if k != "type"}) + obj_type = cfg["type"] + if isinstance(obj_type, str): + obj_cls = registry.get(obj_type) + if obj_cls is None: + raise KeyError(f"{obj_type!r} is not registered in {registry.name}.") + else: + obj_cls = obj_type + return obj_cls(**args) + + +def mkdir_or_exist(path: str | os.PathLike[str]) -> None: + Path(path).mkdir(parents=True, exist_ok=True) + + +def dump(obj: Any, file: str | os.PathLike[str]) -> None: + with open(file, "wb") as handle: + pickle.dump(obj, handle) + + +def load(file: str | os.PathLike[str]) -> Any: + with open(file, "rb") as handle: + return pickle.load(handle) + + +__all__ = [ + "Config", + "Registry", + "build_from_cfg", + "dump", + "load", + "mkdir_or_exist", +] diff --git a/integrations/sana/tests/parity_check/compat/mmcv/runner.py b/integrations/sana/tests/parity_check/compat/mmcv/runner.py new file mode 100644 index 000000000..3e82685a2 --- /dev/null +++ b/integrations/sana/tests/parity_check/compat/mmcv/runner.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Small ``mmcv.runner`` compatibility layer for SANA inference imports.""" + +from __future__ import annotations + +import torch.distributed as dist + + +def get_dist_info() -> tuple[int, int]: + if dist.is_available() and dist.is_initialized(): + return dist.get_rank(), dist.get_world_size() + return 0, 1 + + +class DefaultOptimizerConstructor: + pass + + +OPTIMIZERS = {} +OPTIMIZER_BUILDERS = {} + + +def build_optimizer(*_args, **_kwargs): + raise NotImplementedError("Optimizer construction is not available in the parity harness.") + + +__all__ = [ + "DefaultOptimizerConstructor", + "OPTIMIZER_BUILDERS", + "OPTIMIZERS", + "build_optimizer", + "get_dist_info", +] diff --git a/integrations/sana/tests/parity_check/compat/mmcv/utils/__init__.py b/integrations/sana/tests/parity_check/compat/mmcv/utils/__init__.py new file mode 100644 index 000000000..3acb82c06 --- /dev/null +++ b/integrations/sana/tests/parity_check/compat/mmcv/utils/__init__.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Minimal ``mmcv.utils`` namespace for SANA inference imports.""" + +from __future__ import annotations + +import torch.nn as nn + +_BatchNorm = nn.modules.batchnorm._BatchNorm +_InstanceNorm = nn.modules.instancenorm._InstanceNorm + +__all__ = ["_BatchNorm", "_InstanceNorm"] diff --git a/integrations/sana/tests/parity_check/compat/mmcv/utils/logging.py b/integrations/sana/tests/parity_check/compat/mmcv/utils/logging.py new file mode 100644 index 000000000..bb2396088 --- /dev/null +++ b/integrations/sana/tests/parity_check/compat/mmcv/utils/logging.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Logging globals expected by upstream SANA.""" + +logger_initialized: dict[str, bool] = {} + +__all__ = ["logger_initialized"] diff --git a/integrations/sana/tests/parity_check/diff_parity.py b/integrations/sana/tests/parity_check/diff_parity.py new file mode 100644 index 000000000..db9d395e2 --- /dev/null +++ b/integrations/sana/tests/parity_check/diff_parity.py @@ -0,0 +1,79 @@ +# 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. + +"""Compute SANA-WM frame parity as mean absolute uint8 delta.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import numpy as np + + +def _load_frames(path: Path) -> np.ndarray: + frames = np.load(path) + if frames.ndim != 4 or frames.shape[-1] != 3: + raise ValueError(f"{path} must contain [T,H,W,3] frames; got {frames.shape}.") + if frames.dtype != np.uint8: + frames = np.clip(frames, 0, 255).astype(np.uint8) + return frames + + +def _summary(upstream: np.ndarray, flashdreams: np.ndarray) -> dict[str, Any]: + if upstream.shape != flashdreams.shape: + raise ValueError( + "shape mismatch: " + f"upstream={tuple(upstream.shape)} flashdreams={tuple(flashdreams.shape)}" + ) + diff = np.abs(upstream.astype(np.int16) - flashdreams.astype(np.int16)) + per_frame = diff.reshape(diff.shape[0], -1).mean(axis=1) + return { + "shape": list(upstream.shape), + "mean_abs_delta": float(diff.mean()), + "mean_abs_delta_over_255": float(diff.mean() / 255.0), + "max_abs_delta": int(diff.max()), + "per_frame_mean_abs_delta": [float(v) for v in per_frame], + "per_frame_max_abs_delta": [ + int(v) for v in diff.reshape(diff.shape[0], -1).max(axis=1) + ], + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--upstream", type=Path, required=True) + parser.add_argument("--flashdreams", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + summary = _summary(_load_frames(args.upstream), _load_frames(args.flashdreams)) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") + + print( + "mean |Delta|: " + f"{summary['mean_abs_delta']:.4f} / 255 " + f"({summary['mean_abs_delta_over_255']:.6f})" + ) + print(f"max |Delta|: {summary['max_abs_delta']} / 255") + for idx, value in enumerate(summary["per_frame_mean_abs_delta"]): + print(f"frame {idx:04d}: mean |Delta| {value:.4f} / 255") + + +if __name__ == "__main__": + main() diff --git a/integrations/sana/tests/parity_check/pyproject.toml b/integrations/sana/tests/parity_check/pyproject.toml new file mode 100644 index 000000000..f92ac2738 --- /dev/null +++ b/integrations/sana/tests/parity_check/pyproject.toml @@ -0,0 +1,55 @@ +[project] +name = "sana-wm-parity-check" +version = "0.0.0" +description = "Isolated venv for the SANA-WM upstream parity and benchmark harness." +requires-python = ">=3.12,<3.13" +dependencies = [ + "flashdreams", + "flashdreams-sana-wm", + "accelerate>=1.3", + "diffusers>=0.37", + "einops>=0.7", + "flash-linear-attention>=0.4.2", + "ftfy>=6.0", + "huggingface-hub>=0.36", + "imageio[ffmpeg]>=2.31", + "numpy>=1.24,<2.5", + "omegaconf>=2.3", + "opencv-python-headless>=4.8", + "pillow>=10", + "protobuf>=7.35.0,<8", + "pyrallis>=0.3", + "pytz>=2024.0", + "pyyaml>=6.0", + "qwen-vl-utils>=0.0.8", + "safetensors>=0.5", + "scipy>=1.11", + "sentencepiece>=0.2", + "termcolor>=2.4", + "timm>=0.6.13", + "torch>=2.11", + "torchvision>=0.26", + "tqdm>=4.60", + "transformers>=5.0,<6", + "triton>=3.6; sys_platform == 'linux'", +] + +[project.optional-dependencies] +# Transformer Engine is only needed for the FP8/FP4 quant backends. It is gated +# behind an opt-in extra so the default `uv sync` (BF16 benchmarking) does not +# try to build/import it. Install with `uv sync --extra quant` once a +# transformer-engine-torch build matching the local torch/CUDA ABI is available. +quant = [ + "transformer-engine[pytorch,core-cu13]>=2.12; sys_platform != 'win32'", +] + +[tool.uv.sources] +flashdreams = { path = "../../../../flashdreams", editable = true } +flashdreams-sana-wm = { path = "../..", editable = true } + +[tool.uv] +managed = true +override-dependencies = [ + "nvidia-cublas>=13.4", +] +no-build-isolation-package = ["transformer-engine-torch"] diff --git a/integrations/sana/tests/parity_check/run.sh b/integrations/sana/tests/parity_check/run.sh new file mode 100644 index 000000000..a5460210d --- /dev/null +++ b/integrations/sana/tests/parity_check/run.sh @@ -0,0 +1,185 @@ +#!/usr/bin/env bash +# 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. + +# Patch pinned upstream SANA-WM with instrumentation, run it and FlashDreams on +# the same demo input, dump decoded uint8 frames, and compute mean |Delta| / 255. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" + +_abspath() { + case "$1" in + /*) printf '%s\n' "$1" ;; + *) printf '%s/%s\n' "${PWD}" "$1" ;; + esac +} + +SANA_REPO="$(_abspath "${SANA_REPO:-${HOME}/dev/Sana}")" +PATCH_FILE="${SCRIPT_DIR}/changes.patch" +REPO_URL="https://github.com/NVlabs/Sana.git" +PIN_COMMIT="6298508" + +OUTPUT_DIR="$(_abspath "${OUTPUT_DIR:-${SCRIPT_DIR}/outputs/parity}")" +UPSTREAM_OUT="${OUTPUT_DIR}/upstream" +NATIVE_OUT="${OUTPUT_DIR}/flashdreams" +IMAGE_PATH="$(_abspath "${IMAGE_PATH:-${SANA_REPO}/asset/sana_wm/demo_0.png}")" +PROMPT_PATH="$(_abspath "${PROMPT_PATH:-${SANA_REPO}/asset/sana_wm/demo_0.txt}")" +CAMERA_PATH="$(_abspath "${CAMERA_PATH:-${SANA_REPO}/asset/sana_wm/demo_0_pose.npy}")" +INTRINSICS_PATH="$(_abspath "${INTRINSICS_PATH:-${SANA_REPO}/asset/sana_wm/demo_0_intrinsics.npy}")" +NUM_FRAMES="${NUM_FRAMES:-121}" +FPS="${FPS:-16}" +STEP="${STEP:-60}" +CFG_SCALE="${CFG_SCALE:-5.0}" +SEED="${SEED:-42}" +NO_REFINER="${NO_REFINER:-1}" +FORCE_CUDNN_SDPA="${FORCE_CUDNN_SDPA:-0}" +COMPILE_STAGE1="${COMPILE_STAGE1:-0}" +STAGE1_PRECISION="${STAGE1_PRECISION:-bf16}" +REFINER_PRECISION="${REFINER_PRECISION:-${STAGE1_PRECISION}}" +QUANT_BACKEND="${QUANT_BACKEND:-auto}" + +_is_true() { + case "${1,,}" in + 1|true|yes|on) return 0 ;; + *) return 1 ;; + esac +} + +if [[ ! -d "${SANA_REPO}/.git" ]]; then + echo "[setup] cloning ${REPO_URL} -> ${SANA_REPO}" + git clone "${REPO_URL}" "${SANA_REPO}" +else + echo "[setup] repo already present at ${SANA_REPO}, skipping clone" +fi + +cd "${SANA_REPO}" +CURRENT_COMMIT="$(git rev-parse --short HEAD)" +if [[ "${CURRENT_COMMIT}" != "${PIN_COMMIT}" ]]; then + echo "[setup] checking out pinned commit ${PIN_COMMIT}" + git checkout "${PIN_COMMIT}" +else + echo "[setup] already at pinned commit ${PIN_COMMIT}, skipping checkout" +fi + +if git apply --reverse --check "${PATCH_FILE}" >/dev/null 2>&1; then + echo "[setup] patch already applied, skipping" +elif git apply --check "${PATCH_FILE}" >/dev/null 2>&1; then + echo "[setup] applying ${PATCH_FILE}" + git apply "${PATCH_FILE}" +else + echo "[setup] ERROR: ${PATCH_FILE} neither cleanly applies nor is already applied." >&2 + exit 1 +fi + +echo "[setup] ensuring Python deps via uv sync (isolated venv)" +( cd "${SCRIPT_DIR}" && uv sync ) +UPSTREAM_PYTHONPATH="${SCRIPT_DIR}/compat:${SANA_REPO}" +if [[ -n "${PYTHONPATH:-}" ]]; then + UPSTREAM_PYTHONPATH="${UPSTREAM_PYTHONPATH}:${PYTHONPATH}" +fi + +mkdir -p "${UPSTREAM_OUT}" "${NATIVE_OUT}" + +UPSTREAM_FRAMES="${UPSTREAM_OUT}/frames.npy" +NATIVE_FRAMES="${NATIVE_OUT}/frames.npy" +UPSTREAM_STATS="${UPSTREAM_OUT}/stats.json" +NATIVE_STATS="${NATIVE_OUT}/stats.json" + +UPSTREAM_REFINER_ARGS=() +NATIVE_REFINER_ARGS=() +if _is_true "${NO_REFINER}"; then + UPSTREAM_REFINER_ARGS+=(--no_refiner) + NATIVE_REFINER_ARGS+=(--no-refiner) +fi +UPSTREAM_BACKEND_ARGS=() +NATIVE_BACKEND_ARGS=() +if _is_true "${FORCE_CUDNN_SDPA}"; then + UPSTREAM_BACKEND_ARGS+=(--force_cudnn_sdpa) + NATIVE_BACKEND_ARGS+=(--force-cudnn-sdpa) +fi +UPSTREAM_COMPILE_ARGS=() +NATIVE_COMPILE_ARGS=() +if _is_true "${COMPILE_STAGE1}"; then + UPSTREAM_COMPILE_ARGS+=(--compile_stage1) + NATIVE_COMPILE_ARGS+=(--compile-stage1) +fi +UPSTREAM_PRECISION_ARGS=( + --stage1_precision "${STAGE1_PRECISION}" + --refiner_precision "${REFINER_PRECISION}" +) +NATIVE_PRECISION_ARGS=( + --stage1-precision "${STAGE1_PRECISION}" + --refiner-precision "${REFINER_PRECISION}" + --quant-backend "${QUANT_BACKEND}" +) + +echo "[run] upstream SANA-WM -> ${UPSTREAM_OUT}" +( cd "${SCRIPT_DIR}" && \ + PYTHONPATH="${UPSTREAM_PYTHONPATH}" \ + uv run python "${SANA_REPO}/inference_video_scripts/wm/inference_sana_wm.py" \ + --image "${IMAGE_PATH}" \ + --prompt "${PROMPT_PATH}" \ + --camera "${CAMERA_PATH}" \ + --intrinsics "${INTRINSICS_PATH}" \ + --output_dir "${UPSTREAM_OUT}" \ + --name upstream \ + --num_frames "${NUM_FRAMES}" \ + --fps "${FPS}" \ + --step "${STEP}" \ + --cfg_scale "${CFG_SCALE}" \ + --seed "${SEED}" \ + --no_action_overlay \ + --dump_frames "${UPSTREAM_FRAMES}" \ + --stats_json "${UPSTREAM_STATS}" \ + "${UPSTREAM_PRECISION_ARGS[@]}" \ + "${UPSTREAM_BACKEND_ARGS[@]}" \ + "${UPSTREAM_COMPILE_ARGS[@]}" \ + "${UPSTREAM_REFINER_ARGS[@]}" ) + +echo "[run] FlashDreams SANA-WM -> ${NATIVE_OUT}" +( cd "${SCRIPT_DIR}" && \ + uv run python "${SCRIPT_DIR}/run_native.py" \ + --image-path "${IMAGE_PATH}" \ + --prompt-path "${PROMPT_PATH}" \ + --camera-path "${CAMERA_PATH}" \ + --intrinsics-path "${INTRINSICS_PATH}" \ + --output-dir "${NATIVE_OUT}" \ + --name flashdreams \ + --num-frames "${NUM_FRAMES}" \ + --fps "${FPS}" \ + --step "${STEP}" \ + --cfg-scale "${CFG_SCALE}" \ + --seed "${SEED}" \ + --dump-frames "${NATIVE_FRAMES}" \ + --stats-json "${NATIVE_STATS}" \ + "${NATIVE_PRECISION_ARGS[@]}" \ + "${NATIVE_BACKEND_ARGS[@]}" \ + "${NATIVE_COMPILE_ARGS[@]}" \ + "${NATIVE_REFINER_ARGS[@]}" ) + +echo "[diff] summarising parity -> ${OUTPUT_DIR}/parity.json" +( cd "${SCRIPT_DIR}" && \ + uv run python "${SCRIPT_DIR}/diff_parity.py" \ + --upstream "${UPSTREAM_FRAMES}" \ + --flashdreams "${NATIVE_FRAMES}" \ + --output "${OUTPUT_DIR}/parity.json" ) + +echo "[run] done." +echo " upstream frames : ${UPSTREAM_FRAMES}" +echo " flashdreams frames: ${NATIVE_FRAMES}" +echo " parity JSON : ${OUTPUT_DIR}/parity.json" diff --git a/integrations/sana/tests/parity_check/run_native.py b/integrations/sana/tests/parity_check/run_native.py new file mode 100644 index 000000000..d49601ab7 --- /dev/null +++ b/integrations/sana/tests/parity_check/run_native.py @@ -0,0 +1,322 @@ +# 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. + +"""Run the FlashDreams SANA-WM pipeline with frame dumps and profiler JSON.""" + +from __future__ import annotations + +import argparse +from contextlib import nullcontext +import json +import os +import time +from pathlib import Path +from typing import Any + +import imageio.v3 as iio +import numpy as np +import torch +from PIL import Image + +from flashdreams.infra.config import derive_config +from sana_wm.camera import ( + default_intrinsics_vec4, + load_intrinsics, + resize_center_crop_geometry, + snap_num_frames, + transform_intrinsics_for_crop, +) +from sana_wm.conditioning import SanaWMI2VConditioningRequest +from sana_wm.config import RUNNER_SANA_WM_BIDIRECTIONAL +from sana_wm.constants import DEFAULT_VIDEO_HEIGHT, DEFAULT_VIDEO_WIDTH +from sana_wm.decoder import SanaWMDecodedVideo +from sana_wm.runner import ( + _active_quantized_precisions, + _pipeline_config, + _resolve_quant_backend, + _validate_precision_request, +) + + +def _apply_backend_defaults() -> None: + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + if torch.cuda.is_available(): + torch.backends.cuda.enable_flash_sdp(False) + torch.backends.cuda.enable_math_sdp(True) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_cudnn_sdp(True) + try: + import torch._inductor.config as inductor_config + + inductor_config.coordinate_descent_tuning = True + inductor_config.epilogue_fusion = True + except Exception: + pass + + +def _resolve_device(device: str) -> torch.device: + if device == "auto": + if torch.cuda.is_available(): + return torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}") + return torch.device("cpu") + if device == "cuda" and torch.cuda.is_available(): + return torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}") + return torch.device(device) + + +def _prepare_inputs( + *, + image_path: Path, + camera_path: Path, + intrinsics_path: Path | None, + intrinsics_hfov_deg: float, + num_frames_requested: int, +) -> tuple[Image.Image, np.ndarray, np.ndarray, int]: + image = Image.open(image_path).convert("RGB") + c2w_full = np.load(camera_path).astype(np.float32) + if c2w_full.ndim != 3 or c2w_full.shape[1:] != (4, 4): + raise ValueError(f"--camera-path must be [F,4,4]; got {c2w_full.shape}.") + + num_frames = min(num_frames_requested, c2w_full.shape[0]) + num_frames = snap_num_frames(num_frames, stride=8, upper_bound=c2w_full.shape[0]) + c2w = c2w_full[:num_frames] + + resized_size, crop_offset = resize_center_crop_geometry( + image.size, + target_h=DEFAULT_VIDEO_HEIGHT, + target_w=DEFAULT_VIDEO_WIDTH, + ) + resized = image.resize(resized_size, Image.LANCZOS) + left, top = crop_offset + cropped = resized.crop( + (left, top, left + DEFAULT_VIDEO_WIDTH, top + DEFAULT_VIDEO_HEIGHT) + ) + if intrinsics_path is None: + intrinsics_src = default_intrinsics_vec4( + image.size, + num_frames, + hfov_deg=intrinsics_hfov_deg, + ) + else: + intrinsics_src = load_intrinsics(intrinsics_path, num_frames) + intrinsics_vec4 = transform_intrinsics_for_crop( + intrinsics_src, + image.size, + resized_size, + crop_offset, + ) + return cropped, c2w, intrinsics_vec4, num_frames + + +def _install_stage1_compile_hook(pipeline: Any) -> None: + transformer = pipeline.diffusion_model.transformer + original_ensure_model = transformer._ensure_model + compiled = {"done": False} + + def _ensure_model_with_compile() -> None: + original_ensure_model() + if compiled["done"]: + return + transformer.model = torch.compile( + transformer.model, + mode="max-autotune-no-cudagraphs", + ) + compiled["done"] = True + + transformer._ensure_model = _ensure_model_with_compile + + +def _write_video(output_dir: Path, name: str, frames: np.ndarray, fps: int) -> Path: + output_dir.mkdir(parents=True, exist_ok=True) + path = output_dir / f"{name}_generated.mp4" + iio.imwrite(path, frames, fps=fps) + return path + + +def _write_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2, default=str) + "\n", encoding="utf-8") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--image-path", type=Path, required=True) + parser.add_argument("--prompt-path", type=Path, required=True) + parser.add_argument("--camera-path", type=Path, required=True) + parser.add_argument("--intrinsics-path", type=Path, default=None) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--name", default="flashdreams") + parser.add_argument("--dump-frames", type=Path, default=None) + parser.add_argument("--stats-json", type=Path, default=None) + parser.add_argument("--num-frames", type=int, default=161) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--step", type=int, default=60) + parser.add_argument("--cfg-scale", type=float, default=5.0) + parser.add_argument("--flow-shift", type=float, default=None) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--negative-prompt", default="") + parser.add_argument("--no-refiner", action="store_true") + parser.add_argument("--save-stage1", action="store_true") + parser.add_argument("--device", default="auto") + parser.add_argument("--intrinsics-hfov-deg", type=float, default=90.0) + parser.add_argument("--stage1-precision", choices=["bf16", "fp8", "fp4"], default="bf16") + parser.add_argument("--refiner-precision", choices=["bf16", "fp8", "fp4"], default="bf16") + parser.add_argument( + "--quant-backend", + choices=["auto", "torch", "torch-fp8", "torch-fp4"], + default="auto", + ) + parser.add_argument("--force-cudnn-sdpa", action="store_true") + parser.add_argument("--compile-stage1", action="store_true") + args = parser.parse_args() + + if args.force_cudnn_sdpa: + _apply_backend_defaults() + + prompt = args.prompt_path.read_text(encoding="utf-8", errors="replace").strip() + if not prompt: + raise ValueError(f"Prompt file is empty: {args.prompt_path}") + + device = _resolve_device(args.device) + quantized = _active_quantized_precisions( + stage1_precision=args.stage1_precision, + refiner_precision=args.refiner_precision, + refiner_enabled=not args.no_refiner, + ) + quant_backend = _resolve_quant_backend(args.quant_backend, quantized) + _validate_precision_request( + device=device, + stage1_precision=args.stage1_precision, + refiner_precision=args.refiner_precision, + refiner_enabled=not args.no_refiner, + quant_backend=args.quant_backend, + ) + with nullcontext(): + image, c2w, intrinsics_vec4, num_frames = _prepare_inputs( + image_path=args.image_path, + camera_path=args.camera_path, + intrinsics_path=args.intrinsics_path, + intrinsics_hfov_deg=args.intrinsics_hfov_deg, + num_frames_requested=args.num_frames, + ) + runner_cfg = derive_config( + RUNNER_SANA_WM_BIDIRECTIONAL, + output_dir=args.output_dir, + runner_name=args.name, + num_frames=num_frames, + fps=args.fps, + step=args.step, + cfg_scale=args.cfg_scale, + flow_shift=args.flow_shift, + seed=args.seed, + negative_prompt=args.negative_prompt, + no_refiner=args.no_refiner, + save_stage1=args.save_stage1, + refiner_seed=args.seed, + stage1_precision=args.stage1_precision, + refiner_precision=args.refiner_precision, + quant_backend=args.quant_backend, + ) + pipeline_cfg = _pipeline_config(runner_cfg, quant_backend=quant_backend) + pipeline_cfg = derive_config(pipeline_cfg, enable_sync_and_profile=True) + pipeline = pipeline_cfg.setup().to(device).eval() + if args.compile_stage1: + _install_stage1_compile_hook(pipeline) + encoder = getattr(pipeline, "encoder", None) + if encoder is not None: + text_encoder = getattr(encoder, "text_encoder", None) + if text_encoder is not None: + text_encoder._ensure_text_encoder() + first_frame_encoder = getattr(encoder, "first_frame_encoder", None) + if first_frame_encoder is not None: + first_frame_encoder._ensure_vae() + pipeline.diffusion_model.transformer._ensure_model() + decoder = getattr(pipeline, "decoder", None) + vae_decoder = getattr(decoder, "vae_decoder", None) + if vae_decoder is not None: + vae_decoder._ensure_vae() + + request = SanaWMI2VConditioningRequest( + image=image, + prompt=prompt, + poses_c2w=c2w, + intrinsics_vec4=intrinsics_vec4, + num_frames=num_frames, + fps=args.fps, + steps=args.step, + cfg_scale=args.cfg_scale, + flow_shift=args.flow_shift, + seed=args.seed, + negative_prompt=args.negative_prompt, + ) + + cache = pipeline.initialize_cache( + decoder_context={ + "prompt": prompt, + "fps": args.fps, + "save_stage1": args.save_stage1, + "refiner_seed": args.seed, + "sink_size": runner_cfg.sink_size, + } + ) + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats(device) + torch.cuda.synchronize(device) + generation_start = time.perf_counter() + with torch.inference_mode(): + decoded = pipeline.generate( + 0, + cache, + input=request, + ) + if torch.cuda.is_available(): + torch.cuda.synchronize(device) + wall_s = time.perf_counter() - generation_start + stats = pipeline.finalize(0, cache) or {} + + if not isinstance(decoded, SanaWMDecodedVideo): + raise TypeError(f"expected SanaWMDecodedVideo, got {type(decoded).__name__}") + + frames = np.asarray(decoded.video_hwc, dtype=np.uint8) + video_path = _write_video(args.output_dir, args.name, frames, args.fps) + if args.dump_frames is not None: + args.dump_frames.parent.mkdir(parents=True, exist_ok=True) + np.save(args.dump_frames, frames) + if args.stats_json is not None: + _write_json( + args.stats_json, + { + "backend": "flashdreams", + "runner": "sana-wm-bidirectional", + "video_path": str(video_path), + "video_shape": list(frames.shape), + "num_frames": num_frames, + "seed": args.seed, + "fps": args.fps, + "step": args.step, + "cfg_scale": args.cfg_scale, + "no_refiner": args.no_refiner, + "compile_stage1": args.compile_stage1, + "force_cudnn_sdpa": args.force_cudnn_sdpa, + "wall_s": wall_s, + "stats_ms": stats, + }, + ) + print(f"[native] wrote {video_path}") + + +if __name__ == "__main__": + main() diff --git a/integrations/sana/tests/parity_check/uv.lock b/integrations/sana/tests/parity_check/uv.lock new file mode 100644 index 000000000..81c9f57a0 --- /dev/null +++ b/integrations/sana/tests/parity_check/uv.lock @@ -0,0 +1,1681 @@ +version = 1 +revision = 3 +requires-python = "==3.12.*" +resolution-markers = [ + "sys_platform == 'win32'", + "sys_platform != 'win32'", +] + +[manifest] +overrides = [{ name = "nvidia-cublas", specifier = ">=13.4" }] + +[[package]] +name = "accelerate" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "psutil" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, + { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8d/75/94cd5d389649578aca399e5aa822637eec18319a1dadc400ffe2f9a7493f/accelerate-1.14.0.tar.gz", hash = "sha256:41b9c4377a54e0b460a959b0defa1b736e4ca0a2373252d9a539964c2afe3c8d", size = 412167, upload-time = "2026-06-11T13:45:52.326Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/db/253133d7e7cb40d3af384bb2f5c0b4a2b7fdcffbc95c688cc67a20a3c103/accelerate-1.14.0-py3-none-any.whl", hash = "sha256:e94390c2863b873be18f623f9df48a0d8fe5eff13ea7f1a00092b0a7904888c6", size = 389246, upload-time = "2026-06-11T13:45:50.477Z" }, +] + +[[package]] +name = "annotated-doc" +version = "0.0.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/57/ba/046ceea27344560984e26a590f90bc7f4a75b06701f653222458922b558c/annotated_doc-0.0.4.tar.gz", hash = "sha256:fbcda96e87e9c92ad167c2e53839e57503ecfda18804ea28102353485033faa4", size = 7288, upload-time = "2025-11-10T22:07:42.062Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/d3/26bf1008eb3d2daa8ef4cacc7f3bfdc11818d111f7e2d0201bc6e3b49d45/annotated_doc-0.0.4-py3-none-any.whl", hash = "sha256:571ac1dc6991c450b25a9c2d84a3705e2ae7a53467b5d111c24fa8baabbed320", size = 5303, upload-time = "2025-11-10T22:07:40.673Z" }, +] + +[[package]] +name = "annotated-types" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5f/56/a8120250d128bed162cd73c76d45f6ef9991f3e068f62a8ee060afa3104a/annotated_types-0.8.0.tar.gz", hash = "sha256:13b2beaad985e05e2d6407ee4c4f35590b11f8d693a258a561055cac8f64cab7", size = 15893, upload-time = "2026-07-23T20:16:13.995Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/91/8acff4f5e50511b911bbccb72b8628a49c68ce14148cd9f6431094859a90/annotated_types-0.8.0-py3-none-any.whl", hash = "sha256:f072f4d804ea359e4eaf198b1af7a8b0943881a87f31bb764f8bf219bb9419e0", size = 13427, upload-time = "2026-07-23T20:16:12.938Z" }, +] + +[[package]] +name = "antlr4-python3-runtime" +version = "4.9.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" } + +[[package]] +name = "anyio" +version = "4.14.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "idna" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/61/cc/a381afa6efea9f496eff839d4a6a1aed3bfafc7b3ab4b0d1b243a12573dd/anyio-4.14.2.tar.gz", hash = "sha256:cfa139f3ed1a23ee8f88a145ddb5ac7605b8bbfd8592baacd7ce3d8bb4313c7f", size = 260176, upload-time = "2026-07-12T20:29:07.082Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/35/f2287558c17e29fafc8ef3daf819bb9834061cfa43bff8014f7df7f63bdc/anyio-4.14.2-py3-none-any.whl", hash = "sha256:9f505dda5ac9f0c8309b5e8bd445a8c2bf7246f3ce950121e45ea15bc41d1494", size = 125813, upload-time = "2026-07-12T20:29:05.763Z" }, +] + +[[package]] +name = "av" +version = "18.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ae/a4/570a5a35c8638aba01e739925846c35fdd6b0756a15526766d0a4dd3b7df/av-18.0.0.tar.gz", hash = "sha256:4ef7e72c3d3a872584a1215173b16e0226811037f40dcdbf75992631098df1ba", size = 4340222, upload-time = "2026-07-02T06:37:58.907Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/4a/9e3463df030e063d757fa12f0f39be6541b45b06b5bad48c2ce361b924bf/av-18.0.0-cp311-abi3-macosx_11_0_x86_64.whl", hash = "sha256:149289d40e732a6e49c9530bc245b49d9964cfd1c8c9e06778703b7d5bba6b25", size = 22499354, upload-time = "2026-07-02T06:36:58.751Z" }, + { url = "https://files.pythonhosted.org/packages/77/b3/2576a44b4f39c7462ced4c17fec04c756f7b0f3c5cb940d124173e417d6a/av-18.0.0-cp311-abi3-macosx_14_0_arm64.whl", hash = "sha256:35274c20d2ad3b4774fe632bcef2e34af79858ddf899352339cc3babbc13a484", size = 18175248, upload-time = "2026-07-02T06:37:01.741Z" }, + { url = "https://files.pythonhosted.org/packages/84/74/6732f17b96dc23fd23b876b2805435855abdc8a3b397142be4e581165de8/av-18.0.0-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:4d683b7747a0ba9222b8a5f81e41db5f796e7f64473454ec4fe2548e083c2fa0", size = 33387843, upload-time = "2026-07-02T06:37:05.097Z" }, + { url = "https://files.pythonhosted.org/packages/6d/b9/7708c43fed7ae28b4a1bad060b4221e3334cd827cec24f7165902a6ac1f4/av-18.0.0-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ae56b40b6f8b067a8ad2dac664fbfbabac7f7a55b9a7bb031eb99289252bc017", size = 35536910, upload-time = "2026-07-02T06:37:08.806Z" }, + { url = "https://files.pythonhosted.org/packages/5a/94/eba99691d184f6a395a242d54dc370e2fd2265e95bbc98e2963a0fdbdd6c/av-18.0.0-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:ea2e8ebbce521f21b55df9400e00d721623c9020ef158f5a188a96130be0743f", size = 38984619, upload-time = "2026-07-02T06:37:11.861Z" }, + { url = "https://files.pythonhosted.org/packages/c9/cf/0d7aee07fe16aa9ffdf96043c14bed5485a52c0dea4259de87aa306ecab4/av-18.0.0-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ef96dabb3e50dac249913145dff5424b302b257fd95dcb64be3c7b7a8aef16d1", size = 34451176, upload-time = "2026-07-02T06:37:15.154Z" }, + { url = "https://files.pythonhosted.org/packages/76/92/810da80b12680d4c4fe235bd1b4003289be9213ac7f114b77b8ecf0e3b3e/av-18.0.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:0f65518a184613e41536f29e8758c8e3d8293e46bf5bef108f04f925bbfa3f44", size = 36619869, upload-time = "2026-07-02T06:37:18.495Z" }, + { url = "https://files.pythonhosted.org/packages/11/85/0f121ff43dc5a70696676c98a8f1674e2fa787614c2abaacb15fa1a9bc99/av-18.0.0-cp311-abi3-win_amd64.whl", hash = "sha256:aaf4d354d2beaa6651e4f92e54409a578bde64f79c0beef9a30b388d06f7c629", size = 27556236, upload-time = "2026-07-02T06:37:21.388Z" }, + { url = "https://files.pythonhosted.org/packages/8b/f6/2509754d4d2356abc6fc0ea3d57c12ade29bac23a1fb7fc215a53ca518fb/av-18.0.0-cp311-abi3-win_arm64.whl", hash = "sha256:adac2b3833b6cb9bd6cb52664a522b94db453615b3675b1dbb26e13fe1c80da6", size = 20221133, upload-time = "2026-07-02T06:37:23.88Z" }, +] + +[[package]] +name = "boto3" +version = "1.43.53" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, + { name = "jmespath" }, + { name = "s3transfer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6c/da/0e90eab875f2eb4b8708fef2c198f2559ed1e451a1016f1cd4fcdcfbfbe3/boto3-1.43.53.tar.gz", hash = "sha256:c80425acab314d7af09609562053f565139e1fe49108eacfcc1601ebfaee235b", size = 112678, upload-time = "2026-07-21T19:28:53.002Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/45/27/7e72d25fdde77668b7bd4fa47381192dd2aa64fb77265e4bab786fd9fe2a/boto3-1.43.53-py3-none-any.whl", hash = "sha256:5383e705d8a976a14f23bb8c113c07a396931a019db98fce4cdc68650ec6e4d8", size = 140025, upload-time = "2026-07-21T19:28:50.957Z" }, +] + +[[package]] +name = "botocore" +version = "1.43.53" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jmespath" }, + { name = "python-dateutil" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/db/6b/ebcefacc4de3cd4f1c449540d86877f76c2f5e586a620831012decbb2b2c/botocore-1.43.53.tar.gz", hash = "sha256:36d93dd8db68ee75f6b61ca9f775161b8168844e4601698701530e6efdded141", size = 15720336, upload-time = "2026-07-21T19:28:41.547Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/e5/1b60e394f0fff97ee70dd16913382b7dffda85b98e639c0c9e8ff56cfaa7/botocore-1.43.53-py3-none-any.whl", hash = "sha256:b7ee9a70d187e5348883c820990ccd9436ab14e2bd6622741fc96fe561e816b8", size = 15404628, upload-time = "2026-07-21T19:28:37.475Z" }, +] + +[[package]] +name = "certifi" +version = "2026.7.22" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a3/c2/24167ea9858356b47a87a50d39908bfdb72ceeefe0041586e704e5376b3a/certifi-2026.7.22.tar.gz", hash = "sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55", size = 138112, upload-time = "2026-07-22T03:35:12.644Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/a7/71ac2cff56fec219ed242bb11b8efb69fcc4bec75db06fb7bfe35de520e6/certifi-2026.7.22-py3-none-any.whl", hash = "sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775", size = 136983, upload-time = "2026-07-22T03:35:11.276Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.4.9" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bd/2a/23f34ec9d04624958e137efdc394888716353190e75f25dd22c7a2c7a8aa/charset_normalizer-3.4.9.tar.gz", hash = "sha256:673611bbd43f0810bec0b0f028ddeaaa501190339cac411f347ac76917c3ae7b", size = 152439, upload-time = "2026-07-07T14:34:58.454Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/70/4a/ecbd131485c07fcdfad54e28946d513e3da22ef3b4bd854dcafae54ec739/charset_normalizer-3.4.9-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:45b0cc4e3556cd875e09102988d1ab8356c998b596c9fced84547c8138b487a0", size = 319300, upload-time = "2026-07-07T14:33:15.666Z" }, + { url = "https://files.pythonhosted.org/packages/ec/96/5d9364e3342d69f3a045e1777bc47c85c383e6e9466d561b33fdb419d1f9/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9b2aff1c7b3884512b9512c3eaadd9bab39fb45042ffaaa1dd08ff2b9f8109d9", size = 215802, upload-time = "2026-07-07T14:33:17.031Z" }, + { url = "https://files.pythonhosted.org/packages/4b/4c/5361f9aa7f2cb58d94f2ab831b3d493f69efb1d239654b4744e3c09527cb/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9104ed0bd76a429d46f9ec0dbc9b08ad1d2dcdf2b00a5a0daa1c145329b35b44", size = 237171, upload-time = "2026-07-07T14:33:18.576Z" }, + { url = "https://files.pythonhosted.org/packages/50/78/ce342ca4ff30b2eb49fe6d9578df85974f90c67d294113e94efdd9664cbd/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:7b86a2b16095d250c6f58b3d9b2eee6f4147754344f3dab0922f7c9bf7d226c9", size = 233075, upload-time = "2026-07-07T14:33:20.084Z" }, + { url = "https://files.pythonhosted.org/packages/01/c4/4fa4c8b3097a11f3c5f09a35b72ed6855fb1d332469504962ab7bafcc702/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5e226f6218febc71f6c1fc2fafb91c226f75bdc1d8fb12d66823716e891608fd", size = 224256, upload-time = "2026-07-07T14:33:21.747Z" }, + { url = "https://files.pythonhosted.org/packages/87/3a/ad914516df7e358a81aae018caa5e0470ba827fa6d763b1d2e87d920a5f6/charset_normalizer-3.4.9-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:90c44bc373b7687f6948b693cceaea1348ae0975d7474746559494468e3c1d84", size = 208784, upload-time = "2026-07-07T14:33:23.313Z" }, + { url = "https://files.pythonhosted.org/packages/d7/74/3c12f9755717dfe5c5c87da63f35d765fa0c00382ec26bf23f7fae34f2ba/charset_normalizer-3.4.9-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:9cdef90ae47919cae358d8ab15797a800ed41da7aba5d72419fb510729e2ed4b", size = 219928, upload-time = "2026-07-07T14:33:24.814Z" }, + { url = "https://files.pythonhosted.org/packages/33/9a/895095b83e7907abd6d3d99aad3a38ad0d9686cc186cb0c94c24320fe63e/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:60f44ade2cf573dad7a277e6f8ca9a51a21dda572b13bd7d8539bb3cd5dbedde", size = 218489, upload-time = "2026-07-07T14:33:26.42Z" }, + { url = "https://files.pythonhosted.org/packages/a1/34/ef5c05f412f42520d7709b7d3784d19640839eb7366ded1755511585429f/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:a1786910334ed46ab1dd73222f2cd1e05c2c3bb39f6dddb4f8b36fc382058a39", size = 210267, upload-time = "2026-07-07T14:33:27.952Z" }, + { url = "https://files.pythonhosted.org/packages/83/dc/9b29fa4412b318bf3bfea985c35d67eb55e04b59a7c3f2237168b0e0be6f/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:03d07803992c6c7bbc976327f34b18b6160327fc81cb82c9d504720ac0be3b62", size = 226030, upload-time = "2026-07-07T14:33:29.397Z" }, + { url = "https://files.pythonhosted.org/packages/0e/42/6dbc00b8cd16011691203e33570fa42ed5746599a2e878112d16eab403a3/charset_normalizer-3.4.9-cp312-cp312-win32.whl", hash = "sha256:78841cccf1af7b40f6f716338d50c0902dbe88d9f800b3c973b7a9a0a693a642", size = 151185, upload-time = "2026-07-07T14:33:30.781Z" }, + { url = "https://files.pythonhosted.org/packages/80/cc/f920afd1a23c58ccd53c1d36085a71893a4737ff5e66e0371efab6809850/charset_normalizer-3.4.9-cp312-cp312-win_amd64.whl", hash = "sha256:4b3dac63058cc36820b0dd072f89898604e2d39686fe05321729d00d8ac185a0", size = 162557, upload-time = "2026-07-07T14:33:32.176Z" }, + { url = "https://files.pythonhosted.org/packages/f0/e6/0386d43a261ff4e4b30c5857af7df877254b46bec7b9d1b74b6bf969a90b/charset_normalizer-3.4.9-cp312-cp312-win_arm64.whl", hash = "sha256:78fa18e436a1a0e58dbd7e02fc4473f3f32cceb12df9dfca542d075961c307d2", size = 152665, upload-time = "2026-07-07T14:33:33.711Z" }, + { url = "https://files.pythonhosted.org/packages/98/2b/f97f1c193fb855c345d678f5077d6926034db0722df74c8f057020e05a25/charset_normalizer-3.4.9-py3-none-any.whl", hash = "sha256:68e5f26a1ad57ded6d1cfb85331d1c1a195314756471d97758c48498bb4dcdf5", size = 64538, upload-time = "2026-07-07T14:34:56.993Z" }, +] + +[[package]] +name = "click" +version = "8.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/76/d4/81420972a676e8ffea40450d8c8c92943e7218a78fe9b64359836cc9876b/click-8.4.2.tar.gz", hash = "sha256:9a6cea6e60b17ebe0a44c5cc636d94f09bd66142c1cd7d8b4cd731c4917a15f6", size = 338000, upload-time = "2026-06-24T17:45:15.148Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/e2/79c688af8b210d232694e31e59da9f6ec747bae31c3f5946e4e9b98860d5/click-8.4.2-py3-none-any.whl", hash = "sha256:e6f9f66136c816745b9d65817da91d61d957fb16e02e4dcd0552553c5a197b76", size = 119243, upload-time = "2026-06-24T17:45:13.73Z" }, +] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + +[[package]] +name = "cuda-bindings" +version = "13.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/67/5e7dba1ba576dd73da5dee894ca076ca5e959450dfff66d6d510a255d1f7/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7855c4868aabc0cfae28abbe83d56734bdfbd08f08fc234ac1912a12858bf49", size = 6025351, upload-time = "2026-05-29T23:11:49.685Z" }, + { url = "https://files.pythonhosted.org/packages/39/2a/6d2e9047d1fb243dbaa364b01e0297534b9ed7fd27dba1c9f361519cf69b/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e32d08f71ebcdf00f0f41eab2eb37e8da94c8ed411cc9f7f7a019ce6b34abe3a", size = 6657965, upload-time = "2026-05-29T23:11:52.227Z" }, +] + +[[package]] +name = "cuda-pathfinder" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/b4/d088047afe39827556df21118cac9ffd20cc3f968c99a7681494d1eb333c/cuda_pathfinder-1.6.0-py3-none-any.whl", hash = "sha256:1503af579d8379c24bdd65528379bc57039b0455be9f5f9686cf8e473a1fce51", size = 54591, upload-time = "2026-07-21T15:03:56.224Z" }, +] + +[[package]] +name = "cuda-toolkit" +version = "13.0.3.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/c7/a79086a62c98befcdb8349656c6f114e2db3b8b2422f6e25c97a7f2a9a3c/cuda_toolkit-13.0.3.0-py2.py3-none-any.whl", hash = "sha256:d693caaa261214ddd7dbb60d68e71cbed884e68c2be7509778f3051da0b91c3f", size = 2512, upload-time = "2026-04-14T00:50:08.173Z" }, +] + +[package.optional-dependencies] +cublas = [ + { name = "nvidia-cublas" }, + { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +cudart = [ + { name = "nvidia-cuda-runtime", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +cufft = [ + { name = "nvidia-cufft", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +cufile = [ + { name = "nvidia-cufile", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +cupti = [ + { name = "nvidia-cuda-cupti", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +curand = [ + { name = "nvidia-curand", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +cusolver = [ + { name = "nvidia-cublas" }, + { name = "nvidia-cusolver", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +cusparse = [ + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +nvjitlink = [ + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +nvrtc = [ + { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] +nvtx = [ + { name = "nvidia-nvtx", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" }, +] + +[[package]] +name = "diffusers" +version = "0.39.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filelock" }, + { name = "httpx" }, + { name = "huggingface-hub" }, + { name = "importlib-metadata" }, + { name = "numpy" }, + { name = "pillow" }, + { name = "regex" }, + { name = "requests" }, + { name = "safetensors" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1a/81/6095237b86a3116c4789f28c4435d5296c00c0fc74ffde99008fd6b3a36c/diffusers-0.39.0.tar.gz", hash = "sha256:14bb1d98c85a0e463d734c99aaa73b480a7bc9bad22af30fbf730ef8f09c1d67", size = 4651240, upload-time = "2026-07-03T08:48:47.904Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/3f/7469c46e9d22307ea686bab687d70e6bf328722952f9d10339f5e913e608/diffusers-0.39.0-py3-none-any.whl", hash = "sha256:912aca51b5787365110806e984d5555735bf8a461073bb8459029d0bca7870ef", size = 5631176, upload-time = "2026-07-03T08:48:45.337Z" }, +] + +[[package]] +name = "docstring-parser" +version = "0.18.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/4d/f332313098c1de1b2d2ff91cf2674415cc7cddab2ca1b01ae29774bd5fdf/docstring_parser-0.18.0.tar.gz", hash = "sha256:292510982205c12b1248696f44959db3cdd1740237a968ea1e2e7a900eeb2015", size = 29341, upload-time = "2026-04-14T04:09:19.867Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/5f/ed01f9a3cdffbd5a008556fc7b2a08ddb1cc6ace7effa7340604b1d16699/docstring_parser-0.18.0-py3-none-any.whl", hash = "sha256:b3fcbed555c47d8479be0796ef7e19c2670d428d72e96da63f3a40122860374b", size = 22484, upload-time = "2026-04-14T04:09:18.638Z" }, +] + +[[package]] +name = "einops" +version = "0.8.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2c/77/850bef8d72ffb9219f0b1aac23fbc1bf7d038ee6ea666f331fa273031aa2/einops-0.8.2.tar.gz", hash = "sha256:609da665570e5e265e27283aab09e7f279ade90c4f01bcfca111f3d3e13f2827", size = 56261, upload-time = "2026-01-26T04:13:17.638Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl", hash = "sha256:54058201ac7087911181bfec4af6091bb59380360f069276601256a76af08193", size = 65638, upload-time = "2026-01-26T04:13:18.546Z" }, +] + +[[package]] +name = "filelock" +version = "3.32.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c0/80/8232b582c4b318b817cf1274ba74976b07b34d35ef439b3eb948f98645a1/filelock-3.32.0.tar.gz", hash = "sha256:7be2ad23a14607ccc71808e68fe30848aeace7058ace17852f68e2a68e310402", size = 213757, upload-time = "2026-07-21T13:17:42.898Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/06/79/b4c714bef36bc4ec2beeae1e0c124f0223888cd8c6feb1cdc56038116920/filelock-3.32.0-py3-none-any.whl", hash = "sha256:d396bea984af47333ef05e50eae7eff88c84256de6112aea0ec48a233c064fe3", size = 97732, upload-time = "2026-07-21T13:17:41.55Z" }, +] + +[[package]] +name = "fla-core" +version = "0.5.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "einops" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4e/62/99e149f19a447ce809d7f4fa64ae61c073da337505b9db7f389502470820/fla_core-0.5.1.tar.gz", hash = "sha256:7f3cf56edfbaa9115f4937d1181372e5c7b11809ad8eb2e411fffc3caf729f48", size = 498800, upload-time = "2026-06-18T18:17:15.377Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/78/a55ee7a62515dcb9220770dd99dfe59ac6599da8af84ad1d20f9f407df4d/fla_core-0.5.1-py3-none-any.whl", hash = "sha256:02150d34aa1e37f6b8ed9b2feec5d29af93680573e5077ed981238179c11fb06", size = 702955, upload-time = "2026-06-18T18:17:12.229Z" }, +] + +[[package]] +name = "flash-linear-attention" +version = "0.5.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "fla-core" }, + { name = "transformers" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4c/e8/8f115be585a046795e4a1f7ade727889219bb66b8d96cc668b8c9e437c0e/flash_linear_attention-0.5.1.tar.gz", hash = "sha256:8840fd4c37de8b0612dc8fd493867f3d330672ba2f17c024a2ce37239634e247", size = 221733, upload-time = "2026-06-18T18:17:16.661Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/68/3c/5819fb19dc071302ca818616a4e64d4454e1f1193929eb8365c8d38e9052/flash_linear_attention-0.5.1-py3-none-any.whl", hash = "sha256:9022862f0a238752372c81290694b8b2ed1cb2d13fc40d340ffeeb554d6bee5c", size = 403096, upload-time = "2026-06-18T18:17:14.025Z" }, +] + +[[package]] +name = "flashdreams" +source = { editable = "../../../../flashdreams" } +dependencies = [ + { name = "boto3" }, + { name = "botocore" }, + { name = "einops" }, + { name = "filelock" }, + { name = "ftfy" }, + { name = "huggingface-hub" }, + { name = "loguru" }, + { name = "numpy" }, + { name = "nvidia-ml-py" }, + { name = "safetensors" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, + { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32'" }, + { name = "tqdm" }, + { name = "transformers" }, + { name = "triton-windows", marker = "sys_platform == 'win32'" }, + { name = "tyro" }, + { name = "urllib3" }, +] + +[package.metadata] +requires-dist = [ + { name = "aiohttp", marker = "extra == 'serving'", specifier = ">=3.9" }, + { name = "aiortc", marker = "extra == 'serving'", specifier = ">=1.9" }, + { name = "boto3", specifier = ">=1.35" }, + { name = "botocore", specifier = ">=1.35" }, + { name = "einops", specifier = ">=0.7" }, + { name = "filelock", specifier = ">=3" }, + { name = "ftfy", specifier = ">=6.0" }, + { name = "huggingface-hub", specifier = ">=0.33" }, + { name = "loguru", specifier = ">=0.7" }, + { name = "mediapy", marker = "extra == 'dev'", specifier = ">=1.1" }, + { name = "mediapy", marker = "extra == 'examples'", specifier = ">=1.1" }, + { name = "mediapy", marker = "extra == 'runners'", specifier = ">=1.1" }, + { name = "numpy", specifier = ">=1.24,<2.5" }, + { name = "nvidia-ml-py", specifier = ">=12.0" }, + { name = "opencv-python-headless", marker = "extra == 'examples'", specifier = ">=4.5" }, + { name = "opencv-python-headless", marker = "extra == 'runners'", specifier = ">=4.5" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, + { name = "pytest-manual-marker", marker = "extra == 'dev'", specifier = ">=2.0" }, + { name = "safetensors", specifier = ">=0.4" }, + { name = "scipy", marker = "extra == 'examples'", specifier = ">=1.11" }, + { name = "scipy", marker = "extra == 'runners'", specifier = ">=1.11" }, + { name = "tomli", marker = "extra == 'dev'", specifier = ">=2.0" }, + { name = "torch", marker = "sys_platform != 'win32'", specifier = ">=2.9" }, + { name = "torch", marker = "sys_platform == 'win32'", specifier = ">=2.9", index = "https://download.pytorch.org/whl/cu130" }, + { name = "tqdm", specifier = ">=4.60" }, + { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "sys_platform != 'win32' and extra == 'dev'", specifier = ">=2.12" }, + { name = "transformers", specifier = ">=5.0,<6" }, + { name = "triton-windows", marker = "sys_platform == 'win32'", specifier = ">=3.5.0" }, + { name = "tyro", specifier = ">=1.0" }, + { name = "urllib3", specifier = ">=2.7.0" }, +] +provides-extras = ["dev", "examples", "runners", "serving"] + +[package.metadata.requires-dev] +cuda12 = [ + { name = "torch", marker = "sys_platform != 'win32'", specifier = ">=2.9", index = "https://download.pytorch.org/whl/cu128", conflict = { package = "flashdreams", group = "cuda12" } }, + { name = "torch", marker = "sys_platform == 'win32'", specifier = ">=2.9", index = "https://download.pytorch.org/whl/cu130" }, + { name = "torchvision", marker = "sys_platform != 'win32'", specifier = ">=0.24", index = "https://download.pytorch.org/whl/cu128", conflict = { package = "flashdreams", group = "cuda12" } }, + { name = "torchvision", marker = "sys_platform == 'win32'", specifier = ">=0.24", index = "https://download.pytorch.org/whl/cu130" }, +] +cuda13 = [ + { name = "torch", marker = "sys_platform != 'win32'", specifier = ">=2.9" }, + { name = "torch", marker = "sys_platform == 'win32'", specifier = ">=2.9", index = "https://download.pytorch.org/whl/cu130" }, + { name = "torchvision", marker = "sys_platform != 'win32'", specifier = ">=0.24" }, + { name = "torchvision", marker = "sys_platform == 'win32'", specifier = ">=0.24", index = "https://download.pytorch.org/whl/cu130" }, +] + +[[package]] +name = "flashdreams-sana-wm" +version = "0.1.0" +source = { editable = "../../" } +dependencies = [ + { name = "accelerate" }, + { name = "diffusers" }, + { name = "flashdreams" }, + { name = "imageio", extra = ["ffmpeg"] }, + { name = "pillow" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torchvision" }, + { name = "transformers" }, +] + +[package.metadata] +requires-dist = [ + { name = "accelerate", specifier = ">=1.0" }, + { name = "diffusers", specifier = ">=0.36" }, + { name = "flashdreams", editable = "../../../../flashdreams" }, + { name = "imageio", extras = ["ffmpeg"], specifier = ">=2.31" }, + { name = "pillow", specifier = ">=10" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, + { name = "pyyaml", specifier = ">=6.0" }, + { name = "safetensors", specifier = ">=0.5" }, + { name = "torchvision", specifier = ">=0.26" }, + { name = "transformers", specifier = ">=5.0,<6" }, +] +provides-extras = ["dev"] + +[[package]] +name = "fsspec" +version = "2026.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/10/a1/ae4e3e5003468d6391d2c77b6fa1cd73bd5d13511d81c642d7b28ac90ed4/fsspec-2026.6.0.tar.gz", hash = "sha256:f5bac145310fe30e16e1471bd6840b2d990d609e872251d7e674241822abf01a", size = 313646, upload-time = "2026-06-16T01:57:28.105Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/22/4222d7ddf3da30f363edaa98e329c2bce6c65497c9cb2810931c8b2c0fbc/fsspec-2026.6.0-py3-none-any.whl", hash = "sha256:02e0b71817df9b2169dc30a16832045764def1191b43dcff5bb85bdee212d2a1", size = 203949, upload-time = "2026-06-16T01:57:26.358Z" }, +] + +[[package]] +name = "ftfy" +version = "6.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a5/d3/8650919bc3c7c6e90ee3fa7fd618bf373cbbe55dff043bd67353dbb20cd8/ftfy-6.3.1.tar.gz", hash = "sha256:9b3c3d90f84fb267fe64d375a07b7f8912d817cf86009ae134aa03e1819506ec", size = 308927, upload-time = "2024-10-26T00:50:35.149Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ab/6e/81d47999aebc1b155f81eca4477a616a70f238a2549848c38983f3c22a82/ftfy-6.3.1-py3-none-any.whl", hash = "sha256:7c70eb532015cd2f9adb53f101fb6c7945988d023a085d127d1573dc49dd0083", size = 44821, upload-time = "2024-10-26T00:50:33.425Z" }, +] + +[[package]] +name = "h11" +version = "0.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/01/ee/02a2c011bdab74c6fb3c75474d40b3052059d95df7e73351460c8588d963/h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1", size = 101250, upload-time = "2025-04-24T03:35:25.427Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, +] + +[[package]] +name = "hf-xet" +version = "1.5.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/63/39/67be8d71f900d9a55761b6022821d6679fb56c64f1b6063d5af2c2606727/hf_xet-1.5.2.tar.gz", hash = "sha256:73044bd31bae33c984af832d19c752a0dffb67518fee9ddbd91d616e1101cf47", size = 903674, upload-time = "2026-07-16T17:29:56.833Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/ba/2b70603c7552db82baeb2623e2336898304a17328845151be4fe1f48d420/hf_xet-1.5.2-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:f922b8f5fb84f1dd3d7ab7a1316354a1bca9b1c73ecfc19c76e51a2a49d29799", size = 4033760, upload-time = "2026-07-16T17:29:43.884Z" }, + { url = "https://files.pythonhosted.org/packages/60/ac/b097a86a1e4a6098f3a79382643ab09d5733d87ccc864877ad1e12b49b70/hf_xet-1.5.2-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:045f84440c55cdeb659cf1a1dd48c77bcd0d2e93632e2fea8f2c3bdee79f38ed", size = 3841438, upload-time = "2026-07-16T17:29:45.539Z" }, + { url = "https://files.pythonhosted.org/packages/d3/35/db860aa3a0780660324a506ad4b3d322ddc6ecbba4b9340aed0942cbf21c/hf_xet-1.5.2-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:db78c39c83d6279daddc98e2238f373ab8980685556d42472b4ec51abcf03e8c", size = 4428006, upload-time = "2026-07-16T17:29:46.996Z" }, + { url = "https://files.pythonhosted.org/packages/af/6b/832dd980af4b0c3ae0660e309285f2ffcdff2faa38129390dbb47aa4a3f9/hf_xet-1.5.2-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:7db73c810500c54c6760be8c39d4b2e476974de85424c50063efc22fdda13025", size = 4221099, upload-time = "2026-07-16T17:29:48.525Z" }, + { url = "https://files.pythonhosted.org/packages/9e/05/ae50f0d34e3254e6c3e208beb2519f6b8673016fc4b3643badaf6450d186/hf_xet-1.5.2-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:6395cfe3c9cbead4f16b31808b0e67eac428b66c656f856e99636adaddea878f", size = 4420766, upload-time = "2026-07-16T17:29:50.092Z" }, + { url = "https://files.pythonhosted.org/packages/07/a9/c050bc2743a2bcd68928bfee157b08681667a164a24ec95fbfcfcd717e08/hf_xet-1.5.2-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:cde8cd167126bb6109b2ceb19b844433a4988643e8f3e01dd9dd0e4a34535097", size = 4636716, upload-time = "2026-07-16T17:29:51.62Z" }, + { url = "https://files.pythonhosted.org/packages/e9/f8/68b01c5c2edb56ac9a67b3d076ffddcb90867abaee923923eb34e7a14e76/hf_xet-1.5.2-cp38-abi3-win_amd64.whl", hash = "sha256:ecf63d1cb69a9a7319910f8f83fcf9b46e7a32dfcf4b8f8eeddb55f647306e65", size = 3988373, upload-time = "2026-07-16T17:29:53.395Z" }, + { url = "https://files.pythonhosted.org/packages/39/c6/988383e9dc17294d536fcbcd6fd16eed882e411ad16c954984a53e47b09c/hf_xet-1.5.2-cp38-abi3-win_arm64.whl", hash = "sha256:1da28519496eb7c8094c11e4d25509b4a468457a0302d58136099db2fd9a671d", size = 3816957, upload-time = "2026-07-16T17:29:54.991Z" }, +] + +[[package]] +name = "httpcore" +version = "1.0.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "h11" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/94/82699a10bca87a5556c9c59b5963f2d039dbd239f25bc2a63907a05a14cb/httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8", size = 85484, upload-time = "2025-04-24T22:06:22.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, +] + +[[package]] +name = "httpx" +version = "0.28.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "certifi" }, + { name = "httpcore" }, + { name = "idna" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406, upload-time = "2024-12-06T15:37:23.222Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, +] + +[[package]] +name = "huggingface-hub" +version = "1.24.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "hf-xet", marker = "platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, + { name = "httpx" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "tqdm" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/9b/d3bb4e7d792835daf34dd7091bbc7d7b4e0437d9388f1ea7239cce49f478/huggingface_hub-1.24.0.tar.gz", hash = "sha256:18431ff4daae0749aa9ba102fc952e314c98e1d30ebdec5319d85ca0a83e1ae5", size = 921848, upload-time = "2026-07-17T09:54:01.022Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5f/c3/aeaaf3911d2529614be18d1c8b5496afc185560e76568063d517283318af/huggingface_hub-1.24.0-py3-none-any.whl", hash = "sha256:6ed4120a84a6beec900640aa7e346bd766a6b7341e41526fef5dc8bd81fb7d59", size = 771904, upload-time = "2026-07-17T09:53:59.106Z" }, +] + +[[package]] +name = "idna" +version = "3.18" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/63/9496c57188a2ee585e0f1db071d75089a11e98aa86eb99d9d7618fc1edce/idna-3.18.tar.gz", hash = "sha256:ffb385a7e039654cef1ab9ef32c6fafe283c0c0467bba1d9029738ce4a14a848", size = 196711, upload-time = "2026-06-02T14:34:07.794Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455, upload-time = "2026-06-02T14:34:06.319Z" }, +] + +[[package]] +name = "imageio" +version = "2.37.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "pillow" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/48/62/aa770a9307508d2a2a2c62d536a49347bffe9e55322db27838d3c93d0b07/imageio-2.37.4.tar.gz", hash = "sha256:e45cbc5e83502047fb138f7f585f7f105a136a57eea5f4b3cfc6ce1b52720bd3", size = 390173, upload-time = "2026-07-20T05:26:11.369Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3e/2d/ca050652104bab2cf55e569db2a178b1b61cb041fef28307f2db383f6d9f/imageio-2.37.4-py3-none-any.whl", hash = "sha256:1ab2e22c8debf700f24c3ac43e8f95f3b3a8110c83b93411e97b4b0b2cd1c7e6", size = 318000, upload-time = "2026-07-20T05:26:09.874Z" }, +] + +[package.optional-dependencies] +ffmpeg = [ + { name = "imageio-ffmpeg" }, + { name = "psutil" }, +] + +[[package]] +name = "imageio-ffmpeg" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/44/bd/c3343c721f2a1b0c9fc71c1aebf1966a3b7f08c2eea8ed5437a2865611d6/imageio_ffmpeg-0.6.0.tar.gz", hash = "sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755", size = 25210, upload-time = "2025-01-16T21:34:32.747Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/58/87ef68ac83f4c7690961bce288fd8e382bc5f1513860fc7f90a9c1c1c6bf/imageio_ffmpeg-0.6.0-py3-none-macosx_10_9_intel.macosx_10_9_x86_64.whl", hash = "sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61", size = 24932969, upload-time = "2025-01-16T21:34:20.464Z" }, + { url = "https://files.pythonhosted.org/packages/40/5c/f3d8a657d362cc93b81aab8feda487317da5b5d31c0e1fdfd5e986e55d17/imageio_ffmpeg-0.6.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b1ae3173414b5fc5f538a726c4e48ea97edc0d2cdc11f103afee655c463fa742", size = 21113891, upload-time = "2025-01-16T21:34:00.277Z" }, + { url = "https://files.pythonhosted.org/packages/33/e7/1925bfbc563c39c1d2e82501d8372734a5c725e53ac3b31b4c2d081e895b/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1d47bebd83d2c5fc770720d211855f208af8a596c82d17730aa51e815cdee6dc", size = 25632706, upload-time = "2025-01-16T21:33:53.475Z" }, + { url = "https://files.pythonhosted.org/packages/a0/2d/43c8522a2038e9d0e7dbdf3a61195ecc31ca576fb1527a528c877e87d973/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:c7e46fcec401dd990405049d2e2f475e2b397779df2519b544b8aab515195282", size = 29498237, upload-time = "2025-01-16T21:34:13.726Z" }, + { url = "https://files.pythonhosted.org/packages/a0/13/59da54728351883c3c1d9fca1710ab8eee82c7beba585df8f25ca925f08f/imageio_ffmpeg-0.6.0-py3-none-win32.whl", hash = "sha256:196faa79366b4a82f95c0f4053191d2013f4714a715780f0ad2a68ff37483cc2", size = 19652251, upload-time = "2025-01-16T21:34:06.812Z" }, + { url = "https://files.pythonhosted.org/packages/2c/c6/fa760e12a2483469e2bf5058c5faff664acf66cadb4df2ad6205b016a73d/imageio_ffmpeg-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02fa47c83703c37df6bfe4896aab339013f62bf02c5ebf2dce6da56af04ffc0a", size = 31246824, upload-time = "2025-01-16T21:34:28.6Z" }, +] + +[[package]] +name = "importlib-metadata" +version = "9.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "zipp" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a9/01/15bb152d77b21318514a96f43af312635eb2500c96b55398d020c93d86ea/importlib_metadata-9.0.0.tar.gz", hash = "sha256:a4f57ab599e6a2e3016d7595cfd72eb4661a5106e787a95bcc90c7105b831efc", size = 56405, upload-time = "2026-03-20T06:42:56.999Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/38/3d/2d244233ac4f76e38533cfcb2991c9eb4c7bf688ae0a036d30725b8faafe/importlib_metadata-9.0.0-py3-none-any.whl", hash = "sha256:2d21d1cc5a017bd0559e36150c21c830ab1dc304dedd1b7ea85d20f45ef3edd7", size = 27789, upload-time = "2026-03-20T06:42:55.665Z" }, +] + +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + +[[package]] +name = "jmespath" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d3/59/322338183ecda247fb5d1763a6cbe46eff7222eaeebafd9fa65d4bf5cb11/jmespath-1.1.0.tar.gz", hash = "sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d", size = 27377, upload-time = "2026-01-22T16:35:26.279Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, +] + +[[package]] +name = "loguru" +version = "0.7.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "win32-setctime", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3a/05/a1dae3dffd1116099471c643b8924f5aa6524411dc6c63fdae648c4f1aca/loguru-0.7.3.tar.gz", hash = "sha256:19480589e77d47b8d85b2c827ad95d49bf31b0dcde16593892eb51dd18706eb6", size = 63559, upload-time = "2024-12-06T11:20:56.608Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0c/29/0348de65b8cc732daa3e33e67806420b2ae89bdce2b04af740289c5c6c8c/loguru-0.7.3-py3-none-any.whl", hash = "sha256:31a33c10c8e1e10422bfd431aeb5d351c7cf7fa671e3c4df004162264b28220c", size = 61595, upload-time = "2024-12-06T11:20:54.538Z" }, +] + +[[package]] +name = "markdown-it-py" +version = "4.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mdurl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/ff/7841249c247aa650a76b9ee4bbaeae59370dc8bfd2f6c01f3630c35eb134/markdown_it_py-4.2.0.tar.gz", hash = "sha256:04a21681d6fbb623de53f6f364d352309d4094dd4194040a10fd51833e418d49", size = 82454, upload-time = "2026-05-07T12:08:28.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/81/4da04ced5a082363ecfa159c010d200ecbd959ae410c10c0264a38cac0f5/markdown_it_py-4.2.0-py3-none-any.whl", hash = "sha256:9f7ebbcd14fe59494226453aed97c1070d83f8d24b6fc3a3bcf9a38092641c4a", size = 91687, upload-time = "2026-05-07T12:08:27.182Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, + { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, + { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, + { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, + { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, +] + +[[package]] +name = "mdurl" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, +] + +[[package]] +name = "ml-dtypes" +version = "0.5.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/b8/3c70881695e056f8a32f8b941126cf78775d9a4d7feba8abcb52cb7b04f2/ml_dtypes-0.5.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a174837a64f5b16cab6f368171a1a03a27936b31699d167684073ff1c4237dac", size = 676927, upload-time = "2025-11-17T22:31:48.182Z" }, + { url = "https://files.pythonhosted.org/packages/54/0f/428ef6881782e5ebb7eca459689448c0394fa0a80bea3aa9262cba5445ea/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900", size = 5028464, upload-time = "2025-11-17T22:31:50.135Z" }, + { url = "https://files.pythonhosted.org/packages/3a/cb/28ce52eb94390dda42599c98ea0204d74799e4d8047a0eb559b6fd648056/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff", size = 5009002, upload-time = "2025-11-17T22:31:52.001Z" }, +] + +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + +[[package]] +name = "mypy-extensions" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, +] + +[[package]] +name = "networkx" +version = "3.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, +] + +[[package]] +name = "numpy" +version = "2.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d0/ad/fed0499ce6a338d2a03ebae59cd15093910c8875328855781952abf6c2fe/numpy-2.4.6.tar.gz", hash = "sha256:f3a3570c4a2a16746ac2c31a7c7c7b0c186b95ce902e33db6f28094ed7387dda", size = 20735807, upload-time = "2026-05-18T23:37:14.07Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/2a/3d7b5ac8aac24feaf9ad7ed58f45b0bbc06d37e4338ae84c9f2298b570f9/numpy-2.4.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:001fbb8e08d942dd57599e781f2472269ee7f2755fae407b4f67b2f0b17da3f1", size = 16689119, upload-time = "2026-05-18T23:33:54.065Z" }, + { url = "https://files.pythonhosted.org/packages/ea/12/92c4c131527599e8288d6918e888d88726f84d805d784b771f32408aeaef/numpy-2.4.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ebfb099f8dcf083deef3ac1ca4c1503f387cf76296fcb3816b66f5ecb5f54fdb", size = 14699246, upload-time = "2026-05-18T23:33:57.621Z" }, + { url = "https://files.pythonhosted.org/packages/ad/fe/c0a6b7b2ca128a8fb228575147073b660656734b8ebe4d76c8fd748dcc79/numpy-2.4.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:3213d622a0283a39a93d188f3cf72b26862df52fbb4ca3697f51705016523d41", size = 5204410, upload-time = "2026-05-18T23:34:00.302Z" }, + { url = "https://files.pythonhosted.org/packages/f3/d4/9770d14ba719432bb90a421bfd443872ed0f70f7264b64bec12ea363d5fd/numpy-2.4.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:357cc07a6d7b0b182ff02249616a03742827ebb1277546b5c7cd7f7620a45698", size = 6551240, upload-time = "2026-05-18T23:34:02.852Z" }, + { url = "https://files.pythonhosted.org/packages/c9/c6/50a46a6205feba2343f1d6d17438107c5dc491ed1c736e6ea68689fd906b/numpy-2.4.6-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5f9fb9157b4ce2971008323afe46053787b526ef624fea915b261468a8421a0f", size = 15671012, upload-time = "2026-05-18T23:34:05.485Z" }, + { url = "https://files.pythonhosted.org/packages/99/60/14115e6364fa676c5397c2ad3004e527e9aa487abf5d0706ec81bbd08529/numpy-2.4.6-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:90f9849678c75fe7afa2d348ac842c168b0a4d3d61919687216dfc547976d853", size = 16645538, upload-time = "2026-05-18T23:34:09.265Z" }, + { url = "https://files.pythonhosted.org/packages/ae/c5/693cbe59e57db94d2231fa519ca3978dc9e19da5a8f088588f5c6e947ff2/numpy-2.4.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:c1a2af6c6ef86344a6b0db6b97834208bf598db514f2b155042439b62605601a", size = 17020706, upload-time = "2026-05-18T23:34:13.053Z" }, + { url = "https://files.pythonhosted.org/packages/ef/fc/85b7c4eff9b4966ade25c2273cf7e7012e92366c032058653934b37de044/numpy-2.4.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e5805d5a22fd19c8ccff10a9561f9df94436b0545619ea579db2d3c35294bce2", size = 18368541, upload-time = "2026-05-18T23:34:17.024Z" }, + { url = "https://files.pythonhosted.org/packages/f6/81/e1b27545deedce7f4a0b348618c6b62d74e36a4dc9ccd42f3eb2f85eee32/numpy-2.4.6-cp312-cp312-win32.whl", hash = "sha256:e3eeb0aabd6bd5ce64faae67e9935203a6991b4bc2a485a767fbafb2c5125f45", size = 5962825, upload-time = "2026-05-18T23:34:20.3Z" }, + { url = "https://files.pythonhosted.org/packages/ab/ca/feab00bd44aa5fe1ad2c18f08b4d3bb92e26484b0b1d1443897809ed528c/numpy-2.4.6-cp312-cp312-win_amd64.whl", hash = "sha256:d8e8286dd7cea7895157318d1b91cdacac64c479f3cbc8dce548331728484751", size = 12321687, upload-time = "2026-05-18T23:34:23.095Z" }, + { url = "https://files.pythonhosted.org/packages/63/cf/5a6d34850a39d1093558564f77ee8e8e0bee5061151b8f05a55711001ec7/numpy-2.4.6-cp312-cp312-win_arm64.whl", hash = "sha256:4081eb135ac24158bd51cdfbef16f1c64df7063b1143f24731387137c092bec8", size = 10221482, upload-time = "2026-05-18T23:34:25.876Z" }, +] + +[[package]] +name = "nvdlfw-inspect" +version = "0.2.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyyaml" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" } }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8a/86/94188e03e5d4dd7b73c390b0cddcde5618b3799c18e327b2bf15763f6137/nvdlfw_inspect-0.2.2-py3-none-any.whl", hash = "sha256:8a4dc2814c5a4cd19ae304170b9bfa514538ef3c3eb243a45a82404ec3cb279d", size = 30964, upload-time = "2025-12-03T10:52:01.933Z" }, +] + +[[package]] +name = "nvidia-cublas" +version = "13.6.0.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-nvrtc" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/4f/22/7c08d8e93f6a2e4879ac83c12696aecaf17a0fc2c9e8d204caceaa3b8426/nvidia_cublas-13.6.0.2-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:946f6a252b1cc72d8de912c75975fd6d8ba44f67d4e5044fe764ddb909f4a688", size = 518599300, upload-time = "2026-06-29T16:54:56.859Z" }, + { url = "https://files.pythonhosted.org/packages/fd/6c/173c7a3db77a6592210f73f194f0f8ed5e51b6ec61cfed7b1eee06ac5fd3/nvidia_cublas-13.6.0.2-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b82c80c886cea6da6e149a5c3bdba274f12b7e4ec4b00a050b916b0446fb4153", size = 410473152, upload-time = "2026-06-29T16:55:54.297Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti" +version = "13.0.85" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/2a/80353b103fc20ce05ef51e928daed4b6015db4aaa9162ed0997090fe2250/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151", size = 10310827, upload-time = "2025-09-04T08:26:42.012Z" }, + { url = "https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8", size = 10715597, upload-time = "2025-09-04T08:26:51.312Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc" +version = "13.0.88" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575", size = 90215200, upload-time = "2025-09-04T08:28:44.204Z" }, + { url = "https://files.pythonhosted.org/packages/b7/dc/6bb80850e0b7edd6588d560758f17e0550893a1feaf436807d64d2da040f/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b", size = 43015449, upload-time = "2025-09-04T08:28:20.239Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime" +version = "13.0.96" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/87/4f/17d7b9b8e285199c58ce28e31b5c5bbaa4d8271af06a89b6405258245de2/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55", size = 2261060, upload-time = "2025-10-09T08:55:15.78Z" }, + { url = "https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548", size = 2243632, upload-time = "2025-10-09T08:55:36.117Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu13" +version = "9.20.0.48" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" }, + { url = "https://files.pythonhosted.org/packages/6e/5e/edb9c0ae051602c3ccaffe424256463636d639e27d7f302dde9975ef9e7a/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0c45dd8eeb50b603f07995b1b300c62ffe6a1980482b82b3bcf94a4ca9d49304", size = 366173588, upload-time = "2026-03-09T19:29:34.474Z" }, +] + +[[package]] +name = "nvidia-cufft" +version = "12.0.0.61" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, + { url = "https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3", size = 214085489, upload-time = "2025-09-04T08:31:56.044Z" }, +] + +[[package]] +name = "nvidia-cufile" +version = "1.15.1.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44", size = 1223672, upload-time = "2025-09-04T08:32:22.779Z" }, + { url = "https://files.pythonhosted.org/packages/ab/73/cc4a14c9813a8a0d509417cf5f4bdaba76e924d58beb9864f5a7baceefbf/nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1", size = 1136992, upload-time = "2025-09-04T08:32:14.119Z" }, +] + +[[package]] +name = "nvidia-curand" +version = "10.4.0.35" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/72/7c2ae24fb6b63a32e6ae5d241cc65263ea18d08802aaae087d9f013335a2/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a", size = 61962106, upload-time = "2025-08-04T10:21:41.128Z" }, + { url = "https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc", size = 59544258, upload-time = "2025-08-04T10:22:03.992Z" }, +] + +[[package]] +name = "nvidia-cusolver" +version = "12.0.4.66" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas" }, + { name = "nvidia-cusparse" }, + { name = "nvidia-nvjitlink" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, + { url = "https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112", size = 200941980, upload-time = "2025-09-04T08:33:22.767Z" }, +] + +[[package]] +name = "nvidia-cusparse" +version = "12.6.3.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, + { url = "https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b", size = 145942937, upload-time = "2025-09-04T08:33:58.029Z" }, +] + +[[package]] +name = "nvidia-cusparselt-cu13" +version = "0.8.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/e1/cdc1797eadf82d3a9a575a19b33fdc871a97edbec42c00b5b5e914f4aff4/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f", size = 221051344, upload-time = "2025-09-05T18:49:51.289Z" }, + { url = "https://files.pythonhosted.org/packages/34/7d/2661f2fb3ac4302f3a246f5fc030213ac60c1fe0bce84f9783dbd831dbb7/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:786ce87568c303fadb5afcc7102d454cd3040d75f6f8626f5db460d1871f4dd0", size = 170148586, upload-time = "2025-09-05T18:50:50.248Z" }, +] + +[[package]] +name = "nvidia-ml-py" +version = "13.610.43" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f0/b5/a8fbc356f768fa5c9cfd646668fd7d34bf55bdd1c6e20754642a64d930d4/nvidia_ml_py-13.610.43.tar.gz", hash = "sha256:65437eb73d68d0c62c931ca4d45038472faff03bd0b8729abba4b899f70d60f2", size = 52109, upload-time = "2026-06-01T18:54:08.829Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/23/45/caa600acfab94560807a20a64b5830d2cd3c3202b7f1328644d70b7d6bd8/nvidia_ml_py-13.610.43-py3-none-any.whl", hash = "sha256:f13c72698edef492f985cc225f14faafe68ae065a2e407f45bdf6f4b9b43fde8", size = 53163, upload-time = "2026-06-01T18:54:07.704Z" }, +] + +[[package]] +name = "nvidia-nccl-cu13" +version = "2.29.7" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/72/0d/daf50d44177ee0cbc7ff0a0c91eb5ff676c82be42f9a970bc7597f440c3a/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:674a12383e3c38a1bcccae7d4f3633b37852230b6047883cb2f4c2d1b36d9bf5", size = 206014712, upload-time = "2026-03-03T05:34:20.843Z" }, + { url = "https://files.pythonhosted.org/packages/67/f4/58e4e91b6919367c7aafb8e36fce9aad1a3047e536bf7e2fd560927d3a4c/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d", size = 205976000, upload-time = "2026-03-03T05:36:24.472Z" }, +] + +[[package]] +name = "nvidia-nvjitlink" +version = "13.3.33" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f0/ee/580ca6f29dcab0221db8706badca1bbbb084f1975c4d4e83329c3a7e31f0/nvidia_nvjitlink-13.3.33-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:26a6de7fb4c8fdaa7703d3dad720d6d427ddfea5c48a528fd97c11733ad830e5", size = 40742423, upload-time = "2026-05-26T16:54:51.613Z" }, + { url = "https://files.pythonhosted.org/packages/69/30/45414e35ff2eee7db3da037e5707037ccf9d2b5218ffbdb055ea4d5aa98a/nvidia_nvjitlink-13.3.33-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ce48b37dfeb3cb1eae4cf85adacb47d7a6539ea2272870c9a3628ce275c2037e", size = 39168635, upload-time = "2026-05-26T16:54:13.906Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu13" +version = "3.4.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/0f/05cc9c720236dcd2db9c1ab97fff629e96821be2e63103569da0c9b72f19/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9", size = 60215947, upload-time = "2025-09-06T00:32:20.022Z" }, + { url = "https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80", size = 60412546, upload-time = "2025-09-06T00:32:41.564Z" }, +] + +[[package]] +name = "nvidia-nvtx" +version = "13.0.85" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4", size = 148047, upload-time = "2025-09-04T08:29:01.761Z" }, + { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, +] + +[[package]] +name = "omegaconf" +version = "2.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ce/3d/e4b57b8d9008c6ebe0d5eff901f91d5700cf7bdb8c8863df817463a7fd5e/omegaconf-2.3.1.tar.gz", hash = "sha256:e5e7de64aeebeddaf8e6d3f7a783b32ac2a01c0fbd9c878012caecb891a1f42a", size = 3298472, upload-time = "2026-06-11T05:05:12.885Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/0e/152509871bf30df6fc38569f52a2db9b55dd41aae957adae50a053ac7778/omegaconf-2.3.1-py3-none-any.whl", hash = "sha256:3d701d14e9a8828f1edd28bb70b725908b34277cdd72cf7d6a83f94dadc6b6a0", size = 79502, upload-time = "2026-06-11T05:05:09.954Z" }, +] + +[[package]] +name = "onnx" +version = "1.22.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "protobuf" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/04/19/8ea73a64b368b75fe339771a20a02bc61ea1f551484c9e3d9d0bfbd0450f/onnx-1.22.0.tar.gz", hash = "sha256:ef40c0aaf0b643857ea9306fc7eddce17eaf9fb0407e4801f1fc5758443a38e0", size = 12024721, upload-time = "2026-06-15T12:50:05.354Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ee/6a/481561f1093834376ed493e4ca42a73e5be0d50031f2969c86593bdc7c96/onnx-1.22.0-cp312-abi3-macosx_12_0_universal2.whl", hash = "sha256:596fbf0490947533c1c1045ba860851dc9fb77471023dac9a71ba5b42ceab103", size = 20167081, upload-time = "2026-06-15T12:49:32.078Z" }, + { url = "https://files.pythonhosted.org/packages/84/55/b34fc2aa30aa54b4a775402d24c4082242c720283a274fe976ac8eb94480/onnx-1.22.0-cp312-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ae5a563f281cd9d2845622cecf6c092a57e4ee1b138f66fdbbdd4200567a5e16", size = 18889249, upload-time = "2026-06-15T12:49:34.7Z" }, + { url = "https://files.pythonhosted.org/packages/09/a6/bd32357e6cc1ecb473afd78193d7231724f284435d2db25696ecfaaa1503/onnx-1.22.0-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:955e02e1f6d385b53d52f9cd7b9cdf5caf417c300bcfe3c64c6d542be763845b", size = 19106514, upload-time = "2026-06-15T12:49:37.424Z" }, + { url = "https://files.pythonhosted.org/packages/5a/9d/3af461ac6c714b8b369cb71499659932f4f12cfb066250b62f7567c3d530/onnx-1.22.0-cp312-abi3-pyemscripten_2025_0_wasm32.whl", hash = "sha256:82e9f27fc1223cb06d68a56bed6f9d3caf3d0dad1b61bce45006d529b15bd94c", size = 16966387, upload-time = "2026-06-15T12:49:40.918Z" }, +] + +[[package]] +name = "onnx-ir" +version = "0.2.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "sympy" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/e6/672fefb2f108d077f58181a7babf4c0f8d1182a30353ffc9c79c63afc5ee/onnx_ir-0.2.1.tar.gz", hash = "sha256:8b8b10a93f43e65962104de6070c43c5dacb0e3cdfefc7c8059dd83c9db64f35", size = 144279, upload-time = "2026-04-20T20:21:47.735Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8c/aa/f7a53321c60b9ad9ee184b6018292ed6b5389947592a2c8c09c736bb7f9e/onnx_ir-0.2.1-py3-none-any.whl", hash = "sha256:c7285da889312f91882de2092e298a9eeeefbfc1d1951c49d983992967eb09a7", size = 166792, upload-time = "2026-04-20T20:21:46.357Z" }, +] + +[[package]] +name = "onnxscript" +version = "0.7.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "onnx-ir" }, + { name = "packaging" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c4/3a/4d79bce3f460e0df7fed54a92ce80827f25da66511da368bb00783ad8d20/onnxscript-0.7.1.tar.gz", hash = "sha256:309fb86484b11fa4ded90dba580e0d63f1a0827588e521cecaf2eeddb46d6e86", size = 618160, upload-time = "2026-06-29T23:33:21.526Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dd/bd/a0c8e737b6afda10e42a597787d53d5b66e00268df6f59184701eeae37d9/onnxscript-0.7.1-py3-none-any.whl", hash = "sha256:544763b7fdef49940cdd9412ff5135cbae96d59ac6bc1921457f21280f40f4b7", size = 721970, upload-time = "2026-06-29T23:33:23.298Z" }, +] + +[[package]] +name = "opencv-python-headless" +version = "5.0.0.93" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1d/99/76b7c80252aa83c1af16393454aafd125a0287101afe8deb0a6821af0e30/opencv_python_headless-5.0.0.93.tar.gz", hash = "sha256:b82f9831daab90b725c7c1ee1b36cb5732c367096ac76d119e64e14eb70d5f3c", size = 81817738, upload-time = "2026-07-02T07:01:06.039Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/53/7c/8c8097891c509d98cd128493835c95631c80be6a8f37ed9d25716c2e16f1/opencv_python_headless-5.0.0.93-cp37-abi3-macosx_13_0_arm64.whl", hash = "sha256:030ca5e0837a2963ab36ef896baa9767eb8d2b83353fb28af5a521e40dd8756f", size = 48322581, upload-time = "2026-07-02T05:50:34.207Z" }, + { url = "https://files.pythonhosted.org/packages/90/8c/eab2ad388c3cbab2a350c10c2ef19ce6bd099240afc31789032c996bab52/opencv_python_headless-5.0.0.93-cp37-abi3-macosx_14_0_x86_64.whl", hash = "sha256:1e55af3abfb462eeeabe5c775f12bdb36216d8a93a3583d69e6bd6e1d6ba7d00", size = 34782894, upload-time = "2026-07-02T05:51:39.856Z" }, + { url = "https://files.pythonhosted.org/packages/ec/78/afca939f40ffe2b2380bfa86f812b2f7d4acc5a27b27dc41b49cad7ce7b4/opencv_python_headless-5.0.0.93-cp37-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:10818d91510e05c04568ae12b5cd120779c70c01bf897b001a6221fe430df80f", size = 36521085, upload-time = "2026-07-02T06:55:24.429Z" }, + { url = "https://files.pythonhosted.org/packages/2b/97/8170e9819764c47e436c130d3ff6cfb73b58f923eae9d3a03d8982b04aec/opencv_python_headless-5.0.0.93-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:09a872a157c1376ab922a69bbf22f9a95bcc7b658a9d8b436a60212b02b2eeb4", size = 56563598, upload-time = "2026-07-02T06:55:47.355Z" }, + { url = "https://files.pythonhosted.org/packages/3a/98/1a28a7101e31801042b3098871a74b76c61581d328ef40774ff4edb53a56/opencv_python_headless-5.0.0.93-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:840bd717c21e5c11cadadc022a823315ea417f961213d06b4df010e019eb16f4", size = 39648433, upload-time = "2026-07-02T06:56:04.255Z" }, + { url = "https://files.pythonhosted.org/packages/9b/21/f6ef335f6e65724aa78b8d792b48d40a48c381715f1e62f5a5049e09d07e/opencv_python_headless-5.0.0.93-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ed709fdf9aa0bd1f2ed8549e71d19449b03a675bb581eb292285f6861953be37", size = 61204038, upload-time = "2026-07-02T06:56:41.823Z" }, + { url = "https://files.pythonhosted.org/packages/d0/8f/b8756467ea991449a293797f6b3fa80fcfdd29598a0a60d1cd5715b96e61/opencv_python_headless-5.0.0.93-cp37-abi3-win32.whl", hash = "sha256:c6bcd96b185975ea240d22cfdb15a1f6d080cc95264cfbe2621f21bb144d89b9", size = 35411237, upload-time = "2026-07-02T05:50:12.901Z" }, + { url = "https://files.pythonhosted.org/packages/b8/88/763b967f7efd7226b82c9fae16d560cba049b1f0c036647e65c610fd636e/opencv_python_headless-5.0.0.93-cp37-abi3-win_amd64.whl", hash = "sha256:829717b6a95554f273e49e357cee3b3a2a26b6f4842fbc1bed2b45bdd8f87e0e", size = 43825962, upload-time = "2026-07-02T05:50:09.627Z" }, +] + +[[package]] +name = "packaging" +version = "26.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d7/f1/e7a6dd94a8d4a5626c03e4e99c87f241ba9e350cd9e6d75123f992427270/packaging-26.2.tar.gz", hash = "sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661", size = 228134, upload-time = "2026-04-24T20:15:23.917Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195, upload-time = "2026-04-24T20:15:22.081Z" }, +] + +[[package]] +name = "pillow" +version = "12.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1c/3d/bb7fca845737cf9d7dbde16ed1843984665ff2e0a518f5db43e77ec540b9/pillow-12.3.0.tar.gz", hash = "sha256:3b8182a766685eaa002637e28b4ec8d6b18819a0c71f579bf0dbaa5830297cce", size = 47025035, upload-time = "2026-07-01T11:56:38.965Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/bf/fb3ebff8ddcb76aac5a01389251bbbb9519922a9b520d8247c1ca864a25d/pillow-12.3.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:ba09209fbe443b4acccebe845d8a138b89a8f4fbaeedd44953490b5315d5e965", size = 5345969, upload-time = "2026-07-01T11:54:06.397Z" }, + { url = "https://files.pythonhosted.org/packages/d8/66/9a386a92561f402389a4fc70c18838bf6d35eb5eb5c6850b4b2dc64f5048/pillow-12.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ffd0c5368496f41b0944be820fcb7a838aa6e623d250b01acf2643939c3f99d7", size = 4780323, upload-time = "2026-07-01T11:54:09.351Z" }, + { url = "https://files.pythonhosted.org/packages/25/27/ac8f99618ffd3dde21db0f4d4b1d2ab00c0880595bfd17df103f7f39fd0c/pillow-12.3.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d9c7f76c0673154f044e9d78c8655fb4213f6ca31a836df48b40fe5d187717b9", size = 6266838, upload-time = "2026-07-01T11:54:11.71Z" }, + { url = "https://files.pythonhosted.org/packages/84/21/a35af28dcc61f37ed850a2d64c65c701321dfbf25085e469d5559360cbbf/pillow-12.3.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:78cb2c6865a35ab8ff8b75fd122f6033b92a62c82801110e48ddd6c936a45d91", size = 6940830, upload-time = "2026-07-01T11:54:13.732Z" }, + { url = "https://files.pythonhosted.org/packages/eb/51/8b08617af3ad95e33ce6d7dd2c99ed6c8298f7fb131636303956be022e25/pillow-12.3.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e491916b378fba47242221bb9ead245211b70d504f495d105d17b14a24b4907c", size = 6344383, upload-time = "2026-07-01T11:54:15.756Z" }, + { url = "https://files.pythonhosted.org/packages/1d/72/cf78ac9780bb93c28328f408973845a309d4d145041665f734572ced1b52/pillow-12.3.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:0dd2064cbc55aaec028ef5fbb60fa47bb6c3e7918e07ff17935284b227a9d2df", size = 7052934, upload-time = "2026-07-01T11:54:17.721Z" }, + { url = "https://files.pythonhosted.org/packages/20/20/25e0f4dc178a6bc0696793720055519a0de89e7661dae886992decbd2f81/pillow-12.3.0-cp312-cp312-win32.whl", hash = "sha256:dbce0b29841537a2fa4a214c2bbf14de3587c9680caa9b4e217568472490b28f", size = 6472684, upload-time = "2026-07-01T11:54:19.839Z" }, + { url = "https://files.pythonhosted.org/packages/45/89/da2f7971a317f83d807fdd4065c0af40208e59e692cc43d315a71a0e96d1/pillow-12.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:a2b55dd6b2a4c4b7d87ffa56bdb33fdc5fdb9a462173861a7bc097f17d91cb09", size = 7227137, upload-time = "2026-07-01T11:54:22.025Z" }, + { url = "https://files.pythonhosted.org/packages/de/47/4845a0a6c0dbf1db8456bd9fc791f13c5ced7ced20606d08a0aacfd25b49/pillow-12.3.0-cp312-cp312-win_arm64.whl", hash = "sha256:331b624368d4f1d069149002f25f44bc61c8919ce8ddb3c45bdad8f6e2d89510", size = 2568267, upload-time = "2026-07-01T11:54:24.051Z" }, +] + +[[package]] +name = "protobuf" +version = "7.35.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/da/01/9ef0afd7999eb9badb3a768b4aedd78c86d4c65cfaf1958ab276199e76b4/protobuf-7.35.1.tar.gz", hash = "sha256:ce115a26fe0c39a2c29973d914d327e516a6455464489fe3cd1e51a1b354f81a", size = 458717, upload-time = "2026-06-11T21:55:40.257Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/03/8aeeb7458d22546bf64b5250ca1daeb5ff757d900e8e4a7476c6f0db843e/protobuf-7.35.1-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:24f857477359a85c0c235261b8ba905fd51b2562f4a64ca1df5473f29850cbf6", size = 433226, upload-time = "2026-06-11T21:55:31.719Z" }, + { url = "https://files.pythonhosted.org/packages/37/4b/dfb89eb0e652a1ff073c39a59fb5e3a83cfe9b57a2c83fa6d78270101767/protobuf-7.35.1-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:11d6b0ec246892d85215b0a13ca6e0233cf5284b68f0ac02646427f4ff88a799", size = 328847, upload-time = "2026-06-11T21:55:34.035Z" }, + { url = "https://files.pythonhosted.org/packages/0f/58/dc12f2cd484951524af6e3382c785869b9b3fb5e52ee95ae23add53ee8f9/protobuf-7.35.1-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:b73f9489a4b8b1c9cb1f8ed951c736392592edb24b9d6819f36d2e10b171d5b4", size = 344030, upload-time = "2026-06-11T21:55:34.941Z" }, + { url = "https://files.pythonhosted.org/packages/e4/be/5b3cfe508bfab6761414ff944e3366eb13be4fd71efcd69450f89ba39f43/protobuf-7.35.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:74758715c53d7158fb76caf4f0cfdacc5329a4b1bb994f865d6cf302d413a1c4", size = 327130, upload-time = "2026-06-11T21:55:35.921Z" }, + { url = "https://files.pythonhosted.org/packages/d8/bc/6d6c7ba8709c85f8f2c390b2b118d6fb08a783676a572271851bf45a7d22/protobuf-7.35.1-cp310-abi3-win32.whl", hash = "sha256:353652e4efd0bca5b5fc2656abf8307ef351f0cf938c9eba09f0e09c20a25c30", size = 428945, upload-time = "2026-06-11T21:55:37.034Z" }, + { url = "https://files.pythonhosted.org/packages/0a/19/8d0cb6f20a1ef7b18f1c8986ad5783f22f84cce39c6ce9a6e645ea55192e/protobuf-7.35.1-cp310-abi3-win_amd64.whl", hash = "sha256:230a75ddfc2de4806e56696ce9640c1cdfdb6543b7cfce98d42a4c0a0e7bdb87", size = 439996, upload-time = "2026-06-11T21:55:38.123Z" }, + { url = "https://files.pythonhosted.org/packages/19/c7/5f7c636ec43e0c545e28d1f1db71990108306f7bdcb89f069ba97e428e7f/protobuf-7.35.1-py3-none-any.whl", hash = "sha256:4bc97768d8fe4ad6743c8a19403e314511ed9f6d13205b687e52421c023ac1b9", size = 171659, upload-time = "2026-06-11T21:55:39.155Z" }, +] + +[[package]] +name = "psutil" +version = "7.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/36/5ee6e05c9bd427237b11b3937ad82bb8ad2752d72c6969314590dd0c2f6e/psutil-7.2.2-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486", size = 129090, upload-time = "2026-01-28T18:15:22.168Z" }, + { url = "https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979", size = 129859, upload-time = "2026-01-28T18:15:23.795Z" }, + { url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" }, + { url = "https://files.pythonhosted.org/packages/63/65/37648c0c158dc222aba51c089eb3bdfa238e621674dc42d48706e639204f/psutil-7.2.2-cp36-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e", size = 156997, upload-time = "2026-01-28T18:15:27.794Z" }, + { url = "https://files.pythonhosted.org/packages/8e/13/125093eadae863ce03c6ffdbae9929430d116a246ef69866dad94da3bfbc/psutil-7.2.2-cp36-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8", size = 148972, upload-time = "2026-01-28T18:15:29.342Z" }, + { url = "https://files.pythonhosted.org/packages/04/78/0acd37ca84ce3ddffaa92ef0f571e073faa6d8ff1f0559ab1272188ea2be/psutil-7.2.2-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc", size = 148266, upload-time = "2026-01-28T18:15:31.597Z" }, + { url = "https://files.pythonhosted.org/packages/b4/90/e2159492b5426be0c1fef7acba807a03511f97c5f86b3caeda6ad92351a7/psutil-7.2.2-cp37-abi3-win_amd64.whl", hash = "sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988", size = 137737, upload-time = "2026-01-28T18:15:33.849Z" }, + { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, +] + +[[package]] +name = "pydantic" +version = "2.13.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-types" }, + { name = "pydantic-core" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/18/a5/b60d21ac674192f8ab0ba4e9fd860690f9b4a6e51ca5df118733b487d8d6/pydantic-2.13.4.tar.gz", hash = "sha256:c40756b57adaa8b1efeeced5c196f3f3b7c435f90e84ea7f443901bec8099ef6", size = 844775, upload-time = "2026-05-06T13:43:05.343Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/7b/122376b1fd3c62c1ed9dc80c931ace4844b3c55407b6fb2d199377c9736f/pydantic-2.13.4-py3-none-any.whl", hash = "sha256:45a282cde31d808236fd7ea9d919b128653c8b38b393d1c4ab335c62924d9aba", size = 472262, upload-time = "2026-05-06T13:43:02.641Z" }, +] + +[[package]] +name = "pydantic-core" +version = "2.46.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9d/56/921726b776ace8d8f5db44c4ef961006580d91dc52b803c489fafd1aa249/pydantic_core-2.46.4.tar.gz", hash = "sha256:62f875393d7f270851f20523dd2e29f082bcc82292d66db2b64ea71f64b6e1c1", size = 471464, upload-time = "2026-05-06T13:37:06.98Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/8c/af022f0af448d7747c5154288d46b5f2bc5f17366eaa0e23e9aa04d59f3b/pydantic_core-2.46.4-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:3245406455a5d98187ec35530fd772b1d799b26667980872c8d4614991e2c4a2", size = 2106158, upload-time = "2026-05-06T13:38:57.215Z" }, + { url = "https://files.pythonhosted.org/packages/19/95/6195171e385007300f0f5574592e467c568becce2d937a0b6804f218bc49/pydantic_core-2.46.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:962ccbab7b642487b1d8b7df90ef677e03134cf1fd8880bf698649b22a69371f", size = 1951724, upload-time = "2026-05-06T13:37:02.697Z" }, + { url = "https://files.pythonhosted.org/packages/8e/bc/f47d1ff9cbb1620e1b5b697eef06010035735f07820180e74178226b27b3/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8233f2947cf85404441fd7e0085f53b10c93e0ee78611099b5c7237e36aacbf7", size = 1975742, upload-time = "2026-05-06T13:37:09.448Z" }, + { url = "https://files.pythonhosted.org/packages/5b/11/9b9a5b0306345664a2da6410877af6e8082481b5884b3ddd78d47c6013ce/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3a233125ac121aa3ffba9a2b59edfc4a985a76092dc8279586ab4b71390875e7", size = 2052418, upload-time = "2026-05-06T13:37:38.234Z" }, + { url = "https://files.pythonhosted.org/packages/f1/b7/a65fec226f5d78fc39f4a13c4cc0c768c22b113438f60c14adc9d2865038/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5b712b53160b79a5850310b912a5ef8e57e56947c8ad690c227f5c9d7e561712", size = 2232274, upload-time = "2026-05-06T13:38:27.753Z" }, + { url = "https://files.pythonhosted.org/packages/68/f0/92039db98b907ef49269a8271f67db9cb78ae2fc68062ef7e4e77adb5f61/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9401557acd873c3a7f3eb9383edef8ac4968f9510e340f4808d427e75667e7b4", size = 2309940, upload-time = "2026-05-06T13:38:05.353Z" }, + { url = "https://files.pythonhosted.org/packages/5f/97/2aab507d3d00ca626e8e57c1eac6a79e4e5fbcc63eb99733ff55d1717f65/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:926c9541b14b12b1681dca8a0b75feb510b06c6341b70a8e500c2fdcff837cce", size = 2094516, upload-time = "2026-05-06T13:39:10.577Z" }, + { url = "https://files.pythonhosted.org/packages/22/37/a8aca44d40d737dde2bc05b3c6c07dff0de07ce6f82e9f3167aeaf4d5dea/pydantic_core-2.46.4-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:56cb4851bcaf3d117eddcef4fe66afd750a50274b0da8e22be256d10e5611987", size = 2136854, upload-time = "2026-05-06T13:40:22.59Z" }, + { url = "https://files.pythonhosted.org/packages/24/99/fcef1b79238c06a8cbec70819ac722ba76e02bc8ada9b0fd66eba40da01b/pydantic_core-2.46.4-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:c68fcd102d71ea85c5b2dfac3f4f8476eff42a9e078fd5faefff6d145063536b", size = 2180306, upload-time = "2026-05-06T13:40:10.666Z" }, + { url = "https://files.pythonhosted.org/packages/ae/6c/fc44000918855b42779d007ae63b0532794739027b2f417321cddbc44f6a/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:b2f69dec1725e79a012d920df1707de5caf7ed5e08f3be4435e25803efc47458", size = 2190044, upload-time = "2026-05-06T13:40:43.231Z" }, + { url = "https://files.pythonhosted.org/packages/6b/65/d9cadc9f1920d7a127ad2edba16c1db7916e59719285cd6c94600b0080ba/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:8d0820e8192167f80d88d64038e609c31452eeca865b4e1d9950a27a4609b00b", size = 2329133, upload-time = "2026-05-06T13:39:57.365Z" }, + { url = "https://files.pythonhosted.org/packages/d0/cf/c873d91679f3a30bcf5e7ac280ce5573483e72295307685120d0d5ad3416/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:fbdb89b3e1c94a30cc5edfce477c6e6a5dc4d8f84665b455c27582f211a1c72c", size = 2374464, upload-time = "2026-05-06T13:38:06.976Z" }, + { url = "https://files.pythonhosted.org/packages/9d/1d/8987ad40f65ae1432753072f214fb5c74fe47ffbd0698bb9cbbb585664f8/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:1d8ba486450b14f3b1d63bc521d410ec7565e52f887b9fb671791886436a42f7", size = 2095527, upload-time = "2026-05-06T13:39:52.283Z" }, + { url = "https://files.pythonhosted.org/packages/64/d3/84c282a7eee1d3ac4c0377546ef5a1ea436ce26840d9ac3b7ed54a377507/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:3009f12e4e90b7f88b4f9adb1b0c4a3d58fe7820f3238c190047209d148026df", size = 1936024, upload-time = "2026-05-06T13:40:15.671Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ca/eac61596cdeb4d7e174d3dc0bd8a6238f14f75f97a24e7b7db4c7e7340a0/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ad785e92e6dc634c21555edc8bd6b64957ab844541bcb96a1366c202951ae526", size = 1990696, upload-time = "2026-05-06T13:38:34.717Z" }, + { url = "https://files.pythonhosted.org/packages/fa/c3/7c8b240552251faf6b3a957db200fcfbbcec36763c050428b601e0c9b83b/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0", size = 2147590, upload-time = "2026-05-06T13:39:29.883Z" }, +] + +[[package]] +name = "pygments" +version = "2.20.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" }, +] + +[[package]] +name = "pyrallis" +version = "0.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyyaml" }, + { name = "typing-inspect" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b3/ff/4d865cd0166b3e0b0b09b1f5399114fe27c836b74f45f580d61efecf28dd/pyrallis-0.3.1.tar.gz", hash = "sha256:ab7298f31c633d4858ec3b045a30e6cc9b9f7ab50f21ea93aa3fb28532ac4b6e", size = 33799, upload-time = "2022-03-14T08:23:25.805Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/21/896ce8acd58566f38ffa179b833fd52b4c285fd4e90842ad6b8362372639/pyrallis-0.3.1-py3-none-any.whl", hash = "sha256:632370b563486495f5f9e7caf86cddffab8351214a3bdc60ae0d23b261ebdb59", size = 33419, upload-time = "2022-03-14T08:23:24.193Z" }, +] + +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, +] + +[[package]] +name = "pytz" +version = "2026.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ff/46/dd499ec9038423421951e4fad73051febaa13d2df82b4064f87af8b8c0c3/pytz-2026.2.tar.gz", hash = "sha256:0e60b47b29f21574376f218fe21abc009894a2321ea16c6754f3cad6eb7cdd6a", size = 320861, upload-time = "2026-05-04T01:35:29.667Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/dd/96da98f892250475bdf2328112d7468abdd4acc7b902b6af23f4ed958ea0/pytz-2026.2-py2.py3-none-any.whl", hash = "sha256:04156e608bee23d3792fd45c94ae47fae1036688e75032eea2e3bf0323d1f126", size = 510141, upload-time = "2026-05-04T01:35:27.408Z" }, +] + +[[package]] +name = "pyyaml" +version = "6.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/33/422b98d2195232ca1826284a76852ad5a86fe23e31b009c9886b2d0fb8b2/pyyaml-6.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196", size = 182063, upload-time = "2025-09-25T21:32:11.445Z" }, + { url = "https://files.pythonhosted.org/packages/89/a0/6cf41a19a1f2f3feab0e9c0b74134aa2ce6849093d5517a0c550fe37a648/pyyaml-6.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc09d0aa354569bc501d4e787133afc08552722d3ab34836a80547331bb5d4a0", size = 173973, upload-time = "2025-09-25T21:32:12.492Z" }, + { url = "https://files.pythonhosted.org/packages/ed/23/7a778b6bd0b9a8039df8b1b1d80e2e2ad78aa04171592c8a5c43a56a6af4/pyyaml-6.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28", size = 775116, upload-time = "2025-09-25T21:32:13.652Z" }, + { url = "https://files.pythonhosted.org/packages/65/30/d7353c338e12baef4ecc1b09e877c1970bd3382789c159b4f89d6a70dc09/pyyaml-6.0.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5fdec68f91a0c6739b380c83b951e2c72ac0197ace422360e6d5a959d8d97b2c", size = 844011, upload-time = "2025-09-25T21:32:15.21Z" }, + { url = "https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc", size = 807870, upload-time = "2025-09-25T21:32:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/05/c0/b3be26a015601b822b97d9149ff8cb5ead58c66f981e04fedf4e762f4bd4/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e", size = 761089, upload-time = "2025-09-25T21:32:17.56Z" }, + { url = "https://files.pythonhosted.org/packages/be/8e/98435a21d1d4b46590d5459a22d88128103f8da4c2d4cb8f14f2a96504e1/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea", size = 790181, upload-time = "2025-09-25T21:32:18.834Z" }, + { url = "https://files.pythonhosted.org/packages/74/93/7baea19427dcfbe1e5a372d81473250b379f04b1bd3c4c5ff825e2327202/pyyaml-6.0.3-cp312-cp312-win32.whl", hash = "sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5", size = 137658, upload-time = "2025-09-25T21:32:20.209Z" }, + { url = "https://files.pythonhosted.org/packages/86/bf/899e81e4cce32febab4fb42bb97dcdf66bc135272882d1987881a4b519e9/pyyaml-6.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b", size = 154003, upload-time = "2025-09-25T21:32:21.167Z" }, + { url = "https://files.pythonhosted.org/packages/1a/08/67bd04656199bbb51dbed1439b7f27601dfb576fb864099c7ef0c3e55531/pyyaml-6.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd", size = 140344, upload-time = "2025-09-25T21:32:22.617Z" }, +] + +[[package]] +name = "qwen-vl-utils" +version = "0.0.14" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "av" }, + { name = "packaging" }, + { name = "pillow" }, + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b6/b1/ad4fc2260a3badd278b38d642f3b987412f1f6682f0ef2b31b0572d5caa8/qwen_vl_utils-0.0.14.tar.gz", hash = "sha256:9c7cad5ae803b3a10f8bb7194deb12aeacdd032f92f4224e880c73587a7346ad", size = 8453, upload-time = "2025-09-23T09:38:57.532Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c4/43/80f67e0336cb2fc725f8e06f7fe35c1d0fe946f4d2b8b2175e797e07349e/qwen_vl_utils-0.0.14-py3-none-any.whl", hash = "sha256:5e28657bfd031e56bd447c5901b58ddfc3835285ed100f4c56580e0ade054e96", size = 8120, upload-time = "2025-09-23T09:38:56.297Z" }, +] + +[[package]] +name = "regex" +version = "2026.7.19" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/20/98/04b13f1ddfb63158025291c02e03eb42fbb7acb51d091d541050eb4e35e8/regex-2026.7.19.tar.gz", hash = "sha256:7e77b324909c1617cbb4c668677e2c6ae13f44d7c1de0d4f15f2e3c10f3315b5", size = 416440, upload-time = "2026-07-19T00:19:48.923Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3b/b9/d11d7e501ac8fd7d617684423ebb9561e0b998481c1e4cbc0cb212c5d74a/regex-2026.7.19-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:2cc3460cedf7579948486eab03bc9ad7089df4d7281c0f47f4afe03e8d13f02d", size = 496778, upload-time = "2026-07-19T00:17:05.677Z" }, + { url = "https://files.pythonhosted.org/packages/3f/a9/a5ab6f312f24318019170dc485d5421fe4f89e43a98640da50d95a8a7041/regex-2026.7.19-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:0e9554c8785eac5cffe6300f69a91f58ba72bc88a5f8d661235ad7c6aa5b8ccd", size = 297122, upload-time = "2026-07-19T00:17:07.59Z" }, + { url = "https://files.pythonhosted.org/packages/b3/63/4cab4d7f2d384a144d420b763d97674cb70619c878ea6fcd7640d0e62143/regex-2026.7.19-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d7da47a0f248977f08e2cb659ff3c17ddc13a4d39b3a7baa0a81bf5b415430f6", size = 292009, upload-time = "2026-07-19T00:17:09.648Z" }, + { url = "https://files.pythonhosted.org/packages/22/85/102a81b218298957d4ea7d2f084fae537a71add9d6ff93c8e67284c5f45e/regex-2026.7.19-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:93db40c8de0815baab96a06e08a984bac71f989d13bab789e382158c5d426797", size = 796708, upload-time = "2026-07-19T00:17:11.542Z" }, + { url = "https://files.pythonhosted.org/packages/78/b5/dc136af5629938a037cd2b304c12240e132ec92f38be8ff9cc89af2a1f2d/regex-2026.7.19-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:66bd62c59a5427746e8c44becae1d9b99d22fb13f30f492083dfb9ad7c45cc18", size = 865651, upload-time = "2026-07-19T00:17:13.312Z" }, + { url = "https://files.pythonhosted.org/packages/e0/75/67402ae3cd9c8c988a4c805d15ee3eef015e7ca4cb112cf3e640fc1f4153/regex-2026.7.19-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:1649eb39fcc9ea80c4d2f110fde2b8ab2aef3877b98f02ab9b14e961f418c511", size = 911756, upload-time = "2026-07-19T00:17:15.015Z" }, + { url = "https://files.pythonhosted.org/packages/2a/8e/096d00c7c480ef2ff4265349b14e2261d4ab787ba1f74e2e80d1c58079c3/regex-2026.7.19-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9dce8ec9695f531a1b8a6f314fd4b393adcccf2ea861db480cdf97a301d01a68", size = 801798, upload-time = "2026-07-19T00:17:17.208Z" }, + { url = "https://files.pythonhosted.org/packages/f0/41/e7ecac6edb5722417f85cc67eaf386322fbe8acf6918ec2fdc37c20dd9d0/regex-2026.7.19-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:3080a7fd38ef049bd489e01c970c97dd84ff446a885b0f1f6b26d9b1ad13ce11", size = 776933, upload-time = "2026-07-19T00:17:19.347Z" }, + { url = "https://files.pythonhosted.org/packages/6f/69/03c9b3f058d66403e0ca2c938696e81d51cd4c6d47ec5265f02f96948d9a/regex-2026.7.19-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1d793a7988e04fcb1e2e135567443d82173225d657419ec09414a9b5a145b986", size = 784338, upload-time = "2026-07-19T00:17:21.057Z" }, + { url = "https://files.pythonhosted.org/packages/f6/f7/b38ab3d43f284afbb618fcd15d0e77eb786ae461ce1f6bc7494619ddc0f2/regex-2026.7.19-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:e8b0abe7d870f53ca5143895fef7d1041a0c831a140d3dc2c760dd7ba25d4a8b", size = 860452, upload-time = "2026-07-19T00:17:23.119Z" }, + { url = "https://files.pythonhosted.org/packages/15/5c/ff60ef0571121714f3cf9920bc183071e384a10b556d042e0fdb06cc07a5/regex-2026.7.19-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:4e5413bd5f13d3a4e3539ca98f70f75e7fca92518dd7f117f030ebedd10b60cb", size = 765958, upload-time = "2026-07-19T00:17:24.81Z" }, + { url = "https://files.pythonhosted.org/packages/aa/0f/bd34021162c0ab47f9a315bd56cd5642e920c8e5668a75ef6c6a6fca590d/regex-2026.7.19-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:73b133a9e6fb512858e7f065e96f1180aa46646bc74a83aea62f1d314f3dd035", size = 851765, upload-time = "2026-07-19T00:17:26.993Z" }, + { url = "https://files.pythonhosted.org/packages/2a/20/a2ca43edade0595cccfdc98636739f536d9e26898e7dbddc2b9e98898953/regex-2026.7.19-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:dbe6493fbd27321b1d1f2dd4f5c7e5bd4d8b1d7cab7f32fd67db3d0b2ed8248a", size = 789714, upload-time = "2026-07-19T00:17:28.699Z" }, + { url = "https://files.pythonhosted.org/packages/5d/47/e02db4015d424fc83c00ea0ac8c5e5ec14397943de9abf909d5ce3a25931/regex-2026.7.19-cp312-cp312-win32.whl", hash = "sha256:ddd67571c10869f65a5d7dde536d1e066e306cc90de57d7de4d5f34802428bb5", size = 267157, upload-time = "2026-07-19T00:17:31.051Z" }, + { url = "https://files.pythonhosted.org/packages/08/8e/c780c131f79b42ed22d1bd7da4096c2c35f813e835acd02ef0f018bd892c/regex-2026.7.19-cp312-cp312-win_amd64.whl", hash = "sha256:e30d40268a28d54ce0437031750497004c22602b8e3ab891f759b795a003b312", size = 277777, upload-time = "2026-07-19T00:17:32.848Z" }, + { url = "https://files.pythonhosted.org/packages/3e/4c/e4d7e086449bdf379d89774bf1f89dc4a41943f3c5a6125a03905b34b5fb/regex-2026.7.19-cp312-cp312-win_arm64.whl", hash = "sha256:de9208bb427130c82a5dbfd104f92c8876fc9559278c880b3002755bbbe9c83d", size = 277136, upload-time = "2026-07-19T00:17:34.803Z" }, +] + +[[package]] +name = "requests" +version = "2.34.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "charset-normalizer" }, + { name = "idna" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ac/c3/e2a2b89f2d3e2179abd6d00ebd70bff6273f37fb3e0cc209f48b39d00cbf/requests-2.34.2.tar.gz", hash = "sha256:f288924cae4e29463698d6d60bc6a4da69c89185ad1e0bcc4104f584e960b9ed", size = 142856, upload-time = "2026-05-14T19:25:27.735Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, +] + +[[package]] +name = "rich" +version = "15.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c0/8f/0722ca900cc807c13a6a0c696dacf35430f72e0ec571c4275d2371fca3e9/rich-15.0.0.tar.gz", hash = "sha256:edd07a4824c6b40189fb7ac9bc4c52536e9780fbbfbddf6f1e2502c31b068c36", size = 230680, upload-time = "2026-04-12T08:24:00.75Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" }, +] + +[[package]] +name = "s3transfer" +version = "0.19.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "botocore" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/65/da/4bef7ce7bb989b222aa4785a413896dbec53306dfc59c6ce7d16a7ffbd6a/s3transfer-0.19.1.tar.gz", hash = "sha256:d3d6371dc3f1e5c5427b2b457bcf13bcf87bec334c95aed18642eae61f6926f3", size = 165354, upload-time = "2026-07-10T19:32:04.849Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/23/e84c64ad0e8bc59cd1b2ef98def848deff0ef3456c542afe74d51e9e8c85/s3transfer-0.19.1-py3-none-any.whl", hash = "sha256:d5fd7005ee39307455ad5f310b5ea67f4b1960d7fed5b3671ee50c249de675de", size = 90072, upload-time = "2026-07-10T19:32:03.673Z" }, +] + +[[package]] +name = "safetensors" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/45/06/f955dbbb1859e3bd23c8ac6141af5106e7ad5fedec4a3a6e3d60f94b7001/safetensors-0.8.0.tar.gz", hash = "sha256:fabaf3e0f18a6618d9b36560682562157f77c2b71fcffc7b432be2baed9d753d", size = 325846, upload-time = "2026-06-09T07:52:25.563Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/a0/f718cda65b05407d228f97602cf60dca269c979867aa5beb25410de26cd3/safetensors-0.8.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c554f85858e05226d3c2828e32395e677434685d6d94594a41643361c5e837f0", size = 473568, upload-time = "2026-06-09T07:52:18.829Z" }, + { url = "https://files.pythonhosted.org/packages/f5/b1/fa7c600e7dceae12e9606c7578cbc9ff1e1ed55844883ee5c92205e86226/safetensors-0.8.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c80201d22cbf405b80647a60ada77bba06c8fba2da2743ba1e89cdcc39a81f25", size = 484562, upload-time = "2026-06-09T07:52:17.518Z" }, + { url = "https://files.pythonhosted.org/packages/09/7d/65a7de0af421317bb36a067241e4235fff194eed60b961ed6d3f59a3fc60/safetensors-0.8.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7a46e5ff292c356d6991e60942ba7f79817682d3a2cef0702136448cb9c4d235", size = 502844, upload-time = "2026-06-09T07:52:07.624Z" }, + { url = "https://files.pythonhosted.org/packages/91/4f/3175c9d75634e0e0dda0082794193521035edd7c70a6f212bf33ca06ddf4/safetensors-0.8.0-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4124502b78f03534117c848f87a39b8f31e577b15eff423bf8bfb95f2a8c30d0", size = 511823, upload-time = "2026-06-09T07:52:09.565Z" }, + { url = "https://files.pythonhosted.org/packages/20/87/846c289e7aa2299eff406335717cf43ce8777194ece8aad75772e0411615/safetensors-0.8.0-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7bc0a787ba8a35be368ee3574edfa2b1ad389eebd0a72e482ae275490e3f6c98", size = 633461, upload-time = "2026-06-09T07:52:11.128Z" }, + { url = "https://files.pythonhosted.org/packages/76/22/8d64d9df2c45d5ded401df889d0ad90882804ca172d79ec4f0df8f727fe0/safetensors-0.8.0-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:040070828e36dc8e122178bbbd5830ff9e97920affb84cbe0f46442497bed358", size = 545148, upload-time = "2026-06-09T07:52:13.603Z" }, + { url = "https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd6f3f93c9a0a7cc2788ee63fb763353d4bd2e89b0751bc78fcf7dda00bea774", size = 516040, upload-time = "2026-06-09T07:52:16.29Z" }, + { url = "https://files.pythonhosted.org/packages/46/fb/cdaed17ceb2948784fd9c36b6fd3e951b608547cea81a48e8ee6f8cfdfcb/safetensors-0.8.0-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:fcdd41ec4628fee5799f807c73c353629130fbd942aa23d83c623dd6c9d52d78", size = 513832, upload-time = "2026-06-09T07:52:12.37Z" }, + { url = "https://files.pythonhosted.org/packages/0d/49/1e15de264dcc3b77943d2d0c56a95809956883b1c2d6d585c792523f180b/safetensors-0.8.0-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:8e9f537aa183a38ace122d27303dcd986b26bd2a7591f9181d7f0c396f4677ca", size = 559930, upload-time = "2026-06-09T07:52:14.743Z" }, + { url = "https://files.pythonhosted.org/packages/2a/43/bf38443278eab4b1be1fce2931e2b012ad9cb7df52ada751d0aab8f7659a/safetensors-0.8.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:87eec7ffed2b809f05a398a8becb7d013f19f7837cd15d9748580d6cf30dbaf4", size = 678670, upload-time = "2026-06-09T07:52:20.032Z" }, + { url = "https://files.pythonhosted.org/packages/72/e3/68cd3fa5b48488e84add63e04cb12f3bc28ae4638c06d4508c6e88823d0e/safetensors-0.8.0-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:4a95ae2b05d7726d751da4ebf626a2ca782b706e101bd894c95bc2450b1cffcc", size = 786679, upload-time = "2026-06-09T07:52:21.322Z" }, + { url = "https://files.pythonhosted.org/packages/29/4b/1c19c509d56e01f4fbb3d0a2e597450f6cc04d1d56cf52defb0a62dfd715/safetensors-0.8.0-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:3ae091f16662658bdc019a4ff6cb4c085bb7d725eb5978b183ffd265863b6d2d", size = 765683, upload-time = "2026-06-09T07:52:22.594Z" }, + { url = "https://files.pythonhosted.org/packages/27/43/41c1621732edd934d868a00d1b891584c892a7b62a9aab82ea5a0a5623ee/safetensors-0.8.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8e080062fcde23be189565e1c3305d16751a218ecf9412c8601e64204eb6f846", size = 722361, upload-time = "2026-06-09T07:52:23.924Z" }, + { url = "https://files.pythonhosted.org/packages/8e/3f/73ccf82579412b4a71c4ca673f10b5f1f888d7cf5af7fe24f27d30307be4/safetensors-0.8.0-cp310-abi3-win32.whl", hash = "sha256:2ddf52eac562eda224f99acfa7889d02968c1fd59a5b011ae7d8137c37e9c02d", size = 342401, upload-time = "2026-06-09T07:52:28.895Z" }, + { url = "https://files.pythonhosted.org/packages/1b/6d/3fba214c1e5e0f69991677ec3bc17023f0421776975e1de0c682dca475e2/safetensors-0.8.0-cp310-abi3-win_amd64.whl", hash = "sha256:096ec1a98435df7beb08853bb5aa9081a84f23d0adc67ed1a0a10550f608373f", size = 355540, upload-time = "2026-06-09T07:52:27.832Z" }, + { url = "https://files.pythonhosted.org/packages/8d/fc/7eedc3510d97878876e32774eebbeb61c43f148a96e915c84229a3e967aa/safetensors-0.8.0-cp310-abi3-win_arm64.whl", hash = "sha256:f7838e5135a406ad3e02efdcb8cf2e5397d368b0154537c4fec682dbc544d452", size = 340500, upload-time = "2026-06-09T07:52:26.745Z" }, +] + +[[package]] +name = "sana-wm-parity-check" +version = "0.0.0" +source = { virtual = "." } +dependencies = [ + { name = "accelerate" }, + { name = "diffusers" }, + { name = "einops" }, + { name = "flash-linear-attention" }, + { name = "flashdreams" }, + { name = "flashdreams-sana-wm" }, + { name = "ftfy" }, + { name = "huggingface-hub" }, + { name = "imageio", extra = ["ffmpeg"] }, + { name = "numpy" }, + { name = "omegaconf" }, + { name = "opencv-python-headless" }, + { name = "pillow" }, + { name = "protobuf" }, + { name = "pyrallis" }, + { name = "pytz" }, + { name = "pyyaml" }, + { name = "qwen-vl-utils" }, + { name = "safetensors" }, + { name = "scipy" }, + { name = "sentencepiece" }, + { name = "termcolor" }, + { name = "timm" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, + { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32'" }, + { name = "torchvision" }, + { name = "tqdm" }, + { name = "transformers" }, + { name = "triton", marker = "sys_platform == 'linux'" }, +] + +[package.optional-dependencies] +quant = [ + { name = "transformer-engine", extra = ["core-cu13", "pytorch"], marker = "sys_platform != 'win32'" }, +] + +[package.metadata] +requires-dist = [ + { name = "accelerate", specifier = ">=1.3" }, + { name = "diffusers", specifier = ">=0.37" }, + { name = "einops", specifier = ">=0.7" }, + { name = "flash-linear-attention", specifier = ">=0.4.2" }, + { name = "flashdreams", editable = "../../../../flashdreams" }, + { name = "flashdreams-sana-wm", editable = "../../" }, + { name = "ftfy", specifier = ">=6.0" }, + { name = "huggingface-hub", specifier = ">=0.36" }, + { name = "imageio", extras = ["ffmpeg"], specifier = ">=2.31" }, + { name = "numpy", specifier = ">=1.24,<2.5" }, + { name = "omegaconf", specifier = ">=2.3" }, + { name = "opencv-python-headless", specifier = ">=4.8" }, + { name = "pillow", specifier = ">=10" }, + { name = "protobuf", specifier = ">=7.35.0,<8" }, + { name = "pyrallis", specifier = ">=0.3" }, + { name = "pytz", specifier = ">=2024.0" }, + { name = "pyyaml", specifier = ">=6.0" }, + { name = "qwen-vl-utils", specifier = ">=0.0.8" }, + { name = "safetensors", specifier = ">=0.5" }, + { name = "scipy", specifier = ">=1.11" }, + { name = "sentencepiece", specifier = ">=0.2" }, + { name = "termcolor", specifier = ">=2.4" }, + { name = "timm", specifier = ">=0.6.13" }, + { name = "torch", specifier = ">=2.11" }, + { name = "torchvision", specifier = ">=0.26" }, + { name = "tqdm", specifier = ">=4.60" }, + { name = "transformer-engine", extras = ["pytorch", "core-cu13"], marker = "sys_platform != 'win32' and extra == 'quant'", specifier = ">=2.12" }, + { name = "transformers", specifier = ">=5.0,<6" }, + { name = "triton", marker = "sys_platform == 'linux'", specifier = ">=3.6" }, +] +provides-extras = ["quant"] + +[[package]] +name = "scipy" +version = "1.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a7/25/c2700dfaf6442b4effaa91af24ebce5dc9d31bb4a69706313aae70d72cd0/scipy-1.18.0.tar.gz", hash = "sha256:67b2ad2ad54c72ca6d04975a9b2df8c3638c34ddd5b28738e94fc2b57929d378", size = 30774447, upload-time = "2026-06-19T15:01:43.456Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6a/19/ca10ead60b0acc80b2b833c2c4a4f2ff753d0f58b811f70d911c7e94a25c/scipy-1.18.0-cp312-cp312-macosx_10_15_x86_64.whl", hash = "sha256:7bd21faaf5a1a3b2eff922d02db5f191b99a6518db9078a8fb23169f6d22259a", size = 31056519, upload-time = "2026-06-19T14:59:45.203Z" }, + { url = "https://files.pythonhosted.org/packages/96/72/1e6442a00cd2924d361aa1b642ab6373ec35c6fabf311a760be9f76e0f13/scipy-1.18.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:265915e79107de9f946b855e50d7470d5893ec3f54b342e1aa6201cbdcd8bb6b", size = 28681889, upload-time = "2026-06-19T14:59:48.103Z" }, + { url = "https://files.pythonhosted.org/packages/9b/2d/11dd93d21e147a73ba22bd75c0b9208d3a2e0ec76d53170ce7d9029b1015/scipy-1.18.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:9ab7b758be6940954a713ee466e2043e9f6e2ed965c1fce5c91039f4be3d90a9", size = 20423580, upload-time = "2026-06-19T14:59:50.665Z" }, + { url = "https://files.pythonhosted.org/packages/9c/01/93552f75e0d2a7dd115a45e59209c51e8d514daff02fc887d2623be06fe1/scipy-1.18.0-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:97b6cddaaee0a779ef6b5ca83c9604b27cc16b2b8fc22c142652df8793319fb8", size = 23054441, upload-time = "2026-06-19T14:59:53.564Z" }, + { url = "https://files.pythonhosted.org/packages/3c/23/21f5e703643d66f21faa6b4c73195bfcad70c55efcb4f1ab327cd7c4101a/scipy-1.18.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:52a96e21517c7292375c0e27dd796a811f03fcea5fd4d108fdfea8145dcf17ab", size = 33968720, upload-time = "2026-06-19T14:59:56.415Z" }, + { url = "https://files.pythonhosted.org/packages/dd/aa/1b939f6c67ed68635bb538e6752d3dacc02f66535182e939a89581a44e9c/scipy-1.18.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1f55797419e16e7f30cf88ffb3113ce0467f00cfe3f70d5c281730b21769bfc2", size = 35287115, upload-time = "2026-06-19T14:59:59.411Z" }, + { url = "https://files.pythonhosted.org/packages/b6/ff/eec46be7e9234208f801062b53e1983085eddebd693f6c9bfb03b459830d/scipy-1.18.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ad033410e2e0672ffdc1042110cef20e1c46f8fd0616cee1d44d8d58fad8fc11", size = 35577989, upload-time = "2026-06-19T15:00:02.235Z" }, + { url = "https://files.pythonhosted.org/packages/84/ca/210d4759c7210bb7d269437421959b39a33434e2776b60c5cb8a763bb30a/scipy-1.18.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:4a55985d54c769c872e64b7f4c8a81cc30ef700cc04296abbbf3705439c126de", size = 37421717, upload-time = "2026-06-19T15:00:05.102Z" }, + { url = "https://files.pythonhosted.org/packages/2b/54/9a9edb45345bd6744da5ddfb6628e5d5185920494c6a67ec45b6381004cb/scipy-1.18.0-cp312-cp312-win_amd64.whl", hash = "sha256:71ccc8faa2dd16ac310233203474a8b5cb67f10dedd54a3116d34943f4b19132", size = 36597428, upload-time = "2026-06-19T15:00:08.112Z" }, + { url = "https://files.pythonhosted.org/packages/99/0e/33f32a2a58987e26aec0f7df252cbbad1e90ae77bdbc76f40dd4ed0cf0ea/scipy-1.18.0-cp312-cp312-win_arm64.whl", hash = "sha256:d88363fd9d8fbd3511bd273f1a49efb2a540773ddf92a91d57498ce7dd7f3e76", size = 24351481, upload-time = "2026-06-19T15:00:11.103Z" }, +] + +[[package]] +name = "sentencepiece" +version = "0.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cc/33/ea3cb3839607eb175da835244a798f797f478c5ddf0e8ecdf57ea85a4c70/sentencepiece-0.2.2.tar.gz", hash = "sha256:3d2b5e824b5622038dc7b490897efe05ebbbb9e7350fc142f3ecc8789ef9bdf6", size = 8218435, upload-time = "2026-07-12T08:39:34.701Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/13/7a562289c8d5b49ebdf3f9c1e8ab67cf14a8743b1d90c8f406bfdec36b72/sentencepiece-0.2.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:1edb10e520e4bddf74d85b0f5ae74cc2d60c2b448885080bfb618bc2b3a49f6b", size = 2188384, upload-time = "2026-07-12T08:38:28.486Z" }, + { url = "https://files.pythonhosted.org/packages/85/d1/912f14fd5eae168aba726ffb6a9a2dc1c71fe7676c53da6f5c442b886d4a/sentencepiece-0.2.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:f7c06c751c19d923435a54bff4f7e66e728fad160e8da28254f133abc9725820", size = 1441553, upload-time = "2026-07-12T08:38:30.552Z" }, + { url = "https://files.pythonhosted.org/packages/bd/44/caa9cab5f261a019e2808bc5046152775dc57352ba9cbae7525e9e7a1ed4/sentencepiece-0.2.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:38111ed1f79268f399c505028023d5eaaf0ab4e5eafceb709468b0d3323e7838", size = 1347176, upload-time = "2026-07-12T08:38:32.211Z" }, + { url = "https://files.pythonhosted.org/packages/19/90/cd798935668cff71d309d8ff10385844ecf216b1fe454f1993ed8bf2cb91/sentencepiece-0.2.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cbce24284f51f71d10a42b7b9c964dcb9048b28f1c8e5db40bcbcb6f428cba6a", size = 1325200, upload-time = "2026-07-12T08:38:33.689Z" }, + { url = "https://files.pythonhosted.org/packages/b6/2d/37e3da037318a70066ded0d51bc2a7f35491ae6338dd993d5eb1503fc3b5/sentencepiece-0.2.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c8a168b040bc61681293f79a949b5d911c8e25086f4260285b8d97ab5f1195da", size = 1397736, upload-time = "2026-07-12T08:38:35.771Z" }, + { url = "https://files.pythonhosted.org/packages/8d/11/753fca2e6b109be3ab7867abf357dfe48677fe726ae5a5363d0b54ca9450/sentencepiece-0.2.2-cp312-cp312-win_amd64.whl", hash = "sha256:7c6e7bf684dc12145bfa685d3060beaea55139134ba848289bee514ed42e7383", size = 1248030, upload-time = "2026-07-12T08:38:37.604Z" }, + { url = "https://files.pythonhosted.org/packages/e2/0a/70efbe861ca182d7d4b6e1a20f58e043400848fa9f2915229f082e221648/sentencepiece-0.2.2-cp312-cp312-win_arm64.whl", hash = "sha256:76ff5814db72e7462dece042d7593cdf102b8ec82c2b1cc201a2add34ee3050d", size = 1187325, upload-time = "2026-07-12T08:38:39.348Z" }, +] + +[[package]] +name = "setuptools" +version = "83.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/34/26/f5d29e25ffdb535afef2d35cdb55b325298f96debd670da4c325e08d70f4/setuptools-83.0.0.tar.gz", hash = "sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef", size = 1154254, upload-time = "2026-07-04T15:31:22.699Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5d/40/e1e72872c6354b306daef1703549e8e83b4d43cfea356311bf722a043752/setuptools-83.0.0-py3-none-any.whl", hash = "sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3", size = 1008090, upload-time = "2026-07-04T15:31:20.885Z" }, +] + +[[package]] +name = "shellingham" +version = "1.5.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/58/15/8b3609fd3830ef7b27b655beb4b4e9c62313a4e8da8c676e142cc210d58e/shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de", size = 10310, upload-time = "2023-10-24T04:13:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, +] + +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, +] + +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + +[[package]] +name = "termcolor" +version = "3.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/46/79/cf31d7a93a8fdc6aa0fbb665be84426a8c5a557d9240b6239e9e11e35fc5/termcolor-3.3.0.tar.gz", hash = "sha256:348871ca648ec6a9a983a13ab626c0acce02f515b9e1983332b17af7979521c5", size = 14434, upload-time = "2025-12-29T12:55:21.882Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/33/d1/8bb87d21e9aeb323cc03034f5eaf2c8f69841e40e4853c2627edf8111ed3/termcolor-3.3.0-py3-none-any.whl", hash = "sha256:cf642efadaf0a8ebbbf4bc7a31cec2f9b5f21a9f726f4ccbb08192c9c26f43a5", size = 7734, upload-time = "2025-12-29T12:55:20.718Z" }, +] + +[[package]] +name = "timm" +version = "1.0.28" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, + { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32'" }, + { name = "torchvision" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/03/e41389ac641747bfec48d016fde8be1eade1901e6f2c1aedcb0c8cb4b5d9/timm-1.0.28.tar.gz", hash = "sha256:3789d313fdd5541a327b60180d70dbb4bdec73db8ff0655e413db3c3d134a9a4", size = 2451413, upload-time = "2026-07-11T17:24:32.615Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/76/de1bfac17d183c49c6d0887903d3064ced51cf1d9ba7a8d611c1a8808c4f/timm-1.0.28-py3-none-any.whl", hash = "sha256:e577b88da96b3a722ea5e2f042455ce6f715d398304d8e63b17d126ed7d89968", size = 2597944, upload-time = "2026-07-11T17:24:30.869Z" }, +] + +[[package]] +name = "tokenizers" +version = "0.22.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/73/6f/f80cfef4a312e1fb34baf7d85c72d4411afde10978d4657f8cdd811d3ccc/tokenizers-0.22.2.tar.gz", hash = "sha256:473b83b915e547aa366d1eee11806deaf419e17be16310ac0a14077f1e28f917", size = 372115, upload-time = "2026-01-05T10:45:15.988Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/92/97/5dbfabf04c7e348e655e907ed27913e03db0923abb5dfdd120d7b25630e1/tokenizers-0.22.2-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:544dd704ae7238755d790de45ba8da072e9af3eea688f698b137915ae959281c", size = 3100275, upload-time = "2026-01-05T10:41:02.158Z" }, + { url = "https://files.pythonhosted.org/packages/2e/47/174dca0502ef88b28f1c9e06b73ce33500eedfac7a7692108aec220464e7/tokenizers-0.22.2-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:1e418a55456beedca4621dbab65a318981467a2b188e982a23e117f115ce5001", size = 2981472, upload-time = "2026-01-05T10:41:00.276Z" }, + { url = "https://files.pythonhosted.org/packages/d6/84/7990e799f1309a8b87af6b948f31edaa12a3ed22d11b352eaf4f4b2e5753/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2249487018adec45d6e3554c71d46eb39fa8ea67156c640f7513eb26f318cec7", size = 3290736, upload-time = "2026-01-05T10:40:32.165Z" }, + { url = "https://files.pythonhosted.org/packages/78/59/09d0d9ba94dcd5f4f1368d4858d24546b4bdc0231c2354aa31d6199f0399/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:25b85325d0815e86e0bac263506dd114578953b7b53d7de09a6485e4a160a7dd", size = 3168835, upload-time = "2026-01-05T10:40:38.847Z" }, + { url = "https://files.pythonhosted.org/packages/47/50/b3ebb4243e7160bda8d34b731e54dd8ab8b133e50775872e7a434e524c28/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bfb88f22a209ff7b40a576d5324bf8286b519d7358663db21d6246fb17eea2d5", size = 3521673, upload-time = "2026-01-05T10:40:56.614Z" }, + { url = "https://files.pythonhosted.org/packages/e0/fa/89f4cb9e08df770b57adb96f8cbb7e22695a4cb6c2bd5f0c4f0ebcf33b66/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1c774b1276f71e1ef716e5486f21e76333464f47bece56bbd554485982a9e03e", size = 3724818, upload-time = "2026-01-05T10:40:44.507Z" }, + { url = "https://files.pythonhosted.org/packages/64/04/ca2363f0bfbe3b3d36e95bf67e56a4c88c8e3362b658e616d1ac185d47f2/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:df6c4265b289083bf710dff49bc51ef252f9d5be33a45ee2bed151114a56207b", size = 3379195, upload-time = "2026-01-05T10:40:51.139Z" }, + { url = "https://files.pythonhosted.org/packages/2e/76/932be4b50ef6ccedf9d3c6639b056a967a86258c6d9200643f01269211ca/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:369cc9fc8cc10cb24143873a0d95438bb8ee257bb80c71989e3ee290e8d72c67", size = 3274982, upload-time = "2026-01-05T10:40:58.331Z" }, + { url = "https://files.pythonhosted.org/packages/1d/28/5f9f5a4cc211b69e89420980e483831bcc29dade307955cc9dc858a40f01/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:29c30b83d8dcd061078b05ae0cb94d3c710555fbb44861139f9f83dcca3dc3e4", size = 9478245, upload-time = "2026-01-05T10:41:04.053Z" }, + { url = "https://files.pythonhosted.org/packages/6c/fb/66e2da4704d6aadebf8cb39f1d6d1957df667ab24cff2326b77cda0dcb85/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:37ae80a28c1d3265bb1f22464c856bd23c02a05bb211e56d0c5301a435be6c1a", size = 9560069, upload-time = "2026-01-05T10:45:10.673Z" }, + { url = "https://files.pythonhosted.org/packages/16/04/fed398b05caa87ce9b1a1bb5166645e38196081b225059a6edaff6440fac/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:791135ee325f2336f498590eb2f11dc5c295232f288e75c99a36c5dbce63088a", size = 9899263, upload-time = "2026-01-05T10:45:12.559Z" }, + { url = "https://files.pythonhosted.org/packages/05/a1/d62dfe7376beaaf1394917e0f8e93ee5f67fea8fcf4107501db35996586b/tokenizers-0.22.2-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:38337540fbbddff8e999d59970f3c6f35a82de10053206a7562f1ea02d046fa5", size = 10033429, upload-time = "2026-01-05T10:45:14.333Z" }, + { url = "https://files.pythonhosted.org/packages/fd/18/a545c4ea42af3df6effd7d13d250ba77a0a86fb20393143bbb9a92e434d4/tokenizers-0.22.2-cp39-abi3-win32.whl", hash = "sha256:a6bf3f88c554a2b653af81f3204491c818ae2ac6fbc09e76ef4773351292bc92", size = 2502363, upload-time = "2026-01-05T10:45:20.593Z" }, + { url = "https://files.pythonhosted.org/packages/65/71/0670843133a43d43070abeb1949abfdef12a86d490bea9cd9e18e37c5ff7/tokenizers-0.22.2-cp39-abi3-win_amd64.whl", hash = "sha256:c9ea31edff2968b44a88f97d784c2f16dc0729b8b143ed004699ebca91f05c48", size = 2747786, upload-time = "2026-01-05T10:45:18.411Z" }, + { url = "https://files.pythonhosted.org/packages/72/f4/0de46cfa12cdcbcd464cc59fde36912af405696f687e53a091fb432f694c/tokenizers-0.22.2-cp39-abi3-win_arm64.whl", hash = "sha256:9ce725d22864a1e965217204946f830c37876eee3b2ba6fc6255e8e903d5fcbc", size = 2612133, upload-time = "2026-01-05T10:45:17.232Z" }, +] + +[[package]] +name = "torch" +version = "2.13.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "sys_platform != 'win32'", +] +dependencies = [ + { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "triton", marker = "sys_platform == 'linux'" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c4/3a/ed0f4d4d1dcde03bced7aac9a28e800abcdc0cbd06b6775044c9fbd877b7/torch-2.13.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:2fe228aba290d14b9f31b049be550dbd469c3fd3013d7a19705b30454da97027", size = 111213045, upload-time = "2026-07-08T16:05:22.997Z" }, + { url = "https://files.pythonhosted.org/packages/df/a9/f6a2a4d763ff1df02e9a64c477029db614295bc9367f4131223791ccc243/torch-2.13.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:572df8be8ffb4599c88cbd6a0726f1f854f4da65d2e3c09f0e2c2283333cd6d4", size = 427210998, upload-time = "2026-07-08T16:04:37.708Z" }, + { url = "https://files.pythonhosted.org/packages/f3/82/fea946351658e6534db52d2cc12bc53087cbf87f9440c5f180f367c1950b/torch-2.13.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:796633c4cdf0fe2cdced72d8f88f22e73dbcfce83132763162f6d4bff13b820b", size = 526605292, upload-time = "2026-07-08T16:04:22.81Z" }, +] + +[[package]] +name = "torch" +version = "2.13.0+cu130" +source = { registry = "https://download.pytorch.org/whl/cu130" } +resolution-markers = [ + "sys_platform == 'win32'", +] +dependencies = [ + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.13.0%2Bcu130-cp312-cp312-win_amd64.whl", hash = "sha256:2efab1e83604ca628c6d85b9e188c153690980498d1297081a9dad704919303c", upload-time = "2026-07-08T20:26:27Z" }, +] + +[[package]] +name = "torchvision" +version = "0.28.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "pillow" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, + { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/15/49/c1cab1ecbb3ff1a380a3f99283db1dee61b8afe354f6352c643b65937130/torchvision-0.28.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:e9f54c30cd52e3ef7fd034cc69b7bb7e0964e1c8f8743e018ab92e95b40f9eee", size = 1856020, upload-time = "2026-07-08T16:07:52.182Z" }, + { url = "https://files.pythonhosted.org/packages/f0/4c/95233776e2def960e5abb7a07931230a545f43717a56a1e1140162033598/torchvision-0.28.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5cf78ebc401ce64ae19b8c55de866bb836797d559a4de9c25ccbe74cfa642d3a", size = 7842127, upload-time = "2026-07-08T16:07:53.446Z" }, + { url = "https://files.pythonhosted.org/packages/93/e4/e9b2495d0d57b9f60d63c57d0a910410a81b4b073bf70917bef815291119/torchvision-0.28.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:028a3d481b37d785605620d7cdad897064c5a55bae2aa1f2658766333e291940", size = 7675040, upload-time = "2026-07-08T16:07:58.017Z" }, + { url = "https://files.pythonhosted.org/packages/7c/9c/55ed9cb6dfe3ee9c837df5cd0e758372e5829aa38b8dd71343aa632cc4e2/torchvision-0.28.0-cp312-cp312-win_amd64.whl", hash = "sha256:87dc16b2df427c1318ad335f1e2be2b3b15b2cf20f7934c83b0505a48425ee5d", size = 4085785, upload-time = "2026-07-08T16:07:50.928Z" }, +] + +[[package]] +name = "tqdm" +version = "4.69.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8c/69/40407dfc835517f058b603dbf37a6df094d8582b015a51eddc988febbcb7/tqdm-4.69.0.tar.gz", hash = "sha256:700c5e85dcd5f009dd6222588a29180a193a748247a5d855b4d67db93d79a53b", size = 792569, upload-time = "2026-07-17T18:09:06.2Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/21/99a0cdaf54eb35e77623c41b5a2c9472ee4404bba687052791fe2aba6773/tqdm-4.69.0-py3-none-any.whl", hash = "sha256:9979978912be667a6ef21fd5d8abf54e324e63d82f7f43c360792ebc2bc4e622", size = 676680, upload-time = "2026-07-17T18:09:04.172Z" }, +] + +[[package]] +name = "transformer-engine" +version = "2.17.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f4/26/e501a474a3f36c2d439561fa4f2094a3c0344df439500dda55937a8fccf7/transformer_engine-2.17.0-py3-none-any.whl", hash = "sha256:c274fd74ea2e4caa7132921e1ecfa3847931abbed9f6b8cab61346cdb833bc76", size = 1001710, upload-time = "2026-07-09T00:36:15.453Z" }, +] + +[package.optional-dependencies] +core-cu13 = [ + { name = "transformer-engine-cu13" }, +] +pytorch = [ + { name = "transformer-engine-torch" }, +] + +[[package]] +name = "transformer-engine-cu13" +version = "2.17.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "importlib-metadata" }, + { name = "packaging" }, + { name = "pydantic" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/21/5e0ec562539798ffd375c218b8bd7612b387455ccc79ae14576f411c9609/transformer_engine_cu13-2.17.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:1c8e6627cf02358201f513ab902c7b4f82f179214bd775edc97ce0f2001d82da", size = 244570183, upload-time = "2026-07-09T00:36:35.507Z" }, +] + +[[package]] +name = "transformer-engine-torch" +version = "2.17.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "einops" }, + { name = "nvdlfw-inspect" }, + { name = "onnx" }, + { name = "onnxscript" }, + { name = "packaging" }, + { name = "pydantic" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" } }, + { name = "transformer-engine-cu13" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9d/f1/ad613157261ac71262f5f260acf07fe22ae5ee0bb0675622c4f19a3e8afb/transformer_engine_torch-2.17.0.tar.gz", hash = "sha256:52ca109cdbc987ca02f87735fa375da8a02e117acb87e25dcb7e27dca37aa87f", size = 369712, upload-time = "2026-07-09T00:36:28.065Z" } + +[[package]] +name = "transformers" +version = "5.14.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "regex" }, + { name = "safetensors" }, + { name = "tokenizers" }, + { name = "tqdm" }, + { name = "typer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5a/fb/2a2ba88f325e68a921d8b69ff63b477830b2e73ade9a3c8c8cab2f06d741/transformers-5.14.1.tar.gz", hash = "sha256:60d196c27781eacf8637e2b533f517582907ad6f9ae142046d6b69431a5b2173", size = 9295927, upload-time = "2026-07-16T09:41:57.773Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6f/67/8d85ca2323233ae3c0365a659c4e52ee1f587b440e4bc577e7d8e4416d0f/transformers-5.14.1-py3-none-any.whl", hash = "sha256:9db974c4079ede2d1a3ea7ca5a240df33f2cc26fc2b36ba64c5f2a4f43b6e725", size = 11625234, upload-time = "2026-07-16T09:41:54.143Z" }, +] + +[[package]] +name = "triton" +version = "3.7.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/fa/f856e24deb462d5f18bd4b5a746957862ab9b6ee5834bda60605ec348366/triton-3.7.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9497f2e696ee368862a181a90b2dcc03ca978cc4f602abd67c7d81022a6988e1", size = 184692359, upload-time = "2026-06-17T20:03:48.288Z" }, + { url = "https://files.pythonhosted.org/packages/c4/6f/fb96d15db6f36d6eae4cafb998c2e0353bf59d7c4ea1662d7497f269134a/triton-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7e40869937a68206ec70d7f25bb7ec6433cb083f9135e1f36dbd318dc449a728", size = 197719725, upload-time = "2026-06-17T19:53:20.419Z" }, +] + +[[package]] +name = "triton-windows" +version = "3.7.1.post27" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/76/30/325b420efd0047e119679c646a9a410db216069800ec009fae3da26c69a3/triton_windows-3.7.1.post27-cp312-cp312-win_amd64.whl", hash = "sha256:f5406230d7dbf6965bc4051fcad27b81c39ba4a4bfde06f494dd7ff4eb325a9e", size = 49683004, upload-time = "2026-06-21T16:48:14.02Z" }, +] + +[[package]] +name = "typeguard" +version = "4.5.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/67/1c/dfba5c4633cafc4c701f237d2ba63b416805047fd6d96aab4cfc40969f98/typeguard-4.5.2.tar.gz", hash = "sha256:5a16dcac23502039299c97c8941651bc33d7ea8cc4b2f7d6bbb1b528f6eea423", size = 80240, upload-time = "2026-05-14T12:59:40.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/29/74eeb4d3f3ae61ca096b018ad486b3b3c74b17bec09ab4edab721cbefec3/typeguard-4.5.2-py3-none-any.whl", hash = "sha256:fcf9de18bd945cdb4c7b996e12b4c51ce83f92f191314a6d7cf1739586ec98cf", size = 36748, upload-time = "2026-05-14T12:59:39.473Z" }, +] + +[[package]] +name = "typer" +version = "0.27.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-doc" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "rich" }, + { name = "shellingham" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/37/78/fda3361b56efc27944f24225f6ecd13d96d6fcfe37bd0eb34e2f4c63f9fc/typer-0.27.0.tar.gz", hash = "sha256:629bd12ea5d13a17148125d9a264f949eb171fb3f120f9b04d85873cab054fa5", size = 203430, upload-time = "2026-07-15T19:21:07.007Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/40/03/26a383c9e58c213199d1aad1c3d353cfc22d4444ec6d2c0bf8ad02523843/typer-0.27.0-py3-none-any.whl", hash = "sha256:6f4b27631e47f077871b7dc30e933ec0131c1390fbe0e387ea5574b5bac9ccf1", size = 122716, upload-time = "2026-07-15T19:21:05.553Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/cc/6253133b5bb138fc3306cebfbda2c520f545d36b5be2c7255cc528bb45d6/typing_extensions-4.16.0.tar.gz", hash = "sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5", size = 113555, upload-time = "2026-07-02T08:40:05.92Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" }, +] + +[[package]] +name = "typing-inspect" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mypy-extensions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/dc/74/1789779d91f1961fa9438e9a8710cdae6bd138c80d7303996933d117264a/typing_inspect-0.9.0.tar.gz", hash = "sha256:b23fc42ff6f6ef6954e4852c1fb512cdd18dbea03134f91f856a95ccc9461f78", size = 13825, upload-time = "2023-05-24T20:25:47.612Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/65/f3/107a22063bf27bdccf2024833d3445f4eea42b2e598abfbd46f6a63b6cb0/typing_inspect-0.9.0-py3-none-any.whl", hash = "sha256:9ee6fc59062311ef8547596ab6b955e1b8aa46242d854bfc78f4f6b0eff35f9f", size = 8827, upload-time = "2023-05-24T20:25:45.287Z" }, +] + +[[package]] +name = "typing-inspection" +version = "0.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, +] + +[[package]] +name = "tyro" +version = "1.0.15" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "docstring-parser" }, + { name = "typeguard" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/12/78/a5749a6c1ee9abc2999e294f339f8f72476d1a60bb95fc0e86156aafed3b/tyro-1.0.15.tar.gz", hash = "sha256:3f1d60887723eecb9c489f195d11f079c4a1f33df74b723552ad31ec57c667bb", size = 593822, upload-time = "2026-06-20T08:48:28.364Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5d/28/d607636187cf6c18eb72efb5d65c1d1b9e451db03d84676f835dd488fdbd/tyro-1.0.15-py3-none-any.whl", hash = "sha256:982da1d566005f1b2a6b56f6be6c6929c96f0e5fab9d61bc6097573ddfaf8d13", size = 215045, upload-time = "2026-06-20T08:48:27.104Z" }, +] + +[[package]] +name = "urllib3" +version = "2.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, +] + +[[package]] +name = "wcwidth" +version = "0.8.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/34/74/c6428f875774288bec1396f5bfcbc2d925700a4dad61727fd5f2b12f249d/wcwidth-0.8.2.tar.gz", hash = "sha256:91fbef97204b96a3d4d421609b80340b760cf33e26da123ff243d76b1fda8dda", size = 1466253, upload-time = "2026-06-29T18:11:11.601Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/96/42/3e5985a0a7e57de470b320c6d6a1a67c844f6737a587f3d44dd13d1819e7/wcwidth-0.8.2-py3-none-any.whl", hash = "sha256:d63947694a0539a1d51e01eda7caf800c291020e6cdd7e28ad7b14dd33ad4f85", size = 323166, upload-time = "2026-06-29T18:11:09.888Z" }, +] + +[[package]] +name = "win32-setctime" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b3/8f/705086c9d734d3b663af0e9bb3d4de6578d08f46b1b101c2442fd9aecaa2/win32_setctime-1.2.0.tar.gz", hash = "sha256:ae1fdf948f5640aae05c511ade119313fb6a30d7eabe25fef9764dca5873c4c0", size = 4867, upload-time = "2024-12-07T15:28:28.314Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/07/c6fe3ad3e685340704d314d765b7912993bcb8dc198f0e7a89382d37974b/win32_setctime-1.2.0-py3-none-any.whl", hash = "sha256:95d644c4e708aba81dc3704a116d8cbc974d70b3bdb8be1d150e36be6e9d1390", size = 4083, upload-time = "2024-12-07T15:28:26.465Z" }, +] + +[[package]] +name = "zipp" +version = "4.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b9/d8/eab98a517c14134c0b2eb4e2387bc5f457334293ec5d2dd3857ec2966802/zipp-4.1.0.tar.gz", hash = "sha256:4cb57381f544315db7688e976e922a2b18cdb513d21cc194eb42232ba2a3e602", size = 26214, upload-time = "2026-05-18T20:08:57.967Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/13/547360d81e6d88d58492968ffda9f9542854f11310ee556fef14260cc886/zipp-4.1.0-py3-none-any.whl", hash = "sha256:25ad4e16390cd314347dd8f1de67a2ac538ae658ed4ab9db16029c07c188e97f", size = 10238, upload-time = "2026-05-18T20:08:57.045Z" }, +] diff --git a/integrations/sana/tests/test_camera.py b/integrations/sana/tests/test_camera.py new file mode 100644 index 000000000..e54d14eb9 --- /dev/null +++ b/integrations/sana/tests/test_camera.py @@ -0,0 +1,195 @@ +# 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 tests for SANA-WM camera and intrinsics helpers.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import torch + +from sana_wm.camera import ( + action_string_to_c2w, + default_intrinsics_vec4, + fit_camera_trajectory, + fit_intrinsics_sequence, + load_intrinsics, + prepare_camera, + resize_center_crop_geometry, + snap_num_frames, + transform_intrinsics_for_crop, +) + +pytestmark = pytest.mark.ci_cpu + + +def test_default_intrinsics_vec4_centers_and_matches_hfov() -> None: + vec = default_intrinsics_vec4((1000, 500), num_frames=3, hfov_deg=90.0) + assert vec.shape == (3, 4) + fx, fy, cx, cy = vec[0] + # 90 deg hfov with square pixels -> fx = 0.5 * W / tan(45) = 0.5 * W. + assert fx == pytest.approx(500.0) + assert fy == pytest.approx(500.0) + # Principal point at the image center. + assert cx == pytest.approx(500.0) + assert cy == pytest.approx(250.0) + # Same intrinsics for every frame. + assert np.allclose(vec, vec[0]) + + +def test_default_intrinsics_vec4_rejects_bad_fov() -> None: + for bad in (0.0, 180.0, 200.0): + with pytest.raises(ValueError): + default_intrinsics_vec4((640, 360), num_frames=1, hfov_deg=bad) + + +def test_action_string_rolls_out_identity_plus_motion() -> None: + """Expand ``w-3`` to an identity frame plus three motion frames.""" + c2w = action_string_to_c2w("w-3", smooth=False) + + assert c2w.shape == (4, 4, 4) + np.testing.assert_allclose(c2w[0], np.eye(4), atol=1e-6) + assert np.all(np.diff(c2w[:, 2, 3]) > 0) + + +def test_action_string_repeats_to_requested_frame_count() -> None: + """Use action prompts to derive the requested number of camera poses.""" + c2w = action_string_to_c2w("w-1", smooth=False, num_frames=5) + + assert c2w.shape == (5, 4, 4) + np.testing.assert_allclose(c2w[0], np.eye(4), atol=1e-6) + np.testing.assert_allclose( + c2w[:, 2, 3], + [0.0, 0.025, 0.05, 0.075, 0.1], + atol=1e-6, + ) + + +def test_action_string_rejects_unknown_keys() -> None: + """Reject invalid action DSL tokens.""" + with pytest.raises(ValueError, match="unknown keys"): + action_string_to_c2w("q-3") + + +def test_snap_num_frames_matches_ltx2_stride() -> None: + """Snap requested frames to nearest ``8k + 1`` value.""" + assert snap_num_frames(321) == 321 + assert snap_num_frames(322) == 321 + assert snap_num_frames(325) == 329 + assert snap_num_frames(325, upper_bound=326) == 321 + + +def test_fit_intrinsics_sequence_interpolates_short_sequences() -> None: + """Interpolate per-frame intrinsics when the source sequence is shorter.""" + source = np.array([[10.0, 20.0, 1.0, 2.0], [30.0, 40.0, 3.0, 4.0]]) + + fitted = fit_intrinsics_sequence(source, 3) + + np.testing.assert_allclose( + fitted, + np.array( + [ + [10.0, 20.0, 1.0, 2.0], + [20.0, 30.0, 2.0, 3.0], + [30.0, 40.0, 3.0, 4.0], + ], + dtype=np.float32, + ), + ) + + +def test_fit_camera_trajectory_interpolates_short_sequences() -> None: + """Fit explicit camera paths to the requested rollout length.""" + source = np.broadcast_to(np.eye(4, dtype=np.float32), (2, 4, 4)).copy() + source[1, 2, 3] = 2.0 + + fitted = fit_camera_trajectory(source, 5) + + assert fitted.shape == (5, 4, 4) + np.testing.assert_allclose(fitted[:, 2, 3], [0.0, 0.5, 1.0, 1.5, 2.0]) + np.testing.assert_allclose( + fitted[:, 3], + np.broadcast_to(np.array([0.0, 0.0, 0.0, 1.0]), (5, 4)), + ) + for rotation in fitted[:, :3, :3]: + np.testing.assert_allclose(rotation.T @ rotation, np.eye(3), atol=1e-5) + assert np.linalg.det(rotation) == pytest.approx(1.0, abs=1e-5) + + +@pytest.mark.parametrize( + ("array", "expected"), + [ + ( + np.array([100.0, 110.0, 50.0, 55.0], dtype=np.float32), + np.array([[100.0, 110.0, 50.0, 55.0]] * 3, dtype=np.float32), + ), + ( + np.array( + [[100.0, 0.0, 50.0], [0.0, 110.0, 55.0], [0.0, 0.0, 1.0]], + dtype=np.float32, + ), + np.array([[100.0, 110.0, 50.0, 55.0]] * 3, dtype=np.float32), + ), + ], +) +def test_load_intrinsics_accepts_static_shapes( + tmp_path: Path, array: np.ndarray, expected: np.ndarray +) -> None: + """Load static vector and matrix intrinsics as per-frame vectors.""" + path = tmp_path / "intrinsics.npy" + np.save(path, array) + + loaded = load_intrinsics(path, num_frames=3) + + np.testing.assert_allclose(loaded, expected) + + +def test_resize_crop_geometry_and_intrinsics_transform() -> None: + """Map source intrinsics through SANA-WM resize and center-crop geometry.""" + src_size = (640, 480) + resized_size, crop_offset = resize_center_crop_geometry(src_size) + intrinsics = np.array([[400.0, 420.0, 320.0, 240.0]], dtype=np.float32) + + transformed = transform_intrinsics_for_crop( + intrinsics, src_size, resized_size, crop_offset + ) + + assert resized_size == (1280, 960) + assert crop_offset == (0, 128) + np.testing.assert_allclose( + transformed, + np.array([[800.0, 840.0, 640.0, 352.0]], dtype=np.float32), + ) + + +def test_prepare_camera_shapes_for_sana_wm_resolution() -> None: + """Build raymap and chunk-Plucker tensors at 704x1280 SANA-WM shape.""" + num_frames = 17 + poses = np.broadcast_to(np.eye(4, dtype=np.float32), (num_frames, 4, 4)).copy() + poses[:, 2, 3] = np.linspace(0.0, 1.0, num_frames) + intrinsics = np.broadcast_to( + np.array([900.0, 900.0, 640.0, 352.0], dtype=np.float32), + (num_frames, 4), + ).copy() + + camera = prepare_camera(poses, intrinsics) + + assert camera["raymap"].shape == (3, 20) + assert camera["chunk_plucker"].shape == (48, 3, 22, 40) + assert camera["raymap"].dtype == torch.float32 + assert camera["chunk_plucker"].dtype == torch.float32 diff --git a/integrations/sana/tests/test_parity_benchmark_summary.py b/integrations/sana/tests/test_parity_benchmark_summary.py new file mode 100644 index 000000000..2a10e18b6 --- /dev/null +++ b/integrations/sana/tests/test_parity_benchmark_summary.py @@ -0,0 +1,264 @@ +# 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 tests for the SANA-WM parity benchmark summary.""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from types import ModuleType + +import pytest + +pytestmark = pytest.mark.ci_cpu + + +def _load_bench_summary() -> ModuleType: + path = Path("integrations/sana/tests/parity_check/bench_summary.py") + spec = importlib.util.spec_from_file_location("sana_wm_bench_summary", path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _load_bench_sweep_summary() -> ModuleType: + path = Path("integrations/sana/tests/parity_check/bench_sweep_summary.py") + spec = importlib.util.spec_from_file_location("sana_wm_bench_sweep_summary", path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_benchmark_summary_uses_generation_ms_per_clip_for_chart() -> None: + module = _load_bench_summary() + summary = { + "inputs": { + "image_path": "image.png", + "prompt_path": "prompt.txt", + "camera_path": "pose.npy", + "intrinsics_path": "intrinsics.npy", + "num_frames": 20, + "seed": 42, + "no_refiner": True, + "stage1_precision": "fp8", + "refiner_precision": "fp8", + "quant_backend": "torch-fp8", + "compile_stage1": False, + "force_cudnn_sdpa": True, + "warmup_runs": 1, + }, + "upstream": { + "runs_measured": 3, + "wall_median_s": 2.0, + "wall_p90_s": 2.2, + "stage1_total_median_ms": 1600.0, + "dit_median_ms": 1500.0, + "vae_decode_median_ms": 300.0, + "mem_peak_median_gib": 10.0, + }, + "flashdreams": { + "runs_measured": 3, + "wall_median_s": 1.5, + "wall_p90_s": 1.8, + "encode_median_ms": 100.0, + "dit_median_ms": 1100.0, + "vae_decode_median_ms": 250.0, + "mem_peak_median_gib": 12.0, + }, + } + + report = module._render_markdown(summary) + chart = module._render_chart_markdown(summary, "Test GPU") + + assert "## Benchmark metric" in report + assert "- stage1_precision: `fp8`" in report + assert "generation median / clip | 2000.00 ms | 1500.00 ms" in report + assert "generation median / frame, diagnostic | 100.00 ms | 75.00 ms" in report + assert "## Timing breakdown" in report + assert "| Stage-1 DiT median | 1500.00 ms | 1100.00 ms |" in report + assert "Stage-1 DiT metric" not in report + assert "Generation after model load" not in report + assert chart == ( + "# SANA-WM Benchmark Data (ms)\n" + "\n" + "| device | official | flashdreams |\n" + "| --- | ---: | ---: |\n" + "| Test GPU | 2000.00 | 1500.00 |\n" + ) + + +def test_benchmark_summary_writes_chart_data(tmp_path: Path) -> None: + module = _load_bench_summary() + upstream = tmp_path / "upstream" / "run_0" + flashdreams = tmp_path / "flashdreams" / "run_0" + upstream.mkdir(parents=True) + flashdreams.mkdir(parents=True) + (upstream / "stats.json").write_text( + json.dumps( + { + "wall_s": 4.0, + "mem_peak_gib": 10.0, + "timings_s": { + "stage1_sample_s": 3.5, + "stage1_dit_s": 3.0, + "vae_decode_s": 0.5, + }, + } + ), + encoding="utf-8", + ) + (flashdreams / "stats.json").write_text( + json.dumps( + { + "wall_s": 2.0, + "stats_ms": { + "encode_ms": 100.0, + "diffuse_ms": 1500.0, + "decode_ms": 400.0, + "mem_peak_gib": 12.0, + }, + } + ), + encoding="utf-8", + ) + out_json = tmp_path / "bench.json" + out_md = tmp_path / "bench.md" + out_chart = tmp_path / "perf.md" + + module.main( + [ + "--upstream-dir", + str(tmp_path / "upstream"), + "--flashdreams-dir", + str(tmp_path / "flashdreams"), + "--warmup-runs", + "0", + "--image-path", + "image.png", + "--prompt-path", + "prompt.txt", + "--camera-path", + "pose.npy", + "--intrinsics-path", + "intrinsics.npy", + "--num-frames", + "40", + "--seed", + "42", + "--device-label", + "Test GPU", + "--stage1-precision", + "bf16", + "--refiner-precision", + "bf16", + "--output-json", + str(out_json), + "--output-md", + str(out_md), + "--output-chart-md", + str(out_chart), + ] + ) + + assert "| Test GPU | 4000.00 | 2000.00 |" in out_chart.read_text(encoding="utf-8") + report = out_md.read_text(encoding="utf-8") + assert "generation median / clip | 4000.00 ms | 2000.00 ms" in report + assert "generation median / frame, diagnostic | 100.00 ms | 50.00 ms" in report + assert "generation median / frame |" not in report + assert "per generated clip" in report + assert "not as independently timed frames" in report + assert "per generated frame" not in report + summary = json.loads(out_json.read_text(encoding="utf-8")) + assert summary["benchmark"] == { + "metric": "generation_ms_per_clip", + "unit": "ms", + "timing_boundary": ( + "pipeline.generate after model setup; excludes model construction, " + "checkpoint loading, video writing, and frame dumps" + ), + "device_label": "Test GPU", + "chart_label": "Test GPU", + "official": 4000.0, + "flashdreams": 2000.0, + } + + +def test_benchmark_summary_keeps_frame_normalized_diagnostics() -> None: + module = _load_bench_summary() + assert module._generation_ms_per_frame( + { + "inputs": {"num_frames": 40}, + "upstream": {"wall_median_s": 4.0}, + }, + "upstream", + ) == 100.0 + assert module._generation_ms_per_clip( + { + "inputs": {"num_frames": 40}, + "upstream": {"wall_median_s": 4.0}, + }, + "upstream", + ) == 4000.0 + + +def test_benchmark_sweep_summary_writes_precision_chart_data(tmp_path: Path) -> None: + module = _load_bench_sweep_summary() + bf16_json = tmp_path / "bf16.json" + fp8_json = tmp_path / "fp8.json" + bf16_json.write_text( + json.dumps({"benchmark": {"official": 1000.0, "flashdreams": 900.0}}), + encoding="utf-8", + ) + fp8_json.write_text( + json.dumps({"benchmark": {"official": 800.0, "flashdreams": 700.0}}), + encoding="utf-8", + ) + out_json = tmp_path / "bench.json" + out_md = tmp_path / "bench.md" + out_chart = tmp_path / "perf.md" + + module.main( + [ + "--item", + f"BF16:{bf16_json}", + "--item", + f"FP8:{fp8_json}", + "--output-json", + str(out_json), + "--output-md", + str(out_md), + "--output-chart-md", + str(out_chart), + ] + ) + + assert out_chart.read_text(encoding="utf-8") == ( + "# SANA-WM Precision Sweep Summary (ms)\n" + "\n" + "| precision | official | flashdreams |\n" + "| --- | ---: | ---: |\n" + "| BF16 | 1000.00 | 900.00 |\n" + "| FP8 | 800.00 | 700.00 |\n" + ) + payload = json.loads(out_json.read_text(encoding="utf-8")) + assert payload["benchmark"]["metric"] == "generation_ms_per_clip" + assert payload["benchmark"]["unit"] == "ms" + assert [row["label"] for row in payload["rows"]] == ["BF16", "FP8"] diff --git a/integrations/sana/tests/test_quant_cuda.py b/integrations/sana/tests/test_quant_cuda.py new file mode 100644 index 000000000..ff6265760 --- /dev/null +++ b/integrations/sana/tests/test_quant_cuda.py @@ -0,0 +1,111 @@ +# 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. + +"""CUDA smoke tests for SANA-WM low-precision linear replacements.""" + +from __future__ import annotations + +import pytest +import torch + +from sana_wm.quant import TorchScaledMMFP4Linear, TorchScaledMMFP8Linear + +pytestmark = pytest.mark.ci_gpu + + +def _require_quant_cuda() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for quantization smokes") + missing = [ + name + for name in ("_scaled_mm", "float8_e4m3fn", "float4_e2m1fn_x2") + if not hasattr(torch, name) + ] + if missing: + pytest.skip(f"PyTorch lacks quantization primitive(s): {missing}") + + +def test_fp8_linear_runs_scaled_mm_on_cuda() -> None: + """Exercise the TE-free E4M3 FP8 Linear replacement on CUDA.""" + _require_quant_cuda() + source = torch.nn.Linear(32, 64, bias=True, device="cuda", dtype=torch.bfloat16) + quantized = TorchScaledMMFP8Linear.from_linear( + source, + out_dtype=torch.bfloat16, + ) + inputs = torch.randn(4, 5, 32, device="cuda", dtype=torch.bfloat16) + + output = quantized(inputs) + + assert quantized.weight.shape == source.weight.shape + assert quantized.bias is not None + assert output.shape == (4, 5, 64) + assert output.dtype == torch.bfloat16 + assert torch.isfinite(output.float()).all() + + +def test_fp4_linear_runs_scaled_mm_on_blackwell() -> None: + """Exercise the TE-free E2M1 NVFP4 Linear replacement on Blackwell CUDA.""" + _require_quant_cuda() + major, minor = torch.cuda.get_device_capability() + if major < 10: + pytest.skip(f"NVFP4 requires Blackwell-class CUDA, got sm_{major}{minor}") + + source = torch.nn.Linear(32, 64, bias=True, device="cuda", dtype=torch.bfloat16) + quantized = TorchScaledMMFP4Linear.from_linear( + source, + out_dtype=torch.bfloat16, + ) + inputs = torch.randn(4, 5, 32, device="cuda", dtype=torch.bfloat16) + + output = quantized(inputs) + + assert quantized.weight.shape == source.weight.shape + assert quantized.bias is not None + assert output.shape == (4, 5, 64) + assert output.dtype == torch.bfloat16 + assert torch.isfinite(output.float()).all() + + +def test_quantized_linears_roughly_track_source_layer() -> None: + """Catch gross scale/layout regressions beyond expected FP4/FP8 noise.""" + _require_quant_cuda() + major, minor = torch.cuda.get_device_capability() + if major < 10: + pytest.skip(f"NVFP4 requires Blackwell-class CUDA, got sm_{major}{minor}") + + torch.manual_seed(123) + source = torch.nn.Linear(32, 64, bias=True, device="cuda", dtype=torch.bfloat16) + inputs = torch.randn(4, 5, 32, device="cuda", dtype=torch.bfloat16) + reference = source(inputs).float() + + fp8 = TorchScaledMMFP8Linear.from_linear( + source, + out_dtype=torch.bfloat16, + ) + fp4 = TorchScaledMMFP4Linear.from_linear( + source, + out_dtype=torch.bfloat16, + ) + + fp8_error = (fp8(inputs).float() - reference).abs() + fp4_error = (fp4(inputs).float() - reference).abs() + fp8_rel_mae = fp8_error.mean() / reference.abs().mean().clamp_min(1e-6) + fp4_rel_mae = fp4_error.mean() / reference.abs().mean().clamp_min(1e-6) + + assert fp8_error.max().item() <= 0.10 + assert fp8_rel_mae.item() <= 0.06 + assert fp4_error.max().item() <= 0.35 + assert fp4_rel_mae.item() <= 0.20 diff --git a/integrations/sana/tests/test_smoke.py b/integrations/sana/tests/test_smoke.py new file mode 100644 index 000000000..c04a30158 --- /dev/null +++ b/integrations/sana/tests/test_smoke.py @@ -0,0 +1,1758 @@ +# 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 smoke tests for the SANA-WM configs.""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +import sana_wm.conditioning as conditioning_module +import sana_wm.decoder as decoder_module +import sana_wm.refiner as refiner_module + +try: + import tomllib +except ModuleNotFoundError: # pragma: no cover - Python < 3.11 fallback + import tomli as tomllib + +from sana_wm.config import ( + PIPELINE_SANA_WM_BIDIRECTIONAL, + RUNNER_CONFIGS, + RUNNER_SANA_WM_BIDIRECTIONAL, +) +from sana_wm.constants import ( + SANA_WM_CONFIG_PATH, + SANA_WM_HF_REPO, + SANA_WM_MODEL_PATH, +) +from sana_wm.runner import ( + SanaWMRunner, + SanaWMRunnerConfig, + _pipeline_config, + _resolve_quant_backend, + _validate_precision_request, +) +from sana_wm.conditioning import ( + SanaWMCameraConditioningEncoderConfig, + SanaWMCameraRequest, + SanaWMConditioningEncoderConfig, + SanaWMFirstFrameEncoderConfig, + SanaWMTextPromptEncoderConfig, + SanaWMTextPromptRequest, +) +from sana_wm.decoder import ( + SanaWMDecodedVideo, + SanaWMLTX2LatentRefinerConfig, + SanaWMLTX2VAEDecoderConfig, + SanaWMVideoDecoderConfig, +) +from sana_wm.diffusion import SanaWMDiffusionModelConfig +from sana_wm.scheduler import ( + SanaWMLTXEulerScheduler, + SanaWMLTXEulerSchedulerConfig, +) +from sana_wm.transformer import ( + SanaWMTransformerCache, + SanaWMTransformerConfig, + SanaWMStage1Conditioning, + _avoid_degenerate_tile_tail, + _load_inference_config, + _stage1_quant_include_patterns, +) +from sana_wm.refiner import SanaWMLTX2Refiner, _pack_latents, _unpack_latents +from sana_wm.quant import ( + apply_rht16, + nvfp4_global_scale, + TorchScaledMMFP4Linear, + TorchScaledMMFP8Linear, + replace_linear_with_torch_fp4, + replace_linear_with_torch_fp8, +) +from sana_wm.stage1_model import ( + GLUMBConvTemp, + SANA_WM_STAGE1_SPEC, + SanaWMStage1Model, + SanaWMStage1Spec, + Stage1SelfAttention, + linearize_stage1_ffn_for_quant, +) + +from flashdreams.infra.config import derive_config +from flashdreams.infra.decoder import StreamingVideoDecoder +from flashdreams.infra.diffusion.model import DiffusionModel + +pytestmark = pytest.mark.ci_cpu + +ENTRY_POINT_GROUP = "flashdreams.runner_configs" + + +def test_runner_config_is_registered() -> None: + """Expose the SANA-WM runner slug.""" + assert RUNNER_CONFIGS == { + "sana-wm-bidirectional": RUNNER_SANA_WM_BIDIRECTIONAL, + } + + +def test_runner_name_mirrors_pipeline_name() -> None: + """Keep ``flashdreams-run `` aligned with the wrapped config slug.""" + assert RUNNER_SANA_WM_BIDIRECTIONAL.runner_name == ( + PIPELINE_SANA_WM_BIDIRECTIONAL.name + ) + + +def test_runner_has_description() -> None: + """Provide non-empty CLI help text for the runner registry.""" + assert RUNNER_SANA_WM_BIDIRECTIONAL.description.strip() + + +def test_pipeline_uses_sana_diffusion_model() -> None: + """Keep the public runner wired to explicit SANA-WM boundaries.""" + pipeline = PIPELINE_SANA_WM_BIDIRECTIONAL + transformer = PIPELINE_SANA_WM_BIDIRECTIONAL.diffusion_model.transformer + + assert isinstance(pipeline.encoder, SanaWMConditioningEncoderConfig) + assert isinstance(pipeline.decoder, SanaWMVideoDecoderConfig) + assert isinstance(pipeline.diffusion_model, SanaWMDiffusionModelConfig) + assert pipeline.diffusion_model._target is DiffusionModel + assert isinstance(transformer, SanaWMTransformerConfig) + assert isinstance(pipeline.diffusion_model.scheduler, SanaWMLTXEulerSchedulerConfig) + assert isinstance(pipeline.decoder.refiner, SanaWMLTX2LatentRefinerConfig) + + +def test_sana_decoder_uses_video_decoder_contract() -> None: + """Expose SANA-WM VAE temporal sizing through the FlashDreams decoder API.""" + decoder = SanaWMVideoDecoderConfig().setup() + + assert isinstance(decoder, StreamingVideoDecoder) + assert decoder.spatial_compression_ratio == 32 + assert decoder.temporal_compression_ratio == 8 + assert decoder.get_output_temporal_size(0, 21) == 161 + assert decoder.get_input_temporal_size(0, 161) == 21 + with pytest.raises(ValueError, match="8k\\+1"): + decoder.get_input_temporal_size(0, 160) + with pytest.raises(ValueError, match="one AR step"): + decoder.get_output_temporal_size(1, 21) + + +def test_sana_diffusion_config_instantiates_base_model() -> None: + """Use FlashDreams' shared diffusion model rather than a Sana-only runner.""" + model = PIPELINE_SANA_WM_BIDIRECTIONAL.diffusion_model.setup() + + assert type(model) is DiffusionModel + assert isinstance(model.transformer.config, SanaWMTransformerConfig) + assert isinstance(model.scheduler, SanaWMLTXEulerScheduler) + + +def test_runner_pipeline_config_routes_runtime_fields_to_components() -> None: + """Apply CLI overrides to the component that owns each runtime field.""" + cfg = derive_config( + RUNNER_SANA_WM_BIDIRECTIONAL, + config_path="local_config.yaml", + model_path="local_model.safetensors", + stage1_precision="fp8", + step=7, + flow_shift=6.5, + no_refiner=True, + offload_vae=True, + offload_text_encoder=True, + offload_stage1=True, + ) + + pipeline = _pipeline_config(cfg, quant_backend="torch-fp8") + + assert isinstance(pipeline.encoder, SanaWMConditioningEncoderConfig) + assert isinstance(pipeline.decoder, SanaWMVideoDecoderConfig) + assert pipeline.encoder.config_path == "local_config.yaml" + assert pipeline.encoder.text_encoder.config_path == "local_config.yaml" + assert pipeline.encoder.text_encoder.stage1_precision == "fp8" + assert pipeline.encoder.text_encoder.quant_backend == "torch-fp8" + assert pipeline.encoder.text_encoder.offload_text_encoder is True + assert pipeline.encoder.first_frame_encoder.offload_vae is True + assert pipeline.decoder.vae_decoder.config_path == "local_config.yaml" + assert pipeline.decoder.vae_decoder.offload_vae is True + assert pipeline.decoder.refiner is None + assert pipeline.diffusion_model.transformer.checkpoint_path == ( + "local_model.safetensors" + ) + assert pipeline.diffusion_model.transformer.offload_stage1 is True + assert pipeline.diffusion_model.scheduler.num_inference_steps == 7 + assert pipeline.diffusion_model.scheduler.shift == 6.5 + + +def test_sana_ltx_scheduler_step_pins_zero_timestep_tokens() -> None: + """Keep first-frame tokens fixed in the per-token LTX Euler step.""" + scheduler = SanaWMLTXEulerSchedulerConfig(num_inference_steps=4).setup() + + timesteps = scheduler.timesteps( + num_inference_steps=4, + shift=5.0, + device=torch.device("cpu"), + ) + sample = torch.ones((1, 2, 1)) + model_output = torch.ones_like(sample) + stepped = scheduler.step_ltx( + model_output=model_output, + timestep=torch.tensor(1000.0), + next_timestep=torch.tensor(500.0), + sample=sample, + per_token_timesteps=torch.tensor([[1000.0, 0.0]]), + schedule_timesteps=torch.tensor([1000.0, 500.0, 0.0]), + ) + + assert isinstance(scheduler, SanaWMLTXEulerScheduler) + assert timesteps.shape == (5,) + assert float(timesteps[-1]) == 0.0 + torch.testing.assert_close( + stepped, + torch.tensor([[[1.5], [1.0]]]), + ) + + +def test_sana_ltx_scheduler_matches_diffusers_per_token_step() -> None: + """Match diffusers FlowMatch Euler for SANA-WM's per-token branch.""" + from diffusers import FlowMatchEulerDiscreteScheduler + + ours = SanaWMLTXEulerSchedulerConfig(num_inference_steps=4).setup() + schedule_timesteps = ours.timesteps( + num_inference_steps=4, + shift=5.0, + device=torch.device("cpu"), + ) + upstream = FlowMatchEulerDiscreteScheduler(shift=5.0) + upstream.set_timesteps(4, device=torch.device("cpu")) + torch.testing.assert_close(schedule_timesteps[:-1], upstream.timesteps) + torch.testing.assert_close(schedule_timesteps / 1000.0, upstream.sigmas) + + sample = torch.tensor([[[1.0, -0.5], [0.25, 0.5], [-1.0, 0.75]]]) + model_output = torch.tensor([[[0.2, 0.4], [-0.3, 0.1], [0.5, -0.2]]]) + per_token_timesteps = torch.stack( + [ + torch.tensor( + [ + float(upstream.timesteps[0]), + 0.0, + float(upstream.timesteps[0]), + ] + ) + ] + ) + + expected = upstream.step( + model_output, + upstream.timesteps[0], + sample, + per_token_timesteps=per_token_timesteps, + return_dict=False, + )[0] + actual = ours.step_ltx( + model_output=model_output, + timestep=schedule_timesteps[0], + next_timestep=schedule_timesteps[1], + sample=sample, + per_token_timesteps=per_token_timesteps, + schedule_timesteps=schedule_timesteps, + ) + + torch.testing.assert_close(actual, expected) + + +def test_sana_transformer_keeps_conditioned_frame_fixed_with_generic_scheduler() -> None: + """Keep SANA conditioning out of the scheduler and inside the transformer.""" + scheduler = SanaWMLTXEulerSchedulerConfig(num_inference_steps=1).setup() + transformer = SanaWMTransformerConfig().setup() + initial_noise = torch.zeros((1, 1, 2, 1, 1), dtype=torch.float32) + initial_noise[:, :, 0] = 5.0 + conditioning = SanaWMStage1Conditioning( + condition=torch.ones((1, 1, 1, 1)), + uncondition=None, + model_kwargs={"data_info": {"condition_frame_info": {0: 0.0}}}, + first_latent=torch.empty((1, 1, 1, 1, 1)), + latent_shape=tuple(initial_noise.shape), + cfg_scale=1.0, + flow_shift=5.0, + steps=1, + seed=0, + ) + + class DummyModel(torch.nn.Module): + def forward( + self, + noisy_latent: torch.Tensor, + timestep: torch.Tensor, + _prompt_embeds: torch.Tensor, + **_kwargs: object, + ) -> torch.Tensor: + assert timestep.shape == (1, 1, 2) + torch.testing.assert_close(timestep[:, :, 0], torch.zeros((1, 1))) + assert torch.all(timestep[:, :, 1] > 0) + return -torch.ones_like(noisy_latent) + + transformer.model = DummyModel() + transformer._model_built = True + + def predict_flow(noisy_latent: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + return transformer.predict_flow( + noisy_latent=noisy_latent, + timestep=timestep, + cache=SanaWMTransformerCache(), + input=conditioning, + ) + + sampled = scheduler.sample( + initial_noise=initial_noise, + predict_flow=predict_flow, + ) + + torch.testing.assert_close(sampled[:, :, 0], initial_noise[:, :, 0]) + assert torch.all(sampled[:, :, 1] > initial_noise[:, :, 1]) + + +def test_transformer_contract_shape_and_conditioning_guard() -> None: + """Pin the SANA-WM Stage-1 boundary to the public model layout.""" + transformer_cfg = SanaWMTransformerConfig() + transformer = transformer_cfg.setup() + + assert transformer.latent_shape == (1, 128, 21, 22, 40) + cache = transformer.initialize_autoregressive_cache() + + with pytest.raises(RuntimeError, match="without conditioning"): + transformer.predict_flow( + noisy_latent=torch.empty(transformer.latent_shape), + timestep=torch.tensor(1000.0), + cache=cache, + ) + + +def test_transformer_initial_noise_uses_conditioning_payload() -> None: + """Generate first-frame-pinned noise through the shared transformer hook.""" + transformer = SanaWMTransformerConfig().setup() + transformer.weight_dtype = torch.float32 + first_latent = torch.full((1, 1, 1, 1, 1), 7.0) + conditioning = SanaWMStage1Conditioning( + condition=torch.empty((1, 1, 1, 1)), + uncondition=None, + model_kwargs={}, + first_latent=first_latent, + latent_shape=(1, 1, 2, 1, 1), + cfg_scale=1.0, + flow_shift=1.0, + steps=1, + seed=123, + ) + cache = transformer.initialize_autoregressive_cache() + + noise = transformer.initial_noise( + latent_shape=(1, 1, 2, 1, 1), + rng=None, + cache=cache, + input=conditioning, + ) + + assert cache.conditioning is conditioning + assert noise.shape == (1, 1, 2, 1, 1) + torch.testing.assert_close(noise[:, :, :1], first_latent) + + +def test_transformer_predict_flow_applies_cfg_from_conditioning_input() -> None: + """Keep CFG inside the transformer boundary used by base diffusion.""" + + class DummyModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.calls: list[dict[str, torch.Tensor]] = [] + + def forward( + self, + noisy_latent: torch.Tensor, + timestep: torch.Tensor, + prompt_embeds: torch.Tensor, + *, + mask: torch.Tensor, + camera_conditions: torch.Tensor, + chunk_plucker: torch.Tensor, + **_kwargs: object, + ) -> torch.Tensor: + assert "negative_mask" not in _kwargs + self.calls.append( + { + "noisy_latent": noisy_latent, + "timestep": timestep, + "prompt_embeds": prompt_embeds, + "mask": mask, + "camera_conditions": camera_conditions, + "chunk_plucker": chunk_plucker, + "data_info": _kwargs["data_info"], + } + ) + branch_values = prompt_embeds.flatten(1).mean(dim=1).reshape( + -1, + 1, + 1, + 1, + 1, + ) + return torch.ones_like(noisy_latent) * branch_values + + transformer = SanaWMTransformerConfig().setup() + dummy_model = DummyModel() + transformer.model = dummy_model + transformer._model_built = True + cond_mask = torch.ones((1, 1)) + neg_mask = torch.zeros((1, 1)) + camera = torch.zeros((1, 3, 2, 1, 1)) + chunk_plucker = torch.ones((1, 6, 2, 1, 1)) + conditioning = SanaWMStage1Conditioning( + condition=torch.ones((1, 1, 1, 1)), + uncondition=torch.zeros((1, 1, 1, 1)), + model_kwargs={ + "mask": cond_mask, + "negative_mask": neg_mask, + "camera_conditions": camera, + "chunk_plucker": chunk_plucker, + "data_info": {"condition_frame_info": {0: 0.0}}, + }, + first_latent=torch.empty((1, 1, 1, 1, 1)), + latent_shape=(1, 1, 2, 1, 1), + cfg_scale=2.0, + flow_shift=1.0, + steps=1, + seed=0, + ) + + out = transformer.predict_flow( + noisy_latent=torch.zeros((1, 1, 2, 1, 1)), + timestep=torch.full((1, 1, 2), 1000.0), + cache=SanaWMTransformerCache(), + input=conditioning, + ) + + torch.testing.assert_close( + out, + torch.tensor([[[[[0.0]], [[2.0]]]]]), + ) + assert len(dummy_model.calls) == 1 + call = dummy_model.calls[0] + assert call["noisy_latent"].shape == (2, 1, 2, 1, 1) + assert call["timestep"].shape == (2, 1, 2) + assert call["prompt_embeds"].shape == (2, 1, 1, 1) + torch.testing.assert_close(call["mask"], torch.cat([neg_mask, cond_mask], dim=0)) + torch.testing.assert_close(call["camera_conditions"], camera) + torch.testing.assert_close( + call["chunk_plucker"], + torch.cat([chunk_plucker, chunk_plucker], dim=0), + ) + assert call["data_info"] == {"condition_frame_info": {0: 0.0}} + + +def test_transformer_predict_flow_caches_static_plucker_embedding() -> None: + """Avoid re-projecting rollout-constant Plucker conditioning each step.""" + + class DummyModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.prepare_calls = 0 + self.prepare_camera_calls = 0 + self.forward_calls: list[dict[str, object]] = [] + self.camera_cache = object() + + def prepare_plucker_embedding( + self, + chunk_plucker: torch.Tensor, + ) -> torch.Tensor: + self.prepare_calls += 1 + return chunk_plucker.flatten(2).transpose(1, 2) + + def prepare_camera_projection_cache( + self, + camera_conditions: torch.Tensor, + *, + frames: int, + height: int, + width: int, + ) -> tuple[torch.Tensor, object]: + del camera_conditions + self.prepare_camera_calls += 1 + return torch.empty(frames, height, width), self.camera_cache + + def forward( + self, + noisy_latent: torch.Tensor, + timestep: torch.Tensor, + prompt_embeds: torch.Tensor, + *, + chunk_plucker_emb: torch.Tensor, + rotary_emb: torch.Tensor, + camera_cache: object, + **_kwargs: object, + ) -> torch.Tensor: + del timestep, prompt_embeds + assert "chunk_plucker" not in _kwargs + self.forward_calls.append( + { + "chunk_plucker_emb": chunk_plucker_emb, + "rotary_emb": rotary_emb, + "camera_cache": camera_cache, + } + ) + return torch.ones_like(noisy_latent) + + transformer = SanaWMTransformerConfig().setup() + dummy_model = DummyModel() + transformer.model = dummy_model + transformer._model_built = True + chunk_plucker = torch.ones((1, 2, 2, 1, 1)) + conditioning = SanaWMStage1Conditioning( + condition=torch.ones((1, 1, 1)), + uncondition=None, + model_kwargs={ + "chunk_plucker": chunk_plucker, + "camera_conditions": torch.zeros((1, 2, 20)), + "data_info": {"condition_frame_info": {0: 0.0}}, + }, + first_latent=torch.empty((1, 1, 1, 1, 1)), + latent_shape=(1, 1, 2, 1, 1), + cfg_scale=1.0, + flow_shift=1.0, + steps=1, + seed=0, + ) + cache = SanaWMTransformerCache() + + for _ in range(2): + transformer.predict_flow( + noisy_latent=torch.zeros((1, 1, 2, 1, 1)), + timestep=torch.full((1, 1, 2), 1000.0), + cache=cache, + input=conditioning, + ) + + assert dummy_model.prepare_calls == 1 + assert dummy_model.prepare_camera_calls == 1 + assert "chunk_plucker" not in conditioning.model_kwargs + assert "camera_cache" in conditioning.model_kwargs + assert len(dummy_model.forward_calls) == 2 + cached = conditioning.model_kwargs["chunk_plucker_emb"] + assert isinstance(cached, torch.Tensor) + assert all( + call["chunk_plucker_emb"] is cached for call in dummy_model.forward_calls + ) + assert all( + call["camera_cache"] is dummy_model.camera_cache + for call in dummy_model.forward_calls + ) + + +def test_inference_config_loads_yaml( + tmp_path: Path, +) -> None: + """Parse SANA-WM YAML with local config objects.""" + config_path = tmp_path / "config.yaml" + config_path.write_text( + """ +model: + model: SanaMSVideoCamCtrl_1600M_P1_D20 + mixed_precision: bf16 + chunk_split_strategy: first_chunk_plus_one +vae: + vae_type: LTX2VAE_diffusers + vae_pretrained: hf://example/model + weight_dtype: bfloat16 + vae_latent_dim: 128 + vae_stride: [8, 32, 32] +text_encoder: + text_encoder_name: gemma-2-2b-it + model_max_length: 300 + chi_prompt: ["prefix"] +scheduler: + flow_shift: 9.95 + inference_flow_shift: 9.8 +""", + encoding="utf-8", + ) + + cfg = _load_inference_config(str(config_path)) + + assert cfg.model.model == "SanaMSVideoCamCtrl_1600M_P1_D20" + assert cfg.model.get("missing", "fallback") == "fallback" + assert cfg.vae.vae_stride == [8, 32, 32] + assert cfg.text_encoder.chi_prompt == ["prefix"] + assert cfg.scheduler.inference_flow_shift == 9.8 + assert cfg.work_dir == "" + + +def test_text_prompt_encoder_outputs_padded_prompt_schema( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Encode prompts behind the explicit text component boundary.""" + + class DummyTokens: + def __init__(self, max_length: int) -> None: + self.input_ids = torch.arange(max_length).reshape(1, max_length) + self.attention_mask = torch.ones((1, max_length), dtype=torch.long) + + def to(self, _device: torch.device) -> "DummyTokens": + return self + + class DummyTokenizer: + def encode(self, _text: str) -> list[int]: + return [1, 2] + + def __call__( + self, + *_args: object, + max_length: int, + **_kwargs: object, + ) -> DummyTokens: + return DummyTokens(max_length) + + class DummyTextEncoder(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.empty(0)) + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor]: + del attention_mask + hidden = torch.ones((1, input_ids.shape[1], 12), dtype=torch.float32) + return (hidden,) + + monkeypatch.setenv("SANA_WM_STAGE1_NVFP4_TEXT_PAD_MULTIPLE", "8") + monkeypatch.setattr( + conditioning_module, + "_load_inference_config", + lambda _path: SimpleNamespace( + model=SimpleNamespace(mixed_precision="bf16"), + text_encoder=SimpleNamespace( + text_encoder_name="T5-small", + model_max_length=5, + chi_prompt=["prefix"], + ), + ), + ) + monkeypatch.setattr( + conditioning_module, + "_get_tokenizer_and_text_encoder", + lambda **_kwargs: (DummyTokenizer(), DummyTextEncoder()), + ) + encoder = SanaWMTextPromptEncoderConfig( + config_path="dummy.yaml", + stage1_precision="fp4", + ).setup() + + encoded = encoder( + SanaWMTextPromptRequest( + prompt="drive forward", + negative_prompt="low quality", + ) + ) + + assert encoded.condition.shape == (1, 1, 8, 12) + assert encoded.condition_mask.shape == (1, 8) + assert encoded.negative.shape == (1, 1, 8, 12) + assert encoded.negative_mask.shape == (1, 8) + + +def test_first_frame_encoder_outputs_latent_shape( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Encode the first frame behind the named I2V component.""" + from PIL import Image + + class DummyLatentDist: + def mode(self) -> torch.Tensor: + return torch.ones((1, 4, 1, 2, 3), dtype=torch.float32) + + class DummyPosterior: + latent_dist = DummyLatentDist() + + class DummyVAE(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.empty(0)) + self.latents_mean = torch.zeros(4) + self.latents_std = torch.ones(4) + self.config = SimpleNamespace(scaling_factor=1.0) + + def encode(self, _images: torch.Tensor) -> DummyPosterior: + return DummyPosterior() + + monkeypatch.setattr( + conditioning_module, + "_load_inference_config", + lambda _path: SimpleNamespace( + model=SimpleNamespace(mixed_precision="bf16"), + vae=SimpleNamespace( + weight_dtype="float32", + vae_type="LTX2VAE_diffusers", + vae_pretrained="local", + ), + ), + ) + monkeypatch.setattr( + conditioning_module, + "_get_vae", + lambda *_args, **_kwargs: DummyVAE(), + ) + encoder = SanaWMFirstFrameEncoderConfig(config_path="dummy.yaml").setup() + + latent = encoder(Image.new("RGB", (4, 4))) + + assert latent.shape == (1, 4, 1, 2, 3) + assert latent.dtype == torch.bfloat16 + + +def test_camera_conditioning_encoder_outputs_sana_schema() -> None: + """Build raymap/chunk-Plucker tensors behind the camera component.""" + poses = np.broadcast_to(np.eye(4, dtype=np.float32), (17, 4, 4)).copy() + intrinsics = np.broadcast_to( + np.array([900.0, 900.0, 640.0, 352.0], dtype=np.float32), + (17, 4), + ).copy() + encoder = SanaWMCameraConditioningEncoderConfig().setup() + + camera = encoder( + SanaWMCameraRequest( + poses_c2w=poses, + intrinsics_vec4=intrinsics, + ) + ) + + assert camera["raymap"].shape == (3, 20) + assert camera["chunk_plucker"].shape == (48, 3, 22, 40) + + +def test_video_decoder_returns_structured_video(monkeypatch: pytest.MonkeyPatch) -> None: + """Decode Stage-1 latents through the explicit decoder component.""" + decoder = SanaWMVideoDecoderConfig(refiner=None).setup() + monkeypatch.setattr( + decoder.vae_decoder, + "decode_latents", + lambda _latents: np.zeros((1, 2, 2, 3), dtype=np.uint8), + ) + + decoded = decoder( + torch.zeros((1, 4, 1, 1, 1)), + autoregressive_index=0, + cache=None, + ) + + assert isinstance(decoded, SanaWMDecodedVideo) + assert decoded.video_hwc.shape == (1, 2, 2, 3) + assert decoded.stage1_video_hwc is None + + +def test_stage1_model_matches_checkpoint_schema() -> None: + """Pin the Stage-1 module to the public checkpoint schema.""" + state = SanaWMStage1Model().state_dict() + + assert SANA_WM_STAGE1_SPEC.chunk_size is None + assert SANA_WM_STAGE1_SPEC.chunk_split_strategy == "first_chunk_plus_one" + assert len(state) == 872 + assert tuple(state["x_embedder.proj.weight"].shape) == (2240, 128, 1, 1, 1) + assert tuple(state["raymap_embedder.proj.weight"].shape) == (2240, 3, 1, 1, 1) + assert tuple(state["plucker_embedder.proj.weight"].shape) == ( + 2240, + 48, + 1, + 1, + 1, + ) + assert tuple(state["pos_embed"].shape) == (1, 484, 2240) + assert tuple(state["y_embedder.y_embedding"].shape) == (300, 2304) + assert tuple(state["blocks.0.attn.qkv.weight"].shape) == (6720, 2240) + assert tuple(state["blocks.0.attn.conv_k.weight"].shape) == (2240, 1, 4) + assert tuple(state["blocks.0.cross_attn.kv_linear.weight"].shape) == ( + 4480, + 2240, + ) + assert tuple(state["blocks.19.mlp.t_conv.weight"].shape) == (2240, 2240, 3, 1) + assert tuple(state["final_layer.linear.weight"].shape) == (128, 2240) + + for block_index in range(SANA_WM_STAGE1_SPEC.depth): + has_gdn_conv = f"blocks.{block_index}.attn.conv_k.weight" in state + assert has_gdn_conv is SANA_WM_STAGE1_SPEC.block_uses_gdn(block_index) + + +def test_stage1_forward_preserves_latent_shape() -> None: + """Exercise the Stage-1 forward path on a small CPU-safe spec.""" + spec = SanaWMStage1Spec( + latent_channels=4, + 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, + conv_kernel_size=3, + temporal_kernel_size=3, + plucker_channels=6, + raymap_channels=3, + softmax_every_n=2, + ) + model = SanaWMStage1Model(spec) + latents = torch.randn(1, 4, 3, 2, 2) + timesteps = torch.ones(1, 1, 3) + text = torch.randn(1, 1, 5, 12) + mask = torch.ones(1, 5) + plucker = torch.randn(1, 6, 3, 2, 2) + + out = model(latents, timesteps, text, mask=mask, chunk_plucker=plucker) + + assert out.shape == latents.shape + + +def test_stage1_forward_accepts_precomputed_camera_cache() -> None: + """Precomputed camera cache must match the raw camera-conditioning path.""" + torch.manual_seed(0) + spec = SanaWMStage1Spec( + latent_channels=4, + hidden_size=16, + text_dim=12, + timestep_dim=8, + depth=1, + num_heads=2, + head_dim=8, + max_text_length=5, + latent_grid_size=(2, 2), + mlp_ratio=1, + conv_kernel_size=3, + temporal_kernel_size=3, + plucker_channels=6, + raymap_channels=3, + softmax_every_n=1, + ) + model = SanaWMStage1Model(spec).eval() + with torch.no_grad(): + for param in model.parameters(): + param.normal_(mean=0.0, std=0.02) + latents = torch.randn(1, 4, 3, 2, 2) + timesteps = torch.ones(1, 1, 3) + text = torch.randn(1, 1, 5, 12) + mask = torch.ones(1, 5) + plucker = torch.randn(1, 6, 3, 2, 2) + camera = torch.zeros(1, 3, 20) + camera[..., :16] = torch.eye(4).flatten() + camera[..., 16:] = torch.tensor([1.0, 1.0, 0.5, 0.5]) + rotary_emb, camera_cache = model.prepare_camera_projection_cache( + camera, + frames=3, + height=2, + width=2, + ) + + raw = model( + latents, + timesteps, + text, + mask=mask, + chunk_plucker=plucker, + camera_conditions=camera, + ) + cached = model( + latents, + timesteps, + text, + mask=mask, + chunk_plucker=plucker, + camera_conditions=camera, + rotary_emb=rotary_emb, + camera_cache=camera_cache, + ) + + torch.testing.assert_close(cached, raw) + + +def test_stage1_self_attention_uses_camera_conditions() -> None: + """Changing camera conditions must affect camera attention.""" + torch.manual_seed(0) + spec = SanaWMStage1Spec( + latent_channels=4, + hidden_size=16, + text_dim=12, + timestep_dim=8, + depth=1, + num_heads=2, + head_dim=8, + max_text_length=5, + latent_grid_size=(2, 2), + mlp_ratio=1, + conv_kernel_size=3, + temporal_kernel_size=3, + plucker_channels=6, + raymap_channels=3, + softmax_every_n=2, + ) + attn = Stage1SelfAttention(spec, use_gdn_convs=True).eval() + hidden = torch.randn(1, 8, 16) + camera = torch.zeros(1, 2, 20) + camera[..., :16] = torch.eye(4).flatten() + camera[..., 16:] = torch.tensor([1.0, 1.0, 0.5, 0.5]) + shifted_camera = camera.clone() + shifted_camera[..., 3] = 0.25 + + base = attn(hidden, HW=(2, 2, 2), camera_conditions=camera) + shifted = attn(hidden, HW=(2, 2, 2), camera_conditions=shifted_camera) + + assert base.shape == hidden.shape + assert not torch.allclose(base, shifted) + + +def test_stage1_self_attention_broadcasts_single_camera_batch() -> None: + """CFG can reuse one camera projection cache for both prompt branches.""" + torch.manual_seed(1) + spec = SanaWMStage1Spec( + latent_channels=4, + hidden_size=16, + text_dim=12, + timestep_dim=8, + depth=1, + num_heads=2, + head_dim=8, + max_text_length=5, + latent_grid_size=(2, 2), + mlp_ratio=1, + conv_kernel_size=3, + temporal_kernel_size=3, + plucker_channels=6, + raymap_channels=3, + softmax_every_n=2, + ) + attn = Stage1SelfAttention(spec, use_gdn_convs=True).eval() + hidden = torch.randn(2, 8, 16) + camera = torch.zeros(1, 2, 20) + camera[..., :16] = torch.eye(4).flatten() + camera[..., 16:] = torch.tensor([1.0, 1.0, 0.5, 0.5]) + + broadcast = attn(hidden, HW=(2, 2, 2), camera_conditions=camera) + duplicated = attn( + hidden, + HW=(2, 2, 2), + camera_conditions=torch.cat([camera, camera], dim=0), + ) + + assert broadcast.shape == hidden.shape + torch.testing.assert_close(broadcast, duplicated) + + +def test_stage1_self_attention_uses_gdn_main_without_gdn_convs() -> None: + """Blocks without GDN conv weights still route main attention through GDN.""" + spec = SanaWMStage1Spec( + latent_channels=4, + hidden_size=16, + text_dim=12, + timestep_dim=8, + depth=1, + num_heads=2, + head_dim=8, + max_text_length=5, + latent_grid_size=(2, 2), + mlp_ratio=1, + conv_kernel_size=3, + temporal_kernel_size=3, + plucker_channels=6, + raymap_channels=3, + softmax_every_n=1, + ) + attn = Stage1SelfAttention(spec, use_gdn_convs=False).eval() + hidden = torch.randn(1, 4, 16) + called = {"gdn_main": False} + + def fake_gdn_main( + x: torch.Tensor, + *, + HW: tuple[int, int, int], + rotary_emb: torch.Tensor | None, + precomputed_gates: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + del HW, rotary_emb + called["gdn_main"] = True + assert precomputed_gates[0].shape == (1, 2, 1, 4) + return torch.zeros_like(x) + + def fail_softmax_main(*_args: object, **_kwargs: object) -> torch.Tensor: + raise AssertionError("main attention should not use softmax") + + attn._forward_gdn_main = fake_gdn_main # type: ignore[method-assign] + attn._forward_softmax_main = fail_softmax_main # type: ignore[method-assign] + + out = attn(hidden, HW=(1, 2, 2), apply_output_gate=False) + + assert called["gdn_main"] + torch.testing.assert_close(out, torch.zeros_like(hidden)) + + +def test_transformer_releases_stage1_runtime() -> None: + """Free Stage-1-only modules and conditioning before decode/refine.""" + transformer = SanaWMTransformerConfig().setup() + transformer.model = torch.nn.Linear(1, 1) + transformer._model_built = True + transformer._stage1_quantized = True + cache = SanaWMTransformerCache( + conditioning=SanaWMStage1Conditioning( + condition=torch.empty(1), + uncondition=None, + model_kwargs={}, + first_latent=torch.empty(1), + latent_shape=(1, 1, 1, 1, 1), + cfg_scale=1.0, + flow_shift=1.0, + steps=1, + seed=0, + ) + ) + + transformer.release_stage1_runtime(cache) + + assert cache.conditioning is None + assert transformer.model is None + assert transformer._model_built is False + assert transformer._stage1_quantized is False + + +def _stage1_conditioning_cache() -> SanaWMTransformerCache: + return SanaWMTransformerCache( + conditioning=SanaWMStage1Conditioning( + condition=torch.empty(1), + uncondition=None, + model_kwargs={}, + first_latent=torch.empty(1), + latent_shape=(1, 1, 1, 1, 1), + cfg_scale=1.0, + flow_shift=1.0, + steps=1, + seed=0, + ) + ) + + +def test_postprocess_keeps_stage1_resident_by_default() -> None: + """Default keeps the Stage-1 DiT resident so a reused pipeline never reloads.""" + transformer = SanaWMTransformerConfig().setup() + assert transformer.config.offload_stage1 is False + model = torch.nn.Linear(1, 1) + transformer.model = model + transformer._model_built = True + transformer._stage1_quantized = True + cache = _stage1_conditioning_cache() + latent = torch.zeros(1) + + result = transformer.postprocess_clean_latent(latent, cache) + + assert result is latent + assert transformer.model is model + assert transformer._model_built is True + assert transformer._stage1_quantized is True + # Per-rollout conditioning is still dropped even when the model stays warm. + assert cache.conditioning is None + + +def test_postprocess_releases_stage1_when_offloading() -> None: + """offload_stage1 restores the free-before-decode behavior.""" + transformer = SanaWMTransformerConfig(offload_stage1=True).setup() + transformer.model = torch.nn.Linear(1, 1) + transformer._model_built = True + transformer._stage1_quantized = True + cache = _stage1_conditioning_cache() + + transformer.postprocess_clean_latent(torch.zeros(1), cache) + + assert transformer.model is None + assert transformer._model_built is False + assert cache.conditioning is None + + +def test_vae_tiling_defaults_match_upstream_fast_path() -> None: + """Keep VAE decode on the upstream-style fast tiles by default.""" + + class DummyVAE: + def __init__(self) -> None: + self.calls: list[dict[str, int]] = [] + self.tile_sample_min_height = 512 + self.tile_sample_stride_height = 448 + self.tile_sample_min_width = 512 + self.tile_sample_stride_width = 448 + self.tile_sample_min_num_frames = 96 + self.tile_sample_stride_num_frames = 64 + self.use_framewise_encoding = False + self.use_framewise_decoding = False + self.spatial_compression_ratio = 32 + + def enable_tiling(self, **kwargs: int) -> None: + self.calls.append(kwargs) + + decoder = SanaWMLTX2VAEDecoderConfig().setup() + decoder.vae = DummyVAE() + + decoder._configure_vae_tiling() + + assert decoder.vae.calls == [ + { + "tile_sample_min_height": 512, + "tile_sample_stride_height": 448, + "tile_sample_min_width": 512, + "tile_sample_stride_width": 448, + "tile_sample_min_num_frames": 96, + "tile_sample_stride_num_frames": 64, + } + ] + assert decoder.vae.tile_sample_min_height == 512 + assert decoder.vae.tile_sample_stride_height == 448 + assert decoder.vae.tile_sample_min_width == 512 + assert decoder.vae.tile_sample_stride_width == 448 + assert decoder.vae.tile_sample_min_num_frames == 96 + assert decoder.vae.tile_sample_stride_num_frames == 64 + assert decoder.vae.use_framewise_encoding is True + assert decoder.vae.use_framewise_decoding is True + + +def test_decode_retries_vae_oom_with_smaller_tiles( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Retry VAE decode with smaller tiles after a CUDA OOM.""" + + class DummyVAE: + def __init__(self) -> None: + self.tile_sample_min_height = 256 + self.tile_sample_stride_height = 224 + self.tile_sample_min_width = 256 + self.tile_sample_stride_width = 224 + self.tile_sample_min_num_frames = 24 + self.tile_sample_stride_num_frames = 8 + self.use_framewise_encoding = False + self.use_framewise_decoding = False + self.spatial_compression_ratio = 32 + self.temporal_compression_ratio = 8 + + def to(self, *_args: object, **_kwargs: object) -> "DummyVAE": + return self + + def enable_tiling(self, **kwargs: int) -> None: + for name, value in kwargs.items(): + setattr(self, name, value) + + decoder = SanaWMLTX2VAEDecoderConfig().setup() + decoder.vae = DummyVAE() + decoder.vae_dtype = torch.float32 + monkeypatch.setattr(decoder, "_ensure_vae", lambda: None) + calls = 0 + inference_modes: list[bool] = [] + + def decode_once(_samples: torch.Tensor) -> torch.Tensor: + nonlocal calls + calls += 1 + inference_modes.append(torch.is_inference_mode_enabled()) + if calls == 1: + raise torch.OutOfMemoryError("test OOM") + return torch.zeros((1, 3, 1, 2, 2), dtype=torch.float32) + + monkeypatch.setattr(decoder, "_vae_decode", decode_once) + + video = decoder.decode_latents(torch.zeros((1, 4, 1, 1, 1))) + + assert calls == 2 + assert inference_modes == [True, True] + assert video.shape == (1, 2, 2, 3) + assert decoder.vae.tile_sample_min_height == 128 + assert decoder.vae.tile_sample_stride_height == 64 + assert decoder.vae.tile_sample_min_width == 128 + assert decoder.vae.tile_sample_stride_width == 64 + assert decoder.vae.tile_sample_min_num_frames == 16 + assert decoder.vae.tile_sample_stride_num_frames == 8 + + +def test_vae_tiling_avoids_degenerate_latent_tails() -> None: + """Avoid last spatial VAE tiles with size one after compression.""" + assert ( + _avoid_degenerate_tile_tail( + sample_extent=704, + sample_tile_min=256, + sample_stride=224, + compression_ratio=32, + ) + == 192 + ) + assert ( + _avoid_degenerate_tile_tail( + sample_extent=1280, + sample_tile_min=128, + sample_stride=112, + compression_ratio=32, + ) + == 64 + ) + + +def test_hf_defaults_point_at_bidirectional_release() -> None: + """Pin every default SANA-WM artefact to the bidirectional HF repo.""" + assert SANA_WM_HF_REPO == "Efficient-Large-Model/SANA-WM_bidirectional" + assert SANA_WM_MODEL_PATH == ( + "hf://Efficient-Large-Model/SANA-WM_bidirectional/" + "dit/sana_wm_1600m_720p.safetensors" + ) + assert SANA_WM_CONFIG_PATH == ( + "hf://Efficient-Large-Model/SANA-WM_bidirectional/config.yaml" + ) + + +def test_runner_setup_preserves_cli_fields() -> None: + """Construct the runner and preserve CLI override fields.""" + cfg = derive_config( + RUNNER_SANA_WM_BIDIRECTIONAL, + image_path=Path("missing.png"), + prompt="demo", + ) + + runner = cfg.setup() + + assert isinstance(runner, SanaWMRunner) + assert runner.config.image_path == Path("missing.png") + + +def test_runner_derives_intrinsics_when_omitted(tmp_path: Path) -> None: + """Omitting --intrinsics-path derives per-frame intrinsics from the frame.""" + from PIL import Image + + image_path = tmp_path / "frame.png" + Image.new("RGB", (640, 360), color=(10, 20, 30)).save(image_path) + cfg = derive_config( + RUNNER_SANA_WM_BIDIRECTIONAL, + image_path=image_path, + prompt="demo", + intrinsics_path=None, + action="w-4", + num_frames=25, + ) + runner = cfg.setup() + + _image, c2w, intrinsics_vec4, num_frames = runner._prepare_inputs() + + assert num_frames == 25 + assert c2w.shape[0] == num_frames + assert intrinsics_vec4.shape == (num_frames, 4) + # Derived focal lengths are finite and positive after the crop transform. + assert np.all(np.isfinite(intrinsics_vec4)) + assert np.all(intrinsics_vec4[:, :2] > 0) + + +def test_runner_fits_camera_path_to_requested_frames(tmp_path: Path) -> None: + """Do not let short explicit trajectories cap the output frame count.""" + from PIL import Image + + image_path = tmp_path / "frame.png" + camera_path = tmp_path / "camera.npy" + Image.new("RGB", (640, 360), color=(10, 20, 30)).save(image_path) + c2w = np.broadcast_to(np.eye(4, dtype=np.float32), (3, 4, 4)).copy() + c2w[:, 2, 3] = [0.0, 1.0, 2.0] + np.save(camera_path, c2w) + cfg = derive_config( + RUNNER_SANA_WM_BIDIRECTIONAL, + image_path=image_path, + prompt="demo", + intrinsics_path=None, + camera_path=camera_path, + num_frames=25, + ) + runner = cfg.setup() + + _image, fitted_c2w, intrinsics_vec4, num_frames = runner._prepare_inputs() + + assert num_frames == 25 + assert fitted_c2w.shape[0] == num_frames + assert intrinsics_vec4.shape == (num_frames, 4) + np.testing.assert_allclose(fitted_c2w[[0, -1], 2, 3], [0.0, 2.0]) + + +def test_runner_config_type() -> None: + """Keep the exported literal on the SANA-WM runner config subclass.""" + assert isinstance(RUNNER_SANA_WM_BIDIRECTIONAL, SanaWMRunnerConfig) + + +def test_runner_defaults_to_bf16_precision() -> None: + """Keep the default runner on BF16 precision.""" + assert RUNNER_SANA_WM_BIDIRECTIONAL.stage1_precision == "bf16" + assert RUNNER_SANA_WM_BIDIRECTIONAL.refiner_precision == "bf16" + assert RUNNER_SANA_WM_BIDIRECTIONAL.quant_backend == "auto" + assert RUNNER_SANA_WM_BIDIRECTIONAL.no_refiner is False + + +def test_stage1_quant_scope_matches_upstream_precision_cli() -> None: + """Keep FP8 and FP4 on upstream's self-attn + cross-attn + FFN scope.""" + patterns = _stage1_quant_include_patterns() + + assert patterns is not None + assert r"^blocks\.\d+\.attn\.qkv$" in patterns + assert r"^blocks\.\d+\.attn\.output_gate$" in patterns + assert r"^blocks\.\d+\.cross_attn\." in patterns + assert r"^blocks\.\d+\.mlp\.inverted_conv\.linear$" in patterns + assert r"^blocks\.\d+\.mlp\.point_conv\.linear$" in patterns + + +def test_stage1_ffn_linearization_preserves_forward() -> None: + """Expose pointwise FFN convs as Linear modules without changing math.""" + torch.manual_seed(0) + spec = SanaWMStage1Spec( + latent_channels=4, + hidden_size=16, + text_dim=12, + timestep_dim=8, + depth=1, + num_heads=2, + head_dim=8, + max_text_length=5, + latent_grid_size=(2, 2), + mlp_ratio=1, + conv_kernel_size=3, + temporal_kernel_size=3, + plucker_channels=6, + raymap_channels=3, + softmax_every_n=2, + ) + mlp = GLUMBConvTemp(spec).eval() + inputs = torch.randn(1, 8, 16) + reference = mlp(inputs, frames=2, height=2, width=2) + + converted, skipped = linearize_stage1_ffn_for_quant(mlp) + output = mlp(inputs, frames=2, height=2, width=2) + + assert converted == 2 + assert skipped == 0 + assert hasattr(mlp.inverted_conv, "linear") + assert hasattr(mlp.point_conv, "linear") + torch.testing.assert_close(output, reference) + + +def test_fp4_rht16_is_orthogonal() -> None: + """The tiled Hadamard rotation must preserve dot products before quantization.""" + inputs = torch.randn(3, 4, 32) + weights = torch.randn(5, 32) + + rotated_inputs = apply_rht16(inputs) + rotated_weights = apply_rht16(weights) + + torch.testing.assert_close( + rotated_inputs.reshape(-1, 32) @ rotated_weights.t(), + inputs.reshape(-1, 32) @ weights.t(), + atol=1.0e-5, + rtol=1.0e-5, + ) + + +def test_fp4_rht16_rejects_unaligned_last_dimension() -> None: + """NVFP4 RHT is tiled in groups of sixteen values.""" + with pytest.raises(ValueError, match="divisible by 16"): + apply_rht16(torch.randn(2, 15)) + + +def test_fp4_global_scale_uses_nvfp4_dynamic_range() -> None: + """The hierarchical FP4 global scale should target E4M3 x E2M1 range.""" + inputs = torch.tensor([-2688.0, 0.0, 1344.0]) + + assert nvfp4_global_scale(inputs).item() == pytest.approx(1.0) + + +def test_refiner_is_flashdreams_owned( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Build the refiner adapter.""" + calls: dict[str, object] = {} + + class DummyRefiner: + def __init__(self, **kwargs: object) -> None: + calls.update(kwargs) + + monkeypatch.setattr(refiner_module, "SanaWMLTX2Refiner", DummyRefiner) + monkeypatch.setattr( + decoder_module, + "resolve_hf_path", + lambda value: f"/resolved/{value}", + ) + refiner = SanaWMLTX2LatentRefinerConfig().setup() + + refiner._ensure_refiner() + + assert refiner._refiner_built is True + assert isinstance(refiner.refiner, DummyRefiner) + assert calls["refiner_root"] == ( + "/resolved/hf://Efficient-Large-Model/SANA-WM_bidirectional/refiner" + ) + assert calls["gemma_root"] == ( + "/resolved/hf://Efficient-Large-Model/SANA-WM_bidirectional/refiner/text_encoder" + ) + assert calls["dtype"] is torch.bfloat16 + assert calls["precision"] == "bf16" + assert calls["quant_backend"] == "torch" + + +def test_refiner_latent_pack_round_trips() -> None: + """Preserve LTX-2 latent layout across token packing and unpacking.""" + latents = torch.arange(1 * 4 * 3 * 2 * 2, dtype=torch.float32).reshape( + 1, + 4, + 3, + 2, + 2, + ) + + packed = _pack_latents(latents, patch_size=1, patch_size_t=1) + unpacked = _unpack_latents( + packed, + num_frames=3, + height=2, + width=2, + patch_size=1, + patch_size_t=1, + ) + + torch.testing.assert_close(unpacked, latents) + + +def test_refiner_sink_bidirectional_path_runs_under_inference_mode( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Keep the supported LTX-2 refiner path explicit and inference-only.""" + + class DummyTransformer(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = type("Config", (), {"patch_size": 1, "patch_size_t": 1})() + + refiner = SanaWMLTX2Refiner.__new__(SanaWMLTX2Refiner) + torch.nn.Module.__init__(refiner) + refiner.device = torch.device("cpu") + refiner.dtype = torch.float32 + refiner.transformer = DummyTransformer() + inference_modes: list[bool] = [] + monkeypatch.setattr(refiner, "_prepare_quantization", lambda: None) + monkeypatch.setattr( + refiner, + "_encode_prompt", + lambda _prompt: (torch.zeros((1, 1, 1)), torch.ones((1, 1))), + ) + + def predict_current_x0( + *, + sink: torch.Tensor, + noisy_current: torch.Tensor, + prompt_embeds: torch.Tensor, + prompt_attention_mask: torch.Tensor, + sigma: torch.Tensor, + fps: float, + ) -> torch.Tensor: + del sink, prompt_embeds, prompt_attention_mask, sigma, fps + inference_modes.append(torch.is_inference_mode_enabled()) + return torch.zeros_like(_pack_latents(noisy_current)) + + monkeypatch.setattr(refiner, "_predict_current_x0", predict_current_x0) + latents = torch.zeros((1, 2, 3, 1, 1)) + + refined = refiner.refine_latents( + latents, + "demo", + fps=16.0, + progress=False, + sigmas=(0.5, 0.0), + ) + + assert refined.shape == latents.shape + assert inference_modes == [True] + + +def test_auto_quant_backend_resolves_to_torch_backend() -> None: + """Keep default quantization on the Torch backend.""" + assert _resolve_quant_backend("auto", ["fp4"]) == "torch" + assert _resolve_quant_backend("auto", ["fp8", "fp4"]) == "torch" + assert _resolve_quant_backend("auto", []) == "torch" + + +def test_quantized_precision_requires_cuda() -> None: + """Reject fp8/fp4 before loading checkpoints on CPU-only devices.""" + with pytest.raises(ValueError, match="requires a CUDA device"): + _validate_precision_request( + device=torch.device("cpu"), + stage1_precision="fp8", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="torch", + ) + + +def test_bf16_precision_does_not_require_cuda_or_quant_backend() -> None: + """Keep BF16 as the TE-free compatibility path.""" + _validate_precision_request( + device=torch.device("cpu"), + stage1_precision="bf16", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="auto", + ) + + +def test_fp8_precision_requires_hopper(monkeypatch: pytest.MonkeyPatch) -> None: + """Reject FP8 on pre-Hopper GPUs before checking quantization backend.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (8, 9)) + + with pytest.raises(ValueError, match="requires a Hopper or newer GPU"): + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp8", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="torch", + ) + + +def test_fp4_precision_requires_blackwell(monkeypatch: pytest.MonkeyPatch) -> None: + """Reject NVFP4 on Hopper-class GPUs before checking quantization backend.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (9, 0)) + + with pytest.raises(ValueError, match="requires a Blackwell GPU"): + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp4", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="torch", + ) + + +def test_torch_fp8_backend_validates_primitives( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Allow FP8 validation with the required PyTorch primitives.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) + monkeypatch.setattr(torch.version, "cuda", "12.8") + monkeypatch.setattr(torch, "_scaled_mm", object(), raising=False) + monkeypatch.setattr(torch, "float8_e4m3fn", torch.uint8, raising=False) + + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp8", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="torch-fp8", + ) + + +def test_auto_fp8_backend_uses_torch_primitives( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Route default FP8 to the Torch backend.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) + monkeypatch.setattr(torch.version, "cuda", "12.8") + monkeypatch.setattr(torch, "_scaled_mm", object(), raising=False) + monkeypatch.setattr(torch, "float8_e4m3fn", torch.uint8, raising=False) + + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp8", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="auto", + ) + + +def test_runner_reaches_normal_input_validation() -> None: + """Reach normal file input validation.""" + cfg = derive_config( + RUNNER_SANA_WM_BIDIRECTIONAL, + image_path=Path("missing.png"), + prompt="demo", + ) + runner = cfg.setup() + + with pytest.raises(FileNotFoundError, match="missing.png"): + runner.run() + + +def test_torch_fp8_backend_rejects_fp4(monkeypatch: pytest.MonkeyPatch) -> None: + """Do not silently route FP4 to a backend without an NVFP4 kernel.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) + + with pytest.raises(ValueError, match="accepts fp8 only"): + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp4", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="torch-fp8", + ) + + +def test_torch_fp4_backend_validates_primitives(monkeypatch: pytest.MonkeyPatch) -> None: + """Allow FP4 validation with the required PyTorch primitives.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) + monkeypatch.setattr(torch, "_scaled_mm", object(), raising=False) + monkeypatch.setattr(torch, "float4_e2m1fn_x2", torch.uint8, raising=False) + monkeypatch.setattr(torch, "float8_e4m3fn", torch.uint8, raising=False) + + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp4", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="torch-fp4", + ) + + +def test_torch_fp4_backend_rejects_fp8(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep precision-specific backends explicit.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) + + with pytest.raises(ValueError, match="accepts fp4 only"): + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp8", + refiner_precision="bf16", + refiner_enabled=True, + quant_backend="torch-fp4", + ) + + +def test_torch_backend_allows_mixed_fp8_fp4( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Allow mixed FP8/FP4 requests on the Torch backend.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (12, 0)) + monkeypatch.setattr(torch.version, "cuda", "12.8") + monkeypatch.setattr(torch, "_scaled_mm", object(), raising=False) + monkeypatch.setattr(torch, "float4_e2m1fn_x2", torch.uint8, raising=False) + monkeypatch.setattr(torch, "float8_e4m3fn", torch.uint8, raising=False) + + _validate_precision_request( + device=torch.device("cuda:0"), + stage1_precision="fp8", + refiner_precision="fp4", + refiner_enabled=True, + quant_backend="torch", + ) + + +def test_torch_fp8_linear_replaces_matching_modules() -> None: + """Provide a Torch replacement for eligible Linear modules.""" + if not hasattr(torch, "float8_e4m3fn"): + pytest.skip("torch.float8_e4m3fn is required for FP8 replacement") + + module = torch.nn.Sequential( + torch.nn.Linear(16, 32, bias=True), + torch.nn.Sequential(torch.nn.Linear(16, 32, bias=False)), + torch.nn.Linear(15, 32, bias=False), + ) + + converted, skipped = replace_linear_with_torch_fp8( + module, + recipe=None, + params_dtype=torch.bfloat16, + skip_patterns=(), + include_patterns=("^0$", "^1\\.0$", "^2$"), + ) + + assert converted == 2 + assert skipped == 1 + assert isinstance(module[0], TorchScaledMMFP8Linear) + assert isinstance(module[1][0], TorchScaledMMFP8Linear) + assert isinstance(module[2], torch.nn.Linear) + assert module[0].weight.shape == (32, 16) + assert module[0].bias is not None + + +def test_torch_fp4_linear_replaces_matching_modules( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Route eligible FP4 modules through the replacement helper.""" + module = torch.nn.Sequential( + torch.nn.Linear(32, 64, bias=True), + torch.nn.Sequential(torch.nn.Linear(32, 64, bias=False)), + torch.nn.Linear(16, 64, bias=False), + ) + + @classmethod + def from_linear( + cls: type[TorchScaledMMFP4Linear], + source: torch.nn.Linear, + *, + out_dtype: torch.dtype, + use_rht: bool = True, + use_global_scale: bool = True, + weight_scale_2d: bool = True, + ) -> TorchScaledMMFP4Linear: + del out_dtype, use_rht, use_global_scale, weight_scale_2d + instance = cls.__new__(cls) + torch.nn.Module.__init__(instance) + instance.in_features = source.in_features + instance.out_features = source.out_features + return instance + + monkeypatch.setattr(TorchScaledMMFP4Linear, "from_linear", from_linear) + + converted, skipped = replace_linear_with_torch_fp4( + module, + recipe=None, + params_dtype=torch.bfloat16, + skip_patterns=(), + include_patterns=("^0$", "^1\\.0$", "^2$"), + ) + + assert converted == 2 + assert skipped == 1 + assert isinstance(module[0], TorchScaledMMFP4Linear) + assert isinstance(module[1][0], TorchScaledMMFP4Linear) + assert isinstance(module[2], torch.nn.Linear) + + +def test_pyproject_entry_point_matches_runner_literal() -> None: + """Keep the package entry point aligned with ``RUNNER_CONFIGS``.""" + pyproject = tomllib.loads( + Path("integrations/sana/pyproject.toml").read_text(encoding="utf-8") + ) + entry_points = pyproject["project"]["entry-points"][ENTRY_POINT_GROUP] + + assert entry_points == { + "sana-wm-bidirectional": "sana_wm.config:RUNNER_SANA_WM_BIDIRECTIONAL", + } + + +def test_pyproject_package_selection() -> None: + """Keep the package metadata aligned with the integration package.""" + pyproject = tomllib.loads( + Path("integrations/sana/pyproject.toml").read_text(encoding="utf-8") + ) + + packages = pyproject["tool"]["setuptools"]["packages"]["find"] + dependencies = set(pyproject["project"]["dependencies"]) + + assert "accelerate>=1.0" in dependencies + assert packages["include"] == ["sana_wm*"] + assert "diffusers>=0.36" in dependencies + assert "transformers>=5.0,<6" in dependencies diff --git a/integrations/sana/tests/test_stage1_cuda.py b/integrations/sana/tests/test_stage1_cuda.py new file mode 100644 index 000000000..d0f60ada6 --- /dev/null +++ b/integrations/sana/tests/test_stage1_cuda.py @@ -0,0 +1,280 @@ +# 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. + +"""CUDA tests for SANA-WM Stage-1 kernels.""" + +from __future__ import annotations + +import pytest +import torch + +import sana_wm.stage1_model as stage1 + +pytestmark = pytest.mark.ci_gpu + + +def _require_stage1_cuda() -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for Stage-1 kernel tests") + if stage1.triton is None: + pytest.skip("Triton is required for the Stage-1 temporal conv CUDA kernel") + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_bidirectional_temporal_conv_fast_path_matches_eager( + dtype: torch.dtype, +) -> None: + """Compare the CUDA temporal-conv fast path against the eager conv path.""" + _require_stage1_cuda() + torch.manual_seed(1234) + batch, frames, height, width, channels = 2, 5, 3, 4, 16 + x = torch.randn( + batch, + frames * height * width, + channels, + device="cuda", + dtype=dtype, + ) + conv = torch.nn.Conv1d( + channels, + channels, + kernel_size=4, + groups=channels, + bias=False, + device="cuda", + dtype=dtype, + ) + + reference = stage1._apply_bidirectional_temporal_conv_eager( + x, + conv, + (frames, height, width), + ) + with torch.inference_mode(): + actual = stage1._apply_bidirectional_temporal_conv( + x, + conv, + (frames, height, width), + ) + + assert actual.shape == x.shape + assert actual.dtype == dtype + torch.testing.assert_close( + actual.float(), + reference.float(), + rtol=2e-2, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("inverse_rope", [False, True]) +@pytest.mark.parametrize("matrix_batch", [1, 2]) +def test_ucpe_transform_apply_fast_path_matches_eager( + dtype: torch.dtype, + inverse_rope: bool, + matrix_batch: int, +) -> None: + """Compare the CUDA UCPE transform fast path against the eager path.""" + _require_stage1_cuda() + torch.manual_seed(4321) + batch, frames, height, width, heads, dim = 2, 3, 2, 3, 2, 16 + tokens = frames * height * width + x = torch.randn( + batch, + tokens, + heads, + dim, + device="cuda", + dtype=dtype, + ) + matrix = torch.randn( + matrix_batch, + tokens, + 4, + 4, + device="cuda", + dtype=dtype, + ) + rope = stage1._slice_rope_for_camera( + stage1._wan_rope_complex(dim, frames, height, width, x.device), + dim, + ) + + reference = stage1._ucpe_transform_apply_eager( + x, + matrix, + rope, + inverse_rope=inverse_rope, + ) + with torch.inference_mode(): + actual = stage1._ucpe_transform_apply( + x, + matrix, + rope, + inverse_rope=inverse_rope, + ) + + assert actual.shape == x.shape + assert actual.dtype is torch.float32 + torch.testing.assert_close( + actual, + reference, + rtol=2e-2 if dtype is torch.bfloat16 else 1e-5, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("inverse_rope", [False, True]) +@pytest.mark.parametrize("matrix_batch", [1, 2]) +def test_ucpe_transform_fast_path_matches_eager_norms( + dtype: torch.dtype, + inverse_rope: bool, + matrix_batch: int, +) -> None: + """Compare the UCPE fast path including GDN norm buffers.""" + _require_stage1_cuda() + torch.manual_seed(2468) + batch, frames, height, width, heads, dim = 2, 3, 2, 3, 2, 16 + tokens = frames * height * width + x = torch.randn( + batch, + tokens, + heads, + dim, + device="cuda", + dtype=dtype, + ) + matrix = torch.randn( + matrix_batch, + tokens, + 4, + 4, + device="cuda", + dtype=dtype, + ) + rope = stage1._slice_rope_for_camera( + stage1._wan_rope_complex(dim, frames, height, width, x.device), + dim, + ) + + reference = stage1._ucpe_transform_apply_eager( + x, + matrix, + rope, + inverse_rope=inverse_rope, + ) + reference_pre = stage1._ucpe_transform_norms_eager(x) + reference_post = stage1._ucpe_transform_norms_eager(reference) + with torch.inference_mode(): + actual, actual_pre, actual_post = stage1._ucpe_transform( + x, + matrix, + rope, + inverse_rope=inverse_rope, + ) + + assert actual.shape == x.shape + assert actual_pre.shape == (batch, heads, tokens) + assert actual_post.shape == (batch, heads, tokens) + torch.testing.assert_close( + actual, + reference, + rtol=2e-2 if dtype is torch.bfloat16 else 1e-5, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + ) + torch.testing.assert_close( + actual_pre, + reference_pre, + rtol=2e-2 if dtype is torch.bfloat16 else 1e-5, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + ) + torch.testing.assert_close( + actual_post, + reference_post, + rtol=2e-2 if dtype is torch.bfloat16 else 1e-5, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_silu_multiply_fast_path_matches_eager(dtype: torch.dtype) -> None: + """Compare the fused GLU elementwise fast path against PyTorch eager.""" + _require_stage1_cuda() + torch.manual_seed(1357) + value = torch.randn( + 4, + 17, + 5, + 7, + device="cuda", + dtype=dtype, + ).contiguous(memory_format=torch.channels_last) + gate = torch.randn_like(value).contiguous(memory_format=torch.channels_last) + + reference = value * torch.nn.functional.silu(gate) + with torch.inference_mode(): + actual = stage1._silu_multiply(value, gate, inplace=False) + + assert actual.shape == value.shape + assert actual.dtype == value.dtype + assert actual.is_contiguous(memory_format=torch.channels_last) + torch.testing.assert_close( + actual.float(), + reference.float(), + rtol=2e-2 if dtype is torch.bfloat16 else 1e-5, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("scale", [1.0, 0.0375]) +def test_rmsnorm_relu_heads_fast_path_matches_eager( + dtype: torch.dtype, + scale: float, +) -> None: + """Compare fused GDN RMSNorm/ReLU preparation against eager math.""" + _require_stage1_cuda() + torch.manual_seed(9753) + batch, tokens, heads, dim = 2, 13, 4, 8 + x = torch.randn( + batch, + tokens, + heads, + dim, + device="cuda", + dtype=dtype, + ) + weight = torch.randn(heads * dim, device="cuda", dtype=dtype) + eps = 1e-6 + inv_rms = stage1._inv_rms(x, eps) + reference = torch.nn.functional.relu( + x.float() * inv_rms[:, :, None, None] * weight.float().view(heads, dim) + ) + if scale != 1.0: + reference = reference * scale + + with torch.inference_mode(): + actual = stage1._rmsnorm_relu_heads(x, weight, eps, scale=scale) + + assert actual.shape == x.shape + assert actual.dtype is torch.float32 + torch.testing.assert_close( + actual, + reference, + rtol=2e-2 if dtype is torch.bfloat16 else 1e-5, + atol=8e-2 if dtype is torch.bfloat16 else 1e-5, + ) diff --git a/pyproject.toml b/pyproject.toml index 250048b84..49b04c628 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,7 @@ extraPaths = [ "integrations/fastvideo_causal_wan22", "integrations/hy_worldplay", "integrations/lingbot", + "integrations/sana", "integrations/self_forcing", "integrations/wan21", "integrations/wan22", @@ -74,6 +75,7 @@ extra-paths = [ "integrations/fastvideo_causal_wan22", "integrations/hy_worldplay", "integrations/lingbot", + "integrations/sana", "integrations/self_forcing", "integrations/wan21", "integrations/wan22", diff --git a/uv.lock b/uv.lock index 023ec81e9..9b271241a 100644 --- a/uv.lock +++ b/uv.lock @@ -27,6 +27,7 @@ members = [ "flashdreams-hy-worldplay", "flashdreams-lingbot", "flashdreams-omnidreams", + "flashdreams-sana-wm", "flashdreams-self-forcing", "flashdreams-wan21", "flashdreams-wan22", @@ -76,6 +77,27 @@ test = [ { name = "tomli", specifier = ">=2.0" }, ] +[[package]] +name = "accelerate" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "packaging" }, + { name = "psutil" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(sys_platform != 'win32' and extra == 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torch", version = "2.12.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'win32' and extra == 'extra-11-flashdreams-dev') or (sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torch", version = "2.12.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8d/75/94cd5d389649578aca399e5aa822637eec18319a1dadc400ffe2f9a7493f/accelerate-1.14.0.tar.gz", hash = "sha256:41b9c4377a54e0b460a959b0defa1b736e4ca0a2373252d9a539964c2afe3c8d", size = 412167, upload-time = "2026-06-11T13:45:52.326Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/db/253133d7e7cb40d3af384bb2f5c0b4a2b7fdcffbc95c688cc67a20a3c103/accelerate-1.14.0-py3-none-any.whl", hash = "sha256:e94390c2863b873be18f623f9df48a0d8fe5eff13ea7f1a00092b0a7904888c6", size = 389246, upload-time = "2026-06-11T13:45:50.477Z" }, +] + [[package]] name = "accessible-pygments" version = "0.0.5" @@ -1569,6 +1591,44 @@ requires-dist = [ ] provides-extras = ["interactive-drive", "dev"] +[[package]] +name = "flashdreams-sana-wm" +version = "0.1.0" +source = { editable = "integrations/sana" } +dependencies = [ + { name = "accelerate" }, + { name = "diffusers" }, + { name = "flashdreams" }, + { name = "imageio", extra = ["ffmpeg"] }, + { name = "pillow" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torchvision", version = "0.26.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(sys_platform != 'win32' and extra == 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torchvision", version = "0.27.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'win32' and extra == 'extra-11-flashdreams-dev') or (sys_platform != 'win32' and extra != 'group-11-flashdreams-cuda12') or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "torchvision", version = "0.27.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "sys_platform == 'win32' or (extra == 'extra-11-flashdreams-dev' and extra == 'group-11-flashdreams-cuda12') or (extra == 'group-11-flashdreams-cuda12' and extra == 'group-11-flashdreams-cuda13')" }, + { name = "transformers" }, +] + +[package.optional-dependencies] +dev = [ + { name = "pytest" }, +] + +[package.metadata] +requires-dist = [ + { name = "accelerate", specifier = ">=1.0" }, + { name = "diffusers", specifier = ">=0.36" }, + { name = "flashdreams", editable = "flashdreams" }, + { name = "imageio", extras = ["ffmpeg"], specifier = ">=2.31" }, + { name = "pillow", specifier = ">=10" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, + { name = "pyyaml", specifier = ">=6.0" }, + { name = "safetensors", specifier = ">=0.5" }, + { name = "torchvision", specifier = ">=0.26" }, + { name = "transformers", specifier = ">=5.0,<6" }, +] +provides-extras = ["dev"] + [[package]] name = "flashdreams-self-forcing" version = "0.1.0"