如何魔改 diffusers 中的 pipelines

如何魔改 diffusers 中的 pipelines

整个 Stable Diffusion 及其 pipeline 长得就很适合 hack 的样子。不管是通过简单地调整采样过程中的一些参数,还是直接魔改 pipeline 内部甚至 UNet 内部的 Attention,都可以实现很多有趣的功能或采样生图结果。

本文主要介绍两种魔改 diffusers pipelines 的方式,一是通过注册 callback 函数,来在采样生图过程中执行某些操作,二是直接自己写 custom pipelines。

pipeline callbacks

可参考官方文档:Pipelines Callback

通过在 pipe 推理生图时传入自定义的回调函数,不用动底层代码,可以在每个时间步结束时动态地执行一些我们想要的动作,比如在特定的时间步修改特定的采样参数或张量。目前仅支持 callback_on_step_end 在单步结束时执行回调函数。我们通过两个例子来介绍如何通过 callback 函数来魔改 diffusers pipelines。

Dynamic classifier-free guidance

classifier-free guidance (cfg)用于通过 prompt 来引导图像生成的内容。在 diffusers 中,会同时用 CLIP 文本编码器同时编码 prompt 的 embeds 和空 prompt (空字符串或 negative prompt)的 embeds,然后拼接起来,一起通过交叉注意力与 UNet 交互。

通过 callback 函数,我们可以按照自己的需求动态控制 cfg,比如说,我想在特定的时间步之后停止使用 cfg,从而节省计算开销,并且性能不会有很大的损失。 callback 函数需要接收以下参数:

  • pipeline:通过 pipe 可以访问和编辑许多重要的采样参数,如 pipe.num_timestepspipe.guidance_scale 等。在本例中,我们就可以通过将 pipe._guidance_scale 来停用 cfg。

  • timestep 和 step_index:这两个参数可以让我们知道本次采样过程一共有多少时间步,以及当前我们位于哪一步。从而可以根据当前位于整个采样过程中的位置,来选择进行什么操作。在本例中,我们可以设置在整个采样过程 40% 及以后的位置,停用 cfg。

  • callback_kwargs:callback_kwargs 是一个字典,包含了在采样生图的过程中你可以编辑的张量。具体包含哪些张量,需要再调用 pipe 采样生图时通过 callback_on_step_end_tensor_inputs 参数传入。不同的 pipe 可能包含不同的可编辑张量,具体可以通过 pipe 的 _callback_tensor_inputs 属性查看。在本例中,我们需要在停用 cfg 之后调整 prompt_embeds 张量的批尺寸,丢掉空 prompt 部分。这是因为 sd pipe 是根据 _guidance_scale 的值来判断是否进行 cfg,所以我们将这个值改为零,就不会进行 cfg 了,所以需要将 prompt_embeds 不带 prompt 的部分丢掉,只保存带 prompt 的部分。

返回值方面,回调函数必须返回(修改好的) callback_kwargs。

最终,我们的 callback 函数是这样的:

def callback_dynamic_cfg(pipe, step_index, timestep, callback_kwargs, percent):if step_index == int(pipe.num_timesteps * percent):prompt_embeds = callback_kwargs['prompt_embeds']prompt_embeds = prompt_embeds.chunk(2)[-1]pipe._guidance_scale = 0.0callback_kwargs['prompt_embeds'] = prompt_embedsreturn callback_kwargs

然后在推理生图时传入该参数:

pipeline = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipeline = pipeline.to("cuda")prompt = "a photo of an astronaut riding a horse on mars"generator = torch.Generator(device="cuda").manual_seed(1)
out = pipeline(prompt,generator=generator,callback_on_step_end=callback_dynamic_cfg,callback_on_step_end_tensor_inputs=['prompt_embeds']
)out.images[0].save("out_custom_cfg.png")
Display image after each generation step

在搭建生图 UI 时,一般需要支持用户在看到前几个时间步的结果不符合预期时,手动终止采样生图过程。这就需要两个功能,一是展示每一步的生图结果,二是支持在过程中终止本次生图。下面以展示中间步结果为例,介绍回调函数的使用。

