DOCSEARCHbyAlgolia
Submit Docs

Stability Generative Models (SGM)

repository·main·Indexed 12 days ago

https://github.com/stability-ai/generative-models

A repository by Stability AI containing implementations of various generative models, including diffusion architectures. It features a modular, config-driven system using OmegaConf and PyTorch Lightning, providing engines like AutoencodingEngine and DiffusionEngine. The toolkit includes utilities for training via main.py, a Streamlit inference demo for text-to-image and image-to-image sampling, and scripts for detecting invisible watermarks.

Tokens
3.9K
Snippets
12
Records
17
Agent score
98%

What's inside Stability Generative Models

  1. How the modular configuration system works

    main

    The codebase follows a modular, config-driven philosophy. Instead of extensive subclassing, submodules are built and combined by calling instantiate_from_config() on objects defined in YAML configuration files.

    Key components are separated into distinct configuration blocks:

    • conditioner_config: Defines the GeneralConditioner. It uses emb_models (a list of AbstractEmbModel objects) to handle various conditioning inputs (vectors, sequences, spatial). Each embedder specifies is_trainable, ucg_rate (for classifier-free guidance dropout), and an input_key (e.g., txt or cls).
    • network_config: Defines the neural network backbone (formerly unet_config).
    • loss_config: Configures the loss function and sigma_sampler_config for standard diffusion training.
    • sampler_config: Defines the numerical solver, number of steps, discretization type, and guidance wrappers (like classifier-free guidance). The sampler is independent of the model.

  2. Detect invisible watermarks

    main

    Images generated by this code use the invisible-watermark library. You can detect these watermarks using the provided script scripts/demo/detect.py.

    Minimal Setup for Detection: If you don't want a full installation, you can use this minimal environment:

    python -m venv .detect
source .detect/bin/activate
pip install "numpy>=1.17" "PyWavelets>=1.1.1" "opencv-python>=4.1.0.25"
pip install --no-deps invisible-watermark

    Usage (requires a working sgm installation):

    # Test a single file
python scripts/demo/detect.py <your_filename>

# Test multiple files
python scripts/demo/detect.py <file1> <file2>

