Diffusers documentation

QwenImage21Transformer2DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

QwenImage21Transformer2DModel

Однопоточный трансформер, используемый в Qwen-Image 2.1. Латенты текста и изображения разделяют одну последовательность, а единая общая проекция modulation питает каждый блок.

Два поведения отличают трансформеры QwenImage 2.1 от более ранних версий:

  • Блочное каузальное внимание — внимание следует правилу (q_idx >= kv_idx) or same_image_block, поэтому совместная последовательность является каузальной, а каждый блок изображений остаётся внутренне двунаправленным. QwenImage21AttnProcessor реализует это как один вызов внимания на каждый префиксный сегмент и является вариантом по умолчанию. QwenImage21FlexAttnProcessor реализует это как единый вызов flex_attention, управляемый BlockMask, что быстрее после компиляции модели. Оба дают одинаковые результаты.
  • causal_condition — токены текста и условного изображения модулируются из t = 0, а не из сэмплированного шага времени. Их активации не зависят от шага денойзинга, поэтому ключи и значения этого префикса можно кэшировать между шагами с помощью аргумента kv_cache.

Загрузите его с помощью:

import torch
from diffusers import QwenImage21Transformer2DModel

transformer = QwenImage21Transformer2DModel.from_pretrained(
    "Qwen/Qwen-Image-2.1", subfolder="transformer", dtype=torch.bfloat16
)

QwenImage21Transformer2DModel

класс diffusers.QwenImage21Transformer2DModel

< >

( patch_size: int = 1in_channels: int = 64out_channels: int | None = 64num_layers: int = 32attention_head_dim: int = 128num_attention_heads: int = 32context_in_dim: int = 4096mlp_ratio: int = 3axes_dims_rope: tuple[int, int, int] = (16, 56, 56)eps: float = 1e-06causal_condition: bool = True )

Параметры

  • patch_size (int, по умолчанию 1) — Длина стороны латентного патча, свёрнутого в размерность каналов. Версия 2.1 использует латенты без патчинга.
  • in_channels (int, по умолчанию 64) — Каналы латентного представления на входе.
  • out_channels (int, необязательный, по умолчанию 64) — Количество каналов в латентном пространстве на выходе. Если не указано, используется значение in_channels.
  • num_layers (int, по умолчанию 32) — Количество однопоточных блоков.
  • attention_head_dim (int, по умолчанию 128) — Количество каналов на голову внимания.
  • num_attention_heads (int, по умолчанию 32) — Количество головок внимания.
  • context_in_dim (int, по умолчанию 4096) — Размерность каналов encoder_hidden_states.
  • mlp_ratio (int, по умолчанию 3) — Коэффициент расширения в полносвязном слое.
  • axes_dims_rope (tuple[int], по умолчанию (16, 56, 56)) — Размеры ротации для осей кадра, высоты и ширины.
  • eps (float, по умолчанию 1e-6) — Значение эпсилон для слоёв нормализации.
  • causal_condition (bool, по умолчанию True) — Модулирует токены текста и условного изображения начиная с t = 0. Обязательно для кэширования KV.

Однопоточный Transformer, используемый в Qwen-Image 2.1.

Латенты текста и изображения разделяют одну последовательность: токены изображения-условия подставляются в текстовый поток на те позиции, которые для них зарезервировал vision-language-энкодер, а токены целевого изображения добавляются в конец. Единая общая проекция modulation питает каждый блок, поэтому сами блоки не содержат собственных параметров модуляции.

Две особенности выделяют версию 2.1:

  • Блочное каузальное внимание — внимание следует правилу (q_idx >= kv_idx) or same_image_block, поэтому последовательность является каузальной, в то время как каждый блок изображения остаётся внутренне двунаправленным. Бэкенд внимания flex обеспечивает эффективное внимание за один проход; без него модель использует точный многопроходный prefill на основе SDPA, обрабатывая каждый блок отдельно.
  • causal_condition — токены текста и условного изображения модулируются из t = 0, а не из сэмплированного временного шага, что также делает их активации независимыми от временного шага и позволяет кэшировать их между шагами денойзинга.

build_token_metadata

< >

( image_pad_mask: Tensorimg_shapes: list[tuple[int, int, int]] ) → tuple[torch.Tensor, torch.Tensor]

Параметры

  • image_pad_mask (torch.Tensor) — (seq_len,) bool, True в позициях токенов изображений.
  • img_shapes (list[tuple[int, int, int]]) — Для каждого изображения (frame, height, width) в латентных токенах, сначала изображения-условия, затем целевое изображение.

Возвращает

tuple[torch.Tensor, torch.Tensor]

image_ids (seq_len,) со значением -1 в позициях текста и уникальным id для каждого блока изображения, а также target_token_mask (seq_len,), отмечающая токены целевого изображения.

Пометьте каждый токен объединённой последовательности блоком изображения, которому он принадлежит.

Границы блоков определяются количеством токенов в img_shapes, а не последовательностями True в image_pad_mask: два условных изображения, которые случайно оказываются рядом без текста между ними, образуют одну последовательность, но должны оставаться отдельными блоками, иначе они будут обращаться друг к другу двунаправленно.

forward

< >

( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorimg_shapes: list[list[tuple[int, int, int]]]img_mask: Tensorencoder_hidden_states_mask: typing.Optional[torch.Tensor] = Noneattention_kwargs: dict[str, typing.Any] | None = Nonekv_cache: diffusers.models.transformers.transformer_qwenimage21.QwenImage21KVCache | None = Nonekv_cache_mode: str | None = Nonereturn_dict: bool = True ) → Transformer2DModelOutput or tuple

Параметры

  • hidden_states (torch.Tensor формы (batch_size, image_sequence_length, in_channels)) — Упакованные латенты, где сначала идут условные изображения, а целевое изображение — последним.
  • encoder_hidden_states (torch.Tensor формы (batch_size, text_sequence_length, context_in_dim)) — Векторные представления текста из мультимодального энкодера.
  • timestep (torch.Tensor) — Текущий шаг денойзинга, масштабированный до [0, 1].
  • img_shapes (list[list[tuple[int, int, int]]]) — Список по образцам, содержащий (frame, height, width) в латентных токенах, где сначала идут изображения-условия, а последнее — целевое изображение. Все образцы должны иметь одинаковую структуру.
  • img_mask (torch.Tensor формы (batch_size, vlm_sequence_length)) — True в слотах изображений кодировщика vision-language, каждый из которых соответствует группе 2x2 латентных токенов.
  • encoder_hidden_states_mask (torch.Tensor, необязательный) — (batch_size, text_sequence_length) bool, отмечающий валидные токены текста. Позиции с заполнением исключаются из внимания.
  • kv_cache (QwenImage21KVCache, необязательный) — Контейнер кэша. Передавайте вместе с kv_cache_mode, чтобы включить префиксовый KV-кэш.
  • kv_cache_mode (str, необязательный) — "extract" для предварительного заполнения кэша (первый шаг денойзинга), "cached" для декодирования из него (последующие шаги). Требуется causal_condition=True.
  • attention_kwargs (dict, необязательный) — Передаётся в обработчики внимания и содержит scale для слоёв LoRA.
  • return_dict (bool, необязательный, по умолчанию True) — Определяет, следует ли возвращать Transformer2DModelOutput вместо обычного кортежа.

Возвращает

Transformer2DModelOutput или tuple

Transformer2DModelOutput, если return_dict равно True, иначе обычный tuple, первым элементом которого являются денойзанные латенты.

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