new config and merge method api (0.1.0rc) (#39)

.gitignore

@@ -1,9 +1,12 @@
-.idea
+/.idea
 __pycache__
-sd_mecha.egg-info
-build
-dist
-config
-venv
+/sd_mecha.egg-info
+/build
+/dist
+/config
+/venv
+/model_configs/venvs
+/sd_mecha/extensions/builtin/model_configs/**/*.yaml
+/sd_mecha/extensions/builtin/model_configs/**/*.yaml.bak
 
-examples/*_test.py
+/examples/*_test.py

README.md

@@ -3,56 +3,58 @@
 [![PyPI version](https://badge.fury.io/py/sd-mecha.svg)](https://badge.fury.io/py/sd-mecha)
 [![Discord Server](https://dcbadge.vercel.app/api/server/2EPaw6fxxm?style=flat)](https://discord.gg/invite/2EPaw6fxxm)
 
-sd-mecha is a memory-efficient general-purpose model merger. It can merge any model architecture given appropriate configuration:
-- diffusion models
+```python
+import sd_mecha
+
+# create the merge plan
+recipe = sd_mecha.weighted_sum("/path/to/model_a.safetensors", "/path/to/model_b.safetensors", alpha=0.5)
+
+# initialize merge engine
+merger = sd_mecha.RecipeMerger()
+
+# merge!
+merger.merge_and_save(recipe)
+```
+
+sd-mecha is a general memory-efficient model merging library. It can merge *any* model:
+- Diffusion models
 - LLMs
-- Depth models
-- Scorers
-- ...
+- VLMs
+- Aesthetic scorers
+- etc.
 
 ## Features
 
-- Memory efficient model merging -- merge a very large number of models at the same time
-- Mecha recipes as a textual and interpretable format (.mecha)
-- Extension API through python:
-  - add new architectures (i.e. Stable Cascade, Stable Diffusion 3, etc.)
-  - add new model types (i.e. OFT networks, LoKr, etc.)
-  - add new merge methods
-- Recipe variables for general recipe templates
-- Compose recipe templates to create mega recipes
-- Builtin support for popular model architectures:
-  - SD1.5
-  - SDXL
-  - SD3
-- Merge LoRAs together and to checkpoints
-- Block-wise hyperparameters for precise control of blocks
-- Class-wise hyperparameters for precise control of layer types
-- Support arbitrary model architectures and types using the `sd_mecha.extensions` module
-- Merge SDXL LoRAs to models and with other LoRAs
+- 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
+- Merge LyCORIS networks together and to checkpoints
+- Block-wise hyperparameters for precise control of blocks (aka MBW)
 
 ## Install
 
 ```commandline
-pip install sd-mecha torch
+pip install sd-mecha
 ```
 
-sd-mecha depends additionally on:
-
-- `torch>=2.0.1`
-
-The pypi package does not ship with `torch` so that you can install the appropriate version for your system.
+Make sure to install the appropriate release of [`torch`](https://pytorch.org/get-started/locally/) to get the best performance.
 
 ## Usage
 
 ### Merge models
 
-To merge models, mecha needs a recipe as input. There are multiple ways to provide a recipe:
-- using the python merging API
-- using the CLI with .mecha recipes
+To merge models, mecha uses recipes.
+A recipe is a list of instructions that describes the exact steps needed to obtain the final merged model.
 
-#### Using the python merging API
+#### Using python
 
-Here's an example simple sum-twice merge setup:
+Here's an example script that merges three Stable Diffusion 1.5 models:
 
 ```python
 import sd_mecha
@@ -61,11 +63,11 @@ import sd_mecha
 # all builtin merge methods are direct properties of the `sd_mecha` package for convenience
 recipe = sd_mecha.weighted_sum(
     sd_mecha.weighted_sum(
-        "ghostmix_v20Bakedvae",
-        "deliberate_v2",
+        "ghostmix_v20Bakedvae.safetensors",
+        "deliberate_v2.safetensors",
         alpha=0.5,
     ),
-    "dreamshaper_332BakedVaeClipFix",
+    "dreamshaper_332BakedVaeClipFix.safetensors",
     alpha=0.33,
 )
 
@@ -75,130 +77,42 @@ merger = sd_mecha.RecipeMerger(
 )
 
 # perform the entire merge plan and save to output path
-merger.merge_and_save(recipe, output="basic_merge")
+merger.merge_and_save(recipe, output="basic_merge.safetensors")
 ```
 
 See the [examples](/examples) directory for more examples.
 
-#### Using the CLI with .mecha recipes
-
-It is alternatively possible to merge recipes previously serialized to `.mecha`.
-This is only possible if the recipe is concrete. (i.e. all potential parameters have been replaced with actual models)
-
-```commandline
-python -m sd_mecha merge path/to/recipe.mecha
-```
-
-For more information:
-
-```commandline
-python -m sd_mecha merge --help
-```
-
 ### Get Model-Specific Information
 
-The interface for block/class hyperparameters requires prior knowledge of the blocks and classes of the architecture being merged.
-The command `info` was made to discover the names of the blocks and/or classes to use.
-
-To show the registered model architectures:
-
-```commandline
-python -m sd_mecha info
-```
-
-Mecha has builtin support for the SD1.x and the SDXL architectures:
-
-```
-Available architectures:
-- sd1
-- sdxl
-```
-
-To view the available blocks and classes of an architecture, specify the architecture:
+To specify block weights, we need to know the name of the blocks.
 
-```commandline
-python -m sd_mecha info sd1
-```
-```
-Component "txt":
-  Blocks:
-  - in0
-  - in1
-  - in2
-  ...
-  Classes:
-  - final_layer_norm
-  - layer_norm1
-  - layer_norm2
-  - mlp_fc1
-  ...
-Component "unet":
-  Blocks:
-  ...
-  Classes:
-  ...
-```
+This information can be discovered using the `extensions.model_configs` submodule.
 
-Given this information, it is possible to set i.e. the value of block `in2` in the `txt` component specifically:
+Mecha has builtin support for Stable Diffusion 1.X, Stable Diffusion XL and Stable Diffusion 3:
 
 ```python
-import sd_mecha
-recipe = sd_mecha.weighted_sum(
-    "ghostmix_v20Bakedvae",
-    "dreamshaper_332BakedVaeClipFix",
-    alpha=(
-      sd_mecha.default("sd1", "txt", 0.33) |
-      sd_mecha.blocks("sd1", "txt", in2=0.75)
-    ),
-)
-```
+from sd_mecha.extensions.model_configs import get_all
 
-See the [merging API section](#using-the-python-merging-api) above for more info.
+all_configs = get_all()
 
-If run as verbose, it also shows the keys that are associated with each block/class:
-
-```commandline
-python -m sd_mecha info sd1 -v
-```
-```
-Component "txt":
-  Blocks:
-    in0:
-    - model.diffusion_model.input_blocks.0.0.bias
-    - model.diffusion_model.input_blocks.0.0.weight
-    in1:
-    - model.diffusion_model.input_blocks.1.0.emb_layers.1.bias
-    - model.diffusion_model.input_blocks.1.0.emb_layers.1.weight
-    - model.diffusion_model.input_blocks.1.0.in_layers.0.bias
-    - model.diffusion_model.input_blocks.1.0.in_layers.0.weight
-    ...
-  ...
-...
-```
-
-### Compose recipes
-
-It is possible to compose recipes together to create more complex recipes.
-For this to work, the base recipe must be general: (i.e. the parameters to replace must exist in the base recipe)
-
-```commandline
-python -m sd_mecha compose path/to/recipe.mecha [options]
+print([config.identifier for config in all_configs])
+# ["sd1-ldm-base", "sdxl-sgm-base", "sd3-sgm-base", ...]
 ```
 
-For example, here we compose the recipe [incompatible_fusion.mecha](examples/recipes/incompatible_fusion.mecha)
-with another recipe for parameter "a" and
-SD1.5 base for parameter "c":
+To view the available components of a model:
 
-```commandline
-python -m sd_mecha compose examples/recipes/incompatible_fusion.mecha \
-  -p a examples/recipes/weighted_sum.mecha \
-  -p c v1-5-pruned.safetensors
-```
+```python
+from sd_mecha.extensions import model_configs
 
-For more information:
+config = model_configs.resolve("sd1-ldm")
+for component_id, component in config.components.items():
+      # block_keys contains the state dict keys that the block controls
+      print(f"{component_id}")
 
-```shell
-python -m sd_mecha compose --help
+# this prints:
+#   clip_l
+#   vae
+#   diffusers
 ```
 
 ## Motivation

examples/add_perpendicular.py

@@ -3,13 +3,13 @@
 
 
 recipe = sd_mecha.add_perpendicular(
-    "ghostmix_v20Bakedvae",
-    "dreamshaper_332BakedVaeClipFix",
-    "pure/v1-5-pruned",
+    sd_mecha.model("ghostmix_v20Bakedvae.safetensors"),
+    sd_mecha.model("dreamshaper_332BakedVaeClipFix.safetensors"),
+    sd_mecha.model("pure/v1-5-pruned.safetensors"),
 )
 
 merger = sd_mecha.RecipeMerger(
     models_dir=r"E:\sd\models\Stable-diffusion",
 )
 
-merger.merge_and_save(recipe, output="basic_merge")
+merger.merge_and_save(recipe)

examples/basic_merge.py

@@ -3,10 +3,12 @@
 
 
 # plan a simple weighted sum
-recipe = sd_mecha.weighted_sum("ghostmix_v20Bakedvae", "dreamshaper_332BakedVaeClipFix")
+a = sd_mecha.model("ghostmix_v20Bakedvae.safetensors")
+b = sd_mecha.model("dreamshaper_332BakedVaeClipFix.safetensors")
+recipe = sd_mecha.weighted_sum(a, b)
 
 # merger provides global defaults for methods
 merger = sd_mecha.RecipeMerger(models_dir=r"E:\sd\models\Stable-diffusion")
 
 # perform the entire merge plan and save to output path
-merger.merge_and_save(recipe, output="basic_merge")
+merger.merge_and_save(recipe)

examples/binomial_dropout_merge.py

@@ -1,6 +1,11 @@
 import sd_mecha
 sd_mecha.set_log_level()
 
-recipe = sd_mecha.dropout("ghostmix_v20Bakedvae", "dreamshaper_332BakedVaeClipFix", probability=0.9, alpha=0.5, seed=0)
+recipe = sd_mecha.dropout(
+    sd_mecha.model("pure/v1-5-pruned.safetensors"),
+    sd_mecha.model("ghostmix_v20Bakedvae.safetensors"),
+    sd_mecha.model("dreamshaper_332BakedVaeClipFix.safetensors"),
+    probability=0.9, alpha=0.5, seed=0,
+)
 merger = sd_mecha.RecipeMerger(models_dir=r"E:\sd\models\Stable-diffusion")
-merger.merge_and_save(recipe, output="basic_merge")
+merger.merge_and_save(recipe)