bizyengine 0.4.2__py3-none-any.whl

bizyengine/core/common/client.py

@@ -0,0 +1,262 @@
+import json
+import os
+import pprint
+import urllib.error
+import urllib.request
+import warnings
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import Any, Union
+
+import aiohttp
+import comfy
+
+from .caching import CacheManager
+
+__all__ = ["send_request"]
+
+from dataclasses import dataclass, field
+
+from .env_var import BIZYAIR_API_KEY, BIZYAIR_DEBUG, BIZYAIR_SERVER_ADDRESS
+
+IS_API_KEY_VALID = None
+
+
+@dataclass
+class APIKeyState:
+    current_api_key: str = field(default=None)
+    is_valid: bool = field(default=None)
+
+
+api_key_state = APIKeyState()
+
+
+def set_api_key(api_key: str = "YOUR_API_KEY", override: bool = False):
+    global BIZYAIR_API_KEY, api_key_state
+    if api_key_state.is_valid is not None and not override:
+        warnings.warn("API key has already been set and will not be overridden.")
+        return
+    if validate_api_key(api_key):
+        BIZYAIR_API_KEY = api_key
+        api_key_state.is_valid = True
+        print("\033[92mAPI key is set successfully.\033[0m")
+    else:
+        api_key_state.is_valid = False
+        warnings.warn("Invalid API key provided.")
+
+
+def validate_api_key(api_key: str = None) -> bool:
+    global api_key_state
+    if not api_key or not isinstance(api_key, str):
+        warnings.warn("API key is not set.")
+        return False
+    # if api_key_state.current_api_key == api_key and api_key_state.is_valid is not None:
+    #     return api_key_state.is_valid
+    api_key_state.current_api_key = api_key
+    url = f"{BIZYAIR_SERVER_ADDRESS}/user/info"
+    headers = {"accept": "application/json", "authorization": f"Bearer {api_key}"}
+    try:
+        response_data = send_request(
+            method="GET", url=url, headers=headers, callback=None
+        )
+        if "message" not in response_data or response_data["message"] != "Ok":
+            api_key_state.is_valid = False
+            raise ValueError(
+                f"\033[91mAPI key validation failed. API Key: {api_key}\033[0m"
+            )
+        else:
+            api_key_state.is_valid = True
+    except ConnectionError as ce:
+        api_key_state.is_valid = False
+        raise ValueError(f"\033[91mConnection error: {ce}\033[0m")
+    except PermissionError as pe:
+        api_key_state.is_valid = False
+        raise ValueError(
+            f"\033[91mError validating API key: {api_key}, error: {pe}\033[0m"
+        )
+    except Exception as e:
+        api_key_state.is_valid = False
+        raise ValueError(f"\033[91mOther error: {e}\033[0m")
+    return api_key_state.is_valid
+
+
+def get_api_key() -> str:
+    global BIZYAIR_API_KEY
+    try:
+        validate_api_key(BIZYAIR_API_KEY)
+    except Exception as e:
+        print(str(e))
+        raise ValueError(str(e))
+    return BIZYAIR_API_KEY
+
+
+def _headers():
+    headers = {
+        "accept": "application/json",
+        "content-type": "application/json",
+        "authorization": f"Bearer {get_api_key()}",
+    }
+    return headers
+
+
+def process_response_data(response_data: dict) -> dict:
+    # Check if 'result' key exists, indicating a cloud response
+    if "result" in response_data:
+        try:
+            msg = json.loads(response_data["result"])
+        except json.JSONDecodeError:
+            raise ValueError(f"Failed to decode JSON from response. {response_data=}")
+    else:
+        # Handle local response directly
+        msg = response_data
+
+    return msg  # Return processed data or modify as needed
+
+
+def send_request(
+    method: str = "POST",
+    url: str = None,
+    data: bytes = None,
+    verbose=False,
+    callback: callable = process_response_data,
+    response_handler: callable = json.loads,
+    cache_manager: CacheManager = None,
+    **kwargs,
+) -> Union[dict, Any]:
+    try:
+        headers = kwargs.pop("headers") if "headers" in kwargs else _headers()
+        headers["User-Agent"] = "BizyAir Client"
+
+        req = urllib.request.Request(
+            url, data=data, headers=headers, method=method, **kwargs
+        )
+        with urllib.request.urlopen(req) as response:
+            response_data = response.read().decode("utf-8")
+    except urllib.error.URLError as e:
+        error_message = str(e)
+        response_body = e.read().decode("utf-8") if hasattr(e, "read") else "N/A"
+        if verbose:
+            print(f"URLError encountered: {error_message}")
+            print(f"Response Body: {response_data}")
+        code, message = "N/A", "N/A"
+        try:
+            response_dict = json.loads(response_body)
+            if isinstance(response_dict, dict):
+                code = response_dict.get("code", "N/A")
+                message = response_dict.get("message", "N/A")
+
+        except json.JSONDecodeError:
+            if verbose:
+                print("Failed to decode response body as JSON.")
+
+        if "Unauthorized" in error_message:
+            raise PermissionError(
+                "Key is invalid, please refer to https://cloud.siliconflow.cn to get the API key.\n"
+                "If you have the key, please click the 'API Key' button at the bottom right to set the key."
+            )
+        elif code != "N/A" and message != "N/A":
+            raise ConnectionError(
+                f"Failed to handle your request: {error_message}.\n"
+                + f"    Error code: {code}.\n"
+                + f"    Error message: {message}\n."
+                + "The cause of this issue may be incorrect parameter status or ongoing background tasks. \n"
+                + "If retrying after waiting for a while still does not resolve the issue, please report it to "
+                + "Bizyair's official support."
+            )
+        else:
+            common_sites = [
+                "https://www.baidu.com",
+                "https://www.bing.com",
+                "https://www.alibaba.com",
+            ]
+            results = {}
+            for site in common_sites:
+                success = False
+                try:
+
+                    class NoRedirectHandler(urllib.request.HTTPRedirectHandler):
+                        def redirect_request(self, req, fp, code, msg, headers, newurl):
+                            return None
+
+                    opener = urllib.request.build_opener(NoRedirectHandler())
+                    response = opener.open(site, timeout=5)
+                    success = 200 <= response.getcode() < 400
+                except urllib.error.HTTPError as e:
+                    success = 200 <= e.code < 400
+                except (urllib.error.URLError, TimeoutError):
+                    pass
+                results[site] = "Success" if success else "Failed"
+            raise ConnectionError(
+                f"Failed to connect to the server: {url}.\n"
+                + "The connection attempts to the public sites return the following results:\n"
+                + "\n".join(
+                    [f"    {site}: {result}" for site, result in results.items()]
+                )
+                + "\nPlease check the network connection."
+            )
+
+    if response_handler:
+        response_data = response_handler(response_data)
+    if callback:
+        return callback(response_data)
+    return response_data
+
+
+async def async_send_request(
+    method: str = "POST",
+    url: str = None,
+    data: bytes = None,
+    verbose=False,
+    callback: callable = process_response_data,
+    **kwargs,
+) -> dict:
+    headers = kwargs.pop("headers") if "headers" in kwargs else _headers()
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.request(
+                method, url, data=data, headers=headers, **kwargs
+            ) as response:
+                response_data = await response.text()
+                if response.status != 200:
+                    error_message = f"HTTP Status {response.status}"
+                    if verbose:
+                        print(f"Error encountered: {error_message}")
+                    if response.status == 401:
+                        raise PermissionError(
+                            "Key is invalid, please refer to https://cloud.siliconflow.cn to get the API key.\n"
+                            "If you have the key, please click the 'BizyAir Key' button at the bottom right to set the key."
+                        )
+                    else:
+                        raise ConnectionError(
+                            f"Failed to connect to the server: {error_message}.\n"
+                            + "Please check your API key and ensure the server is reachable.\n"
+                            + "Also, verify your network settings and disable any proxies if necessary.\n"
+                            + "After checking, please restart the ComfyUI service."
+                        )
+                if callback:
+                    return callback(json.loads(response_data))
+                return json.loads(response_data)
+    except aiohttp.ClientError as e:
+        print(f"Error fetching data: {e}")
+        return {}
+    except Exception as e:
+        print(f"Error fetching data: {str(e)}")
+        return {}
+
+
+def fetch_models_by_type(
+    url: str, model_type: str, *, method="GET", verbose=False
+) -> dict:
+    if not validate_api_key(BIZYAIR_API_KEY):
+        return {}
+
+    payload = {"type": model_type}
+    if BIZYAIR_DEBUG:
+        pprint.pprint(payload)
+    msg = send_request(
+        method=method,
+        url=url,
+        data=json.dumps(payload).encode("utf-8"),
+        verbose=verbose,
+    )
+    return msg

