SANA-WM

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

推理分为两个阶段运行：

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

两个阶段都通过 AutoencoderKLLTX2Video VAE 进行解码。

可用模型：

模型	推荐的 dtype
`Efficient-Large-Model/SANA-WM_bidirectional-diffusers`	`torch.bfloat16`
`Efficient-Large-Model/SANA-WM_bidirectional-diffusers-refiner`	`torch.bfloat16`

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

推理

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 以独立仓库的形式发布，类似于 SDXL 将 base 和 refiner 分开。请传入 base pipeline 的 VAE，以便共享权重，而不是加载两次：

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 将返回精炼后的潜变量而非视频，这在您希望自行进行解码时非常有用。

在两个流水线中均启用 offloading，或者在 stage 2 之前释放 stage 1（pipe.transformer.to("cpu")）。这两个阶段在 torch.bfloat16 下合计约 45 GB，若将两者同时驻留在单张 80 GB 显卡上，留给激活值的空间将过于有限。此外请注意，stage 2 会遵循 torch_dtype：以 torch.float32 加载 refiner 大约会使内存占用翻倍，并导致输出略有变化。

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

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 快照），请运行一次转换脚本：

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（LTX-2，空间 ×32 / 时间 ×8）
transformer — SanaWMTransformer3DModel，16 亿参数的双向 DiT
scheduler — FlowMatchEulerDiscreteScheduler

Stage-2 refiner 是一个独立的仓库，拥有自己的 transformer（SanaWMLTX2RefinerTransformer3DModel）、connectors（LTX2TextConnectors）、tokenizer、Gemma-3 text_encoder 和 scheduler。它本身没有 vae——请传入基础 pipeline 的。

SanaWMPipeline

class diffusers.SanaWMPipeline

< source >

( tokenizer: GemmaTokenizer | GemmaTokenizerFasttext_encoder: Gemma2PreTrainedModelvae: AutoencoderKLLTX2Videotransformer: SanaWMTransformer3DModelscheduler: FlowMatchEulerDiscreteScheduler )

参数

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 进行解码。

call

< source >

( image: PIL.Image.Image | str | Pathprompt: strc2w: np.ndarray | None = Noneaction: str | None = Noneintrinsics: np.ndarray | list[float] | None = Noneheight: int = 704width: int = 1280num_frames: int = 161fps: int = 16num_inference_steps: int = 60guidance_scale: float = 5.0negative_prompt: str = ''generator: torch.Generator | list[torch.Generator] | None = Noneprompt_embeds: torch.Tensor | None = Noneprompt_attention_mask: torch.Tensor | None = Nonenegative_prompt_embeds: torch.Tensor | None = Nonenegative_prompt_attention_mask: torch.Tensor | None = Nonemax_sequence_length: int = 300chi_prompt: list[str] | None = Noneoutput_type: Literal['np', 'pil', 'pt', 'latent'] = 'np'callback_on_step_end: Callable[[DiffusionPipeline, int, int, dict], dict] | None = Nonecallback_on_step_end_tensor_inputs: list[str] = ['latents']return_dict: bool = True )

参数

image (PIL.Image.Image 或 str) — 首帧图像（PIL 或路径）。
prompt (str) — 文本提示。
c2w (np.ndarray, 可选) — (F, 4, 4) 相机到世界位姿。与 action 互斥。
action (str, 可选) — 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) — 无分类器引导的缩放系数。
negative_prompt (str，默认为 "") — 可选的负面提示词。
generator (torch.Generator 或 list[torch.Generator]，可选) — 一个或多个 torch 生成器，用于使噪声采样具有确定性。如果为 None，则采样是非确定性的。
prompt_embeds (torch.Tensor, optional) — 预计算好的文本嵌入，用于跳过文本编码器。必须与 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] 还是元组。

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

示例：

>>> 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

< source >

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

encode_prompt

< source >

