bizyengine 1.2.45__py3-none-any.whl → 1.2.71__py3-none-any.whl

bizyengine/core/image_utils.py

@@ -1,292 +1,21 @@
-import base64
-import io
-import json
-import os
-import pickle
-import zlib
 from enum import Enum
-from functools import singledispatch
-from typing import Any, List, Union
 
-import numpy as np
-import torch
-from PIL import Image
+from bizyairsdk import (
+    base64_to_tensor,
+    decode_base64_to_image,
+    decode_base64_to_np,
+    decode_comfy_image,
+    decode_data,
+    encode_comfy_image,
+    encode_data,
+    encode_image_to_base64,
+    numpy_to_base64,
+)
 
-from .common.env_var import BIZYAIR_DEBUG
-
-# Marker to identify base64-encoded tensors
-TENSOR_MARKER = "TENSOR:"
-IMAGE_MARKER = "IMAGE:"
+# from .common.env_var import BIZYAIR_DEBUG
 
 
 class TaskStatus(Enum):
     PENDING = "pending"
     PROCESSING = "processing"
     COMPLETED = "completed"
-
-
-def convert_image_to_rgb(image: Image.Image) -> Image.Image:
-    if image.mode != "RGB":
-        return image.convert("RGB")
-    return image
-
-
-def encode_image_to_base64(
-    image: Image.Image, format: str = "png", quality: int = 100, lossless=False
-) -> str:
-    image = convert_image_to_rgb(image)
-    with io.BytesIO() as output:
-        image.save(output, format=format, quality=quality, lossless=lossless)
-        output.seek(0)
-        img_bytes = output.getvalue()
-        if BIZYAIR_DEBUG:
-            print(f"encode_image_to_base64: {format_bytes(len(img_bytes))}")
-    return base64.b64encode(img_bytes).decode("utf-8")
-
-
-def decode_base64_to_np(img_data: str, format: str = "png") -> np.ndarray:
-    img_bytes = base64.b64decode(img_data)
-    if BIZYAIR_DEBUG:
-        print(f"decode_base64_to_np: {format_bytes(len(img_bytes))}")
-    with io.BytesIO(img_bytes) as input_buffer:
-        img = Image.open(input_buffer)
-        # https://github.com/comfyanonymous/ComfyUI/blob/a178e25912b01abf436eba1cfaab316ba02d272d/nodes.py#L1511
-        img = img.convert("RGB")
-        return np.array(img)
-
-
-def decode_base64_to_image(img_data: str) -> Image.Image:
-    img_bytes = base64.b64decode(img_data)
-    with io.BytesIO(img_bytes) as input_buffer:
-        img = Image.open(input_buffer)
-        if BIZYAIR_DEBUG:
-            format_info = img.format.upper() if img.format else "Unknown"
-            print(f"decode image format: {format_info}")
-        return img
-
-
-def format_bytes(num_bytes: int) -> str:
-    """
-    Converts a number of bytes to a human-readable string with units (B, KB, or MB).
-
-    :param num_bytes: The number of bytes to convert.
-    :return: A string representing the number of bytes in a human-readable format.
-    """
-    if num_bytes < 1024:
-        return f"{num_bytes} B"
-    elif num_bytes < 1024 * 1024:
-        return f"{num_bytes / 1024:.2f} KB"
-    else:
-        return f"{num_bytes / (1024 * 1024):.2f} MB"
-
-
-def _legacy_encode_comfy_image(image: torch.Tensor, image_format="png") -> str:
-    input_image = image.cpu().detach().numpy()
-    i = 255.0 * input_image[0]
-    input_image = np.clip(i, 0, 255).astype(np.uint8)
-    base64ed_image = encode_image_to_base64(
-        Image.fromarray(input_image), format=image_format
-    )
-    return base64ed_image
-
-
-def _legacy_decode_comfy_image(
-    img_data: Union[List, str], image_format="png"
-) -> torch.tensor:
-    if isinstance(img_data, List):
-        decoded_imgs = [decode_comfy_image(x, old_version=True) for x in img_data]
-
-        combined_imgs = torch.cat(decoded_imgs, dim=0)
-        return combined_imgs
-
-    out = decode_base64_to_np(img_data, format=image_format)
-    out = np.array(out).astype(np.float32) / 255.0
-    output = torch.from_numpy(out)[None,]
-    return output
-
-
-def _new_encode_comfy_image(images: torch.Tensor, image_format="WEBP", **kwargs) -> str:
-    """https://docs.comfy.org/essentials/custom_node_snippets#save-an-image-batch
-    Encode a batch of images to base64 strings.
-
-    Args:
-        images (torch.Tensor): A batch of images.
-        image_format (str, optional): The format of the images. Defaults to "WEBP".
-
-    Returns:
-        str: A JSON string containing the base64-encoded images.
-    """
-    results = {}
-    for batch_number, image in enumerate(images):
-        i = 255.0 * image.cpu().numpy()
-        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
-        base64ed_image = encode_image_to_base64(img, format=image_format, **kwargs)
-        results[batch_number] = base64ed_image
-
-    return json.dumps(results)
-
-
-def _new_decode_comfy_image(img_datas: str, image_format="WEBP") -> torch.tensor:
-    """
-    Decode a batch of base64-encoded images.
-
-    Args:
-        img_datas (str): A JSON string containing the base64-encoded images.
-        image_format (str, optional): The format of the images. Defaults to "WEBP".
-
-    Returns:
-        torch.Tensor: A tensor containing the decoded images.
-    """
-    img_datas = json.loads(img_datas)
-
-    decoded_imgs = []
-    for img_data in img_datas.values():
-        decoded_image = decode_base64_to_np(img_data, format=image_format)
-        decoded_image = np.array(decoded_image).astype(np.float32) / 255.0
-        decoded_imgs.append(torch.from_numpy(decoded_image)[None,])
-
-    return torch.cat(decoded_imgs, dim=0)
-
-
-def encode_comfy_image(
-    image: torch.Tensor, image_format="WEBP", old_version=False, lossless=False
-) -> str:
-    if old_version:
-        return _legacy_encode_comfy_image(image, image_format)
-    return _new_encode_comfy_image(image, image_format, lossless=lossless)
-
-
-def decode_comfy_image(
-    img_data: Union[List, str], image_format="WEBP", old_version=False
-) -> torch.tensor:
-    if old_version:
-        return _legacy_decode_comfy_image(img_data, image_format)
-    return _new_decode_comfy_image(img_data, image_format)
-
-
-def tensor_to_base64(tensor: torch.Tensor, compress=True) -> str:
-    tensor_np = tensor.cpu().detach().numpy()
-
-    tensor_bytes = pickle.dumps(tensor_np)
-    if compress:
-        tensor_bytes = zlib.compress(tensor_bytes)
-
-    tensor_b64 = base64.b64encode(tensor_bytes).decode("utf-8")
-    return tensor_b64
-
-
-def base64_to_tensor(tensor_b64: str, compress=True) -> torch.Tensor:
-    tensor_bytes = base64.b64decode(tensor_b64)
-
-    if compress:
-        tensor_bytes = zlib.decompress(tensor_bytes)
-
-    tensor_np = pickle.loads(tensor_bytes)
-
-    tensor = torch.from_numpy(tensor_np)
-    return tensor
-
-
-@singledispatch
-def decode_data(input, old_version=False):
-    raise NotImplementedError(f"Unsupported type: {type(input)}")
-
-
-@decode_data.register(int)
-@decode_data.register(float)
-@decode_data.register(bool)
-def _(input, **kwargs):
-    return input
-
-
-@decode_data.register(type(None))
-def _(input, **kwargs):
-    return [None]
-
-
-@decode_data.register(dict)
-def _(input, **kwargs):
-    return {k: decode_data(v, **kwargs) for k, v in input.items()}
-
-
-@decode_data.register(list)
-def _(input, **kwargs):
-    return [decode_data(x, **kwargs) for x in input]
-
-
-@decode_data.register(str)
-def _(input: str, **kwargs):
-    if input.startswith(TENSOR_MARKER):
-        tensor_b64 = input[len(TENSOR_MARKER) :]
-        return base64_to_tensor(tensor_b64)
-    elif input.startswith(IMAGE_MARKER):
-        tensor_b64 = input[len(IMAGE_MARKER) :]
-        old_version = kwargs.get("old_version", False)
-        return decode_comfy_image(tensor_b64, old_version=old_version)
-    return input
-
-
-@singledispatch
-def encode_data(output, disable_image_marker=False, old_version=False):
-    raise NotImplementedError(f"Unsupported type: {type(output)}")
-
-
-@encode_data.register(dict)
-def _(output, **kwargs):
-    return {k: encode_data(v, **kwargs) for k, v in output.items()}
-
-
-@encode_data.register(list)
-def _(output, **kwargs):
-    return [encode_data(x, **kwargs) for x in output]
-
-
-def is_image_tensor(tensor) -> bool:
-    """https://docs.comfy.org/essentials/custom_node_datatypes#image
-
-    Check if the given tensor is in the format of an IMAGE (shape [B, H, W, C] where C=3).
-
-    `Args`:
-        tensor (torch.Tensor): The tensor to check.
-
-    `Returns`:
-        bool: True if the tensor is in the IMAGE format, False otherwise.
-    """
-    try:
-        if not isinstance(tensor, torch.Tensor):
-            return False
-
-        if len(tensor.shape) != 4:
-            return False
-
-        B, H, W, C = tensor.shape
-        if C != 3:
-            return False
-
-        return True
-    except:
-        return False
-
-
-@encode_data.register(torch.Tensor)
-def _(output, **kwargs):
-    if is_image_tensor(output) and not kwargs.get("disable_image_marker", False):
-        old_version = kwargs.get("old_version", False)
-        lossless = kwargs.get("lossless", True)
-        return IMAGE_MARKER + encode_comfy_image(
-            output, image_format="WEBP", old_version=old_version, lossless=lossless
-        )
-    return TENSOR_MARKER + tensor_to_base64(output)
-
-
-@encode_data.register(int)
-@encode_data.register(float)
-@encode_data.register(bool)
-@encode_data.register(type(None))
-def _(output, **kwargs):
-    return output
-
-
-@encode_data.register(str)
-def _(output, **kwargs):
-    return output

