Hugging Face's logo

Diffusers documentation

SanaWMLTX2RefinerTransformer3DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

SanaWMLTX2RefinerTransformer3DModel

SanaWMPipeline의 2단계로 사용되며 SanaWMLTX2Refiner에 의해 구동되는 chunk-causal autoregressive refiner DiT입니다.

구조적으로 LTX2VideoTransformer3DModel과 동일합니다. 설정 인자, 서브모듈, 파라미터 이름이 모두 같으므로, 공개된 LTX-2 체크포인트를 그대로 로드할 수 있습니다. 포워드 패스만 다릅니다:

  • 비디오 스트림만 실행됩니다(오디오 및 오디오/비디오 교차 주의력 브랜치는 건너뜁니다),
  • self-attention은 attention sink와 최근 정제된 히스토리를 포함하는 명시적 슬라이딩 윈도우 KV 캐시(SanaWMRefinerKVCache)에 대해 실행되므로 블록당 계산량이 제한되고 전체 정제 비용은 비디오 길이에 따라 선형으로 증가합니다.
  • 호출자가 비디오 RoPE를 제공하여, 각 자기회귀 윈도우가 소스 비디오에서 모든 프레임의 절대 인덱스를 유지할 수 있게 합니다(자세한 내용은 SanaWMLTX2RefinerTransformer3DModel.build_rotary_emb_for_absolute_positions() 참조).

모델은 다음으로 로드할 수 있습니다:

import torch
from diffusers import SanaWMLTX2RefinerTransformer3DModel

transformer = SanaWMLTX2RefinerTransformer3DModel.from_pretrained(
    "Efficient-Large-Model/SANA-WM_bidirectional-diffusers",
    subfolder="refiner/transformer",
    torch_dtype=torch.bfloat16,
)

SanaWMLTX2RefinerTransformer3DModel

클래스 diffusers.SanaWMLTX2RefinerTransformer3DModel

< >

( in_channels: int = 128out_channels: int | None = 128patch_size: int = 1patch_size_t: int = 1num_attention_heads: int = 32attention_head_dim: int = 128cross_attention_dim: int = 4096vae_scale_factors: tuple[int, int, int] = (8, 32, 32)pos_embed_max_pos: int = 20base_height: int = 2048base_width: int = 2048gated_attn: bool = Falsecross_attn_mod: bool = Falseaudio_in_channels: int = 128audio_out_channels: int | None = 128audio_patch_size: int = 1audio_patch_size_t: int = 1audio_num_attention_heads: int = 32audio_attention_head_dim: int = 64audio_cross_attention_dim: int = 2048audio_scale_factor: int = 4audio_pos_embed_max_pos: int = 20audio_sampling_rate: int = 16000audio_hop_length: int = 160audio_gated_attn: bool = Falseaudio_cross_attn_mod: bool = Falsenum_layers: int = 48activation_fn: str = 'gelu-approximate'qk_norm: str = 'rms_norm_across_heads'norm_elementwise_affine: bool = Falsenorm_eps: float = 1e-06caption_channels: int = 3840attention_bias: bool = Trueattention_out_bias: bool = Truerope_theta: float = 10000.0rope_double_precision: bool = Truecausal_offset: int = 1timestep_scale_multiplier: int = 1000cross_attn_timestep_scale_multiplier: int = 1000rope_type: str = 'interleaved'use_prompt_embeddings = Trueperturbed_attn: bool = Falseff_bias: bool = Trueaudio_ff_bias: bool = Trueuse_prompt_adaln_single: bool = Trueuse_keyframes_abs_pos_embedding: bool = False )

매개변수

  • in_channels (int, 기본값 128) — 입력의 채널 수입니다.
  • out_channels (int, 기본값 128) — 출력의 채널 수입니다.
  • patch_size (int, 기본값 1) — 패치 임베딩 레이어에서 사용할 공간 패치의 크기입니다.
  • patch_size_t (int, 기본값 1) — 패치 임베딩 레이어에서 사용할 시간적 패치의 크기입니다.
  • num_attention_heads (int, 기본값 32) — 멀티헤드 어텐션에 사용할 헤드의 수입니다.
  • attention_head_dim (int, 기본값 128) — 각 헤드당 채널 수입니다.
  • cross_attention_dim (int, 기본값 4096) — 크로스 어텐션 헤드의 채널 수입니다.
  • num_layers (int, 기본값 48) — 사용할 Transformer 블록의 레이어 수입니다.
  • activation_fn (str, 기본값 "gelu-approximate") — 피드포워드에서 사용할 활성화 함수입니다.
  • qk_norm (str, 기본값 "rms_norm_across_heads") — 사용할 정규화 레이어입니다.
  • rope_type (str, 기본값 "interleaved") — 사용할 RoPE 적용 방식("interleaved" 또는 "split").

