Hugging Face's logo

Diffusers documentation

潜在一貫性蒸留

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

潜在一貫性蒸留

Latent Consistency Models (LCMs) は、わずか数ステップで高品質な画像を生成できるため、多くのパイプラインが少なくとも25ステップ以上を必要とする現状において、大きな飛躍を意味します。LCMは、任意のStable Diffusionモデルにlatent consistency distillation手法を適用して生成されます。この手法は、latent空間に対してone-stage guided distillationを適用し、skipping-step手法を組み込むことでタイムステップを一貫してスキップし、蒸留プロセスを高速化します（詳細は論文のセクション4.1、4.2、4.3を参照してください）。

vRAM が限られた GPU で学習を行う場合は、gradient_checkpointing、gradient_accumulation_steps、mixed_precision を有効にして、メモリ使用量を削減し、学習を高速化してください。xFormers と bitsandbytes の 8-bit オプティマイザーによるメモリ効率の高いアテンションを有効にすることで、メモリ使用量をさらに削減できます。

このガイドでは、train_lcm_distill_sd_wds.py スクリプトを詳しく解説し、その仕組みを理解したうえで、ご自身のユースケースに合わせてどのようにカスタマイズできるかをご紹介します。

スクリプトを実行する前に、ソースからライブラリをインストールしていることを確認してください。

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

次に、トレーニングスクリプトが含まれるサンプルフォルダに移動し、使用するスクリプトに必要な依存関係をインストールします:

cd examples/consistency_distillation
pip install -r requirements.txt

🤗 Accelerate は、複数 GPU/TPU での学習や混合精度学習を支援するためのライブラリです。ハードウェアや環境に基づいて、学習環境を自動的に設定します。詳細については、🤗 Accelerate の クイックツアー をご覧ください。

🤗 Accelerate の環境を初期化します（torch.compile を有効にすると、トレーニングが大幅に高速化されます）:

accelerate config

設定を選択せずにデフォルトの🤗 Accelerate環境を設定するには:

accelerate config default

または、ノートブックのようにインタラクティブシェルをサポートしていない環境では、以下を使用できます。

from accelerate.utils import write_basic_config

write_basic_config()

最後に、独自のデータセットでモデルをトレーニングしたい場合は、トレーニング用のデータセットを作成するガイドを参照して、トレーニングスクリプトと互換性のあるデータセットの作成方法を確認してください。

スクリプトのパラメータ

以下のセクションでは、スクリプトの修正方法を理解するために重要な部分に焦点を当てていますが、スクリプトのすべての側面を詳細に説明するものではありません。さらに詳しく知りたい場合は、スクリプトをぜひお読みいただき、ご質問やご懸念事項があればお知らせください。

トレーニングスクリプトには、トレーニング実行をカスタマイズするための多くのパラメータが用意されています。すべてのパラメータとその説明は、parse_args()関数に記載されています。この関数は、トレーニングのバッチサイズや学習率などの各パラメータにデフォルト値を設定しますが、必要に応じてトレーニングコマンドで独自の値を設定することもできます。

例えば、fp16 形式の混合精度でトレーニングを高速化するには、トレーニングコマンドに --mixed_precision パラメータを追加します。

accelerate launch train_lcm_distill_sd_wds.py \
  --mixed_precision="fp16"

ほとんどのパラメータはText-to-imageのトレーニングガイドと同じですが、本ガイドでは、潜在一貫性蒸留に関連するパラメータに焦点を当てます。

  • --pretrained_teacher_model: 教師モデルとして使用する事前学習済み潜在拡散モデルへのパス
  • --pretrained_vae_model_name_or_path: 事前学習済み VAE のパス。SDXL の VAE は数値的不安定性を生じることが知られているため、このパラメータにより、fp16 で動作する madebyollin による代替 VAE（この VAE など）を指定できます。
  • --w_min と --w_max: guidance scale サンプリングにおける guidance scale の最小値と最大値
  • --num_ddim_timesteps: DDIMサンプリングにおけるタイムステップ数
  • --loss_type: 潜在一貫性蒸留で計算する損失の種類（L2またはHuber）; 外れ値に対してより頑健であるため、通常はHuber損失が推奨されます
  • --huber_c: Huber損失のパラメータ

学習スクリプト

学習スクリプトは、画像の前処理と学習用データセットの作成を行うデータセットクラス Text2ImageDataset を作成することから始まります。

def transform(example):
    image = example["image"]
    image = TF.resize(image, resolution, interpolation=transforms.InterpolationMode.BILINEAR)

    c_top, c_left, _, _ = transforms.RandomCrop.get_params(image, output_size=(resolution, resolution))
    image = TF.crop(image, c_top, c_left, resolution, resolution)
    image = TF.to_tensor(image)
    image = TF.normalize(image, [0.5], [0.5])

    example["image"] = image
    return example

クラウドに保存された大規模なデータセットの読み書きにおけるパフォーマンス向上のため、このスクリプトはWebDataset形式を使用して、変換を適用し、トレーニング用のデータセットとデータローダーを構築する前処理パイプラインを作成します。これにより、まずデータセット全体をダウンロードする必要なく、画像を処理してトレーニングループに供給できます。

processing_pipeline = [
    wds.decode("pil", handler=wds.ignore_and_continue),
    wds.rename(image="jpg;png;jpeg;webp", text="text;txt;caption", handler=wds.warn_and_continue),
    wds.map(filter_keys({"image", "text"})),
    wds.map(transform),
    wds.to_tuple("image", "text"),
]

