【Diffusers库】第五篇加载pipeline、model、schedulers

Robin C

已于 2024-04-02 23:08:05 修改

阅读量6.2k

点赞数 20

文章标签：人工智能 AIGC stable diffusion

于 2024-04-02 23:06:27 首次发布

本文链接：https://blog.csdn.net/qq_38423499/article/details/137158458

版权

写在前面的话

这是我们研发的用于 消费决策的AI助理 ，我们会持续优化，欢迎体验与反馈。微信扫描二维码，添加即可。
官方链接：https://ailab.smzdm.com/

************************************************************** 分割线 ****************************************************************

官网的这一章节，我觉得有点繁琐，部分内容在前面的博客中有提到，所以这篇梳理总结一下，省些时间。

从不同的位置加载模型

from diffusers import DiffusionPipeline, StableDiffusionPipeline，StableDiffusionImg2ImgPipeline

repo_id = "runwayml/stable-diffusion-v1-5"

# 方法一， 第一次加载会自动下载，随后再次加载会从缓存中加载
pipe = DiffusionPipeline.from_pretrained(repo_id)

# 方法二， 第一次加载会自动下载，随后再次加载会从缓存中加载
pipe = StableDiffusionPipeline.from_pretrained(repo_id)

# 方法三， 第一次加载会自动下载，随后再次加载会从缓存中加载，用于图生图
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(repo_id)

# 方法四， 从指定路径加载
repo_id = "./stable-diffusion-v1-5"
stable_diffusion = DiffusionPipeline.from_pretrained(repo_id)

加载模型的不同部分

调度器

查看可用调度器

在进行图像生成的时候，pipeline有好几个部分，不同的部分产生的作用不一样，因而我们可以替换一下不同的组件，来提升图像生成的质量和效率。

改变调度器，并且权衡图像生成时的速度和质量非常重要。
模型的不同组件通常是独立训练的，您可以将组件换成性能更好的组件。
在微调过程中，通常只训练一些组件，如UNet或文本编码器。

为了检查调度器与加载模型的兼容性，可以使用compatibles函数：

from diffusers import DiffusionPipeline

repo_id = "runwayml/stable-diffusion-v1-5"
stable_diffusion = DiffusionPipeline.from_pretrained(repo_id)
stable_diffusion.scheduler.compatibles

[<class 'diffusers.schedulers.scheduling_lms_discrete.LMSDiscreteScheduler'>,
 <class 'diffusers.schedulers.scheduling_unipc_multistep.UniPCMultistepScheduler'>, 
 <class 'diffusers.utils.dummy_torch_and_torchsde_objects.DPMSolverSDEScheduler'>, 
 <class 'diffusers.schedulers.scheduling_heun_discrete.HeunDiscreteScheduler'>, 
 <class 'diffusers.schedulers.scheduling_ddim.DDIMScheduler'>, 
 <class 'diffusers.schedulers.scheduling_dpmsolver_singlestep.DPMSolverSinglestepScheduler'>, 
 <class 'diffusers.schedulers.scheduling_deis_multistep.DEISMultistepScheduler'>, 
 <class 'diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler'>, 
 <class 'diffusers.schedulers.scheduling_pndm.PNDMScheduler'>, 
 <class 'diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler'>, 
 <class 'diffusers.schedulers.scheduling_k_dpm_2_discrete.KDPM2DiscreteScheduler'>, 
 <class 'diffusers.schedulers.scheduling_ddpm.DDPMScheduler'>, 
 <class 'diffusers.schedulers.scheduling_k_dpm_2_ancestral_discrete.KDPM2AncestralDiscreteScheduler'>, 
 <class 'diffusers.schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteScheduler'>]

这里可以看到这个pipeline可以使用的调度器

更换可用调度器

from diffusers import DiffusionPipeline, EulerDiscreteScheduler, DPMSolverMultistepScheduler
repo_id = "runwayml/stable-diffusion-v1-5"

# 使用EulerDiscreteScheduler中的调度器，替换默认的调度器，但是需要指定文件夹
scheduler = EulerDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
stable_diffusion = DiffusionPipeline.from_pretrained(repo_id, scheduler=scheduler)

安全模型

主要通过from_pretrained方法中的safety_checker参数来设定的。这个主要是防止色情和暴力之类的。

from diffusers import DiffusionPipeline

repo_id = "runwayml/stable-diffusion-v1-5"
stable_diffusion = DiffusionPipeline.from_pretrained(repo_id, safety_checker=None)

重复使用组件

还可以在多个pipeline中，使用相同的组件，以避免重复加载：

from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline

model_id = "runwayml/stable-diffusion-v1-5"
stable_diffusion_txt2img = StableDiffusionPipeline.from_pretrained(model_id)
components = stable_diffusion_txt2img.components

for i in components:
     print(i)

输出了这个pipeline里的不同的组件，这些组件中记录了模型参数、网络结构等等：

vae
text_encoder
tokenizer
unet
scheduler
safety_checker
feature_extractor

然后，可以将组件传递到另一个pipeline，无需将模型重新加载到内存中：

stable_diffusion_img2img = StableDiffusionImg2ImgPipeline(**components)

当然，也可以分开加载：

from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline

model_id = "runwayml/stable-diffusion-v1-5"
stable_diffusion_txt2img = StableDiffusionPipeline.from_pretrained(model_id)
stable_diffusion_img2img = StableDiffusionImg2ImgPipeline(
    vae=stable_diffusion_txt2img.vae,
    text_encoder=stable_diffusion_txt2img.text_encoder,
    tokenizer=stable_diffusion_txt2img.tokenizer,
    unet=stable_diffusion_txt2img.unet,
    scheduler=stable_diffusion_txt2img.scheduler,
    safety_checker=None,
    feature_extractor=None,
    requires_safety_checker=False,
)

