bizyengine 0.4.2__py3-none-any.whl

bizyengine/misc/route_sam.py

@@ -0,0 +1,60 @@
+import json
+import os
+from enum import Enum
+
+from aiohttp import web
+
+BIZYAIR_DEBUG = os.getenv("BIZYAIR_DEBUG", False)
+CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
+
+SHOW_CASES = {}
+SAM_COORDINATE = {}
+IS_RESET_SAM = False
+
+
+class EDIT_MODE(Enum):
+    box = 0
+    point = 1
+
+
+from server import PromptServer
+
+
+@PromptServer.instance.routes.post("/bizyair/postsam")
+async def save_sam(request):
+    global IS_RESET_SAM
+    IS_RESET_SAM = False
+    post = await request.post()
+    SAM_COORDINATE["nums"] = post.get("nums")
+    SAM_COORDINATE["mode"] = json.loads(post.get("mode"))
+    if SAM_COORDINATE["mode"] == EDIT_MODE.point.value:
+        SAM_COORDINATE["point_coords"] = json.loads(post.get("coords"))
+    elif SAM_COORDINATE["mode"] == EDIT_MODE.box.value:
+        SAM_COORDINATE["box_coords"] = json.loads(post.get("coords"))
+    SAM_COORDINATE["filename"] = post.get("filename")
+
+    return web.Response(status=200)
+
+
+@PromptServer.instance.routes.get("/bizyair/getsam")
+async def get_sam(request):
+    return web.Response(
+        text=json.dumps(SAM_COORDINATE, ensure_ascii=False),
+        content_type="application/json",
+    )
+
+
+@PromptServer.instance.routes.get("/bizyair/resetsam")
+async def reset_sam(request):
+    global IS_RESET_SAM
+    IS_RESET_SAM = True
+    return web.Response(status=200)
+
+
+@PromptServer.instance.routes.get("/bizyair/isresetsam")
+async def isreset_sam(request):
+    status = {"isresetsam": IS_RESET_SAM}
+    return web.Response(
+        text=json.dumps(status),
+        content_type="application/json",
+    )

bizyengine/misc/segment_anything.py

