bizyengine 0.4.2__py3-none-any.whl

bizyengine/core/nodes_base.py

@@ -0,0 +1,159 @@
+import importlib
+import logging
+import warnings
+from functools import wraps
+from typing import List
+
+from .data_types import is_send_request_datatype
+from .nodes_io import BizyAirNodeIO, create_node_data
+
+try:
+    comfy_nodes = importlib.import_module("nodes")
+except ModuleNotFoundError:
+    warnings.warn("Importing comfyui.nodes failed!")
+    comfy_nodes = type("nodes", (object,), {"NODE_DISPLAY_NAME_MAPPINGS": {}})
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+LOGO = "☁️"
+PREFIX = f"BizyAir"
+NODE_CLASS_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS = {}
+
+
+def to_camel_case(s):
+    return "".join(word.capitalize() for word in s.split("_"))
+
+
+def validate_category(cls, prefix):
+    assert cls.CATEGORY.startswith(f"☁️{prefix}")
+
+
+def register_node(cls, prefix):
+    class_name = (
+        f"{prefix}_{cls.__name__}"
+        if not cls.__name__.startswith(prefix)
+        else cls.__name__
+    )
+    logger.debug(
+        f"Class: {cls}, Name: {class_name}, Has DISPLAY_NAME: {hasattr(cls, 'NODE_DISPLAY_NAME')}"
+    )
+
+    if hasattr(cls, "NODE_DISPLAY_NAME"):
+        display_name = cls.NODE_DISPLAY_NAME
+        if not display_name.startswith(f"{LOGO}{prefix}"):
+            display_name = f"{LOGO}{prefix} {display_name}"
+    else:
+        base_name = class_name[len(prefix) + 1 :]
+        if base_name in comfy_nodes.NODE_DISPLAY_NAME_MAPPINGS:
+            display_name = (
+                f"{LOGO}{prefix} {comfy_nodes.NODE_DISPLAY_NAME_MAPPINGS[base_name]}"
+            )
+        else:
+            display_name = f"{LOGO}{prefix} {base_name}"
+            logger.warning(
+                f"Display name '{display_name}' might differ from the native display name."
+            )
+
+    NODE_CLASS_MAPPINGS[class_name] = cls
+    NODE_DISPLAY_NAME_MAPPINGS[class_name] = display_name
+
+
+def ensure_unique_id(org_func, original_has_unique_id=False):
+    @wraps(org_func)
+    def new_func(self, **kwargs):
+        if original_has_unique_id:
+            self._assigned_id = kwargs.get("unique_id", "UNIQUE_ID")
+        elif "unique_id" in kwargs:
+            self._assigned_id = kwargs.pop("unique_id")
+        return org_func(self, **kwargs)
+
+    return new_func
+
+
+def ensure_hidden_unique_id(org_input_types_func):
+    original_has_unique_id = False
+
+    @wraps(org_input_types_func)
+    def new_input_types_func():
+        nonlocal original_has_unique_id
+
+        result = org_input_types_func()
+        if "hidden" not in result:
+            result["hidden"] = {"unique_id": "UNIQUE_ID"}
+        elif "unique_id" not in result["hidden"]:
+            result["hidden"]["unique_id"] = "UNIQUE_ID"
+        else:
+            original_has_unique_id = True
+        return result
+
+    # Ensure original_has_unique_id is correctly set before returning
+    new_input_types_func()
+    return new_input_types_func, original_has_unique_id
+
+
+class BizyAirBaseNode:
+    FUNCTION = "default_function"
+
+    def __init_subclass__(cls, **kwargs):
+        if not cls.CATEGORY.startswith(f"{LOGO}{PREFIX}"):
+            cls.CATEGORY = f"{LOGO}{PREFIX}/{cls.CATEGORY}"
+        register_node(cls, PREFIX)
+        cls.setup_input_types()
+
+    @classmethod
+    def setup_input_types(cls):
+        # https://docs.comfy.org/essentials/custom_node_more_on_inputs#hidden-inputs
+        new_input_types_func, original_has_unique_id = ensure_hidden_unique_id(
+            cls.INPUT_TYPES
+        )
+        cls.INPUT_TYPES = new_input_types_func
+        setattr(
+            cls,
+            cls.FUNCTION,
+            ensure_unique_id(getattr(cls, cls.FUNCTION), original_has_unique_id),
+        )
+
+    @property
+    def assigned_id(self):
+        assert self._assigned_id is not None and isinstance(self._assigned_id, str)
+        return str(self._assigned_id)
+
+    def default_function(self, **kwargs):
+        class_type = self._determine_class_type()
+
+        node_ios = self._process_non_send_request_types(class_type, kwargs)
+        # TODO: add processing for send_request_types
+        send_request_datatype_list = self._get_send_request_datatypes()
+        if len(send_request_datatype_list) == len(self.RETURN_TYPES):
+            return self._process_all_send_request_types(node_ios)
+        return node_ios
+
+    def _get_send_request_datatypes(self):
+        return [
+            return_type
+            for return_type in self.RETURN_TYPES
+            if is_send_request_datatype(return_type)
+        ]
+
+    def _determine_class_type(self):
+        class_type = getattr(self, "CLASS_TYPE_NAME", type(self).__name__)
+        if class_type.startswith(f"{PREFIX}_"):
+            class_type = class_type[len(PREFIX) + 1 :]
+        return class_type
+
+    def _process_non_send_request_types(self, class_type, kwargs):
+        outs = []
+        for slot_index, _ in enumerate(self.RETURN_TYPES):
+            node = BizyAirNodeIO(node_id=self.assigned_id, nodes={})
+            node.add_node_data(
+                class_type=class_type, inputs=kwargs, outputs={"slot_index": slot_index}
+            )
+            outs.append(node)
+        return tuple(outs)
+
+    def _process_all_send_request_types(self, node_ios: List[BizyAirNodeIO]):
+        out = node_ios[0].send_request()
+        assert len(out) == len(self.RETURN_TYPES)
+        return out

