Hugging Face's logo

Diffusers documentation

AutoencoderDC

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

AutoencoderDC

Модель 2D-автоэнкодера, используемая в SANA и представленная в DCAE авторами Junyu Chen*, Han Cai*, Junsong Chen, Enze Xie, Shang Yang, Haotian Tang, Muyang Li, Yao Lu, Song Han из MIT HAN Lab.

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

Мы представляем Deep Compression Autoencoder (DC-AE) — новое семейство моделей автоэнкодеров для ускорения высокодетализированных диффузионных моделей. Существующие модели автоэнкодеров демонстрируют впечатляющие результаты при умеренном пространственном коэффициенте сжатия (например, 8x), но не способны поддерживать приемлемое качество реконструкции при высоких пространственных коэффициентах сжатия (например, 64x). Мы решаем эту задачу, вводя две ключевые техники: (1) Residual Autoencoding, в рамках которой мы проектируем модели так, чтобы они обучались на остатках, основанных на признаках, преобразованных из пространства в каналы, что снижает сложность оптимизации автоэнкодеров с высоким пространственным сжатием; (2) Decoupled High-Resolution Adaptation — эффективная стратегия обучения в три этапа с разделением процессов, предназначенная для смягчения штрафа за обобщение у автоэнкодеров с высоким пространственным сжатием. Благодаря этим решениям мы повышаем пространственный коэффициент сжатия автоэнкодера до 128, сохраняя при этом качество реконструкции. Применяя наш DC-AE к латентным диффузионным моделям, мы достигаем значительного ускорения без потери точности. Например, на ImageNet 512x512 наш DC-AE обеспечивает ускорение инференса в 19,1 раза и ускорение обучения в 17,9 раз на GPU H100 для UViT-H, при этом достигая лучшего FID по сравнению с широко используемым автоэнкодером SD-VAE-f8. Наш код доступен по адресу этот URL.

Следующие модели DCAE выпущены и поддерживаются в Diffusers.

Формат DiffusersИсходный формат
mit-han-lab/dc-ae-f32c32-sana-1.0-diffusersmit-han-lab/dc-ae-f32c32-sana-1.0
mit-han-lab/dc-ae-f32c32-in-1.0-diffusersmit-han-lab/dc-ae-f32c32-in-1.0
mit-han-lab/dc-ae-f32c32-mix-1.0-diffusersmit-han-lab/dc-ae-f32c32-mix-1.0
mit-han-lab/dc-ae-f64c128-in-1.0-diffusersmit-han-lab/dc-ae-f64c128-in-1.0
mit-han-lab/dc-ae-f64c128-mix-1.0-diffusersmit-han-lab/dc-ae-f64c128-mix-1.0
mit-han-lab/dc-ae-f128c512-in-1.0-diffusersmit-han-lab/dc-ae-f128c512-in-1.0
mit-han-lab/dc-ae-f128c512-mix-1.0-diffusersmit-han-lab/dc-ae-f128c512-mix-1.0

Эта модель была добавлена пользователем lawrence-cj.

Загрузите модель в формате Diffusers с помощью from_pretrained().

from diffusers import AutoencoderDC

ae = AutoencoderDC.from_pretrained("mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers", dtype=torch.float32).to("cuda")  # or "mps", "xpu", "cpu"

Загрузка модели в Diffusers с помощью from_single_file

from difusers import AutoencoderDC

ckpt_path = "https://huggingface.co/mit-han-lab/dc-ae-f32c32-sana-1.0/blob/main/model.safetensors"
model = AutoencoderDC.from_single_file(ckpt_path)

Модель AutoencoderDC имеет варианты контрольных точек в формате одного файла in и mix, которые имеют совпадающие ключи контрольных точек, но используют разные коэффициенты масштабирования. Diffusers не может автоматически определить правильный файл конфигурации для использования с моделью на основе только контрольной точки и по умолчанию настраивает модель, используя файл конфигурации варианта mix. Чтобы переопределить автоматически определенный конфиг, используйте аргумент config при загрузке из одного файла с контрольными точками варианта in.

from diffusers import AutoencoderDC

ckpt_path = "https://huggingface.co/mit-han-lab/dc-ae-f128c512-in-1.0/blob/main/model.safetensors"
model = AutoencoderDC.from_single_file(ckpt_path, config="mit-han-lab/dc-ae-f128c512-in-1.0-diffusers")

AutoencoderDC

класс diffusers.AutoencoderDC

< >

( in_channels: int = 3latent_channels: int = 32attention_head_dim: int = 32encoder_block_types: str | tuple[str] = 'ResBlock'decoder_block_types: str | tuple[str] = 'ResBlock'encoder_block_out_channels: tuple = (128, 256, 512, 512, 1024, 1024)decoder_block_out_channels: tuple = (128, 256, 512, 512, 1024, 1024)encoder_layers_per_block: tuple = (2, 2, 2, 3, 3, 3)decoder_layers_per_block: tuple = (3, 3, 3, 3, 3, 3)encoder_qkv_multiscales: tuple = ((), (), (), (5,), (5,), (5,))decoder_qkv_multiscales: tuple = ((), (), (), (5,), (5,), (5,))upsample_block_type: str = 'pixel_shuffle'downsample_block_type: str = 'pixel_unshuffle'decoder_norm_types: str | tuple[str] = 'rms_norm'decoder_act_fns: str | tuple[str] = 'silu'encoder_out_shortcut: bool = Truedecoder_in_shortcut: bool = Truedecoder_conv_act_fn: str = 'relu'scaling_factor: float = 1.0 )

