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

Чанково-каузальный авто-регрессивный рефинер DiT, используемый как второй этап SanaWMPipeline, управляемый SanaWMLTX2Refiner.

Архитектурно он идентичен LTX2VideoTransformer3DModel — те же аргументы конфигурации, те же подмодули, те же имена параметров — поэтому опубликованный чекпоинт LTX-2 загружается в него без изменений. Прямой проход (forward pass) отличается:

  • обрабатывается только видеопоток (ветви кросс-внимания аудио и аудио/видео пропускаются),
  • self-attention выполняется с использованием явного скользящего окна KV-кэша (SanaWMRefinerKVCache), содержащего attention sink и недавнюю уточнённую историю, поэтому вычисления на блок ограничены, а общая стоимость уточнения линейно масштабируется с длиной видео,
  • вызывающая сторона предоставляет видео 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").

Chunk-causal autoregressive refiner transformer, используемый на этапе 2 SANA-WM.

Архитектурно идентичен LTX2VideoTransformer3DModel — те же аргументы конфигурации, те же подмодули, те же имена параметров — поэтому опубликованный чекпоинт LTX-2 загружается в него без изменений. Отличается только прямой проход:

  • обрабатывается только видеопоток (ветви кросс-внимания аудио и аудио/видео пропускаются),
  • self-attention выполняется с использованием явного скользящего окна KV-кэша (SanaWMRefinerKVCache), содержащего attention sink и недавнюю уточнённую историю, поэтому стоимость уточнения ограничена для каждого AR-блока и линейно масштабируется с длиной видео,
  • вызывающая сторона предоставляет видео RoPE, что позволяет каждому AR-окну сохранять абсолютный индекс каждого кадра в исходном видео (см. build_rotary_emb_for_absolute_positions()).

Остальные аргументы полностью соответствуют аргументам LTX2VideoTransformer3DModel. Аргументы и подмодули, связанные с аудио, сохраняются исключительно для того, чтобы аудио-веса в чекпоинте корректно восстанавливались; они не используются в прямом проходе рефинера.

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) — Частота кадров видео, определяющая масштабирование временного RoPE в LTX-2.

Возвращает

tuple[torch.Tensor, torch.Tensor]

кортеж RoPE (cos, sin).

Создайте видео-RoPE для явного списка абсолютных индексов латентных кадров.

SanaWMLTX2AudioVideoRotaryPosEmbed.prepare_video_coords предполагает непрерывный torch.arange(num_frames), что приемлемо для двунаправленного вывода. Уточнитель с AR-окном скользящего типа, напротив, должен сохранять абсолютный индекс каждого кадра в исходном видео, чтобы RoPE корректно учитывал временную фазу в окне [sink + recent + active].

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) — Патчифицированные латенты видео активного блока, с формой (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), уже масштабированный множителем self.config.timestep_scale_multiplier.
  • video_rotary_emb (tuple[torch.Tensor, torch.Tensor]) — (cos, sin) RoPE для абсолютных позиций кадров активного блока, как возвращаемые функцией build_rotary_emb_for_absolute_positions().
  • encoder_attention_mask (torch.Tensor, необязательный) — Маска внимания для текста с мультипликативным применением, форма (batch_size, text_seq_len).
  • kv_cache (SanaWMRefinerKVCache, optional) — скользящее окно KV-кэша, содержащее attention sink и недавнюю уточнённую историю для каждого слоя, а также общий sink_pe RoPE. Обязателен, если задан kv_cache_mode.
  • kv_cache_mode (str, необязательный) — Одно из значений:
    • "inject": внимание к [sink + history + current] K/V (шаги денойзинга).
    • "capture_pre_rope": без префикса; записывает K/V до применения RoPE для данного прохода в кэш (используется один раз для инициализации attention sink из исходных латентов первого этапа).
    • "inject_and_capture_post_rope": обращается к [sink + history + current] K/V и записывает post-RoPE K/V этого блока в кэш, чтобы его можно было добавить в историю.

    Если None, блок выполняет обычное полное само-внимание только по текущим токенам.

  • attention_kwargs (dict, необязательный) — Необязательные аргументы, передаваемые в обработку масштаба LoRA.
  • return_dict (bool, необязательный, по умолчанию True) — Определяет, следует ли возвращать Transformer2DModelOutput вместо обычного кортежа.