我们知道,在 SD 中,去噪采样生图过程是发生在隐空间的,以 SDXL 为例,隐空间的特征图尺寸为 128 × 128 × 3 128\times 128\times 3 128×128×3 。我们需要将其转换到像素空间,才能看到这一步的实际生图结果。SDXL 还有点特殊,需要先将四通道的隐空间转换为 RGB 三通道,详情见 Explaining the SDXL latent space 。

def latents_to_rgb(latents):weights = ((60, -60, 25, -70),(60,  -5, 15, -50),(60,  10, -5, -35))weights_tensor = torch.t(torch.tensor(weights, dtype=latents.dtype).to(latents.device))biases_tensor = torch.tensor((150, 140, 130), dtype=latents.dtype).to(latents.device)rgb_tensor = torch.einsum("...lxy,lr -> ...rxy", latents, weights_tensor) + biases_tensor.unsqueeze(-1).unsqueeze(-1)image_array = rgb_tensor.clamp(0, 255)[0].byte().cpu().numpy()image_array = image_array.transpose(1, 2, 0)return Image.fromarray(image_array)def decode_tensors(pipe, step, timestep, callback_kwargs):latents = callback_kwargs["latents"]image = latents_to_rgb(latents)image.save(f"{step}.png")return callback_kwargs

然后再推理生图时传入该参数回调函数:

from diffusers import AutoPipelineForText2Image
import torch
from PIL import Imagepipeline = AutoPipelineForText2Image.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",torch_dtype=torch.float16,variant="fp16",use_safetensors=True
).to("cuda")image = pipeline(prompt = "A croissant shaped like a cute bear.",negative_prompt = "Deformed, ugly, bad anatomy",callback_on_step_end=decode_tensors,callback_on_step_end_tensor_inputs=["latents"],
).images[0]

在这里插入图片描述

Custom Pipelines

可参考官方文档:Custome Pipelines、contribute-pipeline

在 diffusers 中,我们可以很方便的自定义并加载自己的定制化 pipelines。在实现自己的自定义 pipelines 时,需要继承基类 DiffusionPipeline

加载 custom pipelines

1 从 diffusers 仓库中加载自定义的 pipeline

从 hf hub 加载自定义的 pipeline 非常简单,只需要将 model_id 传入 custom_pipeline 参数,然后就会加载该仓库中对应的 pipeline.py。

from diffusers import DiffusionPipelinepipeline = DiffusionPipeline.from_pretrained("google/ddpm-cifar10-32", custom_pipeline="hf-internal-testing/diffusers-dummy-pipeline"
)

2 从本地文件加载自定义的 pipeline

如果想从本地文件加载 pipeline,需要将 pipeline.py 所在的文件目录(注意是目录名)传给 custom_pipeline 参数。

from diffusers import DiffusionPipelinepipeline = DiffusionPipeline.from_pretrained("google/ddpm-cifar10-32", custom_pipeline="path/to/dir"
)

3 加载官方收录的社区 custom pipelines

这里是合入 diffuses 官方仓库的一些社区的自定义 pipelines。我们只需要将对应文件名(不含 py 后缀,如 clip_guided_stable_diffusion)传给 custom_pipeline 参数。

由于自定义 pipelines 的通常比较复杂,所以我们也可以通过官方 pipeline 来加载模型,再将模型传入自定义 pipelines。

from diffusers import DiffusionPipeline
from transformers import CLIPFeatureExtractor, CLIPModelclip_model_id = "laion/CLIP-ViT-B-32-laion2B-s34B-b79K"feature_extractor = CLIPFeatureExtractor.from_pretrained(clip_model_id)
clip_model = CLIPModel.from_pretrained(clip_model_id)pipeline = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4",custom_pipeline="clip_guided_stable_diffusion",clip_model=clip_model,feature_extractor=feature_extractor,
)
实现 custom pipelines

我们可以继承 DiffusionPipeline 基类并实现自己的 custom pipeline,这样所有人就都可以加载我们实现的 pipeline。一个 custom pipeline 的框架大致如下:

