302 lines
12 KiB
Python
302 lines
12 KiB
Python
import os
|
|
import subprocess
|
|
import json
|
|
import logging
|
|
from functools import partial
|
|
import re
|
|
import glob
|
|
import sys
|
|
from folder_paths import models_dir as MODELS_DIR
|
|
from folder_paths import base_path as BASE_PATH
|
|
from folder_paths import get_full_path
|
|
|
|
|
|
from .utils.utils import compute_sha256, windows_to_linux_path
|
|
from .utils.pytree import tree_map
|
|
from .file_upload import collect_local_file, process_local_file_path_async
|
|
|
|
|
|
model_list_json = json.load(open(os.path.join(os.path.dirname(__file__), "model_info.json")))
|
|
model_loaders_info = json.load(open(os.path.join(os.path.dirname(__file__), "model_loader_info.json")))
|
|
node_deps_info = json.load(open(os.path.join(os.path.dirname(__file__), "node_deps_info.json")))
|
|
node_blacklist = json.load(open(os.path.join(os.path.dirname(__file__), "node_blacklist.json")))
|
|
|
|
model_suffix = [".ckpt", ".safetensors", ".bin", ".pth", ".pt", ".onnx", ".gguf"]
|
|
extra_packages = ["transformers", "timm", "diffusers", "accelerate"]
|
|
|
|
|
|
def get_full_path_or_raise(folder_name: str, filename: str) -> str:
|
|
full_path = get_full_path(folder_name, filename)
|
|
if full_path is None:
|
|
raise FileNotFoundError(f"Model in folder '{folder_name}' with filename '{filename}' not found.")
|
|
return full_path
|
|
|
|
|
|
def handle_model_info(ckpt_path, filename, rel_save_path):
|
|
ckpt_path = windows_to_linux_path(ckpt_path)
|
|
metadata_path = ckpt_path + ".json"
|
|
if os.path.isfile(metadata_path):
|
|
metadata = json.load(open(metadata_path))
|
|
model_id = metadata["id"]
|
|
else:
|
|
logging.info(f"computing sha256 of {ckpt_path}")
|
|
if not os.path.isfile(ckpt_path):
|
|
raise NotImplementedError(f"please install {ckpt_path} first!")
|
|
model_id = compute_sha256(ckpt_path)
|
|
data = {
|
|
"id": model_id,
|
|
"save_path": rel_save_path,
|
|
"filename": filename,
|
|
}
|
|
json.dump(data, open(metadata_path, "w"))
|
|
if model_id in model_list_json:
|
|
urls = [item["url"] for item in model_list_json[model_id]["links"]][:10] # use the top 10
|
|
else:
|
|
urls = []
|
|
|
|
item = {
|
|
"filename": windows_to_linux_path(filename),
|
|
"save_path": windows_to_linux_path(rel_save_path),
|
|
"urls": urls,
|
|
}
|
|
return model_id, item
|
|
|
|
|
|
def inspect_repo_version(module_path):
|
|
# Create and return the JSON result
|
|
result = {
|
|
"name": os.path.basename(module_path),
|
|
"repo": "",
|
|
"commit": ""
|
|
}
|
|
|
|
if not os.path.isdir(os.path.join(module_path, ".git")):
|
|
return result
|
|
|
|
# Get the remote repository URL
|
|
try:
|
|
remote_url = subprocess.check_output(
|
|
['git', 'config', '--get', 'remote.origin.url'],
|
|
cwd=module_path
|
|
).strip().decode()
|
|
except Exception:
|
|
return result
|
|
|
|
# Get the latest commit hash
|
|
try:
|
|
commit_hash = subprocess.check_output(
|
|
['git', 'rev-parse', 'HEAD'],
|
|
cwd=module_path
|
|
).strip().decode()
|
|
except Exception:
|
|
return result
|
|
|
|
# Create and return the JSON result
|
|
result = {
|
|
"name": os.path.basename(module_path),
|
|
"repo": remote_url,
|
|
"commit": commit_hash
|
|
}
|
|
return result
|
|
|
|
def fetch_model_searcher_results(model_ids):
|
|
import requests
|
|
url = "https://models-searcher.myshell.life/search_urls"
|
|
headers = {
|
|
"Content-Type": "application/json"
|
|
}
|
|
data = {
|
|
"sha256": model_ids
|
|
}
|
|
|
|
response = requests.post(url, headers=headers, json=data)
|
|
if response.status_code == 200:
|
|
results = [item[:10] for item in response.json()]
|
|
else:
|
|
results = None
|
|
return results
|
|
|
|
def split_package_version(require_line):
|
|
require_line = require_line.strip()
|
|
|
|
pattern = r"^([a-zA-Z0-9_\-\[\]]+)(.*)$"
|
|
match = re.match(pattern, require_line.strip())
|
|
|
|
if match:
|
|
package_name = match.group(1) # First capturing group is the package name
|
|
version_specifier = match.group(2) if match.group(2) else "" # Second group is the version, if present
|
|
return package_name, version_specifier
|
|
else:
|
|
assert len(require_line) == 0 or require_line.strip()[0] == "#", require_line
|
|
return None, None
|
|
|
|
def get_package_version(package_name):
|
|
try:
|
|
if sys.version_info >= (3, 8):
|
|
from importlib.metadata import version, PackageNotFoundError
|
|
return version(package_name)
|
|
else:
|
|
from pkg_resources import get_distribution, DistributionNotFound
|
|
return get_distribution(package_name).version
|
|
except Exception:
|
|
return None
|
|
|
|
def resolve_dependencies(prompt, custom_dependencies): # resolve custom nodes and models at the same time
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
import folder_paths
|
|
|
|
custom_nodes = []
|
|
ckpt_paths = {}
|
|
|
|
file_mapping_dict = {}
|
|
|
|
|
|
def collect_unknown_models(filename):
|
|
if type(filename) != str:
|
|
return
|
|
is_model = False
|
|
for possible_suffix in model_suffix:
|
|
if filename.endswith(possible_suffix):
|
|
is_model = True
|
|
if is_model:
|
|
print(f"find {filename}, is_model=True")
|
|
# find possible paths
|
|
matching_files = []
|
|
# Walk through all subdirectories and files in the directory
|
|
rel_save_path = None
|
|
for possible_folder_name in folder_paths.folder_names_and_paths:
|
|
full_path = folder_paths.get_full_path(possible_folder_name, filename)
|
|
if full_path is None:
|
|
continue
|
|
rel_save_path = os.path.relpath(folder_paths.folder_names_and_paths[possible_folder_name][0][0], folder_paths.models_dir)
|
|
matching_files.append(full_path)
|
|
break
|
|
print(f"matched files: {matching_files}")
|
|
if len(set(matching_files)) == 1:
|
|
assert rel_save_path is not None
|
|
ckpt_paths[matching_files[0]] = {
|
|
"filename": filename,
|
|
"rel_save_path": rel_save_path
|
|
}
|
|
|
|
for node_id, node_info in prompt.items():
|
|
node_class_type = node_info.get("class_type")
|
|
if node_class_type is None:
|
|
raise NotImplementedError(f"Missing nodes founded, please first install the missing nodes using ComfyUI Manager")
|
|
node_cls = NODE_CLASS_MAPPINGS[node_class_type]
|
|
if hasattr(node_cls, "RELATIVE_PYTHON_MODULE") and node_cls.RELATIVE_PYTHON_MODULE.startswith("custom_nodes."):
|
|
print(node_cls.RELATIVE_PYTHON_MODULE)
|
|
custom_nodes.append(node_cls.RELATIVE_PYTHON_MODULE)
|
|
if node_class_type in model_loaders_info:
|
|
for field_name, filename in node_info["inputs"].items():
|
|
if type(filename) != str:
|
|
continue
|
|
for item in model_loaders_info[node_class_type]:
|
|
pattern = item["field_name"]
|
|
if re.match(f"^{pattern}$", field_name) and any([filename.endswith(possible_suffix) for possible_suffix in model_suffix]):
|
|
ckpt_path = get_full_path_or_raise(item["save_path"], filename)
|
|
if hasattr(folder_paths, "map_legacy"):
|
|
save_folder = folder_paths.map_legacy(item["save_path"])
|
|
else:
|
|
save_folder = item["save_path"]
|
|
rel_save_path = os.path.relpath(folder_paths.folder_names_and_paths[save_folder][0][0], folder_paths.models_dir)
|
|
ckpt_paths[ckpt_path] = {
|
|
"filename": filename,
|
|
"rel_save_path": rel_save_path
|
|
}
|
|
else:
|
|
tree_map(collect_unknown_models, node_info["inputs"])
|
|
|
|
list(map(partial(collect_local_file, mapping_dict=file_mapping_dict), node_info["inputs"].values()))
|
|
|
|
print("ckpt_paths:", ckpt_paths)
|
|
custom_nodes = list(set(custom_nodes))
|
|
# step 0: comfyui version
|
|
repo_info = inspect_repo_version(BASE_PATH)
|
|
if repo_info["repo"] == "":
|
|
repo_info["require_recheck"] = True
|
|
if repo_info["name"] in custom_dependencies["custom_nodes"]:
|
|
repo_info["repo"] = custom_dependencies["custom_nodes"][repo_info["name"]].get("repo", "")
|
|
repo_info["commit"] = custom_dependencies["custom_nodes"][repo_info["name"]].get("commit", "")
|
|
comfyui_version = repo_info
|
|
|
|
# step 1: custom nodes
|
|
custom_nodes_list = []
|
|
custom_nodes_names = []
|
|
requirements_lines = []
|
|
for custom_node in custom_nodes:
|
|
try:
|
|
repo_info = inspect_repo_version(os.path.join(BASE_PATH, custom_node.replace(".", "/")))
|
|
custom_nodes_list.append(repo_info)
|
|
if repo_info["repo"] == "":
|
|
repo_info["require_recheck"] = True
|
|
if repo_info["name"] in custom_dependencies["custom_nodes"]:
|
|
repo_info["repo"] = custom_dependencies["custom_nodes"][repo_info["name"]].get("repo", "")
|
|
repo_info["commit"] = custom_dependencies["custom_nodes"][repo_info["name"]].get("commit", "")
|
|
custom_nodes_names.append(repo_info["name"])
|
|
except:
|
|
print(f"failed to resolve repo info of {custom_node}")
|
|
requirement_file = os.path.join(BASE_PATH, custom_node.replace(".", "/"), "requirements.txt")
|
|
if os.path.isfile(requirement_file):
|
|
requirements_lines += open(requirement_file).readlines()
|
|
requirements_lines = list(set(requirements_lines))
|
|
requirements_packages = [package_name for package_name, version_specifier in map(split_package_version, requirements_lines) if package_name is not None]
|
|
package_names = set(requirements_packages + extra_packages)
|
|
pypi_deps = {
|
|
package_name: get_package_version(package_name)
|
|
for package_name in package_names
|
|
}
|
|
|
|
for repo_name in custom_nodes_names:
|
|
if repo_name in node_deps_info:
|
|
for deps_node in node_deps_info[repo_name]:
|
|
if deps_node["name"] not in custom_nodes_names:
|
|
repo_info = inspect_repo_version(os.path.join(BASE_PATH, "custom_nodes", deps_node["name"]))
|
|
deps_node["commit"] = repo_info["commit"]
|
|
custom_nodes_list.append(deps_node)
|
|
custom_nodes_names.append(deps_node["name"])
|
|
|
|
black_list_nodes = []
|
|
for repo_name in custom_nodes_names:
|
|
if repo_name in node_blacklist:
|
|
black_list_nodes.append({"name": repo_name, "reason": node_blacklist[repo_name]["reason"]})
|
|
|
|
# step 2: models
|
|
models_dict = {}
|
|
missing_model_ids = []
|
|
for ckpt_path, ckpt_info in ckpt_paths.items():
|
|
model_id, item = handle_model_info(ckpt_path, ckpt_info["filename"], ckpt_info["rel_save_path"])
|
|
models_dict[model_id] = item
|
|
if len(item["urls"]) == 0:
|
|
item["require_recheck"] = True
|
|
if model_id in custom_dependencies["models"]:
|
|
item["urls"] = custom_dependencies["models"][model_id].get("urls", [])
|
|
missing_model_ids.append(model_id)
|
|
|
|
# try to fetch from myshell model searcher
|
|
missing_model_results_myshell = fetch_model_searcher_results(missing_model_ids)
|
|
if missing_model_results_myshell is not None:
|
|
for missing_model_id, missing_model_urls in zip(missing_model_ids, missing_model_results_myshell):
|
|
if len(missing_model_urls) > 0:
|
|
models_dict[missing_model_id]["require_recheck"] = False
|
|
models_dict[missing_model_id]["urls"] = missing_model_urls
|
|
print("successfully fetch results from myshell", models_dict[missing_model_id])
|
|
|
|
# step 3: handle local files
|
|
process_local_file_path_async(file_mapping_dict, max_workers=20)
|
|
files_dict = {v[0]: {"filename": windows_to_linux_path(os.path.relpath(v[2], BASE_PATH)), "urls": [v[1]]} for v in file_mapping_dict.values()}
|
|
|
|
depencencies = {
|
|
"comfyui_version": comfyui_version,
|
|
"custom_nodes": custom_nodes_list,
|
|
"models": models_dict,
|
|
"files": files_dict,
|
|
"pypi": pypi_deps
|
|
}
|
|
|
|
return_dict = {
|
|
"dependencies": depencencies,
|
|
"black_list_nodes": black_list_nodes,
|
|
}
|
|
return return_dict
|