bizyengine/misc/llm.py

@@ -381,22 +381,39 @@ class BizyAirJoyCaption2(BizyAirMiscBaseNode):
         )
 
         try:
-            if "result" in ret:
-                ret = json.loads(ret["result"])
-            if ret["type"] == "error":
-                raise Exception(ret["message"])
+            if isinstance(ret, dict) and "result" in ret:
+                if isinstance(ret["result"], str):
+                    ret = json.loads(ret["result"])
+
+            if not isinstance(ret, dict):
+                raise ValueError(
+                    f"Response is not a dictionary after processing: {ret}"
+                )
+
+            if ret.get("type") == "error":
+                error_message = ret.get("message", "Unknown error from API")
+                raise Exception(f"API returned an error: {error_message}")
+
+            msg = ret["data"]
+            if not isinstance(msg, dict) or msg.get("type") not in (
+                "comfyair",
+                "bizyair",
+            ):
+                raise Exception(f"Unexpected response type: {msg}")
+
+            caption = msg["data"]
+            return (caption,)
+
+        except json.JSONDecodeError as e:
+            raise Exception(f"Failed to parse JSON response: {ret}. Error: {e}")
+        except (KeyError, TypeError) as e:
+            raise Exception(
+                f"Invalid response format, missing key or wrong type. Response: {ret}. Error: {e}"
+            )
         except Exception as e:
