SANA-WM

SANA-WM es un modelo de mundo de imagen a video controlado por cámara, construido sobre SANA. Dada una imagen de primer fotograma, un prompt de texto y una trayectoria de cámara (ya sea una pose c2w explícita o una cadena de acciones WASD/IJKL), genera un video cuyo movimiento sigue la trayectoria de cámara solicitada.

La inferencia se ejecuta en dos etapas:

1. Etapa 1 — SANA-WM DiT. Un DiT bidireccional de 1.6B parámetros con atención lineal GDN-Triton y una rama de control de cámara UCPE. El muestreo utiliza un planificador Euler de flow-matching estilo LTX con pasos de tiempo por token; el primer marco latente es la ancla de condicionamiento.
2. Etapa 2 — Refinador LTX-2 (opcional). Un pipeline de refinamiento Euler bidireccional con sink independiente (SanaWMLTX2Refiner) que envuelve SanaWMLTX2RefinerTransformer3DModel + LTX2TextConnectors y un codificador de texto Gemma-3, ejecutado durante 3 pasos de sigma destilados.

Ambas etapas decodifican a través del VAE AutoencoderKLLTX2Video.

Modelos disponibles:

SANA-WM se entrena a una resolución fija de 704×1280. El dtype recomendado para los pesos del transformer es torch.bfloat16; mantenga el codificador de texto en torch.bfloat16 y el VAE en torch.float32 para obtener la mejor precisión numérica. El pipeline espera las intrínsecas de la cámara [fx, fy, cx, cy] en coordenadas de píxeles de la imagen original; la transformación de redimensionado y recorte centrado se aplica internamente.

Inferencia

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
)

Pasa action=None y proporciona tus propias poses c2w (un array numpy de (F, 4, 4)) para controlar la trayectoria de la cámara de forma explícita. Omita output_type="latent" para obtener el video directamente de la etapa 1 y saltarse el refinador.

Etapa 2 — el refinador LTX-2

SanaWMLTX2Refiner se distribuye como su propio repositorio, de la misma manera que SDXL separa el modelo base y el refiner. Pasa el VAE del pipeline base para que los pesos se compartan en lugar de cargarse dos veces:

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)

Sin un vae, el refiner devuelve latentes refinados en lugar de video, lo cual es útil si deseas decodificarlos tú mismo.

Habilite el offloading en ambos pipelines, o libere la etapa 1 antes de la etapa 2 (pipe.transformer.to("cpu")). Las dos etapas juntas ocupan alrededor de 45 GB en torch.bfloat16, y mantener ambas residentes en una tarjeta de 80 GB deja demasiado poco espacio para las activaciones. Tenga en cuenta también que la etapa 2 respeta torch_dtype: cargar el refiner en torch.float32 duplica aproximadamente su memoria y cambia ligeramente la salida.

Si no tienes las intrínsecas de la cámara, un bloque modular alojado puede estimarlas a partir de un solo fotograma. Se encuentra fuera de diffusers porque incorpora Pi3X, una dependencia adicional y un segundo punto de control, por lo que no se descarga nada hasta que lo solicites:

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

Convertir el punto de control publicado

Si tienes la versión de origen de SANA-WM (no la instantánea de diffusers preconvertida), ejecuta el script de conversión una vez:

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

Esto escribe dos directorios: el pipeline base en --dst y el refiner de la etapa 2 junto a él en ./SANA-WM_bidirectional-diffusers-refiner (se puede sobrescribir con --dst-refiner). Luego, carga cada uno desde su ruta local como de costumbre.

Componentes

• tokenizer — GemmaTokenizerFast
• text_encoder — Gemma-2 (devuelve los estados ocultos del decodificador)
• vae — AutoencoderKLLTX2Video (LTX-2, espacial ×32 / temporal ×8)
• transformer — SanaWMTransformer3DModel, DiT bidireccional de 1.600 millones de parámetros
• scheduler — FlowMatchEulerDiscreteScheduler

El refiner de la etapa 2 es un repositorio independiente con su propio transformer (SanaWMLTX2RefinerTransformer3DModel), connectors (LTX2TextConnectors), tokenizer, text_encoder Gemma-3 y scheduler. No tiene un vae propio: pase el de la 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

El refinador de la etapa 2 de LTX-2 es un DiffusionPipeline independiente que recibe los latentes de la etapa 1. Si se le proporciona un vae (el de la pipeline base, para compartir los pesos), decodifica a video; si no, devuelve los latentes refinados.

SanaWMPipelineOutput