Nitro-T is a family of text-to-image diffusion models focused on highly efficient training. Our models achieve competitive scores on image generation benchmarks compared to previous models focused on efficient training while requiring less than 1 day of training from scratch on 32 AMD Instinct™ MI300X GPUs. The release consists of:
⚡️
Open-source code! Our GitHub provides training and data preparation scripts to reproduce our results. We hope this codebase for efficient diffusion model training enables researchers to iterate faster on ideas and lowers the barrier for independent developers to build custom models.
📝 Read our
technical blog post for more details on the techniques we used to achieve fast training and for results and evaluations.
1import torch
2from diffusers import DiffusionPipeline
3from transformers import AutoModelForCausalLM
4
5torch.set_grad_enabled(False)
6
7device = torch.device('cuda:0')
8dtype = torch.bfloat16
9resolution = 1024
10MODEL_NAME = "amd/Nitro-T-1.2B"
11
12text_encoder = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", torch_dtype=dtype)
13pipe = DiffusionPipeline.from_pretrained(
14 MODEL_NAME,
15 text_encoder=text_encoder,
16 torch_dtype=dtype,
17 trust_remote_code=True,
18)
19pipe.to(device)
20
21image = pipe(
22 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.",
23 height=resolution, width=resolution,
24 num_inference_steps=20,
25 guidance_scale=4.0,
26).images[0]
27
28image.save("output.png")
For more details on training and evaluation please visit the
GitHub repo and read our
technical blog post.
Copyright (c) 2018-2025 Advanced Micro Devices, Inc. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.