bizyengine 1.2.6__py3-none-any.whl → 1.2.8__py3-none-any.whl

bizyengine/core/nodes_base.py

@@ -4,8 +4,11 @@ import warnings
 from functools import wraps
 from typing import List
 
+from bizyengine.core.common.env_var import BIZYAIR_DEBUG, BIZYAIR_SERVER_MODE
 from bizyengine.core.configs.conf import config_manager
+from server import PromptServer
 
+from .common.client import get_api_key
 from .data_types import is_send_request_datatype
 from .nodes_io import BizyAirNodeIO, create_node_data
 
@@ -23,6 +26,37 @@ PREFIX = f"BizyAir"
 NODE_CLASS_MAPPINGS = {}
 NODE_DISPLAY_NAME_MAPPINGS = {}
 
+BIZYAIR_PROMPT_KEY = "bizyair_prompt"
+BIZYAIR_PARAM_MAGIC_NODE_ID = "bizyair_magic_node"
+
+
+def pop_api_key_and_prompt_id(kwargs):
+    extra_data = {}
+    prompt = None
+    if BIZYAIR_PROMPT_KEY in kwargs:
+        prompt = kwargs.pop(BIZYAIR_PROMPT_KEY)
+    if BIZYAIR_SERVER_MODE:
+        if BIZYAIR_PARAM_MAGIC_NODE_ID in prompt:
+            extra_data["api_key"] = prompt[BIZYAIR_PARAM_MAGIC_NODE_ID]["_meta"][
+                "api_key"
+            ]
+            extra_data["prompt_id"] = prompt[BIZYAIR_PARAM_MAGIC_NODE_ID]["_meta"][
+                "prompt_id"
+            ]
+            logging.debug(
+                "Using server mode passed in prompt_id: " + extra_data["prompt_id"]
+            )
+    else:
+        extra_data["api_key"] = get_api_key()
+        if (
+            PromptServer.instance is not None
+            and PromptServer.instance.last_prompt_id is not None
+        ):
+            extra_data["prompt_id"] = PromptServer.instance.last_prompt_id
+            print("Processing prompt with ID: " + PromptServer.instance.last_prompt_id)
+
+    return extra_data
+
 
 def process_kwargs(kwargs):
     possibleWidgetNames = [
@@ -91,7 +125,7 @@ def ensure_unique_id(org_func, original_has_unique_id=False):
     return new_func
 
 
-def ensure_hidden_unique_id(org_input_types_func):
+def ensure_hidden_unique_id_and_prompt(org_input_types_func):
     original_has_unique_id = False
 
     @wraps(org_input_types_func)
@@ -105,6 +139,9 @@ def ensure_hidden_unique_id(org_input_types_func):
             result["hidden"]["unique_id"] = "UNIQUE_ID"
         else:
             original_has_unique_id = True
+        # Also set prompt, but to avoid naming conflict prefix with 'bizyair'
+        if BIZYAIR_PROMPT_KEY not in result["hidden"]:
+            result["hidden"][BIZYAIR_PROMPT_KEY] = "PROMPT"
         return result
 
     # Ensure original_has_unique_id is correctly set before returning
@@ -121,11 +158,19 @@ class BizyAirBaseNode:
         register_node(cls, PREFIX)
         cls.setup_input_types()
 
+        # 验证FUNCTION接受**kwargs
+        if BIZYAIR_DEBUG:
+            import inspect
+
+            sig = inspect.signature(getattr(cls, cls.FUNCTION))
+            params = sig.parameters.values()
+            assert any([True for p in params if p.kind == p.VAR_KEYWORD])
+
     @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
+        # https://docs.comfy.org/custom-nodes/backend/more_on_inputs#hidden-and-flexible-inputs
+        new_input_types_func, original_has_unique_id = (
+            ensure_hidden_unique_id_and_prompt(cls.INPUT_TYPES)
         )
         cls.INPUT_TYPES = new_input_types_func
         setattr(
@@ -140,13 +185,14 @@ class BizyAirBaseNode:
         return str(self._assigned_id)
 
     def default_function(self, **kwargs):
+        extra_data = pop_api_key_and_prompt_id(kwargs)
         class_type = self._determine_class_type()
         kwargs = process_kwargs(kwargs)
         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 self._process_all_send_request_types(node_ios, **extra_data)
         return node_ios
 
     def _get_send_request_datatypes(self):
@@ -172,7 +218,39 @@ class BizyAirBaseNode:
             outs.append(node)
         return tuple(outs)
 
-    def _process_all_send_request_types(self, node_ios: List[BizyAirNodeIO]):
-        out = node_ios[0].send_request()
+    def _process_all_send_request_types(self, node_ios: List[BizyAirNodeIO], **kwargs):
+        out = node_ios[0].send_request(**kwargs)
         assert len(out) == len(self.RETURN_TYPES)
         return out
+
+
+class BizyAirMiscBaseNode:
+    # 作为Misc节点基类来保证hidden prompt, unique id存在
+    def __init_subclass__(cls, **kwargs):
+        if hasattr(cls, "CATEGORY"):
+            if not cls.CATEGORY.startswith(f"{LOGO}{PREFIX}"):
+                cls.CATEGORY = f"{LOGO}{PREFIX}/{cls.CATEGORY}"
+        cls.setup_input_types()
+
+        # 验证FUNCTION接受**kwargs
+        if BIZYAIR_DEBUG:
+            import inspect
+
+            sig = inspect.signature(getattr(cls, cls.FUNCTION))
+            params = sig.parameters.values()
+            assert any([True for p in params if p.kind == p.VAR_KEYWORD])
+
+    @classmethod
+    def setup_input_types(cls):
+        if not hasattr(cls, "INPUT_TYPES"):
+            cls.INPUT_TYPES = lambda: {}
+        # https://docs.comfy.org/custom-nodes/backend/more_on_inputs#hidden-and-flexible-inputs
+        new_input_types_func, original_has_unique_id = (
+            ensure_hidden_unique_id_and_prompt(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),
+        )

bizyengine/core/nodes_io.py

@@ -82,10 +82,10 @@ class BizyAirNodeIO:
                 self.nodes.update(other.nodes)
 
     def send_request(
-        self, url=None, headers=None, *, progress_callback=None, stream=False
+        self, url=None, headers=None, *, progress_callback=None, stream=False, **kwargs
     ) -> any:
         out = invoker.prompt_server.execute(
-            prompt=self.nodes, last_node_ids=[self.node_id]
+            prompt=self.nodes, last_node_ids=[self.node_id], **kwargs
         )
         return out
 

bizyengine/misc/llm.py

@@ -3,13 +3,14 @@ import json
 
 import aiohttp
 from aiohttp import web
+from bizyengine.core import BizyAirMiscBaseNode, pop_api_key_and_prompt_id
+from bizyengine.core.common import client
 from bizyengine.core.common.env_var import BIZYAIR_SERVER_ADDRESS
 from bizyengine.core.image_utils import decode_data, encode_comfy_image, encode_data
 from server import PromptServer
 
 from .utils import (
     decode_and_deserialize,
-    get_api_key,
     get_llm_response,
     get_vlm_response,
     send_post_request,
@@ -17,52 +18,7 @@ from .utils import (
 )
 
 
-async def fetch_all_models(api_key):
-    url = f"{BIZYAIR_SERVER_ADDRESS}/llm/models"
-    headers = {"accept": "application/json", "authorization": f"Bearer {api_key}"}
-    params = {"type": "text", "sub_type": "chat"}
-
-    try:
-        async with aiohttp.ClientSession() as session:
-            async with session.get(
-                url, headers=headers, params=params, timeout=10
-            ) as response:
-                if response.status == 200:
-                    data = await response.json()
-                    all_models = [model["id"] for model in data["data"]]
-                    return all_models
-                else:
-                    print(f"Error fetching models: HTTP Status {response.status}")
-                    return []
-    except aiohttp.ClientError as e:
-        print(f"Error fetching models: {e}")
-        return []
-    except asyncio.exceptions.TimeoutError:
-        print("Request to fetch models timed out")
-        return []
-
-
-@PromptServer.instance.routes.post("/bizyair/get_silicon_cloud_llm_models")
-async def get_silicon_cloud_llm_models_endpoint(request):
-    data = await request.json()
-    api_key = data.get("api_key", get_api_key())
-    all_models = await fetch_all_models(api_key)
-    llm_models = [model for model in all_models if "vl" not in model.lower()]
-    llm_models.append("No LLM Enhancement")
-    return web.json_response(llm_models)
-
-
-@PromptServer.instance.routes.post("/bizyair/get_silicon_cloud_vlm_models")
-async def get_silicon_cloud_vlm_models_endpoint(request):
-    data = await request.json()
-    api_key = data.get("api_key", get_api_key())
-    all_models = await fetch_all_models(api_key)
-    vlm_models = [model for model in all_models if "vl" in model.lower()]
-    vlm_models.append("No VLM Enhancement")
-    return web.json_response(vlm_models)
-
-
-class SiliconCloudLLMAPI:
+class SiliconCloudLLMAPI(BizyAirMiscBaseNode):
     def __init__(self):
         pass
 
@@ -93,7 +49,7 @@ class SiliconCloudLLMAPI:
                     "FLOAT",
                     {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.01},
                 ),
-            }
+            },
         }
 
     RETURN_TYPES = ("STRING",)
