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

bizyengine/misc/nodes.py

@@ -55,10 +55,11 @@ class BizyAir_KSampler(BizyAirBaseNode):
         }
 
     RETURN_TYPES = ("LATENT",)
-    FUNCTION = "sample"
+    # FUNCTION = "sample"
     RETURN_NAMES = (f"LATENT",)
     CATEGORY = f"{PREFIX}/sampling"
 
+    # deprecated
     def sample(
         self,
         model,
@@ -127,7 +128,7 @@ class KSamplerAdvanced(BizyAirBaseNode):
         }
 
     RETURN_TYPES = ("LATENT",)
-    FUNCTION = "sample"
+    # FUNCTION = "sample"
 
     CATEGORY = "sampling"
 
@@ -159,7 +160,7 @@ class BizyAir_CheckpointLoaderSimple(BizyAirBaseNode):
         }
 
     RETURN_TYPES = (data_types.MODEL, data_types.CLIP, data_types.VAE)
-    FUNCTION = "load_checkpoint"
+    # FUNCTION = "load_checkpoint"
     CATEGORY = f"{PREFIX}/loaders"
     RETURN_NAMES = (
         f"model",
@@ -224,7 +225,7 @@ class BizyAir_CLIPTextEncode(BizyAirBaseNode):
 
     RETURN_TYPES = (data_types.CONDITIONING,)
     RETURN_NAMES = ("CONDITIONING",)
-    FUNCTION = "encode"
+    # FUNCTION = "encode"
 
     CATEGORY = f"{PREFIX}/conditioning"
 
@@ -249,7 +250,7 @@ class BizyAir_VAEDecode(BizyAirBaseNode):
 
     RETURN_TYPES = ("IMAGE",)
     RETURN_NAMES = (f"IMAGE",)
-    FUNCTION = "decode"
+    # FUNCTION = "decode"
 
     CATEGORY = f"{PREFIX}/latent"
 
@@ -298,6 +299,7 @@ class BizyAir_LoraLoader(BizyAirBaseNode):
     RETURN_TYPES = (data_types.MODEL, data_types.CLIP)
     RETURN_NAMES = ("MODEL", "CLIP")
 
+    # 不能使用default_function
     FUNCTION = "load_lora"
     CATEGORY = f"{PREFIX}/loaders"
 
@@ -309,6 +311,7 @@ class BizyAir_LoraLoader(BizyAirBaseNode):
         strength_model,
         strength_clip,
         model_version_id: str = None,
+        **kwargs,
     ):
         assigned_id = self.assigned_id
         new_model: BizyAirNodeIO = model.copy(assigned_id)
@@ -371,7 +374,9 @@ class BizyAir_LoraLoader_Legacy(BizyAirBaseNode):
 
     CATEGORY = f"{PREFIX}/loaders"
 
-    def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
+    def load_lora(
+        self, model, clip, lora_name, strength_model, strength_clip, **kwargs
+    ):
         assigned_id = self.assigned_id
         new_model: BizyAirNodeIO = model.copy(assigned_id)
         new_clip: BizyAirNodeIO = clip.copy(assigned_id)
@@ -401,9 +406,10 @@ class BizyAir_VAEEncode(BizyAirBaseNode):
 
     RETURN_TYPES = ("LATENT",)
     RETURN_NAMES = (f"LATENT",)
-    FUNCTION = "encode"
+    # FUNCTION = "encode"
     CATEGORY = f"{PREFIX}/latent"
 
