AuroraAI-Research/AuroraIMG-6M · Hugging Face

来源：https://huggingface.co/AuroraAI-Research/AuroraIMG-6M

AuroraIMG-6M 是一个文本到图像模型，仅有682万参数，这个数字包含了全部组件：文本编码器、去噪器和 VAE。它可以根据简短的英文描述生成 64×64 的图像，且所有组件均在 CC12M 数据集上从零开始训练，未使用任何预训练的文本编码器或 VAE。

在 1 万张 CC12M 留出集图像上，它达到了 FID 33 和 CLIP 24.8 的指标。

概览

分辨率 64×64 RGB 总参数量 6.82M（文本编码器 1.32M / 去噪器 4.42M / VAE 1.08M）方法 潜空间流匹配（rectified flow），使用无分类器引导 潜空间维度 8 × 16 × 16 文本编码器 4 层 transformer，4096 词表 BPE，最长 64 个 token 采样器 Euler，50 步，CFG 6.0（CFG 扫描中最优）数据 CC12M，约 988 万训练对

快速开始

from diffusers import DiffusionPipeline
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32

pipe = DiffusionPipeline.from_pretrained(
    "AuroraAI-Research/AuroraIMG-6M",
    custom_pipeline="AuroraAI-Research/AuroraIMG-6M",
    trust_remote_code=True,
    torch_dtype=dtype
).to(device)

prompt = "a sunset over the ocean"
print(f"Generating: '{prompt}'...")

generator = torch.Generator(device=device).manual_seed(42)
image = pipe(
    prompt=prompt,
    num_inference_steps=50,
    guidance_scale=6.0,
    generator=generator
).images[0]

filename = "Aurora-Image.png"
image.save(filename)
print(f"\n Saved to {filename}!")

该模型足够小，可以在 CPU 上运行。

示例

使用相同的八个提示词和相同的随机种子，逐步增加引导强度（CFG 1、2、3、5，从上到下）：