@@ -102,23 +58,18 @@ class SiliconCloudLLMAPI:
     CATEGORY = "☁️BizyAir/AI Assistants"
 
     def get_llm_model_response(
-        self, model, system_prompt, user_prompt, max_tokens, temperature
+        self, model, system_prompt, user_prompt, max_tokens, temperature, **kwargs
     ):
         if model == "No LLM Enhancement":
             return {"ui": {"text": (user_prompt,)}, "result": (user_prompt,)}
-        response = get_llm_response(
-            model,
-            system_prompt,
-            user_prompt,
-            max_tokens,
-            temperature,
+        ret = get_llm_response(
+            model, system_prompt, user_prompt, max_tokens, temperature, **kwargs
         )
-        ret = json.loads(response)
         text = ret["choices"][0]["message"]["content"]
         return (text,)  # if update ui:  {"ui": {"text": (text,)}, "result": (text,)}
 
 
-class SiliconCloudVLMAPI:
+class SiliconCloudVLMAPI(BizyAirMiscBaseNode):
     def __init__(self):
         pass
 
@@ -148,7 +99,7 @@ class SiliconCloudVLMAPI:
                     {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.01},
                 ),
                 "detail": (["auto", "low", "high"], {"default": "auto"}),
-            }
+            },
         }
 
     RETURN_TYPES = ("STRING",)
