From 1a552c4c8ca53e3402b0425da066cd57d37ab52d Mon Sep 17 00:00:00 2001 From: sn0w12 Date: Tue, 25 Jun 2024 03:04:09 +0200 Subject: [PATCH] add get_model_type_simple --- sn0w.py | 13 +++++++++++++ src/load_lora_from_folder.py | 4 ++-- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/sn0w.py b/sn0w.py index 1e25b22..49217a0 100644 --- a/sn0w.py +++ b/sn0w.py @@ -175,6 +175,19 @@ class Utility: def get_model_type(model_patcher): return model_patcher.model.__class__.__name__ + @staticmethod + def get_model_type_simple(model_patcher): + model_type = Utility.get_model_type(model_patcher) + + if model_type == "BaseModel": + return "SD15" + elif model_type == "SDXL": + return "SDXL" + elif model_type == "SD3": + return "SD3" + + return None + @classmethod def put_favourite_on_top(cls, setting, arr): # Convert to a list if its a dictionary diff --git a/src/load_lora_from_folder.py b/src/load_lora_from_folder.py index 1ca83ff..7ba4ba9 100644 --- a/src/load_lora_from_folder.py +++ b/src/load_lora_from_folder.py @@ -64,13 +64,13 @@ class LoadLoraFolderNode: # Clean and split the prompt into parts prompt_parts = [self.clean_string(part) for part in prompt.split(separator)] - model_type = Utility.get_model_type(model) + model_type = Utility.get_model_type_simple(model) # Retrieve and filter Lora file paths full_lora_paths = folder_paths.get_filename_list("loras") try: - if model_type == "BaseModel": + if model_type == "SD15": filtered_lora_paths = folder_paths.get_filename_list("loras_15") elif model_type == "SDXL": filtered_lora_paths = folder_paths.get_filename_list("loras_xl")