CFG 扫描 (https://huggingface.co/AuroraAI-Research/AuroraIMG-6M/blob/main/guidance_sweep.png)

最终检查点在 CFG 3 下对固定提示词集的生成结果（“a red car on a street”、“a dog running on green grass”、“a sunset over the ocean”、“a snowy mountain with blue sky”、“a plate of food on a table”、“a man standing in front of a building”、“a small boat on a calm lake”、“a vase of flowers on a wooden table”）：

最终样本 (https://huggingface.co/AuroraAI-Research/AuroraIMG-6M/blob/main/samples_0300000.png)

COCO val2014 真实图像（上排）与其描述生成的图像（下排）对比：

COCO 真实 vs 生成 (https://huggingface.co/AuroraAI-Research/AuroraIMG-6M/blob/main/coco_real_vs_gen.png)

单独展示 VAE 的效果：上为原始图像，下为重建图像。

VAE 重建 (https://huggingface.co/AuroraAI-Research/AuroraIMG-6M/blob/main/vae_recon.png)

工作原理

**VAE（108万参数）。**一个小型卷积 VAE，使用残差块和 GroupNorm，将 64×64×3 的图像压缩为 8 通道 16×16 的潜空间表示。训练损失包括 L1 损失、基于梯度的清晰度损失、Gram 矩阵损失和一个很小的 KL 项（5e-5），不使用对抗损失或 LPIPS 损失。在留出集图像上，重建 PSNR 达到 31.9 dB。

**文本编码器（132万参数）。**一个 4 层 pre-LN transformer（宽度 128，4 个注意力头），使用旋转位置编码（RoPE），底层是基于 CC12M 描述训练的字节级 BPE 分词器。它与去噪器联合训练，因此能学习到图像模型认为最有用的表示。

**去噪器（442万参数）。**一个紧凑的 U-Net，在 16×16、8×8 和 4×4 三个尺度上处理潜空间。残差块使用 AdaLN-Zero 条件化（scale、shift 和 gate 由时间嵌入生成，零初始化使得每个块初始时等价于恒等映射）。时间嵌入使用随机傅里叶特征，并叠加文本嵌入的均值池化结果。每个层级包含一个注意力模块，包括自注意力、对文本 token 的交叉注意力和一个前馈层。

训练

**数据：**整个 CC12M 数据集（pixparse/cc12m-wds）以流式方式读取并过滤。描述文本转为小写，清除 HTML 和 URL，过滤掉过于笼统的描述（如“image“、“stock photo“等）或少于 3 个词的描述。短边小于 128 像素或非 RGB 的图像也被过滤。剩余图像按短边 64 像素缩放（双三次插值），中心裁剪，并按像素哈希去重。最终得到 9,882,283 个训练对和 5 万张留出图像。唯一的数据增强是随机水平翻转。

**阶段一：VAE。**训练 25k 步，批量大小 2048，使用 AdamW（学习率 5e-4，余弦衰减）。最终损失 0.0880（L1 0.0343，KL 5.91）。

**阶段二：流匹配。**训练 300k 步，批量大小 2048，约 6.14 亿个样本，约合 62 个 epoch。模型以 MSE 损失预测速度（ε − z），时间步采样自 logit-normal 分布并混合 15% 的均匀分布。训练时，缓存的潜空间表示会使用 VAE 自身的后验方差重新加噪。权重矩阵和卷积核使用 Muon 优化器（学习率 0.02），其余参数（嵌入层、输入/输出层、归一化层、偏置、时间投影和池化投影）使用 AdamW（学习率 5e-4）。此外：10k 步预热，余弦衰减至 1%，梯度裁剪 1.0，EMA 0.9998，10% 描述词丢弃率用于 CFG，bf16 autocast 以及 torch.compile。最终训练损失：0.6208。以下所有样本和指标均使用 EMA 权重。

结果

AuroraIMG-6M 的所有指标均在原生 64×64 分辨率下测量。FID 使用 torchmetrics（2048 维 Inception 特征），CLIP score 使用 openai/clip-vit-base-patch32。两者都会在内部上采样图像，因此应将其视为追踪本模型趋势的参考值，而非与 256px 以上模型直接对比。

CC12M 留出集（1 万张图像）

我将 CFG 从 1.0 扫描到 9.0，6.0 的效果最佳——在不牺牲 FID 的情况下几乎达到了该模型所能达到的最高 CLIP，因此推荐使用该设置。

FID ↓ CLIP score ↑ AuroraIMG-6M（CFG 6.0） 33.01 24.80 真实图像 n/a 28.25

与对应描述对比时，生成图像的 CLIP 分数约为真实照片的 88%（24.80 vs 28.25）。

仅 VAE

FID ↓ PSNR ↑ CLIP score ↑ 8.97 31.92 dB 26.95

GenEval

总分（各任务平均）：0.1387。12.88% 的图像和 20.61% 的提示词完全正确（2212 张图像，553 个提示词）。

任务 准确率 single_object 50.00%（160 / 320） colors 26.86%（101 / 376） two_object 3.79%（15 / 396） counting 1.56%（5 / 320） position 0.75%（3 / 400） color_attr 0.25%（1 / 400）

与 Stable Diffusion 1.5 对比

模型 参数量 总分 ↑ 单物体 双物体 计数 颜色 位置 颜色属性 AuroraIMG-6M 6.82M（总计） 0.14 0.50 0.04 0.02 0.27 0.01 0.00 Stable Diffusion v1.5 0.86B（仅 U-Net） 0.43 0.97 0.38 0.35 0.76 0.04 0.06

GenEval 2

AuroraIMG-6M 使用 CFG 6.0、50 步 Euler 采样、每个提示词生成 4 张图像。

模型 参数量 Soft-TIFA GM ↑ Soft-TIFA AM ↑ AuroraIMG-6M 6.82M 2.21 21.84 Stable Diffusion 2.1 0.87B（仅 U-Net） 5.24 0.8 Supra2-IMG 104.1M（仅去噪器） 6.80 36.38 Surjo-Image-Preview 108.3M（仅去噪器） 8.07 42.30

各技能维度拆解（Soft-TIFA AM，原子级别）：

技能 AuroraIMG-6M Supra2-IMG Surjo-Image-Preview 物体 29.64 46.77 61.03 属性 26.01 43.56 45.80 位置 15.33 31.07 28.91 计数 11.90 22.14 22.30 动作 0.29 1.51 2.03

**提示词技巧。**该模型在 CC12M 的替代文本上训练，因此简短、小写、描述风格的提示词效果最好（如 “a dog running on green grass”、“a snowy mountain with blue sky”）。

适用场景

• 实验研究小型生成模型：Muon、流匹配、AdaLN-Zero、小型 VAE、文本编码器联合训练。
• 在资源极其有限的硬件上运行文本到图像生成。

注意事项

• **分辨率。**输出为 64×64，因此人脸、手部和文字等小细节无法清晰呈现。
• **提示词复杂度。**它最适合单一主体和简单场景。计数、位置和多属性提示词是其短板（参见 GenEval 和 GenEval 2）。
• **语言。**仅支持英文，小写，最长 62 个描述 token。
• **数据偏差。**CC12M 是网络抓取的替代文本，因此模型会继承其偏差和图库照片风格。仅应用了上述质量过滤，未对数据或输出做额外的安全过滤。
• 不适用于生成真实人物或事件的写实图像，或任何对准确性和公平性有要求的场景。

许可证

模型权重以 Apache 2.0 许可证发布。模型在 CC12M 上训练，该数据集以图像 URL 和描述的形式分发；底层图像有各自的许可证，请根据您的使用场景查阅数据集条款。

引用

@misc{auroraimg6m,
  title        = {AuroraIMG-6M: a 6.8M-parameter text-to-image model},
  author       = {AuroraAI-Research},
  year         = {2026},
  howpublished = {\url{https://huggingface.co/AuroraAI-Research/AuroraIMG-6M}}
}