Package for Mixing Stable Diffusiion XL Models by Segmind
Project description
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}}
}
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.