StableCascade
基于42倍压缩潜空间的图像生成模型,推理速度比Stable Diffusion快16倍
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
基于42倍压缩潜空间的图像生成模型,推理速度比Stable Diffusion快16倍
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
2023 年底,Stability AI 发布了 Stable Cascade,这是基于 Würstchen 架构的官方代码库。和大家熟知的 Stable Diffusion 不同,Stable Cascade 做了一个极其大胆的设计决策:将图像压缩到前所未有的 24×24 像素潜在空间——比 SD 的 128×128 小了整整 42 倍。
这个数字意味着什么?打个比方:如果你把 AI 图像生成比作一次长途运输,Stable Diffusion 是一次运一车货,而 Stable Cascade 是一次运一个集装箱。运输成本和时间的差距,不是一点半点。根据论文数据,Würstchen v2 相比 SD 1.5 实现了 16 倍的计算成本降低。
Stable Diffusion 的 VAE 将 1024×1024 图像压缩到 128×128(压缩率 8),而 Stable Cascade 通过 Stage A 和 Stage B 两阶段级联压缩,把同一张图像压缩到 24×24(压缩率 42)。Stage C(文本条件扩散模型)只在 24×24 的极小潜在空间内运行,大幅降低计算量。生成完成后,Stage B 和 Stage A 负责将 24×24 重建回 1024×1024 高质量图像。
更令人惊讶的是,即使压缩率如此之高,重建质量依然保持极高水准。代码中的 Stage A 使用 Vector Quantized VAE(VQ-VAE),包含 8192 个码本的向量量化;Stage B 则是一个 4 通道输入/输出的扩散模型,包含多个 Attention Block 和 ResBlock。两者的配合实现了"小体积输入 + 高质量输出"的理想效果。
Stable Cascade 的名字"Cascade"(级联)来源于它的三阶段设计:
Stage A(20M 参数):轻量级 VQ-VAE 编解码器,负责最基础的图像量化。使用 VectorQuantize 将图像特征映射到 8192 个离散码本,scale_factor 为 0.43。它是固定组件,由于体积小(仅 20M 参数)不需要微调。
Stage B(700M / 1.5B 参数):扩散解码器,接收 Stage C 生成的 24×24 潜空间图像和 EfficientNet 编码的图像条件,恢复为 96×96 预览图像。1.5B 版本在重建细节上明显优于 700M 版本。配置中支持 self_attn、dropout、c_hidden 等灵活参数。
Stage C(1B / 3.6B 参数):文本条件扩散模型,是整个系统的核心。在 24×24 潜在空间内根据文本提示生成图像。3.6B 版本经过了大量微调,效果远优于 1B 版本。配置使用 CLIP text encoder(text 和 pooled 两个分支)进行文本编码,支持多种条件注入方式(CTA=Condition+Text+Attention)。
Stage B 和 Stage C 都支持通过配置文件灵活切换 Attention 类型(Self-Attn / Cross-Attn / None),这为后续的 ControlNet 扩展提供了基础。
代码库中包含一个自研的 GDF(Generic Diffusion Framework),这是一个灵活到令人印象深刻的扩散模型框架。它通过组合不同组件(Schedule、Scaler、Target、NoiseCondition、LossWeight)来支持多种扩散范式:DDPM、DDIM、EDM(Karras 论文噪声调度)、Rectified Flows。GDF 的设计哲学是"一切皆可组合":只需几行代码就可以在不同的训练 Schedule、采样 Schedule、损失权重之间切换。这种模块化设计让 Stable Cascade 的代码极其优雅,也为后续研究提供了友好的扩展基础。
Stable Cascade 对 ControlNet 和 LoRA 的支持相当完善:
ControlNet 支持:提供了 4 种 ControlNet 类型——Canny 边缘检测、Inpainting/Outpainting(使用显著性检测模型)、Face Identity(人脸识别,使用 ArcFace)、Super Resolution(超分辨率)。ControlNet 的实现使用了 EfficientNet 作为 backbone(支持 effnet 和 simple 两种 bottleneck 模式),从主模型的各层引出控制信号。
LoRA 支持:通过 modules/lora.py 中的 apply_lora() 函数实现,使用 PyTorch 的 parametrization 机制,无需修改原模型结构即可注入 LoRA 层。同时还支持 ReToken 技术,可以在模型中新增可学习的 token(类似 Textual Inversion)。
和其他图像生成模型不同,Stable Cascade 不只使用文本条件,还大量使用 EfficientNet-V2-S 作为图像编码器。这个编码器将输入图像转换为 16 通道的潜在表示,用于引导生成过程。这使得图生图、图像变体、ControlNet 引导等功能成为可能。
训练代码展现了对大规模训练的充分考虑:
Bucketeer 数据分桶优化效率use_fsdp=False 选项支持 PyTorch FSDP,但默认关闭multi_aspect_ratio 参数支持不同宽高比的图像训练虽然没有 Docker 支持,但 Stable Cascade 提供了开箱即用的 Gradio Web UI,且已集成到 Hugging Face diffusers 库中——任何熟悉 diffusers API 的开发者都可以用 3 行代码完成部署。Gradio App 支持所有核心功能:文生图、图像变体、图生图、LoRA 推理。
Stage C 3.6B + Stage B 1.5B 的组合大约需要 15GB+ 的模型权重,加上推理激活值,推荐 20GB+ VRAM 的 GPU(RTX 3090 / A5000 及以上)。部署关键注意事项:
bash models/download_models.sh 脚本下载,支持分模块下载Stable Cascade 代表了 AI 图像生成领域的效率转向。当行业普遍在追求更大模型时,Stability AI 选择了一条相反的路——通过架构创新让推理和训练都变得更便宜。42 倍压缩率的潜在空间意味着更快的迭代、更低的门槛,让更多研究者可以在消费级硬件上探索高质量图像生成的可能性。三阶段级联的设计也提供了一种可插拔的架构:你可以只替换 Stage C 来实验新的文本条件方法,而不需要改变整个系统。