Upload 13 files

SwarmRemBg.py

@@ -0,0 +1,38 @@
+from PIL import Image
+import numpy as np
+import torch
+
+class SwarmRemBg:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+            }
+        }
+
+    CATEGORY = "SwarmUI/images"
+    RETURN_TYPES = ("IMAGE", "MASK",)
+    FUNCTION = "rem"
+
+    def rem(self, images):
+        from rembg import remove
+
+        output = []
+        masks = []
+        for image in images:
+            i = 255.0 * image.cpu().numpy()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+            img = img.convert("RGBA")
+            img = remove(img, post_process_mask=True)
+            output.append(np.array(img).astype(np.float32) / 255.0)
+            if 'A' in img.getbands():
+                mask = np.array(img.getchannel('A')).astype(np.float32) / 255.0
+                masks.append(1. - mask)
+            else:
+                masks.append(np.zeros((64,64), dtype=np.float32))
+        return (torch.from_numpy(np.array(output)), torch.from_numpy(np.array(masks)))
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmRemBg": SwarmRemBg,
+}

SwarmSaveAnimationWS.py

@@ -0,0 +1,118 @@
+import comfy, folder_paths, io, struct, subprocess, os, random, sys, time
+from PIL import Image
+import numpy as np
+from server import PromptServer, BinaryEventTypes
+from imageio_ffmpeg import get_ffmpeg_exe
+
+SPECIAL_ID = 12345
+VIDEO_ID = 12346
+FFMPEG_PATH = get_ffmpeg_exe()
+
+
+class SwarmSaveAnimationWS:
+    methods = {"default": 4, "fastest": 0, "slowest": 6}
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE", ),
+                "fps": ("FLOAT", {"default": 6.0, "min": 0.01, "max": 1000.0, "step": 0.01}),
+                "lossless": ("BOOLEAN", {"default": True}),
+                "quality": ("INT", {"default": 80, "min": 0, "max": 100}),
+                "method": (list(s.methods.keys()),),
+                "format": (["webp", "gif", "gif-hd", "h264-mp4", "h265-mp4", "webm", "prores"],),
+            },
+        }
+
+    CATEGORY = "SwarmUI/video"
+    RETURN_TYPES = ()
+    FUNCTION = "save_images"
+    OUTPUT_NODE = True
+
+    def save_images(self, images, fps, lossless, quality, method, format):
+        method = self.methods.get(method)
+        if images.shape[0] == 0:
+            return { }
+        if images.shape[0] == 1:
+            pbar = comfy.utils.ProgressBar(SPECIAL_ID)
+            i = 255.0 * images[0].cpu().numpy()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+            pbar.update_absolute(0, SPECIAL_ID, ("PNG", img, None))
+            return { }
+
+        out_img = io.BytesIO()
+        if format in ["webp", "gif"]:
+            if format == "webp":
+                type_num = 3
+            else:
+                type_num = 4
+            pil_images = []
+            for image in images:
+                i = 255. * image.cpu().numpy()
+                img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+                pil_images.append(img)
+            pil_images[0].save(out_img, save_all=True, duration=int(1000.0 / fps), append_images=pil_images[1 : len(pil_images)], lossless=lossless, quality=quality, method=method, format=format.upper(), loop=0)
+        else:
+            i = 255. * images.cpu().numpy()
+            raw_images = np.clip(i, 0, 255).astype(np.uint8)
+            args = [FFMPEG_PATH, "-v", "error", "-f", "rawvideo", "-pix_fmt", "rgb24",
+                    "-s", f"{len(raw_images[0][0])}x{len(raw_images[0])}", "-r", str(fps), "-i", "-", "-n" ]
+            if format == "h264-mp4":
+                args += ["-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "19"]
+                ext = "mp4"
+                type_num = 5
+            elif format == "h265-mp4":
+                args += ["-c:v", "libx265", "-pix_fmt", "yuv420p"]
+                ext = "mp4"
+                type_num = 5
+            elif format == "webm":
+                args += ["-pix_fmt", "yuv420p", "-crf", "23"]
+                ext = "webm"
+                type_num = 6
+            elif format == "prores":
+                args += ["-c:v", "prores_ks", "-profile:v", "3", "-pix_fmt", "yuv422p10le"]
+                ext = "mov"
+                type_num = 7
+            elif format == "gif-hd":
+                args += ["-filter_complex", "split=2 [a][b]; [a] palettegen [pal]; [b] [pal] paletteuse"]
+                ext = "gif"
+                type_num = 4
+            path = folder_paths.get_save_image_path("swarm_tmp_", folder_paths.get_temp_directory())[0]
+            rand = '%016x' % random.getrandbits(64)
+            file = os.path.join(path, f"swarm_tmp_{rand}.{ext}")
+            result = subprocess.run(args + [file], input=raw_images.tobytes(), capture_output=True)
+            if result.returncode != 0:
+                print(f"ffmpeg failed with return code {result.returncode}", file=sys.stderr)
+                f_out = result.stdout.decode("utf-8").strip()
+                f_err = result.stderr.decode("utf-8").strip()
+                if f_out:
+                    print("ffmpeg out: " + f_out, file=sys.stderr)
+                if f_err:
+                    print("ffmpeg error: " + f_err, file=sys.stderr)
+                raise Exception(f"ffmpeg failed: {f_err}")
+            # TODO: Is there a way to get ffmpeg to operate entirely in memory?
+            with open(file, "rb") as f:
+                out_img.write(f.read())
+            os.remove(file)
+
+        out = io.BytesIO()
+        header = struct.pack(">I", type_num)
+        out.write(header)
+        out.write(out_img.getvalue())
+        out.seek(0)
+        preview_bytes = out.getvalue()
+        server = PromptServer.instance
+        server.send_sync("progress", {"value": 12346, "max": 12346}, sid=server.client_id)
+        server.send_sync(BinaryEventTypes.PREVIEW_IMAGE, preview_bytes, sid=server.client_id)
+
+        return { }
+
+    @classmethod
+    def IS_CHANGED(s, images, fps, lossless, quality, method, format):
+        return time.time()
+
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmSaveAnimationWS": SwarmSaveAnimationWS,
+}

