Diffusers documentation
SanaWMLTX2RefinerTransformer3DModel
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
< source >( 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 = (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
< source >( batch_size: intframe_positions: listheight: intwidth: intdevice: devicefps: float ) → tuple[torch.Tensor, torch.Tensor]
명시된 절대 잠재 프레임 인덱스 목록에 대해 비디오 RoPE를 구성합니다.
SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords는 연속적인 torch.arange(num_frames)을 가정하며, 이는 양방향 추론에는 문제가 없습니다. 슬라이딩 윈도우 AR 리파이너는 각 프레임의 원본 비디오 내 절대 인덱스를 유지해야 하므로, [sink + recent + active] 윈도우에 걸쳐 RoPE가 올바른 시간적 위상을 포착할 수 있습니다.
forward
< source >( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorvideo_rotary_emb: tupleencoder_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_peRoPE를 포함합니다.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
< source >( num_layers: int )
각 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
< source >( )
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 슬롯입니다.
마지막 forward에서 포착한 post-RoPE K/V를 pop합니다.
마지막 forward에서 캡처된 pre-RoPE K/V를 pop합니다.
RoPE 적용 후 최근 히스토리 K/V를 반환하거나, 비어 있으면 None을 반환합니다.
pre-RoPE sink K/V를 반환합니다. 캡처되지 않았거나(또는 비어 있으면) None을 반환합니다.
현재 forward에서 생성된 post-RoPE K/V를 저장합니다.
현재 forward에서 생성된 pre-RoPE K/V를 저장합니다.
RoPE 적용 후 최근 히스토리의 K/V를 저장합니다.
RoPE 적용 전 sink K/V를 저장합니다.
Transformer2DModelOutput
클래스 diffusers.models.modeling_outputs.Transformer2DModelOutput
< source >( 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의 출력입니다.