-            raise Exception(f"Unexpected response: {ret} {e=}")
-
-        msg = ret["data"]
-        if msg["type"] not in (
-            "comfyair",
-            "bizyair",
-        ):
-            raise Exception(f"Unexpected response type: {msg}")
-
-        caption = msg["data"]
-        return (caption,)
+            raise Exception(
+                f"An unexpected error occurred while processing the response: {ret}. Error: {e}"
+            )
 
 
 class BizyAirJoyCaption3(BizyAirJoyCaption2):

bizyengine/misc/utils.py

@@ -1,14 +1,21 @@
+import asyncio
 import base64
+import concurrent.futures
 import json
+import logging
 import os
 import pickle
+import re
+import threading
+import time
 import urllib.parse
 import urllib.request
 import zlib
-from typing import List, Tuple, Union
+from typing import Callable, Dict, Generic, List, Optional, Tuple, TypeVar, Union
 
 import numpy as np
 
+from bizyengine.bizy_server.api_client import APIClient
 from bizyengine.core import pop_api_key_and_prompt_id
 from bizyengine.core.common import client
 from bizyengine.core.common.env_var import BIZYAIR_SERVER_ADDRESS
@@ -36,7 +43,7 @@ def send_post_request(api_url, payload, headers):
     try:
         data = json.dumps(payload).encode("utf-8")
         req = urllib.request.Request(api_url, data=data, headers=headers, method="POST")
