#

プロンプト付きの安定した拡散を利用したペイント中の画像

11import argparse
12from pathlib import Path
13from typing import Optional
14
15import torch
16
17from labml import lab, monit
18from labml_nn.diffusion.stable_diffusion.latent_diffusion import LatentDiffusion
19from labml_nn.diffusion.stable_diffusion.sampler import DiffusionSampler
20from labml_nn.diffusion.stable_diffusion.sampler.ddim import DDIMSampler
21from labml_nn.diffusion.stable_diffusion.util import load_model, save_images, load_img, set_seed

#

イメージ・イン・ペインティング・クラス

24class InPaint:

#

28    model: LatentDiffusion
29    sampler: DiffusionSampler

#

checkpoint_path チェックポイントのパスです
ddim_steps はサンプリングステップの数
ddim_eta DDIM サンプリング定数です $η$

31    def __init__(self, *, checkpoint_path: Path,
32                 ddim_steps: int = 50,
33                 ddim_eta: float = 0.0):

#

39        self.ddim_steps = ddim_steps

#

負荷潜伏拡散モデル

42        self.model = load_model(checkpoint_path)

#

デバイスを取得

44        self.device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")

#

モデルをデバイスに移動

46        self.model.to(self.device)

#

DDIM サンプラーを初期化

49        self.sampler = DDIMSampler(self.model,
50                                   n_steps=ddim_steps,
51                                   ddim_eta=ddim_eta)

#

dest_path 生成された画像を保存するパスです
orig_img 変換する画像です
strength 元の画像のどの程度保存しないかを指定します
batch_size はバッチで生成する画像の数です
prompt で画像を生成するプロンプトです
uncond_scale 無条件ガイダンススケールです $s$ これは次の用途に使用されます $ϵ_{θ} (x_{t}, c) = s ϵ_{cond} (x_{t}, c) + (s - 1) ϵ_{cond} (x_{t}, c_{u})$

53    @torch.no_grad()
54    def __call__(self, *,
55                 dest_path: str,
56                 orig_img: str,
57                 strength: float,
58                 batch_size: int = 3,
59                 prompt: str,
60                 uncond_scale: float = 5.0,
61                 mask: Optional[torch.Tensor] = None,
62                 ):

#

プロンプトを一括作成

73        prompts = batch_size * [prompt]

#

画像を読み込む

75        orig_image = load_img(orig_img).to(self.device)

#

潜在空間に画像をエンコードし、batch_size そのコピーを作成します

77        orig = self.model.autoencoder_encode(orig_image).repeat(batch_size, 1, 1, 1)

#

mask 提供されていない場合は、画像の下半分を保存するようにサンプルマスクを設定します

80        if mask is None:
81            mask = torch.zeros_like(orig, device=self.device)
82            mask[:, :, mask.shape[2] // 2:, :] = 1.
83        else:
84            mask = mask.to(self.device)

#

ノイズは元の画像を拡散します

86        orig_noise = torch.randn(orig.shape, device=self.device)

#

オリジナルを拡散させるまでのステップ数を求める

89        assert 0. <= strength <= 1., 'can only work with strength in [0.0, 1.0]'
90        t_index = int(strength * self.ddim_steps)

#

AMP オートキャスティング

93        with torch.cuda.amp.autocast():

#

無条件スケーリングでは、 $1$ 空のプロンプトでは埋め込みは取得されません (条件なし)。

95            if uncond_scale != 1.0:
96                un_cond = self.model.get_text_conditioning(batch_size * [""])
97            else:
98                un_cond = None

#

プロンプトの埋め込みを入手

100            cond = self.model.get_text_conditioning(prompts)

#

元の画像にノイズを追加

102            x = self.sampler.q_sample(orig, t_index, noise=orig_noise)

#

マスクされた領域はそのままに、ノイズの多い画像から再構成します

104            x = self.sampler.paint(x, cond, t_index,
105                                   orig=orig,
106                                   mask=mask,
107                                   orig_noise=orig_noise,
108                                   uncond_scale=uncond_scale,
109                                   uncond_cond=un_cond)

#

オートエンコーダから画像をデコードします

111            images = self.model.autoencoder_decode(x)

#

画像を保存

114        save_images(images, dest_path, 'paint_')

#

CLI

117def main():

#

121    parser = argparse.ArgumentParser()
122
123    parser.add_argument(
124        "--prompt",
125        type=str,
126        nargs="?",
127        default="a painting of a cute monkey playing guitar",
128        help="the prompt to render"
129    )
130
131    parser.add_argument(
132        "--orig-img",
133        type=str,
134        nargs="?",
135        help="path to the input image"
136    )
137
138    parser.add_argument("--batch_size", type=int, default=4, help="batch size", )
139    parser.add_argument("--steps", type=int, default=50, help="number of sampling steps")
140
141    parser.add_argument("--scale", type=float, default=5.0,
142                        help="unconditional guidance scale: "
143                             "eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))")
144
145    parser.add_argument("--strength", type=float, default=0.75,
146                        help="strength for noise: "
147                             " 1.0 corresponds to full destruction of information in init image")
148
149    opt = parser.parse_args()
150    set_seed(42)
151
152    in_paint = InPaint(checkpoint_path=lab.get_data_path() / 'stable-diffusion' / 'sd-v1-4.ckpt',
153                       ddim_steps=opt.steps)
154
155    with monit.section('Generate'):
156        in_paint(dest_path='outputs',
157                 orig_img=opt.orig_img,
158                 strength=opt.strength,
159                 batch_size=opt.batch_size,
160                 prompt=opt.prompt,
161                 uncond_scale=opt.scale)

#

165if __name__ == "__main__":
166    main()