diff --git a/py/libs/api/bizyair.py b/py/libs/api/bizyair.py new file mode 100644 index 0000000..afbb429 --- /dev/null +++ b/py/libs/api/bizyair.py @@ -0,0 +1,369 @@ +import yaml +import pathlib +import base64 +import io +import json +import os +import pickle +import zlib +import urllib.parse +import urllib.request +import urllib.error +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 wandb.env import API_KEY + +root_path = pathlib.Path(__file__).parent.parent.parent.parent +config_path = os.path.join(root_path,'config.yaml') + +class BizyAIRAPI: + def __init__(self): + self.base_url = 'https://bizyair-api.siliconflow.cn/x/v1' + self.api_key = None + + + def getAPIKey(self): + if self.api_key is None: + if os.path.isfile(config_path): + with open(config_path, 'r') as f: + data = yaml.load(f, Loader=yaml.FullLoader) + if 'BIZYAIR_API_KEY' not in data: + raise Exception("Please add BIZYAIR_API_KEY to config.yaml") + self.api_key = data['BIZYAIR_API_KEY'] + else: + raise Exception("Please add config.yaml to root path") + return self.api_key + + def send_post_request(self, url, payload, headers): + try: + data = json.dumps(payload).encode("utf-8") + req = urllib.request.Request(url, data=data, headers=headers, method="POST") + with urllib.request.urlopen(req) as response: + response_data = response.read().decode("utf-8") + return response_data + except urllib.error.URLError as e: + if "Unauthorized" in str(e): + raise Exception( + "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 Exception( + f"Failed to connect to the server: {e}, if you have no key, " + ) + + # joycaptionTwo + def joyCaption2(self, payload, image): + api_key = self.getAPIKey() + url = f"{self.base_url}/supernode/joycaption2" + 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 = self.send_post_request(url=url, payload=payload, headers=headers) + ret = json.loads(ret) + + try: + if "result" in ret: + ret = json.loads(ret["result"]) + except Exception as e: + raise Exception(f"Unexpected response: {ret} {e=}") + + if ret["type"] == "error": + raise Exception(ret["message"]) + + msg = ret["data"] + if msg["type"] not in ("comfyair", "bizyair",): + raise Exception(f"Unexpected response type: {msg}") + + caption = msg["data"] + + return caption + +bizyairAPI = BizyAIRAPI() + + + +BIZYAIR_DEBUG = True +# Marker to identify base64-encoded tensors +TENSOR_MARKER = "TENSOR:" +IMAGE_MARKER = "IMAGE:" + + +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) +@decode_data.register(type(None)) +def _(input, **kwargs): + return input + + +@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 diff --git a/py/libs/fluxai.py b/py/libs/api/fluxai.py similarity index 96% rename from py/libs/fluxai.py rename to py/libs/api/fluxai.py index 3e41d12..eaaa4d4 100644 --- a/py/libs/fluxai.py +++ b/py/libs/api/fluxai.py @@ -5,7 +5,7 @@ import requests import pathlib from aiohttp import web -root_path = pathlib.Path(__file__).parent.parent.parent +root_path = pathlib.Path(__file__).parent.parent.parent.parent config_path = os.path.join(root_path,'config.yaml') class FluxAIAPI: def __init__(self): diff --git a/py/libs/stability.py b/py/libs/api/stability.py similarity index 97% rename from py/libs/stability.py rename to py/libs/api/stability.py index b0354ff..3b40db2 100644 --- a/py/libs/stability.py +++ b/py/libs/api/stability.py @@ -5,11 +5,11 @@ import requests import pathlib from aiohttp import web from server import PromptServer -from .image import tensor2pil, pil2tensor, image2base64, pil2byte -from .log import log_node_error +from ..image import tensor2pil, pil2tensor, image2base64, pil2byte +from ..log import log_node_error -root_path = pathlib.Path(__file__).parent.parent.parent +root_path = pathlib.Path(__file__).parent.parent.parent.parent config_path = os.path.join(root_path,'config.yaml') default_key = [{'name':'Default', 'key':''}] diff --git a/py/nodes/api.py b/py/nodes/api.py index 0d5da5b..578a51a 100644 --- a/py/nodes/api.py +++ b/py/nodes/api.py @@ -1,6 +1,7 @@ - - -from ..libs.fluxai import fluxaiAPI +import comfy.utils +from ..libs.api.fluxai import fluxaiAPI +from ..libs.api.bizyair import bizyairAPI, encode_data +from nodes import NODE_CLASS_MAPPINGS as ALL_NODE_CLASS_MAPPINGS class fluxPromptGenAPI: @@ -32,10 +33,129 @@ class fluxPromptGenAPI: prompt = fluxaiAPI.promptGenerate(describe, cookie_override) return (prompt,) +class joyCaption2API: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "do_sample": ([True, False],), + "temperature": ( + "FLOAT", + { + "default": 0.5, + "min": 0.0, + "max": 2.0, + "step": 0.01, + "round": 0.001, + "display": "number", + }, + ), + "max_tokens": ( + "INT", + { + "default": 256, + "min": 16, + "max": 512, + "step": 16, + "display": "number", + }, + ), + "caption_type": ( + [ + "Descriptive", + "Descriptive (Informal)", + "Training Prompt", + "MidJourney", + "Booru tag list", + "Booru-like tag list", + "Art Critic", + "Product Listing", + "Social Media Post", + ], + ), + "caption_length": ( + ["any", "very short", "short", "medium-length", "long", "very long"] + + [str(i) for i in range(20, 261, 10)], + ), + "extra_options": ( + "STRING", + { + "placeholder": "Extra options(e.g):\nIf there is a person/character in the image you must refer to them as {name}.", + "tooltip": "Extra options for the model", + "multiline": True, + }, + ), + "name_input": ( + "STRING", + { + "default": "", + "tooltip": "Name input is only used if an Extra Option is selected that requires it.", + }, + ), + "custom_prompt": ( + "STRING", + { + "default": "", + "multiline": True, + }, + ), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("caption",) + + FUNCTION = "joycaption2" + OUTPUT_NODE = False + + CATEGORY = "EasyUse/API" + + def joycaption2( + self, + image, + do_sample, + temperature, + max_tokens, + caption_type, + caption_length, + extra_options, + name_input, + custom_prompt, + ): + pbar = comfy.utils.ProgressBar(100) + pbar.update_absolute(10) + SIZE_LIMIT = 1536 + _, w, h, c = image.shape + if w > SIZE_LIMIT or h > SIZE_LIMIT: + node_class = ALL_NODE_CLASS_MAPPINGS['easy imageScaleDownToSize'] + image, = node_class().image_scale_down_to_size(image, SIZE_LIMIT, True) + + payload = { + "image": None, + "do_sample": do_sample == True, + "temperature": temperature, + "max_new_tokens": max_tokens, + "caption_type": caption_type, + "caption_length": caption_length, + "extra_options": [extra_options], + "name_input": name_input, + "custom_prompt": custom_prompt, + } + + pbar.update_absolute(30) + caption = bizyairAPI.joyCaption2(payload, image) + + pbar.update_absolute(100) + return (caption,) + NODE_CLASS_MAPPINGS = { "easy fluxPromptGenAPI": fluxPromptGenAPI, + "easy joyCaption2API": joyCaption2API, } NODE_DISPLAY_NAME_MAPPINGS = { - "easy fluxPromptGenAPI": "Flux Prompt Gen (API)", + "easy fluxPromptGenAPI": "Prompt Gen (FluxAI)", + "easy joyCaption2API": "JoyCaption2 (BizyAIR)", } \ No newline at end of file diff --git a/py/nodes/deprecated.py b/py/nodes/deprecated.py index dbe516a..8017254 100644 --- a/py/nodes/deprecated.py +++ b/py/nodes/deprecated.py @@ -336,7 +336,7 @@ class injectNoiseToLatent: return (samples,) -from ..libs.stability import stableAPI +from ..libs.api.stability import stableAPI class stableDiffusion3API: @classmethod