Diffusers documentation
理念
理念
🧨 Diffusers 提供了跨多种模态的最先进的预训练扩散模型。 其目标是作为用于推理和训练的模块化工具箱。
我们致力于构建一个经得起时间考验的库,因此对 API 设计极为重视。
简而言之,Diffusers 旨在成为 PyTorch 的自然扩展。因此,我们的大多数设计选择都基于 PyTorch 的设计原则。下面介绍其中最重要的几点:
可用性优先于性能
- 尽管 Diffusers 内置了许多性能增强功能(参见 内存与速度),但模型始终以最高精度和最低优化级别加载。因此,除非用户另行指定,扩散流水线默认始终在 CPU 上以 float32 精度实例化。这确保了在不同平台和加速器上的可用性,也意味着运行该库无需进行复杂的安装配置。
- Diffusers 旨在成为一个轻量级包,因此必需的依赖项很少,但有许多可选依赖项可以提升性能(例如
accelerate、safetensors、onnx等)。我们努力将该库保持尽可能轻量,以便其他包可以毫无顾虑地将其作为依赖项添加。 - Diffusers 更倾向于使用简洁、自解释的代码,而非高度压缩、难以理解的代码。这意味着,诸如 lambda 函数等简写语法以及高级 PyTorch 运算符通常不被推荐。
简单
正如 PyTorch 所言,显式优于隐式,简单优于复杂。这一设计理念体现在库的多个部分中:
- 我们遵循 PyTorch 的 API,提供诸如
DiffusionPipeline.to之类的方法,以便用户自行处理设备管理。 - 抛出简洁的错误信息优于静默地纠正错误输入。Diffusers 旨在引导用户,而非一味追求库的易用性。
- 复杂的模型与调度器逻辑被显式暴露,而非在内部以“魔法”方式处理。调度器/采样器与扩散模型相互独立,彼此依赖极少。这要求用户自行编写展开的去噪循环。然而,这种分离使得调试更加容易,并赋予用户对去噪过程的调整或替换扩散模型、调度器更大的控制权。
- 扩散流水线中独立训练的各个组件,例如文本编码器、unet 和变分自编码器,各自拥有独立的模型类。这要求用户自行处理不同模型组件之间的交互,且序列化格式会将各模型组件拆分到不同的文件中。不过,这种方式更便于调试和自定义。得益于 Diffusers 能够分离扩散流水线中的单个组件,DreamBooth 或 Textual Inversion 的训练变得非常简单。
可调整、对贡献者友好的抽象层
在库的许多部分,Diffusers 采用了 Transformers 库 的一项重要设计原则,即优先使用复制粘贴的代码,而非仓促的抽象。这一设计原则立场鲜明,与诸如 不要重复自己(DRY) 等主流设计原则形成鲜明对比。 简而言之,就像 Transformers 对建模文件所做的那样,Diffusers 倾向于为 pipeline 和 scheduler 保持极低的抽象层级以及高度自包含的代码。 函数、较长的代码块甚至类都可能在多个文件之间被复制,乍一看这似乎是一种糟糕、草率的设计选择,会让库变得难以维护。然而,这种设计已被证明对 Transformers 极为成功,并且对于社区驱动、开源的机器学习库来说非常有意义,因为:
- 机器学习是一个发展极为迅速的领域,其中的范式、模型架构和算法都在快速变化,因此很难定义能够长期适用的代码抽象。
- 机器学习从业者通常希望快速调整现有代码以进行构思和研究,因此他们更倾向于使用自包含的代码,而非包含大量抽象的代码。
- 开源库依赖社区贡献,因此必须构建一个易于贡献的库。代码越抽象,依赖越多,可读性就越差,贡献的难度也就越大。由于担心破坏关键功能,贡献者往往会停止向高度抽象的库提交代码。如果向某个库提交代码不会破坏其他基础代码,那么不仅更容易吸引潜在的新贡献者,也更容易进行代码审查,并支持对多个部分进行并行贡献。
在 Hugging Face,我们将这种设计称为单文件策略,这意味着某个类的几乎所有代码都应编写在一个自包含的单一文件中。如需了解更多关于这一理念的内容,可以参阅这篇博客文章。
在 Diffusers 中,我们对流水线、调度器和模型均遵循这一设计原则。一些较旧的模型早于该约定,因此保持原样;所有新的模型架构都位于各自独立的自包含文件中。有关详细信息,请参阅下方的 Models 部分。
很好,现在你应该已经大致理解了 🧨 Diffusers 为何采用这样的设计 🤗。 我们力求在整个库中一致地应用这些设计原则。当然,也存在一些偏离该理念的小例外,或是一些不太理想的设计选择。如果你对设计有任何反馈,我们非常 ❤️ 希望能在 GitHub 上直接听到你的意见。
设计哲学详解
接下来是设计哲学的具体细节。Diffusers 提供了两种将 模型 和 调度器 组合成可运行工作流的方式:标准的 pipeline,它是单体式的,每个 pipeline 类对应一个任务;以及 Modular Diffusers,它是可组合的、基于模块的。以下章节将依次介绍 pipeline、模块化 pipeline、模型和调度器。
Pipelines
标准或模块化的 Pipelines 仅用于推理。它们的设计目标是易于使用(因此并不 100% 遵循 Simple over easy):可读性强、自解释、易于调整,并且最好被视为如何使用 models 和 schedulers 的示例。它们的功能并不完整。要在 Diffusers 之上构建功能完整的用户界面,请使用 Modular Diffusers。
遵循以下设计原则:
- Pipelines follow the single-file policy. All pipelines can be found in individual directories under src/diffusers/pipelines. One pipeline folder corresponds to one diffusion paper/project/release. Multiple pipeline files can be gathered in one pipeline folder, as it’s done for
src/diffusers/pipelines/stable-diffusion. If pipelines share similar functionality, one can make use of the # Copied from mechanism. - 所有 Pipelines 都继承自 DiffusionPipeline。
- 每个流水线都由不同的模型和调度器组件组成,这些组件在
model_index.json文件 中有文档说明,可以作为流水线的属性以相同的名称访问,并可以通过DiffusionPipeline.components函数在流水线之间共享。 - 每个 pipeline 都应可通过
DiffusionPipeline.from_pretrained函数加载。 - 每个 pipeline 都应通过
__call__方法提供且仅提供一种运行方式。__call__参数的命名应在所有 pipeline 中保持一致。 - Pipelines 的命名应反映其旨在解决的任务。
模块化 Diffusers
Modular Diffusers 是标准流水线的可组合替代方案。您可以使用可复用的流水线模块来构建工作流,这些模块可以自由混合、匹配、替换和共享。标准流水线只是关于如何使用模型和调度器的松散参考示例。Modular Diffusers 是构建功能完整的 Diffusers 用户界面的推荐路径,也是社区以去中心化方式构建和共享新流水线的途径。
遵循以下设计原则:
- 模块化流水线遵循单文件策略。每个流水线都位于 src/diffusers/modular_pipelines/ 下的独立文件夹中,该文件夹将工作流拆分为每个阶段一个文件:
encoders.py、before_denoise.py、denoise.py和decoders.py。另外两个文件完善了流水线。modular_blocks_<model>.py负责组装各个阶段,而modular_pipeline.py则定义了针对每个模型的 ModularPipeline 子类。模块化流水线之间不会相互导入。 - 每个模块化流水线被定义为一组 ModularPipelineBlocks。叶子块位于阶段文件中(
encoders.py、before_denoise.py、denoise.py、decoders.py),而modular_blocks_<model>.py则通过 SequentialPipelineBlocks 和 AutoPipelineBlocks 等容器类将它们组装成完整的工作流。这将DiffusionPipeline combines所合并的两个概念分离开来。块是一个纯定义:它声明输入、输出和组件依赖,但不持有权重,也无法运行。通过.init_pipeline(repo_id)创建的 ModularPipeline 是可运行的对应实体。保持块无状态且不含权重,正是使其能够自由组合、替换并在不同工作流间共享的关键。 - 要支持新任务,请编写任务特定的代码块,将其与现有代码块组合,并在顶层代码块组装中将该工作流注册到
_workflow_map。单个ModularPipeline可以支持许多工作流,例如文本生成图像、图像生成图像和图像修复,而DiffusionPipeline仅运行一个工作流。
请参阅 Modular Diffusers 文档 以获取完整的设计和使用指南。
模型
模型被设计为可配置的工具箱,是 PyTorch 的 Module 类 的自然扩展。它们应遵循 单文件策略。一些较旧的模型早于该约定,因此保持原样。请将它们视为遗留例外,而非新模型应遵循的模式。例如,原始的 UNet2DConditionModel 类曾用于多种 UNet 变体。
遵循以下设计原则:
- 每种模型架构类型都位于 src/diffusers/models 下的独立文件夹中,例如 transformers/、autoencoders/ 或 unets/。每个模型家族在该文件夹中都有自己独立的文件,例如 transformer_flux.py 和 transformer_wan.py。
- 模型遵循单文件策略。每个模型文件都应是自包含的,但少数所有模型都相同使用的标准模块除外,例如时间步嵌入和归一化层。请从 embeddings.py 和 normalization.py 中导入这些模块。
- 模型旨在暴露复杂性,就像 PyTorch 的
Module类一样,并提供清晰的错误信息。 - 所有模型都继承自
ModelMixin和ConfigMixin。 - 当模型优化无需对代码进行重大改动、能保持向后兼容性,并且能带来显著的内存或计算收益时,就可以对其进行性能优化。
- 模型默认应采用最高精度和最低性能设置。
- 若要集成一个与现有模型架构相似的新架构,请复制现有文件作为起点并进行相应调整。对于保持不变的层,使用
# Copied from注释,以便make fix-copies能够保持它们同步。
调度器
Scheduler 负责在推理阶段引导去噪过程,并在训练阶段定义噪声调度。它们被设计为独立的类,具有可加载的配置文件,并严格遵循 单文件策略。
遵循以下设计原则:
- 所有调度器都位于
src/diffusers/schedulers中。 - 调度器不允许从大型 utils 文件中导入,并且应保持高度自包含。
- 一个调度器 Python 文件对应一个调度器算法(例如论文中定义的算法)。
- 如果调度器具有相似的功能,我们可以利用
# Copied from机制。 - 所有调度器都继承自
SchedulerMixin和ConfigMixin。 - 如此处所述,调度器可以通过
ConfigMixin.from_config方法轻松替换。 - 每个调度器都必须具备
set_num_inference_steps和step函数。set_num_inference_steps(...)必须在每次去噪过程之前调用,即 在调用step(...)之前。 - 每个调度器都通过
timesteps属性暴露需要“循环遍历”的时间步,该属性是一个时间步数组,模型将针对这些时间步被调用。 step(...)函数接收模型的预测输出和“当前”样本 (x_t),并返回“前一步”的、去噪程度略高的样本 (x_t-1)。- 鉴于扩散调度器的复杂性,
step函数并未暴露所有细节,因此可能显得有些“黑盒”。 - 在几乎所有情况下,新的调度器都应实现在一个新的调度文件中。