main() 関数では、ノイズスケジューラ、トークナイザー、テキストエンコーダー、VAE といった必要なコンポーネントがすべて読み込まれます。教師 UNet もここで読み込まれ、その後、教師 UNet から生徒 UNet を作成できます。生徒 UNet は、学習中にオプティマイザーによって更新されます。

teacher_unet = UNet2DConditionModel.from_pretrained(
    args.pretrained_teacher_model, subfolder="unet", revision=args.teacher_revision
)

unet = UNet2DConditionModel(**teacher_unet.config)
unet.load_state_dict(teacher_unet.state_dict(), strict=False)
unet.train()

これで、UNetのパラメータを更新するためのoptimizerを作成できます:

optimizer = optimizer_class(
    unet.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

datasetを作成します:

dataset = Text2ImageDataset(
    train_shards_path_or_url=args.train_shards_path_or_url,
    num_train_examples=args.max_train_samples,
    per_gpu_batch_size=args.train_batch_size,
    global_batch_size=args.train_batch_size * accelerator.num_processes,
    num_workers=args.dataloader_num_workers,
    resolution=args.resolution,
    shuffle_buffer_size=1000,
    pin_memory=True,
    persistent_workers=True,
)
train_dataloader = dataset.train_dataloader

次に、training loopのセットアップと、latent consistency distillation手法の実装（詳細は論文のAlgorithm 1を参照）を行います。このスクリプトのセクションでは、latentsにノイズを追加し、サンプリングとguidance scale embeddingの作成、そしてノイズから元の画像を予測する処理を行います。

pred_x_0 = predicted_origin(
    noise_pred,
    start_timesteps,
    noisy_model_input,
    noise_scheduler.config.prediction_type,
    alpha_schedule,
    sigma_schedule,
)

model_pred = c_skip_start * noisy_model_input + c_out_start * pred_x_0

次に、教師モデルの予測とLCMの予測を取得し、損失を計算したうえで、LCMへ逆伝播します。

if args.loss_type == "l2":
    loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
elif args.loss_type == "huber":
    loss = torch.mean(
        torch.sqrt((model_pred.float() - target.float()) ** 2 + args.huber_c**2) - args.huber_c
    )

トレーニングループの仕組みについて詳しく知りたい場合は、デノイジングプロセスの基本的なパターンを解説しているパイプライン、モデル、スケジューラーの理解をご覧ください。

スクリプトを実行します

これで学習スクリプトを起動して、蒸留を開始する準備が整いました!

このガイドでは、こちらの Hub に保存されている Conceptual Captions 12M データセットへのパスを指定するために --train_shards_path_or_url を使用します。環境変数 MODEL_DIR に教師モデルの名前を、OUTPUT_DIR にモデルを保存する場所を設定してください。

export MODEL_DIR="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="path/to/saved/model"

accelerate launch train_lcm_distill_sd_wds.py \
    --pretrained_teacher_model=$MODEL_DIR \
    --output_dir=$OUTPUT_DIR \
    --mixed_precision=fp16 \
    --resolution=512 \
    --learning_rate=1e-6 --loss_type="huber" --ema_decay=0.95 --adam_weight_decay=0.0 \
    --max_train_steps=1000 \
    --max_train_samples=4000000 \
    --dataloader_num_workers=8 \
    --train_shards_path_or_url="pipe:curl -L -s https://huggingface.co/datasets/laion/conceptual-captions-12m-webdataset/resolve/main/data/{00000..01099}.tar?download=true" \
    --validation_steps=200 \
    --checkpointing_steps=200 --checkpoints_total_limit=10 \
    --train_batch_size=12 \
    --gradient_checkpointing --enable_xformers_memory_efficient_attention \
    --gradient_accumulation_steps=1 \
    --use_8bit_adam \
    --resume_from_checkpoint=latest \
    --report_to=wandb \
    --seed=453645634 \
    --push_to_hub

トレーニングが完了すると、新しいLCMを使って推論を実行できます。

from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
import torch

unet = UNet2DConditionModel.from_pretrained("your-username/your-model", dtype=torch.float16, variant="fp16")
pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", unet=unet, dtype=torch.float16, variant="fp16")

pipeline.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipeline.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "sushi rolls in the form of panda heads, sushi platter"

image = pipeline(prompt, num_inference_steps=4, guidance_scale=1.0).images[0]

LoRA

LoRAは、学習可能なパラメータ数を大幅に削減するための学習手法です。その結果、学習が高速になり、生成される重みは非常に小さい（約100MB）ため、保存しやすくなります。LoRAで学習するには、train_lcm_distill_lora_sd_wds.pyまたはtrain_lcm_distill_lora_sdxl.wds.pyスクリプトを使用してください。

LoRAの学習スクリプトについては、LoRAの学習ガイドで詳しく説明しています。

Stable Diffusion XL

Stable Diffusion XL (SDXL) は高解像度の画像を生成する強力なテキストから画像へのモデルであり、そのアーキテクチャに第2のテキストエンコーダーを追加しています。train_lcm_distill_sdxl_wds.py スクリプトを使用して、LoRA で SDXL モデルを学習します。

SDXLの学習スクリプトについては、SDXLの学習ガイドで詳しく説明しています。

次のステップ

LCMモデルの蒸留が完了しました。おめでとうございます。LCMについてさらに詳しく知るために、以下のリソースが役立つかもしれません:

  • テキストから画像への変換、画像から画像への変換、およびLoRAチェックポイントとの組み合わせで、推論にLCMを使用する方法を学びましょう。
  • SDXL LCM-LoRAによる超高速推論、品質比較、ベンチマークなどについて詳しく知るには、SDXL in 4 steps with Latent Consistency LoRAsのブログ記事をお読みください。
Update on GitHub

←Custom Diffusion DDPOによる強化学習トレーニング→