@@ -0,0 +1,276 @@
+import hashlib
+import json
+import os
+from enum import Enum
+
+import folder_paths
+import numpy as np
+import torch
+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):
+    auto = 0
+    text = 1
+    points_box = 2
+    batched_boxes = 3
+
+
+class EDIT_MODE(Enum):
+    box = 0
+    point = 1
+
+
+class BizyAirSegmentAnythingText:
+    API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/sam"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "prompt": ("STRING", {}),
+                "box_threshold": (
+                    "FLOAT",
+                    {"default": 0.3, "min": 0, "max": 1.0, "step": 0.01},
+                ),
+                "text_threshold": (
+                    "FLOAT",
+                    {"default": 0.3, "min": 0, "max": 1.0, "step": 0.01},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE", "MASK")
+    FUNCTION = "text_sam"
+
+    CATEGORY = "☁️BizyAir/segment-anything"
+
+    def text_sam(self, image, prompt, box_threshold, text_threshold):
+        API_KEY = get_api_key()
+        SIZE_LIMIT = 1536
+        device = image.device
+        _, w, h, c = image.shape
+        assert (
+            w <= SIZE_LIMIT and h <= SIZE_LIMIT
+        ), f"width and height must be less than {SIZE_LIMIT}x{SIZE_LIMIT}, but got {w} and {h}"
+
+        payload = {
+            "image": None,
+            "mode": 1,  # 文本分割模式
+            "params": {
+                "prompt": prompt,
+                "box_threshold": box_threshold,
+                "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)
+
+        try:
+            if "result" in ret:
+                ret = json.loads(ret["result"])
+        except Exception as e:
+            raise Exception(f"Unexpected response: {ret} {e=}")
+
+        if ret["status"] == "error":
+            raise Exception(ret["message"])
+
+        msg = ret["data"]
+        if msg["type"] not in ("bizyair",):
+            raise Exception(f"Unexpected response type: {msg}")
+
+        if "error" in msg:
+            raise Exception(f"Error happens: {msg}")
+
+        img = msg["image"]
+        mask_image = msg["mask_image"]
+
+        img = (
+            (torch.from_numpy(decode_base64_to_np(img)).float() / 255.0)
+            .unsqueeze(0)
+            .to(device)
+        )
+        img_mask = (
+            torch.from_numpy(decode_base64_to_np(mask_image)).float() / 255.0
+        ).to(device)
+        img_mask = img_mask.mean(dim=-1)
+        img_mask = img_mask.unsqueeze(0)
+
+        return (img, img_mask)
+
+
+class BizyAirSegmentAnythingPointBox:
+    API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/sam"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        input_dir = folder_paths.get_input_directory()
+        files = [
+            f
+            for f in os.listdir(input_dir)
+            if os.path.isfile(os.path.join(input_dir, f))
+        ]
+        return {
+            "required": {
+                "image": (sorted(files), {"image_upload": True}),
+                "is_point": ("BOOLEAN", {"default": True}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE", "MASK", "IMAGE")
+    RETURN_NAMES = ("processed_image", "mask", "original_image")
+    FUNCTION = "apply"
+
+    CATEGORY = "☁️BizyAir/segment-anything"
+
+    def apply(self, image, is_point):
+        API_KEY = get_api_key()
+        SIZE_LIMIT = 1536
+
+        # 加载原始图像
+        original_image, _ = LoadImage().load_image(image)
+        # 直接克隆原始图像用于处理和透传
+        image_to_process = original_image.clone()
+
+        device = image_to_process.device
+        _, w, h, c = image_to_process.shape
+        assert (
+            w <= SIZE_LIMIT and h <= SIZE_LIMIT
+        ), f"width and height must be less than {SIZE_LIMIT}x{SIZE_LIMIT}, but got {w} and {h}"
+
+        if is_point:
+            coordinates = [
+                eval(SAM_COORDINATE["point_coords"][key])
+                for key in SAM_COORDINATE["point_coords"]
+            ]
+
+            input_points = [
+                [float(coord["startx"]), float(coord["starty"])]
+                for coord in coordinates
+            ]
+
+            input_label = [coord["pointType"] for coord in coordinates]
+            payload = {
+                "image": None,
+                "mode": INFER_MODE.points_box.value,
+                "params": {
+                    "input_points": json.dumps(input_points),
+                    "input_label": json.dumps(input_label),
+                    "input_boxes": None,
+                },
+            }
+        else:
+            coordinates = [
+                eval(SAM_COORDINATE["box_coords"][key])
+                for key in SAM_COORDINATE["box_coords"]
+            ]
+            input_box = [
+                [
+                    float(coord["startx"]),
+                    float(coord["starty"]),
+                    float(coord["endx"]),
+                    float(coord["endy"]),
+                ]
+                for coord in coordinates
+            ]
+
+            payload = {
+                "image": None,
+                "mode": INFER_MODE.batched_boxes.value,
+                "params": {
+                    "input_points": None,
+                    "input_label": None,
+                    "input_boxes": json.dumps(input_box),
+                },
+            }
+
+        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
+
+        ret: str = send_post_request(self.API_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["status"] == "error":
+            raise Exception(ret["message"])
+
+        msg = ret["data"]
+        if msg["type"] not in ("bizyair",):
+            raise Exception(f"Unexpected response type: {msg}")
+
+        if "error" in msg:
+            raise Exception(f"Error happens: {msg}")
+
+        img = msg["image"]
+        mask_image = msg["mask_image"]
+
+        processed_img = (
+            (torch.from_numpy(decode_base64_to_np(img)).float() / 255.0)
+            .unsqueeze(0)
+            .to(device)
+        )
+        img_mask = (
+            torch.from_numpy(decode_base64_to_np(mask_image)).float() / 255.0
+        ).to(device)
+        img_mask = img_mask.mean(dim=-1)
+        img_mask = img_mask.unsqueeze(0)
+
+        # 直接返回克隆的原始图像,无需转换
+        return (processed_img, img_mask, image_to_process)
+
+    @classmethod
+    def IS_CHANGED(s, image, is_point):
+        image_path = folder_paths.get_annotated_filepath(image)
+        m = hashlib.sha256()
+        with open(image_path, "rb") as f:
+            m.update(f.read())
+        return m.digest().hex()
+
+    @classmethod
+    def VALIDATE_INPUTS(s, image, is_point):
+        if not folder_paths.exists_annotated_filepath(image):
+            return "Invalid image file: {}".format(image)
+
+        return True
+
+
+NODE_CLASS_MAPPINGS = {
+    "BizyAirSegmentAnythingText": BizyAirSegmentAnythingText,
+    "BizyAirSegmentAnythingPointBox": BizyAirSegmentAnythingPointBox,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "BizyAirSegmentAnythingText": "☁️BizyAir Text Guided SAM",
+    "BizyAirSegmentAnythingPointBox": "☁️BizyAir Point-Box Guided SAM",
+}

bizyengine/misc/supernode.py

@@ -0,0 +1,182 @@
+import hashlib
+import json
+import os
+import uuid
+from enum import Enum
+
+import folder_paths
+import node_helpers
+import numpy as np
+import torch
+from bizyengine.core.common.env_var import BIZYAIR_SERVER_ADDRESS
+from bizyengine.core.image_utils import (
+    decode_base64_to_np,
+    decode_data,
+    encode_data,
+    encode_image_to_base64,
+)
+from nodes import LoadImage
+from PIL import Image, ImageOps, ImageSequence
+
+from .utils import (
+    decode_and_deserialize,
+    get_api_key,
+    send_post_request,
+    serialize_and_encode,
+)
+
+
+class RemoveBackground:
+    API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/removebg"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE", "MASK")
+    FUNCTION = "remove_background"
+
+    CATEGORY = "☁️BizyAir"
+
+    def remove_background(self, image):
+        API_KEY = get_api_key()
+        device = image.device
+        _, h, w, _ = image.shape
+        assert (
+            w <= 1536 and h <= 1536
+        ), f"width and height must be less than 1536, but got {w} and {h}"
+
+        payload = {
+            "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
+        )
+        tensors = decode_and_deserialize(response)
+        t_images = tensors["images"].to(device)
+        t_mask = tensors["mask"].to(device)
+        return (t_images, t_mask)
+
+
+class GenerateLightningImage:
+    API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/realvis4lightning"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": (
+                    "STRING",
+                    {"multiline": True, "dynamicPrompts": True, "default": "a dog"},
+                ),
+                "seed": ("INT", {"default": 1, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
+                "width": ("INT", {"default": 1024, "min": 16, "max": 1024, "step": 8}),
+                "height": ("INT", {"default": 1024, "min": 16, "max": 1024, "step": 8}),
+                "cfg": (
+                    "FLOAT",
+                    {
+                        "default": 1.5,
+                        "min": 0.0,
+                        "max": 10.0,
+                        "step": 0.1,
+                        "round": 0.01,
+                    },
+                ),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 4}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "generate_image"
+
+    CATEGORY = "☁️BizyAir"
+
+    def generate_image(self, prompt, seed, width, height, cfg, batch_size):
+        API_KEY = get_api_key()
+        assert (
+            width <= 1024 and height <= 1024
+        ), f"width and height must be less than 1024, but got {width} and {height}"
+
+        payload = {
+            "batch_size": batch_size,
+            "width": width,
+            "height": height,
+            "prompt": prompt,
+            "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
+        )
+        tensors_np = decode_and_deserialize(response)
+        tensors = torch.from_numpy(tensors_np)
+
+        return (tensors,)
+
+
+class ToggleServerEndpoint:
+    BIZYAIR_SERVER_ENDPOINTS = [
+        "https://bizyair-api-st.siliconflow.cn/x/v1",
+        "https://bizyair-api.siliconflow.cn/x/v1",
+    ]
+
+    def __init__(self):
+        self.current_index = 0
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "endpoint": (s.BIZYAIR_SERVER_ENDPOINTS,),
+            }
+        }
+
+    RETURN_TYPES = ()
+    FUNCTION = "toggle_endpoint"
+    OUTPUT_NODE = True
+
+    CATEGORY = "☁️BizyAir"
+
+    def toggle_endpoint(self, endpoint):
+        BIZYAIR_SERVER_ADDRESS.address = endpoint
+        from server import PromptServer
+
+        PromptServer.instance.send_sync(
+            "bizyair.server.endpoint.switch",
+            {"message": f"Switch server endpoint to {endpoint}"},
+        )
+        return ()
+
+
+NODE_CLASS_MAPPINGS = {
+    "BizyAirRemoveBackground": RemoveBackground,
+    "BizyAirGenerateLightningImage": GenerateLightningImage,
+    "BizyAirToggleServerEndpoint": ToggleServerEndpoint,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "BizyAirRemoveBackground": "☁️BizyAir Remove Image Background",
+    "BizyAirGenerateLightningImage": "☁️BizyAir Generate Photorealistic Images",
+    "BizyAirToggleServerEndpoint": "☁️BizyAir Switch Server Endpoint",
+}