# Kandinsky 2.2

> [!WARNING]
> このスクリプトは実験的なものであり、過学習やカタストロフィック・フォーゲッティングなどの問題に陥りやすいため、注意が必要です。データセットで最適な結果を得るために、さまざまなハイパーパラメータの探索をお試しください。

Kandinsky 2.2 は、よりフォトリアリスティックな画像を生成できる多言語テキストから画像への生成モデルです。このモデルには、テキストプロンプトから画像埋め込みを作成する image prior model と、prior model の埋め込みに基づいて画像を生成する decoder model が含まれています。そのため、Diffusers には Kandinsky 2.2 用の2つの独立したスクリプトがあり、1つは prior model の学習用、もう1つは decoder model の学習用です。両方のモデルを個別に学習させることも可能ですが、最適な結果を得るためには、prior model と decoder model の両方を学習させることをお勧めします。

GPUの種類によっては、モデルをメモリに収め、トレーニングを高速化するために、`gradient_checkpointing`（⚠️ 以前のモデルではサポートされていません！）、`mixed_precision`、および`gradient_accumulation_steps`を有効にする必要があります。[xFormers](../optimization/attention_backends)（バージョン[v0.0.16](https://github.com/huggingface/diffusers/issues/2234#issuecomment-1416931212)では一部のGPUでトレーニングに失敗するため、代わりに開発版をインストールする必要がある場合があります）を使用してメモリ効率の高いアテンションを有効にすることで、メモリ使用量をさらに削減できます。

This guide explores the [train_text_to_image_prior.py](https://github.com/huggingface/diffusers/blob/main/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py) and the [train_text_to_image_decoder.py](https://github.com/huggingface/diffusers/blob/main/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py) scripts to help you become more familiar with it, and how you can adapt it for your own use-case.

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

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

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

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

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

🤗 Accelerate の環境を初期化します:

```bash
accelerate config
```

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

```bash
accelerate config default
```

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

```py
from accelerate.utils import write_basic_config

write_basic_config()
```

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

> [!TIP]
> 以下のセクションでは、スクリプトの修正方法を理解するために重要なトレーニングスクリプトの一部を強調しています。ただし、スクリプトのすべての側面を詳細に説明するものではありません。さらに詳しく知りたい場合は、スクリプトを自由に読み進めていただき、ご質問や懸念事項があればお知らせください。

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

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

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

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

ほとんどのパラメータは[テキストから画像へ](text2image#script-parameters)のトレーニングガイドと同じです。それでは、Kandinskyのトレーニングスクリプトの解説に入っていきましょう!

### Min-SNR重み付け

[Min-SNR](https://huggingface.co/papers/2303.09556)重み付け戦略は、損失を再バランス化して収束を高速化することで、トレーニングを支援できます。トレーニングスクリプトは`epsilon`（ノイズ）または`v_prediction`の予測をサポートしていますが、Min-SNRはどちらの予測タイプにも対応しています。この重み付け戦略はPyTorchでのみサポートされています。

`--snr_gamma` パラメータを追加し、推奨値である 5.0 に設定します:

```bash
accelerate launch train_text_to_image_prior.py \
  --snr_gamma=5.0
```

## 学習スクリプト

学習スクリプトは、[テキストから画像へ](text2image#training-script)の学習ガイドと同様ですが、priorモデルとdecoderモデルの学習をサポートするように修正されています。本ガイドでは、Kandinsky 2.2の学習スクリプトに固有のコードに焦点を当てます。

[`main()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L441) 関数には、データセットの準備とモデルの学習を行うコードが含まれています。

すぐに気づく主な違いの一つとして、トレーニングスクリプトは、スケジューラやトークナイザーに加えて、画像の前処理用に[CLIPImageProcessor](https://huggingface.co/docs/transformers/main/en/model_doc/clip#transformers.CLIPImageProcessor)を読み込み、画像のエンコード用に[CLIPVisionModelWithProjection](https://huggingface.co/docs/transformers/main/en/model_doc/clip#transformers.CLIPVisionModelWithProjection)モデルを読み込む点です。

```py
noise_scheduler = DDPMScheduler(beta_schedule="squaredcos_cap_v2", prediction_type="sample")
image_processor = CLIPImageProcessor.from_pretrained(
    args.pretrained_prior_model_name_or_path, subfolder="image_processor"
)
tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="tokenizer")

with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="image_encoder", dtype=weight_dtype
    ).eval()
    text_encoder = CLIPTextModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="text_encoder", dtype=weight_dtype
    ).eval()
```

Kandinsky は画像埋め込みを生成するために [PriorTransformer](/docs/diffusers/main/ja/api/models/prior_transformer#diffusers.PriorTransformer) を使用するため、オプティマイザーを設定して、prior モードのパラメータを学習するようにする必要があります。

```py
prior = PriorTransformer.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder="prior")
prior.train()
optimizer = optimizer_cls(
    prior.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)
```

次に、入力キャプションがトークン化され、画像は[CLIPImageProcessor](https://huggingface.co/docs/transformers/main/en/model_doc/clip#transformers.CLIPImageProcessor)によって[前処理](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L632)されます:

```py
def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
    examples["text_input_ids"], examples["text_mask"] = tokenize_captions(examples)
    return examples
```

最後に、[training loop](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_prior.py#L718)は入力画像を潜在変数に変換し、画像埋め込みにノイズを加え、予測を行います。

```py
model_pred = prior(
    noisy_latents,
    timestep=timesteps,
    proj_embedding=prompt_embeds,
    encoder_hidden_states=text_encoder_hidden_states,
    attention_mask=text_mask,
).predicted_image_embedding
```

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

[`main()`](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L440) 関数には、データセットの準備とモデルの学習を行うコードが含まれています。

従来のモデルとは異なり、デコーダーは [VQModel](/docs/diffusers/main/ja/api/models/vq#diffusers.VQModel) を初期化して潜在変数を画像にデコードし、[UNet2DConditionModel](/docs/diffusers/main/ja/api/models/unet2d-cond#diffusers.UNet2DConditionModel) を使用します:

```py
with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
    vae = VQModel.from_pretrained(
        args.pretrained_decoder_model_name_or_path, subfolder="movq", dtype=weight_dtype
    ).eval()
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        args.pretrained_prior_model_name_or_path, subfolder="image_encoder", dtype=weight_dtype
    ).eval()
unet = UNet2DConditionModel.from_pretrained(args.pretrained_decoder_model_name_or_path, subfolder="unet")
```

次に、スクリプトには複数の画像変換と、これらの変換を画像に適用してピクセル値を返すための[前処理](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L622)関数が含まれています。

```py
def preprocess_train(examples):
    images = [image.convert("RGB") for image in examples[image_column]]
    examples["pixel_values"] = [train_transforms(image) for image in images]
    examples["clip_pixel_values"] = image_processor(images, return_tensors="pt").pixel_values
    return examples
```

最後に、[training loop](https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/kandinsky2_2/text_to_image/train_text_to_image_decoder.py#L706)は、画像を潜在変数に変換し、ノイズを追加し、ノイズ残差を予測する処理を担当します。

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

```py
model_pred = unet(noisy_latents, timesteps, None, added_cond_kwargs=added_cond_kwargs).sample[:, :4]
```

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

すべての変更が完了したか、デフォルトの設定で問題ない場合は、トレーニングスクリプトの起動準備が整いました! 🚀

[Naruto BLIP captions](https://huggingface.co/datasets/lambdalabs/naruto-blip-captions) データセットで学習し、独自のナルトキャラクターを生成します。ただし、[学習用データセットの作成](create_dataset) ガイドに従って、独自のデータセットを作成して学習することも可能です。Hub 上のデータセット名を環境変数 `DATASET_NAME` に設定してください。独自のファイルで学習する場合は、環境変数 `TRAIN_DIR` にデータセットへのパスを設定してください。

複数の GPU で学習を行う場合は、`accelerate launch` コマンドに `--multi_gpu` パラメータを追加してください。

> [!TIP]
> Weights & Biases でトレーニングの進捗を監視するには、トレーニングコマンドに `--report_to=wandb` パラメータを追加します。また、結果を追跡するために、トレーニングコマンドに `--validation_prompt` も追加する必要があります。これにより、モデルのデバッグや中間結果の確認に非常に役立ちます。

```bash
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_prior.py \
  --dataset_name=$DATASET_NAME \
  --resolution=768 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --checkpoints_total_limit=3 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --validation_prompts="A robot naruto, 4k photo" \
  --report_to="wandb" \
  --push_to_hub \
  --output_dir="kandi2-prior-naruto-model"
```

```bash
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_decoder.py \
  --dataset_name=$DATASET_NAME \
  --resolution=768 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \
  --max_train_steps=15000 \
  --learning_rate=1e-05 \
  --max_grad_norm=1 \
  --checkpoints_total_limit=3 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --validation_prompts="A robot naruto, 4k photo" \
  --report_to="wandb" \
  --push_to_hub \
  --output_dir="kandi2-decoder-naruto-model"
```

トレーニングが完了したら、新しく学習したモデルを推論に使用できます!

```py
from diffusers import AutoPipelineForText2Image, DiffusionPipeline
import torch

prior_pipeline = DiffusionPipeline.from_pretrained(output_dir, dtype=torch.float16)
prior_components = {"prior_" + k: v for k,v in prior_pipeline.components.items()}
pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", **prior_components, dtype=torch.float16)

pipe.enable_model_cpu_offload()
prompt="A robot naruto, 4k photo"
image = pipeline(prompt=prompt, negative_prompt=negative_prompt).images[0]
```

> [!TIP]
> ご自身で学習したデコーダーのチェックポイントに`kandinsky-community/kandinsky-2-2-decoder`を置き換えても構いません。

```py
from diffusers import AutoPipelineForText2Image
import torch

pipeline = AutoPipelineForText2Image.from_pretrained("path/to/saved/model", dtype=torch.float16)
pipeline.enable_model_cpu_offload()

prompt="A robot naruto, 4k photo"
image = pipeline(prompt=prompt).images[0]
```

デコーダーモデルでは、保存されたチェックポイントから推論を実行することもできます。これは中間結果を確認する場合に便利です。この場合、チェックポイントをUNetに読み込みます。

```py
from diffusers import AutoPipelineForText2Image, UNet2DConditionModel

unet = UNet2DConditionModel.from_pretrained("path/to/saved/model" + "/checkpoint-<N>/unet")

pipeline = AutoPipelineForText2Image.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", unet=unet, dtype=torch.float16)
pipeline.enable_model_cpu_offload()

image = pipeline(prompt="A robot naruto, 4k photo").images[0]
```

## 次のステップ

Kandinsky 2.2 モデルの学習が完了しました。新しいモデルの使用方法について詳しく知るために、以下のガイドが役立つ可能性があります:

- [Kandinsky](../api/pipelines/kandinsky)のガイドを読み、テキストから画像への変換、画像から画像への変換、インペインティング、補間など、さまざまなタスクでの使用方法や、ControlNetとの組み合わせ方法について学びましょう。
- [DreamBooth](dreambooth)と[LoRA](lora)のトレーニングガイドを参照し、わずか数枚のサンプル画像でパーソナライズされたKandinskyモデルをトレーニングする方法を学びましょう。これらの2つのトレーニング手法は組み合わせることも可能です。