Diffusers documentation

SANA-WM

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

SANA-WM

SANA-WMは、SANAを基盤として構築された、カメラ制御による画像から動画への変換を行うワールドモデルです。最初のフレームの画像、テキストプロンプト、カメラ軌道(明示的なc2wポーズまたはWASD/IJKLアクション文字列)を指定すると、要求されたカメラパスに沿った動きを持つ動画を生成します。

推論は2つの段階で実行されます:

  1. ステージ 1 — SANA-WM DiT。 GDN-Triton 線形アテンションと UCPE カメラ制御ブランチを備えた、1.6B パラメータの双方向 DiT です。サンプリングには、トークンごとのタイムステップを持つ LTX スタイルのフローマッチング・オイラースケジューラを使用します。最初の潜在フレームが条件付けアンカーとなります。
  2. ステージ 2 — LTX-2 リファイナー(オプション)。 SanaWMLTX2RefinerTransformer3DModel + LTX2TextConnectors および Gemma-3 テキストエンコーダーをラップする、独立した sink-bidirectional Euler リファイナーパイプライン(SanaWMLTX2Refiner)であり、蒸留済み sigma ステップ 3 ステップで実行されます。

両方のステージは、AutoencoderKLLTX2Video VAE を通じてデコードされます。

利用可能なモデル:

モデル推奨されるdtype
Efficient-Large-Model/SANA-WM_bidirectional-diffuserstorch.bfloat16
Efficient-Large-Model/SANA-WM_bidirectional-diffusers-refinertorch.bfloat16

SANA-WM は固定解像度 704×1280 で学習されています。transformer の重みには推奨 dtype を使用し、テキストエンコーダーは torch.bfloat16、VAE は torch.float32 に設定することで、最適な数値精度が得られます。パイプラインは、元画像のピクセル座標系におけるカメラ内パラメータ [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" を外すことで、ステージ 1 から直接動画を出力し、リファイナーをスキップできます。

ステージ 2 — LTX-2 リファイナー

SanaWMLTX2Refiner は、SDXL がベースモデルとリファイナーを分離しているのと同様に、独立したリポジトリとして提供されています。重みを2回読み込むのではなく共有できるように、ベースパイプラインの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 は動画ではなく精錬された latents を返します。これは、自分でデコードを行いたい場合に便利です。

両方のパイプラインでオフローディングを有効にするか、ステージ 2 の前にステージ 1 を解放します(pipe.transformer.to("cpu"))。2 つのステージを合計すると torch.bfloat16 で約 45 GB となり、両方を 1 枚の 80 GB カード上に常駐させると、アクティベーションに割り当てるメモリが不足します。また、ステージ 2 は torch_dtype を尊重します。torch.float32 でリファイナを読み込むと、メモリ使用量が約 2 倍になり、出力がわずかに変化します。

カメラの内部パラメータがない場合、ホストされたモジュラーブロックを使用して、単一のフレームからそれらを推定できます。このブロックは、Pi3Xという追加の依存関係と2番目のチェックポイントを必要とするため、diffusersの外に配置されています。そのため、明示的に要求するまで何もダウンロードされません。

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

2つのディレクトリが書き出されます。ベースのpipelineは--dstに、ステージ2のリファイナーは./SANA-WM_bidirectional-diffusers-refinerに書き出されます(--dst-refinerで上書き可能)。その後、それぞれをローカルパスから通常どおり読み込みます。

コンポーネント

ステージ2のリファイナーは独立したリポジトリであり、独自のtransformer(SanaWMLTX2RefinerTransformer3DModel)、connectors(LTX2TextConnectors)、tokenizer、Gemma-3のtext_encoderおよびschedulerを備えています。独自のvaeは持たないため、ベースのパイプラインのものを渡してください。

SanaWMPipeline

class diffusers.SanaWMPipeline

< >

( 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 風のトークンごとのタイムステップ)。latents を直接扱います。

SANA-WMのカメラ制御による画像から動画へのパイプライン。

c2w ポーズまたは WASD/IJKL アクション文字列で指定されたカメラ軌跡、テキストプロンプト、および最初のフレームの画像から動画を生成します。ステージ1のサンプリングには1600Mの双方向 SANA DiT を、ステージ2の仕上げには LTX-2 sink-bidirectional Euler refiner を使用し、どちらも LTX-2 VAE を通じてデコードされます。

__call__

< >

( 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、optional) — (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) — ターゲットのフレーム数。LTX-2 VAE の制約により、8k+1 にスナップされます。
  • fps (int、デフォルト値: 16) — 出力フレームレート(リファイナーにも渡されます)。
  • num_inference_steps (int、デフォルト値は 60) — ステージ 1 の DiT サンプリングステップ数。
  • guidance_scale (float、デフォルト値は 5.0) — Classifier-free guidance のスケール。
  • negative_prompt (str、デフォルト値は "") — 任意のネガティブプロンプト。
  • generator (torch.Generator または list[torch.Generator]、省略可) — ノイズサンプリングを決定論的にするための1つ以上の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) — プロンプトの最大トークン数。
  • chi_prompt (list[str]、optional) — chi-prompt のプレフィックスを上書きします(デフォルトは公開リリースと同一です)。
  • output_type(“np”、“pil”、“pt”、または “latent”。デフォルトは “np”)— 出力形式。
  • callback_on_step_end (Callable, optional) — 各デノイジングステップの終了時に、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

