用纯 Mojo编写了这个迷你稳定扩散模型和助手（例如 DDPM 采样器、CLIP ）
开源地址 ：https://github.com/maniartech/mojo-stringbuilder
我使用 Stable Diffusion 和 Stable Diffusion XL 生成的一些图像...如果我们可以使用 Mojo 生成它们会怎样？
概述💡
这是迷你稳定扩散模型前向传播的 100% Mojo 实现。该模型的每个组件都是从头开始实现的，从基本整数和浮点数一直到矩阵乘法、卷积、图像数组以及 PyTorch 中存在的操作，例如线性层、上采样和广播。要在 Mojo 中查看这些操作并在您的项目中使用它们，请检查我们的 helpers/utils.mojo。这将节省您的时间，因为您不需要从头开始实施它们。然而，在这个项目中，这些基本构建块的主要用途是构建交叉注意模块、编码器、解码器、扩散模块、CLIP 和稳定扩散管道的其他组件。
该项目的目标是提供模型的基本实现，该模型虽然较小，但遵循与稳定扩散非常相似的架构。我希望任何想要修改给定代码以将自己的稳定扩散权重加载到该模型中的人都可以使用此实现。代码划分如下：
vae.mojo：变分自动编码器（VAE），包括编码器和解码器。编码器用于对预先存在的图像进行编码，以便稳定扩散模型可以使用它们。所有图像生成前向传递都会使用解码器，无论是否提供初始图像
Sampler.mojo：DDPM 采样器的实现
pipeline.mojo：创建一个 Mojo 管道，它接受文本提示（以及可选的图像）并计算通过模型的前向传递
fusion.mojo：稳定扩散模型的扩散部分的代码（由 UNet、时间嵌入和输出层组成）
Clip.mojo：由剪辑嵌入和 CLIP 播放器组成，它实现了生成推理管道使用的 CLIP 文本编码器所需的结构
helpers 文件夹：包含 utils.mojo 文件中代码中各处使用的重要低级函数（例如卷积、图像调整大小、矩阵乘法等）。还包含attention.mojo 文件中的注意力模块（自注意力、交叉注意力）。
demo.mojo：如何运行前向传递的简单示例
tokenizer_creation.py：一个 Python 文件，用于从 Hugging Face 检索 CLIP 分词器并将其存储为 tokenizer_clip.bin。该 .bin 文件将在图像生成期间被读取以加载 CLIP 标记器值。因此，应在运行任何前向传递之前执行此 Python 文件，因为这将允许您使用计算机中的实际值重新创建 CLIP 标记生成器。
用法🔨
首先，检索 CLIP 分词器文件
python tokenizer_creation.py
接下来，编译“helpers”包
mojo package helpers -o "helpers.mojopkg"
然后，修改vae.mojo和diffusion.mojo，使这些文件中定义的结构体与您正在加载其权重的模型具有相同的架构。由于此文件中有各种版本的Stable Diffusion（XL、mini等）我提供了 VAE 和 Diffusion 组件的一个非常小的基本实现。
接下来，将权重加载到模型中。为此，请查看 helpers/utils.mojo 中的“Tokenizer”结构。此结构中的init函数显示了如何将权重从 .bin 文件加载到结构中，因此相同的过程可以应用于您想要加载权重的任何其他结构（CLIP、编码器、扩散等）。只需复制并粘贴该代码并根据需要进行修改即可。该文件输入/输出逻辑是从令人惊叹的 Llama2.mojo 项目中检索的，链接在“感谢”部分。
最后，修改 demo.mojo 文件以设置您想要用于模型的参数并运行前向传递
mojo demo.mojo
此外，请随意更改 pipeline.mojo 文件中的“image_size”别名以适合您想要使用的图像宽度/长度！