From 4ee8d104f07ea519959d18075dd8ebb597cdb44c Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sat, 21 Jan 2023 18:39:13 -0500 Subject: [PATCH] working, but needs diffusers PR to be accepted --- scripts/merge_fe.py | 111 +++++++++++++++++++++++++++++++++++++++----- 1 file changed, 100 insertions(+), 11 deletions(-) mode change 100644 => 100755 scripts/merge_fe.py diff --git a/scripts/merge_fe.py b/scripts/merge_fe.py old mode 100644 new mode 100755 index df01235599..4f24961a18 --- a/scripts/merge_fe.py +++ b/scripts/merge_fe.py @@ -3,15 +3,17 @@ import npyscreen import os import sys -import re -import shutil import traceback import argparse -from ldm.invoke.globals import Globals, global_set_root +import safetensors.torch +from ldm.invoke.globals import Globals, global_set_root, global_cache_dir +from ldm.invoke.model_manager import ModelManager from omegaconf import OmegaConf from pathlib import Path from typing import List +CONFIG_FILE = None + class FloatSlider(npyscreen.Slider): # this is supposed to adjust display precision, but doesn't def translate_value(self): @@ -30,6 +32,14 @@ class mergeModelsForm(npyscreen.FormMultiPageAction): 'inv_sigmoid', 'add_difference'] + def __init__(self, parentApp, name): + self.parentApp = parentApp + super().__init__(parentApp, name) + + @property + def model_manager(self): + return self.parentApp.model_manager + def afterEditing(self): self.parentApp.setNextForm(None) @@ -83,6 +93,11 @@ class mergeModelsForm(npyscreen.FormMultiPageAction): lowest=0, value=0.5, ) + self.force = self.add_widget_intelligent( + npyscreen.Checkbox, + name='Force merge of incompatible models', + value=False, + ) self.merged_model_name = self.add_widget_intelligent( npyscreen.TitleText, name='Name for merged model', @@ -108,17 +123,44 @@ class mergeModelsForm(npyscreen.FormMultiPageAction): if self.validate_field_values(): self.parentApp.setNextForm(None) self.editing = False + self.parentApp.merge_arguments = self.marshall_arguments() + npyscreen.notify('Starting the merge...') + import diffusers # this keeps the message up while diffusers loads else: self.editing = True def ok_cancel(self): sys.exit(0) + def marshall_arguments(self)->dict: + model_names = self.model_names + models = [ + model_names[self.model1.value[0]], + model_names[self.model2.value[0]], + ] + if self.model3.value[0] > 0: + models.append(model_names[self.model3.value[0]-1]) + + models = [self.model_manager.model_name_or_path(x) for x in models] + + args = dict( + pretrained_model_name_or_path_list=models, + alpha = self.alpha.value, + interp = self.interpolations[self.merge_method.value[0]], + force = self.force.value, + cache_dir = global_cache_dir('diffusers'), + merged_model_name = self.merged_model_name.value, + ) + return args + def validate_field_values(self)->bool: bad_fields = [] - selected_models = set((self.model1.value[0],self.model2.value[0],self.model3.value[0])) - if len(selected_models) < 3: - bad_fields.append('Please select two or three DIFFERENT models to compare') + model_names = self.model_names + selected_models = set((model_names[self.model1.value[0]],model_names[self.model2.value[0]])) + if self.model3.value[0] > 0: + selected_models.add(model_names[self.model3.value[0]-1]) + if len(selected_models) < 2: + bad_fields.append(f'Please select two or three DIFFERENT models to compare. You selected {selected_models}') if len(bad_fields) > 0: message = 'The following problems were detected and must be corrected:' for problem in bad_fields: @@ -129,13 +171,15 @@ class mergeModelsForm(npyscreen.FormMultiPageAction): return True def get_model_names(self)->List[str]: - conf = OmegaConf.load(os.path.join(Globals.root,'configs/models.yaml')) - model_names = [name for name in conf.keys() if conf[name].get('format',None)=='diffusers'] + model_names = [name for name in self.model_manager.model_names() if self.model_manager.model_info(name).get('format') == 'diffusers'] + print(model_names) return sorted(model_names) -class MyApplication(npyscreen.NPSAppManaged): +class Mergeapp(npyscreen.NPSAppManaged): def __init__(self): super().__init__() + conf = OmegaConf.load(Path(Globals.root) / 'configs' / 'models.yaml') + self.model_manager = ModelManager(conf,'cpu','float16') # precision doesn't really matter here def onStart(self): npyscreen.setTheme(npyscreen.Themes.DefaultTheme) @@ -151,6 +195,51 @@ if __name__ == '__main__': ) args = parser.parse_args() global_set_root(args.root_dir) + + CONFIG_FILE = os.path.join(Globals.root,'configs/models.yaml') + os.environ['HF_HOME'] = str(global_cache_dir('diffusers')) + + mergeapp = Mergeapp() + mergeapp.run() + from diffusers import DiffusionPipeline + args = mergeapp.merge_arguments + merged_model_name = args['merged_model_name'] + merged_pipe = None + print(args) + + try: + print(f'DEBUG: {args["pretrained_model_name_or_path_list"][0]}') + pipe = DiffusionPipeline.from_pretrained(args['pretrained_model_name_or_path_list'][0], + custom_pipeline='checkpoint_merger' + ) + merged_pipe = pipe.merge(**args) + dump_path = Path(Globals.root) / 'models' / 'merged_diffusers' + os.makedirs(dump_path,exist_ok=True) + dump_path = dump_path / merged_model_name + merged_pipe.save_pretrained ( + dump_path, + safe_serialization=1 + ) + except Exception as e: + print(f'** An error occurred while merging the pipelines: {str(e)}') + print('** DETAILS:') + print(traceback.format_exc()) + sys.exit(-1) + + print(f'>> Merged model is saved to {dump_path}') + response = input('Import this model into InvokeAI? [y]').strip() or 'y' + if response.startswith(('y','Y')): + try: + mergeapp.model_manager.import_diffuser_model( + dump_path, + model_name = merged_model_name, + description = f'Merge of models {args["pretrained_model_name_or_path_list"]}' + ) + mergeapp.model_manager.commit(CONFIG_FILE) + print('>> Merged model imported.') + except Exception as e: + print(f'** New model could not be committed to config.yaml: {str(e)}') + print('** DETAILS:') + print(traceback.format_exc()) + - myapplication = MyApplication() - myapplication.run()