AutoencoderDC

El modelo de Autoencoder 2D utilizado en SANA e introducido en DCAE por los autores Junyu Chen*, Han Cai*, Junsong Chen, Enze Xie, Shang Yang, Haotian Tang, Muyang Li, Yao Lu, Song Han del MIT HAN Lab.

El resumen del artículo es:

Presentamos Deep Compression Autoencoder (DC-AE), una nueva familia de modelos autoencoder para acelerar los modelos de difusión de alta resolución. Los modelos autoencoder existentes han demostrado resultados impresionantes con una relación de compresión espacial moderada (p. ej., 8x), pero no logran mantener una precisión de reconstrucción satisfactoria para relaciones de compresión espacial altas (p. ej., 64x). Abordamos este desafío introduciendo dos técnicas clave: (1) Residual Autoencoding, donde diseñamos nuestros modelos para aprender residuos basados en las características transformadas de espacio a canal, con el fin de aliviar la dificultad de optimización de los autoencoders de alta compresión espacial; (2) Decoupled High-Resolution Adaptation, una estrategia de entrenamiento desacoplada y eficiente en tres fases para mitigar la penalización de generalización de los autoencoders de alta compresión espacial. Con estos diseños, mejoramos la relación de compresión espacial del autoencoder hasta 128 manteniendo la calidad de reconstrucción. Al aplicar nuestro DC-AE a los modelos de difusión latente, logramos una aceleración significativa sin pérdida de precisión. Por ejemplo, en ImageNet 512x512, nuestro DC-AE proporciona una aceleración de inferencia de 19.1x y una aceleración de entrenamiento de 17.9x en GPU H100 para UViT-H, logrando al mismo tiempo un FID mejor en comparación con el autoencoder SD-VAE-f8 ampliamente utilizado. Nuestro código está disponible en this https URL.

Los siguientes modelos de DCAE se han publicado y están soportados en Diffusers.

Formato de Diffusers	Formato original
`mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers`	`mit-han-lab/dc-ae-f32c32-sana-1.0`
`mit-han-lab/dc-ae-f32c32-in-1.0-diffusers`	`mit-han-lab/dc-ae-f32c32-in-1.0`
`mit-han-lab/dc-ae-f32c32-mix-1.0-diffusers`	`mit-han-lab/dc-ae-f32c32-mix-1.0`
`mit-han-lab/dc-ae-f64c128-in-1.0-diffusers`	`mit-han-lab/dc-ae-f64c128-in-1.0`
`mit-han-lab/dc-ae-f64c128-mix-1.0-diffusers`	`mit-han-lab/dc-ae-f64c128-mix-1.0`
`mit-han-lab/dc-ae-f128c512-in-1.0-diffusers`	`mit-han-lab/dc-ae-f128c512-in-1.0`
`mit-han-lab/dc-ae-f128c512-mix-1.0-diffusers`	`mit-han-lab/dc-ae-f128c512-mix-1.0`

Este modelo fue aportado por lawrence-cj.

Carga un modelo en formato Diffusers con 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"

Cargar un modelo en Diffusers mediante 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)

El modelo AutoencoderDC tiene variantes de checkpoint de archivo único in y mix que tienen claves de checkpoint coincidentes, pero utilizan factores de escala diferentes. No es posible que Diffusers infiera automáticamente el archivo de configuración correcto para usar con el modelo basándose solo en el checkpoint y, por defecto, configurará el modelo utilizando el archivo de configuración de la variante mix. Para anular la configuración determinada automáticamente, use el argumento config al usar la carga de archivo único con checkpoints de la variante 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

clase diffusers.AutoencoderDC

< source >

( 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 )

Parámetros

