add model merging documentation and launcher script menu entries

docs/features/MODEL_MERGING.md

@@ -0,0 +1,77 @@
+---
+title: Model Merging
+---
+
+# :material-image-off: Model Merging
+
+## How to Merge Models
+
+As of version 2.3, InvokeAI comes with a script that allows you to
+merge two or three diffusers-type models into a new merged model. The
+resulting model will combine characteristics of the original, and can
+be used to teach an old model new tricks.
+
+You may run the merge script by starting the invoke launcher
+(`invoke.sh` or `invoke.bat`) and choosing the option for _merge
+models_. This will launch a text-based interactive user interface that
+prompts you to select the models to merge, how to merge them, and the
+merged model name.
+
+Alternatively you may activate InvokeAI's virtual environment from the
+command line, and call the script via `merge_models_fe.py` (the "fe"
+stands for "front end"). There is also a version that accepts
+command-line arguments, which you can run with the command
+`merge_models.py`.
+
+The user interface for the text-based interactive script is
+straightforward. It shows you a series of setting fields. Use control-N (^N)
+to move to the next field, and control-P (^P) to move to the previous
+one. You can also use TAB and shift-TAB to move forward and
+backward. Once you are in a multiple choice field, use the up and down
+cursor arrows to move to your desired selection, and press <SPACE> or
+<ENTER> to select it. Change text fields by typing in them, and adjust
+scrollbars using the left and right arrow keys.
+
+Once you are happy with your settings, press the OK button. Note that
+there may be two pages of settings, depending on the height of your
+screen, and the OK button may be on the second page. Advance past the
+last field of the first page to get to the second page, and reverse
+this to get back.
+
+If the merge runs successfully, it will create a new diffusers model
+under the selected name and register it with InvokeAI.
+
+## The Settings
+
+* Model Selection -- there are three multiple choice fields that
+  display all the diffusers-style models that InvokeAI knows about.
+  If you do not see the model you are looking for, then it is probably
+  a legacy checkpoint model and needs to be converted using the
+  `invoke.py` command-line client and its `!optimize` command. You
+  must select at least two models to merge. The third can be left at
+  "None" if you desire.
+
+* Alpha -- This is the ratio to use when combining models. It ranges
+  from 0 to 1. The higher the value, the more weight is given to the
+  2d and (optionally) 3d models. So if you have two models named "A"
+  and "B", an alpha value of 0.25 will give you a merged model that is
+  25% A and 75% B.
+
+* Interpolation Method -- This is the method used to combine
+  weights. The options are "weighted_sum" (the default), "sigmoid",
+  "inv_sigmoid" and "add_difference". Each produces slightly different
+  results. When three models are in use, only "add_difference" is
+  available. (TODO: cite a reference that describes what these
+  interpolation methods actually do and how to decide among them).
+
+* Force -- Not all models are compatible with each other. The merge
+  script will check for compatibility and refuse to merge ones that
+  are incompatible. Set this checkbox to try merging anyway.
+
+* Name for merged model - This is the name for the new model. Please
+  use InvokeAI conventions - only alphanumeric letters and the
+  characters ".+-".
+
+## Caveats
+
+This is a new script and may contain bugs.

docs/index.md

@@ -157,6 +157,8 @@ images in full-precision mode:
 <!-- seperator -->
 - [Prompt Engineering](features/PROMPTS.md)
 <!-- seperator -->
+- [Model Merging](features/MODEL_MERGING.md)
+<!-- seperator -->
 - Miscellaneous
   - [NSFW Checker](features/NSFW.md)
   - [Embiggen upscaling](features/EMBIGGEN.md)

installer/templates/invoke.bat.in

@@ -10,8 +10,9 @@ echo Do you want to generate images using the
 echo 1. command-line
 echo 2. browser-based UI
 echo 3. run textual inversion training
-echo 4. open the developer console
-echo 5. re-run the configure script to download new models
+echo 4. merge models (diffusers type only)
+echo 5. open the developer console
+echo 6. re-run the configure script to download new models
 set /P restore="Please enter 1, 2, 3, 4 or 5: [5] "
 if not defined restore set restore=2
 IF /I "%restore%" == "1" (
@@ -24,6 +25,9 @@ IF /I "%restore%" == "1" (
     echo Starting textual inversion training..
     python .venv\Scripts\textual_inversion_fe.py --web %*
 ) ELSE IF /I "%restore%" == "4" (
+    echo Starting model merging script..
+    python .venv\Scripts\merge_models_fe.py --web %*
+) ELSE IF /I "%restore%" == "5" (
     echo Developer Console
     echo Python command is:
     where python
@@ -35,7 +39,7 @@ IF /I "%restore%" == "1" (
     echo *************************
     echo *** Type `exit` to quit this shell and deactivate the Python virtual environment ***
     call cmd /k
-) ELSE IF /I "%restore%" == "5" (
+) ELSE IF /I "%restore%" == "6" (
     echo Running configure_invokeai.py...
     python .venv\Scripts\configure_invokeai.py --web %*
 ) ELSE (

installer/templates/invoke.sh.in

@@ -20,16 +20,18 @@ if [ "$0" != "bash" ]; then
     echo "1. command-line"
     echo "2. browser-based UI"
     echo "3. run textual inversion training"
-    echo "4. open the developer console"
+    echo "4. merge models (diffusers type only)"
     echo "5. re-run the configure script to download new models"
+    echo "6. open the developer console"
     read -p "Please enter 1, 2, 3, 4 or 5: [1] " yn
     choice=${yn:='2'}
     case $choice in
         1 ) printf "\nStarting the InvokeAI command-line..\n"; .venv/bin/python .venv/bin/invoke.py  $*;;
         2 ) printf "\nStarting the InvokeAI browser-based UI..\n"; .venv/bin/python .venv/bin/invoke.py --web $*;;
         3 ) printf "\nStarting Textual Inversion:\n"; .venv/bin/python .venv/bin/textual_inversion_fe.py $*;;
