From 5089ac5ad102b6d21edcef7205175a933014a80c Mon Sep 17 00:00:00 2001 From: cmdr2 Date: Mon, 16 Jan 2023 18:30:22 +0530 Subject: [PATCH] Fix a bug where the .vae.pt extension wouldn't get picked up. Thanks Madrang, rbertus2000 and JeLuf --- ui/easydiffusion/model_manager.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/ui/easydiffusion/model_manager.py b/ui/easydiffusion/model_manager.py index 8fbd6aa3..7ed69a8c 100644 --- a/ui/easydiffusion/model_manager.py +++ b/ui/easydiffusion/model_manager.py @@ -190,7 +190,11 @@ def getModels(): nonlocal models_scanned tree = [] for entry in os.scandir(directory): - if entry.is_file() and True in [entry.name.endswith(s) for s in suffixes]: + if entry.is_file(): + matching_suffix = list(filter(lambda s: entry.name.endswith(s), suffixes)) + if len(matching_suffix) == 0: continue + matching_suffix = matching_suffix[0] + mtime = entry.stat().st_mtime mod_time = known_models[entry.path] if entry.path in known_models else -1 if mod_time != mtime: @@ -198,7 +202,7 @@ def getModels(): if is_malicious_model(entry.path): raise MaliciousModelException(entry.path) known_models[entry.path] = mtime - tree.append(entry.name.rsplit('.',1)[0]) + tree.append(entry.name[:-len(matching_suffix)]) elif entry.is_dir(): scan=scan_directory(entry.path, suffixes) if len(scan) != 0: