Add SANA-WM integration#392
Draft
aidanfnv wants to merge 30 commits into
Draft
Conversation
Add a workspace integration package for the upstream NVlabs/Sana SANA-WM bidirectional world model. The package registers the `sana-wm-bidirectional` runner, resolves upstream Sana assets and Hugging Face checkpoints, loads images, prompts, camera trajectories, and intrinsics, then writes generated videos through the upstream pipeline. Port CPU-safe camera trajectory and intrinsics helpers so the integration can be checked without running GPU inference. Add smoke and camera tests covering runner registration, config derivation, action parsing, and geometry transforms. Signed-off-by: Aidan Foster <aidanf@nvidia.com>
Document that the SANA-WM integration is expected to run inside the upstream Sana environment created by `environment_setup.sh sana`. The setup notes keep Sana's pinned PyTorch and model dependencies intact, then install FlashDreams and the local integration in editable mode without dependency resolution. Include the expected checkout layout and call out Transformer Engine as an optional dependency for quantized inference while BF16 remains the default validated path. Signed-off-by: Aidan Foster <aidanf@nvidia.com>
Update the SANA-WM README examples to pass explicit values for boolean Tyro options such as `--no-refiner True`. This matches the generated CLI contract and avoids the confusing "Missing value" error when users treat those fields as flag-style booleans. Signed-off-by: Aidan Foster <aidanf@nvidia.com>
Record the BF16 commands used to validate both the Stage-1-only and refined SANA-WM paths, including the expected output locations under `outputs/sana_wm*/sana_wm_generated.mp4`. Note the upstream warnings that appeared during successful runs so they are not mistaken for failures: the null-embed load warning, missing `pos_embed`, and the optional Apex RMSNorm fallback. These commands establish BF16 as the known-good baseline for the integration. Signed-off-by: Aidan Foster <aidanf@nvidia.com>
Add first-class `stage1_precision` and `refiner_precision` runner options for BF16, FP8, and FP4. The runner maps those options onto upstream SANA-WM Transformer Engine environment selectors inside a scoped environment so external NVFP4 variables cannot accidentally change the default BF16 path. Validate quantized precision requests before loading checkpoints. FP8 now requires a CUDA device, Hopper-or-newer GPU capability, CUDA 12.9 or newer as reported by PyTorch, and Transformer Engine's Float8BlockScaling recipe. FP4 requires a CUDA device, Blackwell GPU capability, and Transformer Engine's NVFP4 recipe. Extend the SANA README with precision requirements and an FP4 example command, and add CPU-safe tests for the env mapping, cleanup behavior, and early capability errors. Signed-off-by: Aidan Foster <aidanf@nvidia.com>
Route the public SANA-WM runner through FlashDreams pipeline, diffusion-model, transformer, and scheduler boundaries instead of the upstream inference harness. Keep the upstream harness only as an explicit reference runner with separate documentation. Replace the vendored upstream diffusion package with a FlashDreams-owned Stage-1 DiT module, direct checkpoint loading, local YAML config parsing, direct diffusers LTX2 VAE encode/decode, and native runner wiring for Stage-1 BF16 execution. Add Transformer Engine-free native FP8 and FP4 Stage-1 Linear replacements using PyTorch scaled-mm and Triton NVFP4 activation quantization, with precision/backend validation before checkpoint loading. Keep the native scope honest: require no-refiner for native runs, keep action overlays and upstream TE quantization on the separate upstream-reference harness, and document the remaining refiner and GDN parity gaps separately from the native runner docs. Expand CPU-safe SANA tests to pin native/reference runner separation, config loading without upstream types, native Stage-1 checkpoint schema, Stage-1 forward shape, runtime release, VAE tiling safeguards, and native quant replacement behavior. Signed-off-by: Aidan Foster <aidanf@nvidia.com>
Run classifier-free guidance as sequential negative and positive Stage-1 forwards instead of doubling the full model batch. This reduces peak activation memory for long native SANA-WM rollouts. Release large Stage-1 block temporaries earlier and make the VAE decode OOM retry reduce temporal tile size as well as spatial tile size. Add CPU coverage for sequential CFG kwargs and the lower-memory VAE retry configuration. Signed-off-by: Aidan Foster <aidanf@nvidia.com>
Keep retry temporal tiling compatible with LTX2's temporal compression ratio so diffusers does not derive a zero latent stride during OOM retry. Update the SANA smoke coverage for the clamped retry stride and sync the uv lock metadata for the integration dependencies declared by the native SANA package. Signed-off-by: Aidan Foster <aidanf@nvidia.com>
Wrap native SANA conditioning and VAE decode in torch.inference_mode so full-length runs do not retain autograd activations during first-frame encode, prompt encode, or LTX2 VAE decode. Extend the VAE OOM retry smoke test to assert both decode attempts run under inference mode. Signed-off-by: Aidan Foster <aidanf@nvidia.com>
Replace the Stage-1 placeholder attention path with checkpoint-compatible GDN, softmax, camera UCPE, output-gate, and FFN behavior. Add a native LTX-2 refiner that loads diffusers transformer/connectors and Gemma text conditioning directly from the SANA-WM checkpoint, then wire the native runner to refine by default before VAE decode. Keep BF16/FP8/FP4 on FlashDreams/PyTorch native paths, add accelerate for low-memory refiner loading, update docs, and cover the new refiner construction plus latent packing in CPU tests. Validation: uv run --package flashdreams-sana-wm pytest integrations/sana/tests/test_smoke.py integrations/sana/tests/test_camera.py; uv run --package flashdreams-sana-wm flashdreams-run --no-instantiate sana-wm-bidirectional; no-instantiate checks for fp8, fp4, and --no-refiner True.
The native LTX-2 refiner prompt encoder returns inference tensors. Keep the whole refiner denoise path inside torch.inference_mode so diffusers Linear calls do not try to save those tensors for autograd. Add a CPU regression assertion that the refiner prediction callback executes with inference mode enabled. Validation: uv run --package flashdreams-sana-wm pytest integrations/sana/tests/test_smoke.py integrations/sana/tests/test_camera.py.
Signed-off-by: Aidan Foster <aidanf@nvidia.com>
Move SANA-WM prompt, first-frame, camera, decode, and refiner orchestration into explicit pipeline components instead of keeping it inside the transformer or runner. The public runner now drives the normal StreamInferencePipeline contract, with a Sana conditioning encoder, Stage-1 transformer, LTX-style scheduler, and video decoder/refiner boundary. Keep the native Stage-1 DiT implementation but make SanaWMDiffusionModelConfig instantiate the shared FlashDreams DiffusionModel. Add generic infra hooks for transformer-provided initial noise and scheduler access to per-step encoder context so Sana can pin the first-frame latent and run per-token LTX Euler sampling without owning a custom diffusion loop. Replace direct diffusers FlowMatch Euler use with a SanaWMLTXEulerScheduler that matches the LTX per-token behavior and keeps the scheduler-specific timestep handling isolated behind the FlashDreams scheduler interface. Restore batched CFG as the default Stage-1 path and remove manual early temporary deletions added during the earlier OOM investigation. Keep inference-mode execution and default-false offload/tiling fallback paths where they remain useful inference hygiene or explicit user controls. Expand CPU coverage for runner registration, component config routing, prompt/first-frame/camera conditioning shapes, scheduler parity, base DiffusionModel instantiation, CFG batching, VAE decode/refiner boundaries, and Stage-1 checkpoint shape coverage. Signed-off-by: Aidan Foster <aidanf@nvidia.com>
Make the SANA-WM pipeline reuse-friendly: the Stage-1 DiT, LTX-2 refiner, and Gemma text encoder are no longer torn down and rebuilt on every generate(). Previously each rollout unconditionally released Stage-1 before decode, released the refiner after refine, and reloaded the ~20 GB Gemma encoder from disk per prompt, so a reused pipeline re-paid a full checkpoint reload and re-quantization (~10 s bf16, ~17 s fp8) plus an ~8-15 s Gemma reload on every call. - transformer: gate the Stage-1 teardown behind a new offload_stage1 flag (default keeps it resident); still drop per-rollout conditioning. - decoder: drop the unconditional refiner release (refine_latents already releases when offload_refiner is set). - refiner: cache the Gemma tokenizer/encoder in CPU RAM instead of reloading per call, and move it to the GPU only for the encode forward so it never inflates peak memory during denoise/decode. Fix deprecated torch_dtype and drop dead _empty_cuda_cache/gc. - runner: expose --offload-stage1 and route it to the transformer. - Add [timing] instrumentation for the Stage-1 build, refiner build, and Gemma load. Measured on Blackwell: warm-reuse pass drops from ~26 s to ~5 s (fp8); full-res fp8 demo_0 peaks at ~54/98 GB with no OOM. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Make --intrinsics-path optional. When omitted, intrinsics are derived from the first-frame size with a centered principal point and a focal length set by a horizontal field of view (new --intrinsics-hfov-deg, default 90 to match the public demo intrinsics). This lets a run be driven by just an image and a prompt (with --action supplying the camera trajectory). - camera: add default_intrinsics_vec4(), returning [F, 4] intrinsics in source-image pixels so they flow through transform_intrinsics_for_crop identically to a loaded file. - runner: derive intrinsics when intrinsics_path is None, log the derivation, and expose --intrinsics-hfov-deg; explicit files behave as before. - README: document optional intrinsics and a minimal image+prompt run. - tests: cover the helper (centering, FOV math, bad-FOV rejection) and the runner derive path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Move SANA-WM first-frame timestep pinning and zero-flow behavior out of the scheduler context path and into the SANA transformer. The shared diffusion model now calls schedulers through the generic initial_noise/predict_flow/rng interface again, and the SANA LTX Euler scheduler only owns timestep construction and update math. Make the decoder/refiner boundary conform to FlashDreams video decoder contracts by subclassing StreamingVideoDecoder and publishing the SANA 8k+1 temporal sizing. Remove runner, cache, and refiner AR block-size/KV-frame knobs that were accepted but ignored, leaving the supported sink-bidirectional LTX-2 refiner path explicit. Route runner step and flow_shift overrides into the scheduler config instead of through SANA conditioning, and update SANA smoke coverage for scheduler isolation, transformer first-frame pinning, decoder sizing, CFG, and refiner behavior. Tests:\n- uv run --package flashdreams-sana-wm --extra dev pytest integrations/sana/tests Signed-off-by: Aidan Foster <aidanf@nvidia.com>
Make the SANA-WM runner choose the snapped rollout length from --num-frames before preparing camera conditioning. Action-derived trajectories now repeat or truncate to that snapped length, so a short action prompt no longer caps the generated video. Fit explicit --camera-path trajectories to the same frame count instead of treating the file length as an upper bound. This matches the existing intrinsics fitting behavior and prevents a 321-frame pose file from truncating a longer requested rollout. Document the updated behavior and add CPU coverage for action repetition, explicit camera-path fitting, and runner input preparation. Tests: - uv run --package flashdreams-sana-wm --extra dev pytest integrations/sana/tests Signed-off-by: Aidan Foster <aidanf@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds FlashDreams integration for SANA-WM, the 2.6B bidirectional camera-controlled world model from NVlabs/Sana.
integrations/sana/— packageflashdreams-sana-wmsana-wm-bidirectional— Stage-1 DiT + LTX-2 refiner"w-100,dw-60,w-101"); trajectories and intrinsics are fitted to the requested frame count so short inputs never cap video length--intrinsics-pathis omitted--offload-stage1releases it before the refiner pass to reduce peak VRAMDifferences from native SANA (
inference_sana_wm.py)Stage-1 DiT: Native SANA loads via
build_model()from the SANA repo (requires a SANA checkout). This integration has a self-contained reimplementation instage1_model.pythat loads the same public checkpoint with no SANA repo dependency.FP8/FP4 backend: Native SANA quantizes with TransformerEngine (
te.Linear,te.fp8_autocast,NVFP4BlockScaling,Float8BlockScaling). This integration uses PyTorchtorch._scaled_mmdirectly (TorchScaledMMFP8Linear,TorchScaledMMFP4Linear) with no TransformerEngine dependency. Precision is selected by the--stage1-precision/--refiner-precisionconfig fields, not by SANA'sSANA_WM_*environment variables.Intrinsics estimation: When
--intrinsics-pathis omitted, native SANA estimates intrinsics using Pi3X (a neural estimator fromyyfz233/Pi3X). This integration instead derives intrinsics from a fixed horizontal-FoV assumption (--intrinsics-hfov-deg, default 90°; the public demo intrinsics are ≈93°). On images whose true HFoV differs from the assumption, the derived intrinsics will be less accurate than Pi3X.No prepared-module cache: Native SANA can save pre-quantized Stage-1 models to disk (
SANA_WM_PREPARED_MODULE_CACHE) to skip re-quantization on subsequent runs. This integration does not implement that cache; quantization runs from scratch each time.Core changes
DiffusionModel.run()inflashdreams.infra.diffusion.modelis refactored to delegate initial latent sampling to a newTransformer.initial_noise()virtual method (signature:latent_shape,rng,cache,input) instead of callingtorch.randndirectly. The default implementation is identical to the previous behaviour. This allows SANA-WM's I2Vfirst-frame pinning to live in
SanaWMTransformerwithout requiring a customDiffusionModelsubclass.Test plan
uv sync --package flashdreams-sana-wm --extra devuv run --extra dev pytest integrations/sana/tests/test_smoke.py integrations/sana/tests/test_camera.py -v(64 passed, CPU)pytest integrations/sana/tests/test_quant_cuda.py -v(requiressm_90+)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 --num-frames 161 --output-dir outputs/sana_wm_demo_0_bf16on RTX 6000 Pro Blackwell and GB300--stage1-precision fp8 --refiner-precision fp8on RTX 6000 Pro Blackwell and GB300--stage1-precision fp4 --refiner-precision fp4on RTX 6000 Pro Blackwell and GB300