KOALA-Lightning-700M Model Card

Summary

Trained using a self-attention-based knowledge distillation method
Teacher model: SDXL-Lightning
Training dataset: a subset of LAION-POP dataset
Training iteration: 500K with a batch size of 128
Training GPUs: 4 x NVIDIA A100 (80GB)

KOALA-Lightning Models

Model	link
koala-lightning-1b	https://huggingface.co/etri-vilab/koala-lightning-1b
koala-lightning-700m	https://huggingface.co/etri-vilab/koala-lightning-700m

Abstract

TL;DR

We propose a fast text-to-image model, called KOALA, by compressing SDXL's U-Net and distilling knowledge from SDXL into our model. KOALA-Lightning-700M can generate a 1024x1024 image in 0.66 seconds on an NVIDIA 4090 GPU, which is more than 4x faster than SDXL. KOALA-700M can be used as a cost-effective alternative between SDM and SDXL in limited resources.

FULL abstract

As text-to-image (T2I) synthesis models increase in size, they demand higher inference costs due to the need for more expensive GPUs with larger memory, which makes it challenging to reproduce these models in addition to the restricted access to training datasets. Our study aims to reduce these inference costs and explores how far the generative capabilities of T2I models can be extended using only publicly available datasets and open-source models. To this end, by using the de facto standard text-to-image model, Stable Diffusion XL (SDXL), we present three key practices in building an efficient T2I model: (1) Knowledge distillation: we explore how to effectively distill the generation capability of SDXL into an efficient U-Net and find that self-attention is the most crucial part. (2) Data: despite fewer samples, high-resolution images with rich captions are more crucial than a larger number of low-resolution images with short captions. (3) Teacher: Step-distilled Teacher allows T2I models to reduce the noising steps. Based on these findings, we build two types of efficient text-to-image models, called KOALA-Turbo &-Lightning, with two compact U-Nets (1B & 700M), reducing the model size up to 54% and 69% of the SDXL U-Net. In particular, the KOALA-Lightning-700M is 4x faster than SDXL while still maintaining satisfactory generation quality. Moreover, unlike SDXL, our KOALA models can generate 1024px high-resolution images on consumer-grade GPUs with 8GB of VRAMs (3060Ti). We believe that our KOALA models will have a significant practical impact, serving as cost-effective alternatives to SDXL for academic researchers and general users in resource-constrained environments.

These 1024x1024 samples were generated by KOALA-Lightning-700M using 10 denoising steps in 0.66 seconds on an NVIDIA 4090 GPU.

Architecture

There are two two types of compressed U-Net, KOALA-1B and KOALA-700M, which are realized by reducing residual blocks and transformer blocks.

U-Net comparison

U-Net	SDM-v2.0	SDXL-Base-1.0	KOALA-1B	KOALA-700M
Param.	865M	2,567M	1,161M	782M
CKPT size	3.46GB	10.3GB	4.4GB	3.0GB
Tx blocks	[1, 1, 1, 1]	[0, 2, 10]	[0, 2, 6]	[0, 2, 5]
Mid block	✓	✓	✓	✗
Latency	1.131s	3.133s	1.604s	1.257s

Tx menans transformer block and CKPT means the trained checkpoint file.
We measured latency with FP16-precision, and 25 denoising steps in NVIDIA 4090 GPU (24GB).
SDM-v2.0 uses 768x768 resolution, while SDXL and KOALA models uses 1024x1024 resolution.

Latency and memory usage comparison on different GPUs

We measured the inference time of SDXL-Turbo and KOALA-Turbo models at a resolution of 512x512, and other models at 1024x1024, using a variety of consumer-grade GPUs: NVIDIA 3060Ti (8GB), 2080Ti (11GB), and 4090 (24GB). 'OOM' indicates Out-of-Memory. Note that SDXL models cannot operate on the 3060Ti with 8GB VRAM, whereas our KOALA models can run on all GPU types.

Key Features

Efficient U-Net Architecture: KOALA models use a simplified U-Net architecture that reduces the model size by up to 54% and 69% respectively compared to its predecessor, Stable Diffusion XL (SDXL).
Self-Attention-Based Knowledge Distillation: The core technique in KOALA focuses on the distillation of self-attention features, which proves crucial for maintaining image generation quality.

Model Description

Developed by ETRI Visual Intelligence Lab
Developer: Youngwan Lee, Kwanyong Park, Yoorhim Cho, Young-Ju Lee, Sung Ju Hwang
Model Description: Latent Diffusion based text-to-image generative model. KOALA models uses the same text encoders as SDXL-Base-1.0 and only replace the denoising U-Net with the compressed U-Nets.
Teacher model: SDXL-Lightning
Training dataset: a subset of LAION-POP dataset
Training iteration: 500K with a batch size of 128
GPUs: 4 x NVIDIA A100 (80GB)
Resources for more information: Check out KOALA report on arXiv and project page.

Usage with 🤗Diffusers library

The inference code with denoising step 25

import torch
from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler

pipe = StableDiffusionXLPipeline.from_pretrained("etri-vilab/koala-lightning-700m", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

# Ensure sampler uses "trailing" timesteps and "sample" prediction type.
pipe.scheduler = EulerDiscreteScheduler.from_config(
    pipe.scheduler.config, timestep_spacing="trailing"
  )


prompt = "A portrait painting of a Golden Retriever like Leonard da Vinci"
negative = "worst quality, low quality, illustration, low resolution"
image = pipe(prompt=prompt, negative_prompt=negative, guidance_scale=3.5, num_inference_steps=10).images[0]

Uses

Direct Use

The model is intended for research purposes only. Possible research areas and tasks include

Generation of artworks and use in design and other artistic processes.
Applications in educational or creative tools.
Research on generative models.
Safe deployment of models which have the potential to generate harmful content.
Probing and understanding the limitations and biases of generative models.
Excluded uses are described below.

Out-of-Scope Use

The model was not trained to be factual or true representations of people or events, and therefore using the model to generate such content is out-of-scope for the abilities of this model.

Limitations and Bias

Text Rendering: The models face challenges in rendering long, legible text within images.
Complex Prompts: KOALA sometimes struggles with complex prompts involving multiple attributes.
Dataset Dependencies: The current limitations are partially attributed to the characteristics of the training dataset (LAION-aesthetics-V2 6+).

Citation

@misc{Lee@koala,
    title={KOALA: Empirical Lessons Toward Memory-Efficient and Fast Diffusion Models for Text-to-Image Synthesis}, 
    author={Youngwan Lee and Kwanyong Park and Yoorhim Cho and Yong-Ju Lee and Sung Ju Hwang},
    year={2023},
    eprint={2312.04005},
    archivePrefix={arXiv},
    primaryClass={cs.CV}
}

etri-vilab
/

koala-lightning-700m