なにメモ

コンピュータビジョンや機械学習関係の話題を書き綴ると思うブログです。

Cool Japan Diffusionの開発体制

はじめに

今回はCool Japan Diffusionの開発体制を説明します。 ただし、あくまで表に出せる範囲でしか出しません。

Cool Japan Diffusionのマシンパワー(表側)

Cool Japan Diffusionは主に我が家のMAGIによって開発されています。

Cool Japan Diffusionの開発体制の概要図(矢印はデータの流れ)

MELCHIOR (メルキオール) : Cool Japan Diffusionの学習を担当する。

  • GPU: RTX 4090 (実測上A100よりも速い) 、GTX 1080 (明日RTX 3060に替える予定)
  • CPU: 第12世代Core i5
  • メインメモリ: 128GB
  • 不揮発性メモリ: 1TB
  • ストレージ (SSD): 1TB
  • ネットワーク: 2.5 GbE

BALTHASAR (バルタザール) : Cool Japan Diffusionの学習に使う画像のキャプション付けを担当する。(Clean Diffusionの開発も担当している。)

  • GPU: RTX 3090、RTX 3080
  • CPU: 第12世代Core i5
  • メインメモリ: 128GB
  • ストレージ (SSD): 2TB
  • ネットワーク: 2.5 GbE

CASPER (カスパー) : Stable Diffusionを追加学習することで生成器を作成したり、学習用データを収集したりすることを担当する。

  • GPU: RTX 3090、GTX 1660
  • CPU: 第12世代Core i5
  • メインメモリ: 64GB
  • 不揮発性メモリ : 1TB
  • ストレージ (SSD): 2TB
  • ストレージ (HDD): 16TB
  • ネットワーク: 2.5 GbE

ARTABAN (アルタバン) : ありとあらゆるデータを確実に保管し、各マシンに供給することを担当する。いわゆるNAS

  • ストレージ: 約20TB (RAID5、スナップショット付き。元々は32TB)
  • ネットワーク: 2.5 GbE

MAGI以外にもクラウドで活躍しているインスタンスたちもいます。

illustrator (イラストレーター) : メルキオールでは対応しきれない学習を担当する。

  • GPU: A100 x 1以上 (x2がコスパいいという仕事上の経験がある)
  • CPU: 12 vCPU以上
  • メインメモリ: 85GB以上
  • ストレージ (SSHD): 100GB以上
  • ストレージ (HDD): 1TB以上

generator n (生成器 n) : Stable Diffusionの知識を抽出するために画像を生成している。

MBA (MacBook Air) : 主に事務処理や裏側に接続することを担当している。

あとは裏側にもいます。これは秘密ですね。

Cool Japan Diffusionのデータ・リソース(表側)

Cool Japan Diffusionのデータ・リソースは倫理的な判断により自身でリミットを課しています。 基本的に絵画やイラスト、アニメ、漫画、ゲーム画面などが用意されています。

  1. パブリックドメインのデータ
  2. 無償公開されている違法ダウンロードではないインターネットのデータ、追加学習したStable Diffusionにより生成されたデータ
  3. 第一段階、第二段階を除く、情報解析として合法的に、かつ、倫理的に利用できるデータ
  4. その他、情報解析として合法的に利用できるデータ
  5. (おまけの画像)

現在、第2段階まで開放されています。現状は、合計1650万枚(おまけを入れると1億枚)以上で、各段階の枚数は次のとおりです。

  1. 200万枚以上
  2. 400万枚以上
  3. 250万枚以上
  4. 800万枚以上
  5. (8000万枚以上)

合計は第4段階までで1億枚以上まで増える予定です。特に第2, 3段階がもりもり増えていきます。 現状でも引っ張り出せばまだあると思いますが、把握しきれていません。

Cool Japan Diffusionの開発

Cool Japan Diffusionは、解明された拡散モデル [1] によるカスケードモデル [2] を採用しています。 このモデルを採用した画像生成AIはまだ見たことがありません。 おそらく現状の実用レベルではこれが限界だと思われます。

  1. 拡散モデルによる128x128 (超解像による512x512)
  2. 拡散モデルによる512x512 (超解像による2048x2048)
  3. 拡散モデルによる2048x2048 (超解像による8196x8196)

実用上は第2段階でサービス提供可能になります。 第1段階の出力は次のとおりです。

この細かさで第2段階は出力されます。 これが純粋な拡散モデルの力です。 1ピクセル1ピクセルで出力されるのです。 しかし、まだ学習途中なので、顔が崩れがちになっていますね。 また、現状超解像で次のとおりになります。

超解像アルゴリズムとしてManga109データセットではSOTAであるHATを採用しています。 この手元にあるHATのモデルには上記のデータ・リソースがふんだんに使われており、 イラストに対しては現状最高性能の超解像だと自負しています。 このクオリティ程度で512x512 (2048x2048) でまずは提供する予定です。 超解像を含めて生成時間は20秒程度です。 期待に答えられるようにがんばります。

Cool Japan Diffusionの課題

マシンパワーが足りていません。あとクラウドがあんまり割り当ててくれません。 H100の登場が待ち遠しいです。マシンパワーをください。

以上です。

参考文献

[1] Tero Karras et al., "Elucidating the Design Space of Diffusion-Based Generative Models," NeurIPS 2022

[2] Chitwan Saharia et al., "Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding," NeurIPS 2022