Hugging Face's logo

Diffusers documentation

UNet2DConditionModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

UNet2DConditionModel

Модель UNet была первоначально представлена Ronneberger et al. для сегментации биомедицинских изображений, однако она также широко используется в 🤗 Diffusers, поскольку на выходе она генерирует изображения того же размера, что и входные. Это один из ключевых компонентов диффузионной системы, так как именно он обеспечивает непосредственный процесс диффузии. В 🤗 Diffusers существует несколько вариантов модели UNet, которые различаются количеством измерений и тем, является ли модель условной. Здесь представлена условная 2D-модель UNet.

Аннотация из статьи:

Существует широкое согласие в том, что успешное обучение глубоких сетей требует тысяч размеченных обучающих примеров. В данной статье мы представляем сеть и стратегию обучения, которые опираются на активное использование аугментации данных для более эффективного использования доступных размеченных примеров. Архитектура состоит из сжимающего пути для захвата контекста и симметричного расширяющего пути, обеспечивающего точную локализацию. Мы показываем, что такая сеть может быть обучена end-to-end по очень малому количеству изображений и превосходит предыдущий лучший метод (сверточная сеть со скользящим окном) на конкурсе ISBI по сегментации нейронных структур в электронных микроскопических стопках. Используя ту же сеть, обученную на изображениях, полученные с помощью световой микроскопии (фазовый контраст и DIC), мы с большим отрывом выиграли конкурс ISBI по отслеживанию клеток 2015 года в этих категориях. Кроме того, сеть работает быстро. Сегментация изображения 512x512 занимает менее секунды на современном GPU. Полная реализация (на основе Caffe) и обученные сети доступны по адресу http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net.

UNet2DConditionModel

класс diffusers.UNet2DConditionModel

< >

