refactor: extract remove_special_loras

This commit is contained in:
Manuel Schmid 2024-03-23 16:11:41 +01:00
parent a6d4a27bb0
commit 2f52cab75d
No known key found for this signature in database
GPG Key ID: 32C4F7569B40B84B
2 changed files with 9 additions and 6 deletions

View File

@ -539,6 +539,7 @@ wildcard_filenames = []
sdxl_lcm_lora = 'sdxl_lcm_lora.safetensors'
sdxl_lightning_lora = 'sdxl_lightning_4step_lora.safetensors'
loras_metadata_remove = [sdxl_lcm_lora, sdxl_lightning_lora]
def get_model_filenames(folder_paths, extensions=None, name_filter=None):

View File

@ -290,6 +290,12 @@ class MetadataParser(ABC):
lora_hash = get_sha256(lora_path)
self.loras.append((Path(lora_name).stem, lora_weight, lora_hash))
@staticmethod
def remove_special_loras(lora_filenames):
for lora_to_remove in modules.config.loras_metadata_remove:
if lora_to_remove in lora_filenames:
lora_filenames.remove(lora_to_remove)
class A1111MetadataParser(MetadataParser):
def get_scheme(self) -> MetadataScheme:
@ -407,9 +413,7 @@ class A1111MetadataParser(MetadataParser):
if lora_data != '':
lora_filenames = modules.config.lora_filenames.copy()
for lora_to_remove in [modules.config.sdxl_lcm_lora, modules.config.sdxl_lightning_lora]:
if lora_to_remove in lora_filenames:
lora_filenames.remove(lora_to_remove)
self.remove_special_loras(lora_filenames)
for li, lora in enumerate(lora_data.split(', ')):
lora_split = lora.split(': ')
lora_name = lora_split[0]
@ -495,9 +499,7 @@ class FooocusMetadataParser(MetadataParser):
def parse_json(self, metadata: dict) -> dict:
model_filenames = modules.config.model_filenames.copy()
lora_filenames = modules.config.lora_filenames.copy()
for lora_to_remove in [modules.config.sdxl_lcm_lora, modules.config.sdxl_lightning_lora]:
if lora_to_remove in lora_filenames:
lora_filenames.remove(lora_to_remove)
self.remove_special_loras(lora_filenames)
for key, value in metadata.items():
if value in ['', 'None']:
continue