Stable Diffusion のコードを斜め読みする


ここ最近DeepLearning関連のコードを書いていなかったので、
最新のDeepLearningの実装やお作法、Stable Diffusionの実装などをStable Diffusionのレポジトリ通して、この際お勉強.

# ソースコード

CompVis/stable-diffusion

https://github.com/CompVis/stable-diffusion

# 参考文献

仕組みなどは下記で簡単に抑えた上でコードを確認.

*本家サイト*

Stable Diffusion with 🧨 Diffusers

https://huggingface.co/blog/stable_diffusion

*日本語解説系*

日本語
【概要速修】Stable Diffusion(テキストから画像生成)はどうやって実現するのかざっくり仕組みを知る(DiffusionModel,Deep Learninig)【機械学習解説動画】

下記で概要を押さえてからコード見たから、大分読みやすかった.

https://youtu.be/aNdb_RUxtOM 

<iframe width="560" height="315" src="https://www.youtube.com/embed/aNdb_RUxtOM" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>

# Top Directory (ざっと)

まず一番上のフォルダで目についたものを確認.

## Model Cardとは

GithubのRepositoryを見ていると目に付くModelCardとは下記に説明がある

https://huggingface.co/docs/hub/models-cards

> Model cards are files that accompany the models and provide handy information. Under the hood, model cards are simple Markdown files with additional metadata. Model cards are essential for discoverability, reproducibility, and sharing! You can find a model card as the README.md file in any model repo.

とのことで、つまり

> モデルカードは、モデルに付随する便利な情報を提供するファイル。内部的には、モデルカードは追加のメタデータを含む単純なMarkdownファイルです。モデルカードは、見つけやすさ、再現性、および共有に不可欠です。モデルカードは、どのモデルリポジトリでも README.md ファイルとして見つけることができます。

ということ. 中身は

- 学習に使ったデータ
- 評価の結果
- 倫理的なポイント、制約

などが記載

## environment.yaml

conda用の環境設定ファイル

## scriptsフォルダ

ここに、デモを実行するためのスクリプトが入っている.

![スクリーンショット 2022-09-10 17.13.37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/849090/caf893e0-b5fb-bf52-9daf-4d6543cb2a48.png)

| ファイル名 | 内容 |
|:-----------|------------:|
| img2img       | 画像から画像への変換        |
| inpaint     | マスク部分の画像補完      |
| txt2img | 文字列から画像の生成 |
| knn2img | 文字列から画像検索 |

(knn2img自信なし. 詳しくは下記参照)

https://github.com/CompVis/latent-diffusion#rdm-with-text-to-image-retrieval

# コードを読んでく

## txt2img.py

いよいよこっからがソースコードの中身の確認.

### 引数
```python
    parser.add_argument(
        "--ckpt",
        type=str,
        default="models/ldm/stable-diffusion-v1/model.ckpt",
        help="path to checkpoint of model",
    )
```
ckptでモデルのファイルを指定. 通常はv1のモデル.

```python
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="the seed (for reproducible sampling)",
    )
```

シードを与えて、同じ画像をサンプリングできるようにする.

```python
    parser.add_argument(
        "--ddim_steps",
        type=int,
        default=50,
        help="number of ddim sampling steps",
    )
```

Diffusion Model のStep数を変更する

などあり.

### メイン処理

```python
    if opt.laion400m:
        print("Falling back to LAION 400M model...")
        opt.config = "configs/latent-diffusion/txt2img-1p4B-eval.yaml"
        opt.ckpt = "models/ldm/text2img-large/model.ckpt"
        opt.outdir = "outputs/txt2img-samples-laion400m"
```
LAION 400Mはオープンなデータセットでそれをもとに学習したモデルを使うかどうか

```python
    seed_everything(opt.seed)
```
Pytorchで再現性を制御するためにseedをセット.

```python
    config = OmegaConf.load(f"{opt.config}")
    model = load_model_from_config(config, f"{opt.ckpt}")
```
OmegaConfは設定系を扱いやすくするためのライブラリ. 
modelのロードはckptファイルからtorch.loadで読み込み

```python
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    model = model.to(device)
```

cudaかデバイスかの選択

```python
    if opt.plms:
        sampler = PLMSSampler(model)
    else:
        sampler = DDIMSampler(model)
```
DiffuseModelのサンプルの選択肢、PLMかDDIM. 何もしなければDDIMが選択.


```python
    print("Creating invisible watermark encoder (see https://github.com/ShieldMnt/invisible-watermark)...")
    wm = "StableDiffusionV1"
    wm_encoder = WatermarkEncoder()
    wm_encoder.set_watermark('bytes', wm.encode('utf-8'))
```
どうやらウォーターマークを画像に書き込んでいる.しかもinvisible watermarkなるもので、リンク先によれば、
離散ウェーブレット変換と離散コサイン変換を使って画像に書き込んでいる模様.
クリップされたりするとダメみたいだが、それ以外の加工なら耐えられるらしい.(明るくするなどしてもOK)

https://github.com/ShieldMnt/invisible-watermark

完全に防げるわけではなく、あくまでちょっと対策できるよという形.

```python
                for n in trange(opt.n_iter, desc="Sampling"):
                    for prompts in tqdm(data, desc="data"):
                        uc = None
                        if opt.scale != 1.0:
                            uc = model.get_learned_conditioning(batch_size * [""])
                        if isinstance(prompts, tuple):
                            prompts = list(prompts)
                        c = model.get_learned_conditioning(prompts)
                        shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
                        samples_ddim, _ = sampler.sample(S=opt.ddim_steps,
                                                         conditioning=c,
                                                         batch_size=opt.n_samples,
                                                         shape=shape,
                                                         verbose=False,
                                                         unconditional_guidance_scale=opt.scale,
                                                         unconditional_conditioning=uc,
                                                         eta=opt.ddim_eta,
                                                         x_T=start_code)
```
与えた複数の文字列を使って、文字列ごとにencodeして、Diffusion Modelのサンプリングを行なっている.
cと書かれているのが条件付けのために渡す特徴.

get_learned_conditioningで生成されるが、この関数は、
```python
    def get_learned_conditioning(self, c):
        if self.cond_stage_forward is None:
            if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode):
                c = self.cond_stage_model.encode(c)
                if isinstance(c, DiagonalGaussianDistribution):
                    c = c.mode()
            else:
                c = self.cond_stage_model(c)
        else:
            assert hasattr(self.cond_stage_model, self.cond_stage_forward)
            c = getattr(self.cond_stage_model, self.cond_stage_forward)(c)
        return c
```
cond_stage_modelは設定から読み込まれるが、設定しているyamlを確認すると
FrozenCLIPTextEncoderが指定されている.
やはり、ここが文字列を特徴量に変換している部分.


```python
                        x_samples_ddim = model.decode_first_stage(samples_ddim)
                        x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
                        x_samples_ddim = x_samples_ddim.cpu().permute(0, 2, 3, 1).numpy()
```

画像にするために処理している. decode_first_stageはencode_first_stageと同様に最初に画像と潜在変数を変換する処理の部分.
VQModelInterface, 何も変換しないIdentityFirstStage, AutoencoderKLの3つの選択肢がある.

コードを見てきた感じ以下のようになっていた.

```python
class VQModel(pl.LightningModule):
    def encode(self, x):
        h = self.encoder(x)
        h = self.quant_conv(h)
        quant, emb_loss, info = self.quantize(h)
        return quant, emb_loss, info

    def decode(self, quant):
        quant = self.post_quant_conv(quant)
        dec = self.decoder(quant)
        return dec
```

VQモデルはベクトル量子化VAEのこと. 近年自然画像を作るのに使われる

```python
class AutoencoderKL(pl.LightningModule):
    def forward(self, input, sample_posterior=True):
        posterior = self.encode(input)
        if sample_posterior:
            z = posterior.sample()
        else:
            z = posterior.mode()
        dec = self.decode(z)
        return dec, posterior
```

AutoEncoder KLって何のことだろうと思ったが、通常のVAEのこと.

```python
                        x_checked_image, has_nsfw_concept = check_safety(x_samples_ddim)
```
これは、画像が問題ないかを弾くための部分.

### sampler.sampleの中身

主に

https://github.com/CompVis/stable-diffusion/blob/ce05de28194041e030ccfc70c635fe3707cdfc30/ldm/models/diffusion/ddim.py#L171-L172

らへん

まずは下記. DDIMの実装を確認.

```python
        self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
```

これはDiffusion ModelのStepごとのノイズの掛け方などの設定をおこなっている部分.

その後に,ddim_samplingを呼び出して実際の複数ステップのsamplingをしている.

