跳到正文
原文
Hugging Face Blog·· 2024-02-03精选AI 评分60

Segmind 发布 SegMoE 扩散专家混合框架

SegMoE: Segmind Mixture of Diffusion Experts

AI 导读

Segmind 发布 SegMoE 框架，可从零构建混合专家（MoE）扩散模型，并已集成 diffusers，同时开源 SegMoE-4x2、SegMoE-2x1 和 SegMoE-SD4x2 三个模型及建模代码（Apache 2.0）。

推荐理由

原文给出了从配置到生产的完整流程和显存占用数据，读者可以据此判断是否用 MoE 方式拼装扩散模型。

正文 · AI 翻译

SegMoE 是一个令人兴奋的框架，可以从零开始创建混合专家扩散模型！SegMoE 全面集成在 HF Mirror 生态系统中，并支持 diffusers 🔥！

今天发布的功能和集成包括：

目录

什么是 SegMoE？

SegMoE 模型遵循与 Stable Diffusion 相同的架构。像 Mixtral 8x7b 一样，SegMoE 模型将多个模型合而为一。其工作方式是将一些前馈层替换为稀疏 MoE 层。MoE 层包含一个路由网络，用于选择哪些专家最有效地处理哪些 token。 你可以使用 segmoe 包创建自己的 MoE 模型！整个过程只需几分钟。如需更多信息，请访问 Github 仓库。我们从流行的库 mergekit 中汲取灵感来设计 segmoe。我们感谢 mergekit 的贡献者提供了如此有用的库。

有关 MoE 的更多细节，请参阅 HF Mirror 🤗 文章：hf.co/blog/moe。

SegMoE 发布 TL;DR；

  • 发布 SegMoE-4x2、SegMoE-2x1 和 SegMoE-SD4x2 版本
  • 发布自定义 MoE 制作代码

关于名字

SegMoE MoE 被称为 SegMoE-AxB，其中 A 指 MoE 合并在一起的专家模型数量，而第二个数字指生成每张图像时涉及的专家数量。根据配置设置，模型的仅部分层（前馈块、注意力或全部）会被复制；其余参数与 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 包。确保你安装了最新版本的 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")

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-v0、SegMoE-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 仓库

注意 同时支持 HF Mirror 和 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_*"],
)

详细用法可在此处找到

免责声明与正在进行的工作

  • 速度较慢：如果每个 token 的专家数量大于 1，MoE 会在多个专家模型之间进行计算。这使得它比单个 SD 1.5 或 SDXL 模型更慢。

  • 显存占用高：MoE 的推理速度非常快，但仍需要大量显存（因此需要昂贵的 GPU）。这使得它们在本地环境中难以使用，但它们非常适合多 GPU 部署。作为参考，SegMoE-4x2 在半精度下需要 24GB 显存。

结论

我们构建 SegMoE 是为了给社区提供一个新工具，只需组合预训练模型，同时保持较低的推理时间，就有可能轻松创建 SOTA 扩散模型。我们很期待看到你能用它构建出什么！

其他资源

来源：Hugging Face Blog · hf-mirror.com

#产品更新#图像生成#开源生态#Hugging Face