Параметры

  • in_channels (int, по умолчанию 3) — Количество входных каналов в образцах.
  • latent_channels (int, по умолчанию 32) — Количество каналов в представлении латентного пространства.
  • encoder_block_types (str | tuple[str], по умолчанию "ResBlock") — Тип(ы) блоков, используемых в кодере.
  • decoder_block_types (str | tuple[str], по умолчанию "ResBlock") — Тип(ы) блоков, используемых в декодере.
  • encoder_block_out_channels (tuple[int, ...], по умолчанию (128, 256, 512, 512, 1024, 1024)) — Количество выходных каналов для каждого блока в энкодере.
  • decoder_block_out_channels (tuple[int, ...], по умолчанию (128, 256, 512, 512, 1024, 1024)) — Количество выходных каналов для каждого блока в декодере.
  • encoder_layers_per_block (tuple[int], по умолчанию (2, 2, 2, 3, 3, 3)) — Количество слоев в каждом блоке энкодера.
  • decoder_layers_per_block (tuple[int], по умолчанию (3, 3, 3, 3, 3, 3)) — Количество слоев в каждом блоке декодера.
  • encoder_qkv_multiscales (tuple[tuple[int, ...], ...], по умолчанию ((), (), (), (5,), (5,), (5,))) — Масштабируемые конфигурации для преобразований QKV (query-key-value) в энкодере.
  • decoder_qkv_multiscales (tuple[tuple[int, ...], ...], по умолчанию ((), (), (), (5,), (5,), (5,))) — Масштабируемые конфигурации для преобразований QKV (query-key-value) декодера.
  • upsample_block_type (str, по умолчанию "pixel_shuffle") — Тип блока, используемого для апсемплинга в декодере.
  • downsample_block_type (str, по умолчанию "pixel_unshuffle") — Тип блока, используемого для понижения разрешения в кодировщике.
  • decoder_norm_types (str | tuple[str], по умолчанию "rms_norm") — Тип(ы) нормализации, используемые в декодере.
  • decoder_act_fns (str | tuple[str], по умолчанию "silu") — Функция(и) активации, используемые в декодере.
  • encoder_out_shortcut (bool, по умолчанию True) — Использовать ли шорткат в конце энкодера.
  • decoder_in_shortcut (bool, по умолчанию True) — Использовать ли shortcut в начале декодера.
  • decoder_conv_act_fn (str, по умолчанию "relu") — Функция активации, используемая в конце декодера.
  • scaling_factor (float, по умолчанию 1.0) — Взаимно-множительная величина среднеквадратичного корня латентных признаков. Используется для масштабирования латентного пространства до единичной дисперсии при обучении диффузионной модели. Перед передачей в диффузионную модель латенты масштабируются по формуле z = z * scaling_factor. При декодировании латенты возвращаются к исходному масштабу по формуле: z = 1 / scaling_factor * z.

Модель автоэнкодера, представленная в DCAE и используемая в SANA.

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

encode

< >

( x: Tensorreturn_dict: bool = True )

Параметры

  • x (torch.Tensor) — Входная партия изображений.
  • return_dict (bool, по умолчанию True) — Определяет, следует ли возвращать ~models.vae.EncoderOutput вместо обычного кортежа.

Кодирует пакет изображений в латенты.

decode

< >

( z: Tensorreturn_dict: bool = True ) → ~models.vae.DecoderOutput or tuple

Параметры

  • z (torch.Tensor) — Входная партия латентных векторов.
  • return_dict (bool, по умолчанию True) — Определяет, следует ли возвращать ~models.vae.DecoderOutput вместо обычного кортежа.

Возвращает

~models.vae.DecoderOutput или tuple

Если return_dict равно True, возвращается ~models.vae.DecoderOutput, иначе возвращается обычный tuple.

Декодировать пакет изображений.

enable_tiling

< >

( tile_sample_min_height: int | None = Nonetile_sample_min_width: int | None = Nonetile_sample_stride_height: float | None = Nonetile_sample_stride_width: float | None = None )

Параметры

  • tile_sample_min_height (int, необязательный) — Минимальная высота, необходимая для того, чтобы образец был разделён на плитки по размерности высоты.
  • tile_sample_min_width (int, необязательный) — Минимальная ширина, необходимая для того, чтобы образец был разделён на плитки по размерности ширины.
  • tile_sample_stride_height (int, необязательный) — Минимальное перекрытие между двумя последовательными вертикальными тайлами. Это обеспечивает отсутствие артефактов тайлинга вдоль вертикальной оси.
  • tile_sample_stride_width (int, необязательный) — Шаг между двумя соседними горизонтальными тайлами. Это обеспечивает отсутствие артефактов тайлинга вдоль ширины.

Включить тайловое декодирование AE. При включении этого параметра AE разбивает входной тензор на тайлы для вычисления декодирования и кодирования в несколько шагов. Это полезно для экономии значительного объема памяти и для обработки изображений большего размера.

forward

< >

( sample: Tensorreturn_dict: bool = True ) → ~models.vae.DecoderOutput or tuple

Параметры

  • sample (torch.Tensor) — Входной образец.
  • return_dict (bool, необязательный, по умолчанию True) — Определяет, следует ли возвращать DecoderOutput вместо обычного кортежа.

Возвращает

~models.vae.DecoderOutput или tuple

Если return_dict равно True, возвращается ~models.vae.DecoderOutput, в противном случае возвращается обычный tuple.

DecoderOutput

класс diffusers.models.autoencoders.vae.DecoderOutput

< >

( sample: Tensorcommit_loss: typing.Optional[torch.FloatTensor] = None )

Параметры

  • sample (torch.Tensor формы (batch_size, num_channels, height, width)) — Декодированный вывод из последнего слоя модели.

Результат метода декодирования.

Update on GitHub

←AsymmetricAutoencoderKL AutoencoderKL→