/diffusion-evolver

Evolve diffusion models by merging.

Primary LanguagePython

Diffusion Evolver

This project is an evolutionary framework for optimizing Stable Diffusion XL models. It allows you to evolve a population of models through crossover, mutation, and selection based on their performance judged by a VLM.

Examples

After Evolution Before Evolution Model
Model A *
Model B *
Model C *

After 10 cycles with prompt "T-Rex wearing aviator sunglasses, posing in front of a diffusion-generated Jurassic landscape, 80s vaporwave style, 4K"

* Images are paired with the closest model output from the initial population. All images share the same seed and diffusion settings. The VLM was not shown the prompt and default settings were used.

* Demo models are available here https://huggingface.co/martyn/sdxl-evolved-demo-models

Models

More models are available here https://huggingface.co/collections/martyn/evolved-sdxl-models-660b9185df88d3dbac68c052

and on Civit.ai https://civitai.com/user/chandro/models

Installation

  1. Clone the repository:

    git clone

  2. Install the required dependencies:

    pip install -r requirements.txt

  3. Set up the necessary API credentials:

    • Obtain an API key from Anthropic.
    • Set the ANTHROPIC_API_KEY environment variable with your API key.

System requirements

Minimum

  • A system capable of running inference on Stable Diffusion XL
  • Hard disk space(7GB*(population_size+initial_population))

Recommended

  • A modern GPU
  • A lot of ram(> 128 GB). More allows for larger population size
  • A ramdisk to store the evolved population

Usage

To run the evolutionary framework, use the following command:

python sdxl-evolve.py model_list.yml [options]

The model_list.yml file should contain a list of initial model candidates in YAML format.

Available options:

  • -seed: Random seed for reproducibility.
  • -cycles: Number of evolutionary cycles to run (default: 10).
  • -elite: Number of elite candidates to keep in each iteration (default: 10).
  • -parents: Number of parents for each child (default: 2).
  • -population: Size of the population (default: 50).
  • -mutation: Chance of mutation (default: 0.05).
  • -output_path: Directory to save the results (default: "evolve_output").
  • -eval_file: Text file containing prompts for evaluation (default: "evals.txt").
  • -eval_samples: Number of samples to evaluate between candidates (default: 3).
  • -vlm: The VLM to use, claude(default) or llava
  • -append_prompt: Adds to the end of eval prompts
  • -negative_prompt: Negative prompt in diffusion sampling
  • -guidance_scale: Guidance scale for diffusion sampling
  • -diffusion_steps: Number of iterations to diffuse with the candidate during eval
  • -width: Generation width
  • -height: Generation height
  • -resize_width: Width to resize images after generation
  • -resize_height: Height to resize images after generation

Documentation

The framework consists of the following main components:

  • evolve.py: Defines the core evolutionary algorithm, including candidate representation, selection, crossover, mutation, and population management.
  • merge.py: Provides functions for merging SafeTensor files to create new model candidates using DARE.
  • sdxl-evolve.py: The main script that orchestrates the evolutionary process, including image generation, evaluation, and comparison using the VLM.

Details

  • Candidate: Represents a model candidate with its file path, initial population flag, p-value, and lambda value.
  • selection: Selects a subset of candidates as parents for breeding.
  • mutation: Applies random mutations to an offspring candidate.
  • breed: Performs crossover and mutation to create a new offspring candidate.
  • evolve: Evolves the population by selecting parents, breeding offspring, and updating the population.
  • run_evolution: Runs the evolutionary process for a specified number of cycles.
  • load_candidates: Loads initial model candidates from a YAML file.
  • write_yaml: Writes the population to a YAML file.
  • generate_images: Generates images using a Stable Diffusion XL pipeline for evaluation.
  • vlm_judge: Uses the VLM to compare and judge the quality of generated images.

References

Feel free to contribute, report issues, or suggest improvements!