bizyengine/core/common/env_var.py

@@ -0,0 +1,101 @@
+import configparser
+import os
+from os import environ
+from pathlib import Path
+
+BIZYAIR_COMFYUI_PATH = Path(os.environ.get("BIZYAIR_COMFYUI_PATH", "./"))
+print(f"\033[92m[BizyAir]\033[0m BizyAir ComfyUI Plugin: {str(BIZYAIR_COMFYUI_PATH)}")
+
+
+class ServerAddress:
+    _instance = None
+    _address = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super(ServerAddress, cls).__new__(cls)
+        return cls._instance
+
+    def __init__(self, address):
+        self._address = address
+
+    @property
+    def address(self):
+        return self._address
+
+    @address.setter
+    def address(self, new_address):
+        self._address = new_address
+
+    def __str__(self):
+        return self._address
+
+
+def env(key, type_, default=None):
+    if key not in environ:
+        return default
+
+    val = environ[key]
+
+    if type_ == str:
+        return val
+    if type_ == bool:
+        if val.lower() in ["1", "true", "yes", "y", "ok", "on"]:
+            return True
+        if val.lower() in ["0", "false", "no", "n", "nok", "off"]:
+            return False
+        raise ValueError(
+            "Invalid environment variable '%s' (expected a boolean): '%s'" % (key, val)
+        )
+    if type_ == int:
+        try:
+            return int(val)
+        except ValueError:
+            raise ValueError(
+                "Invalid environment variable '%s' (expected an integer): '%s'"
+                % (key, val)
+            ) from None
+    raise ValueError("The requested type '%r' is not supported" % type_)
+
+
+def load_api_key():
+
+    file_path = BIZYAIR_COMFYUI_PATH / "api_key.ini"
+
+    if file_path.is_file() and file_path.exists():
+        config = configparser.ConfigParser()
+        config.read(file_path)
+        api_key: str = config.get("auth", "api_key", fallback="").strip().strip("'\"")
+        has_key = api_key.startswith("sk-")
+        return has_key, api_key
+    else:
+        return False, None
+
+
+def create_api_key_file(api_key):
+    config = configparser.ConfigParser()
+    config["auth"] = {"api_key": api_key}
+    file_path = BIZYAIR_COMFYUI_PATH / "api_key.ini"
+    try:
+        with open(file_path, "w", encoding="utf-8") as configfile:
+            config.write(configfile)
+    except Exception as e:
+        raise Exception(f"An error occurred when save the key: {e}")
+
+
+#   production:
+#     service_address: https://bizyair-api.siliconflow.cn/x/v1
+#   uat:
+#     service_address: https://uat-bizyair-api.siliconflow.cn/x/v1
+_BIZYAIR_SERVER_ADDRESS = os.getenv(
+    "BIZYAIR_SERVER_ADDRESS", "https://bizyair-api.siliconflow.cn/x/v1"
+)
+BIZYAIR_SERVER_ADDRESS = ServerAddress(_BIZYAIR_SERVER_ADDRESS)
+BIZYAIR_API_KEY = env("BIZYAIR_API_KEY", str, load_api_key()[1])
+# Development Settings
+BIZYAIR_DEV_REQUEST_URL = env("BIZYAIR_DEV_REQUEST_URL", str, None)
+BIZYAIR_DEBUG = env("BIZYAIR_DEBUG", bool, False)
+BIZYAIR_DEV_GET_TASK_RESULT_SERVER = env(
+    "BIZYAIR_DEV_GET_TASK_RESULT_SERVER", str, None
+)
+BIZYAIR_PRODUCTION_TEST = env("BIZYAIR_PRODUCTION_TEST", str, None)

