SANA-WM

SANA-WM là mô hình thế giới từ ảnh sang video có điều khiển bởi camera, được xây dựng dựa trên SANA. Với một hình ảnh khung đầu tiên, một prompt văn bản và một quỹ đạo camera (dạng các tư thế c2w tường minh hoặc chuỗi hành động WASD/IJKL), nó tạo ra một video có chuyển động tuân theo đường đi camera được yêu cầu.

Việc suy luận được thực hiện qua hai giai đoạn:

1. Giai đoạn 1 — SANA-WM DiT. Một DiT hai chiều có 1,6 tỷ tham số, sử dụng chú ý tuyến tính GDN-Triton và một nhánh điều khiển camera UCPE. Quá trình lấy mẫu sử dụng trình lập lịch Euler flow-matching kiểu LTX với các bước thời gian cho từng token; khung tiềm ẩn đầu tiên đóng vai trò là neo điều kiện.
2. Giai đoạn 2 — bộ tinh chỉnh LTX-2 (tuỳ chọn). Một pipeline Euler hai chiều kiểu sink riêng biệt (SanaWMLTX2Refiner) bao bọc SanaWMLTX2RefinerTransformer3DModel + LTX2TextConnectors và một bộ mã hoá văn bản Gemma-3, chạy trong 3 bước sigma đã chưng cất.

Cả hai giai đoạn đều giải mã thông qua VAE AutoencoderKLLTX2Video.

Các model có sẵn:

SANA-WM được huấn luyện ở độ phân giải cố định 704×1280. dtype được khuyến nghị cho trọng số của transformer — giữ text encoder ở torch.bfloat16 và VAE ở torch.float32 để có số học tốt nhất. Pipeline kỳ vọng các tham số nội tại của camera [fx, fy, cx, cy] trong tọa độ pixel của ảnh gốc; phép biến đổi resize-and-center-crop được áp dụng bên trong.

Suy luận

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
)

Truyền action=None và cung cấp các tư thế c2w của riêng bạn ((F, 4, 4) mảng numpy) để điều khiển quỹ đạo camera một cách tường minh. Bỏ output_type="latent" để nhận video trực tiếp từ giai đoạn 1 và bỏ qua bộ tinh chỉnh.

Giai đoạn 2 — bộ tinh chỉnh LTX-2

SanaWMLTX2Refiner được phát hành dưới dạng một kho lưu trữ riêng, tương tự cách SDXL tách riêng base và refiner. Hãy truyền VAE của pipeline base để các trọng số được chia sẻ thay vì được nạp hai lần:

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)

Nếu không có vae, refiner sẽ trả về các latents đã tinh chỉnh thay vì video, điều này hữu ích nếu bạn muốn tự giải mã.

Bật offloading cho cả hai pipeline, hoặc giải phóng stage 1 trước khi chạy stage 2 (pipe.transformer.to("cpu")). Tổng dung lượng của hai stage này vào khoảng 45 GB ở torch.bfloat16, và việc giữ cả hai trên một card 80 GB sẽ không còn đủ chỗ cho activations. Lưu ý rằng stage 2 tuân theo torch_dtype: nạp refiner ở torch.float32 sẽ làm tăng bộ nhớ lên gấp đôi và thay đổi đầu ra một chút.

Nếu bạn không có thông số nội tại của camera, một khối module được lưu trữ sẵn có thể ước lượng chúng từ một khung hình duy nhất. Khối này nằm bên ngoài diffusers vì nó kéo theo Pi3X — một phụ thuộc bổ sung và một checkpoint thứ hai — nên không có gì được tải xuống cho đến khi bạn yêu cầu:

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

Chuyển đổi checkpoint đã phát hành

Nếu bạn có bản phát hành SANA-WM nguồn (không phải bản snapshot diffusers đã chuyển đổi sẵn), hãy chạy script chuyển đổi một lần:

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

Việc này sẽ ghi hai thư mục: pipeline cơ bản tại --dst, và refiner giai đoạn 2 nằm cạnh nó tại ./SANA-WM_bidirectional-diffusers-refiner (có thể ghi đè bằng --dst-refiner). Sau đó, hãy nạp từng thư mục từ đường dẫn cục bộ của chúng như bình thường.

Các thành phần

• tokenizer — GemmaTokenizerFast
• text_encoder — Gemma-2 (trả về trạng thái ẩn của bộ giải mã)
• vae — AutoencoderKLLTX2Video (LTX-2, không gian ×32 / thời gian ×8)
• transformer — SanaWMTransformer3DModel, DiT hai chiều có 1,6 tỷ tham số
• scheduler — FlowMatchEulerDiscreteScheduler

Refiner giai đoạn 2 là một kho lưu trữ riêng với transformer (SanaWMLTX2RefinerTransformer3DModel), connectors (LTX2TextConnectors), tokenizer, text_encoder Gemma-3 và scheduler riêng. Nó không có vae riêng — hãy truyền vae của pipeline cơ sở.

SanaWMPipeline

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

SanaWMLTX2Refiner

Refiner giai đoạn 2 của LTX-2 là một DiffusionPipeline độc lập, nhận latents từ giai đoạn 1. Hãy cung cấp cho nó một vae (của pipeline cơ sở, để các trọng số được chia sẻ) nhằm giải mã thành video; nếu không có, nó sẽ trả về các latents đã tinh chỉnh.

SanaWMPipelineOutput