( prompt: strnegative_prompt: str = ''device: torch.device | None = Nonemax_sequence_length: int = 300chi_prompt: list[str] | None = Noneprompt_embeds: torch.Tensor | None = Noneprompt_attention_mask: torch.Tensor | None = Nonenegative_prompt_embeds: torch.Tensor | None = Nonenegative_prompt_attention_mask: torch.Tensor | None = None )

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

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

prepare_camera_trajectory

< source >

( c2w: np.ndarray | Noneaction: str | Noneintrinsics: np.ndarray | list[float]num_frames: int )

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

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

prepare_latents

< source >

( first_latent: torch.Tensornum_frames: intheight: intwidth: intdtype: torch.dtypedevice: torch.devicegenerator: torch.Generator )

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

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

SanaWMLTX2Refiner

LTX-2 第二阶段精调器是一个独立的 DiffusionPipeline，用于接收第一阶段潜变量。为其提供一个 vae（即基础流水线中的 VAE，以便共享权重）即可将其解码为视频；若不提供，则返回精调后的潜变量。

class diffusers.SanaWMLTX2Refiner

< source >

( transformer: SanaWMLTX2RefinerTransformer3DModelconnectors: LTX2TextConnectorstokenizer: GemmaTokenizer | GemmaTokenizerFasttext_encoder: Gemma3ForConditionalGenerationscheduler: FlowMatchEulerDiscreteSchedulervae: AutoencoderKLLTX2Video | None = Nonetext_max_sequence_length: int = 1024 )

参数

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] 的可选组件，也可以单独用于精化第一阶段潜变量。

call

< source >

( sana_latent: torch.Tensorprompt: strfps: floatsink_size: int = 1generator: torch.Generator | None = Noneblock_size: int = 3kv_max_frames: int = 11sigmas: tuple[float, ...] = (0.909375, 0.725, 0.421875, 0.0)output_type: str = 'np' ) → torch.Tensor

参数

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 帧原样保留第一阶段的 sink 潜变量，其余帧则包含细化后的输出。

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

使用模型训练时所采用的分块因果 AR 配方（block_size=3、kv_max_frames=11）：将 [source_sink + recent_history + active_block] 的 K/V 滑动窗口按块逐次输入 Transformer，从而将每块的计算量控制在有限范围内，并使总精化成本随视频长度线性增长。

SanaWMPipelineOutput

class diffusers.SanaWMPipelineOutput

< source >

( frames: typing.Union[torch.Tensor, numpy.ndarray, list[list[PIL.Image.Image]]]c2w: numpy.ndarray | None = Nonelatent: typing.Optional[torch.Tensor] = None )

参数

frames (torch.Tensor、np.ndarray 或 list[PIL.Image.Image]) — 生成的视频：当 output_type="np" 时，为形状 (T, H, W, 3)、取值范围在 [0, 1] 的浮点型 np.ndarray；当 "pt" 时，为形状 (T, 3, H, W) 的张量；当 "pil" 时，为包含 T 个 `PIL.Image` 的列表；当 "latent" 时，为潜变量张量 (B, C, T_lat, H_lat, W_lat)。
c2w (np.ndarray) — 与 frames 对齐的相机到世界位姿 (T, 4, 4)（refiner 会丢弃 sink 锚点帧；如果运行了 refiner，此数组会相应地重新对齐）。
latent (torch.Tensor, 可选) — LTX-2 VAE 空间中的潜在张量，形状为 (B, C, T_lat, H_lat, W_lat)。当 `output_type=“latent”` 时返回。

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

Update on GitHub

Diffusers

SANA-WM

推理

阶段 2 — LTX-2 精调器

转换已发布的检查点

组件

SanaWMPipeline

class diffusers.SanaWMPipeline

__call__

check_inputs

encode_prompt

prepare_camera_trajectory

prepare_latents

SanaWMLTX2Refiner

class diffusers.SanaWMLTX2Refiner

__call__

SanaWMPipelineOutput

class diffusers.SanaWMPipelineOutput

call

call