Add SANA-WM integration

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 |

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.
 

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

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 <https://github.com/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/<output-dir>/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 <https://github.com/NVlabs/Sana>`_ 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
+
+   <figure class="benchmark-figure-wrap">
+     <div
+       id="sana-wm-bf16-benchmark-chart"
+       class="benchmark-figure"
+       data-benchmark-md-url="../_static/performance/sana_wm/perf-bf16-0727.md"
+       data-benchmark-series="official:Official Impl:#3b82f6;flashdreams:FlashDreams:#76B900"
+       data-chart-aria-label="SANA-WM BF16 benchmark chart"
+     ></div>
+     <figcaption>
+       <p class="model-footnote">
+         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
+         <a href="https://github.com/NVIDIA/flashdreams/tree/main/integrations/sana/tests/parity_check">this instruction</a>.
+       </p>
+     </figcaption>
+   </figure>
+  <script src="../_static/js/benchmark_chart.js"></script>
+
+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}
+   }

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:

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,