| import json
|
| import os
|
| import folder_paths
|
| import server
|
| from .utils import find_tags
|
|
|
| class easyModelManager:
|
|
|
| def __init__(self):
|
| self.img_suffixes = [".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff", ".svg", ".tif", ".tiff"]
|
| self.default_suffixes = [".ckpt", ".pt", ".bin", ".pth", ".safetensors"]
|
| self.models_config = {
|
| "checkpoints": {"suffix": self.default_suffixes},
|
| "loras": {"suffix": self.default_suffixes},
|
| "unet": {"suffix": self.default_suffixes},
|
| }
|
| self.model_lists = {}
|
|
|
| def find_thumbnail(self, model_type, name):
|
| file_no_ext = os.path.splitext(name)[0]
|
| for ext in self.img_suffixes:
|
| full_path = folder_paths.get_full_path(model_type, file_no_ext + ext)
|
| if os.path.isfile(str(full_path)):
|
| return full_path
|
| return None
|
|
|
| def get_model_lists(self, model_type):
|
| if model_type not in self.models_config:
|
| return []
|
| filenames = folder_paths.get_filename_list(model_type)
|
| model_lists = []
|
| for name in filenames:
|
| model_suffix = os.path.splitext(name)[-1]
|
| if model_suffix not in self.models_config[model_type]["suffix"]:
|
| continue
|
| else:
|
| cfg = {
|
| "name": os.path.basename(os.path.splitext(name)[0]),
|
| "full_name": name,
|
| "remark": '',
|
| "file_path": folder_paths.get_full_path(model_type, name),
|
| "type": model_type,
|
| "suffix": model_suffix,
|
| "dir_tags": find_tags(name),
|
| "cover": self.find_thumbnail(model_type, name),
|
| "metadata": None,
|
| "sha256": None
|
| }
|
| model_lists.append(cfg)
|
|
|
| return model_lists
|
|
|
| def get_model_info(self, model_type, model_name):
|
| pass
|
|
|
|
|
|
|
| |
