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 액션 문자열)를 입력으로 받아, 요청된 카메라 경로를 따르는 움직임이 포함된 비디오를 생성합니다.

Inference는 두 단계로 실행됩니다:

  1. 1단계 — SANA-WM DiT. GDN-Triton 선형 어텐션과 UCPE 카메라 제어 분기를 갖춘 1.6B 파라미터 양방향 DiT입니다. 샘플링에는 토큰별 타임스텝을 사용하는 LTX 스타일 플로우 매칭 오일러 스케줄러를 사용하며, 첫 번째 잠재 프레임이 조건부 앵커 역할을 합니다.
  2. Stage 2 — LTX-2 refiner (선택 사항). SanaWMLTX2RefinerTransformer3DModel + LTX2TextConnectors와 Gemma-3 텍스트 인코더를 래핑하는 별도의 sink-bidirectional Euler refiner 파이프라인(SanaWMLTX2Refiner)으로, 3개의 distilled sigma 단계 동안 실행됩니다.

두 단계 모두 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이며, 최적의 수치 안정성을 위해 텍스트 인코더는 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단계에서 바로 비디오를 얻고 리파이너를 건너뛸 수 있습니다.

Stage 2 — LTX-2 리파이너

SanaWMLTX2Refiner는 SDXL이 base와 refiner를 분리하는 방식과 마찬가지로 자체 저장소로 제공됩니다. 가중치가 두 번 로드되지 않고 공유되도록 base 파이프라인의 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를 반환합니다. 이는 직접 디코딩을 원하는 경우 유용합니다.

두 파이프라인 모두에서 offloading을 활성화하거나, stage 2를 실행하기 전에 stage 1을 해제(pipe.transformer.to("cpu"))할 수 있습니다. 두 stage의 합계는 torch.bfloat16에서 약 45 GB이며, 두 stage를 모두 80 GB 카드 하나에 상주시키면 activations를 위한 여유 공간이 부족합니다. 또한 stage 2는 torch_dtype을 따르므로, refiner를 torch.float32로 로드하면 메모리가 약 2배로 증가하고 출력이 약간 달라집니다.

카메라 내재적 매개변수(camera intrinsics)가 없는 경우, 호스팅된 모듈형 블록(modular block)을 사용하여 단일 프레임에서 이를 추정할 수 있습니다. 이 블록은 Pi3X라는 추가 의존성과 두 번째 체크포인트를 가져오기 때문에 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

이 명령은 두 개의 디렉터리를 작성합니다. 기본 파이프라인은 --dst에, stage-2 refiner는 ./SANA-WM_bidirectional-diffusers-refiner에 함께 생성됩니다(--dst-refiner로 덮어쓸 수 있음). 이후 각 모델을 로컬 경로에서 평소와 같이 로드합니다.

구성 요소

stage-2 refiner는 자체 transformer(SanaWMLTX2RefinerTransformer3DModel), connectors(LTX2TextConnectors), tokenizer, Gemma-3 text_encoder 및 scheduler를 갖춘 별도의 저장소입니다. 자체 vae가 없으므로 base pipeline의 것을 전달해야 합니다.

SanaWMPipeline

클래스 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 스타일의 토큰별 timesteps). 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, 선택 사항) — (F, 4, 4) 카메라-월드 포즈. action과 상호 배타적입니다.
  • action (str, 선택 사항) — Action-DSL 문자열, 예: "w-80,jw-40,w-40". c2w와 상호 배타적입니다.
  • intrinsics (np.ndarray 또는 list[float]) — [fx, fy, cx, cy]는 원본 이미지 픽셀 좌표계입니다. 파이프라인은 내부적으로 resize+crop 변환을 적용합니다.
  • height (int, 기본값 704) — 출력 프레임 높이 (공개 모델의 경우 고정됨).
  • width (int, 기본값 1280) — 출력 프레임의 너비 (공개 모델의 경우 고정됨).
  • num_frames (int, 기본값 161) — 목표 프레임 수; 8k+1로 정렬됩니다 (LTX-2 VAE 제약 조건).
  • fps (int, 기본값 16) — 출력 프레임 레이트(리파이너에도 전달됨).
  • num_inference_steps (int, 기본값 60) — 1단계 DiT 샘플링 단계 수.
  • guidance_scale (float, 기본값 5.0) — Classifier-free guidance scale입니다.
  • negative_prompt (str, 기본값 "") — 선택 사항인 negative prompt입니다.
  • generator (torch.Generator 또는 list[torch.Generator], 선택 사항) — 노이즈 샘플링을 결정적으로 만들기 위한 하나 이상의 torch 생성기입니다. None인 경우 샘플링은 비결정적입니다.
  • prompt_embeds (torch.Tensor, 선택) — 텍스트 인코더를 건너뛰기 위한 사전 계산된 텍스트 임베딩입니다. prompt_attention_mask와 함께 전달해야 합니다.
  • prompt_attention_mask (torch.Tensor, 선택 사항) — prompt_embeds에 대한 어텐션 마스크입니다.
  • negative_prompt_embeds (torch.Tensor, 선택 사항) — 사전 계산된 negative text embeddings입니다. 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, 선택 사항) — 각 노이즈 제거 단계의 끝에서 callback_on_step_end(self, step, timestep, callback_kwargs)로 호출되며, callback_kwargs에는 callback_on_step_end_tensor_inputs에 명시된 이름의 텐서가 포함됩니다. 반환된 dict의 텐서는 파이프라인의 텐서를 대체합니다.
  • callback_on_step_end_tensor_inputs (list[str], 기본값 [“latents”]) — callback_on_step_end에 전달할 텐서를 지정합니다. ._callback_tensor_inputs의 부분 집합이어야 합니다.
  • return_dict (bool, defaults to True) — Return [SanaWMPipelineOutput] vs tuple.

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를 샘플링하고 첫 번째 프레임을 conditioning anchor로 고정합니다.

