From 3125366edafe353b20060599e1c44e0fbabcf6e2 Mon Sep 17 00:00:00 2001 From: doctorpangloss <@hiddenswitch.com> Date: Thu, 30 May 2024 16:50:34 -0700 Subject: [PATCH] Improve compatibility with comfyui-extra-models, improve API --- comfy/cmd/folder_paths.py | 5 ++- comfy/cmd/server.py | 4 +-- comfy/component_model/files.py | 52 ++++++++++++++++++++++++++++ comfy/model_downloader.py | 5 ++- comfy/nodes/base_nodes.py | 17 +++++---- comfy/nodes/package.py | 4 +-- comfy/sd1_clip.py | 32 ++--------------- comfy/sd2_clip.py | 4 +-- comfy/sdxl_clip.py | 6 ++-- comfy/utils.py | 1 + comfy_extras/nodes/nodes_language.py | 12 ++++--- comfy_extras/nodes/nodes_open_api.py | 6 ++-- 12 files changed, 92 insertions(+), 56 deletions(-) create mode 100644 comfy/component_model/files.py diff --git a/comfy/cmd/folder_paths.py b/comfy/cmd/folder_paths.py index 5c12f2726..2d7af8e26 100644 --- a/comfy/cmd/folder_paths.py +++ b/comfy/cmd/folder_paths.py @@ -7,9 +7,8 @@ import sys import time from typing import Optional, List, Set, Dict, Any, Iterator, Sequence -from pkg_resources import resource_filename - from ..cli_args import args +from ..component_model.files import get_package_as_path supported_pt_extensions = frozenset(['.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl']) @@ -94,7 +93,7 @@ else: models_dir = os.path.join(base_path, "models") folder_names_and_paths = FolderNames(models_dir) folder_names_and_paths["checkpoints"] = FolderPathsTuple("checkpoints", [os.path.join(models_dir, "checkpoints")], set(supported_pt_extensions)) -folder_names_and_paths["configs"] = FolderPathsTuple("configs", [os.path.join(models_dir, "configs"), resource_filename("comfy", "configs/")], {".yaml"}) +folder_names_and_paths["configs"] = FolderPathsTuple("configs", [os.path.join(models_dir, "configs"), get_package_as_path("comfy.configs")], {".yaml"}) folder_names_and_paths["loras"] = FolderPathsTuple("loras", [os.path.join(models_dir, "loras")], set(supported_pt_extensions)) folder_names_and_paths["vae"] = FolderPathsTuple("vae", [os.path.join(models_dir, "vae")], set(supported_pt_extensions)) folder_names_and_paths["clip"] = FolderPathsTuple("clip", [os.path.join(models_dir, "clip")], set(supported_pt_extensions)) diff --git a/comfy/cmd/server.py b/comfy/cmd/server.py index 5a5d18c49..20b968f0a 100644 --- a/comfy/cmd/server.py +++ b/comfy/cmd/server.py @@ -23,7 +23,6 @@ from PIL import Image from PIL.PngImagePlugin import PngInfo from aiohttp import web from can_ada import URL, parse as urlparse -from pkg_resources import resource_filename from typing_extensions import NamedTuple import comfy.interruption @@ -38,6 +37,7 @@ from ..cmd import folder_paths from ..component_model.abstract_prompt_queue import AbstractPromptQueue, AsyncAbstractPromptQueue from ..component_model.executor_types import ExecutorToClientProgress from ..component_model.file_output_path import file_output_path +from ..component_model.files import get_package_as_path from ..component_model.queue_types import QueueItem, HistoryEntry, BinaryEventTypes, TaskInvocation, ExecutionError, \ ExecutionStatus from ..digest import digest @@ -115,7 +115,7 @@ class PromptServer(ExecutorToClientProgress): self.sockets = dict() web_root_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../web") if not os.path.exists(web_root_path): - web_root_path = resource_filename('comfy', 'web/') + web_root_path = get_package_as_path('comfy', 'web/') self.web_root = web_root_path routes = web.RouteTableDef() self.routes: web.RouteTableDef = routes diff --git a/comfy/component_model/files.py b/comfy/component_model/files.py new file mode 100644 index 000000000..4bddc7163 --- /dev/null +++ b/comfy/component_model/files.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import json +import os +from importlib import resources as resources +from typing import Optional + + +def get_path_as_dict(config_dict_or_path: str | dict | None, config_path_inside_package: str, package: str = 'comfy') -> dict: + """ + Given a package and a filename inside the package, returns it as a JSON dict; or, returns the file pointed to by + config_dict_or_path, when it is not None and when it exists + + :param config_dict_or_path: a file path or dict pointing to a JSON file. If it exists, it is parsed and returned. Otherwise, when None, falls back to other defaults + :param config_path_inside_package: a filename inside a package + :param package: a package containing the file + :return: + """ + config: dict | None = None + + if config_dict_or_path is None: + config_dict_or_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), config_path_inside_package) + + if isinstance(config_dict_or_path, str): + if config_dict_or_path.startswith("{"): + config = json.loads(config_dict_or_path) + else: + if not os.path.exists(config_dict_or_path): + with resources.as_file(resources.files(package) / config_path_inside_package) as config_path: + with open(config_path) as f: + config = json.load(f) + else: + with open(config_dict_or_path) as f: + config = json.load(f) + elif isinstance(config_dict_or_path, dict): + config = config_dict_or_path + + assert config is not None + return config + + +def get_package_as_path(package: str, subdir: Optional[str] = None) -> str: + """ + Gets the path on the file system to a package. This unpacks it completely. + :param package: the package containing the files + :param subdir: if specified, a subdirectory containing files (and not python packages), such as a web/ directory inside a package + :return: + """ + traversable = resources.files(package) + if subdir is not None: + traversable = traversable / subdir + return os.path.commonpath(list(map(str, traversable.iterdir()))) diff --git a/comfy/model_downloader.py b/comfy/model_downloader.py index da680d429..ce4aebf45 100644 --- a/comfy/model_downloader.py +++ b/comfy/model_downloader.py @@ -167,7 +167,6 @@ KNOWN_CHECKPOINTS = [ CivitFile(133005, 357609, filename="juggernautXL_v9Rundiffusionphoto2.safetensors"), CivitFile(112902, 351306, filename="dreamshaperXL_v21TurboDPMSDE.safetensors"), CivitFile(139562, 344487, filename="realvisxlV40_v40Bakedvae.safetensors"), - ] KNOWN_UNCLIP_CHECKPOINTS = [ @@ -304,6 +303,10 @@ KNOWN_HUGGINGFACE_MODEL_REPOS = { "microsoft/Phi-3-mini-4k-instruct", } +KNOWN_UNET_MODELS: List[Union[CivitFile | HuggingFile]] = [] + +KNOWN_CLIP_MODELS: List[Union[CivitFile | HuggingFile]] = [] + def add_known_models(folder_name: str, symbol: List[Union[CivitFile, HuggingFile]], *models: Union[CivitFile, HuggingFile]) -> List[Union[CivitFile, HuggingFile]]: if args.disable_known_models: diff --git a/comfy/nodes/base_nodes.py b/comfy/nodes/base_nodes.py index ec521edef..b9a84dad3 100644 --- a/comfy/nodes/base_nodes.py +++ b/comfy/nodes/base_nodes.py @@ -26,7 +26,7 @@ from ..cli_args import args from ..cmd import folder_paths, latent_preview from ..execution_context import current_execution_context from ..images import open_image -from ..model_downloader import get_filename_list_with_downloadable, get_or_download, KNOWN_CHECKPOINTS, KNOWN_CLIP_VISION_MODELS, KNOWN_GLIGEN_MODELS, KNOWN_UNCLIP_CHECKPOINTS, KNOWN_LORAS, KNOWN_CONTROLNETS, KNOWN_DIFF_CONTROLNETS, KNOWN_VAES, KNOWN_APPROX_VAES, huggingface_repos +from ..model_downloader import get_filename_list_with_downloadable, get_or_download, KNOWN_CHECKPOINTS, KNOWN_CLIP_VISION_MODELS, KNOWN_GLIGEN_MODELS, KNOWN_UNCLIP_CHECKPOINTS, KNOWN_LORAS, KNOWN_CONTROLNETS, KNOWN_DIFF_CONTROLNETS, KNOWN_VAES, KNOWN_APPROX_VAES, huggingface_repos, KNOWN_CLIP_MODELS, KNOWN_UNET_MODELS from ..nodes.common import MAX_RESOLUTION from .. import controlnet from ..open_exr import load_exr @@ -799,7 +799,7 @@ class ControlNetApplyAdvanced: class UNETLoader: @classmethod def INPUT_TYPES(s): - return {"required": { "unet_name": (folder_paths.get_filename_list("unet"),), + return {"required": { "unet_name": (get_filename_list_with_downloadable("unet", KNOWN_UNET_MODELS),), }} RETURN_TYPES = ("MODEL",) FUNCTION = "load_unet" @@ -807,14 +807,14 @@ class UNETLoader: CATEGORY = "advanced/loaders" def load_unet(self, unet_name): - unet_path = folder_paths.get_full_path("unet", unet_name) + unet_path = get_or_download("unet", unet_name, KNOWN_UNET_MODELS) model = sd.load_unet(unet_path) return (model,) class CLIPLoader: @classmethod def INPUT_TYPES(s): - return {"required": { "clip_name": (folder_paths.get_filename_list("clip"),), + return {"required": { "clip_name": (get_filename_list_with_downloadable("clip", KNOWN_CLIP_MODELS),), "type": (["stable_diffusion", "stable_cascade"], ), }} RETURN_TYPES = ("CLIP",) @@ -823,11 +823,14 @@ class CLIPLoader: CATEGORY = "advanced/loaders" def load_clip(self, clip_name, type="stable_diffusion"): - clip_type = sd.CLIPType.STABLE_DIFFUSION - if type == "stable_cascade": + if type == "stable_diffusion": + clip_type = sd.CLIPType.STABLE_DIFFUSION + elif type == "stable_cascade": clip_type = sd.CLIPType.STABLE_CASCADE + else: + logging.warning(f"Unknown clip type argument passed: {type} for model {clip_name}") - clip_path = folder_paths.get_full_path("clip", clip_name) + clip_path = get_or_download("clip", clip_name, KNOWN_CLIP_MODELS) clip = sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type) return (clip,) diff --git a/comfy/nodes/package.py b/comfy/nodes/package.py index bc0841f61..253d76a65 100644 --- a/comfy/nodes/package.py +++ b/comfy/nodes/package.py @@ -10,10 +10,10 @@ from functools import reduce from importlib.metadata import entry_points from opentelemetry.trace import Span, Status, StatusCode -from pkg_resources import resource_filename from .package_typing import ExportedNodes from ..cmd.main_pre import tracer +from ..component_model.files import get_package_as_path _comfy_nodes: ExportedNodes = ExportedNodes() @@ -28,7 +28,7 @@ def _import_nodes_in_module(exported_nodes: ExportedNodes, module: types.ModuleT exported_nodes.NODE_DISPLAY_NAME_MAPPINGS.update(node_display_names) if web_directory: # load the extension resources path - abs_web_directory = os.path.abspath(resource_filename(module.__name__, web_directory)) + abs_web_directory = os.path.abspath(get_package_as_path(module.__name__, web_directory)) if not os.path.isdir(abs_web_directory): abs_web_directory = os.path.abspath(os.path.join(os.path.dirname(module.__file__), web_directory)) if not os.path.isdir(abs_web_directory): diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py index ef134cc59..8f8f39ea4 100644 --- a/comfy/sd1_clip.py +++ b/comfy/sd1_clip.py @@ -1,8 +1,6 @@ from __future__ import annotations import copy -import importlib.resources as resources -import json import logging import os import traceback @@ -10,36 +8,12 @@ import zipfile from typing import Tuple, Sequence, TypeVar import torch -from pkg_resources import resource_filename from transformers import CLIPTokenizer, PreTrainedTokenizerBase from . import clip_model from . import model_management from . import ops - - -def get_clip_config_dict(text_model_config_or_path: str | dict | None, text_model_config_path_in_comfy: str, package: str = 'comfy') -> dict: - config: dict | None = None - - if text_model_config_or_path is None: - text_model_config_or_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), text_model_config_path_in_comfy) - - if isinstance(text_model_config_or_path, str): - if text_model_config_or_path.startswith("{"): - config = json.loads(text_model_config_or_path) - else: - if not os.path.exists(text_model_config_or_path): - with resources.as_file(resources.files(package) / text_model_config_path_in_comfy) as config_path: - with open(config_path) as f: - config = json.load(f) - else: - with open(text_model_config_or_path) as f: - config = json.load(f) - elif isinstance(text_model_config_or_path, dict): - config = text_model_config_or_path - - assert config is not None - return config +from .component_model.files import get_path_as_dict, get_package_as_path def gen_empty_tokens(special_tokens, length): @@ -109,7 +83,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder): special_tokens = {"start": 49406, "end": 49407, "pad": 49407} assert layer in self.LAYERS - config = get_clip_config_dict(textmodel_json_config, "sd1_clip_config.json") + config = get_path_as_dict(textmodel_json_config, "sd1_clip_config.json") self.transformer = model_class(config, dtype, device, ops.manual_cast) self.num_layers = self.transformer.num_layers @@ -402,7 +376,7 @@ class SDTokenizer: tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_tokenizer") if not os.path.exists(os.path.join(tokenizer_path, "tokenizer_config.json")): # package based - tokenizer_path = resource_filename('comfy', 'sd1_tokenizer/') + tokenizer_path = get_package_as_path('comfy.sd1_tokenizer') self.tokenizer_class = tokenizer_class self.tokenizer_path = tokenizer_path self.tokenizer: PreTrainedTokenizerBase = tokenizer_class.from_pretrained(tokenizer_path) diff --git a/comfy/sd2_clip.py b/comfy/sd2_clip.py index 6fbbedb69..663e7f98f 100644 --- a/comfy/sd2_clip.py +++ b/comfy/sd2_clip.py @@ -1,6 +1,6 @@ from . import sd1_clip -from .sd1_clip import get_clip_config_dict +from .component_model.files import get_path_as_dict class SD2ClipHModel(sd1_clip.SDClipModel): @@ -9,7 +9,7 @@ class SD2ClipHModel(sd1_clip.SDClipModel): layer = "hidden" layer_idx = -2 - textmodel_json_config = get_clip_config_dict(textmodel_json_config, "sd2_clip_config.json") + textmodel_json_config = get_path_as_dict(textmodel_json_config, "sd2_clip_config.json") super().__init__(device=device, freeze=freeze, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"start": 49406, "end": 49407, "pad": 0}) diff --git a/comfy/sdxl_clip.py b/comfy/sdxl_clip.py index 5478c60ad..40988dc25 100644 --- a/comfy/sdxl_clip.py +++ b/comfy/sdxl_clip.py @@ -3,7 +3,7 @@ import copy import torch from . import sd1_clip -from .sd1_clip import get_clip_config_dict +from .component_model.files import get_path_as_dict class SDXLClipG(sd1_clip.SDClipModel): @@ -12,7 +12,7 @@ class SDXLClipG(sd1_clip.SDClipModel): layer = "hidden" layer_idx = -2 - textmodel_json_config = get_clip_config_dict(textmodel_json_config, "clip_config_bigg.json") + textmodel_json_config = get_path_as_dict(textmodel_json_config, "clip_config_bigg.json") super().__init__(device=device, freeze=freeze, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"start": 49406, "end": 49407, "pad": 0}, layer_norm_hidden_state=False) @@ -91,7 +91,7 @@ class StableCascadeTokenizer(sd1_clip.SD1Tokenizer): class StableCascadeClipG(sd1_clip.SDClipModel): def __init__(self, device="cpu", max_length=77, freeze=True, layer="hidden", layer_idx=-1, dtype=None, textmodel_json_config=None): - textmodel_json_config = get_clip_config_dict(textmodel_json_config, "clip_config_bigg.json") + textmodel_json_config = get_path_as_dict(textmodel_json_config, "clip_config_bigg.json") super().__init__(device=device, freeze=freeze, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"start": 49406, "end": 49407, "pad": 49407}, layer_norm_hidden_state=False, enable_attention_masks=True) diff --git a/comfy/utils.py b/comfy/utils.py index 5f7697f3f..0775fa098 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -3,6 +3,7 @@ from __future__ import annotations import contextlib import logging import math +import os import os.path import random import struct diff --git a/comfy_extras/nodes/nodes_language.py b/comfy_extras/nodes/nodes_language.py index 05f37c549..c059a8aec 100644 --- a/comfy_extras/nodes/nodes_language.py +++ b/comfy_extras/nodes/nodes_language.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, List, Dict +from typing import Any, Dict, Optional import torch from fastchat.model import get_conversation_template @@ -31,16 +31,20 @@ class TransformersLoader(CustomNode): def INPUT_TYPES(cls) -> InputTypes: return { "required": { - "ckpt_name": (huggingface_repos(),) + "ckpt_name": (huggingface_repos(),), + "subfolder": ("STRING", {}) } } RETURN_TYPES = "MODEL", FUNCTION = "execute" - def execute(self, ckpt_name: str): + def execute(self, ckpt_name: str, subfolder: Optional[str] = None): + hub_kwargs = {} + if subfolder is not None and subfolder != "": + hub_kwargs["subfolder"] = subfolder with comfy_tqdm(): - model = AutoModelForCausalLM.from_pretrained(ckpt_name, torch_dtype=unet_dtype(), device_map=get_torch_device_name(unet_offload_device()), low_cpu_mem_usage=True, trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained(ckpt_name, torch_dtype=unet_dtype(), device_map=get_torch_device_name(unet_offload_device()), low_cpu_mem_usage=True, trust_remote_code=True, **hub_kwargs) tokenizer = AutoTokenizer.from_pretrained(ckpt_name) model_managed = TransformersManagedModel(ckpt_name, model, tokenizer) return model_managed, diff --git a/comfy_extras/nodes/nodes_open_api.py b/comfy_extras/nodes/nodes_open_api.py index 32501dfb2..b71792a72 100644 --- a/comfy_extras/nodes/nodes_open_api.py +++ b/comfy_extras/nodes/nodes_open_api.py @@ -215,7 +215,7 @@ class StringPosixPathJoin(CustomNode): def INPUT_TYPES(cls) -> InputTypes: return { "required": { - f"value{i}": ("STRING", {"default": "", "multiline": False}) for i in range(5) + f"value{i}": ("STRING", {"default": "", "multiline": False, "forceInput": True}) for i in range(5) } } @@ -284,7 +284,7 @@ class DevNullUris(CustomNode): class StringJoin(CustomNode): @classmethod def INPUT_TYPES(cls) -> InputTypes: - required = {f"value{i}": ("STRING", {"default": "", "multiline": True}) for i in range(5)} + required = {f"value{i}": ("STRING", {"default": "", "multiline": True, "forceInput": True}) for i in range(5)} required["separator"] = ("STRING", {"default": "_"}) return { "required": required @@ -304,7 +304,7 @@ class StringToUri(CustomNode): def INPUT_TYPES(cls) -> InputTypes: return { "required": { - "value": ("STRING", {"default": "", "multiline": True}), + "value": ("STRING", {"default": "", "multiline": True, "forceInput": True}), "batch": ("INT", {"default": 1}) } }