bk-sdm-v2-tiny

BK-SDM-v2 Model Card

BK-SDM-{v2-Base, v2-Small, v2-Tiny} are obtained by compressing SD-v2.1-base.

• Block-removed Knowledge-distilled Stable Diffusion Models (BK-SDMs) are developed for efficient text-to-image (T2I) synthesis:
  • Certain residual & attention blocks are eliminated from the U-Net of SD.
  • Despite the use of very limited data, distillation retraining remains surprisingly effective.
• Resources for more information: Paper, GitHub.

Examples with 🤗Diffusers library.

An inference code with the default PNDM scheduler and 50 denoising steps is as follows.

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("nota-ai/bk-sdm-v2-tiny", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

prompt = "a black vase holding a bouquet of roses"
image = pipe(prompt).images[0]  
    
image.save("example.png")

Compression Method

Based on the U-Net architecture and distillation retraining of BK-SDM, a reduced batch size (from 256 to 128) is used in BK-SDM-v2 for faster training speeds.

• Training Data: 212,776 image-text pairs (i.e., 0.22M pairs) from LAION-Aesthetics V2 6.5+.
• Hardware: A single NVIDIA A100 80GB GPU
• Gradient Accumulations: 4
• Batch: 128 (=4×32)
• Optimizer: AdamW
• Learning Rate: a constant learning rate of 5e-5 for 50K-iteration retraining

Experimental Results

The following table shows the zero-shot results on 30K samples from the MS-COCO validation split. After generating 512×512 images with the PNDM scheduler and 25 denoising steps, we downsampled them to 256×256 for evaluating generation scores.

• Our models were drawn at the 50K-th training iteration.

Compression of SD-v2.1-base

Compression of SD-v1.4

Visual Analysis: Image Areas Affected By Each Word

KD enables our models to mimic the SDM, yielding similar per-word attribution maps. The model without KD behaves differently, causing dissimilar maps and inaccurate generation (e.g., two sheep and unusual bird shapes).

<center> <img alt="cross-attn-maps" img src="https://netspresso-research-code-release.s3.us-east-2.amazonaws.com/assets-bk-sdm/fig_cross-attn-maps_bk-sd-v2.png" width="100%"> </center>

Uses

Please follow the usage guidelines of Stable Diffusion v1.

Acknowledgments

• Microsoft for Startups Founders Hub and Gwangju AICA for generously providing GPU resources.
• CompVis, Runway, and Stability AI for the pioneering research on Stable Diffusion.
• LAION, Diffusers, PEFT, DreamBooth, Gradio, and Core ML Stable Diffusion for their valuable contributions.

Citation

@article{kim2023architectural,
  title={BK-SDM: A Lightweight, Fast, and Cheap Version of Stable Diffusion},
  author={Kim, Bo-Kyeong and Song, Hyoung-Kyu and Castells, Thibault and Choi, Shinkook},
  journal={arXiv preprint arXiv:2305.15798},
  year={2023},
  url={https://arxiv.org/abs/2305.15798}
}

@article{kim2023bksdm,
  title={BK-SDM: Architecturally Compressed Stable Diffusion for Efficient Text-to-Image Generation},
  author={Kim, Bo-Kyeong and Song, Hyoung-Kyu and Castells, Thibault and Choi, Shinkook},
  journal={ICML Workshop on Efficient Systems for Foundation Models (ES-FoMo)},
  year={2023},
  url={https://openreview.net/forum?id=bOVydU0XKC}
}

This model card is based on the Stable Diffusion v1 model card.