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")