Diffusers documentation
AutoencoderDC
AutoencoderDC
SANAで使用され、MIT HAN Lab の Junyu Chen*、Han Cai*、Junsong Chen、Enze Xie、Shang Yang、Haotian Tang、Muyang Li、Yao Lu、Song Han によって DCAE で紹介された 2D Autoencoder モデル。
論文の抄録は以下の通りです：
高解像度拡散モデルの高速化を目的とした新しいオートエンコーダーモデル群である Deep Compression Autoencoder (DC-AE) を提案します。既存のオートエンコーダーモデルは、中程度の空間圧縮率（例：8x）において優れた結果を示していますが、高い空間圧縮率（例：64x）では満足できる再構成精度を維持できません。この課題に対処するため、2つの重要な技術を導入します。(1) Residual Autoencoding：space-to-channel 変換された特徴量に基づいた残差を学習するようにモデルを設計し、高空間圧縮オートエンコーダーの最適化の難しさを軽減します。(2) Decoupled High-Resolution Adaptation：高空間圧縮オートエンコーダーの汎化ペナルティを軽減するための、効率的な分離型3フェーズ学習戦略です。これらの設計により、再構成品質を維持しながら、オートエンコーダーの空間圧縮率を最大 128 まで向上させます。DC-AE を潜在拡散モデルに適用することで、精度の低下なしに大幅な高速化を実現します。例えば、ImageNet 512x512 において、広く使用されている SD-VAE-f8 オートエンコーダーと比較して、UViT-H において H100 GPU 上で 19.1 倍の推論高速化と 17.9 倍の学習高速化を実現し、さらに優れた FID を達成しました。コードは この https URL で利用可能です。
Diffusersでは、以下のDCAEモデルがリリースされ、サポートされています。
このモデルは lawrence-cj によって提供されました。
from_pretrained() を使用して、Diffusers 形式のモデルを読み込みます。
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 バリアントの設定ファイルを使用してモデルを設定します。自動決定された設定を上書きするには、in バリアントのチェックポイントで単一ファイル読み込みを使用する際に、config 引数を使用してください。
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
class 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[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,))decoder_qkv_multiscales: tuple[tuple[int, ...], ...] = ((), (), (), (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（クエリ・キー・バリュー）変換用のマルチスケール設定。 - decoder_qkv_multiscales (
tuple[tuple[int, ...], ...]、デフォルトは((), (), (), (5,), (5,), (5,))) — デコーダーの QKV（クエリ・キー・バリュー）変換用のマルチスケール設定です。 - 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) — デコーダーの先頭でショートカットを使用するかどうか。 - decoder_conv_act_fn (
str、デフォルトは"relu") — デコーダーの末尾で使用する活性化関数です。 - scaling_factor (
float、デフォルトは1.0) — 潜在特徴の二乗平均平方根の乗法的逆数です。拡散モデルの学習時に、潜在空間を分散 1 にスケーリングするために使用されます。潜在変数は、拡散モデルに渡される前にz = z * scaling_factorの式でスケーリングされます。デコード時には、z = 1 / scaling_factor * zの式で元のスケールに戻されます。
DCAE で提案され、SANA で使用されている Autoencoder モデルです。
このモデルは ModelMixin を継承しています。すべてのモデルに実装されている汎用的なメソッド（ダウンロードや保存など）については、スーパークラスのドキュメントを参照してください。
encode
< source >( x: Tensorreturn_dict: bool = True )
バッチの画像を潜在変数にエンコードします。
decode
< source >( z: Tensorreturn_dict: bool = True ) → ~models.vae.DecoderOutput or tuple
画像のバッチをデコードします。
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 )
パラメータ
- tile_sample_min_height (
int, 任意) — サンプルを高さ方向にタイル分割するために必要な最小高さ。 - tile_sample_min_width (
int, 任意) — 幅方向にサンプルをタイルに分割するために必要な最小幅です。 - tile_sample_stride_height (
int, 省略可) — 縦方向の連続する2つのタイル間の最小重複量です。これにより、高さ方向にタイル化によるアーティファクトが発生しないことを保証します。 - tile_sample_stride_width (
int, 任意) — 連続する2つの水平タイル間のストライドです。これにより、幅方向にタイル化によるアーティファクトが発生しないことを保証します。
タイル化AEデコーディングを有効にします。このオプションを有効にすると、AEは入力テンソルをタイルに分割して、デコーディングとエンコーディングを複数ステップで計算します。これにより、大量のメモリを節約し、より大きな画像の処理を可能にします。
forward
< source >( sample: Tensorreturn_dict: bool = True ) → ~models.vae.DecoderOutput or tuple
DecoderOutput
class diffusers.models.autoencoders.vae.DecoderOutput
< source >( sample: Tensorcommit_loss: typing.Optional[torch.FloatTensor] = None )
デコードメソッドの出力。