-	4 ) printf "\nDeveloper Console:\n"; file_name=$(basename "${BASH_SOURCE[0]}"); bash --init-file "$file_name";;
-	5 ) printf "\nRunning configure_invokeai.py:\n"; .venv/bin/python .venv/bin/configure_invokeai.py $*;;
+	4 ) printf "\nMerging Models:\n"; .venv/bin/python .venv/bin/merge_models_fe.py $*;;
+	5 ) printf "\nDeveloper Console:\n"; file_name=$(basename "${BASH_SOURCE[0]}"); bash --init-file "$file_name";;
+	6 ) printf "\nRunning configure_invokeai.py:\n"; .venv/bin/python .venv/bin/configure_invokeai.py $*;;
         * ) echo "Invalid selection"; exit;;
     esac
 else # in developer console

ldm/invoke/merge_diffusers.py

@@ -29,6 +29,9 @@ def merge_diffusion_models(models:List['str'],
     '''
     config_file = global_config_file()
     model_manager = ModelManager(OmegaConf.load(config_file))
+    for mod in models:
+        assert (mod in model_manager.model_names()), f'** Unknown model "{mod}"'
+        assert (model_manager.model_info(mod).get('format',None) == 'diffusers'), f'** {mod} is not a diffusers model. It must be optimized before merging.'
     model_ids_or_paths = [model_manager.model_name_or_path(x) for x in models]
 
     pipe = DiffusionPipeline.from_pretrained(model_ids_or_paths[0],

scripts/merge_models.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python
+
+import argparse
+import os
+import sys
+import traceback
+from pathlib import Path
+
+from omegaconf import OmegaConf
+
+from ldm.invoke.globals import (Globals, global_cache_dir, global_config_file,
+                                global_set_root)
+from ldm.invoke.model_manager import ModelManager
+
+parser = argparse.ArgumentParser(description="InvokeAI textual inversion training")
+parser.add_argument(
+    "--root_dir",
+    "--root-dir",
+    type=Path,
+    default=Globals.root,
+    help="Path to the invokeai runtime directory",
+)
+parser.add_argument(
+    "--models",
+    required=True,
+    type=str,
+    nargs="+",
+    help="Two to three model names to be merged",
+)
+parser.add_argument(
+    "--merged_model_name",
+    "--destination",
+    dest="merged_model_name",
+    type=str,
+    help="Name of the output model. If not specified, will be the concatenation of the input model names.",
+)
+parser.add_argument(
+    "--alpha",
+    type=float,
+    default=0.5,
+    help="The interpolation parameter, ranging from 0 to 1. It affects the ratio in which the checkpoints are merged. Higher values give more weight to the 2d and 3d models",
+)
+parser.add_argument(
+    "--interpolation",
+    dest="interp",
+    type=str,
+    choices=["weighted_sum", "sigmoid", "inv_sigmoid", "add_difference"],
+    default="weighted_sum",
+    help='Interpolation method to use. If three models are present, only "add_difference" will work.',
+)
+parser.add_argument(
+    "--force",
+    action="store_true",
+    help="Try to merge models even if they are incompatible with each other",
+)
+parser.add_argument(
+    "--clobber",
+    "--overwrite",
+    dest='clobber',
+    action="store_true",
+    help="Overwrite the merged model if --merged_model_name already exists",
+)
+
+args = parser.parse_args()
+global_set_root(args.root_dir)
+
+assert args.alpha >= 0 and args.alpha <= 1.0, "alpha must be between 0 and 1"
+assert len(args.models) >= 1 and len(args.models) <= 3, "provide 2 or 3 models to merge"
+
+if not args.merged_model_name:
+    args.merged_model_name = "+".join(args.models)
+    print(
+        f'>> No --merged_model_name provided. Defaulting to "{args.merged_model_name}"'
+    )
+
+model_manager = ModelManager(OmegaConf.load(global_config_file()))
+assert (args.clobber or args.merged_model_name not in model_manager.model_names()), f'A model named "{args.merged_model_name}" already exists. Use --clobber to overwrite.'
+
+# It seems that the merge pipeline is not honoring cache_dir, so we set the
+# HF_HOME environment variable here *before* we load diffusers.
+cache_dir = str(global_cache_dir("diffusers"))
+os.environ["HF_HOME"] = cache_dir
+from ldm.invoke.merge_diffusers import merge_diffusion_models
+
+try:
+    merge_diffusion_models(**vars(args))
+    print(f'>> Models merged into new model: "{args.merged_model_name}".')
+except Exception as e:
+    print(f"** An error occurred while merging the pipelines: {str(e)}")
+    print("** DETAILS:")
+    print(traceback.format_exc())
+    sys.exit(-1)