fix the merge script to correctly display models sorted by base

This commit is contained in:
Lincoln Stein 2023-10-23 20:37:33 -04:00
parent 8e948d3f17
commit c14aa30956

View File

@ -131,6 +131,7 @@ class mergeModelsForm(npyscreen.FormMultiPageAction):
values=[ values=[
"Models Built on SD-1.x", "Models Built on SD-1.x",
"Models Built on SD-2.x", "Models Built on SD-2.x",
"Models Built on SDXL",
], ],
value=[self.current_base], value=[self.current_base],
columns=4, columns=4,
@ -309,7 +310,7 @@ class mergeModelsForm(npyscreen.FormMultiPageAction):
else: else:
return True return True
def get_model_names(self, base_model: Optional[BaseModelType] = None) -> List[str]: def get_model_names(self, base_model: BaseModelType = BaseModelType.StableDiffusion1) -> List[str]:
model_names = [ model_names = [
info["model_name"] info["model_name"]
for info in self.model_manager.list_models(model_type=ModelType.Main, base_model=base_model) for info in self.model_manager.list_models(model_type=ModelType.Main, base_model=base_model)
@ -318,7 +319,8 @@ class mergeModelsForm(npyscreen.FormMultiPageAction):
return sorted(model_names) return sorted(model_names)
def _populate_models(self, value=None): def _populate_models(self, value=None):
base_model = tuple(BaseModelType)[value[0]] bases = ["sd-1", "sd-2", "sdxl"]
base_model = BaseModelType(bases[value[0]])
self.model_names = self.get_model_names(base_model) self.model_names = self.get_model_names(base_model)
models_plus_none = self.model_names.copy() models_plus_none = self.model_names.copy()