# Test all files in a folder
python scripts/demo/detect.py <folder_name>/*
    python scripts/demo/detect.py <your_filename>

  3. Build a distributable wheel with Hatch

    main

    The repository uses PEP 517 compliant packaging with [Hatch]. To build a wheel package:

    1. Install hatch:
      pip install hatch
    2. Build the wheel:
      hatch build -t wheel

    The output will be located in dist/. You can install it via pip install dist/*.whl.

    Note: The package does not specify dependencies; you must install required packages (like PyTorch) manually based on your specific use case.

    pip install hatch
hatch build -t wheel

  4. Run the Streamlit inference demo

    main

    The repository provides a Streamlit-based demo for text-to-image and image-to-image sampling located in scripts/demo/sampling.py.

    Prerequisites:

    1. Download the model weights (e.g., SDXL-base-1.0 or SDXL-refiner-1.0) from Hugging Face.
    2. Place the weights into the checkpoints/ directory.

    Usage: Run the following command to start the server:

    streamlit run scripts/demo/sampling.py --server.port <your_port>

  5. Install generative-models and dependencies

    main

    To use this repository, follow these steps to clone the repo, set up a Python 3.10 virtual environment, and install the required packages. This guide specifically covers the PyTorch 2.0 setup.

    1. Clone the repository
    git clone https://github.com/Stability-AI/generative-models.git
cd generative-models
    1. Set up a virtual environment (PyTorch 2.0) Note: Python 3.10 is recommended to avoid version conflicts.
    python3 -m venv .pt2
source .pt2/bin/activate
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip3 install -r requirements/pt2.txt
    1. Install the sgm package
    pip3 install .
    1. Install sdata for training
    pip3 install -e git+https://github.com/Stability-AI/datapipelines.git@main#egg=sdata
    git clone https://github.com/Stability-AI/generative-models.git
cd generative-models
python3 -m venv .pt2
source .pt2/bin/activate
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip3 install -r requirements/pt2.txt
pip3 install .
pip3 install -e git+https://github.com/Stability-AI/datapipelines.git@main#egg=sdata

  6. Launch model training

    main

    Training is performed using main.py. You can pass one or more configuration files which are merged from left to right (later configs overwrite values from earlier ones).

    Basic Command:

    python main.py --base configs/<config1.yaml> configs/<config2.yaml>

    Example (MNIST):

    python main.py --base configs/example_training/toy/mnist_cond.yaml

    Important Notes:

    • Large Datasets: For configs like imagenet-f8_cond.yaml, you must edit the config to match your dataset path (search for USER: comments in the YAML). Datasets are expected in webdataset format.
    • Latent Models: For latent generative models, you must manually replace the CKPT_PATH placeholder in the config with a valid path to a VAE checkpoint from Hugging Face.
    • PyTorch Compatibility: While the repo supports both pytorch1.13 and pytorch2, autoencoder training (e.g., kl-f4 configs) currently only supports pytorch1.13.

  7. Install the Generative Models environment

    main

    To set up the environment for running Stability AI models, create a Python 3.10 virtual environment and install the required dependencies including PyTorch (with CUDA support), the project requirements, and the sdata package from the Stability AI datapipelines repository.

    python3.10 -m venv .generativemodels
source .generativemodels/bin/activate
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # check CUDA version
pip3 install -r requirements/pt2.txt
pip3 install .
pip3 install -e git+https://github.com/Stability-AI/datapipelines.git@main#egg=sdata
    python3.10 -m venv .generativemodels
source .generativemodels/bin/activate
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # check CUDA version
pip3 install -r requirements/pt2.txt
pip3 install .
pip3 install -e git+https://github.com/Stability-AI/datapipelines.git@main#egg=sdata

  8. How configuration merging works

    main

    The project uses OmegaConf to manage complex hierarchical configurations. When running main.py, the configuration is constructed in this order:

    1. Base Configs: Files provided via -b/--base are loaded and merged from left to right.
    2. CLI Dot-List: Any unknown arguments provided in the format --key.subkey=value are parsed into a dot-list and merged into the configuration.
    3. Lightning Config: A specific lightning key in the config is extracted to configure the PyTorch Lightning Trainer, Logger, Callbacks, and Strategy.
    4. Trainer Overrides: Arguments that match the Trainer.__init__ signature are extracted from the command line and given priority over the configuration-defined trainer settings.

  9. Configure training and sampling via CLI arguments

    main

    The main.py script provides a comprehensive CLI for training and testing generative models. It uses a hierarchical configuration system where base YAML configs are loaded from left-to-right, and can be overridden by command-line arguments.

    Configuration Hierarchy:

    1. Base YAML files (passed via -b/--base).
    2. Command-line arguments (e.g., --key value).
    3. Nested configuration overrides using dot notation (e.g., model.params.key=value).
    4. PyTorch Lightning Trainer arguments (passed directly as --arg value).

    Key CLI Flags:

    • -b, --base: Paths to base .yaml config files. Multiple files are merged left-to-right.
    • -r, --resume: Path to a log directory or a specific checkpoint file to resume training.
    • -n, --name: A postfix for the log directory name.
    • -l, --logdir: The base directory for logging (defaults to logs).
    • -p, --project: Name of the WandB project or path to an existing project.
    • -s, --seed: Random seed for reproducibility.
    • --wandb: Enables logging to Weights & Biases.
    • --scale_lr: Scales the base learning rate by ngpu * batch_size * n_accumulate.
    python main.py -b configs/base_config.yaml -b configs/extra_params.yaml --model.params.new_param=10 --trainer.devices=2 --wandb

  10. Instantiate models using AutoencodingEngine and DiffusionEngine

    main

    The sgm package provides two primary engine classes for working with generative models: AutoencodingEngine and DiffusionEngine. These are imported from the .models module and serve as the main entry points for model-based tasks.

    from sgm import AutoencodingEngine, DiffusionEngine

  11. Use LambdaLinearScheduler for multi-cycle linear decay scheduling

    main

    The LambdaLinearScheduler inherits from LambdaWarmUpCosineScheduler2 but replaces the cosine decay phase with a linear decay phase. It supports multiple cycles configured via lists.

    Important: This scheduler is designed to be used with a base_lr of 1.0.

    Parameters:

    • warm_up_steps: List of warm-up steps per cycle.
    • f_min: List of minimum multiplier values per cycle.
    • f_max: List of maximum multiplier values per cycle.
    • f_start: List of starting multiplier values per cycle.
    • cycle_lengths: List of total steps per cycle.
    • verbosity_interval: If greater than 0, prints status updates every n steps.
    from sgm.lr_scheduler import LambdaLinearScheduler

scheduler = LambdaLinearScheduler(
    warm_up_steps=[500],
    f_min=[0.1],
    f_max=[1.0],
    f_start=[0.0],
    cycle_lengths=[5000],
    verbosity_interval=100
)

# Get multiplier for a step in the linear decay phase
multiplier = scheduler(2500)

  12. Load configurations and instantiate objects with instantiate_from_config

    main

    To build models or components from configuration files, use instantiate_from_config. This utility is typically used in conjunction with get_configs_path to locate the necessary YAML configuration files within the package.

    from sgm import instantiate_from_config, get_configs_path

# Example usage pattern:
# config_path = get_configs_path("path/to/config.yaml")
# model = instantiate_from_config(config_path)