From d132481058f8a827cd407f2121f128a2bb862f7a Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Sun, 2 Apr 2023 17:41:55 -0500 Subject: Embed model merge metadata in .safetensors file --- modules/extras.py | 44 ++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 42 insertions(+), 2 deletions(-) (limited to 'modules/extras.py') diff --git a/modules/extras.py b/modules/extras.py index d8ece955..77d88592 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -1,6 +1,7 @@ import os import re import shutil +import json import torch @@ -71,7 +72,7 @@ def to_half(tensor, enable): return tensor -def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights): +def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_model_name, interp_method, multiplier, save_as_half, custom_name, checkpoint_format, config_source, bake_in_vae, discard_weights, save_metadata): shared.state.begin() shared.state.job = 'model-merge' @@ -241,13 +242,52 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ shared.state.textinfo = "Saving" print(f"Saving to {output_modelname}...") + metadata = {"format": "pt", "models": {}, "merge_recipe": None} + + if save_metadata: + merge_recipe = { + "primary_model_hash": primary_model_info.sha256, + "secondary_model_hash": secondary_model_info.sha256 if secondary_model_info else None, + "tertiary_model_hash": tertiary_model_info.sha256 if tertiary_model_info else None, + "interp_method": interp_method, + "multiplier": multiplier, + "save_as_half": save_as_half, + "custom_name": custom_name, + "config_source": config_source, + "bake_in_vae": bake_in_vae, + "discard_weights": discard_weights, + "is_inpainting": result_is_inpainting_model, + "is_instruct_pix2pix": result_is_instruct_pix2pix_model + } + metadata["merge_recipe"] = json.dumps(merge_recipe) + + def add_model_metadata(checkpoint_info): + metadata["models"][checkpoint_info.sha256] = { + "name": checkpoint_info.name, + "legacy_hash": checkpoint_info.hash, + "merge_recipe": checkpoint_info.metadata.get("merge_recipe", None) + } + + metadata["models"].update(checkpoint_info.metadata.get("models", {})) + + add_model_metadata(primary_model_info) + if secondary_model_info: + add_model_metadata(secondary_model_info) + if tertiary_model_info: + add_model_metadata(tertiary_model_info) + + metadata["models"] = json.dumps(metadata["models"]) + _, extension = os.path.splitext(output_modelname) if extension.lower() == ".safetensors": - safetensors.torch.save_file(theta_0, output_modelname, metadata={"format": "pt"}) + safetensors.torch.save_file(theta_0, output_modelname, metadata=metadata) else: torch.save(theta_0, output_modelname) sd_models.list_models() + created_model = next((ckpt for ckpt in sd_models.checkpoints_list.values() if ckpt.name == filename), None) + if created_model: + created_model.calculate_shorthash() create_config(output_modelname, config_source, primary_model_info, secondary_model_info, tertiary_model_info) -- cgit v1.2.1 From afc349c2c0d7c7543e8cc085cde2beef8549fffc Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Sun, 2 Apr 2023 18:40:33 -0500 Subject: Add field for model merge type Incase this is supported by other merge extensions --- modules/extras.py | 1 + 1 file changed, 1 insertion(+) (limited to 'modules/extras.py') diff --git a/modules/extras.py b/modules/extras.py index 77d88592..9a00c9a3 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -246,6 +246,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ if save_metadata: merge_recipe = { + "type": "webui", # indicate this model was merged with webui's built-in merger "primary_model_hash": primary_model_info.sha256, "secondary_model_hash": secondary_model_info.sha256 if secondary_model_info else None, "tertiary_model_hash": tertiary_model_info.sha256 if tertiary_model_info else None, -- cgit v1.2.1 From 7c016dd642cc29e064715ac04ab3d83c5451b45e Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Sun, 2 Apr 2023 19:06:39 -0500 Subject: Calculate shorthash on merge if not exist --- modules/extras.py | 1 + 1 file changed, 1 insertion(+) (limited to 'modules/extras.py') diff --git a/modules/extras.py b/modules/extras.py index 9a00c9a3..97d14e5a 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -263,6 +263,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ metadata["merge_recipe"] = json.dumps(merge_recipe) def add_model_metadata(checkpoint_info): + checkpoint_info.calculate_shorthash() metadata["models"][checkpoint_info.sha256] = { "name": checkpoint_info.name, "legacy_hash": checkpoint_info.hash, -- cgit v1.2.1 From fbaf6e4fd897fa1f3e3f747f1d699c240cad76a0 Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Sun, 2 Apr 2023 21:41:23 -0500 Subject: Namespace metadata fields --- modules/extras.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) (limited to 'modules/extras.py') diff --git a/modules/extras.py b/modules/extras.py index 97d14e5a..ff4e9c4e 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -242,7 +242,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ shared.state.textinfo = "Saving" print(f"Saving to {output_modelname}...") - metadata = {"format": "pt", "models": {}, "merge_recipe": None} + metadata = {"format": "pt", "sd_merge_models": {}, "sd_merge_recipe": None} if save_metadata: merge_recipe = { @@ -260,17 +260,17 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ "is_inpainting": result_is_inpainting_model, "is_instruct_pix2pix": result_is_instruct_pix2pix_model } - metadata["merge_recipe"] = json.dumps(merge_recipe) + metadata["sd_merge_recipe"] = json.dumps(merge_recipe) def add_model_metadata(checkpoint_info): checkpoint_info.calculate_shorthash() - metadata["models"][checkpoint_info.sha256] = { + metadata["sd_merge_models"][checkpoint_info.sha256] = { "name": checkpoint_info.name, "legacy_hash": checkpoint_info.hash, - "merge_recipe": checkpoint_info.metadata.get("merge_recipe", None) + "sd_merge_recipe": checkpoint_info.metadata.get("sd_merge_recipe", None) } - metadata["models"].update(checkpoint_info.metadata.get("models", {})) + metadata["sd_merge_models"].update(checkpoint_info.metadata.get("sd_merge_models", {})) add_model_metadata(primary_model_info) if secondary_model_info: @@ -278,7 +278,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ if tertiary_model_info: add_model_metadata(tertiary_model_info) - metadata["models"] = json.dumps(metadata["models"]) + metadata["sd_merge_models"] = json.dumps(metadata["sd_merge_models"]) _, extension = os.path.splitext(output_modelname) if extension.lower() == ".safetensors": -- cgit v1.2.1 From 762265eab58cdb8f2d6398769bab43d8b8db0075 Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Wed, 10 May 2023 07:52:45 +0300 Subject: autofixes from ruff --- modules/extras.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) (limited to 'modules/extras.py') diff --git a/modules/extras.py b/modules/extras.py index ff4e9c4e..eb4f0b42 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -136,14 +136,14 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ result_is_instruct_pix2pix_model = False if theta_func2: - shared.state.textinfo = f"Loading B" + shared.state.textinfo = "Loading B" print(f"Loading {secondary_model_info.filename}...") theta_1 = sd_models.read_state_dict(secondary_model_info.filename, map_location='cpu') else: theta_1 = None if theta_func1: - shared.state.textinfo = f"Loading C" + shared.state.textinfo = "Loading C" print(f"Loading {tertiary_model_info.filename}...") theta_2 = sd_models.read_state_dict(tertiary_model_info.filename, map_location='cpu') -- cgit v1.2.1 From 49a55b410b66b7dd9be9335d8a2e3a71e4f8b15c Mon Sep 17 00:00:00 2001 From: Aarni Koskela Date: Thu, 11 May 2023 18:28:15 +0300 Subject: Autofix Ruff W (not W605) (mostly whitespace) --- modules/extras.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'modules/extras.py') diff --git a/modules/extras.py b/modules/extras.py index eb4f0b42..bdf9b3b7 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -199,7 +199,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ result_is_inpainting_model = True else: theta_0[key] = theta_func2(a, b, multiplier) - + theta_0[key] = to_half(theta_0[key], save_as_half) shared.state.sampling_step += 1 -- cgit v1.2.1 From 95cb492e4106646450480ec74014c3ec9a679c1f Mon Sep 17 00:00:00 2001 From: Weiming Date: Wed, 17 May 2023 21:25:50 +0800 Subject: Fixed: #10460 --- modules/extras.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) (limited to 'modules/extras.py') diff --git a/modules/extras.py b/modules/extras.py index bdf9b3b7..aabb3b26 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -242,9 +242,10 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ shared.state.textinfo = "Saving" print(f"Saving to {output_modelname}...") - metadata = {"format": "pt", "sd_merge_models": {}, "sd_merge_recipe": None} + metadata = None if save_metadata: + metadata = {"format": "pt", "sd_merge_models": {}} merge_recipe = { "type": "webui", # indicate this model was merged with webui's built-in merger "primary_model_hash": primary_model_info.sha256, -- cgit v1.2.1 From 76ebf750a463bc24434048dc60e289b0b6198598 Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Wed, 17 May 2023 17:44:07 +0300 Subject: use a local variable instead of dictionary entry for sd_merge_models in merge model metadata code --- modules/extras.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) (limited to 'modules/extras.py') diff --git a/modules/extras.py b/modules/extras.py index aabb3b26..830b53aa 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -245,7 +245,8 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ metadata = None if save_metadata: - metadata = {"format": "pt", "sd_merge_models": {}} + metadata = {"format": "pt"} + merge_recipe = { "type": "webui", # indicate this model was merged with webui's built-in merger "primary_model_hash": primary_model_info.sha256, @@ -263,15 +264,17 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ } metadata["sd_merge_recipe"] = json.dumps(merge_recipe) + sd_merge_models = {} + def add_model_metadata(checkpoint_info): checkpoint_info.calculate_shorthash() - metadata["sd_merge_models"][checkpoint_info.sha256] = { + sd_merge_models[checkpoint_info.sha256] = { "name": checkpoint_info.name, "legacy_hash": checkpoint_info.hash, "sd_merge_recipe": checkpoint_info.metadata.get("sd_merge_recipe", None) } - metadata["sd_merge_models"].update(checkpoint_info.metadata.get("sd_merge_models", {})) + sd_merge_models.update(checkpoint_info.metadata.get("sd_merge_models", {})) add_model_metadata(primary_model_info) if secondary_model_info: @@ -279,7 +282,7 @@ def run_modelmerger(id_task, primary_model_name, secondary_model_name, tertiary_ if tertiary_model_info: add_model_metadata(tertiary_model_info) - metadata["sd_merge_models"] = json.dumps(metadata["sd_merge_models"]) + metadata["sd_merge_models"] = json.dumps(sd_merge_models) _, extension = os.path.splitext(output_modelname) if extension.lower() == ".safetensors": -- cgit v1.2.1