Swift Diffusion

这是一个单文件重新实现的Stable Diffusion模型。它包括CLIP文本分词器、CLIP文本编码器模型、UNet扩散模型和解码器模型。它还包括PLMS推理实现。该实现试图逐层匹配Stable Diffusion的输出，因此，给定相同的起点x_T，本实现和Stable Diffusion将输出相同的图像。

理由

这个重新实现是为了让我自己更好地理解扩散模型。同时，这也是我后续工作的必要步骤，以便在iPad/iPhone等移动设备上启用Stable Diffusion。没有Swift重新实现，用Python进行移动端优化将会很困难，并且无法在App Store上发布。虽然可以采取其他方法，比如导出到ONNX运行时并在移动设备上使用，但这会限制你可以应用的优化类型。显然，在移动设备上运行总内存约8GiB、存储约4GiB的全浮点精度模型并非易事。这可能需要一些非常规的优化，而现有框架可能无法提供。使用我熟悉的东西（我自己构建的框架）将是一个很好的起点。

当前进展

CLIP文本分词器、图像模型、文本模型、UNet扩散模型和自动编码器已经完成移植。examples:txt2img、examples:img2img和examples:inpainting目标是可用的。其他目标，如examples:unet、examples:clip、examples:decoder、examples:encoder和examples:vit是用于将PyTorch权重转换为s4nnc使用的权重的示例程序。

下一步计划

下一步是评估Int8卷积核和Int8 + Float16 GEMM核的可行性。这些应该有助于将运行时内存使用减少到大约2GiB左右。

目前，在运行时，UNet模型除了其3.3GiB的权重外，还使用了约1.5GiB的额外内存。这1.5GiB中的很大一部分是由于注意力层中的点积运算。我已经优化掉了约1GiB，因为之前softmax无法正确地进行原地操作（由于与别名和重塑相关的复杂原因）。我相信这在PyTorch代码中仍然存在，因为没有原地softmax方法。可以将该点积进一步拆分成更小的批次以节省峰值内存使用（沿k的token维度）。如果这些操作正确完成，我们应该能够将UNet的内存使用降低到全浮点精度下约3.8GiB。我还有另一个想法来进一步减少内存使用，那就是压缩UNet中的快捷激活（这些快捷激活将在下采样路径中保存，并在上采样路径中使用，因此占用时间较长）。但我不太确定这能节省多少内存。

将模型转换为FP16将立即节省内存占用。在txt2img、img2img和inpainting中实现了一个开关，可以在Float16和Float32之间切换。仅使用FP16，UNet应该使用约1.9GiB内存。如果我们能卸载UNet模型并在完成时从磁盘加载解码器，这些组合有望最终在移动设备上运行stable diffusion。我们可以使用LLM.int8()transformers技巧进一步将权重量化为int8：https://arxiv.org/pdf/2208.07339.pdf。有关这些技巧的更多信息，请参见：https://github.com/TimDettmers/bitsandbytes。

是否具有可比性

通过在s4nnc中实现permute运算符，我减少了转置流量。当我们比较Swift版本的txt2img与CompVis版本（有更多优化的分支可用，但逐一检查以找到最佳版本需要时间）的扩散过程+解码过程的执行时间时。在GPU上，Swift的txt2img大约需要15秒，而CompVis大约需要11秒（均使用一张2080 Ti）。我使用的LayerNorm / GroupNorm核心存在一些效率问题，同时也存在一些关于为什么选择某些低性能GEMM核心的疑问。不过，我打算转向MPS实现。在CUDA端进一步优化不会转化为MPS端的收益。

MPS后端在M1 Mac Mini上以FP16完成大约需要95秒，峰值内存约4GiB。与PyTorch的约120秒相比，这似乎相当合理。

如何运行

目前需要进行相当多的设置。如果我在s4nnc端支持SwiftPM，会有所帮助。但这不是高优先级。

首先，将此存储库克隆到本地存储后，你需要安装Bazel和s4nnc的各种依赖项。要安装Bazel，请参照：https://bazel.build/install。

你还需要下载模型。我将Stable Diffusion v1.4模型放在https://static.libnnc.org/sd-v1.4.ckpt。请注意，这是一个与s4nnc兼容的文件，而不是你在其他地方下载的PyTorch文件。你可以查看相关示例，了解如何生成此文件。

Linux

其他依赖项包括Swift编译器、CUDA（10.2及以上版本）和clang。对于前两者，你必须自行安装。对于其他依赖项，如果你使用的是类Debian系统，可以通过以下命令安装：

sudo apt install clang llvm libicu-dev libpng-dev libjpeg-dev libatlas-base-dev libblas-dev libgsl-dev libdispatch-dev libomp-dev libfftw3-dev

最后，在Linux上正确设置Bazel：

./bazel/setup_clang.sh /usr/local

完成这些后，你可以运行：

bazel run examples:txt2img --compilation_mode=opt -- /home/包含sd-v1.4.ckpt文件的绝对工作目录 "一位宇航员骑马的照片"

macOS

安装Bazel后，你应该修改此存储库下的WORKSPACE文件。特别是，关于ccv_setting的这一行应修改如下：

ccv_setting(
    name = "local_config_ccv",
    have_accelerate_framework = True,
    have_pthread = True,
)

你还需要在此存储库下添加一个新的.bazelrc.local文件，内容为一行build --config=mps。

之后，你应该能够运行：

bazel run examples:txt2img --compilation_mode=opt -- /Users/包含sd-v1.4.ckpt文件的绝对工作目录 "一位宇航员骑马的照片"

在M1 Mac Mini上大约需要95秒。请耐心等待。

生成的图像将保存在你提供的目录下，文件名为txt2img.png。

对于img2img，它会在你提供的工作目录下查找init_img.png。对于inpainting，它会在工作目录下查找init_inpainting.png。目前，它期望修复部分为绿色。你还需要为inpainting提供提示指导，否则它将无法工作。

现有问题

我似乎无法在没有提示指导的情况下实现修复。我尝试使用CLIP图像模型进行指导。然而，由于Stable Diffusion模型使用文本投影之前的文本嵌入，使得图像模型的嵌入与文本嵌入不可比较。