Koala Lightning 1b

🚀 KOALA-Lightning-1Bモデルカード

このモデルは、SDXLのU-Netを圧縮し、知識蒸留を行うことで高速なテキストから画像への生成を実現します。SDXLよりも大幅に高速で、低リソース環境でも高解像度画像を生成できます。

🚀 クイックスタート

       

概要

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

✨ 主な機能

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

📦 インストール

以下は、🤗Diffusersライブラリを使用した推論コード例です。

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]

💻 使用例

基本的な使用法

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]

📚 ドキュメント

要約

• このモデルは、自己注意ベースの知識蒸留手法を用いて学習されています。
• 教師モデルはSDXL-Lightningです。
• 学習データセットはLAION-POPデータセットのサブセットです。
• 学習反復回数はバッチサイズ128で500K回、学習には4台のNVIDIA A100 (80GB)を使用しています。

KOALA-Lightningモデル

概要

アーキテクチャ

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でのレイテンシーとメモリ使用量の比較

様々な消費者向けGPU（NVIDIA 3060Ti (8GB)、2080Ti (11GB)、4090 (24GB)）を使用して、SDXL-TurboとKOALA-Turboモデルの推論時間を512x512の解像度で、その他のモデルを1024x1024の解像度で測定しました。'OOM'はメモリ不足を示します。SDXLモデルは8GBのVRAMを持つ3060Tiでは動作できませんが、私たちのKOALAモデルはすべてのGPUタイプで動作できます。

モデルの説明

• 開発元: 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レポートとプロジェクトページをご確認ください。

使用方法

直接使用

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

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

範囲外の使用

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

制限とバイアス

• テキストレンダリング：モデルは、画像内に長い読みやすいテキストをレンダリングする際に課題に直面しています。
• 複雑なプロンプト：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}
}