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