Skip to main content

State dict recipe merger

Project description

sd-mecha

PyPI version Discord Server

import sd_mecha

# create the merge plan
a = sd_mecha.model("path/to/model_a.safetensors")
b = sd_mecha.model("path/to/model_b.safetensors")
recipe = sd_mecha.weighted_sum(a, b, alpha=0.5)

# merge!
sd_mecha.merge(recipe, output="path/to/model_out.safetensors")

sd-mecha is a general memory-efficient model merging library. It can merge any model:

  • Diffusion models
  • LLMs
  • VLMs
  • Aesthetic scorers
  • etc.

Features

  • Memory efficient model merging: merge a very large number of models in a single execution
  • Textual and interpretable format for storage and execution (.mecha)
  • Extensible library interface:
    • add custom models
    • add custom merge methods
  • Builtin support for popular diffusion models:
    • Stable Diffusion 1.5
    • Stable Diffusion XL
    • Stable Diffusion 3
    • FLUX Schnell/Dev
  • Merge LyCORIS networks together and to checkpoints
  • Block-wise hyperparameters for precise control of blocks (aka MBW)

Install

pip install sd-mecha

Make sure to install the appropriate release of torch to get the best performance.

Usage

Merge models

To merge models, mecha uses recipes. A recipe is a list of instructions that describes the exact steps needed to derive a state dict from inputs.

Here's an example script that merges three models:

import sd_mecha

# create the merge plan
model_a = sd_mecha.model("path/to/model_a.safetensors")
model_b = sd_mecha.model("path/to/model_b.safetensors")
recipe = sd_mecha.weighted_sum(model_a, model_b, alpha=0.5)

# merge!
sd_mecha.merge(recipe, output="path/to/model_out.safetensors")

See the examples directory for more examples.

Get Model-Specific Information

The library uses a "model config" to designate any specific set of keys of a certain shape.

It is possible to list all available model configs through the sd_mecha.extensions.model_configs module:

from sd_mecha.extensions import model_configs

all_configs = model_configs.get_all()

print([config.identifier for config in all_configs])
# ["sd1-ldm-base", "sdxl-sgm-base", "sd3-sgm-base", ...]

A component of a model config is a subset of keys of the config that belong to the same logical group. For example, all keys starting with "first_stage_model." in Stable Diffusion models belong to the component "vae".

It is possible to query the different components of a model config:

from sd_mecha.extensions import model_configs

config = model_configs.resolve("sd1-ldm")
for component_id, component in config.components().items():
      # component.keys contains the state dict keys that the component owns
      print(f"{component_id}")

# this prints:
#   clip_l
#   vae
#   diffuser

Motivation

Keeping track of full merge recipes has always been a problem for me. I needed something that allows to store merge recipes in a readable format while also being executable. I also needed something that allows to fully merge an entire tree of models without having to save intermediate models to disk.

Typically, mergers load all models in memory before initiating the merge process. This can be very inefficient when the merge focuses on each key individually:

image of typical merge graph

sd-mecha doesn't have this problem as it saves keys as soon as it can:

image of sd-mecha merge graph

This allows to merge a very large number of models simultaneously on low-end hardware.

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.

Source Distribution

sd_mecha-1.1.0rc1.tar.gz (155.3 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

sd_mecha-1.1.0rc1-py3-none-any.whl (176.5 kB view details)

Uploaded Python 3

File details

Details for the file sd_mecha-1.1.0rc1.tar.gz.

File metadata

  • Download URL: sd_mecha-1.1.0rc1.tar.gz
  • Upload date:
  • Size: 155.3 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.10.9

File hashes

Hashes for sd_mecha-1.1.0rc1.tar.gz
Algorithm Hash digest
SHA256 559cbce39bc5ace6cfd04a664a5f8307a7be230a09d3a0f9580f85b6c53eb524
MD5 ceda5c35610ce78355dc3dfd9f3c5ff2
BLAKE2b-256 a32a6210feedbb82528a59498855a38a5a67b9641ebd83d9146251c48bf3ab6c

See more details on using hashes here.

File details

Details for the file sd_mecha-1.1.0rc1-py3-none-any.whl.

File metadata

  • Download URL: sd_mecha-1.1.0rc1-py3-none-any.whl
  • Upload date:
  • Size: 176.5 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.10.9

File hashes

Hashes for sd_mecha-1.1.0rc1-py3-none-any.whl
Algorithm Hash digest
SHA256 41bcaf8b344da2070285a539b233b79607b3b14ecab1c235bf9741a3b1f17c35
MD5 f262d62f2f491333f261fd81e5ff005c
BLAKE2b-256 1b4c546ea43a912fe162a18e1165bea291811bedd9cbe579f0ee9f573a27990f

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page