@@ -157,7 +108,15 @@ class SiliconCloudVLMAPI:
     CATEGORY = "☁️BizyAir/AI Assistants"
 
     def get_vlm_model_response(
-        self, model, system_prompt, user_prompt, images, max_tokens, temperature, detail
+        self,
+        model,
+        system_prompt,
+        user_prompt,
+        images,
+        max_tokens,
+        temperature,
+        detail,
+        **kwargs,
     ):
         if model == "No VLM Enhancement":
             return (user_prompt,)
@@ -171,7 +130,7 @@ class SiliconCloudVLMAPI:
         # 提取所有编码后的图像
         base64_images = list(encoded_images_dict.values())
 
-        response = get_vlm_response(
+        ret = get_vlm_response(
             model,
             system_prompt,
             user_prompt,
@@ -179,13 +138,13 @@ class SiliconCloudVLMAPI:
             max_tokens,
             temperature,
             detail,
+            **kwargs,
         )
-        ret = json.loads(response)
         text = ret["choices"][0]["message"]["content"]
         return (text,)
 
 
-class BizyAirJoyCaption:
+class BizyAirJoyCaption(BizyAirMiscBaseNode):
     # refer to: https://huggingface.co/spaces/fancyfeast/joy-caption-pre-alpha
     API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/joycaption2"
 
@@ -216,7 +175,7 @@ class BizyAirJoyCaption:
                         "display": "number",
                     },
                 ),
-            }
+            },
         }
 
     RETURN_TYPES = ("STRING",)
@@ -224,8 +183,10 @@ class BizyAirJoyCaption:
 
     CATEGORY = "☁️BizyAir/AI Assistants"
 
-    def joycaption(self, image, do_sample, temperature, max_tokens):
-        API_KEY = get_api_key()
+    def joycaption(self, image, do_sample, temperature, max_tokens, **kwargs):
+        extra_data = pop_api_key_and_prompt_id(kwargs)
+        headers = client.headers(api_key=extra_data["api_key"])
+
         SIZE_LIMIT = 1536
         # device = image.device
         _, w, h, c = image.shape
@@ -244,17 +205,18 @@ class BizyAirJoyCaption:
             "name_input": "",
             "custom_prompt": "A descriptive caption for this image:\n",
         }
-        auth = f"Bearer {API_KEY}"
-        headers = {
-            "accept": "application/json",
-            "content-type": "application/json",
-            "authorization": auth,
-        }
         input_image = encode_data(image, disable_image_marker=True)
         payload["image"] = input_image
-
-        ret: str = send_post_request(self.API_URL, payload=payload, headers=headers)
-        ret = json.loads(ret)
+        if "prompt_id" in extra_data:
+            payload["prompt_id"] = extra_data["prompt_id"]
+        data = json.dumps(payload).encode("utf-8")
+
+        ret = client.send_request(
+            url=self.API_URL,
+            data=data,
+            headers=headers,
+            callback=None,
+        )
 
         try:
             if "result" in ret:
@@ -275,7 +237,7 @@ class BizyAirJoyCaption:
         return (caption,)
 
 
-class BizyAirJoyCaption2:
+class BizyAirJoyCaption2(BizyAirMiscBaseNode):
     def __init__(self):
         pass
 