bizyengine/core/nodes_io.py

@@ -0,0 +1,97 @@
+import warnings
+from typing import Any, Dict
+
+from .commands import invoker
+from .image_utils import encode_data
+
+
+def create_node_data(class_type: str, inputs: dict, outputs: dict):
+    assert (
+        outputs.get("slot_index", None) is not None
+    ), "outputs must contain 'slot_index'"
+    assert isinstance(outputs["slot_index"], int), "'slot_index' must be an integer"
+    assert isinstance(class_type, str)
+
+    out = {
+        "class_type": class_type,
+        "inputs": inputs,
+        "outputs": outputs,
+    }
+
+    return out
+
+
+class BizyAirNodeIO:
+    def __init__(
+        self,
+        node_id: int = "0",  # Unique identifier for the current node
+        nodes: Dict[str, Dict[str, any]] = {},
+        *args,
+        **kwargs,
+    ):
+        self._validate_node_id(node_id=node_id)
+        self.node_id = node_id
+        self.nodes = nodes
+
+    def _validate_node_id(self, node_id) -> bool:
+        if node_id is None:
+            raise ValueError("Node ID cannot be None.")
+        if not isinstance(node_id, str):
+            raise ValueError("Node ID must be a string.")
+        if not node_id.isdigit():
+            raise ValueError(
+                "Node ID must be a string that can be converted to an integer."
+            )
+        return True
+
+    def copy(self, new_node_id: str = None):
+        self._validate_node_id(new_node_id)
+        if new_node_id in self.nodes:
+            raise ValueError(f"Node ID '{new_node_id}' already exists.")
+
+        return BizyAirNodeIO(
+            nodes=self.nodes.copy(),
+            node_id=new_node_id,
+        )
+
+    def add_node_data(
+        self,
+        class_type: str,
+        inputs: Dict[str, Any],
+        outputs: Dict[str, Any] = {"slot_index": 0},
+    ):
+        node_data = create_node_data(
+            class_type=class_type,
+            inputs=inputs,
+            outputs=outputs,
+        )
+
+        self.update_nodes_from_others(*inputs.values())
+
+        if self.node_id in self.nodes:
+            warnings.warn(
+                f"Node ID {self.node_id} already exists. Data will be overwritten.",
+                RuntimeWarning,
+            )
+
+        self.nodes[self.node_id] = node_data
+
+    def update_nodes_from_others(self, *others):
+        for other in others:
+            if isinstance(other, BizyAirNodeIO):
+                self.nodes.update(other.nodes)
+
+    def send_request(
+        self, url=None, headers=None, *, progress_callback=None, stream=False
+    ) -> any:
+        out = invoker.prompt_server.execute(
+            prompt=self.nodes, last_node_ids=[self.node_id]
+        )
+        return out
+
+
+@encode_data.register(BizyAirNodeIO)
+def _(output: BizyAirNodeIO, **kwargs):
+    origin_id = output.node_id
+    origin_slot = output.nodes[origin_id]["outputs"]["slot_index"]
+    return [origin_id, origin_slot]

