Segmind-Vega

📣 Read our technical report for more details on our disillation method

Demo

Try out the Segmind-Vega model at Segmind-Vega for ⚡ fastest inference.

Model Description

The Segmind-Vega Model is a distilled version of the Stable Diffusion XL (SDXL), offering a remarkable 70% reduction in size and an impressive 100% speedup while retaining high-quality text-to-image generation capabilities. Trained on diverse datasets, including Grit and Midjourney scrape data, it excels at creating a wide range of visual content based on textual prompts.

Employing a knowledge distillation strategy, Segmind-Vega leverages the teachings of several expert models, including SDXL, ZavyChromaXL, and JuggernautXL, to combine their strengths and produce compelling visual outputs.

Image Comparison (Segmind-Vega vs SDXL)

Speed Comparison (Segmind-Vega vs SD-1.5 vs SDXL)

The tests were conducted on an A100 80GB GPU. (Note: All times are reported with the respective tiny-VAE!)

Parameters Comparison (Segmind-Vega vs SD-1.5 vs SDXL)

Usage:

This model can be used via the 🧨 Diffusers library.

Make sure to install diffusers by running

pip install diffusers

In addition, please install transformers, safetensors, and accelerate:

pip install transformers accelerate safetensors

To use the model, you can run the following:

from diffusers import StableDiffusionXLPipeline
import torch

pipe = StableDiffusionXLPipeline.from_pretrained("segmind/Segmind-Vega", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
pipe.to("cuda")
# if using torch < 2.0
# pipe.enable_xformers_memory_efficient_attention()
prompt = "A cute cat eating a slice of pizza, stunning color scheme, masterpiece, illustration" # Your prompt here
neg_prompt = "(worst quality, low quality, illustration, 3d, 2d, painting, cartoons, sketch)" # Negative prompt here
image = pipe(prompt=prompt, negative_prompt=neg_prompt).images[0]

Please do use negative prompting and a CFG around 9.0 for the best quality!

Model Description

• Developed by: Segmind
• Developers: Yatharth Gupta and Vishnu Jaddipal.
• Model type: Diffusion-based text-to-image generative model
• License: Apache 2.0
• Distilled From: stabilityai/stable-diffusion-xl-base-1.0

Key Features

• Text-to-Image Generation: The Segmind-Vega model excels at generating images from text prompts, enabling a wide range of creative applications.

• Distilled for Speed: Designed for efficiency, this model offers an impressive 100% speedup, making it suitable for real-time applications and scenarios where rapid image generation is essential.

• Diverse Training Data: Trained on diverse datasets, the model can handle a variety of textual prompts and generate corresponding images effectively.

• Knowledge Distillation: By distilling knowledge from multiple expert models, the Segmind-Vega Model combines their strengths and minimizes their limitations, resulting in improved performance.

Model Architecture

The Segmind-Vega Model is a compact version with a remarkable 70% reduction in size compared to the Base SDXL Model.

Training Info

These are the key hyperparameters used during training:

• Steps: 540,000
• Learning rate: 1e-5
• Batch size: 16
• Gradient accumulation steps: 8
• Image resolution: 1024
• Mixed-precision: fp16

Model Sources

For research and development purposes, the Segmind-Vega Model can be accessed via the Segmind AI platform. For more information and access details, please visit Segmind.

Uses

Direct Use

The Segmind-Vega Model is suitable for research and practical applications in various domains, including:

• Art and Design: It can be used to generate artworks, designs, and other creative content, providing inspiration and enhancing the creative process.

• Education: The model can be applied in educational tools to create visual content for teaching and learning purposes.

• Research: Researchers can use the model to explore generative models, evaluate its performance, and push the boundaries of text-to-image generation.

• Safe Content Generation: It offers a safe and controlled way to generate content, reducing the risk of harmful or inappropriate outputs.

• Bias and Limitation Analysis: Researchers and developers can use the model to probe its limitations and biases, contributing to a better understanding of generative models' behavior.

Downstream Use

The Segmind-Vega Model can also be used directly with the 🧨 Diffusers library training scripts for further training, including:

• LoRA:

export MODEL_NAME="segmind/Segmind-Vega"
export VAE_NAME="madebyollin/sdxl-vae-fp16-fix"
export DATASET_NAME="lambdalabs/pokemon-blip-captions"

accelerate launch train_text_to_image_lora_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --pretrained_vae_model_name_or_path=$VAE_NAME \
  --dataset_name=$DATASET_NAME --caption_column="text" \
  --resolution=1024 --random_flip \
  --train_batch_size=1 \
  --num_train_epochs=2 --checkpointing_steps=500 \
  --learning_rate=1e-04 --lr_scheduler="constant" --lr_warmup_steps=0 \
  --mixed_precision="fp16" \
  --seed=42 \
  --output_dir="vega-pokemon-model-lora" \
  --validation_prompt="cute dragon creature" --report_to="wandb" \
  --push_to_hub

• Fine-Tune:

export MODEL_NAME="segmind/Segmind-Vega"
export VAE_NAME="madebyollin/sdxl-vae-fp16-fi