@@ -348,7 +310,7 @@ class BizyAirJoyCaption2:
                         "multiline": True,
                     },
                 ),
-            }
+            },
         }
 
     RETURN_TYPES = ("STRING",)
@@ -367,8 +329,11 @@ class BizyAirJoyCaption2:
         extra_options,
         name_input,
         custom_prompt,
+        **kwargs,
     ):
-        API_KEY = get_api_key()
+        extra_data = pop_api_key_and_prompt_id(kwargs)
+        headers = client.headers(api_key=extra_data["api_key"])
+
         SIZE_LIMIT = 1536
         _, w, h, c = image.shape
         assert (
@@ -386,17 +351,15 @@ class BizyAirJoyCaption2:
             "name_input": name_input,
             "custom_prompt": custom_prompt,
         }
-        auth = f"Bearer {API_KEY}"
-        headers = {
-            "accept": "application/json",
-            "content-type": "application/json",
-            "authorization": auth,
-        }
         input_image = encode_data(image, disable_image_marker=True)
         payload["image"] = input_image
+        if "prompt_id" in extra_data:
+            payload["prompt_id"] = extra_data["prompt_id"]
+        data = json.dumps(payload).encode("utf-8")
 
-        ret: str = send_post_request(self.API_URL, payload=payload, headers=headers)
-        ret = json.loads(ret)
+        ret: str = client.send_request(
+            url=self.API_URL, data=data, headers=headers, callback=None
+        )
 
         try:
             if "result" in ret:

bizyengine/misc/mzkolors.py

@@ -1,23 +1,29 @@
+import json
 import os
 import uuid
 
 import torch
-from bizyengine.core import BizyAirBaseNode, BizyAirNodeIO, create_node_data
+from bizyengine.core import (
+    BizyAirBaseNode,
+    BizyAirMiscBaseNode,
+    BizyAirNodeIO,
+    create_node_data,
+)
+from bizyengine.core.common import client
 from bizyengine.core.common.env_var import BIZYAIR_SERVER_ADDRESS
 from bizyengine.core.data_types import CONDITIONING
 from bizyengine.core.image_utils import encode_data
 
 from .utils import (
     decode_and_deserialize,
-    get_api_key,
-    send_post_request,
+    pop_api_key_and_prompt_id,
     serialize_and_encode,
 )
 
 CATEGORY_NAME = "☁️BizyAir/Kolors"
 
 
-class BizyAirMZChatGLM3TextEncode:
+class BizyAirMZChatGLM3TextEncode(BizyAirMiscBaseNode):
     API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/mzkolorschatglm3"
 
     @classmethod
@@ -25,7 +31,7 @@ class BizyAirMZChatGLM3TextEncode:
         return {
             "required": {
                 "text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
-            }
+            },
         }
 
     RETURN_TYPES = ("CONDITIONING",)
@@ -33,24 +39,26 @@ class BizyAirMZChatGLM3TextEncode:
     FUNCTION = "encode"
     CATEGORY = CATEGORY_NAME
 
-    def encode(self, text):
-        API_KEY = get_api_key()
+    def encode(self, text, **kwargs):
+        extra_data = pop_api_key_and_prompt_id(kwargs)
+        headers = client.headers(api_key=extra_data["api_key"])
+
         assert len(text) <= 4096, f"the prompt is too long, length: {len(text)}"
 
         payload = {
             "text": text,
         }
-        auth = f"Bearer {API_KEY}"
-        headers = {
-            "accept": "application/json",
-            "content-type": "application/json",
-            "authorization": auth,
-        }
-
-        response: str = send_post_request(
-            self.API_URL, payload=payload, headers=headers
+        if "prompt_id" in extra_data:
+            payload["prompt_id"] = extra_data["prompt_id"]
+        data = json.dumps(payload).encode("utf-8")
+
+        tensors_np = client.send_request(
+            url=self.API_URL,
+            data=data,
+            headers=headers,
+            callback=None,
+            response_handler=decode_and_deserialize,
         )
-        tensors_np = decode_and_deserialize(response)
 
         ret_conditioning = []
         for item in tensors_np:
@@ -69,8 +77,8 @@ class BizyAir_MinusZoneChatGLM3TextEncode(BizyAirMZChatGLM3TextEncode, BizyAirBa
 
     FUNCTION = "mz_encode"
 
-    def mz_encode(self, text):
-        out = self.encode(text)[0]
+    def mz_encode(self, text, **kwargs):
+        out = self.encode(text=text, **kwargs)[0]
         node_data = create_node_data(
             class_type="ComfyAirLoadData",
             inputs={"conditioning": {"relay": out}},