< >

( c2w: np.ndarray | Noneaction: str | Noneintrinsics: np.ndarray | list[float] | Nonecallback_on_step_end_tensor_inputs: list[str] | None = None )

__call__ の入力を検証します。無効な入力の場合は例外を発生させ、何も返しません。

encode_prompt

< >

( 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プロンプトをユーザープロンプトの先頭に付加し、その後、select_index = [0, -L+1, ..., -1]のスライスによってBOSトークンと最後のmax_sequence_length - 1トークンを取得します。

prepare_camera_trajectory

< >

( 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

< >

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

初期の潜在変数をサンプリングし、最初のフレームを条件付けアンカーとして固定します。

(latents, condition_mask) を返します。condition_mask は、最初のフレームのトークン(サンプリング全体を通じてクリーンな状態が維持されます)に 1、それ以外の位置に 0 を持ちます。

SanaWMLTX2Refiner

LTX-2 のステージ2リファイナーは、ステージ1の潜在変数を受け取る独立した DiffusionPipeline です。vae(ベースのパイプラインのものを指定し、重みを共有する)を渡すと動画をデコードして返しますが、指定しない場合はリファインされた潜在変数を返します。

class diffusers.SanaWMLTX2Refiner

< >

( 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 で構築されるため、蒸留されたシグマは変更なしでそのまま渡されます。
  • vae ([AutoencoderKLLTX2Video]、任意) — [SanaWMPipeline] で使用されているのと同じ VAE です。vae=pipe.vae を渡すことで重みを共有できます。指定された場合、リファイナは動画をデコードし、指定されなかった場合はリファインされた潜在変数を返します。
  • text_max_sequence_length (int、デフォルト値: 1024) — Gemma-3 トークナイザーに渡すトークンの最大数。

LTX-2 sink-bidirectional Euler refiner — SANA-WM ステージ 2、スタンドアロンパイプラインとして。

LTX-2のコンポーネント(リファイナートランスフォーマー、テキストコネクタ、Gemma-3テキストエンコーダー、トークナイザー)と、蒸留されたシグマスケジュールを保持しEulerステップを実行する[FlowMatchEulerDiscreteScheduler]をラップします。これは[SanaWMPipeline]の任意のコンポーネントとして登録されており、単独で使用してステージ1の潜在変数をリファインすることもできます。

__call__

< >

( 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) ステージ1の潜在変数。
  • prompt — テキストプロンプト。
  • fps — 動画のフレームレート(LTX-2 RoPE の時間方向スケーリングを制御します)。
  • sink_size — 注意力のシンクとしてアンカーする先頭のz_sana生フレームの数(標準値: 1)。
  • generator — FMエンドポイントのノイズに使用するtorch.Generator。デフォルトではシード42で初期化されたジェネレーターが使用され、そのまま再現可能な結果が得られます。
  • block_size — ARブロックあたりの潜在フレーム数(標準値: 3)。
  • kv_max_frames — スライディングウィンドウ内に保持される最大コンテキスト+アクティブフレーム数(標準値: 11 = 1つのシンク + 直近10フレーム)。
  • sigmas — 0.0 で終了する降順の Euler スケジュール(標準的な 3 ステップ蒸留: (0.909375, 0.725, 0.421875, 0.0))。末尾の 0.0 を除いた状態で self.scheduler に渡されます(0.0 はスケジューラ側で自動的に追加されます)。
  • output_type — “latent” を指定すると、精錬された latents が返されます。それ以外の値を指定すると、self.vae を通じてデコードされ、“np”、“pt”、“pil” のいずれかの型に後処理されます。vae が存在しない場合は、いずれの値を指定しても latents が返されます。

戻り値

torch.Tensor

(B, C, F, H, W) の形状を持つ洗練された VAE の潜在変数 — 最初の sink_size フレームは、ステージ 1 のシンク潜在変数をそのまま保持し、残りは洗練された出力を保持します。

LTX-2 リファイナーを実行し、リファインされたVAE潜在変数を返します。

モデルの学習に使用されたチャンク因果 AR レシピ(block_size=3、kv_max_frames=11)を用います。[source_sink + recent_history + active_block] の K/V をスライディングウィンドウとして、ブロックごとに 1 つずつ transformer に供給します。これにより、ブロックごとの計算量は上限が設定され、リファインメントの総コストは動画の長さに比例して線形に増加します。

SanaWMPipelineOutput

クラス diffusers.SanaWMPipelineOutput

< >

( 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" の場合は [0, 1] の範囲の float 型 np.ndarray で、形状は (T, H, W, 3)、"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)(リファイナが実行された場合、リファイナはシンクアンカーフレームを削除するため、この配列はそれに合わせて再整合されます)。
  • latent (torch.Tensor, 省略可) — LTX-2 VAE 空間における潜在テンソル。形状は (B, C, T_lat, H_lat, W_lat) です。 `output_type=“latent”` の場合に返されます。

SANA-WMの画像から動画へのパイプライン用の出力クラスです。

Update on GitHub