OneDiffusion：多彩な拡散モデルによる画像生成・理解タスクへの対応（GoogleColab付）

## はじめに

本ノートブックでは、[OneDiffusion](https://github.com/lehduong/OneDiffusion)と呼ばれる汎用的な大規模拡散モデルについて紹介します。このモデルは、多種多様な画像生成・画像理解タスクに双方向的（テキスト→画像、画像→記述、画像編集など）に対応することを目指して設計されています。

このノートブックでは、以下をカバーします。

- OneDiffusionモデルの概要
- モデルのインストール方法
- 基本的なテキスト→画像生成(T2I)実行例
- Gradioデモの起動方法および説明
- モデルが対応する拡張タスク例（IDカスタマイズ、多視点生成、条件画像からの生成、テキストガイドによる画像編集など）

<blockquote class="twitter-tweet" data-media-max-width="800"><p lang="ja" dir="ltr">One Diffusion動いたぞ！！！ <a href="https://t.co/oBaJnmU40n">https://t.co/oBaJnmU40n</a> <a href="https://t.co/b1zj2UocP2">pic.twitter.com/b1zj2UocP2</a></p>&mdash; Maki@Sunwood AI Labs. (@hAru_mAki_ch) <a href="https://twitter.com/hAru_mAki_ch/status/1867156896798683419?ref_src=twsrc%5Etfw">December 12, 2024</a></blockquote> 

## ニュース

- **2024/12/10**: 重みをリリース
- **2024/12/06**: インストラクション（指示文）からの画像編集機能を追加
- **2024/12/02**: 対象（subject）駆動型生成（Subject-driven generation）を追加



## モデルについて

OneDiffusionは、様々な画像生成・理解タスクを1つの拡散モデルでサポートすることを目標としたフレームワークです。例えば、テキストから画像を生成するだけでなく、単一の入力画像から異なる視点画像を生成したり、画像上でのテキストガイドによる編集を行ったりします。また、対象特定（IDカスタマイズ）やマルチタスク的なゼロショット組み合わせなど、多様な条件付け生成を行うことができます。



## インストール手順

以下はGoogle Colab上でのセットアップ例です。  

```python
!git clone https://github.com/lehduong/OneDiffusion.git
%cd /content/OneDiffusion
```

```python
!apt install g++
```



```python
# PyTorchを2.1.2にダウングレード
!pip install torch==2.3.1
# !pip install flash-attn==2.5.7
!pip install flash-attn
```

```python
#!pip install torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cu118
!pip install "git+https://github.com/facebookresearch/pytorch3d.git"
!pip install -r requirements.txt
```

**注意**:  
- `conda`コマンドは、Colab環境ではそのまま使用できない場合があります。その際は`mamba`や`venv`等で環境を分けたり、あるいはGoogle Colabの標準Python環境に直接`pip install`でパッケージを導入することも可能です。
- 環境によってはPyTorchやpytorch3dのインストールに追加ステップが必要となる場合があります。



## クイックスタート：テキスト→画像生成 (Text-to-Image)

以下は`inference.py`を参考にした最も基本的なテキスト→画像生成例です。

```python
!pip uninstall Pillow -y
!pip install Pillow
```

```python
import torch
from onediffusion.diffusion.pipelines.onediffusion import OneDiffusionPipeline

device = torch.device('cuda:0')

# 事前学習済みモデルの読み込み
pipeline = OneDiffusionPipeline.from_pretrained("lehduong/OneDiffusion").to(device=device, dtype=torch.bfloat16)

# ネガティブプロンプト：望ましくない要素を排除するため
NEGATIVE_PROMPT = "monochrome, greyscale, low-res, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation"

prompt = "[[text2image]] A bipedal black cat wearing a huge oversized witch hat, a wizards robe, casting a spell, in an enchanted forest. The scene is filled with fireflies and moss on surrounding rocks and trees"

# 画像生成
output = pipeline(
    prompt=prompt,
    negative_prompt=NEGATIVE_PROMPT,
    num_inference_steps=50,
    guidance_scale=4,
    height=1024,
    width=1024,
)
image = output.images[0]
image.save('text2image_output.jpg')
image
```

上記セルを実行することで、指定プロンプトに従った1024x1024ピクセルの画像が生成されます。



## Gradioデモの起動

OneDiffusionには、Gradioデモが用意されています。これにより、ブラウザ上でモデルに対するインタラクティブな実験が可能です。以下コマンドでデモが起動できます。

```python
!python gradio_demo.py --captioner molmo
```

実行時に3つの`captioner`モードが選択可能です：

- `molmo`: 高品質なキャプション生成が可能だが、VRAM使用量が大きく(**約40GB以上**必要)
- `llava`: `molmo`より軽量(**約27GB以上**必要)だが、キャプション精度は若干低下
- `disable`: キャプション生成を無効化(**約12GB以上**でOK)。手動でキャプションを与える必要あり。ただし、このモードでも「Caption」ボタンは押す必要がある（空のキャプションを整形するため）。

**注意**:  
GPUメモリ要件は、入力画像数や画像解像度を変更すると増減します。


<blockquote class="twitter-tweet" data-media-max-width="800"><p lang="ja" dir="ltr">One DiffusionのGradioアプリも動いたぞ！！！<br>＊📒ノートブックは後ほど公開します！ <a href="https://t.co/zkT01ejdqh">https://t.co/zkT01ejdqh</a> <a href="https://t.co/skXDr1dDuc">pic.twitter.com/skXDr1dDuc</a></p>&mdash; Maki@Sunwood AI Labs. (@hAru_mAki_ch) <a href="https://twitter.com/hAru_mAki_ch/status/1867157232816939303?ref_src=twsrc%5Etfw">December 12, 2024</a></blockquote> 



## 対応タスク例

1. **テキストから画像生成 (Text-to-Image)**  
   任意のテキストプロンプトを入力するだけで、拡散モデルを通じて高品質な画像を生成可能です。

2. **IDカスタマイズ (Subject-driven generation)**  
   モデルは、特定の対象（ペットや人物など）を指し示して、その対象の視覚的特徴を反映しながら新たなシーンやポーズで生成することができます。

3. **マルチビュー生成 (Multiview Generation)**  
   単一の画像を入力として、異なる視点（カメラアングル、方向）から見た画像を生成します。また、テキストからもマルチビュー画像を直接生成可能です。

4. **条件付け画像生成 (Condition-to-Image) と逆タスク (Image-to-Condition)**  
   条件画像（例：スケッチ、セグメンテーションマップ、ポーズ画像）をもとに写実的あるいは創造的な画像生成を行えます。さらに、画像から条件へのマッピングも可能です。

5. **テキストガイドによる画像編集 (Text-guide image editing)**  
   既存の画像に対し、テキストで指示を与えることで、一部を修正・編集した新たな画像を生成します。

6. **ゼロショット・タスク組み合わせ**  
   モデルは、明示的な追訴学習なしに複数タスクを組み合わせた指示に対応することも可能です（例：IDカスタマイズ＋インペインティング）。ただし、この組み合わせ的な使用はパフォーマンスが不安定な場合があり、必要に応じてモデルの追加微調整を推奨します。



## ライセンス

OneDiffusionは、非商用ライセンス（CC BY-NC）で提供されます。詳細は`LICENSE`ファイルおよび[該当リポジトリ](https://github.com/lehduong/OneDiffusion)を参照してください。



## 論文引用

```bibtex
@misc{le2024diffusiongenerate,
      title={One Diffusion to Generate Them All},
      author={Duong H. Le and Tuan Pham and Sangho Lee and Christopher Clark and Aniruddha Kembhavi and Stephan Mandt and Ranjay Krishna and Jiasen Lu},
      year={2024},
      eprint={2411.16318},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2411.16318},
}
```



## まとめ

本ノートブックでは、OneDiffusionモデルの概要、インストール手順、クイックスタート用コードサンプル、そして対応タスク例を紹介しました。OneDiffusionは様々な画像生成・理解タスクを一元的に扱う可能性を持ち、拡張性・柔軟性に富んだモデルです。ぜひ独自タスクや応用例で試してみてください。


## 📒ノートブック

https://colab.research.google.com/drive/1ERYQrObg9qhxorbdCntYn2uYLTbaK1Dt?usp=sharing

<script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>