import torch
from diffusers import DiffusionPipelineclass MyPipeline(DiffusionPipeline):def __init__(self, unet, scheduler):super().__init__()self.register_modules(unet=unet, scheduler=scheduler)@torch.no_grad()def __call__(self, batch_size: int = 1, num_inference_steps: int = 50):# Sample gaussian noise to begin loopimage = torch.randn((batch_size, self.unet.in_channels, self.unet.sample_size, self.unet.sample_size))image = image.to(self.device)# set step valuesself.scheduler.set_timesteps(num_inference_steps)for t in self.progress_bar(self.scheduler.timesteps):# 1. predict noise model_outputmodel_output = self.unet(image, t).sample# 2. predict previous mean of image x_t-1 and add variance depending on eta# eta corresponds to η in paper and should be between [0, 1]# do x_t -> x_t-1image = self.scheduler.step(model_output, t, image, eta).prev_sampleimage = (image / 2 + 0.5).clamp(0, 1)image = image.cpu().permute(0, 2, 3, 1).numpy()return image

确保 xxxPipeline 这个类的相关实现都在这一个文件,而且该文件中只包含 xxxPipeline 一个类。因为 pipeline 的识别加载是自动的。

接下来,我们以一个最简单的 one-step pipeline 为例,简单介绍自己实现 custom pipeline 的过程。在这个one-step pipeline 中,只会用到 UNet 一个模型,并将 timestep 固定为 1,只进行一次模型前向。

首先新建一个 one_step_unet.py 文件,然后在其中继承 DiffusionPipeline 基类并实现 UnetSchedulerOneForwardPipeline 类。

初始化:定义 __init__ 方法

在初始化方法中,我们简单地的 one-step pipeline 只需要接收 unet 和 scheduler 两个初始化参数,我们在初始化方法中将变量定义好。注意,为了使得 save_pretrained 方法能够将我们的模型完整地保存下来,需要通过 register_modules 方法将我们想要保存的 unet 和 scheduler 注册进来。

from diffusers import DiffusionPipeline
import torchclass UnetSchedulerOneForwardPipeline(DiffusionPipeline):def __init__(self, unet, scheduler):super().__init__()self.register_modules(unet=unet, scheduler=scheduler)

前向推理:定义 __call__ 方法

定义好初始化方法 __init__ 之后,再来实现 pipeline 推理生图的 __call__ 方法。在这里,我们可以任意发挥,任意组合,魔改扩散模型的采样过程,实现自己想要的功能。在我们的 one-step pipeline 中,这里要做的非常简单:采样一个噪声图,UNet 进行一次前向。

@torch.no_grad()
def __call__(self):image = torch.randn(1, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.sanple_size)timestep = 1model_output = self.unet(image, timestep).samplescheduler_output = self.scheduler.step(model_output, timestep, image).prev_samplereturn scheduler_output

这就 ok 了,我们已经实现好了自定义的 one-step pipeline。

推理

我们传入 unet 和 scheduler,实例化一个刚刚自定义好的 UnetSchedulerOneForwardPipeline,然后进行推理生图:

from diffusers import DDPMScheduler, UNet2DModelscheduler = DDPMScheduler()
unet = UNet2DModel()pipeline = UnetSchedulerOneForwardPipeline(unet=unet, scheduler=scheduler)output = pipeline()

如果我们的 custom pipeline 结果如果跟某个已有的 pipeline 的预训练权重是完全一样的,我们还可以直接通过 from_pretrained 方法来加载它们的权重。比如说,我们的 UnetSchedulerOneForwardPipeline 就可以直接加载 google/ddpm-cifar10-32 的权重:

pipeline = UnetSchedulerOneForwardPipeline.from_pretrained("google/ddpm-cifar10-32", use_safetensors=True)output = pipeline()
分享 custom pipelines

要共享自己的 custom pipeline 有三个方法:

  1. 将自己实现的 custom pipeline 推送到 hf hub 仓库,需要将文件名命名为 pipeline.py

  2. 像 diffusers 官方仓库提交 PR,合并之后可以我们的 pipeline 就会出现在这里,别人可以通过文件名来加载

  3. 共享出自己的源码文件,如 clip_guided_stable_diffusion.py,别人也可以导入我们的 pipeline

