SANA-WM

SANA-WM é um modelo de mundo de imagem para vídeo controlado por câmera, construído sobre o SANA. Dada uma imagem do primeiro quadro, um prompt de texto e uma trajetória de câmera (seja poses c2w explícitas ou uma string de ações WASD/IJKL), ele gera um vídeo cujo movimento segue o caminho de câmera solicitado.

A inferência é executada em duas etapas:

1. Etapa 1 — SANA-WM DiT. Um DiT bidirecional com 1,6 bilhão de parâmetros, com atenção linear GDN-Triton e um ramo de controle de câmera UCPE. A amostragem usa um agendador Euler de correspondência de fluxo no estilo LTX com tempos por token; o primeiro frame latente é a âncora de condicionamento.
2. Etapa 2 — refinador LTX-2 (opcional). Um pipeline de refinamento Euler bidirecional com sink separado (SanaWMLTX2Refiner) que encapsula SanaWMLTX2RefinerTransformer3DModel + LTX2TextConnectors e um codificador de texto Gemma-3, executado por 3 passos de sigma destilados.

Ambas as etapas decodificam por meio do VAE AutoencoderKLLTX2Video.

Modelos disponíveis:

SANA-WM é treinado em uma resolução fixa de 704×1280. O dtype recomendado é para os pesos do transformer — mantenha o text encoder em torch.bfloat16 e o VAE em torch.float32 para a melhor numérica. O pipeline espera intrínsecas da câmera [fx, fy, cx, cy] em coordenadas de pixels da imagem original; a transformação de redimensionamento e recorte centralizado é aplicada internamente.

Inferência

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
)

Pass action=None e forneça suas próprias poses c2w (array numpy (F, 4, 4)) para controlar explicitamente a trajetória da câmera. Remova output_type="latent" para obter o vídeo diretamente da etapa 1 e pular o refinador.

Etapa 2 — o refinador LTX-2

SanaWMLTX2Refiner é distribuído como um repositório próprio, da mesma forma que o SDXL separa o modelo base do refiner. Passe o VAE do pipeline base para que os pesos sejam compartilhados em vez de serem carregados duas vezes:

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)

Sem um vae, o refiner retorna latentes refinados em vez de vídeo, o que é útil se você quiser decodificar por conta própria.

Ative o offloading nos dois pipelines, ou libere a etapa 1 antes da etapa 2 (pipe.transformer.to("cpu")). As duas etapas juntas ocupam cerca de 45 GB em torch.bfloat16, e manter ambas residentes em uma única GPU de 80 GB deixa pouco espaço para ativações. Observe também que a etapa 2 respeita torch_dtype: carregar o refiner em torch.float32 praticamente dobra seu uso de memória e altera levemente a saída.

Se você não tiver as intrínsecas da câmera, um bloco modular hospedado pode estimá-las a partir de um único quadro. Ele fica fora do diffusers porque ele traz o Pi3X — uma dependência extra e um segundo checkpoint —, então nada é baixado até que você solicite:

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

Convertendo o checkpoint liberado

Se você tiver a liberação original do SANA-WM (não o snapshot pré-convertido do diffusers), execute o script de conversão uma vez:

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

Isso grava dois diretórios: o pipeline base em --dst e o refinador da etapa 2 ao lado, em ./SANA-WM_bidirectional-diffusers-refiner (sobrescreva com --dst-refiner). Em seguida, carregue cada um de seu caminho local como de costume.

Componentes

• tokenizer — GemmaTokenizerFast
• text_encoder — Gemma-2 (retorna os estados ocultos do decodificador)
• vae — AutoencoderKLLTX2Video (LTX-2, espacial ×32 / temporal ×8)
• transformer — SanaWMTransformer3DModel, DiT bidirecional com 1,6 bilhão de parâmetros
• scheduler — FlowMatchEulerDiscreteScheduler

O refinador da etapa 2 é um repositório separado com seu próprio transformer (SanaWMLTX2RefinerTransformer3DModel), connectors (LTX2TextConnectors), tokenizer, text_encoder Gemma-3 e scheduler. Ele não possui um vae próprio — passe o do pipeline base.

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

O refinador de estágio 2 do LTX-2 é um DiffusionPipeline independente que recebe os latentes do estágio 1. Forneça um vae (o do pipeline base, para que os pesos sejam compartilhados) para que ele decodifique para vídeo; sem um, ele retorna os latentes refinados.

SanaWMPipelineOutput