bizyengine/core/path_utils/__init__.py

@@ -0,0 +1,9 @@
+from .path_manager import (
+    convert_prompt_label_path_to_real_path,
+    disable_refresh_options,
+    enable_refresh_options,
+    get_filename_list,
+    guess_config,
+    guess_url_from_node,
+)
+from .utils import filter_files_extensions

bizyengine/core/path_utils/path_manager.py

@@ -0,0 +1,276 @@
+import copy
+import json
+import os
+import pprint
+import re
+import warnings
+from collections import defaultdict
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Any, Collection, Dict, List, Union
+
+from bizyengine.core.common import fetch_models_by_type
+from bizyengine.core.common.env_var import BIZYAIR_DEBUG, BIZYAIR_SERVER_ADDRESS
+from bizyengine.core.configs.conf import ModelRule, config_manager
+from bizyengine.core.path_utils.utils import (
+    filter_files_extensions,
+    get_service_route,
+    load_yaml_config,
+)
+
+supported_pt_extensions: set[str] = {
+    ".ckpt",
+    ".pt",
+    ".bin",
+    ".pth",
+    ".safetensors",
+    ".pkl",
+    ".sft",
+}
+ScanPathType = list[str]
+folder_names_and_paths: dict[str, ScanPathType] = defaultdict(list)
+filename_path_mapping: dict[str, dict[str, str]] = {}
+
+
+@dataclass
+class RefreshSettings:
+    loras: bool = True
+    controlnet: bool = True
+
+    def get(self, folder_name: str, default: bool = True):
+        return getattr(self, folder_name, default)
+
+    def set(self, folder_name: str, value: bool):
+        setattr(self, folder_name, value)
+
+
+refresh_settings = RefreshSettings()
+
+
+def enable_refresh_options(folder_names: Union[str, list[str]]):
+    if isinstance(folder_names, str):
+        folder_names = [folder_names]
+    for folder_name in folder_names:
+        refresh_settings.set(folder_name, True)
+
+
+def disable_refresh_options(folder_names: Union[str, list[str]]):
+    if isinstance(folder_names, str):
+        folder_names = [folder_names]
+    for folder_name in folder_names:
+        refresh_settings.set(folder_name, False)
+
+
+def _get_config_path():
+    src_bizyair_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    configs_path = os.path.join(src_bizyair_path, "configs")
+    return configs_path
+
+
+configs_path = _get_config_path()
+
+models_config: Dict[str, Dict[str, Any]] = load_yaml_config(
+    os.path.join(configs_path, "models.yaml")
+)
+
+
+def guess_url_from_node(
+    node: Dict[str, Dict[str, Any]], class_type_table: Dict[str, bool]
+) -> Union[List[ModelRule], None]:
+    rules: List[ModelRule] = config_manager.get_rules(node["class_type"])
+    out = [
+        rule
+        for rule in rules
+        if len(rule.inputs) == 0
+        or all(
+            any(re.search(p, node["inputs"][key]) is not None for p in patterns)
+            for key, patterns in rule.inputs.items()
+        )
+    ]
+    return out
+
+
+def guess_config(
+    *,
+    ckpt_name: str = None,
+    unet_name: str = None,
+    vae_name: str = None,
+    clip_name: str = None,
+) -> str:
+    warnings.warn("The interface has changed, please do not use it", DeprecationWarning)
+
+
+def get_config_file_list(base_path=None) -> list:
+    if base_path is None:
+        base_path = os.path.dirname(os.path.abspath(__file__))
+    path = os.path.join(configs_path)
+    extensions = ".yaml"
+    config_files = []
+    for root, _, files in os.walk(path):
+        for file in files:
+            if file.endswith(extensions):
+                file_path = os.path.join(root, file)
+                config_files.append(file_path)
+    return config_files
+
+
+def cached_filename_list(
+    folder_name: str, *, share_id: str = None, verbose=False, refresh=False
+) -> list[str]:
+    global filename_path_mapping
+    if refresh or folder_name not in filename_path_mapping:
+        model_types: Dict[str, str] = models_config["model_types"]
+        if share_id:
+            url = f"{BIZYAIR_SERVER_ADDRESS}/{share_id}/models/files"
+        else:
+            url = get_service_route(models_config["model_hub"]["find_model"])
+        msg = fetch_models_by_type(
+            url=url, method="GET", model_type=model_types[folder_name]
+        )
+        if verbose:
+            pprint.pprint({"cached_filename_list": msg})
+
+        try:
+            if not msg or "data" not in msg or msg["data"] is None:
+                return []
+
+            filename_path_mapping[folder_name] = {
+                x["label_path"]: x["real_path"]
+                for x in msg["data"]["files"]
+                if x["label_path"]
+            }
+        except Exception as e:
+            warnings.warn(f"Failed to get filename list: {e}")
+            return []
+        finally:
+            # TODO fix share_id vaild refresh settings
+            if share_id is None:
+                disable_refresh_options(folder_name)
+
+    return list(
+        filter_files_extensions(
+            filename_path_mapping[folder_name].keys(),
+            extensions=supported_pt_extensions,
+        )
+    )
+
+
+def convert_prompt_label_path_to_real_path(prompt: dict[str, dict[str, any]]) -> dict:
+    # TODO fix Temporarily write dead
+    new_prompt = {}
+    for unique_id in prompt:
+        new_prompt[unique_id] = copy.copy(prompt[unique_id])
+        inputs = copy.copy(prompt[unique_id]["inputs"])
+
+        for key, folder_name in [
+            ("lora_name", "loras"),
+            ("control_net_name", "controlnet"),
+        ]:
+            if key in inputs:
+                value = inputs[key]
+                new_value = filename_path_mapping.get(folder_name, {}).get(value, None)
+                if new_value:
+                    inputs[key] = new_value
+                else:
+                    file_list = get_filename_list(folder_name)
+                    if value not in file_list:
+                        raise ValueError(
+                            f"{key} '{value}' not found in file list. Available {key} names: {', '.join(file_list)}"
+                        )
+
+        new_prompt[unique_id]["inputs"] = inputs
+    return new_prompt
+
+
+def get_share_filename_list(folder_name, share_id, *, verbose=BIZYAIR_DEBUG):
+    assert share_id is not None and isinstance(share_id, str)
+    # TODO fix share_id vaild refresh settings
+    return cached_filename_list(
+        folder_name, share_id=share_id, verbose=verbose, refresh=True
+    )
+
+
+def get_filename_list(folder_name, *, verbose=BIZYAIR_DEBUG):
+
+    global folder_names_and_paths
+    results = folder_names_and_paths.get(folder_name, [])
+    # 社区node上线后移除
+    # if folder_name in models_config["model_types"]:
+    #     refresh = refresh_settings.get(folder_name, True)
+    #     results.extend(
+    #         cached_filename_list(folder_name, verbose=verbose, refresh=refresh)
+    #     )
+    # if folder_name in folder_names_and_paths:
+    #     results.extend(folder_names_and_paths[folder_name])
+    # if BIZYAIR_DEBUG:
+    #     try:
+    #         import folder_paths
+    #
+    #         results.extend(folder_paths.get_filename_list(folder_name))
+    #     except:
+    #         pass
+    return results
+
+
+def filter_files_extensions(
+    files: Collection[str], extensions: Collection[str]
+) -> list[str]:
+    return sorted(
+        list(
+            filter(
+                lambda a: os.path.splitext(a)[-1].lower() in extensions
+                or len(extensions) == 0,
+                files,
+            )
+        )
+    )
+
+
+def recursive_extract_models(data: Any, prefix_path: str = "") -> List[str]:
+    def merge_paths(base_path: str, new_path: Any) -> str:
+        if not isinstance(new_path, str):
+            return base_path
+        return f"{base_path}/{new_path}" if base_path else new_path
+
+    results: List[str] = []
+    if isinstance(data, dict):
+        for key, value in data.items():
+            new_prefix = merge_paths(prefix_path, key)
+            results.extend(recursive_extract_models(value, new_prefix))
+    elif isinstance(data, list):
+        for item in data:
+            new_prefix = merge_paths(prefix_path, item)
+            results.extend(recursive_extract_models(item, new_prefix))
+    elif isinstance(data, str) and prefix_path.endswith(data):
+        return filter_files_extensions([prefix_path], supported_pt_extensions)
+
+    return results
+
+
+def load_json(file_path: str) -> dict:
+    with open(file_path, "r") as file:
+        data = json.load(file)
+    return data
+
+
+def init_config():
+    global folder_names_and_paths
+    for k, filenames in config_manager.model_path_manager.model_paths.items():
+        folder_names_and_paths[k].extend(filenames)
+    if BIZYAIR_DEBUG:
+        pprint.pprint("=" * 20 + "init_config: " + "=" * 20)
+        pprint.pprint(folder_names_and_paths)
+
+
+init_config()
+
+
+if __name__ == "__main__":
+    # print(f"Loaded config from {get_config_file_list()}")
+    # configs = [load_yaml_config(x) for x in get_config_file_list()]
+    # print(get_filename_list("clip_vision"))
+    # print(folder_names_and_paths)
+
+    api_key = os.getenv("BIZYAIR_API_KEY", "")
+    host_ckpts = get_filename_list("loras", verbose=True)
+    print(host_ckpts)

