![見出し画像](https://assets.st-note.com/production/uploads/images/129737777/rectangle_large_type_2_8bdcb6d7bf9cb3a0f41b1e18f1b78005.png?width=800)
SegMoE: Segmind Mixture of Diffusion Experts の概要
以下の記事が面白かったので、かるくまとめました。
1. SegMoE
「SegMoE」は、MoE Diffusionモデルをゼロから作成するためのフレームワークです。「SegMoE」はHuggingFaceエコシステムに統合されており、diffusersでサポートされています。
「SegMoE」は、「Stable Diffusion」と同じアーキテクチャに従います。「Mixtral 8x7b」と同様、「SegMoE」には複数のモデルが1つにまとめられています。これが機能する方法は、一部のフィードフォワードレイヤーを疎なMoEレイヤーに置き換えることです。MoEレイヤーには、どのエキスパートがどのトークンを最も効率的に処理するかを選択するためのルーターネットワークが含まれています。「segmoe」パッケージを使用して独自のMoEモデルを作成できます。このプロセスには数分しかかかりません。詳細については、Githubリポジトリを参照してください。人気のあるライブラリ「mergekit」からインスピレーションを得て「segmoe」を設計しています。
「MoE」の詳細については、HuggingFaceの投稿を参照してください。
2. モデル と カスタムMoE作成コード
2-1. モデル
本日 (2月3日) リリースされるモデルは、次のとおりです。
・SegMoE 4x2 : 4つのエキスパートモデルで構成
・SegMoE 2x1 : 2つのエキスパートモデルで構成
・SegMoE SD 4x2 : 4つのStable Diffusion 1.5 エキスパート モデルで構成
「SegMoE MoE」は「SegMoE-AxB」と呼ばれます。ここで、Aはエキスパートモデル「MoE-d」の数を指し、2番目の数字は各画像の生成に関与するエキスパートの数を指します。構成設定に応じて、モデルの一部のレイヤー (フィードフォワードブロック、アテンション、またはすべて) のみが複製されます。 残りのパラメータはStableDiffusionモデルと同じです。
2-2. カスタムMoe作成コード
カスタムMoE作成コードは以下のリポジトリで公開しています。
3. サンプル
「SegMoe 4x2」のサンプルは、次のとおりです。
![](https://assets.st-note.com/img/1706991695528-4lYTj3dHGc.png?width=800)
「SegMoe 2x1」のサンプルは、次のとおりです。
![](https://assets.st-note.com/img/1706991703290-ZYKtNTEm5o.png?width=800)
「SegMoe SD 4x2」のサンプルは、次のとおりです。
![](https://assets.st-note.com/img/1706991775903-yvXaLCE8R2.png?width=800)
4. diffusersの使用
(1) パッケージのインストール。
pip install -U segmoe diffusers transformers
(2) SegMoE-4x2 による画像生成。
from segmoe import SegMoEPipeline
pipeline = SegMoEPipeline("segmind/SegMoE-4x2-v0", device="cuda")
prompt = "cosmic canvas, orange city background, painting of a chubby cat"
negative_prompt = "nsfw, bad quality, worse quality"
img = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
height=1024,
width=1024,
num_inference_steps=25,
guidance_scale=7.5,
).images[0]
img.save("image.png")
![](https://assets.st-note.com/img/1706991874972-uNsAp0luC2.png?width=800)
ローカルモデルを使用することもできます。
from segmoe import SegMoEPipeline
pipeline = SegMoEPipeline("segmoe_v0", device="cuda")
prompt = "cosmic canvas, orange city background, painting of a chubby cat"
negative_prompt = "nsfw, bad quality, worse quality"
img = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
height=1024,
width=1024,
num_inference_steps=25,
guidance_scale=7.5,
).images[0]
img.save("image.png")
独自の SegMoE を作成する方法は、「Creating your Own SegMoE」を参照してください。
5. 比較
以下の画像に示すように、プロンプトの理解が向上しているようです。各画像は、左からSegMoE-2x1-v0、SegMoE-4x2-v0、ベースモデル (RealVisXL_V3.0) です。
・three green glass bottles (緑色のガラス瓶3本)
![](https://assets.st-note.com/img/1706992134023-iJT5IR7Gtc.png?width=800)
・panda bear with aviator glasses on its head (頭に飛行用の眼鏡をかけたパンダのクマ)
![](https://assets.st-note.com/img/1706992179359-GZ2lL80YS7.png?width=800)
・the statue of Liberty next to the Washington Monument (ワシントン記念塔の隣にある自由の女神像)
![](https://assets.st-note.com/img/1706992230600-n9uHX3VADD.png?width=800)
・Taj Mahal with its reflection. detailed charcoal sketch. (タージ・マハルとその反射。 詳細な木炭スケッチ)
![](https://assets.st-note.com/img/1706992314755-ftqYP9f6AP.png?width=800)
6. 独自の SegMoE の作成
「config.yaml」を準備するだけです。
base_model: 基本モデルのパス、モデル カード、CivitAIダウンロード リンク
num_experts: 使用するエキスパートの数
moe_layers: 混合するレイヤーのタイプ (ff, attn, all : defaultはattn)
num_experts_per_tok: 使用するエキスパートの数
experts:
- source_model: Expert 1 パス、モデル カード、CivitAIダウンロード リンク
positive_prompt: ゲート重みを計算するためのポジティブ プロンプト
negative_prompt: ゲート重みを計算するためのネガティブ プロンプト
- source_model: Expert 2 パス、モデル カード、CivitAIダウンロード リンク
positive_prompt: ゲート重みを計算するためのポジティブ プロンプト
negative_prompt: ゲート重みを計算するためのネガティブ プロンプト
- source_model: Expert 3 パス、モデル カード、CivitAIダウンロード リンク
positive_prompt: ゲート重みを計算するためのポジティブ プロンプト
negative_prompt: ゲート重みを計算するためのネガティブ プロンプト
- source_model: Expert 4 パス、モデル カード、CivitAIダウンロード リンク
positive_prompt: ゲート重みを計算するためのポジティブ プロンプト
negative_prompt: ゲート重みを計算するためのネガティブ プロンプト
任意の数のモデルを組み合わせることができます。「config.yaml」の詳細については、githubリポジトリを参照してください。
HuggingFace モデルとCivitAI モデルの両方がサポートされています。CivitAIモデルの場合は、モデルのダウンロード リンクを貼り付けます。
(例: "https://civitai.com/api/download/models/239306")
次に、コマンドを実行します。
segmoe config.yaml segmoe_v0
これにより、次の構造を持つ segmoe_v0 というフォルダが作成されます。
├── model_index.json
├── scheduler
│ └── scheduler_config.json
├── text_encoder
│ ├── config.json
│ └── model.safetensors
├── text_encoder_2
│ ├── config.json
│ └── model.safetensors
├── tokenizer
│ ├── merges.txt
│ ├── special_tokens_map.json
│ ├── tokenizer_config.json
│ └── vocab.json
├── tokenizer_2
│ ├── merges.txt
│ ├── special_tokens_map.json
│ ├── tokenizer_config.json
│ └── vocab.json
├── unet
│ ├── config.json
│ └── diffusion_pytorch_model.safetensors
└──vae
├── config.json
└── diffusion_pytorch_model.safetensors
あるいは、Python APIを使用して、MoEを作成することもできます。
from segmoe import SegMoEPipeline
pipeline = SegMoEPipeline("config.yaml", device="cuda")
pipeline.save_pretrained("segmoe_v0")
7. HubへのPush
モデルは、huggingface-cli 経由でHubにPushできます。
huggingface-cli upload segmind/segmoe_v0 ./segmoe_v0
モデルは Python から直接HubにPushすることもできます。
from huggingface_hub import create_repo, upload_folder
model_id = "segmind/SegMoE-v0"
repo_id = create_repo(repo_id=model_id, exist_ok=True).repo_id
upload_folder(
repo_id=repo_id,
folder_path="segmoe_v0",
commit_message="Inital Commit",
ignore_patterns=["step_*", "epoch_*"],
)
8. 免責事項と進行中の作業
・速度が遅い : トークンあたりのエキスパートの数が1より大きい場合、MoE は複数のエキスパートモデルにわたって計算を実行します。これにより、単一のSD 1.5またはSDXLよりも遅くなります。
・VRAMの使用量が多い : MoE は推論を非常に高速に実行しますが、それでも大量の VRAM (したがって高価な GPU) が必要です。このため、ローカルで使用するのは困難ですが、複数のGPUを使用する展開には最適です。参考として、SegMoE-4x2には半精度で24GBのVRAMが必要です。
9. おわりに
推論時間を低く抑えながら、事前学習されたモデルを組み合わせるだけで、SOTA diffuserモデルを簡単に作成できる可能性のある新しいツールをコミュニティに提供するため、「SegMoE」は作成されました。これを使って何を構築できるか楽しみです。
この記事が気に入ったらサポートをしてみませんか?