license: apache-2.0 datasets:

• poloclub/diffusiondb base_model:
• stabilityai/stable-diffusion-2-1-base pipeline_tag: text-to-image library_name: diffusers

AMD Nitro扩散模型

简介

AMD Nitro扩散模型是基于AMD Instinct™ GPU从主流扩散模型蒸馏而成的高效文生图模型系列，包含两大核心产品：

• Stable Diffusion 2.1 Nitro：基于UNet架构的单步生成模型，蒸馏自Stable Diffusion 2.1
• PixArt-Sigma Nitro：基于Transformer架构的高分辨率单步模型，蒸馏自PixArt-Sigma

⚡️ 开源代码！本系列模型基于我们对潜在对抗扩散蒸馏技术的重构实现（该技术曾用于开发Stable Diffusion 3 Turbo模型）。由于原作者未公开训练代码，我们发布重构版本以推动该领域研究。

技术细节

• 模型架构：Stable Diffusion 2.1 Nitro与原始SD2.1架构完全兼容，支持diffusers管线调用
• 推理步数：蒸馏后仅需单步推理，训练代码同时支持2/4/8步蒸馏
• 硬件配置：使用4块AMD Instinct™ MI250 GPU单节点完成蒸馏
• 训练数据：采用DiffusionDB的100万条提示词，通过基础模型生成对应图像
• 训练成本：在单节点上2天内即可获得理想蒸馏效果

快速使用

from diffusers import DDPMScheduler, DiffusionPipeline
import torch

scheduler = DDPMScheduler.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="scheduler")
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", scheduler=scheduler)

ckpt_path = '<蒸馏模型路径>'
unet_state_dict = torch.load(ckpt_path)
pipe.unet.load_state_dict(unet_state_dict)
pipe = pipe.to("cuda")

image = pipe(prompt='一只猫的照片',
             num_inference_steps=1,
             guidance_scale=0,
             timesteps=[999]).images[0]

完整训练与评估方法详见GitHub仓库。

性能表现

相较于Stable Diffusion 2.1基础模型，我们在保持CLIP分数仅降低2.5%、FID分数仅上升2.2%的同时，实现了95.9%的FLOPs运算量削减。

许可声明

版权所有 (c) 2018-2024 超威半导体公司 依据Apache License 2.0版本授权（"许可证"），除非符合许可证要求，否则不得使用本文件。您可通过以下链接获取许可证副本： http://www.apache.org/licenses/LICENSE-2.0 除非适用法律要求或书面同意，本软件按"原样"分发，不附带任何明示或暗示的担保或条件。详见许可证中关于权限限制的具体条款。