SwarmYolo.py

@@ -0,0 +1,116 @@
+import torch, folder_paths, comfy
+from PIL import Image
+import numpy as np
+
+class SwarmYoloDetection:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "model_name": (folder_paths.get_filename_list("yolov8"), ),
+                "index": ("INT", { "default": 0, "min": 0, "max": 256, "step": 1 }),
+            },
+            "optional": {
+                "class_filter": ("STRING", { "default": "", "multiline": False }),
+                "sort_order": (["left-right", "right-left", "top-bottom", "bottom-top", "largest-smallest", "smallest-largest"], ),
+                "threshold": ("FLOAT", { "default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01 }),
+            }
+        }
+
+    CATEGORY = "SwarmUI/masks"
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "seg"
+
+    def seg(self, image, model_name, index, class_filter=None, sort_order="left-right", threshold=0.25):
+        # TODO: Batch support?
+        i = 255.0 * image[0].cpu().numpy()
+        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+        # TODO: Cache the model in RAM in some way?
+        model_path = folder_paths.get_full_path("yolov8", model_name)
+        if model_path is None:
+            raise ValueError(f"Model {model_name} not found, or yolov8 folder path not defined")
+        from ultralytics import YOLO
+        model = YOLO(model_path)
+        results = model.predict(img, conf=threshold)
+        boxes = results[0].boxes
+        class_ids = boxes.cls.cpu().numpy() if boxes is not None else []
+        selected_classes = None
+
+        if class_filter and class_filter.strip():
+            class_filter_list = [cls_name.strip() for cls_name in class_filter.split(",") if cls_name.strip()]
+            label_to_id = {name.lower(): id for id, name in model.names.items()}
+            selected_classes = []
+            for cls_name in class_filter_list:
+                if cls_name.isdigit():
+                    selected_classes.append(int(cls_name))
+                else:
+                    class_id = label_to_id.get(cls_name.lower())
+                    if class_id is not None:
+                        selected_classes.append(class_id)
+                    else:
+                        print(f"Class '{cls_name}' not found in the model")
+            selected_classes = selected_classes if selected_classes else None
+
+        masks = results[0].masks
+        if masks is not None and selected_classes is not None:
+            selected_masks = []
+            for i, class_id in enumerate(class_ids):
+                if class_id in selected_classes:
+                    selected_masks.append(masks.data[i].cpu())
+            if selected_masks:
+                masks = torch.stack(selected_masks)
+            else:
+                masks = None
+
+        if masks is None or masks.shape[0] == 0:
+            if boxes is None or len(boxes) == 0:
+                return (torch.zeros(1, image.shape[1], image.shape[2]), )
+            else:
+                if selected_classes:
+                    boxes = [box for i, box in enumerate(boxes) if class_ids[i] in selected_classes]
+            masks = torch.zeros((len(boxes), image.shape[1], image.shape[2]), dtype=torch.float32, device="cpu")
+            for i, box in enumerate(boxes):
+                x1, y1, x2, y2 = box.xyxy[0].tolist()
+                masks[i, int(y1):int(y2), int(x1):int(x2)] = 1.0
+        else:
+            masks = masks.data.cpu()
+        if masks is None or masks.shape[0] == 0:
+            return (torch.zeros(1, image.shape[1], image.shape[2]), )
+
+        masks = torch.nn.functional.interpolate(masks.unsqueeze(1), size=(image.shape[1], image.shape[2]), mode="bilinear").squeeze(1)
+        if index == 0:
+            result = masks[0]
+            for i in range(1, len(masks)):
+                result = torch.max(result, masks[i])
+            return (result.unsqueeze(0), )
+        elif index > len(masks):
+            return (torch.zeros_like(masks[0]).unsqueeze(0), )
+        else:
+            sortedindices = []
+            for mask in masks:
+                match sort_order:
+                    case "left-right":
+                        sum_x = (torch.sum(mask, dim=0) != 0).to(dtype=torch.int)
+                        val = torch.argmax(sum_x).item()
+                    case "right-left":
+                        sum_x = (torch.sum(mask, dim=0) != 0).to(dtype=torch.int)
+                        val = mask.shape[1] - torch.argmax(torch.flip(sum_x, [0])).item() - 1
+                    case "top-bottom":
+                        sum_y = (torch.sum(mask, dim=1) != 0).to(dtype=torch.int)
+                        val = torch.argmax(sum_y).item()
+                    case "bottom-top":
+                        sum_y = (torch.sum(mask, dim=1) != 0).to(dtype=torch.int)
+                        val = mask.shape[0] - torch.argmax(torch.flip(sum_y, [0])).item() - 1
+                    case "largest-smallest" | "smallest-largest":
+                        val = torch.sum(mask).item()
+                sortedindices.append(val)
+            sortedindices = np.argsort(sortedindices)
+            if sort_order in ["right-left", "bottom-top", "largest-smallest"]:
+                sortedindices = sortedindices[::-1].copy()
+            masks = masks[sortedindices]
+            return (masks[index - 1].unsqueeze(0), )
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmYoloDetection": SwarmYoloDetection,
+}

__init__.py

@@ -0,0 +1,25 @@
+import traceback
+
+NODE_CLASS_MAPPINGS = {}
+
+# RemBg doesn't work on all python versions and OS's
+try:
+    from . import SwarmRemBg
+    NODE_CLASS_MAPPINGS.update(SwarmRemBg.NODE_CLASS_MAPPINGS)
+except ImportError:
+    print("Error: [Swarm] RemBg not available")
+    traceback.print_exc()
+# This uses FFMPEG which doesn't install itself properly on Macs I guess?
+try:
+    from . import SwarmSaveAnimationWS
+    NODE_CLASS_MAPPINGS.update(SwarmSaveAnimationWS.NODE_CLASS_MAPPINGS)
+except ImportError:
+    print("Error: [Swarm] SaveAnimationWS not available")
+    traceback.print_exc()
+# Yolo uses Ultralytics, which is cursed
+try:
+    from . import SwarmYolo
+    NODE_CLASS_MAPPINGS.update(SwarmYolo.NODE_CLASS_MAPPINGS)
+except ImportError:
+    print("Error: [Swarm] Yolo not available")
+    traceback.print_exc()

requirements.txt

@@ -0,0 +1,3 @@
+rembg
+dill
+ultralytics