學(xué)習(xí)前言
Inpaint是Stable Diffusion中的常用方法,一起簡單學(xué)習(xí)一下。
源碼下載地址
https://github.com/bubbliiiing/stable-diffusion
喜歡的可以點個star噢。
原理解析
一、先驗知識
txt2img的原理如博文
Diffusion擴散模型學(xué)習(xí)2——Stable Diffusion結(jié)構(gòu)解析-以文本生成圖像(文生圖,txt2img)為例
img2img的原理如博文
Diffusion擴散模型學(xué)習(xí)3——Stable Diffusion結(jié)構(gòu)解析-以圖像生成圖像(圖生圖,img2img)為例
二、什么是inpaint
Inpaint是一項圖片修復(fù)技術(shù),可以從圖片上去除不必要的物體,讓您輕松擺脫照片上的水印、劃痕、污漬、標志等瑕疵。
一般來講,圖片的inpaint過程可以理解為兩步:
1、找到圖片中的需要重繪的部分,比如上述提到的水印、劃痕、污漬、標志等。
2、去掉水印、劃痕、污漬、標志等,自動填充圖片應(yīng)該有的內(nèi)容。
三、Stable Diffusion中的inpaint
Stable Diffusion中的inpaint的實現(xiàn)方式有兩種:
1、開源的inpaint模型
參考鏈接:inpaint_st.py,該模型經(jīng)過特定的訓(xùn)練。需要輸入符合需求的圖片才可以進行inpaint。
需要注意的是,該模型使用的config文件發(fā)生了改變,改為v1-inpainting-inference.yaml。其中最顯著的區(qū)別就是unet_config的in_channels從4變成了9。相比于原來的4,我們增加了4+1(5)個通道的信息。
4+1(5)個通道的信息應(yīng)該是什么呢?一個是被mask后的圖像,對應(yīng)其中的4;一個是mask的圖像,對應(yīng)其中的1。
- 1、我們首先把圖片中需要inpaint的部分給置為0,獲得被mask后的圖像,然后利用VAE編碼,VAE輸出通道為4,假設(shè)被mask的圖像是[512, 512, 3],此時我們獲得了一個[4, 64, 64]的隱含層特征,對應(yīng)其中的4。
- 2、然后需要對mask進行下采樣,采樣到和隱含層特征一樣的高寬,即mask的shape為[1, 512, 512],利用下采樣獲得[1, 64, 64]的mask。本質(zhì)上,我們獲得了隱含層的mask。
- 3、然后我們將 下采樣后的被mask的圖像 和 隱含層的mask 在通道上做一個堆疊,獲得一個[5, 64, 64]的特征,然后將此特征與隨機初始化的高斯噪聲堆疊,則獲得了上述圖片中的9通道特征。
此后采樣的過程與常規(guī)采樣方式一樣,全部采樣完成后,使用VAE解碼,獲得inpaint后的圖像。
可以感受到上述的方式必須基于一個已經(jīng)訓(xùn)練好的unet模型,這要求訓(xùn)練者需要有足夠的算力去完成這一個工作,對大眾開發(fā)者而言并不友好。因此該方法很少在實際中得到使用。
2、基于base模型inpaint
如果我們必須訓(xùn)練一個inpaint模型才能對當(dāng)前的模型進行inpaint,那就太麻煩了,有沒有什么方法可以不需要訓(xùn)練就能inpaint呢?
誒誒,當(dāng)然有哈。
Stable Diffusion就是一個生成模型,如果我們可以做到讓Stable Diffusion只生成指定區(qū)域,并且在生成指定區(qū)域的時候參考其它區(qū)域,那么它自身便是一個天然的inpaint模型。
如何做到這一點呢?我們需要結(jié)合img2img方法,我們首先考慮inpaint的兩個輸入:一個是原圖,另外一個是mask圖。
在img2img中,存在一個denoise參數(shù),假設(shè)我們設(shè)置denoise數(shù)值為0.8,總步數(shù)為20步,那么我們會對輸入圖片進行0.8x20次的加噪聲。如果我們可以在這個加噪聲圖片的基礎(chǔ)上進行重建,那么網(wǎng)絡(luò)必然會考慮加噪聲圖(也就對應(yīng)了原始圖片的特征)。
在圖像重建的20步中,對隱含層特征,我們利用mask將不重建的地方都替換成 原圖按照當(dāng)前步數(shù)加噪后的隱含層特征。此時不重建的地方的特征都由輸入圖片決定。然后不替換需要重建的地方進行,利用unet計算噪聲進行重建。
具體部分,可看下面的循環(huán)與代碼,我已經(jīng)標注出了 替換特征的地方,在這里mask等于1的地方保留原圖,mask等于0的地方不斷的重建。
- 將原圖x0映射到VAE隱空間,得到img_orig;
- 初始化隨機噪聲img(也可以使用img_orig完全加噪后的噪聲);
- 開始循環(huán):
- 對于每一次時間步,根據(jù)時間步生成img_orig對應(yīng)的噪聲特征;
- 一個是基于上個時間步降噪后得到的img,一個是基于原圖得到的img_orig。通過mask將兩者融合, i m g = i m g _ o r i g ? m a s k + ( 1.0 ? m a s k ) ? i m g img = img\_orig * mask + (1.0 - mask) * img img=img_orig?mask+(1.0?mask)?img。即,將原圖中的非mask區(qū)域和噪聲圖中的mask區(qū)域進行融合,得到新的噪聲圖。
- 然后繼續(xù)去噪聲直到結(jié)束。
由于該方法不需要訓(xùn)練新模型,并且重建效果也不錯,所以該方法比較通用。
for i, step in enumerate(iterator):
# index是用來取得對應(yīng)的調(diào)節(jié)參數(shù)的
index = total_steps - i - 1
# 將步數(shù)拓展到bs維度
ts = torch.full((b,), step, device=device, dtype=torch.long)
# --------------------------------------------------------------------------------- #
# 替換特征的地方
# 用于進行局部的重建,對部分區(qū)域的隱向量進行mask。
# 對傳入unet前的隱含層特征,我們利用mask將不重建的地方都替換成 原圖加噪后的隱含層特征
# self.model.q_sample用于對輸入圖片進行ts步數(shù)的加噪
# --------------------------------------------------------------------------------- #
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)
img, pred_x0 = outs
# 回調(diào)函數(shù)
if callback: callback(i)
if img_callback: img_callback(pred_x0, i)
if index % log_every_t == 0 or index == total_steps - 1:
intermediates['x_inter'].append(img)
intermediates['pred_x0'].append(pred_x0)
四、inpaint流程
根據(jù)通用性,本文主要以上述提到的基于base模型inpaint進行解析。
1、輸入圖片到隱空間的編碼
inpaint技術(shù)衍生于圖生圖技術(shù),所以同樣需要指定一張參考的圖像,然后在這個參考圖像上開始工作。
利用VAE編碼器對這張參考圖像進行編碼,使其進入隱空間,只有進入了隱空間,網(wǎng)絡(luò)才知道這個圖像是什么。
此時我們便獲得在隱空間的圖像,后續(xù)會在這個 隱空間加噪后的圖像 的基礎(chǔ)上進行采樣。
2、文本編碼
文本編碼的思路比較簡單,直接使用CLIP的文本編碼器進行編碼就可以了,在代碼中定義了一個FrozenCLIPEmbedder類別,使用了transformers庫的CLIPTokenizer和CLIPTextModel。
在前傳過程中,我們對輸入進來的文本首先利用CLIPTokenizer進行編碼,然后使用CLIPTextModel進行特征提取,通過FrozenCLIPEmbedder,我們可以獲得一個[batch_size, 77, 768]的特征向量。
class FrozenCLIPEmbedder(AbstractEncoder):
"""Uses the CLIP transformer encoder for text (from huggingface)"""
LAYERS = [
"last",
"pooled",
"hidden"
]
def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77,
freeze=True, layer="last", layer_idx=None): # clip-vit-base-patch32
super().__init__()
assert layer in self.LAYERS
# 定義文本的tokenizer和transformer
self.tokenizer = CLIPTokenizer.from_pretrained(version)
self.transformer = CLIPTextModel.from_pretrained(version)
self.device = device
self.max_length = max_length
# 凍結(jié)模型參數(shù)
if freeze:
self.freeze()
self.layer = layer
self.layer_idx = layer_idx
if layer == "hidden":
assert layer_idx is not None
assert 0 <= abs(layer_idx) <= 12
def freeze(self):
self.transformer = self.transformer.eval()
# self.train = disabled_train
for param in self.parameters():
param.requires_grad = False
def forward(self, text):
# 對輸入的圖片進行分詞并編碼,padding直接padding到77的長度。
batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
# 拿出input_ids然后傳入transformer進行特征提取。
tokens = batch_encoding["input_ids"].to(self.device)
outputs = self.transformer(input_ids=tokens, output_hidden_states=self.layer=="hidden")
# 取出所有的token
if self.layer == "last":
z = outputs.last_hidden_state
elif self.layer == "pooled":
z = outputs.pooler_output[:, None, :]
else:
z = outputs.hidden_states[self.layer_idx]
return z
def encode(self, text):
return self(text)
3、采樣流程
a、生成初始噪聲
在inpaint中,我們的初始噪聲獲取于參考圖片,參考第一步獲得Latent特征后,使用該Latent特征基于DDIM Sampler進行加噪,獲得輸入圖片加噪后的特征。
此處先不引入denoise參數(shù),所以直接20步噪聲加到底。在該步,我們執(zhí)行了下面兩個操作:
- 將原圖x0映射到VAE隱空間,得到img_orig;
- 初始化隨機噪聲img(也可以使用img_orig完全加噪后的噪聲);
b、對噪聲進行N次采樣
我們便從上一步獲得的初始特征開始去噪聲。
我們會對ddim_timesteps的時間步取反,因為我們現(xiàn)在是去噪聲而非加噪聲,然后對其進行一個循環(huán),循環(huán)的代碼如下:
循環(huán)中有一個mask,它的作用是用于進行局部的重建,對部分區(qū)域的隱向量進行mask,在此前我們并未用到,這一次我們需要用到了。
- 對于每一次時間步,根據(jù)時間步生成img_orig對應(yīng)的加噪聲特征;
- 一個是基于上個時間步降噪后得到的img;一個是基于原圖得到的img_orig。我們通過mask將兩者融合, i m g = i m g _ o r i g ? m a s k + ( 1.0 ? m a s k ) ? i m g img = img\_orig * mask + (1.0 - mask) * img img=img_orig?mask+(1.0?mask)?img。即,將原圖中的非mask區(qū)域和噪聲圖中的mask區(qū)域進行融合,得到新的噪聲圖。
- 然后繼續(xù)去噪聲直到結(jié)束。
for i, step in enumerate(iterator):
# index是用來取得對應(yīng)的調(diào)節(jié)參數(shù)的
index = total_steps - i - 1
# 將步數(shù)拓展到bs維度
ts = torch.full((b,), step, device=device, dtype=torch.long)
# --------------------------------------------------------------------------------- #
# 替換特征的地方
# 用于進行局部的重建,對部分區(qū)域的隱向量進行mask。
# 對傳入unet前的隱含層特征,我們利用mask將不重建的地方都替換成 原圖加噪后的隱含層特征
# self.model.q_sample用于對輸入圖片進行ts步數(shù)的加噪
# --------------------------------------------------------------------------------- #
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)
img, pred_x0 = outs
# 回調(diào)函數(shù)
if callback: callback(i)
if img_callback: img_callback(pred_x0, i)
if index % log_every_t == 0 or index == total_steps - 1:
intermediates['x_inter'].append(img)
intermediates['pred_x0'].append(pred_x0)
return img, intermediates
c、如何引入denoise
上述代碼是官方自帶的基于base模型的可用于inpaint的代碼,但問題在于并未考慮denoise參數(shù)。
假設(shè)我們對生成圖像的某一區(qū)域不滿意,但是不滿意的不多,其實我們不需要完全進行重建,只需要重建一點點就行了,那么此時我們便需要引入denoise參數(shù),表示我們要重建的強度。
i、加噪的邏輯
同樣,我們的初始噪聲獲取于參考圖片,參考第一步獲得Latent特征后,使用該Latent特征和denoise參數(shù)基于DDIM Sampler進行加噪,獲得輸入圖片加噪后的特征。
加噪的邏輯如下:
- denoise可認為是重建的比例,1代表全部重建,0代表不重建;
- 假設(shè)我們設(shè)置denoise數(shù)值為0.8,總步數(shù)為20步;我們會對輸入圖片進行0.8x20次的加噪聲,剩下4步不加,可理解為80%的特征,保留20%的特征;不過就算加完20步噪聲,原始輸入圖片的信息還是有一點保留的,不是完全不保留。
with torch.no_grad():
if seed == -1:
seed = random.randint(0, 65535)
seed_everything(seed)
# ----------------------- #
# 對輸入圖片進行編碼并加噪
# ----------------------- #
if image_path is not None:
img = HWC3(np.array(img, np.uint8))
img = torch.from_numpy(img.copy()).float().cuda() / 127.0 - 1.0
img = torch.stack([img for _ in range(num_samples)], dim=0)
img = einops.rearrange(img, 'b h w c -> b c h w').clone()
if vae_fp16:
img = img.half()
model.first_stage_model = model.first_stage_model.half()
else:
model.first_stage_model = model.first_stage_model.float()
ddim_sampler.make_schedule(ddim_steps, ddim_eta=eta, verbose=True)
t_enc = min(int(denoise_strength * ddim_steps), ddim_steps - 1)
# 獲得VAE編碼后的隱含層向量
z = model.get_first_stage_encoding(model.encode_first_stage(img))
x0 = z
# 獲得加噪后的隱含層向量
z_enc = ddim_sampler.stochastic_encode(z, torch.tensor([t_enc] * num_samples).to(model.device))
z_enc = z_enc.half() if sd_fp16 else z_enc.float()
ii、mask處理
我們需要對mask進行下采樣,使其和上述獲得的加噪后的特征的shape一樣。
if mask_path is not None:
mask = torch.from_numpy(mask).to(model.device)
mask = torch.nn.functional.interpolate(mask, size=z_enc.shape[-2:])
iii、采樣處理
此時,因為使用到了denoise參數(shù),我們要基于img2img中的decode方法進行采樣。
由于decode方法中不存在mask與x0參數(shù),我們補一下:
@torch.no_grad()
def decode(self, x_latent, cond, t_start, mask, x0, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
use_original_steps=False):
# 使用ddim的時間步
# 這里內(nèi)容看起來很多,但是其實很少,本質(zhì)上就是取了self.ddim_timesteps,然后把它reversed一下
timesteps = np.arange(self.ddpm_num_timesteps) if use_original_steps else self.ddim_timesteps
timesteps = timesteps[:t_start]
time_range = np.flip(timesteps)
total_steps = timesteps.shape[0]
print(f"Running DDIM Sampling with {total_steps} timesteps")
iterator = tqdm(time_range, desc='Decoding image', total=total_steps)
x_dec = x_latent
for i, step in enumerate(iterator):
index = total_steps - i - 1
ts = torch.full((x_latent.shape[0],), step, device=x_latent.device, dtype=torch.long)
# --------------------------------------------------------------------------------- #
# 替換特征的地方
# 用于進行局部的重建,對部分區(qū)域的隱向量進行mask。
# 對傳入unet前的隱含層特征,我們利用mask將不重建的地方都替換成 原圖加噪后的隱含層特征
# self.model.q_sample用于對輸入圖片進行ts步數(shù)的加噪
# --------------------------------------------------------------------------------- #
if mask is not None:
assert x0 is not None
img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass?
x_dec = img_orig * mask + (1. - mask) * x_dec
# 進行單次采樣
x_dec, _ = self.p_sample_ddim(x_dec, cond, ts, index=index, use_original_steps=use_original_steps,
unconditional_guidance_scale=unconditional_guidance_scale,
unconditional_conditioning=unconditional_conditioning)
return x_dec
4、隱空間解碼生成圖片
通過上述步驟,已經(jīng)可以多次采樣獲得結(jié)果,然后我們便可以通過隱空間解碼生成圖片。
隱空間解碼生成圖片的過程非常簡單,將上文多次采樣后的結(jié)果,使用decode_first_stage方法即可生成圖片。
在decode_first_stage方法中,網(wǎng)絡(luò)調(diào)用VAE對獲取到的64x64x3的隱向量進行解碼,獲得512x512x3的圖片。文章來源:http://www.zghlxwxcb.cn/news/detail-692050.html
@torch.no_grad()
def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False):
if predict_cids:
if z.dim() == 4:
z = torch.argmax(z.exp(), dim=1).long()
z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None)
z = rearrange(z, 'b h w c -> b c h w').contiguous()
z = 1. / self.scale_factor * z
# 一般無需分割輸入,所以直接將x_noisy傳入self.model中,在下面else進行
if hasattr(self, "split_input_params"):
......
else:
if isinstance(self.first_stage_model, VQModelInterface):
return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize)
else:
return self.first_stage_model.decode(z)
Inpaint預(yù)測過程代碼
整體預(yù)測代碼如下:文章來源地址http://www.zghlxwxcb.cn/news/detail-692050.html
import os
import random
import cv2
import einops
import numpy as np
import torch
from PIL import Image
from pytorch_lightning import seed_everything
from ldm_hacked import *
# ----------------------- #
# 使用的參數(shù)
# ----------------------- #
# config的地址
config_path = "model_data/sd_v15.yaml"
# 模型的地址
model_path = "model_data/v1-5-pruned-emaonly.safetensors"
# fp16,可以加速與節(jié)省顯存
sd_fp16 = True
vae_fp16 = True
# ----------------------- #
# 生成圖片的參數(shù)
# ----------------------- #
# 生成的圖像大小為input_shape,對于img2img會進行Centter Crop
input_shape = [512, 768]
# 一次生成幾張圖像
num_samples = 1
# 采樣的步數(shù)
ddim_steps = 20
# 采樣的種子,為-1的話則隨機。
seed = 12345
# eta
eta = 0
# denoise強度,for img2img
denoise_strength = 1.00
# ----------------------- #
# 提示詞相關(guān)參數(shù)
# ----------------------- #
# 提示詞
prompt = "a cute dog, with yellow leaf, trees"
# 正面提示詞
a_prompt = "best quality, extremely detailed"
# 負面提示詞
n_prompt = "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
# 正負擴大倍數(shù)
scale = 9
# img2img使用,如果不想img2img這設(shè)置為None。
image_path = "imgs/test_imgs/cat.jpg"
# inpaint使用,如果不想inpaint這設(shè)置為None;inpaint使用需要結(jié)合img2img。
# 注意mask圖和原圖需要一樣大
mask_path = "imgs/test_imgs/cat_mask.jpg"
# ----------------------- #
# 保存路徑
# ----------------------- #
save_path = "imgs/outputs_imgs"
# ----------------------- #
# 創(chuàng)建模型
# ----------------------- #
model = create_model(config_path).cpu()
model.load_state_dict(load_state_dict(model_path, location='cuda'), strict=False)
model = model.cuda()
ddim_sampler = DDIMSampler(model)
if sd_fp16:
model = model.half()
if image_path is not None:
img = Image.open(image_path)
img = crop_and_resize(img, input_shape[0], input_shape[1])
if mask_path is not None:
mask = Image.open(mask_path).convert("L")
mask = crop_and_resize(mask, input_shape[0], input_shape[1])
mask = np.array(mask)
mask = mask.astype(np.float32) / 255.0
mask = mask[None,None]
mask[mask < 0.5] = 0
mask[mask >= 0.5] = 1
with torch.no_grad():
if seed == -1:
seed = random.randint(0, 65535)
seed_everything(seed)
# ----------------------- #
# 對輸入圖片進行編碼并加噪
# ----------------------- #
if image_path is not None:
img = HWC3(np.array(img, np.uint8))
img = torch.from_numpy(img.copy()).float().cuda() / 127.0 - 1.0
img = torch.stack([img for _ in range(num_samples)], dim=0)
img = einops.rearrange(img, 'b h w c -> b c h w').clone()
if vae_fp16:
img = img.half()
model.first_stage_model = model.first_stage_model.half()
else:
model.first_stage_model = model.first_stage_model.float()
ddim_sampler.make_schedule(ddim_steps, ddim_eta=eta, verbose=True)
t_enc = min(int(denoise_strength * ddim_steps), ddim_steps - 1)
# 獲得VAE編碼后的隱含層向量
z = model.get_first_stage_encoding(model.encode_first_stage(img))
x0 = z
# 獲得加噪后的隱含層向量
z_enc = ddim_sampler.stochastic_encode(z, torch.tensor([t_enc] * num_samples).to(model.device))
z_enc = z_enc.half() if sd_fp16 else z_enc.float()
if mask_path is not None:
mask = torch.from_numpy(mask).to(model.device)
mask = torch.nn.functional.interpolate(mask, size=z_enc.shape[-2:])
mask = 1 - mask
# ----------------------- #
# 獲得編碼后的prompt
# ----------------------- #
cond = {"c_crossattn": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)]}
un_cond = {"c_crossattn": [model.get_learned_conditioning([n_prompt] * num_samples)]}
H, W = input_shape
shape = (4, H // 8, W // 8)
if image_path is not None:
samples = ddim_sampler.decode(z_enc, cond, t_enc, mask, x0, unconditional_guidance_scale=scale, unconditional_conditioning=un_cond)
else:
# ----------------------- #
# 進行采樣
# ----------------------- #
samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
shape, cond, verbose=False, eta=eta,
unconditional_guidance_scale=scale,
unconditional_conditioning=un_cond)
# ----------------------- #
# 進行解碼
# ----------------------- #
x_samples = model.decode_first_stage(samples.half() if vae_fp16 else samples.float())
x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
# ----------------------- #
# 保存圖片
# ----------------------- #
if not os.path.exists(save_path):
os.makedirs(save_path)
for index, image in enumerate(x_samples):
cv2.imwrite(os.path.join(save_path, str(index) + ".jpg"), cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
到了這里,關(guān)于AIGC專欄4——Stable Diffusion原理解析-inpaint修復(fù)圖片為例的文章就介紹完了。如果您還想了解更多內(nèi)容,請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章,希望大家以后多多支持TOY模板網(wǎng)!