SegMoE: Segmind 擴散專家混合模型

釋出於 2024 年 2 月 3 日
在 GitHub 上更新

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 個合併模型

  1. SegMoE 2x1 有兩個專家模型。
  2. SegMoE 4x2 有四個專家模型。
  3. SegMoE SD 4x2 有四個 Stable Diffusion 1.5 專家模型。

樣本

使用 SegMoE 4x2 生成的影像

image

使用 SegMoE 2x1 生成的影像

image

使用 SegMoE SD 4x2 生成的影像

image

使用 🤗 Diffusers

請執行以下命令安裝 segmoe 包。請確保您已安裝最新版本的 diffuserstransformers

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")

image

使用本地模型

或者,也可以載入本地模型,這裡 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-v0SegMoE-4x2-v0、基礎模型(RealVisXL_V3.0

image

三隻綠色玻璃瓶

image

戴飛行員眼鏡的熊貓

image

自由女神像旁邊的華盛頓紀念碑

image

泰姬陵及其倒影。精細的木炭素描。

建立您自己的 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 擴散模型。我們很高興看到您能用它構建出什麼!

附加資源

社群

註冊登入 發表評論

© . This site is unofficial and not affiliated with Hugging Face, Inc.