From 238d7fa0eecca6da30d2189f394a8f126ad58f90 Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lstein@gmail.com>
Date: Sun, 17 Sep 2023 16:26:45 -0400
Subject: [PATCH] add models.yaml conversion script

---
 .../backend/model_manager/storage/base.py     |  2 +-
 scripts/convert_models_config_to_3.2.py       | 70 +++++++++++++++++++
 2 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 scripts/convert_models_config_to_3.2.py

diff --git a/invokeai/backend/model_manager/storage/base.py b/invokeai/backend/model_manager/storage/base.py
index f9bd678b0e..1dc5289ecc 100644
--- a/invokeai/backend/model_manager/storage/base.py
+++ b/invokeai/backend/model_manager/storage/base.py
@@ -9,7 +9,7 @@ from typing import List, Optional, Set, Union
 from ..config import BaseModelType, ModelConfigBase, ModelType
 
 # should match the InvokeAI version when this is first released.
-CONFIG_FILE_VERSION = "3.1.1"
+CONFIG_FILE_VERSION = "3.2"
 
 
 class DuplicateModelException(Exception):
diff --git a/scripts/convert_models_config_to_3.2.py b/scripts/convert_models_config_to_3.2.py
new file mode 100644
index 0000000000..fb75631fb3
--- /dev/null
+++ b/scripts/convert_models_config_to_3.2.py
@@ -0,0 +1,70 @@
+# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
+
+"""
+convert_models_config_to_3.2.py.
+
+This script converts a pre-3.2 models.yaml file into the 3.2 format.
+The main difference is that each model is identified by a unique hash,
+rather than the concatenation of base, type and name used previously.
+
+In addition, there are more metadata fields attached to each model.
+These will mostly be empty after conversion, but will be populated
+when new models are downloaded from HuggingFace or Civitae.
+"""
+import argparse
+import sys
+from pathlib import Path
+
+from omegaconf import OmegaConf
+
+from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.backend.model_manager import DuplicateModelException, InvalidModelException, ModelInstall
+from invokeai.backend.model_manager.storage import get_config_store
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert a pre-3.2 models.yaml into the 3.2 version.")
+    parser.add_argument("--root", type=Path, help="Alternate root directory containing the models.yaml to convert")
+    parser.add_argument(
+        "--outfile",
+        type=Path,
+        default=Path("./models-3.2.yaml"),
+        help="File to write to. A file with suffix '.yaml' will use the YAML format. A file with an extension of '.db' will be treated as a SQLite3 database.",
+    )
+    args = parser.parse_args()
+    config_args = ["--root", args.root.as_posix()] if args.root else []
+
+    config = InvokeAIAppConfig.get_config()
+    config.parse_args(config_args)
+    old_yaml_file = OmegaConf.load(config.model_conf_path)
+
+    store = get_config_store(args.outfile)
+    installer = ModelInstall(store=store)
+
+    print(f"Writing 3.2 models configuration into {args.outfile}.")
+
+    for model_key, stanza in old_yaml_file.items():
+        if model_key == "__metadata__":
+            assert (
+                stanza["version"] == "3.0.0"
+            ), f"This script works on version 3.0.0 yaml files, but your configuration points to a {stanza['version']} version"
+            continue
+
+        try:
+            path = config.models_path / stanza["path"]
+            new_key = installer.register_path(path)
+            model_info = store.get_model(new_key)
+            if vae := stanza.get("vae"):
+                model_info.vae = (config.models_path / vae).as_posix()
+            if model_config := stanza.get("config"):
+                model_info.config = (config.root_path / model_config).as_posix()
+            model_info.description = stanza.get("description")
+            store.update_model(new_key, model_info)
+
+            print(f"{model_key} => {new_key}")
+        except (DuplicateModelException, InvalidModelException) as e:
+            print(str(e))
+
+
+if __name__ == "__main__":
+    main()