purohit10saurabh/minFLUX

来源：https://github.com/purohit10saurabh/minFLUX

minFLUX

（非官方）FLUX扩散Transformer的极简PyTorch实现

许可证 欢迎提交PR

一个简化的教育性PyTorch实现，对应Black Forest Labs（https://bfl.ai）的FLUX.1（https://bfl.ai/models/flux-kontext）和FLUX.2（https://bfl.ai/models/flux-2）扩散Transformer（DiT）。旨在通过可逐行验证的官方代码库映射，帮助理解修正流匹配、联合注意力以及FLUX背后的关键设计选择。

扩散模型架构和训练算法根据官方diffusers仓库（https://github.com/huggingface/diffusers/tree/cbf4d9a3c384ef97d6b0e40c9846dd9e0e41886a）推断得出。VAE架构来自官方BFL仓库（flux（https://github.com/black-forest-labs/flux/tree/802fb4713906133fcbd0d8dc5351620ca4773036）和flux2（https://github.com/black-forest-labs/flux2/tree/50fe5162777813d869182b139e83b10743caef15））。每个.py文件都配有对应的.md文件，其中包含每个函数到固定提交版本中确切源代码行的详细映射。

内部内容

• FLUX.1 和 FLUX.2 的 DiT 架构——带有联合注意力的双流和单流Transformer块
• 修正流匹配——使用速度预测和对数正态时间步采样的训练
• 欧拉ODE推理——可配置时间步调度的采样循环
• VAE编码器/解码器——基于残差网络和注意力的架构，含潜在空间归一化
• 可逐行验证的源代码映射——指向官方代码库

扩散方程

训练（修正流匹配）：

\begin{aligned} x_t &= (1 - \sigma(t)) \cdot x_0 + \sigma(t) \cdot \epsilon && \text{(加噪输入)} \\ v &= \epsilon - x_0 && \text{(速度目标)} \\ L &= \left\| model(x_t, t) - v \right\|^2 && \text{(MSE损失)} \end{aligned}

推理（欧拉ODE步）：

x_{t_{\text{next}}} = x_t + (\sigma(t_{\text{next}}) - \sigma(t)) \cdot model(x_t, t)

FLUX.2 架构概览

FLUX.2 架构概览

详细架构：FLUX.2 模型架构

FLUX.1 与 FLUX.2 的主要差异

仓库结构

flux1/                     FLUX.1
  model.py                   DiT（双流 + 单流）
  training.py                流匹配 + 打包/解包
  kontext_training.py        参考图像条件化
  inference.py               欧拉ODE采样
  vae.py                     VAE（缩放/平移）

flux2/                     FLUX.2
  model.py                   DiT（共享调制，SwiGLU）
  training.py                流匹配 + 4D位置ID
  inference.py               欧拉ODE（经验性μ偏移）
  vae.py                     VAE（分块 + 批量归一化）

utils/                     共享
  model.py                   嵌入、RoPE、注意力、归一化
  training.py                噪声、损失、欧拉步、训练循环
  vae_utils.py               残差网络、注意力、上/下采样块

tests/
  test_utils.py              单元测试

flux1/、flux2/ 和 utils/ 中的每个 .py 文件都配有对应的 .md 文件，其中包含到真实源码仓库的逐行映射。

设置

pip install -r requirements.txt
python -m pytest tests/ -v

参与贡献

非常欢迎贡献，尤其是针对：

• 源代码真实性：对照diffusers（https://github.com/huggingface/diffusers）、flux（https://github.com/black-forest-labs/flux）和flux2（https://github.com/black-forest-labs/flux2）交叉验证代码，修正实现上的任何差异。
• 文档：改进附带的 .md 文件，并在 diffusers 变更时更新行映射。
• 组件：添加缺失的FLUX组件或改进现有组件。

请随时打开 issue（https://github.com/purohit10saurabh/minFLUX/issues）或创建拉取请求（https://github.com/purohit10saurabh/minFLUX/pulls）。

免责声明

由于 minFLUX 是根据官方 diffusers 和 BFL 仓库推断得出，代码中可能存在的错误来源包括：

• AI辅助：本仓库为“氛围编程”（vibe-coded），借助AI编写，参考了 diffusers 和 BFL 仓库。部分训练细节依据其他工作（如dreambooth）推断。
• 简化：去除了ControlNet、IP-Adapter、梯度检查点、KV缓存、FSDP/DeepSpeed支持以及注意力处理器分发模式，因此无法与预训练权重兼容。
• 上游代码变更：源代码真值行号引用了特定提交（diffusers（https://github.com/huggingface/diffusers/tree/cbf4d9a3c384ef97d6b0e40c9846dd9e0e41886a）、flux（https://github.com/black-forest-labs/flux/tree/802fb4713906133fcbd0d8dc5351620ca4773036）、flux2（https://github.com/black-forest-labs/flux2/tree/50fe5162777813d869182b139e83b10743caef15））。这些代码库会频繁更新，因此函数可能移动、重命名或改变签名。

引用

如果您使用了本仓库，请按下述方式引用：

@misc{minflux2026,
  author = {Purohit, Saurabh},
  title  = {minFLUX: Minimal Pytorch Implementation of FLUX Diffusion Transformers},
  year   = {2026},
  publisher = {GitHub},
  url    = {https://github.com/purohit10saurabh/minFLUX}
}