# 潜在一貫性蒸留

[Latent Consistency Models (LCMs)](https://hf.co/papers/2310.04378) は、わずか数ステップで高品質な画像を生成できるため、多くのパイプラインが少なくとも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](../optimization/attention_backends) および [bitsandbytes'](https://github.com/TimDettmers/bitsandbytes) の 8-bit オプティマイザーによるメモリ効率の高いアテンションを有効にすることで、メモリ使用量をさらに削減できます。

このガイドでは、[train_lcm_distill_sd_wds.py](https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_sd_wds.py) スクリプトを詳しく解説し、その仕組みを理解したうえで、ご自身のユースケースに合わせてどのようにカスタマイズできるかをご紹介します。

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

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

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

```bash
cd examples/consistency_distillation
pip install -r requirements.txt
```

> [!TIP]
> 🤗 Accelerate は、複数 GPU/TPU での学習や混合精度学習を支援するためのライブラリです。ハードウェアと環境に基づいて、学習設定を自動的に構成します。詳細については、🤗 Accelerate の [クイックツアー](https://huggingface.co/docs/accelerate/quicktour) をご覧ください。

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

```bash
accelerate config
```

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

```bash
accelerate config default
```

または、ノートブックのように対話型シェルをサポートしていない環境では、以下を使用できます:

```py
from accelerate.utils import write_basic_config

write_basic_config()
```

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

## スクリプトのパラメータ

> [!TIP]
> 以下のセクションでは、スクリプトの修正方法を理解するために重要な部分に焦点を当てていますが、スクリプトのすべての側面を詳細に説明するものではありません。さらに詳しく知りたい場合は、[スクリプト](https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_sd_wds.py)をぜひお読みいただき、ご質問やご懸念事項があればお知らせください。

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

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

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

ほとんどのパラメータは[Text-to-image](text2image#script-parameters)トレーニングガイドのパラメータと同一であるため、本ガイドでは潜在一貫性蒸留に関連するパラメータに焦点を当てます。

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

## 学習スクリプト

学習スクリプトは、画像の前処理と学習用データセットの作成を行うデータセットクラス [`Text2ImageDataset`](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L141) を作成することから始まります。

```py
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](https://github.com/webdataset/webdataset)形式を使用して、変換を適用し、トレーニング用のデータセットとデータローダーを構築する前処理パイプラインを作成します。これにより、まずデータセット全体をダウンロードする必要なく、画像を処理してトレーニングループに供給できます。

```py
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()`](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L768) 関数では、ノイズスケジューラ、トークナイザー、テキストエンコーダー、VAE といった必要なコンポーネントがすべて読み込まれます。教師 UNet もここで読み込まれ、その後、教師 UNet から生徒 UNet を作成できます。生徒 UNet は、学習中にオプティマイザーによって更新されます。

```py
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](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L979)を作成できます:

```py
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](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L994)を作成します:

```py
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](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L1049)のセットアップと、latent consistency distillation手法の実装（詳細は論文のAlgorithm 1を参照）を行います。スクリプトのこのセクションでは、latentsにノイズを追加し、サンプリングとguidance scale embeddingの作成、そしてノイズから元の画像を予測する処理を行います。

```py
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
```

次に、[教師モデルの予測](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L1172)と[LCMの予測](https://github.com/huggingface/diffusers/blob/3b37488fa3280aed6a95de044d7a42ffdcb565ef/examples/consistency_distillation/train_lcm_distill_sd_wds.py#L1209)を取得し、損失を計算したうえで、LCMへ逆伝播します。

```py
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
    )
```

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

## スクリプトを実行します

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

このガイドでは、[こちら](https://huggingface.co/datasets/laion/conceptual-captions-12m-webdataset)の Hub に保存されている [Conceptual Captions 12M](https://github.com/google-research-datasets/conceptual-12m) データセットへのパスを指定するために `--train_shards_path_or_url` を使用します。教師モデルの名前を `MODEL_DIR` 環境変数に、モデルを保存する場所を `OUTPUT_DIR` に設定してください。

```bash
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を使って推論を実行できます。

```py
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](https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_lora_sd_wds.py)または[train_lcm_distill_lora_sdxl.wds.py](https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py)スクリプトを使用してください。

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

## Stable Diffusion XL

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

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

## 次のステップ

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

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