AIGC-Stable Diffusion进阶1(附代码)

在上篇文章中对Stable Diffusion举行了开端的认识，也给出了使用案例，这篇文章将进一步的去刨析一下SD模子。
SD模子的构成

在之前也讲过了，SD由CLIP、UNet、Schedule，VAE构成。再来回首一下知识点：
1.CLIP Text Encoder(文本编码器)

作用：将文本信息举行编码天生对应的特性矩阵方便输入到SD模子中。
2.VAE Encoder(变分自编码器)

作用：天生Latent Feature(隐空间特性)和文本特性同时作为模子输入。假如是图生图任务则将图像举行编码天生Latent Feature；假如是文生图则用随机天生的高斯噪声矩阵作为Latent Feature作为输入。【也就是在输入SD模子前有两个输入，文本特性和隐空间特性】
3.U-Net网络

作用：用于不绝地推测噪声，并在每次推测噪声过程中到场文本语义特性。
4.Schedule

作用：对UNet推测的噪声举行优化处置惩罚(动态调解推测的噪声，控制U-Net推测噪声的强度)
5.VAE Decoder(解码器)

作用：将终极得到的Latent Feature颠末解码器天生图像。
在SD的迭代过程中(去噪过程)，噪声会不绝的镌汰，图像信息和文本语义信息会不绝的增长。

---

也就是说Unet的输入有两部分构成，一个是latent、一个是输入文本特性(prompt)。
1.图生图：须要将图像通过VAE encoder天生latent，并与文本特性举行输入
2.文生图：使用随机高斯噪声取代图像latent，并与文本特性举行输入。
在上一篇文章中可以教了各人怎样快速使用diffusers搭建SD模子。但这种方法不敷机动，我们对每个部分举行分析，在这篇文章中会分模块加载Model并手动搭建SD模子。
这里以文生图为例。(文末会附上完备的代码)

---

CLIP Model单独加载

在天生文本特性的时间须要有两个步调：1.prompt->token 2.token->text_embedding

2.     # 1.prompt特征提取
3.     # 1.1 prompt转token
4.     text_tokenizer = CLIPTokenizer.from_pretrained(sd_weight_path, subfolder="tokenizer")
5.     text_token_ids = text_tokenizer(
6.         prompt,
7.         padding='max_length',
8.         max_length=text_tokenizer.model_max_length,
9.         truncation=True,
10.         return_tensors='pt'
11.     ).input_ids
13.    # 1.2提取prompt_embedding
14.     text_encoder = CLIPTextModel.from_pretrained(sd_weight_path,subfolder='text_encoder')
15.     text_embeddings = text_encoder(text_token_ids)[0]

复制代码

Schedule单独加载

Schedule是和UNet相团结的，用于优化差别时间步t时间的latent。
大概可以如许明白，差别时间t对应差别的噪声级别和埋伏表现。缩放输入可以确保输入在每个时间步都是尺度化的，使其对应当前时间t所对应的噪声分布，使模子可以大概稳固地处置惩罚差别程度的噪声。
这里我的明白感觉和CNN中的BN层有点相似，但也有差别，这里我谈一下自己的明白：
扩散模子中的schedule,对输入举行缩放的目的是控制和调解差别时间t对应的噪声程度，使得latent在每个时间t的数据分布和训练时是一样的，从而可以稳固的去噪。也就是和BN雷同，都是对输入举行尺度化大概缩放使数据在特定分布下更稳固。但差别的是，BN层是为相识决梯度消散和爆炸，加快训练速率，而扩散模子的缩放是每个时间t对latent的缩放，为的是控制噪声程度。

2. # 2.单独加载调度算法，生成噪声表
3. scheduler = PNDMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
4. scheduler.set_timesteps(25)
5. timesteps = scheduler.timesteps

复制代码

UNet单独加载

2. # 3.单独加载UNet
3. unet = UNet2DConditionModel.from_pretrained(sd_weight_path,subfolder='unet')

复制代码

UNet推测噪声 

UNet做噪声推测的时间代码如下，此中latents_model_input是颠末Schedule优化后的特性图，t是对应的时间t,text_embeddings是输入的文本特性。
输出是推测噪声，然后对噪声举行分块(chunk操纵)得到两个噪声，一个是无条件噪声，一个是条件噪声。
无条件噪声：
无条件噪声是指模子在没有任何(如文本、标签等)输入环境下，单纯的基于latents得到噪声，这部分噪声推测代表了天生过程中的根本噪声程度。
条件噪声：
条件噪声推测是模子在给定某种条件（如文本嵌入）输入的环境下，基于埋伏变量和这些条件推测的噪声。这部分噪声推测代表了天生过程在特定条件下的噪声程度。