而作为使用者,使用 custom pipeline 的方式有两种:

# 方式 1
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained("google/ddpm-cifar10-32", custom_pipeline="hf-internal-testing/diffusers-dummy-pipeline"
)# 方式 2
from cus_pipe import CustomPipeline    # cus_pipe is copied from hf-internal-testing/diffusers-dummy-pipeline
pipe = CustomPipeline.from_pretrained("google/ddpm-cifar10-32")

这两种使用方式应该是等价的,这可以从源码中看到:

if custom_pipeline is not None:pipeline_class = get_class_from_dynamic_module(custom_pipeline, module_file=CUSTOM_PIPELINE_FILE_NAME, cache_dir=custom_pipeline)
elif cls != DiffusionPipeline:pipeline_class = cls
else:diffusers_module = importlib.import_module(cls.__module__.split(".")[0])pipeline_class = getattr(diffusers_module, config_dict["_class_name"])

总结

diffusers 的 api 设计非常友好,我们可以通过 pipeline callback 和 custom pipeline 等方式定制化实现自己想要的功能,其中前者不用动底层代码,简单优雅,后者则是功能强大,现在最新的 AIGC 相关的论文基本都是通过 custom diffusion 的方式公开自己的源码,非常方便。

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.mzph.cn/news/799897.shtml

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

深入了解图片Base64编码

title: 深入了解图片Base64编码 date: 2024/4/8 10:03:22 updated: 2024/4/8 10:03:22 tags: Base64编码图片转换HTTP请求前端开发移动应用性能优化图片压缩 1. 什么是Base64编码 Base64编码是一种将二进制数据转换为文本字符串的编码方式,通过将数据转换为一种可…

安全威胁情报的漏洞挖掘

前段时间edu上出现了两个网安总队收取安全情报,不收漏洞,下面简单分析一下如何挖掘安全情报。 在发现在edu中新增了两个网安总队收安全情报等漏洞,那威胁情报又会包含哪些内容呢?以前或许会看到各种ss网站、bc网站、yx网站满天飞&…

【数据库总结】

文章目录 1.数据库介绍2.数据库的语法使用数据库语法数据库的执行过程 3.数据库的索引介绍索引的介绍索引创建注意点:索引失效的情况索引不适合哪些场景呢?索引是不是建的越多越好呢?索引的数据结构为什么要用 B 树,而不用普通二叉…

华为 2024 届校园招聘-硬件通⽤/单板开发——第十套

华为 2024 届校园招聘-硬件通⽤/单板开发——第十套 部分题目分享,完整版带答案(共十套)获取(WX:didadidadidida313,加我备注:CSDN huawei硬件单板题目,谢绝白嫖哈) 1、I2 C 总线…

bat批处理命令 获取当前盘符和当前目录和上级目录

echo off echo 当前盘符:%~d0 echo 当前盘符和路径:%~dp0 echo 当前批处理全路径:%~f0 echo 当前盘符和路径的短文件名格式:%~sdp0 echo 当前CMD默认目录:%cd% pause 通过批处理取当前目录的上一级目录1 echo off if …

【御控物联】JavaScript JSON结构转换(22):小结

文章目录 一、结语二、接下来……三、在线转换工具四、技术资料 一、结语 《JavaScript JSON结构转换》主题自24.03.25至24.XX.XX历时XX天,共计编写 XX篇(XX篇功能说明XX篇场景)。 在此我们对《JavaScript JSON结构转换》功能做一下总结! 在JSON结构转…

蓝桥杯 第 9 场 小白入门赛 盖印章

题目: 2.盖印章【算法赛】 - 蓝桥云课 (lanqiao.cn) 思路: 此题主要靠解方程组,但是枚举好像不太行,因为会有负数解, 二元一次方式 设A章需要a次,B章需要b次 第一个方程:a b k; …

Android Apk签名算法使用SHA256

