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 扩散模型。我们很高兴看到您能用它构建出什么！