2. noise_pred = unet(
3.             latents_model_input,
4.             t,
5.             encoder_hidden_states=text_embeddings,
6.             cross_attention_kwargs=None,
7.             return_dict=False,
8.         )[0]  # 预测噪声
9.         noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
10.         # 7.5为相关因子
11.         noise_pred = noise_pred_uncond + 7.5 * (noise_pred_text - noise_pred_uncond)

复制代码

VAE单独加载

2.     # 4.单独加载VAE
3.     # 进行解码
4.     vae = AutoencoderKL.from_pretrained(sd_weight_path, subfolder='vae').to("cuda", dtype=torch.float32)

复制代码

---

去噪可视化

比如这里输入的prompt为：'a girl, beautiful'，neg pormpt为空时，对每次去噪后的latents可视化，代码如下：

1.     # 3.单独加载UNet    unet = UNet2DConditionModel.from_pretrained(sd_weight_path,subfolder='unet')    for i, t in tqdm(enumerate(timesteps)):        latents_model_input = torch.cat([latents] * 2)        latents_model_input = scheduler.scale_model_input(latents_model_input, t)        
2. noise_pred = unet(
3.             latents_model_input,
4.             t,
5.             encoder_hidden_states=text_embeddings,
6.             cross_attention_kwargs=None,
7.             return_dict=False,
8.         )[0]  # 预测噪声
9.         noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
10.         # 7.5为相关因子
11.         noise_pred = noise_pred_uncond + 7.5 * (noise_pred_text - noise_pred_uncond)        # x_t->x_t-1  去噪操纵        latents = scheduler.step(noise_pred, t, latents,return_dict=False)[0]        save_noise_pred_feat(latents, i)

复制代码

可视化结果(部分图)：
latents可视化

latents解码

2. # 进行解码
3. vae = AutoencoderKL.from_pretrained(sd_weight_path, subfolder='vae').to("cuda", dtype=torch.float32)
4. image = vae.decode(latents.to("cuda", dtype=torch.float32) / vae.config.scaling_factor, return_dict=False)[0]

复制代码

终极解码的结果图如下： 
(大要照旧可以，就是可以看到多了一只手，我这里neg prompt为空)

VAE解码 比如这里输入的neg_prompt为：

2. neg_prompt = 'extra limbs, extra arms, extra hands, disfigured, distorted anatomy, unnatural proportions, unrealistic, malformed'

复制代码

天生图像为：

 假云云时输入的neg_prompt为：

2. neg_prompt = 'extra limbs, extra arms, extra hands, disfigured, distorted anatomy, unnatural proportions, unrealistic, malformed, unnatural expression, creepy eyes, unrealistic facial features, disfigured face, distorted expression, awkward smile, disproportionate eyes, unrealistic gaze, malformed eyes'

复制代码

天生图像如下：

---

完备代码

 