SANA-WM의 2단계로 사용되는 chunk-causal autoregressive refiner transformer입니다.

LTX2VideoTransformer3DModel와 아키텍처가 동일합니다. 설정 인자, 서브모듈, 파라미터 이름이 모두 같으므로, 공개된 LTX-2 체크포인트를 그대로 로드할 수 있습니다. 다른 점은 포워드 패스입니다:

  • 비디오 스트림만 실행됩니다(오디오 및 오디오/비디오 교차 주의력 브랜치는 건너뜁니다),
  • self-attention은 attention sink와 최근 정제된 히스토리를 포함하는 명시적 슬라이딩 윈도우 KV 캐시(SanaWMRefinerKVCache)에 대해 실행되므로, 정제 비용은 AR 블록당 상한이 설정되며 비디오 길이에 따라 선형으로 증가합니다.
  • 호출자가 비디오 RoPE를 제공하여 각 AR 윈도우가 소스 비디오에서 모든 프레임의 절대 인덱스를 유지할 수 있게 합니다(build_rotary_emb_for_absolute_positions() 참조).

나머지 인수는 LTX2VideoTransformer3DModel의 인수를 1:1로 그대로 따릅니다. 오디오 쪽 인수와 서브모듈은 체크포인트의 오디오 가중치가 왕복(round-trip)되도록 유지되는 것일 뿐이며, 리파이너 포워드에서는 사용되지 않습니다.

build_rotary_emb_for_absolute_positions

< >

( batch_size: intframe_positions: list[int]height: intwidth: intdevice: devicefps: float ) → tuple[torch.Tensor, torch.Tensor]

매개변수

  • batch_size (int) — 좌표를 브로드캐스트할 배치 크기입니다.
  • frame_positions (list[int]) — 이 윈도우가 포함하는 절대 잠재 프레임 인덱스입니다.
  • height (int), width (int) — 잠재 공간 해상도.
  • device (torch.device) — 좌표를 생성할 디바이스.
  • fps (float) — LTX-2의 시간적 RoPE 스케일링을 결정하는 비디오 프레임 속도입니다.

반환값

tuple[torch.Tensor, torch.Tensor]

(cos, sin) RoPE 튜플

명시된 절대 잠재 프레임 인덱스 목록에 대해 비디오 RoPE를 구성합니다.

SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords는 연속적인 torch.arange(num_frames)을 가정하며, 이는 양방향 추론에는 문제가 없습니다. 슬라이딩 윈도우 AR 리파이너는 각 프레임의 원본 비디오 내 절대 인덱스를 유지해야 하므로, [sink + recent + active] 윈도우에 걸쳐 RoPE가 올바른 시간적 위상을 포착할 수 있습니다.

forward

< >

( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorvideo_rotary_emb: tuple[torch.Tensor, torch.Tensor]encoder_attention_mask: typing.Optional[torch.Tensor] = Nonekv_cache: diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache | None = Nonekv_cache_mode: str | None = Noneattention_kwargs: dict | None = Nonereturn_dict: bool = True ) → Transformer2DModelOutput or tuple

매개변수

  • hidden_states (torch.Tensor) — 활성 블록의 패치화(video)된 비디오 레이턴트(latents)로, 모양은 (batch_size, num_video_tokens, in_channels)입니다.
  • encoder_hidden_states (torch.Tensor) — (batch_size, text_seq_len, caption_channels) 모양의 텍스트 임베딩.
  • timestep (torch.Tensor) — (batch_size, num_video_tokens) 모양의 Timestep으로, 이미 self.config.timestep_scale_multiplier로 스케일링된 값입니다.
  • video_rotary_emb (tuple[torch.Tensor, torch.Tensor]) — build_rotary_emb_for_absolute_positions()에서 반환된 활성 블록의 절대 프레임 위치를 위한 (cos, sin) RoPE입니다.
  • encoder_attention_mask (torch.Tensor, 선택 사항) — (batch_size, text_seq_len) 모양의 곱셈 텍스트 어텐션 마스크입니다.
  • kv_cache (SanaWMRefinerKVCache, 선택) — 레이어별 attention sink와 최근 refined history를 저장하는 sliding-window KV cache이며, 공유 sink_pe RoPE를 포함합니다. kv_cache_mode가 설정된 경우 필수입니다.
  • kv_cache_mode (str, 선택 사항) — 다음 중 하나입니다:
    • "inject": [sink + history + current] K/V에 attend합니다(디노이징 단계).
    • "capture_pre_rope": 접두사가 없습니다. 이 forward의 pre-RoPE K/V를 캐시에 기록합니다(1회 사용되며, raw stage-1 latents에서 attention sink를 시드하는 데 사용됩니다).
    • "inject_and_capture_post_rope": [sink + history + current] K/V에 attend하고, 이 블록의 post-RoPE K/V를 캐시에 기록하여 history에 추가할 수 있도록 합니다.

    None인 경우, 해당 블록은 현재 토큰에 대해서만 일반 전체 셀프 어텐션을 실행합니다.

  • attention_kwargs (dict, 선택 사항) — LoRA scale 처리로 전달되는 선택적 kwargs입니다.
  • return_dict (bool, 선택 사항, 기본값 True) — 일반 튜플 대신 Transformer2DModelOutput을 반환할지 여부.

