哲学
🧨 Diffusers 提供了多个模态的最先进的预训练扩散模型。其目的是作为推理和训练的模块化工具箱。
我们致力于构建一个能够经受时间考验的库,因此非常重视 API 设计。
简而言之,Diffusers 被设计为 PyTorch 的自然扩展。因此,我们的大部分设计选择都基于 PyTorch 的设计原则。让我们来回顾一下最重要的几个原则:
可用性优先于性能
- 尽管 Diffusers 有许多内置的性能增强功能(参见 内存和速度),模型总是以最高精度和最低优化加载。因此,默认情况下,如果用户没有特别定义,扩散管道总是以 CPU 和 float32 精度实例化。这确保了在不同平台和加速器上的可用性,意味着运行库不需要复杂的安装。
- Diffusers 旨在成为一个轻量级的包,因此依赖项非常少,但有许多可选依赖项可以提高性能(如
accelerate
、safetensors
、onnx
等)。我们努力保持库的轻量级,以便可以轻松地将其作为其他包的依赖项添加。 - Diffusers 优先使用简单、自解释的代码,而不是简短、神奇的代码。这意味着通常不希望使用 lambda 函数和高级 PyTorch 操作等简短代码语法。
简单优于容易
正如 PyTorch 所说,显式优于隐式,简单优于复杂。这一设计哲学在库的多个部分都有体现:
- 我们遵循 PyTorch 的 API,如
DiffusionPipeline.to
,让用户处理设备管理。 - 优先使用简洁的错误消息,而不是默默地纠正错误输入。Diffusers 旨在教育用户,而不是使库尽可能容易使用。
- 复杂的模型和调度器逻辑是公开的,而不是在内部神奇地处理。调度器/采样器与扩散模型之间的依赖关系最小。这迫使用户编写展开的去噪循环。然而,这种分离使得调试更容易,并且用户对调整去噪过程或切换扩散模型或调度器有更多控制。
- 扩散管道中单独训练的组件,例如文本编码器、UNet 和变分自编码器,每个都有自己的模型类。这迫使用户处理不同模型组件之间的交互,序列化格式将模型组件分开存储在不同的文件中。然而,这使得调试和自定义更容易。DreamBooth 或 Textual Inversion 训练由于 Diffusers 能够分离扩散管道的单个组件而变得非常简单。
可调整、贡献者友好优于抽象
对于库的大部分内容,Diffusers 采用了 Transformers 库的一个重要设计原则,即优先使用复制粘贴的代码而不是仓促的抽象。这一设计原则非常有争议,与流行的 不要重复自己 (DRY) 设计原则形成鲜明对比。 简而言之,就像 Transformers 在建模文件中所做的那样,Diffusers 优先保持极低的抽象层次和非常自包含的代码,用于管道和调度器。 函数、长代码块,甚至类都可以在多个文件中复制,这在一开始看起来像是一个糟糕的、草率的设计选择,使库难以维护。 然而,这种设计已被证明对 Transformers 非常成功,并且对于社区驱动的开源机器学习库非常有意义,因为:
- 机器学习是一个快速发展的领域,范式、模型架构和算法迅速变化,因此很难定义长期有效的代码抽象。
- 机器学习从业者喜欢能够快速调整现有代码以进行创意和研究,因此更喜欢自包含的代码而不是包含许多抽象的代码。
- 开源库依赖于社区贡献,因此必须构建一个易于贡献的库。代码越抽象,依赖项越多,阅读和贡献的难度就越大。贡献者由于担心破坏关键功能而停止对非常抽象的库进行贡献。如果贡献库不会破坏其他核心代码,不仅对潜在的新贡献者更具吸引力,而且更容易并行审查和贡献多个部分。
在 Hugging Face,我们称这种设计为单文件策略,这意味着某个类的几乎所有代码都应写在一个单独的、自包含的文件中。要了解更多关于这一哲学的信息,可以阅读 这篇博客文章。
在 Diffusers 中,我们对管道和调度器遵循这一哲学,但对扩散模型只部分遵循。我们不完全遵循这一设计的原因是,几乎所有的扩散管道,如 DDPM、Stable Diffusion、unCLIP (DALL·E 2) 和 Imagen 都依赖于相同的扩散模型,即 UNet。
很好,现在你应该大致理解了 🧨 Diffusers 为什么这样设计 🤗。 我们努力在整个库中一致地应用这些设计原则。然而,也有一些小的例外或不幸的设计选择。如果你对设计有反馈,我们非常希望在 GitHub 上直接听到。
设计哲学的细节
现在,让我们深入了解一下设计哲学的细节。Diffusers 主要由三个主要类组成:管道、模型 和 调度器。 让我们详细了解一下每个类的设计决策。
管道
管道设计为易于使用(因此不完全遵循 简单优于容易),功能不完整,应被视为如何使用 模型 和 调度器 进行推理的示例。
遵循以下设计原则:
- 管道遵循单文件策略。所有管道都可以在
src/diffusers/pipelines
下的单独目录中找到。一个管道文件夹对应一个扩散论文/项目/发布。多个管道文件可以聚集在一个管道文件夹中,如src/diffusers/pipelines/stable-diffusion
所示。如果管道具有相似的功能,可以使用 # Copied from 机制。 - 所有管道都继承自 [
DiffusionPipeline
]。 - 每个管道由不同的模型和调度器组件组成,这些组件在
model_index.json
文件 中有文档说明,可以通过与管道同名的属性访问,并且可以通过DiffusionPipeline.components
函数在管道之间共享。 - 每个管道都应通过
DiffusionPipeline.from_pretrained
函数加载。 - 管道仅用于推理。
- 管道应非常易读、自解释且易于调整。
- 管道应设计为可以相互构建,并且易于集成到更高级的 API 中。
- 管道不旨在成为功能完整的用户界面。对于功能完整的用户界面,可以参考 InvokeAI、Diffuzers 和 lama-cleaner。
- 每个管道应通过
__call__
方法以一种且仅有一种方式运行。__call__
参数的命名应在所有管道中共享。 - 管道应以其解决的任务命名。
- 在几乎所有情况下,新的扩散管道应在新的管道文件夹/文件中实现。
模型
模型设计为可配置的工具箱,是 PyTorch 的 Module 类 的自然扩展。它们部分遵循单文件策略。
遵循以下设计原则:
- 模型对应于一种模型架构。例如,
UNet2DConditionModel
类用于所有期望 2D 图像输入并基于某些上下文条件的 UNet 变体。 - 所有模型都可以在
src/diffusers/models
中找到,每个模型架构应在单独的文件中定义,例如unets/unet_2d_condition.py
、transformers/transformer_2d.py
等。 - 模型不遵循单文件策略,应使用较小的模型构建块,如
attention.py
、resnet.py
、embeddings.py
等。注意:这与 Transformers 的建模文件形成鲜明对比,表明模型并不真正遵循单文件策略。 - 模型旨在暴露复杂性,就像 PyTorch 的
Module
类一样,并给出清晰的错误消息。 - 模型都继承自
ModelMixin
和ConfigMixin
。 - 模型可以在不进行重大代码更改、保持向后兼容并显著提高内存或计算性能的情况下进行优化。
- 模型默认应具有最高精度和最低性能设置。
- 要集成新的模型检查点,如果其通用架构可以归类为 Diffusers 中已存在的架构,应调整现有模型架构以使其与新的检查点兼容。只有当模型架构根本不同才应创建新文件。
- 模型应设计为易于扩展以适应未来的变化。这可以通过限制公共函数参数、配置参数并“预见”未来的变化来实现,例如,通常最好添加
string
类型的参数,这些参数可以轻松扩展到新的未来类型,而不是使用布尔类型的is_..._type
参数。对现有架构进行的更改应尽可能少,以使新的模型检查点工作。 - 模型设计是一个在保持代码可读性和简洁性与支持许多模型检查点之间的艰难权衡。对于大部分建模代码,类应适应新的模型检查点,但也有一些例外,为了确保代码长期保持简洁和可读,应添加新类,例如 UNet 块 和 注意力处理器。
调度器
调度器负责指导推理的去噪过程以及定义训练的噪声计划。它们设计为具有可加载配置文件的独立类,并严格遵循单文件策略。
遵循以下设计原则:
- 所有调度器都在
src/diffusers/schedulers
中找到。 - 调度器不允许从大型工具文件中导入,并应保持非常自包含。
- 一个调度器 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
函数不暴露所有复杂性,可能有点像“黑盒”。 - 在几乎所有情况下,新的调度器应在新的调度文件中实现。