(latents, condition_mask)을 반환합니다. 여기서 condition_mask는 첫 번째 프레임의 토큰에 1을, 나머지 부분에 0을 가집니다(첫 번째 프레임의 토큰은 샘플링 전체 과정에서 깨끗한 상태로 유지됩니다).

SanaWMLTX2Refiner

LTX-2 stage-2 refiner는 stage-1 latents를 입력으로 받는 독립형 DiffusionPipeline입니다. vae(베이스 파이프라인의 VAE로, 가중치가 공유됩니다)을 제공하면 비디오로 디코딩하여 반환하고, 제공하지 않으면 정제된 latents를 반환합니다.

클래스 diffusers.SanaWMLTX2Refiner

< >

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

매개변수

  • transformer ([SanaWMLTX2RefinerTransformer3DModel]) — chunk-causal sliding-window 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를 전달하세요. 값이 주어지면 리파이너는 비디오로 디코딩하고, 그렇지 않으면 정제된 latents를 반환합니다.
  • text_max_sequence_length (int, 기본값 1024) — Gemma-3 토크나이저에 전달되는 최대 토큰 수.

LTX-2 sink-bidirectional Euler refiner — SANA-WM stage 2, 독립형 파이프라인으로.

LTX-2 구성 요소(리파이너 트랜스포머 + 텍스트 커넥터 + Gemma-3 텍스트 인코더 + 토크나이저)와 증류된 시그마 스케줄을 포함하고 오일러 스텝을 수행하는 [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) stage-1 latent.
  • prompt — 텍스트 프롬프트.
  • fps — 비디오 프레임 레이트(LTX-2 RoPE의 시간적 스케일링을 결정합니다).
  • sink_size — attention sink로 고정할 앞부분의 원시 z_sana 프레임 수 (기본값: 1).
  • generator — FM 엔드포인트 노이즈에 사용되는 torch.Generator입니다. 기본적으로 42로 시드된 생성기가 사용되어 결과의 재현성이 보장됩니다.
  • block_size — AR 블록당 latent 프레임 수 (기본값: 3).
  • kv_max_frames — 슬라이딩 윈도우에 유지되는 최대 컨텍스트+활성 프레임 수 (기본값: 11 = 1 sink + 10 recent).
  • sigmas — 0.0에서 종료되는 내림차순 Euler 스케줄(정규화된 3단계 증류: (0.909375, 0.725, 0.421875, 0.0)). self.scheduler에 전달되며, 마지막 0.0은 제외됩니다(scheduler가 자체적으로 추가함).
  • output_type — “latent”을 지정하면 정제된 latents를 반환합니다. 그 외의 값은 self.vae를 통해 디코딩한 후 해당 유형(“np”, “pt”, “pil”)으로 후처리하여 반환합니다. vae가 없으면 어떤 경우에도 latents가 반환됩니다.

반환값

torch.Tensor

(B, C, F, H, W) 형태의 정제된 VAE latents — 처음 sink_size 프레임은 stage-1 sink latents를 그대로 유지하고, 나머지 프레임은 정제된 출력을 포함합니다.

LTX-2 리파이너를 실행하고 정제된 VAE 잠재를 반환합니다.

모델 학습에 사용된 chunk-causal AR 레시피(block_size=3, kv_max_frames=11)를 사용합니다. [source_sink + recent_history + active_block] 슬라이딩 윈도우의 K/V를 블록 단위로 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] 범위의 (T, H, W, 3) 형상 float 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) (리파이너가 실행된 경우, sink 앵커 프레임이 제거되고 이 배열도 그에 따라 다시 정렬됩니다).
  • latent (torch.Tensor, 선택 사항) — LTX-2 VAE 공간의 잠재 텐서, 모양은 (B, C, T_lat, H_lat, W_lat)입니다. `output_type=“latent”`일 때 반환됩니다.

SANA-WM 이미지-비디오 파이프라인의 출력 클래스입니다.

Update on GitHub