UNet3DConditionModel

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

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

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

UNet3DConditionModel

класс diffusers.UNet3DConditionModel

< source >

( sample_size: int | None = Nonein_channels: int = 4out_channels: int = 4down_block_types: tuple = ('CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'DownBlock3D')up_block_types: tuple = ('UpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D')block_out_channels: tuple = (320, 640, 1280, 1280)layers_per_block: int = 2downsample_padding: int = 1mid_block_scale_factor: float = 1act_fn: str = 'silu'norm_num_groups: int | None = 32norm_eps: float = 1e-05cross_attention_dim: int = 1024attention_head_dim: int | tuple[int] = 64num_attention_heads: int | tuple[int] | None = Nonetime_cond_proj_dim: int | None = None )

Параметры

sample_size (int или tuple[int, int], необязательный, по умолчанию None) — Высота и ширина входного/выходного образца.
in_channels (int, необязательный, по умолчанию 4) — количество каналов во входном образце.
out_channels (int, необязательный, по умолчанию 4) — количество каналов на выходе.
down_block_types (tuple[str], необязательный, по умолчанию ("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D")) — Кортеж блоков понижающего дискретизирования, которые следует использовать.
up_block_types (tuple[str], необязательный, по умолчанию ("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D")) — Кортеж блоков апсемплинга, которые следует использовать.
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) — Коэффициент масштабирования для среднего блока.
act_fn (str, необязательный, по умолчанию "silu") — Функция активации, которую следует использовать.
norm_num_groups (int, необязательный, по умолчанию 32) — количество групп, используемых для нормализации. Если None, слои нормализации и активации пропускаются на этапе постобработки.
norm_eps (float, необязательный, по умолчанию 1e-5) — значение эпсилон, используемое для нормализации.
cross_attention_dim (int, необязательный, по умолчанию 1024) — размерность признаков кросс-внимания.
attention_head_dim (int, необязательный, по умолчанию 64) — размерность головок внимания.
num_attention_heads (int, необязательный) — Количество головок внимания.
time_cond_proj_dim (int, необязательный, по умолчанию None) — Размерность слоя cond_proj во встраивании временного шага.

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

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

disable_freeu

< source >

( )

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

enable_forward_chunking

< source >

( chunk_size: int | None = Nonedim: int = 0 )

Параметры

chunk_size (int, необязательный) — Размер блока для полносвязных слоёв. Если не указан, полносвязный слой будет выполняться отдельно для каждого тензора с dim=dim.
dim (int, необязательный, по умолчанию 0) — Размерность, по которой вычисления в блоке прямой передачи (feed-forward) должны быть разбиты на части. Выберите dim=0 (размер батча) или dim=1 (длина последовательности).

Устанавливает процессор внимания для использования feed forward chunking.

enable_freeu

< source >

( s1s2b1b2 )

Параметры

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

< source >

( 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 = Nonedown_block_additional_residuals: tuple[torch.Tensor] | None = Nonemid_block_additional_residual: typing.Optional[torch.Tensor] = Nonereturn_dict: bool = True ) → UNet3DConditionOutput or tuple

Параметры

sample (torch.Tensor) — Шумовой входной тензор со следующей формой (batch, num_channels, num_frames, 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.
down_block_additional_residuals — (tuple из torch.Tensor, необязательный): Кортеж тензоров, которые при указании добавляются к остаточным связям блоков down в U-Net.
mid_block_additional_residual — (torch.Tensor, необязательный): Тензор, который при указании добавляется к остаточному соединению среднего блока unet.
return_dict (bool, необязательный, по умолчанию True) — Определяет, следует ли возвращать UNet3DConditionOutput вместо обычного кортежа.
cross_attention_kwargs (dict, необязательный) — Словарь kwargs, который при указании передается в AttnProcessor.

Возвращает

UNet3DConditionOutput или tuple

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

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

fuse_qkv_projections

< source >

( )

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

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

set_attention_slice

< source >

( slice_size: str | int | list[int] )

Параметры

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

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

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

set_default_attn_processor

< source >

( )

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

unfuse_qkv_projections

< source >

( )

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

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

UNet3DConditionOutput

класс diffusers.models.unets.unet_3d_condition.UNet3DConditionOutput

< source >

( sample: Tensor )

Параметры

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

Вывод UNet3DConditionModel.

Update on GitHub