add automatic marking nodes

__init__.py

@@ -16,6 +16,7 @@
     llm,
     mzkolors,
     nodes,
+    nodes_automatic_marking,
     nodes_controlnet_aux,
     nodes_controlnet_union_sdxl,
     segment_anything,
@@ -36,6 +37,7 @@ def update_mappings(module):
 update_mappings(nodes_controlnet_union_sdxl)
 update_mappings(mzkolors)
 update_mappings(segment_anything)
+update_mappings(nodes_automatic_marking)
 
 try:
     import bizy_server

llm.py

@@ -282,6 +282,7 @@ def __init__(self):
 
     # refer to: https://huggingface.co/spaces/fancyfeast/joy-caption-pre-alpha
     API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/joycaption2"
+    # API_URL = "https://maas.platform.oneflow.cloud/supernode/crossing-joycaption2/supernode/crossing-joycaption2"
 
     @classmethod
     def INPUT_TYPES(s):
@@ -353,12 +354,28 @@ def INPUT_TYPES(s):
         }
 
     RETURN_TYPES = ("STRING",)
-    FUNCTION = "joycaption2"
+    FUNCTION = "my_joycaption"
 
     CATEGORY = "☁️BizyAir/AI Assistants"
 
-    def joycaption2(
+    async def send_post_request_async(
+        self, session, url, payload, headers, max_retries=3, retry_delay=2, timeout=100
+    ):
+        try:
+            async with session.post(
+                url, json=payload, headers=headers, timeout=timeout
+            ) as response:
+                response.raise_for_status()
+                ret = await response.json()
+                return ret
+        except Exception as e:
+
+            print(f"Request failed. Error: {e}")
+            return {"data": {"type": "bizyair", "data": ""}}
+
+    async def joycaption2(
         self,
+        session,
         image,
         do_sample,
         temperature,
@@ -371,6 +388,8 @@ def joycaption2(
     ):
         API_KEY = get_api_key()
         SIZE_LIMIT = 1536
+        max_retries = 3
+        retry_delay = 2
         _, w, h, c = image.shape
         assert (
             w <= SIZE_LIMIT and h <= SIZE_LIMIT
@@ -396,17 +415,15 @@ def joycaption2(
         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)
-
-        try:
-            if "result" in ret:
-                ret = json.loads(ret["result"])
-            if ret["type"] == "error":
-                raise Exception(ret["message"])
-        except Exception as e:
-            raise Exception(f"Unexpected response: {ret} {e=}")
-
+        ret = await self.send_post_request_async(
+            session,
+            self.API_URL,
+            payload=payload,
+            headers=headers,
+            max_retries=max_retries,
+            retry_delay=retry_delay,
+        )
+        ret = json.loads(ret["result"])
         msg = ret["data"]
         if msg["type"] not in (
             "comfyair",
@@ -417,6 +434,25 @@ def joycaption2(
         caption = msg["data"]
         return (caption,)
 
+    async def my_joycaption_async(self, image, **kwargs):
+        captions = []
+        tasks = []
+        async with aiohttp.ClientSession() as session:
+            for i in range(image.size(0)):
+                image_file = image[i].unsqueeze(0)
+                tasks.append(self.joycaption2(session, image_file, **kwargs))
+
+            results = await asyncio.gather(*tasks)
+
+        for result in results:
+            captions.append(result[0])
+
+        combined_caption = " | ".join(captions)
+        return {"ui": {"text": (combined_caption,)}, "result": (combined_caption,)}
+
+    def my_joycaption(self, image, **kwargs):
+        return asyncio.run(self.my_joycaption_async(image, **kwargs))
+
 
 NODE_CLASS_MAPPINGS = {
     "BizyAirSiliconCloudLLMAPI": SiliconCloudLLMAPI,

nodes_automatic_marking.py

@@ -0,0 +1,277 @@
+import os
+
+import folder_paths
+import numpy as np
+import torch
+from PIL import Image, ImageOps
+
+
+class SaveCaptionsAndImages:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "captions": ("STRING", {"multiline": True}),
+                "images": ("IMAGE",),
+                "directory_prefix": (
+                    "STRING",
+                    {"default": "lora_dataset", "multiline": False},
+                ),
+            },
+        }
+
+    RETURN_TYPES = ()
+    OUTPUT_NODE = True
+    FUNCTION = "apply"
+
+    def apply(self, captions, images, directory_prefix):
+
+        # Split the captions string into a list using " | " as the delimiter
+        caption_list = captions.split(" | ")
+        full_output_folder = folder_paths.get_output_directory()
+        # Find the next available directory number
+        i = 0
+        while True:
+            dir_path = os.path.join(full_output_folder, f"{directory_prefix}_{i:03d}")
+            if not os.path.exists(dir_path):
+                break
+            i += 1
+        # Validate input
+        if len(caption_list) != len(images):
+            raise ValueError(
+                "The number of captions does not match the number of images."
+            )
+
+        for batch_number, (image, caption) in enumerate(zip(images, caption_list)):
+            # Generate a unique filename for each image
+            filename = f"image_{batch_number:04d}"
+
+            # Generate file paths
+            image_filepath = os.path.join(dir_path, f"{filename}.png")
+            caption_filepath = os.path.join(dir_path, f"{filename}.txt")
+
+            # Ensure directory exists
+            os.makedirs(dir_path, exist_ok=True)
+
+            # Save the image
+            i = 255.0 * image.cpu().numpy()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+            img.save(image_filepath)
+
+            # Write caption to file
+            with open(caption_filepath, "w", encoding="utf-8") as caption_file:
+                caption_file.write(caption)
+
+            print(f"Image saved to: {image_filepath}")
+            print(f"Caption saved to: {caption_filepath}")
+
+        return {}
+
+
+class BizyAirLoadImagesFromFolder:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "folder": ("STRING", {"default": ""}),
+                "width": ("INT", {"default": 1024, "min": 64, "step": 1}),
+                "height": ("INT", {"default": 1024, "min": 64, "step": 1}),
+                "keep_aspect_ratio": (
+                    [
+                        "crop",
+                        "pad",
+                        "stretch",
+                    ],
+                ),
+            },
+            "optional": {
+                "image_load_cap": ("INT", {"default": 100, "min": 0, "step": 1}),
+                "start_index": ("INT", {"default": 0, "min": 0, "step": 1}),
+                "include_subfolders": ("BOOLEAN", {"default": False}),
+            },
+        }
+
+    RETURN_TYPES = (
+        "IMAGE",
+        "MASK",
+        "INT",
+        "STRING",
+    )
+    RETURN_NAMES = (
+        "image",
+        "mask",
+        "count",
+        "image_path",
+    )
+    FUNCTION = "load_images"
+    CATEGORY = "☁️BizyAir/marking"
+    DESCRIPTION = """Loads images from a folder into a batch, images are resized and loaded into a batch."""
+
+    def load_images(
+        self,
+        folder,
+        width,
+        height,
+        image_load_cap,
+        start_index,
+        keep_aspect_ratio,
+        include_subfolders=False,
+    ):
+        if not os.path.isdir(folder):
+            raise FileNotFoundError(f"Folder '{folder} cannot be found.'")
+
+        valid_extensions = [".jpg", ".jpeg", ".png", ".webp"]
+        image_paths = []
+        if include_subfolders:
+            for root, _, files in os.walk(folder):
+                for file in files:
+                    if any(file.lower().endswith(ext) for ext in valid_extensions):
+                        image_paths.append(os.path.join(root, file))
+        else:
+            for file in os.listdir(folder):
+                if any(file.lower().endswith(ext) for ext in valid_extensions):
+                    image_paths.append(os.path.join(folder, file))
+
+        dir_files = sorted(image_paths)
+
+        if len(dir_files) == 0:
+            raise FileNotFoundError(f"No files in directory '{folder}'.")
+
+        # start at start_index
+        dir_files = dir_files[start_index:]
+
+        images = []
+        masks = []
+        image_path_list = []
+
+        limit_images = False
+        if image_load_cap > 0:
+            limit_images = True
+        image_count = 0
+
+        for image_path in dir_files:
+            if os.path.isdir(image_path):
+                continue
+            if limit_images and image_count >= image_load_cap:
+                break
+            i = Image.open(image_path)
+            i = ImageOps.exif_transpose(i)
+
+            # Resize image to maximum dimensions
+            if i.size != (width, height):
+                i = self.resize_with_aspect_ratio(i, width, height, keep_aspect_ratio)
+
+            image = i.convert("RGB")
+            image = np.array(image).astype(np.float32) / 255.0
+            image = torch.from_numpy(image)[None,]
+
+            if "A" in i.getbands():
+                mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0
+                mask = 1.0 - torch.from_numpy(mask)
+                if mask.shape != (height, width):
+                    mask = torch.nn.functional.interpolate(
+                        mask.unsqueeze(0).unsqueeze(0),
+                        size=(height, width),
+                        mode="bilinear",
+                        align_corners=False,
+                    ).squeeze()
+            else:
+                mask = torch.zeros((height, width), dtype=torch.float32, device="cpu")
+
+            images.append(image)
+            masks.append(mask)
+            image_path_list.append(image_path)
+            image_count += 1
+
+        if len(images) == 1:
+            return (images[0], masks[0], 1, image_path_list)
+
+        elif len(images) > 1:
+            image1 = images[0]
+            mask1 = masks[0].unsqueeze(0)
+
+            for image2 in images[1:]:
+                image1 = torch.cat((image1, image2), dim=0)
+
+            for mask2 in masks[1:]:
+                mask1 = torch.cat((mask1, mask2.unsqueeze(0)), dim=0)
+
+            return (image1, mask1, len(images), image_path_list)
+
+    def resize_with_aspect_ratio(self, img, width, height, mode):
+        if mode == "stretch":
+            return img.resize((width, height), Image.Resampling.LANCZOS)
+
+        img_width, img_height = img.size
+        aspect_ratio = img_width / img_height
+        target_ratio = width / height
+
+        if mode == "crop":
+            # Calculate dimensions for center crop
+            if aspect_ratio > target_ratio:
+                # Image is wider - crop width
+                new_width = int(height * aspect_ratio)
+                img = img.resize((new_width, height), Image.Resampling.LANCZOS)
+                left = (new_width - width) // 2
+                return img.crop((left, 0, left + width, height))
+            else:
+                # Image is taller - crop height
+                new_height = int(width / aspect_ratio)
+                img = img.resize((width, new_height), Image.Resampling.LANCZOS)
+                top = (new_height - height) // 2
+                return img.crop((0, top, width, top + height))
+
+        elif mode == "pad":
+            pad_color = self.get_edge_color(img)
+            # Calculate dimensions for padding
+            if aspect_ratio > target_ratio:
+                # Image is wider - pad height
+                new_height = int(width / aspect_ratio)
+                img = img.resize((width, new_height), Image.Resampling.LANCZOS)
+                padding = (height - new_height) // 2
+                padded = Image.new("RGBA", (width, height), pad_color)
+                padded.paste(img, (0, padding))
+                return padded
+            else:
+                # Image is taller - pad width
+                new_width = int(height * aspect_ratio)
+                img = img.resize((new_width, height), Image.Resampling.LANCZOS)
+                padding = (width - new_width) // 2
+                padded = Image.new("RGBA", (width, height), pad_color)
+                padded.paste(img, (padding, 0))
+                return padded
+
+    def get_edge_color(self, img):
+        from PIL import ImageStat
+
+        """Sample edges and return dominant color"""
+        width, height = img.size
+        img = img.convert("RGBA")
+
+        # Create 1-pixel high/wide images from edges
+        top = img.crop((0, 0, width, 1))
+        bottom = img.crop((0, height - 1, width, height))
+        left = img.crop((0, 0, 1, height))
+        right = img.crop((width - 1, 0, width, height))
+
+        # Combine edges into single image
+        edges = Image.new("RGBA", (width * 2 + height * 2, 1))
+        edges.paste(top, (0, 0))
+        edges.paste(bottom, (width, 0))
+        edges.paste(left.resize((height, 1)), (width * 2, 0))
+        edges.paste(right.resize((height, 1)), (width * 2 + height, 0))
+
+        # Get median color
+        stat = ImageStat.Stat(edges)
+        median = tuple(map(int, stat.median))
+        return median
+
+
+NODE_CLASS_MAPPINGS = {
+    "BizyAirLoadImagesFromFolder": BizyAirLoadImagesFromFolder,
+    "SaveCaptionsAndImages": SaveCaptionsAndImages,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "BizyAirLoadImagesFromFolder": "☁️BizyAir LoadImagesFromFolder",
+    "SaveCaptionsAndImages": "☁️BizyAir Save Captions And Images",
+}