Package: segmoe
SegMoE: Segmind Mixture of Diffusion ExpertsSegMoE is a powerful framework for dynamically combining Stable Diffusion Models into a Mixture of Experts within minutes without training. The framework allows for creation of larger models on the fly which offer larger knowledge, better adherence and better image quality. It is inspired bymergekit's mixtral branch but for Stable Diffusion models.InstallationpipinstallsegmoeUsageLoad Checkpoint from Hugging FaceWe release 3 merges on Hugging Face,SegMoE 2x1has two expert models.SegMoE 4x2has four expert models.SegMoE SD 4x2has four Stable Diffusion 1.5 expert models.They can be loaded as follows:fromsegmoeimportSegMoEPipelinepipeline=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")ComparisonThe Prompt Understanding seems to improve as shown in the images below. From Left to Right SegMoE-2x1-v0, SegMoE-4x2-v0, Base Model (RealVisXL_V3.0)three green glass bottlespanda bear with aviator glasses on its headthe statue of Liberty next to the Washington MonumentCreating your Own ModelCreate a yaml config file, config.yaml, with the following structure:base_model:Base Model Path, Model Card or CivitAI Download Linknum_experts:Number of experts to usemoe_layers:Type of Layers to Mix (can be "ff", "attn" or "all"). Defaults to "attn"num_experts_per_tok:Number of Experts to usetype:Type of the individual models (can be "sd" or "sdxl"). Defaults to "sdxl"experts:-source_model:Expert 1 Path, Model Card or CivitAI Download Linkpositive_prompt:Positive Prompt for computing gate weightsnegative_prompt:Negative Prompt for computing gate weights-source_model:Expert 2 Path, Model Card or CivitAI Download Linkpositive_prompt:Positive Prompt for computing gate weightsnegative_prompt:Negative Prompt for computing gate weights-source_model:Expert 3 Path, Model Card or CivitAI Download Linkpositive_prompt:Positive Prompt for computing gate weightsnegative_prompt:Negative Prompt for computing gate weights-source_model:Expert 4 Path, Model Card or CivitAI Download Linkpositive_prompt:Positive Prompt for computing gate weightsnegative_prompt:Negative Prompt for computing gate weightsAny number of models can be combined, An Example config can be foundhere. For detailed information on how to create a config file, please refer to theConfig ParametersNoteBoth 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:segmoeconfig.yamlsegmoe_v0This will create a folder called segmoe_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
└──vae├──config.json└──diffusion_pytorch_model.safetensorsAlternatively, you can also use the following command to create a mixture of experts model:fromsegmoeimportSegMoEPipelinepipeline=SegMoEPipeline("config.yaml",device="cuda")pipeline.save_pretrained("segmoe_v0")Push to HubThe Model can be pushed to the hub via the huggingface-clihuggingface-cliuploadsegmind/segmoe_v0./segmoe_v0Detailed usage can be foundhereSDXL TurboTo use SDXL Turbo style models, just change the scheduler to DPMSolverMultistepScheduler. Example config can be foundhereUsage:fromsegmoeimportSegMoEPipelinepipeline=SegMoETurboPipeline("segomoe_config_turbo.yaml",device="cuda")pipeline.pipe.scheduler=DPMSolverMultistepScheduler.from_config(pipeline.pipe.scheduler.config)prompt="cosmic canvas, orange city background, painting of a chubby cat"image=pipe(prompt=prompt,num_inference_steps=6,guidance_scale=2).images[0]image.save("image.png")Stable Diffusion 1.5 SupportStable Diffusion 1.5 Models are also supported and work natively. Example config can be foundhereNote:Stable Diffusion 1.5 Models can be combined with other SD1.5 Models only.Memory RequirementsSDXL 2xN : 19GBSDXL 4xN : 25GBSD1.5 4xN : 7GBAdvantagesBenefits from The Knowledge of Several Finetuned ExpertsTraining FreeBetter Adaptability to DataModel Can be upgraded by using a better finetuned model as one of the experts.LimitationsThough the Model improves upon the fidelity of images as well as adherence, it does not be drastically better than any one expert without training and relies on the knowledge of the experts.This is not yet optimized for speed.The framework is not yet optimized for memory usage.Research RoadmapOptimize for SpeedOptimize for Memory UsageAdd Support for LoRAsAdd Support for More ModelsAdd Support for TrainingConfig ParametersBase ModelThe 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 or a local path to a safetensors file.Number of ExpertsThe number of experts to use in the mixture of experts model. The number of experts must be greater than 1. The Number of experts can be anything greater than 2 as long as the GPU fits it.MOE LayersThe type of layers to mix. Can be "ff", "attn" or "all". Defaults to "attn". "ff" merges only the feedforward layers, "attn" merges only the attention layers and "all" merges all layers.TypeThe type of the models to mix. Can be "sd" or "sdxl". Defaults to "sdxl".ExpertsThe 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 or a local path to a safetensors file. The positive prompt and negative prompt are the prompts that will be used to compute the gate weights for each expert and impact the quality of the final model, choose these carefully.Citation@misc{segmoe,author={Yatharth Gupta, Vishnu V Jaddipal, Harish Prabhala},title={SegMoE: Segmind Mixture of Diffusion Experts},year={2024},publisher={GitHub},journal={GitHub repository},howpublished={\url{https://github.com/segmind/segmoe}}}