bizyengine/core/path_utils/utils.py

@@ -0,0 +1,34 @@
+import os
+from collections.abc import Collection
+from typing import Dict, Union
+
+import yaml
+from bizyengine.core.common.env_var import BIZYAIR_SERVER_ADDRESS
+
+
+def filter_files_extensions(
+    files: Collection[str], extensions: Collection[str]
+) -> list[str]:
+    return sorted(
+        list(
+            filter(
+                lambda a: os.path.splitext(a)[-1].lower() in extensions
+                or len(extensions) == 0,
+                files,
+            )
+        )
+    )
+
+
+def load_yaml_config(file_path):
+    with open(file_path, "r") as file:
+        config = yaml.safe_load(file)
+    return config
+
+
+def get_service_route(service_config: Dict[str, str]) -> Union[str, None]:
+    if {"route"}.issubset(service_config):
+        return str(
+            service_config.get("service_address", BIZYAIR_SERVER_ADDRESS)
+        ) + service_config.get("route")
+    return None

bizyengine/misc/auth.py

@@ -0,0 +1,83 @@
+import os
+import uuid
+from pathlib import Path
+
+import bizyengine.core
+import server
+from aiohttp import web
+from bizyengine.core.common import create_api_key_file, load_api_key, validate_api_key
+
+API_KEY = None
+# html_file_path = Path(os.path.dirname(os.path.abspath(__file__))) / "set_api_key.html"
+# with open(html_file_path, "r", encoding="utf-8") as htmlfile:
+#     set_api_key_html = htmlfile.read()
+
+
+has_key, api_key = load_api_key()
+if has_key:
+    API_KEY = api_key
+
+
+# @server.PromptServer.instance.routes.get("/bizyair/set-api-key")
+# async def set_api_key_page(request):
+#     return web.Response(text=set_api_key_html, content_type="text/html")
+
+
+@server.PromptServer.instance.routes.post("/bizyair/set_api_key")
+async def set_api_key(request):
+    global API_KEY
+    try:
+        data = await request.post()
+        api_key = data.get("api_key")
+        if api_key:
+            if not validate_api_key(api_key):
+                error_msg = "Wrong API key provided, please refer to cloud.siliconflow.cn to get the key"
+                print("set_api_key:", error_msg)
+                return web.Response(
+                    text=error_msg,
+                    status=400,
+                )
+            create_api_key_file(api_key)
+            API_KEY = api_key
+            bizyengine.core.set_api_key(API_KEY, override=True)
+            print("Set the key sucessfully.")
+            return web.Response(text="ok")
+        else:
+            error_msg = "No API key provided, please refer to cloud.siliconflow.cn to get the key"
+            print("set_api_key:", error_msg)
+            return web.Response(
+                text=error_msg,
+                status=400,
+            )
+    except Exception as e:
+        print(f"set api key error: {str(e)}")
+        return web.Response(text=str(e), status=500)
+
+
+@server.PromptServer.instance.routes.get("/bizyair/get_api_key")
+async def get_api_key(request):
+    global API_KEY
+    try:
+        has_key, api_key = load_api_key()
+        if has_key:
+            API_KEY = api_key
+            bizyengine.core.set_api_key(API_KEY)
+            return web.Response(text="Key has been loaded from the api_key.ini file")
+        else:
+            api_key = request.cookies.get("api_key")
+            if not api_key:
+                print("No api key found in cookies")
+                return web.Response(
+                    text="No api key found in cookies, please refer to cloud.siliconflow.cn to get the key",
+                    status=404,
+                )
+            API_KEY = api_key
+            bizyengine.core.set_api_key(API_KEY)
+            return web.Response(text="Key has been loaded from the cookies")
+
+    except Exception as e:
+        return web.Response(text=str(e), status=500)
+
+
+NODE_CLASS_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS = {}