Hugging Face's logo

Diffusers documentation

SanaWMLTX2RefinerTransformer3DModel

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

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 = (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: listheight: 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: 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) — Патчифицированные латенты видео активного блока, с формой (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→