-        with urllib.request.urlopen(req) as response:
+        with urllib.request.urlopen(req, timeout=3600) as response:
             response_data = response.read().decode("utf-8")
         return response_data
     except urllib.error.URLError as e:
@@ -143,6 +150,18 @@ def get_llm_response(
     extra_data = pop_api_key_and_prompt_id(kwargs)
     headers = client.headers(api_key=extra_data["api_key"])
 
+    # 如果model已不可用，选择第一个可用model
+    if _MODELS_CACHE.get("llm_models") is None:
+        cache_models(extra_data["api_key"])
+    llm_models = _MODELS_CACHE.get("llm_models")
+    if llm_models is None:
+        logging.warning(f"No LLM models available, keeping the original model {model}")
+    elif model not in llm_models:
+        logging.warning(
+            f"Model {model} is not available, using the first available model {llm_models[0]}"
+        )
+        model = llm_models[0]
+
     payload = {
         "model": model,
         "messages": [
@@ -183,6 +202,18 @@ def get_vlm_response(
     extra_data = pop_api_key_and_prompt_id(kwargs)
     headers = client.headers(api_key=extra_data["api_key"])
 
+    # 如果model已不可用，选择第一个可用model
+    if _MODELS_CACHE.get("vlm_models") is None:
+        cache_models(extra_data["api_key"])
+    vlm_models = _MODELS_CACHE.get("vlm_models")
+    if vlm_models is None:
+        logging.warning(f"No VLM models available, keeping the original model {model}")
+    elif model not in vlm_models:
+        logging.warning(
+            f"Model {model} is not available, using the first available model {vlm_models[0]}"
+        )
+        model = vlm_models[0]
+
     messages = [
         {
             "role": "system",
@@ -230,3 +261,149 @@ def get_vlm_response(
         callback=None,
     )
     return response
+
+
+K = TypeVar("K")
+V = TypeVar("V")
+R = TypeVar("R")
+
+
+class TTLCache(Generic[K, V]):
+    """线程安全 TTL 内存缓存（仅依赖标准库）"""
+
+    def __init__(self, ttl_sec: float):
+        self.ttl = ttl_sec
+        self._data: Dict[K, tuple[V, float]] = {}
+        self._lock = threading.RLock()
+        self._stop_evt = threading.Event()
+        # 后台清扫线程
+        self._cleaner = threading.Thread(target=self._cleanup, daemon=True)
+        self._cleaner.start()
+
+    # ---------- 公共 API ----------
+    def set(self, key: K, value: V) -> None:
+        """写入/刷新键值"""
+        with self._lock:
+            self._data[key] = (value, time.time() + self.ttl)
+
+    def get(self, key: K) -> Optional[V]:
+        """读取键值；不存在或已过期返回 None"""
+        with self._lock:
+            val, expire = self._data.get(key, (None, 0))
+            if val is None or time.time() > expire:
+                self._data.pop(key, None)
+                return None
+            return val
+
+    def delete(self, key: K) -> None:
+        """手动删除"""
+        with self._lock:
+            self._data.pop(key, None)
+
+    def stop(self):
+        """停止后台线程（程序退出前调用）"""
+        self._stop_evt.set()
+        self._cleaner.join(timeout=self.ttl + 1)
+
+    # ---------- 内部 ----------
+    def _cleanup(self):
+        """周期清扫过期键"""
+        while not self._stop_evt.wait(self.ttl / 2):
+            with self._lock:
+                now = time.time()
+                for key, (_, expire) in list(self._data.items()):
+                    if now > expire:
+                        self._data.pop(key, None)
+
+
+class SingleFlight(Generic[R]):
+    """Python 版 singleflight.Group（线程安全）"""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._call_map: dict[str, SingleFlight._Call[R]] = {}
+
+    class _Call(Generic[R]):
+        __slots__ = ("mu", "done", "result", "err", "waiters")
+
+        def __init__(self):
+            self.mu = threading.Lock()
+            self.done = False
+            self.result: Optional[R] = None
+            self.err: Optional[BaseException] = None
+            self.waiters = 0
+
+    def do(
+        self, key: str, fn: Callable[[], R]
+    ) -> tuple[R, bool, Optional[BaseException]]:
+        """
+        返回值: (result, shared?, exception)
+        shared=True 表示本次未真正执行 fn，复用了别人结果
+        """
+        with self._lock:
+            call = self._call_map.get(key)
+            if call is None:  # 我是第一个
+                call = self._Call[R]()
+                call.waiters = 1
+                self._call_map[key] = call
+                first = True
+            else:  # 已有并发请求
+                call.waiters += 1
+                first = False
+
+        if first:  # 只有第一个真正执行
+            try:
+                result = fn()
+                with call.mu:
+                    call.result = result
+                    call.done = True
+            except BaseException as e:
+                with call.mu:
+                    call.err = e
+                    call.done = True
+                raise
+            finally:  # 把自己从 map 摘掉
+                with self._lock:
+                    if call.waiters == 0:
+                        self._call_map.pop(key, None)
+        else:  # 其它人阻塞等待
+            with call.mu:
+                while not call.done:
+                    call.mu.wait()
+        # 读取结果
+        with call.mu:
+            if call.err is not None:
+                return call.result, True, call.err
+            return call.result, True, None
+
+
+_MODELS_CACHE = TTLCache[str, list[str]](ttl_sec=600)
+_SF = SingleFlight[None]()
+
+
+def cache_models(request_api_key: str):
+    # TODO: 效果待验证，目前节点只会被ComfyUI串行执行，所以不会出现竞争
+    _SF.do("_cache_models", lambda: _cache_models(request_api_key))
+
+
+def _cache_models(request_api_key: str):
+    # ① 开一条新线程专门跑协程 - 应该不需要在prompt那层上锁，因为并发只有1
+    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
+        api_client = APIClient()
+        all_models = pool.submit(
+            asyncio.run, api_client.fetch_all_llm_models(request_api_key)
+        ).result()
+        if len(all_models) == 0:
+            return
+        llm_models = [
+            model
+            for model in all_models
+            if not (re.search(r"\d+(\.\d+)?v", model.lower()) or "vl" in model.lower())
+        ]
+        vlm_models = [
+            model
+            for model in all_models
+            if re.search(r"\d+(\.\d+)?v", model.lower()) or "vl" in model.lower()
+        ]
+        _MODELS_CACHE.set("llm_models", llm_models)
+        _MODELS_CACHE.set("vlm_models", vlm_models)

bizyengine/version.txt

@@ -1 +1 @@
-1.2.45
+1.2.71

{bizyengine-1.2.45.dist-info → bizyengine-1.2.71.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bizyengine
-Version: 1.2.45
+Version: 1.2.71
 Summary: [a/BizyAir](https://github.com/siliconflow/BizyAir) Comfy Nodes that can run in any environment.
 Author-email: SiliconFlow <yaochi@siliconflow.cn>
 Project-URL: Repository, https://github.com/siliconflow/BizyAir
@@ -13,6 +13,8 @@ Requires-Dist: requests
 Requires-Dist: inputimeout
 Requires-Dist: openai>=1.77.0
 Requires-Dist: pycryptodome
+Requires-Dist: mcp>=1.18.0
+Requires-Dist: bizyairsdk>=0.1.5
 
 ## BizyEngine