Nitro T 1.2B

🚀 AMD Nitro-T

Nitro-T 是一系列专注于高效训练的文本到图像扩散模型。相较于以往注重高效训练的模型，该模型在图像生成基准测试中取得了具有竞争力的成绩，并且在 32 块 AMD Instinct™ MI300X GPU 上从头开始训练不到 1 天即可完成。

🚀 快速开始

你必须使用 diffusers>=0.34 才能从 Huggingface hub 加载模型。

基础用法

import torch
from diffusers import DiffusionPipeline
from transformers import AutoModelForCausalLM

torch.set_grad_enabled(False)

device = torch.device('cuda:0')
dtype = torch.bfloat16
resolution = 1024
MODEL_NAME = "amd/Nitro-T-1.2B"

text_encoder = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", torch_dtype=dtype)
pipe = DiffusionPipeline.from_pretrained(
    MODEL_NAME,
    text_encoder=text_encoder,
    torch_dtype=dtype, 
    trust_remote_code=True,
)
pipe.to(device)

image = pipe(
    prompt="The image is a close-up portrait of a scientist in a modern laboratory. He has short, neatly styled black hair and wears thin, stylish eyeglasses. The lighting is soft and warm, highlighting his facial features against a backdrop of lab equipment and glowing screens.",
    height=resolution, width=resolution,
    num_inference_steps=20,
    guidance_scale=4.0,
).images[0]

image.save("output.png")

更多关于训练和评估的详细信息，请访问 GitHub 仓库 并阅读我们的 技术博客文章。

✨ 主要特性

• 高效训练：在 32 块 AMD Instinct™ MI300X GPU 上从头开始训练不到 1 天即可完成。
• 竞争力表现：在图像生成基准测试中取得了具有竞争力的成绩。

📦 安装指南

文档未提及具体安装步骤，可参考 GitHub 仓库 中的训练和数据准备脚本。

💻 使用示例

基础用法

上述快速开始中的代码示例即为基础用法，可根据提示词生成相应的图像。

高级用法

文档未提供高级用法的代码示例，可参考 GitHub 仓库 中的更多代码。

📚 详细文档

模型架构

Nitro-T-1.2B 是一个文本到图像的多模态扩散变换器（MMDiT），其架构类似于 Stable Diffusion 3。它使用深度压缩自动编码器（DC - AE）的潜在空间，并使用 Llama 3.2 1B 模型进行文本条件处理。此外，还采用了多种技术来减少训练时间。更多详细信息请参阅我们的技术博客文章。

数据集

Nitro-T 模型在一个包含约 3500 万张图像的数据集上进行训练，该数据集由真实和合成数据源组成，这些数据在互联网上公开可用。数据处理脚本可在我们的 GitHub 仓库中找到。

训练成本

Nitro-T-1.2B 在 32 块 AMD Instinct™ MI300X GPU 上从头开始训练不到 1 天即可完成。

模型版本

• Nitro-T-0.6B：一个基于 512px DiT 的模型
• Nitro-T-1.2B：一个基于 1024px MMDiT 的模型

开源代码

⚡️ 开源代码！我们的 GitHub 提供了训练和数据准备脚本，以重现我们的结果。我们希望这个用于高效扩散模型训练的代码库能够让研究人员更快地验证想法，并降低独立开发者构建自定义模型的门槛。

技术博客

📝 阅读我们的 技术博客文章，了解我们实现快速训练的技术细节、结果和评估。

🔧 技术细节

文档未提供超过 50 字的具体技术说明，可参考 技术博客文章 了解更多技术细节。

📄 许可证

版权所有 (c) 2018 - 2025 Advanced Micro Devices, Inc. 保留所有权利。 本项目遵循 Apache 许可证 2.0 版（“许可证”）；除非符合许可证规定，否则不得使用此文件。 你可以在以下地址获取许可证副本： http://www.apache.org/licenses/LICENSE - 2.0 除非适用法律要求或书面同意，否则根据许可证分发的软件按“原样”分发， 不附带任何形式的明示或暗示的保证和条件。 请参阅许可证以了解管理权限和限制的具体语言。

🔍 模型信息

🌟 模型生成图像示例

Nitro - T - 1.2B 在 1024px 分辨率下生成的图像

Nitro - T - 0.6B 在 512px 分辨率下生成的图像