+    # deprecated
     def encode(self, vae, pixels):
         new_vae: BizyAirNodeIO = vae.copy(self.assigned_id)
         new_vae.add_node_data(
@@ -431,9 +437,10 @@ class BizyAir_VAEEncodeForInpaint(BizyAirBaseNode):
 
     RETURN_TYPES = (f"LATENT",)
     RETURN_NAMES = (f"LATENT",)
-    FUNCTION = "encode"
+    # FUNCTION = "encode"
     CATEGORY = f"{PREFIX}/latent/inpaint"
 
+    # deprecated
     def encode(self, vae, pixels, mask, grow_mask_by=6):
         new_vae: BizyAirNodeIO = vae.copy(self.assigned_id)
         new_vae.add_node_data(
@@ -465,7 +472,7 @@ class BizyAir_ControlNetLoader(BizyAirBaseNode):
 
     RETURN_TYPES = (data_types.CONTROL_NET,)
     RETURN_NAMES = ("CONTROL_NET",)
-    FUNCTION = "load_controlnet"
+    # FUNCTION = "load_controlnet"
 
     CATEGORY = f"{PREFIX}/loaders"
 
@@ -506,11 +513,12 @@ class BizyAir_ControlNetLoader_Legacy(BizyAirBaseNode):
 
     RETURN_TYPES = (data_types.CONTROL_NET,)
     RETURN_NAMES = ("CONTROL_NET",)
+    # 似乎不能用default实现
     FUNCTION = "load_controlnet"
 
     CATEGORY = f"{PREFIX}/loaders"
 
-    def load_controlnet(self, control_net_name):
+    def load_controlnet(self, control_net_name, **kwargs):
 
         node_data = create_node_data(
             class_type="ControlNetLoader",
@@ -550,7 +558,7 @@ class BizyAir_ControlNetApplyAdvanced(BizyAirBaseNode):
 
     RETURN_TYPES = (data_types.CONDITIONING, data_types.CONDITIONING)
     RETURN_NAMES = ("positive", "negative")
-    FUNCTION = "apply_controlnet"
+    # FUNCTION = "apply_controlnet"
 
     CATEGORY = "conditioning"
 
@@ -587,7 +595,7 @@ class BizyAir_ControlNetApply(BizyAirBaseNode):
 
     RETURN_TYPES = (data_types.CONDITIONING,)
     RETURN_NAMES = ("CONDITIONING",)
-    FUNCTION = "apply_controlnet"
+    # FUNCTION = "apply_controlnet"
 
     CATEGORY = f"{PREFIX}/conditioning/controlnet"
 
@@ -629,7 +637,7 @@ class BizyAir_CLIPVisionLoader(BizyAirBaseNode):
         }
 
     RETURN_TYPES = ("CLIP_VISION",)
-    FUNCTION = "load_clip"
+    # FUNCTION = "load_clip"
 
     CATEGORY = "loaders"
 
@@ -689,7 +697,7 @@ class VAELoader(BizyAirBaseNode):
 
     RETURN_TYPES = (data_types.VAE,)
     RETURN_NAMES = ("vae",)
-    FUNCTION = "load_vae"
+    # FUNCTION = "load_vae"
 
     CATEGORY = "loaders"
 
@@ -757,7 +765,7 @@ class UNETLoader(BizyAirBaseNode):
         }
 
     RETURN_TYPES = (data_types.MODEL,)
-    FUNCTION = "load_unet"
+    # FUNCTION = "load_unet"
 
     CATEGORY = "advanced/loaders"
 
@@ -851,7 +859,7 @@ class BasicGuider(BizyAirBaseNode):
 
     RETURN_TYPES = ("GUIDER",)
 
-    FUNCTION = "get_guider"
+    # FUNCTION = "get_guider"
     CATEGORY = "sampling/custom_sampling/guiders"
 
     def get_guider(self, model: BizyAirNodeIO, conditioning):
@@ -885,7 +893,7 @@ class BasicScheduler(BizyAirBaseNode):
     RETURN_TYPES = ("SIGMAS",)
     CATEGORY = "sampling/custom_sampling/schedulers"
 
-    FUNCTION = "get_sigmas"
+    # FUNCTION = "get_sigmas"
 
     def get_sigmas(self, **kwargs):
         new_model: BizyAirNodeIO = kwargs["model"].copy(self.assigned_id)
@@ -931,7 +939,7 @@ class DualCLIPLoader(BizyAirBaseNode):
         }
 
     RETURN_TYPES = (data_types.CLIP,)
-    FUNCTION = "load_clip"
+    # FUNCTION = "load_clip"
 
     CATEGORY = "advanced/loaders"
 
@@ -1004,7 +1012,7 @@ class KSamplerSelect(BizyAirBaseNode):
     RETURN_TYPES = ("SAMPLER",)
     CATEGORY = "sampling/custom_sampling/samplers"
 