Android apk签名算法使用SHA256 本文不介绍复杂的签名过程,说一下Android签名算法使用SHA256。 但是SHA1不是相对安全签名算法,SHA256更加安全一些。 一般大公司才会有这种细致的安全要求。 如何查看apk签名是否是SHA1还是SHA256 1、拿到apk文件&…

Notion 开源替代品 AFFINE 部署和使用教程

AFFiNE 是一款完全开源的 Notion Miro 替代品,与 Notion 相比,AFFiNE 更注重隐私安全,优先将笔记内容保存到本地。 GitHub 地址:https://github.com/toeverything/AFFiNE AFFiNE 使用 Rust 和 Typescript 构建,只需…

Vue - 你知道Vue中computed和watch的区别吗

难度级别:中高级及以上 提问概率:70% 二者都是用来监听数据变化的,而且在日常工作中大部分时候都只是局限于简单实用,所以到了面试中很难全面说出二者的区别。接下来我们看一下,二者究竟有哪些区别呢? 先说computed,它的主要用途是监听…

缺省参数讲解

概念: 声明或定义函数时为参数指定一个缺省值,在调用该函数时,如果没有指定实参则采用该形参的缺省值,否则使用指定的实参 就是先给形参一个默认值,若后续传参时你没有给它传参的话,它就使用这个默认值 …

MySQL通用语法解析

SQL通用语法 SQL语句可以单行或多行书写,以分号结尾 SQL语句可以使用空格/缩进来增强语句的可读性 MySQL数据库的SQL语句不区分大小写,关键字建议使用大写 注释 单行注释:-- 注释内容 或 # 注释内容 多行注释:/* 注释内容 */ …

weight-tying探索

在一些领域,将嵌入层和输出层的权重绑定,以达到减少参数量并使得相同token保持统一的embedding空间的作用。 下面的nn.Linear(3, 10)的权重矩阵的尺寸是10*3,即y W x b,因此跟nn.Embedding(10, 3)的权重矩阵大小相等。 impor…

应用运维文档1

统一nginx接入配置指南 Nginx配置规范 1:不带微服务编码上下文至后端,以metadata-ui为例 location段配置信息,location配置中维护微服务编码上下文信息 # app_code: metadata-ui 流水线名称: metadata-ui location ~ ^/metadata-ui/(?P.*) {set $app_code metadata-ui;p…

语音特征的反应——语谱图

语谱图的横坐标为时间,纵坐标为对应时间点的频率。坐标中的每个点用不同颜色表示,颜色越亮表示频率越大,颜色越淡表示频率越小。可以说语谱图是一个在二维平面展示三维信息的图,既能够表示频率信息,又能够表示时间信息。 创建和绘制语谱图的…

卫星遥感监测森林植被健康度

随着地球环境的日益恶化,森林作为地球上最重要的生态系统之一,其变化对全球气候、生态环境和人类社会经济发展产生深远影响。因此,及时、准确地监测森林变化对于保护生态环境、维护生态平衡、推进可持续发展具有重要意义。卫星遥感影像技术因…

若依框架学习——分页查询列表

条件查询【多条件】列表展示【分页】SaCheckPermissionTableName TableId NotBlank Page分页 响应数据封装类

C语言main( ) 函数有什么作⽤?

一、问题 main( ) 函数是C语⾔程序中最重要的组成部分,不可或缺,那么它有什么作⽤? 二、解答 C程序是由⼀个或多个函数组成的,其中必须有⼀个且只有⼀个名为 main( ) 的函数, 该函数是整个程序的⼊口。 既然是程序的⼊…

C#速览入门

C# & .NET C# 程序在 .NET 上运行,而 .NET 是名为公共语言运行时 (CLR) 的虚执行系统和一组类库。 CLR 是 Microsoft 对公共语言基础结构 (CLI) 国际标准的实现。 CLI 是创建执行和开发环境的基础,语言和库可以在其中无缝地协同工作。 用 C# 编写的…

使用Python实现决策树算法

决策树是一种常用的机器学习算法,它可以用于分类和回归任务。在本文中,我们将使用Python来实现一个基本的决策树分类器,并介绍其原理和实现过程。 什么是决策树算法? 决策树是一种基于树形结构的机器学习算法,它通过…