Koala Lightning 700m

🚀 KOALA-Lightning-700Mモデルカード

KOALA-Lightning-700Mは、自己注意ベースの知識蒸留法を用いて訓練されたテキストから画像を生成するモデルです。SDXL-Lightningを教師モデルとし、LAION-POPデータセットの一部を使って訓練されています。このモデルは、SDXLに比べてモデルサイズを大幅に削減し、推論速度を向上させています。

       

📚 詳細ドキュメント

概要

• 訓練手法：自己注意ベースの知識蒸留法を用いて訓練されています。
• 教師モデル：SDXL-Lightning
• 訓練データセット：LAION-POPデータセットの一部
• 訓練反復回数：バッチサイズ128で500K回
• 訓練GPU：4台のNVIDIA A100 (80GB)

KOALA-Lightningモデル

概要説明

これらの1024x1024のサンプルは、NVIDIA 4090 GPU上で0.66秒で10回のノイズ除去ステップを使用してKOALA-Lightning-700Mによって生成されました。

アーキテクチャ

残差ブロックとトランスフォーマーブロックを削減することで実現された2種類の圧縮U-Net、KOALA-1BとKOALA-700Mがあります。

U-Net比較

• Txはトランスフォーマーブロック、CKPTは訓練されたチェックポイントファイルを意味します。
• レイテンシーは、NVIDIA 4090 GPU (24GB)でFP16精度、25回のノイズ除去ステップで測定されました。
• SDM-v2.0は768x768の解像度を使用し、SDXLとKOALAモデルは1024x1024の解像度を使用します。

異なるGPUでのレイテンシーとメモリ使用量の比較

NVIDIA 3060Ti (8GB)、2080Ti (11GB)、および4090 (24GB)などのさまざまなコンシューマーグレードのGPUを使用して、SDXL-TurboとKOALA-Turboモデルの推論時間を512x512の解像度で、その他のモデルを1024x1024の解像度で測定しました。'OOM'はメモリ不足を示します。SDXLモデルは8GBのVRAMを持つ3060Tiでは動作できませんが、我々のKOALAモデルはすべてのGPUタイプで動作できます。

主な特徴

• 効率的なU-Netアーキテクチャ：KOALAモデルは、前作のStable Diffusion XL (SDXL)と比較して、モデルサイズを最大54%と69%削減する簡略化されたU-Netアーキテクチャを使用しています。
• 自己注意ベースの知識蒸留：KOALAの核心技術は、自己注意特徴の蒸留に焦点を当てており、これが画像生成品質を維持するために重要であることが証明されています。

モデルの説明

• 開発者：ETRI Visual Intelligence Lab
• 開発者名：Youngwan Lee、Kwanyong Park、Yoorhim Cho、Young-Ju Lee、Sung Ju Hwang
• モデルの説明：潜在拡散ベースのテキストから画像への生成モデル。KOALAモデルは、SDXL-Base-1.0と同じテキストエンコーダーを使用し、ノイズ除去U-Netのみを圧縮U-Netに置き換えています。
• 教師モデル：SDXL-Lightning
• 訓練データセット：LAION-POPデータセットの一部
• 訓練反復回数：バッチサイズ128で500K回
• GPU：4台のNVIDIA A100 (80GB)
• 詳細情報のリソース：arXivのKOALAレポートとプロジェクトページを参照してください。

💻 使用例

基本的な使用法

ノイズ除去ステップ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]

用途

直接的な使用

このモデルは研究目的のみを想定しています。可能な研究分野やタスクには以下が含まれます。

• アートワークの生成とデザインやその他のアートプロセスでの使用。
• 教育用または創造的なツールでの応用。
• 生成モデルの研究。
• 有害なコンテンツを生成する可能性のあるモデルの安全な展開。
• 生成モデルの制限とバイアスの調査と理解。
• 以下に説明する除外される使用法は除きます。

想定外の使用

このモデルは、人やイベントの事実的または真実の表現として訓練されていないため、このようなコンテンツを生成するためにモデルを使用することは、このモデルの能力範囲外です。

制限とバイアス

• テキストレンダリング：モデルは、画像内に長い読みやすいテキストをレンダリングする際に課題に直面しています。
• 複雑なプロンプト：KOALAは、複数の属性を含む複雑なプロンプトに対処するのに苦労することがあります。
• データセット依存性：現在の制限は、部分的に訓練データセット(LAION-aesthetics-V2 6+)の特性に起因しています。

引用

@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}
}