- 快速上手
- Pipeline 内部执行步骤
- 后续更新计划
diffusers是Hugging Face推出的一个diffusion库,它提供了简单方便的diffusion推理训练pipe,同时拥有一个模型和数据社区,代码可以像torchhub一样直接从指定的仓库去调用别人上传的数据集和pretrain checkpoint。除此之外,安装方便,代码结构清晰,注释齐全,二次开发会十分有效率。
diffusers使用pipeline类封装各个模块,以及安排其中的调用顺序和输出。
快速上手
以下代码调用StableDiffusionPipeline,from_pretrained
是指从特定仓库加载别人预训练好的模型。其中model_id
可以是本地的路径,如果本地没找到对应的文件,则会自动去Hugging Face的Community中去自动下载。
from diffusers import StableDiffusionPipelinemodel_id = "runwayml/stable-diffusion-v1-5"
stable_diffusion_txt2img = StableDiffusionPipeline.from_pretrained(model
如果我们不想用别人的模型,或者我想在模型基础上进行修改,应该怎么办呢?
这里我们可以先定义其中的模块,包括Unet
、VAE
、CLIP
、Scheduler
等,进行编辑和修改后再输入到pipe中去。如下我们自己预定义好Scheduler:
from diffusers import StableDiffusionPipeline, DDPMSchedulermodel_id = "runwayml/stable-diffusion-v1-5"
scheduler = DDPMScheduler.from_pretrained(model_id)
stable_diffusion_txt2img = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler)
ok,不同的inpainting、txt2img、img2img的pipeline,以及里面不同种类的Unet、VAE、scheduler的替换和自定义我们会在以后单独介绍,现在我们先来看看代码怎么进行推理。
from diffusers import StableDiffusionPipeline
model_id = "runwayml/stable-diffusion-v1-5"
stable_diffusion_txt2img = StableDiffusionPipeline.from_pretrained(model_id)
prompt = "A boy wearing white suspenders and playing basketball"
image = stable_diffusion_txt2img (prompt).images[0]
image.save('generated_image.png')
简单来说,我们只需要把文字prompt输入进去,然后就会按照StableDiffusionPipeline中设置好的流程来进行推理。
Pipeline 内部执行步骤
我们看下StableDiffusionPipeline
源码里的具体推理流程,便于更好理解SD,也便于后续进行修改和自定义。这部分建议在理解Latent Diffusion论文原理后,理解记忆!
# 0. 定义unet中的高宽,这里要考虑经过vae后的缩小系数height = height or self.unet.config.sample_size * self.vae_scale_factorwidth = width or self.unet.config.sample_size * self.vae_scale_factor# 1. 检查输入是否合规self.check_inputs(prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds)# 2. 根据prompt个数定义调用的batch_size,device,以及是否做classifier-free guidanceif prompt is not None and isinstance(prompt, str):batch_size = 1elif prompt is not None and isinstance(prompt, list):batch_size = len(prompt)else:batch_size = prompt_embeds.shape[0]device = self._execution_devicedo_classifier_free_guidance = guidance_scale > 1.0# 3. 将输入的文字prompt进行encodingtext_encoder_lora_scale = (cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None)prompt_embeds = self._encode_prompt(prompt,device,num_images_per_prompt,do_classifier_free_guidance,negative_prompt,prompt_embeds=prompt_embeds,negative_prompt_embeds=negative_prompt_embeds,lora_scale=text_encoder_lora_scale,)# 4. 准备scheduler中的时间步数self.scheduler.set_timesteps(num_inference_steps, device=device)timesteps = self.scheduler.timesteps# 5. 生成对应尺寸的初始噪声图latentsnum_channels_latents = self.unet.config.in_channelslatents = self.prepare_latents(batch_size * num_images_per_prompt,num_channels_latents,height,width,prompt_embeds.dtype,device,generator,latents,)# 6. 额外的去噪参数(eta for DDIM)extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)# 7. 循环多个步长进行去噪num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.orderwith self.progress_bar(total=num_inference_steps) as progress_bar:for i, t in enumerate(timesteps):latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latentslatent_model_input = self.scheduler.scale_model_input(latent_model_input, t)noise_pred = self.unet(latent_model_input,t,encoder_hidden_states=prompt_embeds,cross_attention_kwargs=cross_attention_kwargs,return_dict=False,)[0]if do_classifier_free_guidance:noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)if do_classifier_free_guidance and guidance_rescale > 0.0:noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):progress_bar.update()if callback is not None and i % callback_steps == 0:callback(i, t, latents)if not output_type == "latent":image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)else:image = latentshas_nsfw_concept = Noneif has_nsfw_concept is None:do_denormalize = [True] * image.shape[0]else:do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:self.final_offload_hook.offload()if not return_dict:return (image, has_nsfw_concept)return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
其实也能看出来是怎样一个大概的流程,如果我们后续需要改输入模块,那么直接在pipeline的对应部分重写就可以,十分方便。
后续文章会介绍每个模块的原理特点以及应用场景、如何进行训练和finetune、自定义。
后续更新计划
[代码级手把手解diffusers库析] 1、 快速上手
[代码级手把手解diffusers库析] 2、 Scheduler 介绍 & 代码解析
[代码级手把手解diffusers库析] 3、 DiffusionPipeline原理 & 代码解析
[代码级手把手解diffusers库析] 4、 如何自定义Pipeline进行训练推理