Nitro-1 is a series of efficient text-to-image generation models that are distilled from popular diffusion models on AMD Instinct™ GPUs. The release consists of:
⚡️
Open-source code! The models are based on our re-implementation of
Latent Adversarial Diffusion Distillation, the method used to build the popular Stable Diffusion 3 Turbo model. Since the original authors didn't provide training code, we release our re-implementation to help advance further research in the field.
1from diffusers import DDPMScheduler, DiffusionPipeline
2import torch
3
4scheduler = DDPMScheduler.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="scheduler")
5pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", scheduler=scheduler)
6
7ckpt_path = '<path to distilled checkpoint>'
8unet_state_dict = torch.load(ckpt_path)
9pipe.unet.load_state_dict(unet_state_dict)
10pipe = pipe.to("cuda")
11
12image = pipe(prompt='a photo of a cat',
13 num_inference_steps=1,
14 guidance_scale=0,
15 timesteps=[999]).images[0]
For more details on training and evaluation please visit the
GitHub repo.
Compared to the
Stable Diffusion 2.1 base model, we achieve 95.9% reduction in FLOPs at the cost of just 2.5% lower CLIP score and 2.2% higher FID.
Copyright (c) 2018-2024 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.
