模块化扩散

模块化扩散提供了一个易于使用的模块化API，用于设计和训练基于PyTorch的自定义扩散模型。无论您是探索扩散模型的爱好者还是资深的机器学习研究人员，这个框架都适合您。

特点

⚙️ 高度模块化设计：轻松替换扩散过程的不同组件，包括噪声类型、调度类型、去噪网络和损失函数。
📚 不断增长的预构建模块库：通过我们全面的预构建模块选择立即开始使用。
🔨 轻松创建自定义模块：通过继承基类并实现所需方法来创建您自己的原创模块。
🤝 与PyTorch集成：基于PyTorch构建，模块化扩散使您能够使用熟悉的语法开发自定义模块。
🌈 广泛的应用范围：从生成高质量图像到实现非自回归文本合成管道，可能性无穷无尽。

安装

模块化扩散官方支持Python 3.10+，可在PyPI上获得：

pip install modular-diffusion

您还需要为您的系统安装正确的PyTorch发行版。

注意：尽管模块化扩散适用于更高版本的Python，但我们目前建议使用Python 3.10。这是因为torch.compile（显著提高模型速度）目前不适用于Python 3.10以上的版本。

使用

使用模块化扩散，您只需几行代码就可以构建和训练自定义扩散模型。首先，加载并规范化您的数据集。我们正在使用AFHQ中的狗狗图片。

x, _ = zip(*ImageFolder("afhq", ToTensor()))
x = resize(x, [h, w], antialias=False)
x = torch.stack(x) * 2 - 1

接下来，使用模块化扩散的预构建模块或您的自定义模块构建您的自定义模型。

model = diffusion.Model(
   data=Identity(x, batch=128, shuffle=True),
   schedule=Cosine(steps=1000),
   noise=Gaussian(parameter="epsilon", variance="fixed"),
   net=UNet(channels=(1, 64, 128, 256)),
   loss=Simple(parameter="epsilon"),
)

现在，训练并从模型中采样。

losses = [*model.train(epochs=400)]
z = model.sample(batch=10)
z = z[torch.linspace(0, z.shape[0] - 1, 10).long()]
z = rearrange(z, "t b c h w -> c (b h) (t w)")
save_image((z + 1) / 2, "output.png")

最后，惊叹结果。

模块化扩散预览