diff --git a/comfyui_manager/glob/constants.py b/comfyui_manager/glob/constants.py index 87d28cb6..9ec94e54 100644 --- a/comfyui_manager/glob/constants.py +++ b/comfyui_manager/glob/constants.py @@ -33,3 +33,22 @@ model_dir_name_map = { "unet": "diffusion_models", "diffusion_model": "diffusion_models", } + +# List of all model directory names used for checking installed models +MODEL_DIR_NAMES = [ + "checkpoints", + "loras", + "vae", + "text_encoders", + "diffusion_models", + "clip_vision", + "embeddings", + "diffusers", + "vae_approx", + "controlnet", + "gligen", + "upscale_models", + "hypernetworks", + "photomaker", + "classifiers", +] diff --git a/comfyui_manager/glob/manager_server.py b/comfyui_manager/glob/manager_server.py index 0f9615ae..b0842273 100644 --- a/comfyui_manager/glob/manager_server.py +++ b/comfyui_manager/glob/manager_server.py @@ -6,7 +6,6 @@ Handles task queue management, custom node operations, model installation, and s """ import asyncio -import concurrent import copy import heapq import json @@ -1365,82 +1364,7 @@ async def installed_list(request): return web.json_response(res, content_type="application/json") -def check_model_installed(json_obj): - def is_exists(model_dir_name, filename, url): - if filename == "": - filename = os.path.basename(url) - - dirs = folder_paths.get_folder_paths(model_dir_name) - - for x in dirs: - if os.path.exists(os.path.join(x, filename)): - return True - - return False - - model_dir_names = [ - "checkpoints", - "loras", - "vae", - "text_encoders", - "diffusion_models", - "clip_vision", - "embeddings", - "diffusers", - "vae_approx", - "controlnet", - "gligen", - "upscale_models", - "hypernetworks", - "photomaker", - "classifiers", - ] - - total_models_files = set() - for x in model_dir_names: - for y in folder_paths.get_filename_list(x): - total_models_files.add(y) - - def process_model_phase(item): - if ( - "diffusion" not in item["filename"] - and "pytorch" not in item["filename"] - and "model" not in item["filename"] - ): - # non-general name case - if item["filename"] in total_models_files: - item["installed"] = "True" - return - - if item["save_path"] == "default": - model_dir_name = model_dir_name_map.get(item["type"].lower()) - if model_dir_name is not None: - item["installed"] = str( - is_exists(model_dir_name, item["filename"], item["url"]) - ) - else: - item["installed"] = "False" - else: - model_dir_name = item["save_path"].split("/")[0] - if model_dir_name in folder_paths.folder_names_and_paths: - if is_exists(model_dir_name, item["filename"], item["url"]): - item["installed"] = "True" - - if "installed" not in item: - if item["filename"] == "": - filename = os.path.basename(item["url"]) - else: - filename = item["filename"] - - fullpath = os.path.join( - folder_paths.models_dir, item["save_path"], filename - ) - - item["installed"] = "True" if os.path.exists(fullpath) else "False" - - with concurrent.futures.ThreadPoolExecutor(8) as executor: - for item in json_obj["models"]: - executor.submit(process_model_phase, item) +# Function moved to model_utils.py @routes.get("/v2/snapshot/getlist") @@ -1725,28 +1649,7 @@ async def comfyui_switch_version(request): return web.Response(status=400) -async def check_whitelist_for_model(item): - json_obj = await core.get_data_by_mode(ManagerDatabaseSource.cache.value, "model-list.json") - - for x in json_obj.get("models", []): - if ( - x["save_path"] == item["save_path"] - and x["base"] == item["base"] - and x["filename"] == item["filename"] - ): - return True - - json_obj = await core.get_data_by_mode(ManagerDatabaseSource.local.value, "model-list.json") - - for x in json_obj.get("models", []): - if ( - x["save_path"] == item["save_path"] - and x["base"] == item["base"] - and x["filename"] == item["filename"] - ): - return True - - return False +# Function moved to model_utils.py @routes.post("/v2/manager/queue/install_model") diff --git a/comfyui_manager/glob/utils/model_utils.py b/comfyui_manager/glob/utils/model_utils.py index 9fffa3de..2225ec07 100644 --- a/comfyui_manager/glob/utils/model_utils.py +++ b/comfyui_manager/glob/utils/model_utils.py @@ -1,9 +1,10 @@ import os import logging +import concurrent.futures import folder_paths from comfyui_manager.glob import manager_core as core -from comfyui_manager.glob.constants import model_dir_name_map +from comfyui_manager.glob.constants import model_dir_name_map, MODEL_DIR_NAMES def get_model_dir(data, show_log=False): @@ -72,3 +73,89 @@ def get_model_path(data, show_log=False): return os.path.join(base_model, os.path.basename(data["url"])) else: return os.path.join(base_model, data["filename"]) + + +def check_model_installed(json_obj): + def is_exists(model_dir_name, filename, url): + if filename == "": + filename = os.path.basename(url) + + dirs = folder_paths.get_folder_paths(model_dir_name) + + for x in dirs: + if os.path.exists(os.path.join(x, filename)): + return True + + return False + + total_models_files = set() + for x in MODEL_DIR_NAMES: + for y in folder_paths.get_filename_list(x): + total_models_files.add(y) + + def process_model_phase(item): + if ( + "diffusion" not in item["filename"] + and "pytorch" not in item["filename"] + and "model" not in item["filename"] + ): + # non-general name case + if item["filename"] in total_models_files: + item["installed"] = "True" + return + + if item["save_path"] == "default": + model_dir_name = model_dir_name_map.get(item["type"].lower()) + if model_dir_name is not None: + item["installed"] = str( + is_exists(model_dir_name, item["filename"], item["url"]) + ) + else: + item["installed"] = "False" + else: + model_dir_name = item["save_path"].split("/")[0] + if model_dir_name in folder_paths.folder_names_and_paths: + if is_exists(model_dir_name, item["filename"], item["url"]): + item["installed"] = "True" + + if "installed" not in item: + if item["filename"] == "": + filename = os.path.basename(item["url"]) + else: + filename = item["filename"] + + fullpath = os.path.join( + folder_paths.models_dir, item["save_path"], filename + ) + + item["installed"] = "True" if os.path.exists(fullpath) else "False" + + with concurrent.futures.ThreadPoolExecutor(8) as executor: + for item in json_obj["models"]: + executor.submit(process_model_phase, item) + + +async def check_whitelist_for_model(item): + from comfyui_manager.data_models import ManagerDatabaseSource + + json_obj = await core.get_data_by_mode(ManagerDatabaseSource.cache.value, "model-list.json") + + for x in json_obj.get("models", []): + if ( + x["save_path"] == item["save_path"] + and x["base"] == item["base"] + and x["filename"] == item["filename"] + ): + return True + + json_obj = await core.get_data_by_mode(ManagerDatabaseSource.local.value, "model-list.json") + + for x in json_obj.get("models", []): + if ( + x["save_path"] == item["save_path"] + and x["base"] == item["base"] + and x["filename"] == item["filename"] + ): + return True + + return False