KOALA-Lightning-1B 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

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.

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

• 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-1b", 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}
}