( sample_size: int | tuple[int, int] | None = Nonein_channels: int = 4out_channels: int = 4center_input_sample: bool = Falseflip_sin_to_cos: bool = Truefreq_shift: int = 0down_block_types: tuple[str, ...] = ('CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'DownBlock2D')mid_block_type: str | None = 'UNetMidBlock2DCrossAttn'up_block_types: tuple[str, ...] = ('UpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D')only_cross_attention: bool | tuple[bool] = Falseblock_out_channels: tuple[int, ...] = (320, 640, 1280, 1280)layers_per_block: int | tuple[int] = 2downsample_padding: int = 1mid_block_scale_factor: float = 1dropout: float = 0.0act_fn: str = 'silu'norm_num_groups: int | None = 32norm_eps: float = 1e-05cross_attention_dim: int | tuple[int] = 1280transformer_layers_per_block: int | tuple[int] | tuple[tuple] = 1reverse_transformer_layers_per_block: tuple[tuple[int]] | None = Noneencoder_hid_dim: int | None = Noneencoder_hid_dim_type: str | None = Noneattention_head_dim: int | tuple[int] = 8num_attention_heads: int | tuple[int] | None = Nonedual_cross_attention: bool = Falseuse_linear_projection: bool = Falseclass_embed_type: str | None = Noneaddition_embed_type: str | None = Noneaddition_time_embed_dim: int | None = Nonenum_class_embeds: int | None = Noneupcast_attention: bool = Falseresnet_time_scale_shift: str = 'default'resnet_skip_time_act: bool = Falseresnet_out_scale_factor: float = 1.0time_embedding_type: str = 'positional'time_embedding_dim: int | None = Nonetime_embedding_act_fn: str | None = Nonetimestep_post_act: str | None = Nonetime_cond_proj_dim: int | None = Noneconv_in_kernel: int = 3conv_out_kernel: int = 3projection_class_embeddings_input_dim: int | None = Noneattention_type: str = 'default'class_embeddings_concat: bool = Falsemid_block_only_cross_attention: bool | None = Nonecross_attention_norm: str | None = Noneaddition_embed_type_num_heads: int = 64 )

Параметры

  • sample_size (int или tuple[int, int], необязательный, по умолчанию None) — Высота и ширина входного/выходного образца.
  • in_channels (int, необязательный, по умолчанию 4) — Количество каналов во входном образце.
  • out_channels (int, необязательный, по умолчанию 4) — Количество каналов на выходе.
  • center_input_sample (bool, необязательный, по умолчанию False) — Центрировать ли входной образец.
  • flip_sin_to_cos (bool, необязательный, по умолчанию True) — Определяет, следует ли заменить sin на cos во временном эмбеддинге.
  • freq_shift (int, необязательный, по умолчанию 0) — Сдвиг частоты, применяемый к временному эмбеддингу.
  • down_block_types (tuple[str], необязательный, по умолчанию ("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")) — Кортеж блоков понижающего дискретизирования, которые следует использовать.
  • mid_block_type (str, необязательный, по умолчанию "UNetMidBlock2DCrossAttn") — Тип блока для средней части UNet. Может быть одним из UNetMidBlock2DCrossAttn, UNetMidBlock2D или UNetMidBlock2DSimpleCrossAttn. Если None, слой среднего блока пропускается.
  • up_block_types (tuple[str], необязательный, по умолчанию ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")) — Кортеж блоков апсемплинга, которые следует использовать.
  • only_cross_attention(bool or tuple[bool], optional, по умолчанию False) — Включать ли self-attention в базовых блоках трансформера, см. BasicTransformerBlock.
  • block_out_channels (tuple[int], необязательный, по умолчанию (320, 640, 1280, 1280)) — Кортеж количества выходных каналов для каждого блока.
  • layers_per_block (int, необязательный, по умолчанию 2) — Количество слоев в каждом блоке.
  • downsample_padding (int, необязательный, по умолчанию 1) — Размер отступов для свёртки при уменьшении разрешения.
  • mid_block_scale_factor (float, необязательный, по умолчанию 1.0) — Коэффициент масштабирования для среднего блока.
  • dropout (float, необязательный, по умолчанию 0.0) — Вероятность dropout, которую следует использовать.
  • act_fn (str, необязательный, по умолчанию "silu") — Функция активации, которую следует использовать.
  • norm_num_groups (int, необязательный, по умолчанию 32) — количество групп, используемых для нормализации. Если None, слои нормализации и активации пропускаются на этапе постобработки.
  • norm_eps (float, необязательный, по умолчанию 1e-5) — значение эпсилон, используемое для нормализации.
  • cross_attention_dim (int или tuple[int], необязательный, по умолчанию 1280) — Размерность признаков кросс-внимания.
  • transformer_layers_per_block (int, tuple[int] или tuple[tuple], необязательный, по умолчанию 1) — Количество блоков трансформера типа BasicTransformerBlock. Актуально только для CrossAttnDownBlock2D, CrossAttnUpBlock2D, UNetMidBlock2DCrossAttn.
  • reverse_transformer_layers_per_block — (tuple[tuple], необязательный, по умолчанию None): Количество блоков трансформера типа BasicTransformerBlock в блоках апсемплинга U-Net. Актуально только если transformer_layers_per_block имеет тип tuple[tuple] и для CrossAttnDownBlock2D, CrossAttnUpBlock2D, UNetMidBlock2DCrossAttn.
  • encoder_hid_dim (int, необязательный, по умолчанию None) — Если encoder_hid_dim_type определён, encoder_hidden_states будут проецированы из размерности encoder_hid_dim в cross_attention_dim.
  • encoder_hid_dim_type (str, необязательный, по умолчанию None) — Если указано, encoder_hidden_states и, возможно, другие эмбеддинги проецируются в эмбеддинги текста размерности cross_attention в соответствии с encoder_hid_dim_type.
  • attention_head_dim (int, необязательный, по умолчанию 8) — размерность головок внимания.
  • num_attention_heads (int, необязательный) — Количество головок внимания. Если не задано, по умолчанию равно attention_head_dim
  • resnet_time_scale_shift (str, необязательный, по умолчанию "default") — конфигурация сдвига масштаба времени для блоков ResNet (см. ResnetBlock2D). Выберите default или scale_shift.
  • class_embed_type (str, необязательный, по умолчанию None) — Тип классовой эмбеддинга, который в итоге суммируется с временными эмбеддингами. Выберите одно из значений: None, "timestep", "identity", "projection" или "simple_projection".
  • addition_embed_type (str, необязательный, по умолчанию None) — Настраивает необязательное вложение, которое будет складываться с временными вложениями. Выберите None или “text”. Значение “text” использует слой TextTimeEmbedding.
  • addition_time_embed_dim — (int, необязательный, по умолчанию None): Размерность для эмбеддингов временных шагов.
  • num_class_embeds (int, необязательный, по умолчанию None) — Входное измерение обучаемой матрицы эмбеддингов, проецируемой в time_embed_dim, при выполнении условного кодирования по классу, когда class_embed_type равно None.
  • time_embedding_type (str, необязательный, по умолчанию positional) — Тип позиционного эмбеддинга, используемого для временных шагов. Выберите positional или fourier.
  • time_embedding_dim (int, необязательный, по умолчанию None) — Необязательное переопределение размерности проекции временного эмбеддинга.
  • time_embedding_act_fn (str, необязательный, по умолчанию None) — Необязательная функция активации, применяемая один раз к временным эмбеддингам перед их передачей в остальные слои UNet. Выберите одну из silu, mish, gelu или swish.
  • timestep_post_act (str, необязательный, по умолчанию None) — Вторая функция активации, используемая во вложении временного шага. Выберите из silu, mish и gelu.
  • time_cond_proj_dim (int, необязательный, по умолчанию None) — Размерность слоя cond_proj во встраивании временного шага.
  • conv_in_kernel (int, необязательный, по умолчанию 3) — размер ядра слоя conv_in.
  • conv_out_kernel (int, необязательный, по умолчанию 3) — размер ядра слоя conv_out.
  • projection_class_embeddings_input_dim (int, необязательный) — размерность входа class_labels, когда class_embed_type="projection". Обязателен, если class_embed_type="projection".
  • class_embeddings_concat (bool, необязательный, по умолчанию False) — Соединять ли временные эмбеддинги с эмбеддингами классов.
  • mid_block_only_cross_attention (bool, необязательный, по умолчанию None) — Использовать ли cross attention с mid block при использовании UNetMidBlock2DSimpleCrossAttn. Если only_cross_attention задано как одно булево значение, а mid_block_only_cross_attention равно None, то значение only_cross_attention используется в качестве значения для mid_block_only_cross_attention. В противном случае по умолчанию False.

Условная 2D-модель UNet, которая принимает зашумленный образец, условное состояние и шаг времени и возвращает выход с формой образца.

Эта модель наследуется от ModelMixin. См. документацию суперкласса для его универсальных методов, реализованных для всех моделей (например, загрузка или сохранение).

disable_freeu

< >

( )

Отключает механизм FreeU.

enable_freeu

< >

( s1: floats2: floatb1: floatb2: float )

Параметры

  • s1 (float) — Коэффициент масштабирования для этапа 1, ослабляющий вклад признаков пропуска. Это делается для смягчения «эффекта чрезмерного сглаживания» в процессе улучшенного шумоподавления.
  • s2 (float) — Коэффициент масштабирования для этапа 2, ослабляющий вклад признаков пропуска. Это делается для смягчения «эффекта чрезмерного сглаживания» в процессе улучшенного шумоподавления.
  • b1 (float) — Коэффициент масштабирования для этапа 1, усиливающий вклад признаков базовой модели.
  • b2 (float) — Коэффициент масштабирования для этапа 2, усиливающий вклад признаков базовой модели.

Включает механизм FreeU из https://huggingface.co/papers/2309.11497.

Суффиксы после коэффициентов масштабирования обозначают блоки стадий, в которых они применяются.

Список комбинаций значений, которые хорошо работают для различных пайплайнов, таких как Stable Diffusion v1, v2 и Stable Diffusion XL, см. в официальном репозитории.

forward

< >

( sample: Tensortimestep: typing.Union[torch.Tensor, float, int]encoder_hidden_states: Tensorclass_labels: typing.Optional[torch.Tensor] = Nonetimestep_cond: typing.Optional[torch.Tensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Nonecross_attention_kwargs: dict[str, typing.Any] | None = Noneadded_cond_kwargs: dict[str, torch.Tensor] | None = Nonedown_block_additional_residuals: tuple[torch.Tensor] | None = Nonemid_block_additional_residual: typing.Optional[torch.Tensor] = Nonedown_intrablock_additional_residuals: tuple[torch.Tensor] | None = Noneencoder_attention_mask: typing.Optional[torch.Tensor] = Nonereturn_dict: bool = True ) → UNet2DConditionOutput or tuple

Параметры

  • sample (torch.Tensor) — Шумовой входной тензор со следующей формой (batch, channel, height, width).
  • timestep (torch.Tensor или float или int) — Количество шагов для денойзинга входных данных.
  • encoder_hidden_states (torch.Tensor) — Скрытые состояния энкодера с формой (batch, sequence_length, feature_dim).
  • class_labels (torch.Tensor, необязательный, по умолчанию None) — Необязательные метки классов для условного формирования. Их эмбеддинги будут сложены с эмбеддингами временного шага.
  • timestep_cond — (torch.Tensor, необязательный, по умолчанию None): Условные эмбеддинги для временного шага. Если переданы, эмбеддинги будут сложены с сэмплами, прошедшими через слой self.time_embedding, чтобы получить эмбеддинги временного шага.
  • attention_mask (torch.Tensor, необязательный, по умолчанию None) — Маска внимания формы (batch, key_tokens) применяется к encoder_hidden_states. Если значение равно 1, маска сохраняется, иначе, если равно 0, — отбрасывается. Маска преобразуется в смещение (bias), которое добавляет большие отрицательные значения к оценкам внимания, соответствующим токенам, помеченным как «отбросить».
  • cross_attention_kwargs (dict, необязательный) — Словарь kwargs, который при указании передается в AttentionProcessor, определенный в self.processor в модуле diffusers.models.attention_processor.
  • added_cond_kwargs — (dict, необязательный): Словарь kwargs, содержащий дополнительные эмбеддинги, которые при указании добавляются к эмбеддингам, передаваемым в блоки UNet.
  • down_block_additional_residuals — (tuple из torch.Tensor, необязательный): Кортеж тензоров, которые при указании добавляются к остаточным связям блоков down в U-Net.
  • mid_block_additional_residual — (torch.Tensor, необязательный): Тензор, который при указании добавляется к остаточному соединению среднего блока unet.
  • down_intrablock_additional_residuals (tuple из torch.Tensor, необязательный) — дополнительные остаточные связи, добавляемые внутри блоков понижения разрешения UNet, например, из боковых моделей T2I-Adapter
  • encoder_attention_mask (torch.Tensor) — Маска кросс-внимания формы (batch, sequence_length), применяемая к encoder_hidden_states. Если True, маска сохраняется, иначе, если False, она отбрасывается. Маска преобразуется в смещение (bias), которое добавляет большие отрицательные значения к оценкам внимания, соответствующим токенам «отбрасывания».
  • return_dict (bool, необязательный, по умолчанию True) — Определяет, следует ли возвращать UNet2DConditionOutput вместо обычного кортежа.

Возвращает

UNet2DConditionOutput или tuple

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

Метод forward модели UNet2DConditionModel.

fuse_qkv_projections

< >

( )

Включает объединённые проекции QKV. Для модулей само-внимания все матрицы проекций (т. е. запрос, ключ, значение) объединяются. Для модулей кросс-внимания объединяются матрицы проекций ключа и значения.

> Этот API 🧪 экспериментальный.

set_attention_slice

< >

( slice_size: str | int | list[int] = 'auto' )

Параметры

  • slice_size (str или int или list(int), необязательный, по умолчанию "auto") — Если указано "auto", вход для головок внимания делится пополам, и внимание вычисляется в два этапа. Если указано "max", экономится максимальное количество памяти за счёт обработки одного среза за раз. Если предоставлено число, используется столько срезов, сколько получается из attention_head_dim // slice_size. В этом случае attention_head_dim должно быть кратным slice_size.

Включить вычисление внимания с нарезкой.

Когда эта опция включена, модуль внимания разбивает входной тензор на фрагменты, чтобы вычислять внимание в несколько шагов. Это полезно для экономии памяти ценой небольшого снижения скорости.

set_default_attn_processor

< >

( )

Отключает пользовательские обработчики внимания и устанавливает реализацию внимания по умолчанию.

unfuse_qkv_projections

< >

( )

Отключает объединенную проекцию QKV, если она включена.

> Этот API 🧪 экспериментальный.

UNet2DConditionOutput

класс diffusers.models.unets.unet_2d_condition.UNet2DConditionOutput

< >

( sample: Tensor = None )

Параметры

  • sample (torch.Tensor формы (batch_size, num_channels, height, width)) — Скрытые состояния, вычисленные с учетом входных encoder_hidden_states. Выход последнего слоя модели.

Вывод UNet2DConditionModel.

Update on GitHub

←UNet1DModel UNet2DModel→