# SANA-WM

SANA-WM 是一个基于 SANA 构建的、由相机控制的条件图像到视频世界模型。给定首帧图像、文本提示以及相机轨迹（可以是显式的 `c2w` 位姿，也可以是 WASD/IJKL 动作字符串），它可以生成一段运动遵循指定相机路径的视频。

推理分为两个阶段运行：

1. **阶段 1 — SANA-WM DiT。** 这是一个具有 16 亿参数的双向 DiT，采用 GDN-Triton 线性注意力机制和 UCPE 相机控制分支。采样过程使用 LTX 风格的流匹配欧拉调度器，并为每个 token 设置独立的时间步；第一个潜在帧作为条件锚点。
2. **阶段 2 — LTX-2 精修器（可选）。** 一个独立的 sink-bidirectional Euler 精修器流水线（[SanaWMLTX2Refiner](/docs/diffusers/v0.43.0/zh/api/pipelines/sana_wm#diffusers.SanaWMLTX2Refiner)），封装了 [SanaWMLTX2RefinerTransformer3DModel](/docs/diffusers/v0.43.0/zh/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel) + `LTX2TextConnectors` 以及一个 Gemma-3 文本编码器，运行 3 个蒸馏 sigma 步骤。

两个阶段都通过 [AutoencoderKLLTX2Video](/docs/diffusers/v0.43.0/zh/api/models/autoencoderkl_ltx_2#diffusers.AutoencoderKLLTX2Video) VAE 进行解码。

可用模型：

| 模型 | 推荐的 dtype |
|:-----:|:-----------------:|
| [`Efficient-Large-Model/SANA-WM_bidirectional-diffusers`](https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional-diffusers) | `torch.bfloat16` |
| [`Efficient-Large-Model/SANA-WM_bidirectional-diffusers-refiner`](https://huggingface.co/Efficient-Large-Model/SANA-WM_bidirectional-diffusers-refiner) | `torch.bfloat16` |

> [!TIP]
> SANA-WM 在固定的 704×1280 分辨率下训练。建议的 dtype 用于 transformer 权重——将文本编码器保持在 `torch.bfloat16`，将 VAE 保持在 `torch.float32`，以获得最佳数值表现。该 pipeline 期望相机内参 `[fx, fy, cx, cy]` 以 *原始图像* 的像素坐标表示；内部会应用缩放并居中裁剪的变换。

## 推理

```python
import torch
from PIL import Image

from diffusers import SanaWMPipeline
from diffusers.utils import export_to_video

pipe = SanaWMPipeline.from_pretrained(
    "Efficient-Large-Model/SANA-WM_bidirectional-diffusers",
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()  # ~45 GB of weights — offload between stages

# SANA-WM was trained on the LTX-2 VAE in framewise mode with tiling enabled. Without these
# settings the VAE encodes the whole (B, C, T, H, W) clip in one shot, which gives subtly
# different numerics from the released checkpoint.
pipe.vae.enable_tiling()
pipe.vae.use_framewise_encoding = True
pipe.vae.use_framewise_decoding = True
pipe.vae.tile_sample_stride_num_frames = 64
pipe.vae.tile_sample_min_num_frames = 96

prompt = "A car driving across a vast desert plain at golden hour."
output = pipe(
    image=Image.open("input.png").convert("RGB"),
    prompt=prompt,
    action="w-80,jw-40,w-40",        # WASD-style action DSL: forward 80f, jump+forward 40f, forward 40f
    intrinsics=[800.0, 800.0, 845.0, 464.0],  # fx, fy, cx, cy in original-image pixels
    num_frames=161,
    num_inference_steps=60,
    guidance_scale=5.0,
    generator=torch.Generator(device="cuda").manual_seed(42),
    output_type="latent",            # hand the latents to the refiner below
)
```

传入 `action=None` 并提供您自己的 `c2w` 位姿（`(F, 4, 4)` numpy 数组），以显式控制相机轨迹。去掉 `output_type="latent"` 即可直接从第一阶段获取视频，并跳过精调器。

### 阶段 2 — LTX-2 精调器

[SanaWMLTX2Refiner](/docs/diffusers/v0.43.0/zh/api/pipelines/sana_wm#diffusers.SanaWMLTX2Refiner) 以独立仓库的形式发布，类似于 SDXL 将 base 和 refiner 分开的方式。请传入 base pipeline 的 VAE，以便共享权重，而不是加载两次：

```python
from diffusers import SanaWMLTX2Refiner

refiner = SanaWMLTX2Refiner.from_pretrained(
    "Efficient-Large-Model/SANA-WM_bidirectional-diffusers-refiner",
    vae=pipe.vae,
    torch_dtype=torch.bfloat16,
)
refiner.enable_model_cpu_offload()

frames = refiner(output.latent, prompt, fps=16)
export_to_video(list(frames), "sana_wm.mp4", fps=16)
```

如果没有 `vae`，refiner 将返回精炼后的潜变量而非视频，这在您希望自行进行解码时非常有用。

> [!TIP]
> 在**两个**流水线中启用 offloading，或者在 stage 2 之前释放 stage 1（`pipe.transformer.to("cpu")`）。这两个 stage 在 `torch.bfloat16` 下总共约 45 GB，若将两者同时驻留在单张 80 GB 显卡上，留给激活值的空间就太少了。另外请注意，stage 2 会遵循 `torch_dtype`：以 `torch.float32` 加载 refiner 大约会使内存占用翻倍，并略微改变输出结果。

如果您没有相机内参，托管的 [模块化模块](../../modular_diffusers/overview) 可以从单帧图像中对其进行估计。它位于 `diffusers` 之外，因为它会引入 Pi3X —— 这是一个额外的依赖项和第二个检查点 —— 因此，除非您主动请求，否则不会下载任何内容：

```python
from diffusers import ModularPipeline

# One-time per image. Requires `pip install pi3-vision`.
estimator = ModularPipeline.from_pretrained(
    "Efficient-Large-Model/pi3x-intrinsics-estimator", trust_remote_code=True
)
estimator.load_components(dtype=torch.bfloat16)
intrinsics = estimator(image=Image.open("input.png").convert("RGB"), output="intrinsics")
```

## 转换已发布的检查点

如果您拥有 SANA-WM 的源代码发布版本（而非预先转换的 diffusers 快照），请运行一次转换脚本：

```bash
python scripts/convert_sana_wm_to_diffusers.py \
    --src Efficient-Large-Model/SANA-WM_bidirectional \
    --dst ./SANA-WM_bidirectional-diffusers
```

这会写入两个目录：基础流水线位于 `--dst`，第二阶段精化器与其并列位于 `./SANA-WM_bidirectional-diffusers-refiner`（可通过 `--dst-refiner` 覆盖）。然后，像往常一样从各自的本地路径加载它们。

## 组件

- `tokenizer` — `GemmaTokenizerFast`
- `text_encoder` — Gemma-2（返回解码器隐藏状态）
- `vae` — [AutoencoderKLLTX2Video](/docs/diffusers/v0.43.0/zh/api/models/autoencoderkl_ltx_2#diffusers.AutoencoderKLLTX2Video)（LTX-2，空间 ×32 / 时间 ×8）
- `transformer` — [SanaWMTransformer3DModel](/docs/diffusers/v0.43.0/zh/api/models/sana_wm_transformer3d#diffusers.SanaWMTransformer3DModel)，16 亿参数的双向 DiT
- `scheduler` — [FlowMatchEulerDiscreteScheduler](/docs/diffusers/v0.43.0/zh/api/schedulers/flow_match_euler_discrete#diffusers.FlowMatchEulerDiscreteScheduler)

stage-2 refiner 是一个独立的仓库，拥有自己的 `transformer`（[SanaWMLTX2RefinerTransformer3DModel](/docs/diffusers/v0.43.0/zh/api/models/sana_wm_refiner_transformer3d#diffusers.SanaWMLTX2RefinerTransformer3DModel)）、`connectors`（`LTX2TextConnectors`）、`tokenizer`、Gemma-3 `text_encoder` 和 `scheduler`。它本身没有 `vae`——请传入基础 pipeline 的。

## SanaWMPipeline[[diffusers.SanaWMPipeline]]

#### diffusers.SanaWMPipeline[[diffusers.SanaWMPipeline]]

```python
diffusers.SanaWMPipeline(tokenizer: GemmaTokenizer | GemmaTokenizerFast, text_encoder: Gemma2PreTrainedModel, vae: AutoencoderKLLTX2Video, transformer: SanaWMTransformer3DModel, scheduler: FlowMatchEulerDiscreteScheduler)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py#L143)

**参数：**

tokenizer （*GemmaTokenizer* 或 *GemmaTokenizerFast*）：Gemma-2 分词器。

text_encoder （*Gemma2PreTrainedModel*）：Gemma-2 文本编码器。

vae （*AutoencoderKLLTX2Video*）：LTX-2 VAE。

transformer ([*SanaWMTransformer3DModel*])：1600M 双向 SANA-WM DiT。

scheduler ([*FlowMatchEulerDiscreteScheduler*])：Flow-matching Euler 调度器（LTX 风格的逐 token 时间步）。直接处理 latents。

SANA-WM 相机控制图像转视频流水线。

根据首帧图像、文本提示词以及相机轨迹（显式的 `c2w` 位姿或 WASD/IJKL 操作字符串）生成视频。第一阶段采样使用 1600M 双向 SANA DiT，第二阶段精修使用 LTX-2 sink-bidirectional Euler refiner；两者均通过 LTX-2 VAE 进行解码。

#### __调用__[[diffusers.SanaWMPipeline.__call__]]

```python
__call__(image: PIL.Image.Image | str | Path, prompt: str, c2w: np.ndarray | None = None, action: str | None = None, intrinsics: np.ndarray | list[float] | None = None, height: int = 704, width: int = 1280, num_frames: int = 161, fps: int = 16, num_inference_steps: int = 60, guidance_scale: float = 5.0, negative_prompt: str = '', generator: torch.Generator | list[torch.Generator] | None = None, prompt_embeds: torch.Tensor | None = None, prompt_attention_mask: torch.Tensor | None = None, negative_prompt_embeds: torch.Tensor | None = None, negative_prompt_attention_mask: torch.Tensor | None = None, max_sequence_length: int = 300, chi_prompt: list[str] | None = None, output_type: Literal['np', 'pil', 'pt', 'latent'] = 'np', callback_on_step_end: Callable[[DiffusionPipeline, int, int, dict], dict] | None = None, callback_on_step_end_tensor_inputs: list[str] = ['latents'], return_dict: bool = True)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py#L457)

**参数：**

image (*PIL.Image.Image* 或 *str*)：首帧图像（PIL 或路径）。

prompt (*str*)：文本提示。

c2w (*np.ndarray*，*可选*)：`(F, 4, 4)` 相机到世界坐标系的位姿。与 *action* 互斥。

action (*str*, *optional*) : Action-DSL 字符串，例如 `"w-80,jw-40,w-40"`。与 *c2w* 互斥。

intrinsics (*np.ndarray* 或 *list[float]*) : 以**原始图像**像素坐标表示的 `[fx, fy, cx, cy]`。该流水线会在内部应用缩放+裁剪变换。

height (*int*，默认值为 704)：输出帧高度（公开模型中固定）。

width (*int*，默认为 1280)：输出帧宽度（公开模型中固定）。

num_frames (*int*，默认为 161)：目标帧数；会取整到 `8k+1`（LTX-2 VAE 约束）。

fps (*int*，默认为 16)：输出帧率（同时传递给 refiner）。

num_inference_steps (*int*，默认为 60)：第一阶段 DiT 采样步数。

guidance_scale (*float*，默认值为 5.0)：无分类器引导（Classifier-free guidance）的缩放系数。

negative_prompt (*str*，默认为 "") : 可选的负面提示词。

generator (*torch.Generator* 或 *list[torch.Generator]*，*可选*)：一个或多个 torch 生成器，用于使噪声采样具有确定性。如果为 *None*，则采样是非确定性的。

prompt_embeds (*torch.Tensor*, *可选*) : 预计算好的文本嵌入，用于跳过文本编码器。必须与 *prompt_attention_mask* 一起传入。

prompt_attention_mask (*torch.Tensor*, *可选*) : *prompt_embeds* 的注意力掩码。

negative_prompt_embeds (*torch.Tensor*, *可选*) : 预计算的负文本嵌入。必须与 *negative_prompt_attention_mask* 一起传入。

negative_prompt_attention_mask (*torch.Tensor*, *可选*) : *negative_prompt_embeds* 的注意力掩码。

max_sequence_length (*int*，默认为 300)：最大提示词 token 数。

chi_prompt (*list[str]*, *optional*) : 覆盖 chi-prompt 前缀（默认值与公开发布版本一致）。

output_type （*"np"*、*"pil"*、*"pt"* 或 *"latent"*，默认为 *"np"*）：输出格式。

callback_on_step_end (*Callable*，*可选*)：在每个去噪步骤结束时作为 *callback_on_step_end(self, step, timestep, callback_kwargs)* 调用，其中 *callback_kwargs* 包含 *callback_on_step_end_tensor_inputs* 中指定的张量。返回字典中的张量将替换流水线中的对应张量。

callback_on_step_end_tensor_inputs (*list[str]*，默认为 *["latents"]*)：要传递给 *callback_on_step_end* 的张量；必须是 *._callback_tensor_inputs* 的子集。

return_dict (*bool*，默认为 True)：返回 [*SanaWMPipelineOutput*] 还是元组。

**返回值：**

[*SanaWMPipelineOutput*] 包含 *.frames*，其形状为 `(T, H, W, 3)`。当 *output_type="np"* 时，为取值范围在 `[0, 1]` 的浮点型 `np.ndarray`；当 *"pil"* 时，为长度 `T` 的 `PIL.Image.Image` 列表；当 *"pt"* 时，为形状 `(T, 3, H, W)` 的张量；当 *"latent"* 时，为原始潜在张量。

生成 SANA-WM 相机控制视频。

示例：
```py
>>> import torch
>>> from PIL import Image
>>> from diffusers import SanaWMPipeline

>>> pipe = SanaWMPipeline.from_pretrained(
...     "Efficient-Large-Model/SANA-WM_bidirectional-diffusers", torch_dtype=torch.bfloat16
... ).to("cuda")

>>> output = pipe(
...     image=Image.open("input.png").convert("RGB"),
...     prompt="A car driving across a vast desert plain at golden hour.",
...     action="w-80,jw-40,w-40",
...     intrinsics=[800.0, 800.0, 845.0, 464.0],  # fx, fy, cx, cy in original-image pixels
...     num_inference_steps=60,
... )
>>> # output.frames is (T, H, W, 3) float np.ndarray in [0, 1] (diffusers convention).
```

#### check_inputs[[diffusers.SanaWMPipeline.check_inputs]]

```python
check_inputs(c2w: np.ndarray | None, action: str | None, intrinsics: np.ndarray | list[float] | None, callback_on_step_end_tensor_inputs: list[str] | None = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py#L335)

验证 `__call__` 的输入。如果输入无效则抛出异常，且不返回任何值。

#### encode_prompt[[diffusers.SanaWMPipeline.encode_prompt]]

```python
encode_prompt(prompt: str, negative_prompt: str = '', device: torch.device | None = None, max_sequence_length: int = 300, chi_prompt: list[str] | None = None, prompt_embeds: torch.Tensor | None = None, prompt_attention_mask: torch.Tensor | None = None, negative_prompt_embeds: torch.Tensor | None = None, negative_prompt_attention_mask: torch.Tensor | None = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py#L204)

**返回值：**

`(cond, cond_mask, neg, neg_mask)`，其中 `cond` 和 `neg` 是形状为 `(1, 1, L, D)` 的 Gemma 隐藏状态，掩码的形状为 `(1, L)`。

通过 Gemma-2 对提示词和负提示词进行编码。

镜像了 SANA 的 chi-prompt-prefix 技巧：将 chi prompt 前置到用户提示词之前，然后通过 `select_index = [0, -L+1, ..., -1]` 切片提取 BOS token 以及最后 `max_sequence_length - 1` 个 token。

#### prepare_camera_trajectory[[diffusers.SanaWMPipeline.prepare_camera_trajectory]]

```python
prepare_camera_trajectory(c2w: np.ndarray | None, action: str | None, intrinsics: np.ndarray | list[float], num_frames: int)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py#L365)

将相机输入归一化为 `(c2w_(F, 4, 4), intrinsics_(F, 4))`。

同时，将 `num_frames` 对齐到 VAE 友好的 `8k+1`，并裁剪两个数组以使其长度一致。裁剪后的图像和重新缩放的相机内参将在确定目标分辨率后提供。

#### prepare_latents[[diffusers.SanaWMPipeline.prepare_latents]]

```python
prepare_latents(first_latent: torch.Tensor, num_frames: int, height: int, width: int, dtype: torch.dtype, device: torch.device, generator: torch.Generator)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/sana_wm/pipeline_sana_wm.py#L403)

采样初始潜变量，并将第一帧固定为条件锚点。

返回 `(latents, condition_mask)`，其中 `condition_mask` 在第一帧的 token 位置为 1（这些 token 在整个采样过程中保持干净），其余位置为 0。

## SanaWMLTX2Refiner[[diffusers.SanaWMLTX2Refiner]]

LTX-2 的第二阶段精调器是一个独立的 [DiffusionPipeline](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline)，它接收第一阶段的潜变量。为其提供一个 `vae`（即基础流水线中的 VAE，以便共享权重），它便会将潜变量解码为视频；若不提供，则返回精调后的潜变量。

#### diffusers.SanaWMLTX2Refiner[[diffusers.SanaWMLTX2Refiner]]

```python
diffusers.SanaWMLTX2Refiner(transformer: SanaWMLTX2RefinerTransformer3DModel, connectors: LTX2TextConnectors, tokenizer: GemmaTokenizer | GemmaTokenizerFast, text_encoder: Gemma3ForConditionalGeneration, scheduler: FlowMatchEulerDiscreteScheduler, vae: AutoencoderKLLTX2Video | None = None, text_max_sequence_length: int = 1024)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/sana_wm/refiner.py#L47)

**参数：**

transformer ([*SanaWMLTX2RefinerTransformer3DModel*])：采用分块因果滑动窗口 KV 缓存的 LTX-2 视频 DiT。

connectors ([*LTX2TextConnectors*])：LTX-2 文本连接器。

tokenizer : Gemma-3 分词器。

text_encoder : Gemma-3 文本编码器。

scheduler （*FlowMatchEulerDiscreteScheduler*）：Flow-matching Euler 调度器。以 `shift=1.0` 进行构造，因此蒸馏后的 sigmas 会原样通过。

vae ([*AutoencoderKLLTX2Video*]，*可选*)：与 [*SanaWMPipeline*] 使用的 VAE 相同；传入 *vae=pipe.vae* 以共享权重。如果提供了该参数，refiner 将解码为视频，否则返回精炼后的潜变量。

text_max_sequence_length (*int*，默认为 1024)：传递给 Gemma-3 分词器的最大 token 数。

LTX-2 sink-bidirectional Euler refiner — SANA-WM 第二阶段，作为独立流水线。

封装了 LTX-2 组件（refiner transformer + 文本连接器 + Gemma-3 文本编码器 + tokenizer），以及一个携带蒸馏 sigma 调度并执行 Euler 步进的 [*FlowMatchEulerDiscreteScheduler*]。它被注册为 [*SanaWMPipeline*] 的可选组件，也可以单独用于精化第一阶段潜变量。

#### __调用__[[diffusers.SanaWMLTX2Refiner.__call__]]

```python
__call__(sana_latent: torch.Tensor, prompt: str, fps: float, sink_size: int = 1, generator: torch.Generator | None = None, block_size: int = 3, kv_max_frames: int = 11, sigmas: tuple[float, ...] = (0.909375, 0.725, 0.421875, 0.0), output_type: str = 'np')
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/sana_wm/refiner.py#L107)

**参数：**

sana_latent : `(B, C, F, H, W)` 第一阶段潜在表示。

prompt : 文本提示词。

fps : 视频帧率（用于驱动 LTX-2 RoPE 的时间缩放）。

sink_size : 将多少个前导原始 `z_sana` 帧固定为注意力锚点（标准值：1）。

generator : 用于 FM 端点噪声的 torch.Generator。默认使用以 42 为种子的生成器，以便开箱即用地获得可复现的结果。

block_size : 每个 AR 块中的潜在帧数（标准值：3）。

kv_max_frames : 在滑动窗口中保留的最大上下文+活动帧数（标准值：11 = 1 个 sink + 10 个最近帧）。

sigmas : 以 0.0 结束的降序 Euler 调度（标准 3 步蒸馏：`(0.909375, 0.725, 0.421875, 0.0)`）。传入 `self.scheduler`（去掉末尾的 0.0，调度器会自行添加）。

output_type : *"latent"* 返回精炼后的潜在表示。其他任何值都会通过 *self.vae* 进行解码，并后处理为相应类型（*"np"*、*"pt"*、*"pil"*）；如果没有 *vae*，则无论指定什么类型，都会直接返回潜在表示。

**返回：** `*torch.Tensor*`

形状为 `(B, C, F, H, W)` 的细化 VAE 潜变量 — 前 `sink_size` 帧保持原始 stage-1 sink 潜变量不变，其余帧为细化后的输出。

运行 LTX-2 refiner 并返回精炼后的 VAE 潜变量。

使用模型训练时所采用的块因果 AR 配方（`block_size=3`、`kv_max_frames=11`）：将 `[source_sink + recent_history + active_block]` 的 K/V 滑动窗口按块逐次送入 Transformer，因此每个块的计算量是有界的，且总精化成本随视频长度线性增长。

## SanaWMPipelineOutput[[diffusers.SanaWMPipelineOutput]]

#### diffusers.SanaWMPipelineOutput[[diffusers.SanaWMPipelineOutput]]

```python
diffusers.SanaWMPipelineOutput(frames: typing.Union[torch.Tensor, numpy.ndarray, list[list[PIL.Image.Image]]], c2w: numpy.ndarray | None = None, latent: typing.Optional[torch.Tensor] = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/sana_wm/pipeline_output.py#L25)

**参数：**

frames (*torch.Tensor*、*np.ndarray* 或 *list[PIL.Image.Image]*)：生成的视频：当 `output_type="np"` 时，为形状 `(T, H, W, 3)` 的浮点型 `np.ndarray`，取值范围在 `[0, 1]`；当 `"pt"` 时，为形状 `(T, 3, H, W)` 的张量；当 `"pil"` 时，为包含 `T` 个 ``PIL.Image`` 的列表；当 `"latent"` 时，为潜变量张量 `(B, C, T_lat, H_lat, W_lat)`。

c2w (*np.ndarray*) : 相机到世界的位姿 `(T, 4, 4)`，与 `frames` 对齐（refiner 会丢弃 sink 锚点帧；如果运行了 refiner，该数组会相应地重新对齐）。

latent (*torch.Tensor*，可选)：LTX-2 VAE 空间中的潜在张量，形状为 `(B, C, T_lat, H_lat, W_lat)`。当 `output_type="latent"` 时返回。

SANA-WM 图像到视频流水线的输出类。