mindall-e
Kakao Brain 开源的 1.3B 参数文本生成图像模型,两阶段 AR+VQGAN 架构,CLIP 重排序提升生成质量
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
Kakao Brain 开源的 1.3B 参数文本生成图像模型,两阶段 AR+VQGAN 架构,CLIP 重排序提升生成质量
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
你是否曾幻想过:输入一段文字,AI 就能为你生成一幅画?2021年初,OpenAI 发布 DALL-E 震惊了世界——但它没有开源,普通人根本无法体验其背后的魔力。正是这个背景,催生了今天要介绍的项目:minDALL-E,由韩国 Kakao Brain 团队开源的 1.3B 参数文本生成图像模型。
想象一下,你是一个独立游戏开发者,需要为游戏中的 NPC 生成大量风格统一的角色立绘;或者你是一位科普博主,想要为文章配上独特的插图而不是用图库里的老照片。minDALL-E 就是为这类场景设计的:它让你在本地机器上就能运行一个经过 1400万图文对训练的文生图模型,不再依赖云端 API,也不用担心隐私泄露。
minDALL-E 采用了两阶段生成流水线,这是它与初代 DALL-E 最重要的区别之一。
Stage 1 —— VQGAN Tokenizer(图像压缩器)
第一阶段使用 VQGAN(Vector Quantized Generative Adversarial Network)作为图像 tokenizer。它的作用类似于翻译官:把一张 256×256 的 RGB 图片压缩成 16×16=256 个离散 token,每个 token 的取值范围是 0~16383(即 16384 个「词」)。这相当于把图像「文字化」,让后续的 Transformer 能够用处理文本的方式处理图像。
配置文件 dalle-1.3B.yaml 详细记录了 VQGAN 的超参数:嵌入维度 256、通道数 128、4级下采样(ch_mult: [1, 1, 2, 2, 4]),这些数值共同决定了压缩质量和重建保真度。minDALL-E 团队在 ImageNet 和 FFHQ 上对官方 VQGAN 进行了微调,确保 tokenizer 生成的 token 序列能够最大程度保留原始图像的语义信息。
Stage 2 —— 1.3B Transformer(核心生成器)
第二阶段是整个模型的核心:一个参数量达 13亿 的自回归 Transformer。配置显示:嵌入维度 1536、42 层、24 头、文本上下文长度 64 token、图像上下文长度 256 token。
这个 Transformer 同时学习两类 token 的预测:文本 token 和图像 token。训练时,给定文本序列和图像的 VQ token 序列,模型学习预测下一个图像 token——这与 GPT 系列语言模型的训练方式如出一辙。推理时,则从文本 token 出发,逐步自回归地生成 256 个图像 token,再通过 Stage 1 的 VQGAN decoder 将 token 序列解码回像素图像。
图1:minDALL-E 生成的「亚洲绘画风格小鸟」,可见模型对艺术风格的捕捉能力
生成图像的质量不仅取决于模型本身,采样策略同样关键。minDALL-E 提供了两种主流采样方法:
Top-K 采样:每一步从概率最高的 K 个 token 中抽样,K 越小生成的图像越「保守」,K 越大则越多样化。README 推荐 top_k ≤ 256。
Top-P(Nucleus)采样:从累积概率超过 P 的最小 token 集合中抽样,比 Top-K 更动态地适应不同位置的分布差异。
更值得注意的是推理代码中的 CLIP 重排序机制:每次生成 96 张候选图(num_candidates=96),然后用 OpenAI CLIP 计算每张图与文本 prompt 的语义相似度分数,最终按分数从高到低排序输出。这个设计巧妙地弥补了自回归模型在图像质量上的不确定性——模型自己判断哪些图最「像」文本描述。
图2:「金门大桥黑白风格」生成效果,展现了模型对场景结构和光影的把控
从 sampling.py 的实现来看,代码使用了 torch.cuda.amp 混合精度加速,并集成了 tqdm 进度条,让长时间推理过程可视化。最大序列长度固定为 256(即 16×16 的图像 token 网格)。
minDALL-E 的训练数据来自 Google 的 Conceptual Captions 数据集:CC3M(约 330 万图文对)和 CC12M(约 1220 万图文对),合计约 1550 万对,但 README 描述为 1400 万对(可能经过清洗后实际使用了 1400 万对)。这个规模远小于 DALL-E 的 2.5 亿图文对,但作为一个 1.3B 参数的小型模型,已经足够让它学习到丰富的图文对齐能力。
硬件需求是该项目最大的门槛:训练和推理均需要 NVIDIA V100(32GB)或更高规格的 GPU。README 明确指出「在 32GB 显存的 V100 上测试通过」,对于显存更小的 GPU,需要降低 num_candidates 以避免 OOM(内存溢出)。
图3:生成「戴墨镜的猫」的多样性,96个候选中CLIP重排序选出最优
整个代码库结构清晰,模块化程度高:
dalle/models/stage1/:VQGAN tokenizer 实现,包含 vqgan.py(完整 VQGAN 模型)和 layers.py(残差块、注意力层等基础组件)dalle/models/stage2/:Transformer 实现,transformer.py 为主模型(继承自 nn.Module),layers.py 定义了 Transformer Block(多头注意力 + FFN)dalle/utils/:工具函数,包括 sampling.py(采样逻辑)、utils.py(CLIP打分、随机种子设置)和 config.py(OmegaConf配置管理)dalle/models/tokenizer.py:文本 tokenizer,基于 HuggingFace tokenizers 库,词汇表 16384examples/:使用示例,包括 sampling_ex.py(完整采样脚本)和 sampling_interactive_demo.ipynb(Jupyter 交互式 demo)依赖管理使用 requirements.txt,核心依赖为:PyTorch 1.8.0、torchvision、pytorch-lightning、einops(爱因斯坦求和简记法)、omegaconf(YAML配置管理)、CLIP(来自 OpenAI)和 matplotlib。
minDALL-E 发布于 2021 年 12 月,在时间线上早于 Stable Diffusion(2022年8月)。作为 2021 年的研究成果,它的局限性是时代性的:
但换一个角度看,minDALL-E 的历史意义在于它是最早一批将亿级参数文生图模型开源的项目之一。在 Stable Diffusion 出现之前,它为学术研究者提供了宝贵的参考实现——两阶段 AR + VQGAN 的范式直接影响了后续许多文生图项目。
图4:细节丰富的场景生成——粉色大象在海滩漫步
使用 minDALL-E 非常简单:
from dalle.models import Dalle
model = Dalle.from_pretrained('minDALL-E/1.3B')
model.to(device='cuda')
images = model.sampling(prompt="A painting of a monkey with sunglasses",
top_k=256, num_candidates=96)
首次运行会自动下载约 5GB 的模型权重。生成后调用 CLIP 重排序,选出与文本最匹配的图像。
交互式 Jupyter Notebook demo(sampling_interactive_demo.ipynb)提供了滑块控制参数的可视化界面,适合快速探索模型能力。
图5:「沙漠中的埃菲尔铁塔」——检验模型对地标和背景组合的理解能力
minDALL-E 在 GitHub 上获得 631 颗星、65 个 Fork,对于一个 2021 年的研究项目而言,这个关注度相当不错。它证明了:
从技术演进角度看,minDALL-E 处于 GAN→AR→Diffusion 的过渡期,它的许多设计选择为后来的 Stable Diffusion 做了铺垫。它不是一个「能用就行」的玩具,而是一个真正有技术含量的研究级开源实现。
总结:minDALL-E 是一款面向开发者和研究者的开源文生图模型,以 1.3B 参数实现了从文本描述到图像的端到端生成。尽管在分辨率、Web 体验等方面已被后续的 Stable Diffusion 等项目超越,但它作为 2021 年最早开源的亿级参数量文生图模型之一,在技术探索和学术研究领域具有不可替代的参考价值。部署它需要 32GB 以上显存的高端 GPU,适合有深度学习背景的用户作为研究工具或二次开发基础。