stable diffusion文生图代码解读

使用diffusers运行stable diffusion,文生图过程代码解读。

只按照下面这种最简单的运行代码,省略了一些参数的处理步骤。

python 复制代码
from diffusers import DiffusionPipeline
pipeline = DiffusionPipeline.from_pretrained(MODEL_PATH , torch_dtype=torch.float16)
pipeline.to("cuda")
img = pipeline("An image of a squirrel in Picasso style",num_inference_steps=10).images[0]
img.save("result.jpg")

0. 定义height 和 width

如果没有输入参数,默认为unet的采样大小乘以VAE缩放率,64*8=512。

1.检查输入的参数

一些常规检查。

2. 定义参数 batch_size

根据prompt或者prompt_embeds计算batch_size,按照上面的执行代码,默认为1。如果一次输入多个prompt,那么就是prompt的数量

python 复制代码
        if prompt is not None and isinstance(prompt, str):
            batch_size = 1
        elif prompt is not None and isinstance(prompt, list):
            batch_size = len(prompt)
        else:
            batch_size = prompt_embeds.shape[0]
 
#多个prompt
#每个prompt生成的图片数量使用num_images_per_prompt控制
prompt = ["An image of a squirrel in Picasso style","Astronaut in a jungle, cold color palette"]
images = pipeline(prompt,num_images_per_prompt=1,num_inference_steps=10).images

3.对输入的prompt编码

默认使用CLIPTokenizer对输入prompt tokenize,输出为(1,77),CLIP模型默认设置最大文本长度为75,然后还有两个表示开始和结束的特殊字符'<|startoftext|>' '<|endoftext|>',最大长度就是77。

使用openai/clip-vit-large-patch14,对输入进行encoder。CLIP模型的默认embedding dim 为768,那么编码输出的prompt embedding的维度就是(1,77,768)。

如果参数没有输入negative_prompt,那么negative_prompt默认为 '''',仍然可以tokenizer,encoder。

negative prompt embedding的维度也是(1,77,768)。

默认都是有do_classifier_free_guidance(CFG参数),为了避免计算两次,这里把negative prompt 和prompt合并在一起输入。

python 复制代码
        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
            prompt,
            device,
            num_images_per_prompt,
            self.do_classifier_free_guidance,
            negative_prompt,
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            lora_scale=lora_scale,
            clip_skip=self.clip_skip,
        )
        if self.do_classifier_free_guidance:
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) #(2,77,768)

4. 准备timesteps

根据使用的scheduler,计算timesteps。stable diffusion默认使用PNDMScheduler,输入的参数num_inference_steps为10步,那么timesteps的长度就为10。

python 复制代码
        timesteps, num_inference_steps = retrieve_timesteps(
            self.scheduler, num_inference_steps, device, timesteps, sigmas
        )

5.准备latents

SD的主要计算都是在latent space进行,以加快计算速度。简单理解就是在小图计算再放大(并不准确)。

unet.config.in_channels 为4,latents的height 和width分别为输入参数height 和width 整除 VAE的缩放率,也就是 512 // 8 = 64,生成的latents的shape为 (1,4,64,64)。

latents使用了 torch.randn 生成。

python 复制代码
        num_channels_latents = self.unet.config.in_channels
        latents = self.prepare_latents(
            batch_size * num_images_per_prompt,
            num_channels_latents,
            height,
            width,
            prompt_embeds.dtype,
            device,
            generator,
            latents,
        )
latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype, layout=layout).to(device)

6.一些其他参数处理

7.逆扩散,去除噪音

默认使用CFG,那么输入的letents也要复制一遍,和之前的prompt_embeds一起输入到UNet去预测噪声。那么得到的噪声也是两个,分别是无条件(negative_prompt) 噪声和 条件(prompt)噪声。

CFG也是在这里起作用,CFG值越大,那么prompt对预测的噪声影响越大,那么对生成的图像影响也越大。

python 复制代码
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)

然后再根据scheduler的算法,计算前一次的latents(去噪),生成新的latents。

循环执行10次,得到最终的latents。

最后使用VAE的解码部分,将latents还原为图片。

python 复制代码
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                if self.interrupt:
                    continue
                # expand the latents if we are doing classifier free guidance
                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

                # predict the noise residual
                noise_pred = self.unet(
                    latent_model_input,
                    t,
                    encoder_hidden_states=prompt_embeds,
                    timestep_cond=timestep_cond,
                    cross_attention_kwargs=self.cross_attention_kwargs,
                    added_cond_kwargs=added_cond_kwargs,
                    return_dict=False,
                )[0]
                # perform guidance
                if self.do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
                # compute the previous noisy sample x_t -> x_t-1
                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
     if not output_type == "latent":
            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[0]
相关推荐
終不似少年遊*2 分钟前
美国加州房价数据分析01
人工智能·python·机器学习·数据挖掘·数据分析·回归算法
区块链小八歌20 分钟前
链原生 Web3 AI 网络 Chainbase 推出 AVS 主网, 拓展 EigenLayer AVS 场景
人工智能
禾高网络22 分钟前
租赁小程序成品|租赁系统搭建核心功能
java·人工智能·小程序
湫ccc2 小时前
《Opencv》基础操作详解(3)
人工智能·opencv·计算机视觉
Jack_pirate2 小时前
深度学习中的特征到底是什么?
人工智能·深度学习
微凉的衣柜2 小时前
微软在AI时代的战略布局和挑战
人工智能·深度学习·microsoft
GocNeverGiveUp2 小时前
机器学习1-简单神经网络
人工智能·机器学习
Schwertlilien2 小时前
图像处理-Ch2-空间域的图像增强
人工智能
智慧化智能化数字化方案3 小时前
深入解读数据资产化实践指南(2024年)
大数据·人工智能·数据资产管理·数据资产入表·数据资产化实践指南
哦哦~9213 小时前
深度学习驱动的油气开发技术与应用
大数据·人工智能·深度学习·学习