如何魔改 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.0
        callback_kwargs['prompt_embeds'] = prompt_embeds
    return 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 Image

pipeline = 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 DiffusionPipeline

pipeline = 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 DiffusionPipeline

pipeline = 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, CLIPModel

clip_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 DiffusionPipeline


class 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 loop
        image = torch.randn((batch_size, self.unet.in_channels, self.unet.sample_size, self.unet.sample_size))

        image = image.to(self.device)

        # set step values
        self.scheduler.set_timesteps(num_inference_steps)

        for t in self.progress_bar(self.scheduler.timesteps):
            # 1. predict noise model_output
            model_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-1
            image = self.scheduler.step(model_output, t, image, eta).prev_sample

        image = (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 torch

class 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 = 1
    model_output = self.unet(image, timestep).sample
    scheduler_output = self.scheduler.step(model_output, timestep, image).prev_sample
    return scheduler_output

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

推理

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

from diffusers import DDPMScheduler, UNet2DModel

scheduler = 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 的方式公开自己的源码,非常方便。

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

如若内容造成侵权/违法违规/事实不符,请联系我们进行投诉反馈qq邮箱809451989@qq.com,一经查实,立即删除!

相关文章

如何编写一份完整的软件测试报告

软件测试是软件开发过程中一个非常重要的环节,它有助于确保软件的质量和稳定性。编写一份完整的软件测试报告是软件测试工作的重要组成部分,它不仅可以帮助测试团队记录测试结果和发现的问题,还可以为开发团队提供有价值的反馈和改进建议。下…

我独自升级崛起下载慢/下载不了/无法下载的解决方法

我独自升级是一款RGP游戏,以充满独特的画风以及新的类型的挑战得到玩家们的喜爱。通过打怪、完成任务、参加活动等方式获得经验值,合理分配技能点数和属性点数,以适应不同阶段的挑战需求。掌握装备掉落、合成、升级机制,及时更换更…

ubuntu安装nginx以及开启文件服务器

1. 下载源码 下载页面:https://nginx.org/en/download.html 下载地址:https://nginx.org/download/nginx-1.24.0.tar.gz curl -O https://nginx.org/download/nginx-1.24.0.tar.gz2. 依赖配置 sudo apt install gcc make libpcre3-dev zlib1g-dev ope…

中药提取物备案数据库<5000+中药提取物>

NMPA中药提取物备案数据库的建立是确保中药提取物质量安全、规范生产行为、加强监管、保障公众用药安全、促进产业发展和国际化的重要措施。 通过查询中药提取物备案信息我们能了解到中药提取物的实用备案号、药品通用名称、药品生产企业、批准文号、备案日期、备案状态、中药…

Mysql启动报错:本地计算机上的mysql服务启动后停止,某些服务在未由其他服务或程序使用时将自动停止

Mysql启动报错:本地计算机上的mysql服务启动后停止,某些服务在未由其他服务或程序使用时将自动停止 文章目录 Mysql启动报错:本地计算机上的mysql服务启动后停止,某些服务在未由其他服务或程序使用时将自动停止1. 备份mysql的data文件夹2. 重新构建 Wind…

智能照明控制系统的节能优势及其应用

【摘要】:通过智能照明控制系统与传统照明控制的对比,阐明智能照明控制系统的优势,基本组成、在智能建筑中的应用效果以及其实施存在的问题和前景展望。 安科瑞薛瑶瑶18701709087 【关键字】:智能照明控制;系统组成&…

【御控物联】 JavaScript JSON结构转换(21):数组To对象——综合应用

文章目录 一、JSON结构转换是什么?二、术语解释三、案例之《JSON数组 To JSON对象》四、代码实现五、在线转换工具六、技术资料 一、JSON结构转换是什么? JSON结构转换指的是将一个JSON对象或JSON数组按照一定规则进行重组、筛选、映射或转换&#xff0…

医院智慧手术麻醉系统管理源码 C# .net有演示

医院智慧手术麻醉系统管理源码 C# .net有演示 手术麻醉管理系统(DORIS)是应用于医院手术室、麻醉科室的计算机软件系统。该系统针对整个围术期,对病人进行全程跟踪与信息管理,自动集成病人HIS、LIS、RIS、PACS信息,采集监护等设备数据&#x…

【THM】Protocols and Servers 2(协议和服务器 2

介绍 协议和服务器房间涵盖了许多协议: 远程登录HTTP协议文件传输协议邮件传输协议POP3IMAP实现这些协议的服务器会受到不同类型的攻击。仅举几例,请考虑: 嗅探攻击(网络数据包捕获)中间人 ( MITM ) 攻击密码攻击(身份验证攻击)漏洞从安全的角度来看,我们始终需要思考…

基于SSM的宠物管理系统

点击以下链接获取源码: https://download.csdn.net/download/qq_64505944/89076676?spm=1001.2014.3001.5503 技术:SSM(Spring+SpringMVC+MyBatis)+LayUI+Echarts技术栈,分页采用pagehelper插件,EasyExcel进行Excel文件的导入导出。 宠物管理系统 1 CHINER-宠物管理系…

算法刷题Day27 | 39. 组合总和、40.组合总和II、131.分割回文串

目录 0 引言1 组合总和1.1 我的解题 2 组合总和II2.1 解题 3 分割回文串3.1 切割3.2 总结:分割和组合的区别 🙋‍♂️ 作者:海码007📜 专栏:算法专栏💥 标题:算法刷题Day27 | 39. 组合总和、40.…

OpenHarmony实战开发-多设备自适应能力

介绍 本示例是《一次开发,多端部署》的配套示例代码,展示了页面开发的一多能力,包括自适应布局、响应式布局、典型布局场景以及资源文件使用。 说明: 自适应布局能力仅可以保证在外部容器大小在一定范围内变化时,容…

考研数学|张宇《1000题》做不下来,怎么办?

张宇的1000题对于初学者 来说确实有一些难度,我的建议是刷完知能行之后再做1000题。 首先,让我们来看看传统习题册存在的一些问题。虽然传统习题册通常会覆盖考试的各个知识点和题型,但其中一些问题在于它们可能过于注重题目的数量&#xff…

性能分析-nginx

tomcat 像kyj项目请求直接对接 tomcat,tomcat的连接池就会直接影响“并发用户数” 如果这种情况下做性能测试的时候,并发用户数不能满足要求,可以适当加大线程池的配置。 如:项目性能测试发现项目所在机器,资源利用率…

产品推荐 | 瑞苏盈科基于立体帧捕捉和视频处理应用的火星Mars EB1开发板

01 产品概述 火星Mars EB1底板是为火星Mars系列FPGA和SoC核心板设计的通用底板,非常适用于立体帧捕捉和视频处理应用,可以为构建基于FPGA的定制化硬件系统提供一个良好的基础和开端。 02 核心亮点 ■ 与所有火星Mars系列FPGA和SoC核心板兼容 ■ 适用…

Stable Diffusion文生图技术详解:从零基础到掌握CLIP模型、Unet训练和采样器迭代

文章目录 概要Stable Diffusion 底层结构与原理文本编码器(Text Encoder)图片生成器(Image Generator) 那扩散过程发生了什么?stable diffusion 总体架构主要模块分析Unet 网络采样器迭代CLIP 模型 小结 概要 Stable …

windows一键休眠,一键唤醒

1.使windows睡眠不可用,cmd以管理员身份运行: powercfg.exe /hibernate off 2.桌面创建快捷键 Rundll32.exe Powrprof.dll,SetSuspendState Sleep

王权与自由国际服测试资格申请 王权与自由steam国际服预约教程

《王权与自由》是一款由《剑灵》开发商NCsoft公开的旗下大型多人MMORPG新作游戏作品。攻城战,这是游戏中的一个重要玩法,主要分为攻城方和防守方。攻城方需要占领外城的遗址,并利用高轮(包括碎石机、跳跃者和战斗航母)…

蓝桥杯第六届c++大学B组详解

前言: 看了很多博客以及视频讲解,感觉都不是很清楚,比较模棱两可,所以干脆自己一边想,一边写博客,也可帮助到其他人,都是根据自己的逻辑来尽量清楚简单的讲清楚题目,喜欢的不要吝啬三…

【JVM】GC导致的性能问题排查与解决方案,日志、堆分析工具介绍

一、必要性 重要应用程序在使用过程中,忽然无法响应用户请求,排查发现网络联通无问题,gateway能够正常接收分发请求,应用进程正常,正常向注册中心发送请求,但是接收http请求全部返回报错。 添加gc后发现内…