模型变种

这个概念可能比较陌生，我的理解就是在预训练模型的基础上，训练出来的模型，叫做模型变体(Checkpoint variants)

模型中的参数的数据类型在进行存储的时候是有选择的，比如：torch.float16，可以获得较低的精度和较低的存储空间。但是，如果使用的是CPU，则不能使用此变体。
Non-exponential mean averaged (EMA) weights 不能进行图例，只能进行微调模型。

加载

另外，还有2个重要的参数，在加载模型的时候要使用，
一个是指定torch_dtype(数据格式)，就是上文中提到的torch.float16，
一个是指定variant(模型变种)，我感觉，这个参数会明确加载模型时的目的，是用于推理 or 训练。

from diffusers import DiffusionPipeline

# load fp16 variant
stable_diffusion = DiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", variant="fp16", torch_dtype=torch.float16
)
# load non_ema variant
stable_diffusion = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", variant="non_ema")

保存

保存的时候，使用 DiffusionPipeline.save_pretrained() 函数，并且要指明 variant 参数是：fp16 or non_ema，应该将模型保存到与原始检查点相同的文件夹中，这样您就可以从同一文件夹加载这两个模型了：

from diffusers import DiffusionPipeline

# save as fp16 variant
stable_diffusion.save_pretrained("runwayml/stable-diffusion-v1-5", variant="fp16")
# save as non-ema variant
stable_diffusion.save_pretrained("runwayml/stable-diffusion-v1-5", variant="non_ema")

如果不将模型保存到现有文件夹中，则必须指定模型参数，否则它将引发异常，找不到模型：

# 👎 this won't work
stable_diffusion = DiffusionPipeline.from_pretrained("./stable-diffusion-v1-5", torch_dtype=torch.float16)
# 👍 this works
stable_diffusion = DiffusionPipeline.from_pretrained(
    "./stable-diffusion-v1-5", variant="fp16", torch_dtype=torch.float16
)

模型和调度器

模型

ModelMixin.from_pretrained() 函数会下载和缓存最新版本的模型。假如本地已经存在，则不会重复下载。
可以使用subfolder参数从子文件夹加载模型。例如，“runwayml/stable-diffusion-v1-5”的模型权重存储在unet子文件夹中：

from diffusers import UNet2DConditionModel

repo_id = "runwayml/stable-diffusion-v1-5"
model = UNet2DConditionModel.from_pretrained(repo_id, subfolder="unet")

或者直接从工程目录下面加载

from diffusers import UNet2DModel

repo_id = "google/ddpm-cifar10-32"
model = UNet2DModel.from_pretrained(repo_id)

也可以通过ModelMixin来加载和保存模型变量： ModelMixin.from_pretrained() 和 ModelMixin.save_pretrained():

from diffusers import UNet2DConditionModel

model = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet", variant="non-ema")
model.save_pretrained("./local-unet", variant="non-ema")

调度器

调度器的加载可以使用 SchedulerMixin.from_pretrained() ，调度器不会被训练。调度器不会消耗大量的内存。下面的代码可以尝试加载不同的调度器。

from diffusers import StableDiffusionPipeline
from diffusers import (
    DDPMScheduler,
    DDIMScheduler,
    PNDMScheduler,
    LMSDiscreteScheduler,
    EulerDiscreteScheduler,
    EulerAncestralDiscreteScheduler,
    DPMSolverMultistepScheduler,
)

repo_id = "runwayml/stable-diffusion-v1-5"

ddpm = DDPMScheduler.from_pretrained(repo_id, subfolder="scheduler")
ddim = DDIMScheduler.from_pretrained(repo_id, subfolder="scheduler")
pndm = PNDMScheduler.from_pretrained(repo_id, subfolder="scheduler")
lms = LMSDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
euler_anc = EulerAncestralDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
euler = EulerDiscreteScheduler.from_pretrained(repo_id, subfolder="scheduler")
dpm = DPMSolverMultistepScheduler.from_pretrained(repo_id, subfolder="scheduler")

# replace `dpm` with any of `ddpm`, `ddim`, `pndm`, `lms`, `euler_anc`, `euler`
pipeline = StableDiffusionPipeline.from_pretrained(repo_id, scheduler=dpm)

关于model_index.json

一般情况下，在hugging-face下的模型工程，都会存在一个model_index.json，用来说明该模型下的工程文件配置。比如：runwayml/stable-diffusion-v1-5。尤其在实例化之后，pipeline中的各个组件的默认使用情况，就和model_index.json中记载的一样。如下文件就是runwayml/stable-diffusion-v1-5的默认指定组件。

{
  "_class_name": "StableDiffusionPipeline",
  "_diffusers_version": "0.6.0",
  "feature_extractor": [
    "transformers",
    "CLIPImageProcessor"
  ],
  "safety_checker": [
    "stable_diffusion",
    "StableDiffusionSafetyChecker"
  ],
  "scheduler": [
    "diffusers",
    "PNDMScheduler"
  ],
  "text_encoder": [
    "transformers",
    "CLIPTextModel"
  ],
  "tokenizer": [
    "transformers",
    "CLIPTokenizer"
  ],
  "unet": [
    "diffusers",
    "UNet2DConditionModel"
  ],
  "vae": [
    "diffusers",
    "AutoencoderKL"
  ]
}