#

根据给定图像的提示，使用稳定的扩散生成图像

11import argparse
12from pathlib import Path
13
14import torch
15
16from labml import lab, monit
17from labml_nn.diffusion.stable_diffusion.sampler.ddim import DDIMSampler
18from labml_nn.diffusion.stable_diffusion.util import load_model, load_img, save_images, set_seed

#

图像到图像类别

21class Img2Img:

#

checkpoint_path 是检查点的路径
ddim_steps 是采样步骤的数量
ddim_eta 是 DDIM 采样 $η$ 常数

26    def __init__(self, *, checkpoint_path: Path,
27                 ddim_steps: int = 50,
28                 ddim_eta: float = 0.0):

#

34        self.ddim_steps = ddim_steps

#

载荷潜在扩散模型

37        self.model = load_model(checkpoint_path)

#

获取设备

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

#

将模型移至设备

41        self.model.to(self.device)

#

初始化 DDIM 采样器

44        self.sampler = DDIMSampler(self.model,
45                                   n_steps=ddim_steps,
46                                   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})$

48    @torch.no_grad()
49    def __call__(self, *,
50                 dest_path: str,
51                 orig_img: str,
52                 strength: float,
53                 batch_size: int = 3,
54                 prompt: str,
55                 uncond_scale: float = 5.0,
56                 ):

#

做一批提示

67        prompts = batch_size * [prompt]

#

加载图片

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

#

在潜在空间中对图像进行编码并制作batch_size 副本

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

#

获取漫反射原稿的步数

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

#

AMP 自动投射

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

#

在无条件缩放中，无法 $1$ 获取空提示的嵌入值（无条件）。

80            if uncond_scale != 1.0:
81                un_cond = self.model.get_text_conditioning(batch_size * [""])
82            else:
83                un_cond = None

#

获取提示嵌入信息

85            cond = self.model.get_text_conditioning(prompts)

#

向原始图像添加噪点

87            x = self.sampler.q_sample(orig, t_index)

#

从嘈杂的图像中重建

89            x = self.sampler.paint(x, cond, t_index,
90                                   uncond_scale=uncond_scale,
91                                   uncond_cond=un_cond)

#

从自动编码器解码图像

93            images = self.model.autoencoder_decode(x)

#

保存图片

96        save_images(images, dest_path, 'img_')

#

CLI

99def main():

#

103    parser = argparse.ArgumentParser()
104
105    parser.add_argument(
106        "--prompt",
107        type=str,
108        nargs="?",
109        default="a painting of a cute monkey playing guitar",
110        help="the prompt to render"
111    )
112
113    parser.add_argument(
114        "--orig-img",
115        type=str,
116        nargs="?",
117        help="path to the input image"
118    )
119
120    parser.add_argument("--batch_size", type=int, default=4, help="batch size", )
121    parser.add_argument("--steps", type=int, default=50, help="number of ddim sampling steps")
122
123    parser.add_argument("--scale", type=float, default=5.0,
124                        help="unconditional guidance scale: "
125                             "eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))")
126
127    parser.add_argument("--strength", type=float, default=0.75,
128                        help="strength for noise: "
129                             " 1.0 corresponds to full destruction of information in init image")
130
131    opt = parser.parse_args()
132    set_seed(42)
133
134    img2img = Img2Img(checkpoint_path=lab.get_data_path() / 'stable-diffusion' / 'sd-v1-4.ckpt',
135                      ddim_steps=opt.steps)
136
137    with monit.section('Generate'):
138        img2img(
139            dest_path='outputs',
140            orig_img=opt.orig_img,
141            strength=opt.strength,
142            batch_size=opt.batch_size,
143            prompt=opt.prompt,
144            uncond_scale=opt.scale)

#

148if __name__ == "__main__":
149    main()