반환값

Transformer2DModelOutput 또는 tuple

(batch_size, num_video_tokens, out_channels) 모양의 활성 블록에 대한 예측 속도

단일 AR 블록에 대한 비디오 전용 forward pass.

SanaWMRefinerKVCache

클래스 diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache

< >

( num_layers: int )

매개변수

  • num_layers (int) — 레이어별 캐시를 할당할 transformer 블록의 수.

각 transformer 블록마다 [SanaWMRefinerKVLayerCache] 하나를 포함하는 컨테이너이며, 공유 sink RoPE도 함께 포함합니다.

이 구현은 SANA-WM stage-2 refiner가 학습된 rf_shifted_sink KV-cache 계약을 따릅니다. 정제 과정은 chunk-causal 방식으로, block_size개의 latent 프레임을 한 번에 denoise하면서 [attention sink + 최근 히스토리 + 활성 블록]로 구성된 제한된 윈도우의 K/V에 attend합니다.

  • sink_pe: sink 프레임에 대한 (cos, sin) RoPE 튜플입니다. 유한한 워킹 캐시 바로 앞에 sink가 위치하도록, 슬라이딩 sink_rope_offset 기준으로 각 AR 윈도우마다 다시 구성됩니다. RoPE가 레이어에 의존하지 않으므로 레이어 간에 공유됩니다.

SanaWMRefinerKVLayerCache

클래스 diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache

< >

( )

SANA-WM stage-2 chunk-causal AR refiner의 레이어별 KV 캐시.

리파이너의 self-attention이 참조하는 슬라이딩 윈도우 접두사의 두 부분을 보유하며, 마지막 forward에서 캡처한 K/V를 다시 읽기 위한 슬롯도 포함합니다. 모든 텐서는 (batch_size, num_tokens, inner_dim) 형태(즉, 헤드별 unflatten 이전)이며, 리파이너의 self-attention이 연결하는 레이아웃과 일치합니다.

  • sink_k_pre / sink_v: pre-RoPE K/V of the attention-sink frames, raw stage-1 latents에서 한 번만 캡처됩니다. RoPE 적용 전 상태로 저장되어 각 AR 윈도우가 자체적으로 시프트된 sink offset(SanaWMRefinerKVCache.sink_pe)에 RoPE를 다시 적용할 수 있습니다.
  • history_k / history_v: 이미 정제된 최근 프레임의 post-RoPE K/V로, 그대로 연결할 수 있는 상태
  • captured_k_pre / captured_v_pre 및 captured_k_post / captured_v_post: 캡처 kv_cache_mode에 의해 기록되는 readback 슬롯입니다.

get_captured_post_rope

< >

( )

마지막 forward에서 포착한 post-RoPE K/V를 pop합니다.

get_captured_pre_rope

< >

( )

마지막 forward에서 캡처된 pre-RoPE K/V를 pop합니다.

get_history

< >

( )

RoPE 적용 후 최근 히스토리 K/V를 반환하거나, 비어 있으면 None을 반환합니다.

get_sink

< >

( )

pre-RoPE sink K/V를 반환합니다. 캡처되지 않았거나(또는 비어 있으면) None을 반환합니다.

store_captured_post_rope

< >

( key: Tensorvalue: Tensor )

현재 forward에서 생성된 post-RoPE K/V를 저장합니다.

store_captured_pre_rope

< >

( key: Tensorvalue: Tensor )

현재 forward에서 생성된 pre-RoPE K/V를 저장합니다.

store_history

< >

( history_k: Tensorhistory_v: Tensor )

RoPE 적용 후 최근 히스토리의 K/V를 저장합니다.

store_sink

< >

( sink_k_pre: Tensorsink_v: Tensor )

RoPE 적용 전 sink K/V를 저장합니다.

Transformer2DModelOutput

클래스 diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( sample: torch.Tensor )

매개변수

  • sample (torch.Tensor, shape (batch_size, num_channels, height, width) 또는 Transformer2DModel이 discrete인 경우 (batch size, num_vector_embeds - 1, num_latent_pixels)) — encoder_hidden_states 입력에 조건부인 hidden states 출력입니다. discrete인 경우, 노이즈가 제거된 latent 픽셀에 대한 확률 분포를 반환합니다.

Transformer2DModel의 출력입니다.

Update on GitHub

←SanaVideoTransformer3DModel SanaWMTransformer3DModel→