SegMoE: Segmind 擴散專家混合模型
SegMoE 是一個令人興奮的框架,用於從頭開始建立專家混合(Mixture-of-Experts)擴散模型!SegMoE 全面整合到 Hugging Face 生態系統中,並支援
diffusers
🔥!
今日釋出的功能和整合包括:
目錄
什麼是 SegMoE?
SegMoE 模型遵循與 Stable Diffusion 相同的架構。與 Mixtral 8x7b 類似,SegMoE 模型將多個模型合而為一。其工作原理是用稀疏的 MoE 層替換一些前饋層。MoE 層包含一個路由器網路,用於選擇哪些專家最有效地處理哪些標記。您可以使用 segmoe
包建立自己的 MoE 模型!此過程僅需幾分鐘。欲瞭解更多資訊,請訪問Github 倉庫。我們從流行的庫 mergekit
中汲取靈感來設計 segmoe
。我們感謝 mergekit
的貢獻者提供瞭如此有用的庫。
有關 MoE 的更多詳情,請參閱 Hugging Face 🤗 的文章:hf.co/blog/moe。
SegMoE 釋出 TL;DR;
- 釋出 SegMoE-4x2、SegMoE-2x1 和 SegMoE-SD4x2 版本
- 釋出自定義 MoE 建立程式碼
關於命名
SegMoE MoE 被稱為 SegMoE-AxB,其中 A
指的是混合在一起的專家模型的數量,而第二個數字指的是每個影像生成過程中涉及的專家數量。根據配置設定,模型的某些層(前饋塊、注意力層或所有層)被複制;其餘引數與 Stable Diffusion 模型相同。有關 MoE 工作原理的更多詳細資訊,請參閱“專家混合模型解釋”一文。
推理
我們釋出了 Hub 上的 3 個合併模型
- SegMoE 2x1 有兩個專家模型。
- SegMoE 4x2 有四個專家模型。
- SegMoE SD 4x2 有四個 Stable Diffusion 1.5 專家模型。
樣本
使用 SegMoE 4x2 生成的影像
使用 SegMoE 2x1 生成的影像
使用 SegMoE SD 4x2 生成的影像
使用 🤗 Diffusers
請執行以下命令安裝 segmoe
包。請確保您已安裝最新版本的 diffusers
和 transformers
。
pip install -U segmoe diffusers transformers
以下程式碼將載入上述列表中的第二個模型(“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")
使用本地模型
或者,也可以載入本地模型,這裡 segmoe_v0
是包含本地 SegMoE 模型的目錄路徑。請檢視建立您自己的 SegMoE,瞭解如何構建自己的模型!
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-2x1-v0、SegMoE-4x2-v0、基礎模型(RealVisXL_V3.0)
建立您自己的 SegMoE
只需準備一個 config.yaml
檔案,其結構如下:
base_model: Base Model Path, Model Card or CivitAI Download Link
num_experts: Number of experts to use
moe_layers: Type of Layers to Mix (can be "ff", "attn" or "all"). Defaults to "attn"
num_experts_per_tok: Number of Experts to use
experts:
- source_model: Expert 1 Path, Model Card or CivitAI Download Link
positive_prompt: Positive Prompt for computing gate weights
negative_prompt: Negative Prompt for computing gate weights
- source_model: Expert 2 Path, Model Card or CivitAI Download Link
positive_prompt: Positive Prompt for computing gate weights
negative_prompt: Negative Prompt for computing gate weights
- source_model: Expert 3 Path, Model Card or CivitAI Download Link
positive_prompt: Positive Prompt for computing gate weights
negative_prompt: Negative Prompt for computing gate weights
- source_model: Expert 4 Path, Model Card or CivitAI Download Link
positive_prompt: Positive Prompt for computing gate weights
negative_prompt: Negative Prompt for computing gate weights
任意數量的模型都可以組合。有關如何建立配置檔案的詳細資訊,請參閱 github 倉庫
注意:支援 Hugging Face 和 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 建立專家混合模型
from segmoe import SegMoEPipeline
pipeline = SegMoEPipeline("config.yaml", device="cuda")
pipeline.save_pretrained("segmoe_v0")
推送到 Hub
該模型可以透過 huggingface-cli 推送到 hub
huggingface-cli upload segmind/segmoe_v0 ./segmoe_v0
該模型也可以直接從 Python 推送到 Hub
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="Initial Commit",
ignore_patterns=["step_*", "epoch_*"],
)
詳細用法可在此處找到
免責宣告和正在進行的工作
速度較慢:如果每個標記的專家數量大於 1,MoE 會在多個專家模型上執行計算。這使其比單個 SD 1.5 或 SDXL 模型慢。
VRAM 使用量高:MoE 推理速度非常快,但仍需要大量 VRAM(因此需要昂貴的 GPU)。這使得它們在本地設定中難以使用,但它們非常適合多 GPU 部署。作為參考,SegMoE-4x2 在半精度下需要 24GB 的 VRAM。
結論
我們構建 SegMoE 旨在為社群提供一個新工具,透過簡單地組合預訓練模型,同時保持低推理時間,從而輕鬆建立 SOTA 擴散模型。我們很高興看到您能用它構建出什麼!