mirror of
https://git.datalinker.icu/ltdrdata/ComfyUI-Manager
synced 2025-12-09 06:04:31 +08:00
[refactor] Move model utility functions to model_utils module
This commit is contained in:
bymyself
parent
e923434d08
commit
f1b3c6b735
@ -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",
|
||||
]
|
||||
|
||||
@ -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 == "<huggingface>":
|
||||
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"] == "<huggingface>":
|
||||
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")
|
||||
|
||||
@ -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 == "<huggingface>":
|
||||
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"] == "<huggingface>":
|
||||
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
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user