```python
        if x_T is None:
            img = torch.randn(shape, device=device)
        else:
            img = x_T
```
まず処理するx_Tがなければ、ノイズから潜在変数を生成する.
すでに前回の分があればそれをimgにセット.

imgと書かれているが、渡ってきているのは潜在変数で画像としての意味は持たない. (おそらく元々は画像を入れるのが想定されていたため)

```python
        for i, step in enumerate(iterator):
            index = total_steps - i - 1
            ts = torch.full((b,), step, device=device, dtype=torch.long)


            if mask is not None:
                assert x0 is not None
                img_orig = self.model.q_sample(x0, ts)  # TODO: deterministic forward pass?
                img = img_orig * mask + (1. - mask) * img


            outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
                                      quantize_denoised=quantize_denoised, temperature=temperature,
                                      noise_dropout=noise_dropout, score_corrector=score_corrector,
                                      corrector_kwargs=corrector_kwargs,
                                      unconditional_guidance_scale=unconditional_guidance_scale,
                                      unconditional_conditioning=unconditional_conditioning)
```
ここが実際のStepごとのsample処理をしている部分.

でp_sample_ddimの中のapply_model. DDIMSamplerはコンストラクタでmodelを渡されているのでそれで実行.

```python
        if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
            e_t = self.model.apply_model(x, t, c)
```

至る所でcond_stage_keyが使われており、どのような条件付きの学習をするかで分岐が切られている.

今回もmodelを適用しているのは下記だが、事前にcond_stage_keyごとに前処理が行われている.

```python
            # apply model by loop over crops
            output_list = [self.model(z_list[i], t, **cond_list[i]) for i in range(z.shape[-1])]
```

モデルの情報はyamlの次の部分から読み込まれている. これがDDPMに渡されて、modelが作られそれをもとにsampleでのapplyが実行される.

```python
    unet_config:
      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
      params:
        image_size: 32 # unused
        in_channels: 4
        out_channels: 4
        model_channels: 320
        attention_resolutions: [ 4, 2, 1 ]
        num_res_blocks: 2
        channel_mult: [ 1, 2, 4, 4 ]
        num_heads: 8
        use_spatial_transformer: True
        transformer_depth: 1
        context_dim: 768
        use_checkpoint: True
        legacy: False
```

self.modelに入っているのは、DiffusionWrapperクラス. 上記のyamlからデータを起こして指定されたcondで実行.

```python
class DiffusionWrapper(pl.LightningModule):
    def __init__(self, diff_model_config, conditioning_key):
        super().__init__()
        self.diffusion_model = instantiate_from_config(diff_model_config)
        self.conditioning_key = conditioning_key
        assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm']


    def forward(self, x, t, c_concat: list = None, c_crossattn: list = None):
        if self.conditioning_key is None:
            out = self.diffusion_model(x, t)
        elif self.conditioning_key == 'concat':
            xc = torch.cat([x] + c_concat, dim=1)
            out = self.diffusion_model(xc, t)
        elif self.conditioning_key == 'crossattn':
            cc = torch.cat(c_crossattn, 1)
            out = self.diffusion_model(x, t, context=cc)
        elif self.conditioning_key == 'hybrid':
            xc = torch.cat([x] + c_concat, dim=1)
            cc = torch.cat(c_crossattn, 1)
            out = self.diffusion_model(xc, t, context=cc)
        elif self.conditioning_key == 'adm':
            cc = c_crossattn[0]
            out = self.diffusion_model(x, t, y=cc)
        else:
            raise NotImplementedError()


        return out
```

### OpenAIのUnetModelを覗き見る

https://github.com/CompVis/stable-diffusion/blob/ce05de28194041e030ccfc70c635fe3707cdfc30/ldm/modules/diffusionmodules/openaimodel.py

ここに書かれている.

>     The full UNet model with attention and timestep embedding.

タイムステップとAttention機構付きのFULLUnetモデル.

initの中でモデルを作っている. InputとMiddleとOutputの3パート.使っているのは

- ResBlock
- AttentionBlock
- Downsample
- Upsample

ResBlockは言わずと知れたResNetBlock. AttentionBlockがAttention機構. 実装を見るとqkvとかやってる.
Downsample、Upsampleは縮小と拡大.



# ファイル拡張子

## ckpt

DeepLearningのモデルデータの拡張子.