Возвращает

Transformer2DModelOutput или tuple

предсказанная скорость для активного блока, с формой (batch_size, num_video_tokens, out_channels).

Прямой проход только по видео через один блок AR.

SanaWMRefinerKVCache

класс diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVCache

< >

( num_layers: int )

Параметры

  • num_layers (int) — Количество блоков трансформера, для которых выделяется кэш на каждый слой.

Контейнер, содержащий по одному SanaWMRefinerKVLayerCache на каждый блок трансформера, а также общий sink RoPE.

Это реализует контракт KV-кэша rf_shifted_sink, с которым обучался рефинер второго этапа SANA-WM. Рефинирование является блочно-каузальным: block_size латентных кадров денойзируются за один раз, при этом внимание обращается к ограниченному окну [attention sink + recent history + active block] K/V.

  • sink_pe: кортеж (cos, sin) RoPE для sink-кадров, пересобираемый для каждого AR-окна при сдвиге sink_rope_offset, чтобы sink располагался непосредственно перед ограниченным рабочим кэшем. Общий для всех слоёв, поскольку RoPE не зависит от слоя.

SanaWMRefinerKVLayerCache

класс diffusers.models.transformers.transformer_sana_wm_refiner.SanaWMRefinerKVLayerCache

< >

( )

Кэш KV для каждого слоя для этапа 2 SANA-WM, использующего чанк-каузальный AR-рефинер.

Хранит две половины префикса скользящего окна, на которые обращается само-внимание уточнителя, а также слот для чтения обратно K/V, захваченных последним проходом. Все тензоры имеют размерность (batch_size, num_tokens, inner_dim) (то есть до разворачивания по головам), что соответствует раскладке, которую само-внимание уточнителя использует для конкатенации.

  • sink_k_pre / sink_v: pre-RoPE K/V кадров attention-sink, захваченные один раз из исходных латентов первого этапа. Они хранятся в формате pre-RoPE, чтобы каждое AR-окно могло повторно применить RoPE со своим сдвинутым смещением sink (SanaWMRefinerKVCache.sink_pe).
  • history_k / history_v: post-RoPE K/V для уже уточнённых последних кадров, готовые к конкатенации как есть.
  • captured_k_pre / captured_v_pre и captured_k_post / captured_v_post: слоты для чтения, заполняемые режимами захвата kv_cache_mode.

get_captured_post_rope

< >

( )

Снимите K/V после RoPE, полученные последним прямым проходом.

get_captured_pre_rope

< >

( )

Достаньте K/V до применения RoPE, полученные в последнем прямом проходе.

get_history

< >

( )

Возвращает K/V недавней истории после применения RoPE или None, если пусто.

get_sink

< >

( )

Возвращает sink K/V до применения RoPE или None, если он не был захвачен (или пуст).

store_captured_post_rope

< >

( key: Tensorvalue: Tensor )

Сохраните K/V после применения RoPE, полученные в текущем прямом проходе.

store_captured_pre_rope

< >

( key: Tensorvalue: Tensor )

Сохраните K/V до применения RoPE, полученные в текущем прямом проходе.

store_history

< >

( history_k: Tensorhistory_v: Tensor )

Сохраните K/V недавней истории после применения RoPE.

store_sink

< >

( sink_k_pre: Tensorsink_v: Tensor )

Сохраните K/V sink до применения RoPE.

Transformer2DModelOutput

класс diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( sample: torch.Tensor )

Параметры

  • sample (torch.Tensor формы (batch_size, num_channels, height, width) или (batch size, num_vector_embeds - 1, num_latent_pixels), если Transformer2DModel дискретный) — Выходные скрытые состояния, обусловленные входом encoder_hidden_states. Если модель дискретная, возвращает распределения вероятностей для латентных пикселей без шума.

Вывод Transformer2DModel.

Update on GitHub

←SanaVideoTransformer3DModel SanaWMTransformer3DModel→