bizyengine/core/common/utils.py

@@ -0,0 +1,93 @@
+import copy
+import json
+import os
+from typing import Any, List
+
+import torch
+import yaml
+
+
+def truncate_long_strings(obj, max_length=50):
+    if isinstance(obj, str):
+        return obj if len(obj) <= max_length else obj[:max_length] + "..."
+    elif isinstance(obj, dict):
+        return {k: truncate_long_strings(v, max_length) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [truncate_long_strings(v, max_length) for v in obj]
+    elif isinstance(obj, tuple):
+        return tuple(truncate_long_strings(v, max_length) for v in obj)
+    elif isinstance(obj, torch.Tensor):
+        return obj.shape, obj.dtype, obj.device
+    else:
+        return obj
+
+
+def deepcopy_except_tensor(obj, exclude_types=[torch.Tensor]):
+    return deepcopy_except_types(obj=obj, exclude_types=exclude_types)
+
+
+def deepcopy_except_types(obj, exclude_types):
+    """
+    Recursively copy an object, excluding specified data types.
+
+    :param obj: The object to be copied
+    :param exclude_types: A list of data types to be excluded from deep copying
+    :return: The copied object
+    """
+    if any(isinstance(obj, t) for t in exclude_types):
+        return obj  # Return the object directly without deep copying
+    elif isinstance(obj, (list, tuple)):
+        return type(obj)(deepcopy_except_types(item, exclude_types) for item in obj)
+    elif isinstance(obj, dict):
+        return {
+            deepcopy_except_types(key, exclude_types): deepcopy_except_types(
+                value, exclude_types
+            )
+            for key, value in obj.items()
+        }
+    else:
+        return copy.deepcopy(obj)
+
+
+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 [prefix_path]
+
+    return results
+
+
+def _load_yaml_config(file_path):
+    with open(file_path, "r") as file:
+        config = yaml.safe_load(file)
+    return config
+
+
+def _load_json_config(file_path: str) -> dict:
+    with open(file_path, "r") as file:
+        data = json.load(file)
+    return data
+
+
+def load_config_file(file_path: str) -> dict:
+    if not os.path.exists(file_path):
+        raise FileNotFoundError(f"The file {file_path} does not exist.")
+    if file_path.endswith(".json"):
+        return _load_json_config(file_path)
+    elif file_path.endswith(".yaml"):
+        return _load_yaml_config(file_path)
+    else:
+        raise ValueError(f"Unsupported file extension: {file_path}")

bizyengine/core/configs/conf.py

@@ -0,0 +1,112 @@
+import os
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import List, Tuple
+
+from bizyengine.core.common.utils import load_config_file, recursive_extract_models
+
+
+@dataclass
+class ModelRule:
+    mode_type: str
+    base_model: str
+    describe: str
+    score: int
+    route: str
+    class_type: str
+    inputs: dict
+
+
+@dataclass
+class TaskApi:
+    task_result_endpoint: str
+
+
+class ModelRuleManager:
+    def __init__(self, model_rules: list[dict]):
+        self.model_rules = model_rules
+        self.validate()
+        self.gen_model_rule_index_mapping()
+
+    def gen_model_rule_index_mapping(self):
+        self.model_rule_index_mapping = defaultdict(list)
+
+        for idx_1, rule in enumerate(self.model_rules):
+            for idx_2, node in enumerate(rule["nodes"]):
+                self.model_rule_index_mapping[node["class_type"]].append((idx_1, idx_2))
+
+    def validate(self):
+        for rule in self.model_rules:
+            self._validate_rule(rule)
+
+    def _validate_rule(self, rule: dict):
+        if "mode_type" not in rule:
+            raise ValueError("mode_type is required")
+        if "base_model" not in rule:
+            raise ValueError("base_model is required")
+        if "route" not in rule:
+            raise ValueError("route is required")
+        if "nodes" not in rule:
+            raise ValueError("nodes is required")
+
+    def find_rule_indexes(self, class_type: str) -> List[Tuple[int, int]]:
+        return self.model_rule_index_mapping[class_type]
+
+    def find_rules(self, class_type: str) -> List[ModelRule]:
+        rule_indexes = self.find_rule_indexes(class_type)
+        return [
+            ModelRule(
+                mode_type=self.model_rules[idx_1]["mode_type"],
+                base_model=self.model_rules[idx_1]["base_model"],
+                describe=self.model_rules[idx_1]["describe"],
+                score=self.model_rules[idx_1]["score"],
+                route=self.model_rules[idx_1]["route"],
+                class_type=class_type,
+                inputs=self.model_rules[idx_1]["nodes"][idx_2].get("inputs", {}),
+            )
+            for idx_1, idx_2 in rule_indexes
+        ]
+
+
+class ModelPathManager:
+    def __init__(self, config_path: str):
+        model_paths = {}
+        for folder_name, v in load_config_file(config_path).items():
+            model_paths[folder_name] = recursive_extract_models(v)
+        self.model_paths = model_paths
+
+    def get_filenames(self, folder_name: str) -> List[str]:
+        return self.model_paths.get(folder_name, [])
+
+
+class ConfigManager:
+    def __init__(self, model_path_config: str, model_rule_config: str):
+        self.model_path_manager = ModelPathManager(config_path=model_path_config)
+        self.model_rule_config = load_config_file(model_rule_config)
+        self.model_rules = ModelRuleManager(
+            model_rules=self.model_rule_config["model_rules"]
+        )
+
+    def get_filenames(self, folder_name: str) -> List[str]:
+        return self.model_path_manager.get_filenames(folder_name)
+
+    def get_rules(self, class_type: str) -> List[ModelRule]:
+        if class_type.startswith("BizyAir_"):
+            class_type = class_type[8:]
+        return self.model_rules.find_rules(class_type)
+
+    def get_model_version_id_prefix(self):
+        return self.model_rule_config["model_version_config"]["model_version_id_prefix"]
+
+    def get_cache_config(self):
+        return self.model_rule_config.get("cache_config", {})
+
+    def get_task_api(self):
+        return TaskApi(**self.model_rule_config["task_api"])
+
+
+model_path_config = os.path.join(os.path.dirname(__file__), "models.json")
+model_rule_config = os.path.join(os.path.dirname(__file__), "models.yaml")
+config_manager = ConfigManager(
+    model_path_config=model_path_config, model_rule_config=model_rule_config
+)

bizyengine/core/configs/models.json

@@ -0,0 +1,101 @@
+{
+  "checkpoints": [{
+    "sdxl": [
+      "counterfeitxl_v25.safetensors",
+      "dreamshaperXL_lightningDPMSDE.safetensors",
+      "dreamshaperXL_v21TurboDPMSDE.safetensors",
+      "HelloWorldXL_v70.safetensors",
+      "juggernautXL_v9Rdphoto2Lightning.safetensors",
+      "Juggernaut-XL_v9_RunDiffusionPhoto_v2.safetensors",
+      "Juggernaut_X_RunDiffusion_Hyper.safetensors",
+      "mannEDreams_v004.safetensors",
+      "realisticStockPhoto_v20.safetensors",
+      "samaritan3dCartoon_v40SDXL.safetensors"
+    ],
+    "sd15": [
+      "dreamshaper_8.safetensors"
+    ]
+  },
+  "sd3.5_large.safetensors", "sd3.5_large_turbo.safetensors"],
+  "clip_vision": [{
+    "models": [
+      "CLIP-ViT-H-14-laion2B-s32B-b79K.safetensors"
+    ],
+    "kolors": [
+      "pytorch_model.bin"
+    ]
+  }, "sigclip_vision_patch14_384.safetensors"],
+  "controlnet": [{
+    "kolors": [
+      "Kolors-ControlNet-Canny.safetensors",
+      "Kolors-ControlNet-Depth.safetensors"
+    ],
+    "sdxl": [
+      "diffusion_pytorch_model_promax.safetensors"
+    ],
+    "sd15": [
+      "control_v11f1e_sd15_tile.pth"
+    ],
+    "instantid": [
+      "diffusion_pytorch_model.safetensors"
+    ]
+  }, "sd3.5_large_controlnet_blur.safetensors", "sd3.5_large_controlnet_canny.safetensors", "sd3.5_large_controlnet_depth.safetensors"],
+  "ipadapter": {
+    "kolors": [
+      "ip_adapter_plus_general.bin"
+    ]
+  },
+  "loras": {
+    "sdxl": [
+      "Cute_Animals.safetensors",
+      "watercolor_v1_sdxl_lora.safetensors"
+    ],
+    "flux": [
+      "meijia_flux_lora_rank16_bf16.safetensors"
+    ]
+  },
+  "unet": ["shuttle-3.1-aesthetic.safetensors",{
+    "kolors": [
+      "Kolors-Inpainting.safetensors",
+      "Kolors.safetensors"
+    ],
+    "flux": [
+      "flux1-schnell.sft",
+      "flux1-dev.sft",
+      "pixelwave-flux1-dev.safetensors",
+      "flux1-canny-dev.safetensors",
+      "flux1-depth-dev.safetensors",
+      "flux1-fill-dev.safetensors"
+    ]
+  }],
+  "vae": [{
+    "sdxl": [
+      "sdxl_vae.safetensors"
+    ],
+    "flux": [
+      "ae.sft"
+    ]
+  }, "flux.1-canny-vae.safetensors", "flux.1-depth-vae.safetensors", "flux.1-fill-vae.safetensors"],
+  "clip": [
+    "clip_l.safetensors",
+    "clip_g.safetensors",
+    "t5xxl_fp16.safetensors",
+    "t5xxl_fp8_e4m3fn.safetensors"
+  ],
+  "upscale_models": [
+    "4x_NMKD-Siax_200k.pth",
+    "RealESRGAN_x4plus.pth",
+    "RealESRGAN_x4plus_anime_6B.pth",
+    "RealESRGAN_x2plus.pth",
+    "RealESRGAN_x4plus.pth"
+  ],
+  "instantid": [
+    "ip-adapter.bin"
+  ],
+  "pulid": [
+    "pulid_flux.safetensors"
+  ],
+  "style_models": [
+    "flux1-redux-dev.safetensors"
+  ]
+}