-    FUNCTION = "get_sampler"
+    # FUNCTION = "get_sampler"
 
     def get_sampler(self, **kwargs):
         node_data = create_node_data(
@@ -1032,7 +1040,7 @@ class RandomNoise(BizyAirBaseNode):
         }
 
     RETURN_TYPES = ("NOISE",)
-    FUNCTION = "get_noise"
+    # FUNCTION = "get_noise"
     CATEGORY = "sampling/custom_sampling/noise"
 
     def get_noise(self, noise_seed):
@@ -1064,7 +1072,7 @@ class CLIPSetLastLayer(BizyAirBaseNode):
         }
 
     RETURN_TYPES = (data_types.CLIP,)
-    FUNCTION = "set_last_layer"
+    # FUNCTION = "set_last_layer"
 
     CATEGORY = "conditioning"
 
@@ -1141,6 +1149,7 @@ class SharedLoraLoader(BizyAir_LoraLoader_Legacy):
 
     RETURN_TYPES = (data_types.MODEL, data_types.CLIP)
     RETURN_NAMES = ("MODEL", "CLIP")
+    # 似乎不能用default实现
     FUNCTION = "shared_load_lora"
     CATEGORY = f"{PREFIX}/loaders"
     NODE_DISPLAY_NAME = "Shared Lora Loader"
@@ -1422,3 +1431,14 @@ class StyleModelApply(BizyAirBaseNode):
     # FUNCTION = "apply_stylemodel"
 
     CATEGORY = "conditioning/style_model"
+
+
+# 仅用于使用meta传参，故没有输入输出
+class PassParameter(BizyAirBaseNode):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {}
+
+    RETURN_TYPES = ()
+
+    CATEGORY = "☁️BizyAir"

bizyengine/misc/nodes_controlnet_aux.py

@@ -1,22 +1,20 @@
+import json
 import os
 
 import numpy as np
 import torch
+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 .utils import (
-    decode_and_deserialize,
-    get_api_key,
-    send_post_request,
-    serialize_and_encode,
-)
+from .utils import decode_and_deserialize, serialize_and_encode
 
 # Sync with theoritical limit from Comfy base
 # https://github.com/comfyanonymous/ComfyUI/blob/eecd69b53a896343775bcb02a4f8349e7442ffd1/nodes.py#L45
 MAX_RESOLUTION = 1024
 
 
-class BasePreprocessor:
+class BasePreprocessor(BizyAirMiscBaseNode):
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
         if not hasattr(cls, "model_name"):
@@ -24,27 +22,29 @@ class BasePreprocessor:
         cls.API_URL = f"{BIZYAIR_SERVER_ADDRESS}{cls.model_name}"
         cls.CATEGORY = f"☁️BizyAir/{cls.CATEGORY}"
 
-    @staticmethod
-    def get_headers():
-        return {
-            "accept": "application/json",
-            "content-type": "application/json",
-            "authorization": f"Bearer {get_api_key()}",
-        }
-
     RETURN_TYPES = ("IMAGE",)
     FUNCTION = "execute"
 
     def execute(self, **kwargs):
+        extra_data = pop_api_key_and_prompt_id(kwargs)
+        headers = client.headers(api_key=extra_data["api_key"])
+
         compress = True
         image: torch.Tensor = kwargs.pop("image")
         device = image.device
         kwargs["image"] = serialize_and_encode(image, compress)[0]
         kwargs["is_compress"] = compress
