SegMix: Segmind Mixture of Experts

A Framework to combine multiple Stable Diffusion XL models into a mixture of experts model. Functions simialar to mergekit's mixtral branch but for Stable Diffusion XL models.

Installation

pip install segmix

Usage

from segmix import SegMixPipeline

pipeline = SegMixPipeline("segmind/SegMix-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(f"image.png")

Creating your Own Model

Create a yaml config file, config.yaml, with the following structure:

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

For detailed information on how to create a config file, please refer to the Config Parameters

Note Both Huggingface Models and CivitAI Models are supported. For CivitAI models, paste the download link of the model, For Example: "https://civitai.com/api/download/models/239306"

Then run the following command:

segmix config.yaml segmix_v0

This will create a folder called segmix_v0 with the following structure:

├── 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
│   └── segmix_config.json
└──vae
    ├── config.json
    └── diffusion_pytorch_model.safetensors

Alternatively, you can also use the following command to create a mixture of experts model:

from segmix import SegMixPipeline

pipeline = SegMixPipeline("config.yaml", device="cuda")

pipeline.save_pretrained("segmix_v0")

Advantages

• Benefits from The Knowledge of Several Finetuned Experts
• Training Free
• Better Adaptability to Data

Limitations

• This is a very early version of the framework and is not yet optimized for speed.
• The framework is not yet optimized for memory usage.

Research Roadmap

• Optimize for Speed
• Optimize for Memory Usage
• Add Support for LoRAs
• Add Support for More Models
• Add Support for Training

Config Parameters

Base Model

The base model is the model that will be used to generate the initial image. It can be a Huggingface model card, a CivitAI model download link.

Number of Experts

The number of experts to use in the mixture of experts model. The number of experts must be greater than 1.

MOE Layers

The type of layers to mix. Can be "ff", "attn" or "all". Defaults to "attn"

Experts

The experts are the models that will be used to generate the final image. Each expert must have a source model, a positive prompt and a negative prompt. The source model can be a Huggingface model card, a CivitAI model download link. The positive prompt and negative prompt are the prompts that will be used to compute the gate weights for each expert.

Citation

@misc{segmix,
  author = {Yatharth Gupta, Vishnu V Jaddipal, Harish Prabhala},
  title = {SegMix: Segmind Mixture of Experts},
  year = {2024},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/segmind/segmix}}
}