in_channels (int, por defecto 3) — El número de canales de entrada en las muestras.
latent_channels (int, por defecto 32) — El número de canales en la representación del espacio latente.
encoder_block_types (str | tuple[str], por defecto "ResBlock") — El tipo de bloque a utilizar en el codificador.
decoder_block_types (str | tuple[str], por defecto "ResBlock") — El tipo de bloque a utilizar en el decodificador.
encoder_block_out_channels (tuple[int, ...], por defecto (128, 256, 512, 512, 1024, 1024)) — El número de canales de salida de cada bloque en el codificador.
decoder_block_out_channels (tuple[int, ...], por defecto (128, 256, 512, 512, 1024, 1024)) — El número de canales de salida para cada bloque en el decodificador.
encoder_layers_per_block (tuple[int], por defecto (2, 2, 2, 3, 3, 3)) — El número de capas por bloque en el codificador.
decoder_layers_per_block (tuple[int], por defecto (3, 3, 3, 3, 3, 3)) — El número de capas por bloque en el decodificador.
encoder_qkv_multiscales (tuple[tuple[int, ...], ...], por defecto ((), (), (), (5,), (5,), (5,))) — Configuraciones multiescala para las transformaciones QKV (query-key-value) del codificador.
decoder_qkv_multiscales (tuple[tuple[int, ...], ...], por defecto ((), (), (), (5,), (5,), (5,))) — Configuraciones multiescala para las transformaciones QKV (query-key-value) del decodificador.
upsample_block_type (str, por defecto "pixel_shuffle") — El tipo de bloque que se utiliza para el remuestreo ascendente en el decodificador.
downsample_block_type (str, por defecto "pixel_unshuffle") — El tipo de bloque que se utiliza para el muestreo descendente en el codificador.
decoder_norm_types (str | tuple[str], por defecto "rms_norm") — El tipo o tipos de normalización a utilizar en el decodificador.
decoder_act_fns (str | tuple[str], por defecto "silu") — La(s) función(es) de activación a utilizar en el decodificador.
encoder_out_shortcut (bool, por defecto True) — Si se utiliza un atajo al final del codificador.
decoder_in_shortcut (bool, por defecto True) — Si se utiliza un atajo al principio del decodificador.
decoder_conv_act_fn (str, por defecto "relu") — La función de activación que se utiliza al final del decodificador.
scaling_factor (float, por defecto 1.0) — La inversa multiplicativa de la raíz cuadrada de la media cuadrática de las características latentes. Se utiliza para escalar el espacio latente para que tenga varianza unitaria durante el entrenamiento del modelo de difusión. Los latentes se escalan con la fórmula z = z * scaling_factor antes de ser pasados al modelo de difusión. Al decodificar, los latentes se escalan de vuelta a la escala original con la fórmula: z = 1 / scaling_factor * z.

Un modelo Autoencoder introducido en DCAE y utilizado en SANA.

Este modelo hereda de ModelMixin. Consulta la documentación de la clase superior para sus métodos genéricos implementados para todos los modelos (como la descarga o el guardado).

encode

< source >

( x: Tensorreturn_dict: bool = True )

Parámetros

x (torch.Tensor) — Lote de entrada de imágenes.
return_dict (bool, por defecto True) — Si se debe devolver un ~models.vae.EncoderOutput en lugar de una tupla simple.

Codifica un lote de imágenes en latentes.

decode

< source >

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

Parámetros

z (torch.Tensor) — Lote de entrada de vectores latentes.
return_dict (bool, por defecto True) — Si se debe devolver un ~models.vae.DecoderOutput en lugar de una tupla simple.

Devuelve

~models.vae.DecoderOutput o tuple

If return_dict is True, a ~models.vae.DecoderOutput is returned, otherwise a plain tuple is returned.

Decodificar un lote de imágenes.

enable_tiling

< source >

( 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 )

Parámetros

tile_sample_min_height (int, opcional) — La altura mínima requerida para que una muestra se separe en mosaicos a lo largo de la dimensión de altura.
tile_sample_min_width (int, opcional) — El ancho mínimo requerido para que una muestra se separe en mosaicos a lo largo de la dimensión de ancho.
tile_sample_stride_height (int, opcional) — La cantidad mínima de solapamiento entre dos mosaicos verticales consecutivos. Esto garantiza que no se produzcan artefactos de mosaico a lo largo de la dimensión de altura.
tile_sample_stride_width (int, opcional) — La separación entre dos mosaicos horizontales consecutivos. Esto garantiza que no se produzcan artefactos de mosaico a lo largo de la dimensión de ancho.

Habilita la decodificación AE en mosaicos. Cuando esta opción está habilitada, la AE dividirá el tensor de entrada en mosaicos para calcular la decodificación y la codificación en varias etapas. Esto es útil para ahorrar una gran cantidad de memoria y para permitir el procesamiento de imágenes más grandes.

forward

< source >

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

Parámetros

sample (torch.Tensor) — Muestra de entrada.
return_dict (bool, opcional, valor predeterminado True) — Indica si se debe devolver un DecoderOutput en lugar de una tupla simple.

Devuelve

~models.vae.DecoderOutput o tuple

Si return_dict es True, se devuelve un ~models.vae.DecoderOutput, de lo contrario se devuelve una tuple simple.

DecoderOutput

clase diffusers.models.autoencoders.vae.DecoderOutput

< source >

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

Parámetros

sample (torch.Tensor de forma (batch_size, num_channels, height, width)) — La muestra de salida decodificada de la última capa del modelo.

Salida del método de decodificación.

Update on GitHub