-        response: str = send_post_request(
-            self.API_URL, payload=kwargs, headers=self.get_headers()
+        if "prompt_id" in extra_data:
+            kwargs["prompt_id"] = extra_data["prompt_id"]
+        data = json.dumps(kwargs).encode("utf-8")
+
+        image_np = client.send_request(
+            url=self.API_URL,
+            data=data,
+            headers=headers,
+            callback=None,
+            response_handler=decode_and_deserialize,
         )
-        image_np = decode_and_deserialize(response)
         image_torch = torch.from_numpy(image_np).to(device)
         return (image_torch,)
 

bizyengine/misc/nodes_controlnet_union_sdxl.py

@@ -8,13 +8,13 @@ import os
 
 import numpy as np
 import requests
+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_comfy_image, encode_comfy_image
 
-from .utils import get_api_key
 
-
-class StableDiffusionXLControlNetUnionPipeline:
+class StableDiffusionXLControlNetUnionPipeline(BizyAirMiscBaseNode):
     API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/diffusers-v1-stablediffusionxlcontrolnetunionpipeline"
 
     @classmethod
@@ -83,14 +83,6 @@ class StableDiffusionXLControlNetUnionPipeline:
     FUNCTION = "process"
     CATEGORY = "☁️BizyAir/ControlNet"
 
-    @staticmethod
-    def get_headers():
-        return {
-            "accept": "application/json",
-            "content-type": "application/json",
-            "authorization": f"Bearer {get_api_key()}",
-        }
-
     def process(
         self,
         openpose_image=None,
@@ -101,6 +93,7 @@ class StableDiffusionXLControlNetUnionPipeline:
         segment_image=None,
         **kwargs,
     ):
+        extra_data = pop_api_key_and_prompt_id(kwargs)
         controlnet_img = {
             0: openpose_image,
             1: depth_image,
@@ -143,7 +136,7 @@ class StableDiffusionXLControlNetUnionPipeline:
         response = requests.post(
             self.API_URL,
             json=payload,
-            headers=self.get_headers(),
+            headers=client.headers(api_key=extra_data["api_key"]),
         )
 
         result = response.json()

bizyengine/misc/segment_anything.py

@@ -6,13 +6,14 @@ from enum import Enum
 import folder_paths
 import numpy as np
 import torch
+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_base64_to_np, encode_image_to_base64
 from nodes import LoadImage
 from PIL import Image, ImageOps, ImageSequence
 
 from .route_sam import SAM_COORDINATE
-from .utils import get_api_key, send_post_request
 
 
 class INFER_MODE(Enum):
@@ -27,7 +28,7 @@ class EDIT_MODE(Enum):
     point = 1
 
 
-class BizyAirSegmentAnythingText:
+class BizyAirSegmentAnythingText(BizyAirMiscBaseNode):
     API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/sam"
 
     @classmethod
@@ -44,7 +45,7 @@ class BizyAirSegmentAnythingText:
                     "FLOAT",
                     {"default": 0.3, "min": 0, "max": 1.0, "step": 0.01},
                 ),
-            }
+            },
         }
 
     RETURN_TYPES = ("IMAGE", "MASK")
@@ -52,8 +53,10 @@ class BizyAirSegmentAnythingText:
 
     CATEGORY = "☁️BizyAir/segment-anything"
 
-    def text_sam(self, image, prompt, box_threshold, text_threshold):
-        API_KEY = get_api_key()
+    def text_sam(self, image, prompt, box_threshold, text_threshold, **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
@@ -70,19 +73,20 @@ class BizyAirSegmentAnythingText:
                 "text_threshold": text_threshold,
             },
         }
-        auth = f"Bearer {API_KEY}"
-        headers = {
-            "accept": "application/json",
-            "content-type": "application/json",
-            "authorization": auth,
-        }
         image = image.squeeze(0).numpy()
         image_pil = Image.fromarray((image * 255).astype(np.uint8))
         input_image = encode_image_to_base64(image_pil, format="webp")
         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:
@@ -117,7 +121,7 @@ class BizyAirSegmentAnythingText:
         return (img, img_mask)
 
 
-class BizyAirSegmentAnythingPointBox:
+class BizyAirSegmentAnythingPointBox(BizyAirMiscBaseNode):
     API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/sam"
 
     @classmethod
@@ -141,8 +145,10 @@ class BizyAirSegmentAnythingPointBox:
 
     CATEGORY = "☁️BizyAir/segment-anything"
 
-    def apply(self, image, is_point):
-        API_KEY = get_api_key()
+    def apply(self, image, is_point, **kwargs):
+        extra_data = pop_api_key_and_prompt_id(kwargs)
+        headers = client.headers(api_key=extra_data["api_key"])
+
         SIZE_LIMIT = 1536
 
         # 加载原始图像