1. import torchfrom diffusers.pipelines.alt_diffusion.pipeline_alt_diffusion import rescale_noise_cfgfrom diffusers.utils import randn_tensorfrom transformers import CLIPTokenizer, CLIPTextModelfrom diffusers import UNet2DConditionModel, PNDMScheduler, LMSDiscreteScheduler, AutoencoderKLfrom tqdm import tqdmimport numpy as npfrom PIL import Imageimport cv2'''这是一个最小化实现算法快速学习sd过程'''# 用以生存特性图def save_noise_pred_feat(noise_pred, save_name):    noise_pred = (noise_pred / 2 + 0.5).clamp(0, 1)    # 将Tensor转换为 NumPy 数组    noise_pred_np = noise_pred.detach().cpu().numpy()    # 将通道数从第二个维度移动到末了一个维度    # （条件假设通道数是第二个维度）    noise_pred_np = np.moveaxis(noise_pred_np, 1, -1)    # 反归一化    noise_pred_np = (noise_pred_np * 255).round().astype(np.uint8)    # 假如你想将全部特性图归并成一个图像生存，可以使用 PIL 的拼接功能    # 这里仅作示例    combined_img = Image.fromarray(np.concatenate(noise_pred_np, axis=1))    combined_img.save(f'{save_name}_combined_feature_map.png')@torch.no_grad()def UNET():    # 随机种子    generator = torch.Generator("cpu").manual_seed(1024)    # 权重路径    sd_weight_path = 'F:/Stable_webui/stable-diffusion-webui/models/Stable-diffusion/stable-diffusion-v1-5'    # prompt    prompt = 'a girl, beautiful'    # neg prompt   
2. neg_prompt = 'extra limbs, extra arms, extra hands, disfigured, distorted anatomy, unnatural proportions, unrealistic, malformed, unnatural expression, creepy eyes, unrealistic facial features, disfigured face, distorted expression, awkward smile, disproportionate eyes, unrealistic gaze, malformed eyes'    # 1.prompt特性提取    # 1.1 prompt转token    text_tokenizer = CLIPTokenizer.from_pretrained(sd_weight_path, subfolder="tokenizer")    text_token_ids = text_tokenizer(        prompt,        padding='max_length',        max_length=text_tokenizer.model_max_length,        truncation=True,        return_tensors='pt'    ).input_ids    neg_token_ids = text_tokenizer(        neg_prompt,        padding='max_length',        max_length=77,        truncation=True,        return_tensors='pt'    )    # 1.2提取prompt_embedding    text_encoder = CLIPTextModel.from_pretrained(sd_weight_path,subfolder='text_encoder')    text_embeddings = text_encoder(text_token_ids)[0]    text_embeddings = text_embeddings.repeat(1,1,1)    text_embeddings = text_embeddings.view(1*1,77,-1)    neg_prompt_embeds = text_encoder(        neg_token_ids.input_ids,        attention_mask=None,    )[0]    del text_encoder, text_tokenizer    neg_prompt_embeds = neg_prompt_embeds.repeat(1, 1, 1)    neg_prompt_embeds = neg_prompt_embeds.view(1*1, 77, -1)    # 文本特性向量    text_embeddings = torch.cat([neg_prompt_embeds, text_embeddings])    # 2.单独加载调理算法，天生噪声表，在总的时间步数num_train_timesteps中，这个噪声表的值在beta_start到beta_end中，对应公式的αt    scheduler = PNDMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)    scheduler.set_timesteps(25)    timesteps = scheduler.timesteps    print("timesteps :" + str(timesteps))    # 假如是文生图，那么就用一个随机噪声来取代这个图    latents = randn_tensor((1, 4, 64, 64), generator=generator, dtype=torch.float32)    latents = scheduler.init_noise_sigma * latents    # 3.单独加载UNet    unet = UNet2DConditionModel.from_pretrained(sd_weight_path,subfolder='unet')    for i, t in tqdm(enumerate(timesteps)):        latents_model_input = torch.cat([latents] * 2)        latents_model_input = scheduler.scale_model_input(latents_model_input, t)        
3. noise_pred = unet(
4.             latents_model_input,
5.             t,
6.             encoder_hidden_states=text_embeddings,
7.             cross_attention_kwargs=None,
8.             return_dict=False,
9.         )[0]  # 预测噪声
10.         noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
11.         # 7.5为相关因子
12.         noise_pred = noise_pred_uncond + 7.5 * (noise_pred_text - noise_pred_uncond)        # x_t->x_t-1  去噪操纵        latents = scheduler.step(noise_pred, t, latents,return_dict=False)[0]        # save_noise_pred_feat(latents, i)
13.     # 4.单独加载VAE
14.     # 进行解码
15.     vae = AutoencoderKL.from_pretrained(sd_weight_path, subfolder='vae').to("cuda", dtype=torch.float32)    image = vae.decode(latents.to("cuda", dtype=torch.float32) / vae.config.scaling_factor, return_dict=False)[0]    # 后处置惩罚    rec_image = (image / 2 + 0.5).clamp(0, 1)  # 归一化到0~1    rec_image = rec_image.cpu().permute(0, 2, 3, 1).numpy()    # 反归一化    rec_image = (rec_image * 255).round().astype("uint8")    rec_image = rec_image[0]    # 生存重修后图像    cv2.imwrite("mysd.png", cv2.cvtColor(rec_image, cv2.COLOR_RGB2BGR))    torch.cuda.empty_cache()UNET()

复制代码