…xport

transformer_sana_wm.py:
* License header switched to the "HuggingFace Team and SANA-WM Authors"
  style used by merged sana_video.
* Imports rewritten in stdlib -> third-party -> diffusers order; use
  diffusers `from ...utils import logging` instead of stdlib `logging`.
* Fix 9 `Optional[X]` annotations written as `X or None` (Python's `or`
  short-circuits and silently returns `X`).
* Fix two `assert (cond, msg)` tuple-asserts in PatchEmbedMS3D.forward
  that always pass (SyntaxWarning at import time).
* Remove duplicate `__all__` declarations (the second silently overwrote
  the first).
* Remove dead `reset_bn` (imports a nonexistent `packages.apps.utils`,
  would crash on call).
* Remove the duplicate `logger = logging.getLogger(__name__)` further
  down in the file.

transformer_sana_wm_kernels.py:
* License header normalized; collapse three duplicate triton/torch import
  blocks into one.

pipeline_sana_wm.py:
* License header normalized.
* `_decode_latents` now returns `(T, H, W, 3)` float in [0, 1], matching
  the diffusers convention used by `VideoProcessor`. Returning uint8
  silently broke `export_to_video`: it does `frame * 255` assuming float
  input, so uint8 overflows to `(-x) mod 256` and inverts colors.
* `__call__` converts to PIL/uint8 only when `output_type="pil"`.
* Intrinsics argument now accepts (4,), (F, 4), (3, 3), and (F, 3, 3)
  forms (auto-extracts fx, fy, cx, cy from a 3x3 K) and auto-trims to
  `num_frames` when a longer-than-needed trajectory is passed.
* Inline `retrieve_timesteps` with the standard `# Copied from
  diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps`
  marker, matching merged sana_video.
* Docstrings + EXAMPLE_DOC_STRING updated to reflect the new return type.

pipeline_output.py:
* Update `frames` field docstring to describe the new float [0, 1] return.

refiner.py, cam_utils.py, scripts/sana_wm/convert_sana_wm_to_diffusers.py:
* License headers normalized.

Docs:
* New `docs/source/en/api/pipelines/sana_wm.md` and
  `docs/source/en/api/models/sana_wm_transformer3d.md`, modeled on
  sana_video.md / sana_video_transformer3d.md, wired into
  `docs/source/en/_toctree.yml` under Models and Pipelines.

5s end-to-end smoke test (81 frames @ 16fps, 30 stage-1 steps + 3-step
LTX-2 refiner) passes on 1x H100 80GB with `enable_model_cpu_offload`.
Round-trip diff vs raw float frames is 2.06/255 mean (h264 lossy noise),
confirming the export_to_video fix.

…+ KV cache hooks)

The first cleanup pass only kept the legacy single-shot refiner path. That
path is what the model was *not* trained on — its docstring even says
"feeding the full sequence at once is out-of-distribution" — and its cost
is O(T^2) attention over the full latent volume, which made longer videos
unusable (~21 min per refiner step at 321 frames on an H100).

Port the chunk-causal AR mode from the upstream reference so the refiner
matches the training contract:

* `refine_latents` now defaults to `block_size=3, kv_max_frames=11`
  (the canonical AR recipe). Pass `block_size=None` to fall back to the
  legacy single-shot path.
* New `_refine_latents_ar` + `_RefinerChunkRunner` orchestrate the sliding
  window: pre-capture pre-RoPE sink K/V on `z_sana[:source_sink_frames]`
  at sigma=0, then for each `block_size`-frame chunk run a 3-step Euler
  with prefix `{sink_k_pre, sink_v, sink_pe, history_k, history_v}` and
  capture post-RoPE K/V to feed the next window. History is bounded to
  `kv_max_frames - source_sink_frames` so per-block compute is constant.
* New `_predict_x0_active_block` runs the transformer on the active block
  only (Q from active, K/V from prefix+active).
* New `_capture_block_kv` runs sigma=0 forward with a pre_rope/post_rope
  capture flag set on each `attn1`.
* New `_forward_video_only_with_rope` takes a pre-built RoPE so each block
  can use absolute frame positions in the source video.
* `_streaming_self_attention` extended with the `_kv_cache_capture`,
  `_tf_capture_kv`, `_tf_kv_prefix` hook contract that AR mode uses to
  inject and capture K/V on each block.
* New helpers: `_build_rotary_emb_for_absolute_positions`,
  `_set_kv_prefix_on_blocks`, `_clear_kv_prefix_on_blocks`,
  `_set_capture_flag_on_blocks`, `_collect_captured_kv_from_blocks`.
* `_encode_prompt` now also moves the Gemma-3 text encoder back to CPU
  after producing the embeds — otherwise it stays resident through the
  entire AR loop and gates how much GPU memory the refiner transformer
  has left.

Module-level docstring updated to document both modes; existing
single-shot path preserved verbatim.

…eemption)

The AR refiner is expensive (~3-5 min per block) and the refinement loop
ran end-to-end has no in-progress state to recover, so a SLURM preemption
mid-refinement loses all progress. With the canonical
``block_size=3, kv_max_frames=11`` setup, refining a 50s video is 34
blocks of work that has to make it through without preemption on a
backfill queue.

Add per-block atomic checkpointing:

* ``SanaWMLTX2Refiner.refine_latents(checkpoint_dir=Path)`` and
  ``_refine_latents_ar`` accept a directory. After each completed AR
  block, the AR loop writes ``checkpoint_dir/state.pt`` atomically
  (tmp + os.replace).
* The payload is ``{block_idx_done, n_blocks, sink_size, block_size,
  output_shape, output, runner_state}``. ``runner_state`` is a CPU snapshot
  of the runner's ``_sink_kv_pre``, ``_history_kv_post``,
  ``_history_frames`` and ``torch.Generator`` state.
* On entry, if ``state.pt`` exists with a compatible shape signature, the
  AR loop loads the persisted output tensor + runner state and resumes
  from ``block_idx_done + 1`` instead of recomputing from scratch.
* ``SanaWMPipeline.__call__(refiner_checkpoint_dir=...)`` plumbs the
  directory through to the refiner.

Checkpoint size: ~output_volume + sink_KV (~360MB for 50 layers) +
rolling history KV (~3-4GB at full capacity) — saved once per block,
total per-block save overhead ~10s on lustre.