@@ -201,20 +207,18 @@ class BizyAirSegmentAnythingPointBox:
                 },
             }
 
-        auth = f"Bearer {API_KEY}"
-        headers = {
-            "accept": "application/json",
-            "content-type": "application/json",
-            "authorization": auth,
-        }
         # 处理用于API的图像
         api_image = image_to_process.squeeze(0).numpy()
         image_pil = Image.fromarray((api_image * 255).astype(np.uint8))
         input_image = encode_image_to_base64(image_pil, format="webp")
         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 = client.send_request(
+            url=self.API_URL, data=data, headers=headers, callback=None
+        )
 
         try:
             if "result" in ret:

bizyengine/misc/supernode.py

@@ -8,6 +8,8 @@ import folder_paths
 import node_helpers
 import numpy as np
 import torch
+from bizyengine.core import BizyAirMiscBaseNode
+from bizyengine.core.common import client
 from bizyengine.core.common.env_var import BIZYAIR_SERVER_ADDRESS
 from bizyengine.core.image_utils import (
     decode_base64_to_np,
@@ -20,13 +22,12 @@ from PIL import Image, ImageOps, ImageSequence
 
 from .utils import (
     decode_and_deserialize,
-    get_api_key,
-    send_post_request,
+    pop_api_key_and_prompt_id,
     serialize_and_encode,
 )
 
 
-class RemoveBackground:
+class RemoveBackground(BizyAirMiscBaseNode):
     API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/removebg"
 
     @classmethod
@@ -34,7 +35,7 @@ class RemoveBackground:
         return {
             "required": {
                 "image": ("IMAGE",),
-            }
+            },
         }
 
     RETURN_TYPES = ("IMAGE", "MASK")
@@ -42,8 +43,10 @@ class RemoveBackground:
 
     CATEGORY = "☁️BizyAir"
 
-    def remove_background(self, image):
-        API_KEY = get_api_key()
+    def remove_background(self, image, **kwargs):
+        extra_data = pop_api_key_and_prompt_id(kwargs)
+        headers = client.headers(api_key=extra_data["api_key"])
+
         device = image.device
         _, h, w, _ = image.shape
         assert (
@@ -54,26 +57,27 @@ class RemoveBackground:
             "is_compress": True,
             "image": None,
         }
-        auth = f"Bearer {API_KEY}"
-        headers = {
-            "accept": "application/json",
-            "content-type": "application/json",
-            "authorization": auth,
-        }
         input_image, compress = serialize_and_encode(image, compress=True)
         payload["image"] = input_image
         payload["is_compress"] = compress
-
-        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 = client.send_request(
+            url=self.API_URL,
+            data=data,
+            headers=headers,
+            callback=None,
+            response_handler=decode_and_deserialize,
         )
-        tensors = decode_and_deserialize(response)
+
         t_images = tensors["images"].to(device)
         t_mask = tensors["mask"].to(device)
         return (t_images, t_mask)
 
 
-class GenerateLightningImage:
+class GenerateLightningImage(BizyAirMiscBaseNode):
     API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/realvis4lightning"
 
     @classmethod
@@ -98,7 +102,7 @@ class GenerateLightningImage:
                     },
                 ),
                 "batch_size": ("INT", {"default": 1, "min": 1, "max": 4}),
-            }
+            },
         }
 
     RETURN_TYPES = ("IMAGE",)
@@ -106,8 +110,10 @@ class GenerateLightningImage:
 
     CATEGORY = "☁️BizyAir"
 
-    def generate_image(self, prompt, seed, width, height, cfg, batch_size):
-        API_KEY = get_api_key()
+    def generate_image(self, prompt, seed, width, height, cfg, batch_size, **kwargs):
+        extra_data = pop_api_key_and_prompt_id(kwargs)
+        headers = client.headers(api_key=extra_data["api_key"])
+
         assert (
             width <= 1024 and height <= 1024
         ), f"width and height must be less than 1024, but got {width} and {height}"
@@ -120,17 +126,17 @@ class GenerateLightningImage:
             "cfg": cfg,
             "seed": seed,
         }
-        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)
         tensors = torch.from_numpy(tensors_np)
 
         return (tensors,)