Upload 55 files

SwarmBlending.py

@@ -0,0 +1,39 @@
+import torch
+
+class SwarmLatentBlendMasked:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "samples0": ("LATENT",),
+                "samples1": ("LATENT",),
+                "mask": ("MASK",),
+                "blend_factor": ("FLOAT", { "default": 0.5, "min": 0, "max": 1, "step": 0.01, "tooltip": "The blend factor between the two samples. 0 means entirely use sample0, 1 means entirely sample1, 0.5 means 50/50 of each." }),
+            }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "blend"
+    CATEGORY = "SwarmUI/images"
+    DESCRIPTION = "Blends two latent images together within a masked region."
+
+    def blend(self, samples0, samples1, blend_factor, mask):
+        samples_out = samples0.copy()
+        samples0 = samples0["samples"]
+        samples1 = samples1["samples"]
+        while mask.ndim < 4:
+            mask = mask.unsqueeze(0)
+
+        if samples0.shape != samples1.shape:
+            samples1 = torch.nn.functional.interpolate(samples1, size=(samples0.shape[2], samples0.shape[3]), mode="bicubic")
+        if samples0.shape != mask.shape:
+            mask = torch.nn.functional.interpolate(mask, size=(samples0.shape[2], samples0.shape[3]), mode="bicubic")
+
+        samples_blended = samples0 * (1 - mask * blend_factor) + samples1 * (mask * blend_factor)
+        samples_out["samples"] = samples_blended
+        return (samples_out,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmLatentBlendMasked": SwarmLatentBlendMasked,
+}

SwarmClipSeg.py

@@ -0,0 +1,71 @@
+import torch
+from PIL import Image
+import numpy as np
+from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
+import folder_paths
+import os, requests
+
+def get_path():
+    if "clipseg" in folder_paths.folder_names_and_paths:
+        paths = folder_paths.folder_names_and_paths["clipseg"]
+        return paths[0][0]
+    else:
+        # Jank backup path if you're not running properly in Swarm
+        path = os.path.dirname(os.path.realpath(__file__)) + "/models"
+        return path
+
+
+# Manual download of the model from a safetensors conversion.
+# Done manually to guarantee it's only a safetensors file ever and not a pickle
+def download_model(path, urlbase):
+    if os.path.exists(path):
+        return
+    for file in ["config.json", "merges.txt", "model.safetensors", "preprocessor_config.json", "special_tokens_map.json", "tokenizer_config.json", "vocab.json"]:
+        os.makedirs(path, exist_ok=True)
+        filepath = path + file
+        if not os.path.exists(filepath):
+            with open(filepath, "wb") as f:
+                print(f"[SwarmClipSeg] Downloading '{file}'...")
+                f.write(requests.get(f"{urlbase}{file}").content)
+
+
+class SwarmClipSeg:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "match_text": ("STRING", {"multiline": True, "tooltip": "A short description (a few words) to describe something within the image to find and mask."}),
+                "threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step":0.01, "round": False, "tooltip": "Threshold to apply to the mask, higher values will make the mask more strict. Without sufficient thresholding, CLIPSeg may include random stray content around the edges."}),
+            }
+        }
+
+    CATEGORY = "SwarmUI/masks"
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "seg"
+    DESCRIPTION = "Segment an image using CLIPSeg, creating a mask of what part of an image appears to match the given text."
+
+    def seg(self, images, match_text, threshold):
+        # TODO: Batch support?
+        i = 255.0 * images[0].cpu().numpy()
+        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+        # TODO: Cache the model in RAM in some way?
+        path = get_path() + "/clipseg-rd64-refined-fp16-safetensors/"
+        download_model(path, "https://huggingface.co/mcmonkey/clipseg-rd64-refined-fp16/resolve/main/")
+        processor = CLIPSegProcessor.from_pretrained(path)
+        model = CLIPSegForImageSegmentation.from_pretrained(path)
+        with torch.no_grad():
+            mask = model(**processor(text=match_text, images=img, return_tensors="pt", padding=True))[0]
+        mask = torch.nn.functional.threshold(mask.sigmoid(), threshold, 0)
+        mask -= mask.min()
+        max = mask.max()
+        if max > 0:
+            mask /= max
+        while mask.ndim < 4:
+            mask = mask.unsqueeze(0)
+        mask = torch.nn.functional.interpolate(mask, size=(images.shape[1], images.shape[2]), mode="bilinear").squeeze(0)
+        return (mask,)
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmClipSeg": SwarmClipSeg,
+}

SwarmExtractLora.py

@@ -0,0 +1,146 @@
+import comfy.model_management
+import safetensors.torch
+import torch, os, comfy, json
+
+# ATTRIBUTION: This code is a mix of code from kohya-ss, comfy, and Swarm. It would be annoying to disentangle but it's all FOSS and relatively short so it's fine.
+
+CLAMP_QUANTILE = 0.99
+def extract_lora(diff, rank):
+    conv2d = (len(diff.shape) == 4)
+    kernel_size = None if not conv2d else diff.size()[2:4]
+    conv2d_3x3 = conv2d and kernel_size != (1, 1)
+    out_dim, in_dim = diff.size()[0:2]
+    rank = min(rank, in_dim, out_dim)
+
+    if conv2d:
+        if conv2d_3x3:
+            diff = diff.flatten(start_dim=1)
+        else:
+            diff = diff.squeeze()
+
+    U, S, Vh = torch.linalg.svd(diff.float())
+    U = U[:, :rank]
+    S = S[:rank]
+    U = U @ torch.diag(S)
+    Vh = Vh[:rank, :]
+
+    dist = torch.cat([U.flatten(), Vh.flatten()])
+    hi_val = torch.quantile(dist, CLAMP_QUANTILE)
+    low_val = -hi_val
+
+    U = U.clamp(low_val, hi_val)
+    Vh = Vh.clamp(low_val, hi_val)
+    if conv2d:
+        U = U.reshape(out_dim, rank, 1, 1)
+        Vh = Vh.reshape(rank, in_dim, kernel_size[0], kernel_size[1])
+    return (U, Vh)
+
+
+def do_lora_handle(base_data, other_data, rank, prefix, require, do_bias, callback):
+    out_data = {}
+    device = comfy.model_management.get_torch_device()
+    for key in base_data.keys():
+        callback()
+        if key not in other_data:
+            continue
+        base_tensor = base_data[key].float()
+        other_tensor = other_data[key].float()
+        if key.startswith("clip_g"):
+            key = "1." + key[len("clip_g."):]
+        elif key.startswith("clip_l"):
+            key = "0." + key[len("clip_l."):]
+        if require:
+            if not key.startswith(require):
+                print(f"Ignore unmatched key {key} (doesn't match {require})")
+                continue
+            key = key[len(require):]
+        if base_tensor.shape != other_tensor.shape:
+            continue
+        diff = other_tensor.to(device) - base_tensor.to(device)
+        other_tensor = other_tensor.cpu()
+        base_tensor = base_tensor.cpu()
+        max_diff = float(diff.abs().max())
+        if max_diff < 1e-5:
+            print(f"discard unaltered key {key} ({max_diff})")
+            continue
+        if key.endswith(".weight"):
+            fixed_key = key[:-len(".weight")].replace('.', '_')
+            name = f"lora_{prefix}_{fixed_key}"
+            if len(base_tensor.shape) >= 2:
+                print(f"extract key {name} ({max_diff})")
+                out = extract_lora(diff, rank)
+                out_data[f"{name}.lora_up.weight"] = out[0].contiguous().half().cpu()
+                out_data[f"{name}.lora_down.weight"] = out[1].contiguous().half().cpu()
+            else:
+                print(f"ignore valid raw pass-through key {name} ({max_diff})")
+                #out_data[name] = other_tensor.contiguous().half().cpu()
+        elif key.endswith(".bias") and do_bias:
+            fixed_key = key[:-len(".bias")].replace('.', '_')
+            name = f"lora_{prefix}_{fixed_key}"
+            print(f"extract bias key {name} ({max_diff})")
+            out_data[f"{name}.diff_b"] = diff.contiguous().half().cpu()
+
+
+    return out_data
+
+class SwarmExtractLora:
+    def __init__(self):
+        self.loaded_lora = None
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "base_model": ("MODEL", ),
+                "base_model_clip": ("CLIP", ),
+                "other_model": ("MODEL", ),
+                "other_model_clip": ("CLIP", ),
+                "rank": ("INT", {"default": 16, "min": 1, "max": 320}),
+                "save_rawpath": ("STRING", {"multiline": False}),
+                "save_filename": ("STRING", {"multiline": False}),
+                "save_clip": ("BOOLEAN", {"default": True}),
+                "metadata": ("STRING", {"multiline": True}),
+            }
+        }
+
+    CATEGORY = "SwarmUI/models"
+    RETURN_TYPES = ()
+    FUNCTION = "extract_lora"
+    OUTPUT_NODE = True
+    DESCRIPTION = "Internal node, do not use directly - extracts a LoRA from the difference between two models. This is used by SwarmUI Utilities tab."
+
+    def extract_lora(self, base_model, base_model_clip, other_model, other_model_clip, rank, save_rawpath, save_filename, save_clip, metadata):
+        base_data = base_model.model_state_dict()
+        other_data = other_model.model_state_dict()
+        key_count = len(base_data.keys())
+        if save_clip:
+            key_count += len(base_model_clip.get_sd().keys())
+        pbar = comfy.utils.ProgressBar(key_count)
+        class Helper:
+            steps = 0
+            def callback(self):
+                self.steps += 1
+                pbar.update_absolute(self.steps, key_count, None)
+        helper = Helper()
+        out_data = do_lora_handle(base_data, other_data, rank, "unet", "diffusion_model.", True, lambda: helper.callback())
+        if save_clip:
+            # TODO: CLIP keys get wonky, this probably doesn't work? Model-arch-dependent.
+            out_clip = do_lora_handle(base_model_clip.get_sd(), other_model_clip.get_sd(), rank, "te_text_model_encoder_layers", "0.transformer.text_model.encoder.layers.", False, lambda: helper.callback())
+            out_clip = do_lora_handle(base_model_clip.get_sd(), other_model_clip.get_sd(), rank, "te2_text_model_encoder_layers", "1.transformer.text_model.encoder.layers.", False, lambda: helper.callback())
+            out_data.update(out_clip)
+
+        # Can't easily autodetect all the correct modelspec info, but at least supply some basics
+        out_metadata = {
+            "modelspec.title": f"(Extracted LoRA) {save_filename}",
+            "modelspec.description": f"LoRA extracted in SwarmUI"
+        }
+        if metadata:
+            out_metadata.update(json.loads(metadata))
+        path = f"{save_rawpath}{save_filename}.safetensors"
+        print(f"saving to path {path}")
+        safetensors.torch.save_file(out_data, path, metadata=out_metadata)
+        return ()
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmExtractLora": SwarmExtractLora,
+}

SwarmImages.py

@@ -0,0 +1,379 @@
+import torch
+import comfy
+import math
+from nodes import MAX_RESOLUTION
+
+class SwarmImageScaleForMP:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "width": ("INT", {"default": 0, "min": 0, "max": 8192, "tooltip": "The target width of the image."}),
+                "height": ("INT", {"default": 0, "min": 0, "max": 8192, "tooltip": "The target height of the image."}),
+                "can_shrink": ("BOOLEAN", {"default": True, "tooltip": "If true, the image can be shrunk to fit the target size, otherwise it will only be scaled up or left the same."}),
+            }
+        }
+
+    CATEGORY = "SwarmUI/images"
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "scale"
+    DESCRIPTION = "Scales an image to a target width and height, while keeping the aspect ratio."
+
+    def scale(self, image, width, height, can_shrink):
+        mpTarget = width * height
+        oldWidth = image.shape[2]
+        oldHeight = image.shape[1]
+
+        scale = math.sqrt(mpTarget / (oldWidth * oldHeight))
+        if not can_shrink and scale < 1:
+            return (image,)
+        newWid = int(round(oldWidth * scale / 64) * 64)
+        newHei = int(round(oldHeight * scale / 64) * 64)
+        samples = image.movedim(-1, 1)
+        s = comfy.utils.common_upscale(samples, newWid, newHei, "bilinear", "disabled")
+        s = s.movedim(1, -1)
+        return (s,)
+
+
+class SwarmImageCrop:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "x": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 8, "tooltip": "The x coordinate in pixels of the top left corner of the crop."}),
+                "y": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 8, "tooltip": "The y coordinate in pixels of the top left corner of the crop."}),
+                "width": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8, "tooltip": "The width in pixels of the crop."}),
+                "height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8, "tooltip": "The height in pixels of the crop."}),
+            }
+        }
+
+    CATEGORY = "SwarmUI/images"
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "crop"
+    DESCRIPTION = "Crops an image to a specific region."
+
+    def crop(self, image, x, y, width, height):
+        if width <= 0 or height <= 0:
+            return (image,)
+        to_x = width + x
+        to_y = height + y
+        img = image[:, y:to_y, x:to_x, :]
+        return (img,)
+
+
+class SwarmVideoBoomerang:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+            }
+        }
+
+    CATEGORY = "SwarmUI/video"
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "boomerang"
+    DESCRIPTION = "Creates a boomerang effect by having the video play in reverse after the end, as a simple trick to make it appear to loop smoothly forever."
+
+    def boomerang(self, images):
+        # return images followed by  reverse images
+        images = torch.cat((images, images.flip(0)), 0)
+        return (images,)
+
+
+class SwarmImageNoise:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "amount": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 10.0, "step": 0.01, "round": False}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff})
+            },
+            "optional": {
+                "mask": ("MASK",)
+            }
+        }
+
+    CATEGORY = "SwarmUI/images"
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "add_noise"
+    DESCRIPTION = "Adds random noise to an image."
+
+    def add_noise(self, image, amount, seed, mask=None):
+        generator = torch.manual_seed(seed)
+        while image.dim() < 4:
+            image = image.unsqueeze(0)
+        noise = torch.randn(image.size(), dtype=image.dtype, layout=image.layout, generator=generator, device="cpu") * amount
+        if mask is not None:
+            while mask.dim() < 4:
+                mask = mask.unsqueeze(0)
+            mask = torch.nn.functional.interpolate(mask.to(image.device), size=(image.shape[1], image.shape[2]), mode="bicubic")
+            if image.shape[3] == 3 and image.shape[1] > 3: # (channels-last)
+                mask = mask.movedim(1, -1)
+            noise = noise * mask
+        img = image + noise.to(image.device)
+        img = torch.clamp(img, 0, 1)
+        return (img,)
+
+
+class SwarmTrimFrames:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "trim_start": ("INT", {"default": 0, "min": 0, "max": 4096}),
+                "trim_end": ("INT", {"default": 0, "min": 0, "max": 4096})
+            }
+        }
+
+    CATEGORY = "SwarmUI/images"
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "trim"
+    DESCRIPTION = "Trims frames from the start and end of a video."
+
+    def trim(self, image, trim_start, trim_end):
+        if image.shape[0] <= 1:
+            return (image,)
+        s_in = image
+        start = max(0, min(s_in.shape[0], trim_start))
+        end = max(0, min(s_in.shape[0], trim_end))
+        s = s_in[start:s_in.shape[0] - end].clone()
+        return (s,)
+
+
+class SwarmCountFrames:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",)
+            }
+        }
+
+    CATEGORY = "SwarmUI/images"
+    RETURN_TYPES = ("INT",)
+    FUNCTION = "count"
+    DESCRIPTION = "Counts the number of frames in an image."
+
+    def count(self, image):
+        return (image.shape[0],)
+
+
+class SwarmImageWidth:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",)
+            }
+        }
+
+    CATEGORY = "SwarmUI/images"
+    RETURN_TYPES = ("INT",)
+    FUNCTION = "get_width"
+    DESCRIPTION = "Gets the width of an image."
+
+    def get_width(self, image):
+        return (image.shape[-2],)
+
+
+class SwarmImageHeight:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",)
+            }
+        }
+
+    CATEGORY = "SwarmUI/images"
+    RETURN_TYPES = ("INT",)
+    FUNCTION = "get_height"
+    DESCRIPTION = "Gets the height of an image."
+
+    def get_height(self, image):
+        return (image.shape[-3],)
+
+
+class SwarmImageCompositeMaskedColorCorrecting:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "destination": ("IMAGE",),
+                "source": ("IMAGE",),
+                "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
+                "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
+                "mask": ("MASK",),
+                "correction_method": (["None", "Uniform", "Linear", "Linear2"], )
+            }
+        }
+
+    CATEGORY = "SwarmUI/images"
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "composite"
+    DESCRIPTION = "Works like ImageCompositeMasked, but does color correction for inpainted images (ie outside-the-mask areas are expected to be identical)"
+
+    def composite(self, destination, source, x, y, mask, correction_method):
+        destination = destination.clone().movedim(-1, 1)
+        source = source.clone().movedim(-1, 1).to(destination.device)
+        source = comfy.utils.repeat_to_batch_size(source, destination.shape[0])
+
+        x = max(-source.shape[3], min(x, destination.shape[3]))
+        y = max(-source.shape[2], min(y, destination.shape[2]))
+
+        left, top = (x, y)
+        right, bottom = (left + source.shape[3], top + source.shape[2],)
+
+        mask = mask.to(destination.device, copy=True)
+        mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(source.shape[2], source.shape[3]), mode="bilinear")
+        mask = comfy.utils.repeat_to_batch_size(mask, source.shape[0])
+
+        visible_width, visible_height = (destination.shape[3] - left + min(0, x), destination.shape[2] - top + min(0, y),)
+
+        mask = mask[:, :, :visible_height, :visible_width]
+        inverse_mask = torch.ones_like(mask) - mask
+
+        source_section = source[:, :, :visible_height, :visible_width]
+        dest_section = destination[:, :, top:bottom, left:right]
+
+        # Fall through on "None"
+        if correction_method == "Uniform":
+            source_section = color_correct_uniform(source_section, dest_section, inverse_mask)
+        elif correction_method == "Linear":
+            source_section = color_correct_linear(source_section, dest_section, inverse_mask)
+        elif correction_method == "Linear2":
+            source_section = color_correct_linear2(source_section, dest_section, inverse_mask)
+
+        source_portion = mask * source_section
+        destination_portion = inverse_mask * dest_section
+
+        destination[:, :, top:bottom, left:right] = source_portion + destination_portion
+        return (destination.movedim(1, -1),)
+
+
+def color_correct_uniform(source_section: torch.Tensor, dest_section: torch.Tensor, inverse_mask: torch.Tensor) -> torch.Tensor:
+    thresholded = (inverse_mask.clamp(0, 1) - 0.9999).clamp(0, 1) * 10000
+    thresholded_sum = thresholded.sum()
+    if thresholded_sum > 50:
+        source_hsv = rgb2hsv(source_section)
+        dest_hsv = rgb2hsv(dest_section)
+        source_hsv_masked = source_hsv * thresholded
+        dest_hsv_masked = dest_hsv * thresholded
+        diff = dest_hsv_masked - source_hsv_masked
+        diff = diff.sum(dim=[0, 2, 3]) / thresholded_sum
+        diff[0] = 0.0
+        diff = diff.unsqueeze(0).unsqueeze(2).unsqueeze(2)
+        source_hsv = source_hsv + diff
+        source_hsv = source_hsv.clamp(0, 1)
+        source_section = hsv2rgb(source_hsv)
+    return source_section
+
+
+def color_correct_linear(source_section: torch.Tensor, dest_section: torch.Tensor, inverse_mask: torch.Tensor) -> torch.Tensor:
+    thresholded = (inverse_mask.clamp(0, 1) - 0.9999).clamp(0, 1) * 10000
+    thresholded_sum = thresholded.sum()
+    if thresholded_sum > 50:
+        source_hsv = rgb2hsv(source_section)
+        dest_hsv = rgb2hsv(dest_section)
+        source_h = source_hsv[:, 0:1, :, :]
+        source_s = linear_fit(source_hsv[:, 1:2, :, :], dest_hsv[:, 1:2, :, :], thresholded)
+        source_v = linear_fit(source_hsv[:, 2:3, :, :], dest_hsv[:, 2:3, :, :], thresholded)
+        source_hsv = torch.cat([source_h, source_s, source_v], dim=1)
+        source_section = hsv2rgb(source_hsv)
+    return source_section
+
+
+# like color_correct_linear, but fits s*v and v instead of s and v to avoid instability from dark pixels
+def color_correct_linear2(source_section: torch.Tensor, dest_section: torch.Tensor, inverse_mask: torch.Tensor) -> torch.Tensor:
+    thresholded = (inverse_mask.clamp(0, 1) - 0.9999).clamp(0, 1) * 10000
+    thresholded_sum = thresholded.sum()
+    if thresholded_sum > 50:
+        source_hsv = rgb2hsv(source_section)
+        dest_hsv = rgb2hsv(dest_section)
+        source_h = source_hsv[:, 0:1, :, :]
+        source_sv_mul = source_hsv[:, 1:2, :, :] * source_hsv[:, 2:3, :, :]
+        dest_sv_mul = dest_hsv[:, 1:2, :, :] * dest_hsv[:, 2:3, :, :]
+        source_sv_mul = linear_fit(source_sv_mul, dest_sv_mul, thresholded)
+        source_v = linear_fit(source_hsv[:, 2:3, :, :], dest_hsv[:, 2:3, :, :], thresholded)
+        source_s = torch.zeros_like(source_sv_mul)
+        source_s[source_v != 0] = source_sv_mul[source_v != 0] / source_v[source_v != 0]
+        source_s = source_s.clamp(0, 1)
+        source_hsv = torch.cat([source_h, source_s, source_v], dim=1)
+        source_section = hsv2rgb(source_hsv)
+    return source_section
+
+
+def linear_fit(source_component: torch.Tensor, dest_component: torch.Tensor, thresholded: torch.Tensor) -> torch.Tensor:
+    thresholded_sum = thresholded.sum()
+    source_masked = source_component * thresholded
+    dest_masked = dest_component * thresholded
+    # Simple linear regression on dest as a function of source
+    source_mean = source_masked.sum(dim=[0, 2, 3]) / thresholded_sum
+    dest_mean = dest_masked.sum(dim=[0, 2, 3]) / thresholded_sum
+    source_mean = source_mean.unsqueeze(0).unsqueeze(2).unsqueeze(2)
+    dest_mean = dest_mean.unsqueeze(0).unsqueeze(2).unsqueeze(2)
+    source_deviation = (source_component - source_mean) * thresholded
+    dest_deviation = (dest_component - dest_mean) * thresholded
+    numerator = torch.sum(source_deviation * dest_deviation, (0, 2, 3))
+    denominator = torch.sum(source_deviation * source_deviation, (0, 2, 3)) 
+    # When all src the same color, we fall back to assuming m = 1 (uniform offset)
+    m = torch.where(denominator != 0, numerator / denominator, torch.tensor(1.0))
+    m = m.unsqueeze(0).unsqueeze(2).unsqueeze(2)
+    b = dest_mean - source_mean * m
+    source_component = m * source_component + b
+    source_component = source_component.clamp(0, 1)
+    return source_component
+
+
+# from https://github.com/limacv/RGB_HSV_HSL
+def rgb2hsv(rgb: torch.Tensor) -> torch.Tensor:
+    cmax, cmax_idx = torch.max(rgb, dim=1, keepdim=True)
+    cmin = torch.min(rgb, dim=1, keepdim=True)[0]
+    delta = cmax - cmin
+    hsv_h = torch.empty_like(rgb[:, 0:1, :, :])
+    cmax_idx[delta == 0] = 3
+    hsv_h[cmax_idx == 0] = (((rgb[:, 1:2] - rgb[:, 2:3]) / delta) % 6)[cmax_idx == 0]
+    hsv_h[cmax_idx == 1] = (((rgb[:, 2:3] - rgb[:, 0:1]) / delta) + 2)[cmax_idx == 1]
+    hsv_h[cmax_idx == 2] = (((rgb[:, 0:1] - rgb[:, 1:2]) / delta) + 4)[cmax_idx == 2]
+    hsv_h[cmax_idx == 3] = 0.0
+    hsv_h /= 6.0
+    hsv_s = torch.where(cmax == 0, torch.tensor(0.0).type_as(rgb), delta / cmax)
+    hsv_v = cmax
+    return torch.cat([hsv_h, hsv_s, hsv_v], dim=1)
+
+
+def hsv2rgb(hsv: torch.Tensor) -> torch.Tensor:
+    hsv_h, hsv_s, hsv_l = hsv[:, 0:1], hsv[:, 1:2], hsv[:, 2:3]
+    _c = hsv_l * hsv_s
+    _x = _c * (- torch.abs(hsv_h * 6.0 % 2.0 - 1) + 1.0)
+    _m = hsv_l - _c
+    _o = torch.zeros_like(_c)
+    idx = (hsv_h * 6.0).type(torch.uint8)
+    idx = (idx % 6).expand(-1, 3, -1, -1)
+    rgb = torch.empty_like(hsv)
+    rgb[idx == 0] = torch.cat([_c, _x, _o], dim=1)[idx == 0]
+    rgb[idx == 1] = torch.cat([_x, _c, _o], dim=1)[idx == 1]
+    rgb[idx == 2] = torch.cat([_o, _c, _x], dim=1)[idx == 2]
+    rgb[idx == 3] = torch.cat([_o, _x, _c], dim=1)[idx == 3]
+    rgb[idx == 4] = torch.cat([_x, _o, _c], dim=1)[idx == 4]
+    rgb[idx == 5] = torch.cat([_c, _o, _x], dim=1)[idx == 5]
+    rgb += _m
+    return rgb
+
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmImageScaleForMP": SwarmImageScaleForMP,
+    "SwarmImageCrop": SwarmImageCrop,
+    "SwarmVideoBoomerang": SwarmVideoBoomerang,
+    "SwarmImageNoise": SwarmImageNoise,
+    "SwarmTrimFrames": SwarmTrimFrames,
+    "SwarmCountFrames": SwarmCountFrames,
+    "SwarmImageWidth": SwarmImageWidth,
+    "SwarmImageHeight": SwarmImageHeight,
+    "SwarmImageCompositeMaskedColorCorrecting": SwarmImageCompositeMaskedColorCorrecting
+}

SwarmInputNodes.py

@@ -0,0 +1,299 @@
+from . import SwarmLoadImageB64
+import folder_paths
+from nodes import CheckpointLoaderSimple, LoadImage
+from comfy_extras.nodes_video import LoadVideo
+from comfy_api.input_impl import VideoFromFile
+import os, base64, io
+try:
+    from comfy_extras.nodes_audio import LoadAudio
+    from comfy_extras.nodes_audio import load as raw_audio_load
+except ImportError:
+    print("Error: Nodes_Audio failed to import")
+
+INT_MAX = 0xffffffffffffffff
+INT_MIN = -INT_MAX
+
+class SwarmInputGroup:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "title": ("STRING", {"default": "My Group", "tooltip": "The title of the group."}),
+                "open_by_default": ("BOOLEAN", {"default": True, "tooltip": "Whether the group should be open by default."}),
+                "description": ("STRING", {"default": "", "multiline": True, "tooltip": "A description of the group that shows up when you click the '?' button."}),
+                "order_priority": ("FLOAT", {"default": 0, "min": -1024, "max": 1024, "step": 0.5, "round": 0.0000001, "tooltip": "The order priority of the group. Higher values go further down in the list of groups."}),
+                "is_advanced": ("BOOLEAN", {"default": False, "tooltip": "If true, the group will only be visible when 'Display Advanced' is clicked."}),
+                "can_shrink": ("BOOLEAN", {"default": True, "tooltip": "If true, the group can be collapsed by the user. If false, will be forced to remain open."}),
+            },
+        }
+
+    CATEGORY = "SwarmUI/inputs"
+    RETURN_TYPES = ("GROUP",)
+    FUNCTION = "do_input"
+    DESCRIPTION = "SwarmInput nodes let you define custom input controls in Swarm-Comfy Workflows. Group defines a parameter groupging - link this to other nodes to create a collapsible group, all nodes this links to will be inside the group."
+
+    def do_input(self, **kwargs):
+        return (None, )
+
+
+STANDARD_REQ_INPUTS = {
+    "description": ("STRING", {"default": "", "multiline": True, "tooltip": "A description of the input that shows up when you click the '?' button."}),
+    "order_priority": ("FLOAT", {"default": 0, "min": -1024, "max": 1024, "step": 0.5, "round": 0.0000001, "tooltip": "The order priority of the input. Higher values go further down in the list of inputs. This only applies within the group this node is part of."}),
+    "is_advanced": ("BOOLEAN", {"default": False, "tooltip": "If true, the input will only be visible when 'Display Advanced' is clicked."}),
+    "raw_id": ("STRING", {"default": "", "tooltip": "The raw ID of the input. This can be used to customize the input for API usage, or to make use of default SwarmUI parameters. Most of the time, you don't need to touch this. By default this will autogenerate a unique ID based on the title value."}),
+}
+STANDARD_OTHER_INPUTS = {
+    "optional": {
+        "group": ("GROUP", )
+    }
+}
+
+
+class SwarmInputInteger:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "title": ("STRING", {"default": "My Integer", "tooltip": "The name of the input."}),
+                "value": ("INT", {"default": 0, "min": INT_MIN, "max": INT_MAX, "step": 1, "tooltip": "The default value of the input."}),
+                "step": ("INT", {"default": 1, "min": INT_MIN, "max": INT_MAX, "step": 1, "tooltip": "The step size of the input. That is, how much the value changes when you click the up/down arrows or move the slider."}),
+                "min": ("INT", {"default": 0, "min": INT_MIN, "max": INT_MAX, "step": 1, "tooltip": "The minimum value of the input."}),
+                "max": ("INT", {"default": 100, "min": INT_MIN, "max": INT_MAX, "step": 1, "tooltip": "The maximum value of the input."}),
+                "view_max": ("INT", {"default": 100, "min": INT_MIN, "max": INT_MAX, "step": 1, "tooltip": "The maximum value of the input that is displayed in the UI when using a slider. This is useful if you want to allow a higher range of values, but don't want to clutter the UI with a huge slider."}),
+                "view_type": (["big", "small", "seed", "slider", "pot_slider"],{"tooltip": "The type of input control to use. 'big' is a large text input, 'small' is a small text input, 'seed' is a text input with seed-specific controls, 'slider' is a slider, and 'pot_slider' is a Power-Of-Two scaled slider - this is useful for large inputs like resolutions to allow a more natural feeling selection range."}),
+            } | STANDARD_REQ_INPUTS,
+        } | STANDARD_OTHER_INPUTS
+
+    CATEGORY = "SwarmUI/inputs"
+    RETURN_TYPES = ("INT",)
+    FUNCTION = "do_input"
+    DESCRIPTION = "SwarmInput nodes let you define custom input controls in Swarm-Comfy Workflows. Integer lets you input a whole number without a decimal point."
+
+    def do_input(self, value, **kwargs):
+        return (value, )
+
+
+class SwarmInputFloat:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "title": ("STRING", {"default": "My Floating-Point Number", "tooltip": "The name of the input."}),
+                "value": ("FLOAT", {"default": 0, "min": INT_MIN, "max": INT_MAX, "step": 0.01, "round": 0.0000001, "tooltip": "The default value of the input."}),
+                "step": ("FLOAT", {"default": 0.1, "min": INT_MIN, "max": INT_MAX, "step": 0.01, "round": 0.0000001, "tooltip": "The step size of the input. That is, how much the value changes when you click the up/down arrows or move the slider."}),
+                "min": ("FLOAT", {"default": 0, "min": INT_MIN, "max": INT_MAX, "step": 0.01, "round": 0.0000001, "tooltip": "The minimum value of the input."}),
+                "max": ("FLOAT", {"default": 100, "min": INT_MIN, "max": INT_MAX, "step": 0.01, "round": 0.0000001, "tooltip": "The maximum value of the input."}),
+                "view_max": ("FLOAT", {"default": 100, "min": INT_MIN, "max": INT_MAX, "step": 0.01, "round": 0.0000001, "tooltip": "The maximum value of the input that is displayed in the UI when using a slider. This is useful if you want to allow a higher range of values, but don't want to clutter the UI with a huge slider."}),
+                "view_type": (["big", "small", "slider", "pot_slider"], {"tooltip": "The type of input control to use. 'big' is a large text input, 'small' is a small text input, 'slider' is a slider, and 'pot_slider' is a Power-Of-Two scaled slider - this is useful for large inputs like resolutions to allow a more natural feeling selection range."}),
+            } | STANDARD_REQ_INPUTS,
+        } | STANDARD_OTHER_INPUTS
+
+    CATEGORY = "SwarmUI/inputs"
+    RETURN_TYPES = ("FLOAT",)
+    FUNCTION = "do_input"
+    DESCRIPTION = "SwarmInput nodes let you define custom input controls in Swarm-Comfy Workflows. Float lets you input a number with a decimal point."
+
+    def do_input(self, value, **kwargs):
+        return (value, )
+
+
+class SwarmInputText:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "title": ("STRING", {"default": "My Text", "tooltip": "The name of the input."}),
+                "value": ("STRING", {"default": "", "multiline": True, "tooltip": "The default value of the input."}),
+                "view_type": (["normal", "prompt", "big"], {"tooltip": "How to format this text input. 'normal' is a simple single line text input, 'prompt' is a prompt-like text input that has multiple lines and other prompting-specific features, 'big' is an extra large multiline text box."}),
+            } | STANDARD_REQ_INPUTS,
+        } | STANDARD_OTHER_INPUTS
+
+    CATEGORY = "SwarmUI/inputs"
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "do_input"
+    DESCRIPTION = "SwarmInput nodes let you define custom input controls in Swarm-Comfy Workflows. Text lets you input a string of text. This can be simple text inputs, or prompt-like text inputs."
+
+    def do_input(self, value, **kwargs):
+        return (value, )
+
+
+class SwarmInputModelName:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "title": ("STRING", {"default": "My Model Name Input", "tooltip": "The name of the input."}),
+                "value": ("STRING", {"default": "", "multiline": False, "tooltip": "The default value of the input."}),
+                "subtype": (["Stable-Diffusion", "VAE", "LoRA", "Embedding", "ControlNet", "ClipVision"], {"tooltip": "The model subtype to select from."}),
+            } | STANDARD_REQ_INPUTS,
+        } | STANDARD_OTHER_INPUTS
+
+    CATEGORY = "SwarmUI/inputs"
+    RETURN_TYPES = ("",)
+    FUNCTION = "do_input"
+    DESCRIPTION = "SwarmInput nodes let you define custom input controls in Swarm-Comfy Workflows. Model Name lets you have a dropdown select for models of a given model sub-type."
+
+    def do_input(self, value, **kwargs):
+        return (value, )
+
+
+class SwarmInputCheckpoint:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "title": ("STRING", {"default": "My Checkpoint Model Name Input", "tooltip": "The name of the input."}),
+                "value": (folder_paths.get_filename_list("checkpoints"), {"tooltip": "The default value of the input."}),
+            } | STANDARD_REQ_INPUTS,
+        } | STANDARD_OTHER_INPUTS
+
+    CATEGORY = "SwarmUI/inputs"
+    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
+    FUNCTION = "do_input"
+    DESCRIPTION = "SwarmInput nodes let you define custom input controls in Swarm-Comfy Workflows. Checkpoint lets you have a dropdown select for checkpoint models. Acts like the Checkpoint Loader node."
+
+    def do_input(self, value, **kwargs):
+        return CheckpointLoaderSimple().load_checkpoint(value)
+
+
+class SwarmInputDropdown:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "title": ("STRING", {"default": "My Dropdown", "tooltip": "The name of the input."}),
+                "value": ("STRING", {"default": "", "multiline": False, "tooltip": "The default value of the input."}),
+                "values": ("STRING", {"default": "one, two, three", "multiline": True, "tooltip": "A comma-separated list of values to choose from. If you leave this blank, the dropdown will automatically load the value list from the connected node."}),
+            } | STANDARD_REQ_INPUTS,
+        } | STANDARD_OTHER_INPUTS
+
+    CATEGORY = "SwarmUI/inputs"
+    RETURN_TYPES = ("STRING", "",)
+    FUNCTION = "do_input"
+    DESCRIPTION = "SwarmInput nodes let you define custom input controls in Swarm-Comfy Workflows. Dropdown lets you have a dropdown select for a list of values. You can leave the values list empty and attach the blank output to a real dropdown in another node to have it automatically load the value list from that connected node."
+
+    def do_input(self, value, **kwargs):
+        return (value, value, )
+
+
+class SwarmInputBoolean:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "title": ("STRING", {"default": "My Boolean", "tooltip": "The name of the input."}),
+                "value": ("BOOLEAN", {"default": False, "tooltip": "The default value of the input."}),
+            } | STANDARD_REQ_INPUTS,
+        } | STANDARD_OTHER_INPUTS
+
+    CATEGORY = "SwarmUI/inputs"
+    RETURN_TYPES = ("BOOLEAN",)
+    FUNCTION = "do_input"
+    DESCRIPTION = "SwarmInput nodes let you define custom input controls in Swarm-Comfy Workflows. Boolean lets you have a checkbox for a true/false value."
+
+    def do_input(self, value, **kwargs):
+        return (value, )
+
+
+class SwarmInputImage:
+    @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))]
+        files = folder_paths.filter_files_content_types(files, ["image"])
+        return {
+            "required": {
+                "title": ("STRING", {"default": "My Image", "tooltip": "The name of the input."}),
+                "value": ("STRING", {"default": "(Do Not Set Me)", "multiline": True, "tooltip": "Always leave this blank, the SwarmUI server will fill it for you."}),
+                "auto_resize": ("BOOLEAN", {"default": True, "tooltip": "If true, the image will be resized to match the current generation resolution. If false, the image will be kept at whatever size the user input it at."}),
+            } | STANDARD_REQ_INPUTS | {
+                "image": (sorted(files), {"image_upload": True}),
+            },
+        } | STANDARD_OTHER_INPUTS
+
+    CATEGORY = "SwarmUI/inputs"
+    RETURN_TYPES = ("IMAGE","MASK",)
+    FUNCTION = "do_input"
+    DESCRIPTION = "SwarmInput nodes let you define custom input controls in Swarm-Comfy Workflows. Image lets you input an image. Internally this node uses a Base64 string as input when value is set by SwarmUI server (Generate tab), otherwise use select Image (Comfy Workflow tab)."
+
+    def do_input(self, value=None, image=None, **kwargs):
+        if not value or value == "(Do Not Set Me)":
+            return LoadImage().load_image(image)
+        else:
+            return SwarmLoadImageB64.b64_to_img_and_mask(value)
+
+
+class SwarmInputAudio:
+    @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))]
+        files = folder_paths.filter_files_content_types(files, ["audio"])
+        return {
+            "required": {
+                "title": ("STRING", {"default": "My Audio", "tooltip": "The name of the input."}),
+                "value": ("STRING", {"default": "(Do Not Set Me)", "multiline": True, "tooltip": "Always leave this blank, the SwarmUI server will fill it for you."}),
+            } | STANDARD_REQ_INPUTS | {
+                # TODO: This explodes the comfy frontend for some reason
+                #"audio": (sorted(files), {"audio_upload": True}),
+            },
+        } | STANDARD_OTHER_INPUTS
+
+    CATEGORY = "SwarmUI/inputs"
+    RETURN_TYPES = ("AUDIO",)
+    FUNCTION = "do_input"
+    DESCRIPTION = "SwarmInput nodes let you define custom input controls in Swarm-Comfy Workflows. Audio lets you input an audio file. Internally this node uses a Base64 string as input when value is set by SwarmUI server (Generate tab), otherwise use select Audio (Comfy Workflow tab)."
+
+    def do_input(self, value=None, audio=None, **kwargs):
+        if not value or value == "(Do Not Set Me)":
+            return LoadAudio().load_audio(audio)
+        else:
+            audio_data = base64.b64decode(value)
+            audio_bytes = io.BytesIO(audio_data)
+            waveform, sample_rate = raw_audio_load(audio_bytes)
+            audio = {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
+            return (audio, )
+
+
+class SwarmInputVideo:
+    @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))]
+        files = folder_paths.filter_files_content_types(files, ["video"])
+        return {
+            "required": {
+                "title": ("STRING", {"default": "My Video", "tooltip": "The name of the input."}),
+                "value": ("STRING", {"default": "(Do Not Set Me)", "multiline": True, "tooltip": "Always leave this blank, the SwarmUI server will fill it for you."}),
+            } | STANDARD_REQ_INPUTS | {
+                "video": (sorted(files), {"video_upload": True}),
+            },
+        } | STANDARD_OTHER_INPUTS
+
+    CATEGORY = "SwarmUI/inputs"
+    RETURN_TYPES = ("VIDEO",)
+    FUNCTION = "do_input"
+    DESCRIPTION = "SwarmInput nodes let you define custom input controls in Swarm-Comfy Workflows. Video lets you input a video file. Internally this node uses a Base64 string as input when value is set by SwarmUI server (Generate tab), otherwise use select Video (Comfy Workflow tab)."
+
+    def do_input(self, value=None, video=None, **kwargs):
+        if not value or value == "(Do Not Set Me)":
+            return LoadVideo.execute(video)
+        else:
+            video_data = base64.b64decode(value)
+            video_bytes = io.BytesIO(video_data)
+            return (VideoFromFile(video_bytes), )
+
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmInputGroup": SwarmInputGroup,
+    "SwarmInputInteger": SwarmInputInteger,
+    "SwarmInputFloat": SwarmInputFloat,
+    "SwarmInputText": SwarmInputText,
+    "SwarmInputModelName": SwarmInputModelName,
+    "SwarmInputCheckpoint": SwarmInputCheckpoint,
+    "SwarmInputDropdown": SwarmInputDropdown,
+    "SwarmInputBoolean": SwarmInputBoolean,
+    "SwarmInputImage": SwarmInputImage,
+    "SwarmInputAudio": SwarmInputAudio,
+    "SwarmInputVideo": SwarmInputVideo,
+}

SwarmInternalUtil.py

@@ -0,0 +1,111 @@
+import comfy, folder_paths, execution
+from server import PromptServer
+from comfy import samplers
+import functools
+
+# This is purely a hack to provide a list of embeds in the object_info report.
+# Code referenced from Comfy VAE impl. Probably does nothing useful in an actual workflow.
+class SwarmEmbedLoaderListProvider:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "embed_name": (folder_paths.get_filename_list("embeddings"), )
+            }
+        }
+
+    CATEGORY = "SwarmUI/internal"
+    RETURN_TYPES = ("EMBEDDING",)
+    FUNCTION = "load_embed"
+    DESCRIPTION = "Internal node just intended to provide a list of currently known embeddings to Swarm. You can also use it to blindly load an embedding file if you need to."
+
+    def load_embed(self, embed_name):
+        embed_path = folder_paths.get_full_path("embedding", embed_name)
+        sd = comfy.utils.load_torch_file(embed_path)
+        return (sd,)
+
+
+class SwarmJustLoadTheModelPlease:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "clip": ("CLIP,GEMMA",),
+                "vae": ("VAE",),
+            }
+        }
+
+    CATEGORY = "SwarmUI/internal"
+    RETURN_TYPES = ()
+    FUNCTION = "just_load"
+    OUTPUT_NODE = True
+    DESCRIPTION = "Internal node that acts as a final output for a model/clip/vae. This allows swarm to load models when needed without generating anything."
+
+    def just_load(self, model, clip, vae):
+        if model is None:
+            raise ValueError("The model failed to load")
+        if clip is None:
+            raise ValueError("The text encoders (CLIP) failed to load")
+        if vae is None:
+            raise ValueError("The VAE failed to load")
+        return {}
+
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmEmbedLoaderListProvider": SwarmEmbedLoaderListProvider,
+    "SwarmJustLoadTheModelPlease": SwarmJustLoadTheModelPlease
+}
+
+
+# This is a dirty hack to shut up the errors from Dropdown combo mismatch, pending Comfy upstream fix
+ORIG_EXECUTION_VALIDATE = execution.validate_inputs
+async def validate_inputs(prompt_id, prompt, item, validated):
+    raw_result = await ORIG_EXECUTION_VALIDATE(prompt_id, prompt, item, validated)
+    if raw_result is None:
+        return None
+    (did_succeed, errors, unique_id) = raw_result
+    if did_succeed:
+        return raw_result
+    for error in errors:
+        if error['type'] == "return_type_mismatch":
+            o_id = error['extra_info']['linked_node'][0]
+            o_class_type = prompt[o_id]['class_type']
+            if o_class_type == "SwarmInputModelName" or o_class_type == "SwarmInputDropdown":
+                errors.remove(error)
+    did_succeed = len(errors) == 0
+    return (did_succeed, errors, unique_id)
+
+execution.validate_inputs = validate_inputs
+
+# Comfy's app logger has broken terminal compat, so violently force it to auto-flush
+try:
+    from app import logger
+    def patch_interceptor(interceptor):
+        if interceptor:
+            orig = interceptor.write
+            def write(self, data):
+                orig(data)
+                self.flush()
+            interceptor.write = functools.partial(write, interceptor)
+            # Force UTF-8 too, to prevent encoding errors (Comfy will full crash outputting some languages)
+            # (Swarm's C# engine has code to forcibly assume UTF-8, so this is safe. Otherwise it would wonk the terminal if the terminal isn't set to UTF-8)
+            interceptor.reconfigure(encoding='utf-8')
+    patch_interceptor(logger.stdout_interceptor)
+    patch_interceptor(logger.stderr_interceptor)
+except Exception as e:
+    import traceback
+    traceback.print_exc()
+
+# comfy's server/PromptServer is janky with EventID=4, so overwrite send_bytes to interpret EventID=9999123 as 4
+try:
+    server = PromptServer.instance
+    orig = server.send_bytes
+    async def send_bytes(self, event, data, sid=None):
+        if event == 9999123:
+            event = 4
+        await orig(event, data, sid=sid)
+    server.send_bytes = functools.partial(send_bytes, server)
+except Exception as e:
+    import traceback
+    traceback.print_exc()

SwarmKSampler.py

@@ -0,0 +1,356 @@
+import torch, struct, json
+from io import BytesIO
+import latent_preview, comfy
+from server import PromptServer
+from comfy.model_base import SDXL, SVD_img2vid, Flux, WAN21, Chroma
+from comfy import samplers
+import numpy as np
+from math import ceil
+from latent_preview import TAESDPreviewerImpl
+from comfy_execution.utils import get_executing_context
+
+def slerp(val, low, high):
+    low_norm = low / torch.norm(low, dim=1, keepdim=True)
+    high_norm = high / torch.norm(high, dim=1, keepdim=True)
+    dot = (low_norm * high_norm).sum(1)
+    if dot.mean() > 0.9995:
+        return low * val + high * (1 - val)
+    omega = torch.acos(dot)
+    so = torch.sin(omega)
+    res = (torch.sin((1.0 - val) * omega) / so).unsqueeze(1) * low + (torch.sin(val * omega) / so).unsqueeze(1) * high
+    return res
+
+def swarm_partial_noise(seed, latent_image):
+    generator = torch.manual_seed(seed)
+    return torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")
+
+def swarm_fixed_noise(seed, latent_image, var_seed, var_seed_strength):
+    noises = []
+    for i in range(latent_image.size()[0]):
+        if var_seed_strength > 0:
+            noise = swarm_partial_noise(seed, latent_image[i])
+            var_noise = swarm_partial_noise(var_seed + i, latent_image[i])
+            if noise.ndim == 4: # Video models are B C F H W, we're in a B loop already so sub-iterate over F (Frames)
+                for j in range(noise.shape[1]):
+                    noise[:, j] = slerp(var_seed_strength, noise[:, j], var_noise[:, j])
+            else:
+                noise = slerp(var_seed_strength, noise, var_noise)
+        else:
+            noise = swarm_partial_noise(seed + i, latent_image[i])
+        noises.append(noise)
+    return torch.stack(noises, dim=0)
+
+def get_preview_metadata():
+    executing_context = get_executing_context()
+    prompt_id = None
+    node_id = None
+    if executing_context is not None:
+        prompt_id = executing_context.prompt_id
+        node_id = executing_context.node_id
+    if prompt_id is None:
+        prompt_id = PromptServer.instance.last_prompt_id
+    if node_id is None:
+        node_id = PromptServer.instance.last_node_id
+    return {"node_id": node_id, "prompt_id": prompt_id, "display_node_id": node_id, "parent_node_id": node_id, "real_node_id": node_id} # display_node_id, parent_node_id, real_node_id? comfy_execution/progress.py has this.
+
+def swarm_send_extra_preview(id, image):
+    server = PromptServer.instance
+    metadata = get_preview_metadata()
+    metadata["mime_type"] = "image/jpeg"
+    metadata["id"] = id
+    metadata_json = json.dumps(metadata).encode('utf-8')
+    bytesIO = BytesIO()
+    image.save(bytesIO, format="JPEG", quality=90, compress_level=4)
+    image_bytes = bytesIO.getvalue()
+    combined_data = bytearray()
+    combined_data.extend(struct.pack(">I", len(metadata_json)))
+    combined_data.extend(metadata_json)
+    combined_data.extend(image_bytes)
+    server.send_sync(9999123, combined_data, sid=server.client_id)
+
+def swarm_send_animated_preview(id, images):
+    server = PromptServer.instance
+    bytesIO = BytesIO()
+    images[0].save(bytesIO, save_all=True, duration=int(1000.0/6), append_images=images[1 : len(images)], lossless=False, quality=60, method=0, format='WEBP')
+    bytesIO.seek(0)
+    image_bytes = bytesIO.getvalue()
+    metadata = get_preview_metadata()
+    metadata["mime_type"] = "image/webp"
+    metadata["id"] = id
+    metadata_json = json.dumps(metadata).encode('utf-8')
+    combined_data = bytearray()
+    combined_data.extend(struct.pack(">I", len(metadata_json)))
+    combined_data.extend(metadata_json)
+    combined_data.extend(image_bytes)
+    server.send_sync(9999123, combined_data, sid=server.client_id)
+
+def calculate_sigmas_scheduler(model, scheduler_name, steps, sigma_min, sigma_max, rho):
+    model_sampling = model.get_model_object("model_sampling")
+    if scheduler_name == "karras":
+        return comfy.k_diffusion.sampling.get_sigmas_karras(n=steps, sigma_min=sigma_min if sigma_min >= 0 else float(model_sampling.sigma_min), sigma_max=sigma_max if sigma_max >= 0 else float(model_sampling.sigma_max), rho=rho)
+    elif scheduler_name == "exponential":
+        return comfy.k_diffusion.sampling.get_sigmas_exponential(n=steps, sigma_min=sigma_min if sigma_min >= 0 else float(model_sampling.sigma_min), sigma_max=sigma_max if sigma_max >= 0 else float(model_sampling.sigma_max))
+    else:
+        return None
+
+def make_swarm_sampler_callback(steps, device, model, previews):
+    previewer = latent_preview.get_previewer(device, model.model.latent_format) if previews != "none" else None
+    pbar = comfy.utils.ProgressBar(steps)
+    def callback(step, x0, x, total_steps):
+        pbar.update_absolute(step + 1, total_steps, None)
+        if previewer:
+            if (step == 0 or (step < 3 and x0.ndim == 5 and x0.shape[1] > 8)) and not isinstance(previewer, TAESDPreviewerImpl):
+                x0 = x0.clone().cpu() # Sync copy to CPU for first few steps to prevent reading old data, more steps for videos. Future steps allow comfy to do its async non_blocky stuff.
+            if x0.ndim == 5:
+                # video shape is [batch, channels, backwards time, width, height], for previews needs to be swapped to [forwards time, channels, width, height]
+                x0 = x0[0].permute(1, 0, 2, 3)
+                x0 = torch.flip(x0, [0])
+            def do_preview(id, index):
+                preview_img = previewer.decode_latent_to_preview_image("JPEG", x0[index:index+1])
+                swarm_send_extra_preview(id, preview_img[1])
+            if previews == "iterate":
+                do_preview(0, step % x0.shape[0])
+            elif previews == "animate":
+                if x0.shape[0] == 1:
+                    do_preview(0, 0)
+                else:
+                    images = []
+                    for i in range(x0.shape[0]):
+                        preview_img = previewer.decode_latent_to_preview_image("JPEG", x0[i:i+1])
+                        images.append(preview_img[1])
+                    swarm_send_animated_preview(0, images)
+            elif previews == "default":
+                for i in range(x0.shape[0]):
+                    preview_img = previewer.decode_latent_to_preview_image("JPEG", x0[i:i+1])
+                    swarm_send_extra_preview(i, preview_img[1])
+            elif previews == "one":
+                do_preview(0, 0)
+            elif previews == "second":
+                do_preview(0, 1 % x0.shape[0])
+    return callback
+
+
+def loglinear_interp(t_steps, num_steps):
+    """
+    Performs log-linear interpolation of a given array of decreasing numbers.
+    """
+    xs = np.linspace(0, 1, len(t_steps))
+    ys = np.log(t_steps[::-1])
+
+    new_xs = np.linspace(0, 1, num_steps)
+    new_ys = np.interp(new_xs, xs, ys)
+
+    interped_ys = np.exp(new_ys)[::-1].copy()
+    return interped_ys
+
+AYS_NOISE_LEVELS = {
+    "SD1": [14.6146412293, 6.4745760956,  3.8636745985,  2.6946151520, 1.8841921177,  1.3943805092,  0.9642583904,  0.6523686016, 0.3977456272,  0.1515232662,  0.0291671582],
+    "SDXL":[14.6146412293, 6.3184485287,  3.7681790315,  2.1811480769, 1.3405244945,  0.8620721141,  0.5550693289,  0.3798540708, 0.2332364134,  0.1114188177,  0.0291671582],
+    "SVD": [700.00, 54.5, 15.886, 7.977, 4.248, 1.789, 0.981, 0.403, 0.173, 0.034, 0.002],
+    # Flux and Wan from https://github.com/comfyanonymous/ComfyUI/pull/7584
+    "Flux": [0.9968, 0.9886, 0.9819, 0.975, 0.966, 0.9471, 0.9158, 0.8287, 0.5512, 0.2808, 0.001],
+    "Wan": [1.0, 0.997, 0.995, 0.993, 0.991, 0.989, 0.987, 0.985, 0.98, 0.975, 0.973, 0.968, 0.96, 0.946, 0.927, 0.902, 0.864, 0.776, 0.539, 0.208, 0.001],
+    # https://github.com/comfyanonymous/ComfyUI/commit/08ff5fa08a92e0b3f23b9abec979a830a6cffb03#diff-3e4e70e402dcd9e1070ad71ef9292277f10d9faccf36a1c405c0c717a7ee6485R23
+    "Chroma": [0.992, 0.99, 0.988, 0.985, 0.982, 0.978, 0.973, 0.968, 0.961, 0.953, 0.943, 0.931, 0.917, 0.9, 0.881, 0.858, 0.832, 0.802, 0.769, 0.731, 0.69, 0.646, 0.599, 0.55, 0.501, 0.451, 0.402, 0.355, 0.311, 0.27, 0.232, 0.199, 0.169, 0.143, 0.12, 0.101, 0.084, 0.07, 0.058, 0.048, 0.001]
+}
+
+def split_latent_tensor(latent_tensor, tile_size=1024, scale_factor=8):
+    """Generate tiles for a given latent tensor, considering the scaling factor."""
+    latent_tile_size = tile_size // scale_factor  # Adjust tile size for latent space
+    height, width = latent_tensor.shape[-2:]
+
+    # Determine the number of tiles needed
+    num_tiles_x = ceil(width / latent_tile_size)
+    num_tiles_y = ceil(height / latent_tile_size)
+
+    # If width or height is an exact multiple of the tile size, add an additional tile for overlap
+    if width % latent_tile_size == 0:
+        num_tiles_x += 1
+    if height % latent_tile_size == 0:
+        num_tiles_y += 1
+
+    # Calculate the overlap
+    overlap_x = 0 if num_tiles_x == 1 else (num_tiles_x * latent_tile_size - width) / (num_tiles_x - 1)
+    overlap_y = 0 if num_tiles_y == 1 else (num_tiles_y * latent_tile_size - height) / (num_tiles_y - 1)
+    if overlap_x < 32 and num_tiles_x > 1:
+        num_tiles_x += 1
+        overlap_x = (num_tiles_x * latent_tile_size - width) / (num_tiles_x - 1)
+    if overlap_y < 32 and num_tiles_y > 1:
+        num_tiles_y += 1
+        overlap_y = (num_tiles_y * latent_tile_size - height) / (num_tiles_y - 1)
+
+    tiles = []
+
+    for i in range(num_tiles_y):
+        for j in range(num_tiles_x):
+            x_start = j * latent_tile_size - j * overlap_x
+            y_start = i * latent_tile_size - i * overlap_y
+
+            # Correct for potential float precision issues
+            x_start = round(x_start)
+            y_start = round(y_start)
+
+            # Crop the tile from the latent tensor
+            tile_tensor = latent_tensor[..., y_start:y_start + latent_tile_size, x_start:x_start + latent_tile_size]
+            tiles.append(((x_start, y_start, x_start + latent_tile_size, y_start + latent_tile_size), tile_tensor))
+
+    return tiles
+
+def stitch_latent_tensors(original_size, tiles, scale_factor=8):
+    """Stitch tiles together to create the final upscaled latent tensor with overlaps."""
+    result = torch.zeros(original_size)
+
+    # We assume tiles come in the format [(coordinates, tile), ...]
+    sorted_tiles = sorted(tiles, key=lambda x: (x[0][1], x[0][0]))  # Sort by upper then left
+
+    # Variables to keep track of the current row's starting point
+    current_row_upper = None
+
+    for (left, upper, right, lower), tile in sorted_tiles:
+
+        # Check if we're starting a new row
+        if current_row_upper != upper:
+            current_row_upper = upper
+            first_tile_in_row = True
+        else:
+            first_tile_in_row = False
+
+        tile_width = right - left
+        tile_height = lower - upper
+        feather = tile_width // 8  # Assuming feather size is consistent with the example
+
+        mask = torch.ones_like(tile)
+
+        if not first_tile_in_row:  # Left feathering for tiles other than the first in the row
+            for t in range(feather):
+                mask[..., :, t:t+1] *= (1.0 / feather) * (t + 1)
+
+        if upper != 0:  # Top feathering for all tiles except the first row
+            for t in range(feather):
+                mask[..., t:t+1, :] *= (1.0 / feather) * (t + 1)
+
+        # Apply the feathering mask
+        combined_area = tile * mask + result[..., upper:lower, left:right] * (1.0 - mask)
+        result[..., upper:lower, left:right] = combined_area
+
+    return result
+
+class SwarmKSampler:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.5, "round": 0.001}),
+                "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
+                "scheduler": (["turbo", "align_your_steps", "ltxv", "ltxv-image"] + comfy.samplers.KSampler.SCHEDULERS, ),
+                "positive": ("CONDITIONING", ),
+                "negative": ("CONDITIONING", ),
+                "latent_image": ("LATENT", ),
+                "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
+                "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
+                "var_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "var_seed_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, "round": 0.001}),
+                "sigma_max": ("FLOAT", {"default": -1, "min": -1.0, "max": 1000.0, "step":0.01, "round": False}),
+                "sigma_min": ("FLOAT", {"default": -1, "min": -1.0, "max": 1000.0, "step":0.01, "round": False}),
+                "rho": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
+                "add_noise": (["enable", "disable"], ),
+                "return_with_leftover_noise": (["disable", "enable"], ),
+                "previews": (["default", "none", "one", "second", "iterate", "animate"], ),
+                "tile_sample": ("BOOLEAN", {"default": False}),
+                "tile_size": ("INT", {"default": 1024, "min": 256, "max": 4096}),
+            }
+        }
+
+    CATEGORY = "SwarmUI/sampling"
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "run_sampling"
+    DESCRIPTION = "Works like a vanilla Comfy KSamplerAdvanced, but with extra inputs for advanced features such as sigma scale, tiling, previews, etc."
+
+    def sample(self, model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, var_seed, var_seed_strength, sigma_max, sigma_min, rho, add_noise, return_with_leftover_noise, previews):
+        device = comfy.model_management.get_torch_device()
+        latent_samples = latent_image["samples"]
+        latent_samples = comfy.sample.fix_empty_latent_channels(model, latent_samples)
+        disable_noise = add_noise == "disable"
+
+        if disable_noise:
+            noise = torch.zeros(latent_samples.size(), dtype=latent_samples.dtype, layout=latent_samples.layout, device="cpu")
+        else:
+            noise = swarm_fixed_noise(noise_seed, latent_samples, var_seed, var_seed_strength)
+
+        noise_mask = None
+        if "noise_mask" in latent_image:
+            noise_mask = latent_image["noise_mask"]
+
+        sigmas = None
+        if scheduler == "turbo":
+            timesteps = torch.flip(torch.arange(1, 11) * 100 - 1, (0,))[:steps]
+            sigmas = model.model.model_sampling.sigma(timesteps)
+            sigmas = torch.cat([sigmas, sigmas.new_zeros([1])])
+        elif scheduler == "ltx" or scheduler == "ltxv-image":
+            from comfy_extras.nodes_lt import LTXVScheduler
+            sigmas = LTXVScheduler().get_sigmas(steps, 2.05, 0.95, True, 0.1, latent_image if scheduler == "ltxv-image" else None)[0]
+        elif scheduler == "align_your_steps":
+            if isinstance(model.model, SDXL):
+                model_type = "SDXL"
+            elif isinstance(model.model, SVD_img2vid):
+                model_type = "SVD"
+            elif isinstance(model.model, Flux):
+                model_type = "Flux"
+            elif isinstance(model.model, WAN21):
+                model_type = "Wan"
+            elif isinstance(model.model, Chroma):
+                model_type = "Chroma"
+            else:
+                print(f"AlignYourSteps: Unknown model type: {type(model.model)}, defaulting to SD1")
+                model_type = "SD1"
+            sigmas = AYS_NOISE_LEVELS[model_type][:]
+            if (steps + 1) != len(sigmas):
+                sigmas = loglinear_interp(sigmas, steps + 1)
+            sigmas[-1] = 0
+            sigmas = torch.FloatTensor(sigmas)
+        elif sigma_min >= 0 and sigma_max >= 0 and scheduler in ["karras", "exponential"]:
+            if sampler_name in ['dpm_2', 'dpm_2_ancestral']:
+                sigmas = calculate_sigmas_scheduler(model, scheduler, steps + 1, sigma_min, sigma_max, rho)
+                sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
+            else:
+                sigmas = calculate_sigmas_scheduler(model, scheduler, steps, sigma_min, sigma_max, rho)
+            sigmas = sigmas.to(device)
+        
+        out = latent_image.copy()
+        if steps > 0:
+            callback = make_swarm_sampler_callback(steps, device, model, previews)
+
+            samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_samples,
+                                    denoise=1.0, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step,
+                                    force_full_denoise=return_with_leftover_noise == "disable", noise_mask=noise_mask, sigmas=sigmas, callback=callback, seed=noise_seed)
+            out["samples"] = samples
+        return (out, )
+
+    # tiled sample version of sample function
+    def tiled_sample(self, model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, var_seed, var_seed_strength, sigma_max, sigma_min, rho, add_noise, return_with_leftover_noise, previews, tile_size):
+        out = latent_image.copy()
+        # split image into tiles
+        latent_samples = latent_image["samples"]
+        tiles = split_latent_tensor(latent_samples, tile_size=tile_size)
+        # resample each tile using self.sample
+        resampled_tiles = []
+        for coords, tile in tiles:
+            resampled_tile = self.sample(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, {"samples": tile}, start_at_step, end_at_step, var_seed, var_seed_strength, sigma_max, sigma_min, rho, add_noise, return_with_leftover_noise, previews)
+            resampled_tiles.append((coords, resampled_tile[0]["samples"]))
+        # stitch the tiles to get the final upscaled image
+        result = stitch_latent_tensors(latent_samples.shape, resampled_tiles)
+        out["samples"] = result
+        return (out,)
+
+    def run_sampling(self, model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, var_seed, var_seed_strength, sigma_max, sigma_min, rho, add_noise, return_with_leftover_noise, previews, tile_sample,  tile_size):
+        if tile_sample:
+            return self.tiled_sample(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, var_seed, var_seed_strength, sigma_max, sigma_min, rho, add_noise, return_with_leftover_noise, previews, tile_size)
+        else:
+            return self.sample(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, var_seed, var_seed_strength, sigma_max, sigma_min, rho, add_noise, return_with_leftover_noise, previews)
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmKSampler": SwarmKSampler,
+}

SwarmLatents.py

@@ -0,0 +1,37 @@
+import comfy, torch
+
+class SwarmOffsetEmptyLatentImage:
+    def __init__(self):
+        self.device = comfy.model_management.intermediate_device()
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "width": ("INT", {"default": 512, "min": 16, "max": 4096, "step": 8}),
+                "height": ("INT", {"default": 512, "min": 16, "max": 4096, "step": 8}),
+                "off_a": ("INT", {"default": 0, "min": -10, "max": 10, "step": 0.0001}),
+                "off_b": ("INT", {"default": 0, "min": -10, "max": 10, "step": 0.0001}),
+                "off_c": ("INT", {"default": 0, "min": -10, "max": 10, "step": 0.0001}),
+                "off_d": ("INT", {"default": 0, "min": -10, "max": 10, "step": 0.0001}),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})
+            }
+        }
+
+    CATEGORY = "SwarmUI/latents"
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "generate"
+    DESCRIPTION = "Generates a latent image with 4 channels, each channel filled with a different offset value. Designed to allow alternate empty value offsets for SDv1 and SDXL."
+
+    def generate(self, width, height, off_a, off_b, off_c, off_d, batch_size=1):
+        latent = torch.zeros([batch_size, 4, height // 8, width // 8], device=self.device)
+        latent[:, 0, :, :] = off_a
+        latent[:, 1, :, :] = off_b
+        latent[:, 2, :, :] = off_c
+        latent[:, 3, :, :] = off_d
+        return ({"samples":latent}, )
+
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmOffsetEmptyLatentImage": SwarmOffsetEmptyLatentImage
+}

SwarmLoadImageB64.py

@@ -0,0 +1,47 @@
+from PIL import Image, ImageOps
+import numpy as np
+import torch, base64, io
+
+def b64_to_img_and_mask(image_base64):
+    imageData = base64.b64decode(image_base64)
+    i = Image.open(io.BytesIO(imageData))
+    if hasattr(i, 'is_animated') and i.is_animated:
+        images = []
+        for frame in range(i.n_frames):
+            i.seek(frame)
+            images.append(i.convert("RGB"))
+        i.seek(0)
+        image = np.array(images).astype(np.float32) / 255.0
+        image = torch.from_numpy(image)
+    else:
+        i = ImageOps.exif_transpose(i)
+        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. - torch.from_numpy(mask)
+    else:
+        mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
+    return (image, mask.unsqueeze(0))
+
+class SwarmLoadImageB64:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image_base64": ("STRING", {"multiline": True})
+            }
+        }
+
+    CATEGORY = "SwarmUI/images"
+    RETURN_TYPES = ("IMAGE", "MASK")
+    FUNCTION = "load_image_b64"
+    DESCRIPTION = "Loads an image from a base64 string. Works like a regular LoadImage node, but with input format designed to be easier to use through automated calls, including SwarmUI with custom workflows."
+
+    def load_image_b64(self, image_base64):
+        return b64_to_img_and_mask(image_base64)
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmLoadImageB64": SwarmLoadImageB64,
+}

SwarmLoraLoader.py

@@ -0,0 +1,56 @@
+import comfy
+import folder_paths
+
+class SwarmLoraLoader:
+    def __init__(self):
+        self.loaded_lora = None
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL", ),
+                "clip": ("CLIP", ),
+                "lora_names": ("STRING", {"multiline": True, "tooltip": "Comma separated list of lora names to load."}),
+                "lora_weights": ("STRING", {"multiline": True, "tooltip": "Comma separated list of lora weights to apply to each lora. Must match the number of loras."}),
+            }
+        }
+
+    CATEGORY = "SwarmUI/models"
+    RETURN_TYPES = ("MODEL", "CLIP")
+    FUNCTION = "load_loras"
+    DESCRIPTION = "Like a regular LoRA Loader, but designed to take a dynamic list of loras and weights, to allow easier integration with SwarmUI custom workflows."
+
+    def load_loras(self, model, clip, lora_names, lora_weights):
+        if lora_names.strip() == "":
+            return (model, clip)
+
+        lora_names = lora_names.split(",")
+        lora_weights = lora_weights.split(",")
+        lora_weights = [float(x.strip()) for x in lora_weights]
+
+        for i in range(len(lora_names)):
+            lora_name = lora_names[i].strip()
+            weight = lora_weights[i]
+            if weight == 0:
+                continue
+            # This section copied directly from default comfy LoraLoader
+            lora_path = folder_paths.get_full_path("loras", lora_name)
+            lora = None
+            if self.loaded_lora is not None:
+                if self.loaded_lora[0] == lora_path:
+                    lora = self.loaded_lora[1]
+                else:
+                    temp = self.loaded_lora
+                    self.loaded_lora = None
+                    del temp
+            if lora is None:
+                lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
+                self.loaded_lora = (lora_path, lora)
+            model, clip = comfy.sd.load_lora_for_models(model, clip, lora, weight, weight)
+
+        return (model, clip)
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmLoraLoader": SwarmLoraLoader,
+}

SwarmMasks.py

@@ -0,0 +1,287 @@
+import torch, comfy
+
+intermediate_device = comfy.model_management.intermediate_device()
+main_device = comfy.model_management.get_torch_device()
+
+class SwarmSquareMaskFromPercent:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "x": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, "round": 0.0001, "tooltip": "The x position of the mask as a percentage of the image size."}),
+                "y": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, "round": 0.0001, "tooltip": "The y position of the mask as a percentage of the image size."}),
+                "width": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, "round": 0.0001, "tooltip": "The width of the mask as a percentage of the image size."}),
+                "height": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, "round": 0.0001, "tooltip": "The height of the mask as a percentage of the image size."}),
+                "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "tooltip": "The strength of the mask, ie the value of all masked pixels, leaving the rest black ie 0."}),
+            }
+        }
+
+    CATEGORY = "SwarmUI/masks"
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "mask_from_perc"
+    DESCRIPTION = "Creates a simple square mask with the specified dimensions and position, with the specified strength (ie value of all masked pixels, leaving the rest black ie 0)."
+
+    def mask_from_perc(self, x, y, width, height, strength):
+        SCALE = 256
+        mask = torch.zeros((SCALE, SCALE), dtype=torch.float32, device=intermediate_device)
+        mask[int(y*SCALE):int((y+height)*SCALE), int(x*SCALE):int((x+width)*SCALE)] = strength
+        return (mask.unsqueeze(0),)
+
+
+def mask_size_match(mask_a, mask_b):
+    if len(mask_a.shape) == 2:
+        mask_a = mask_a.unsqueeze(0)
+    if len(mask_b.shape) == 2:
+        mask_b = mask_b.unsqueeze(0)
+    height = max(mask_a.shape[1], mask_b.shape[1])
+    width = max(mask_a.shape[2], mask_b.shape[2])
+    if mask_a.shape[1] != height or mask_a.shape[2] != width:
+        mask_a = torch.nn.functional.interpolate(mask_a.unsqueeze(0), size=(height, width), mode="bicubic")[0]
+    if mask_b.shape[1] != height or mask_b.shape[2] != width:
+        mask_b = torch.nn.functional.interpolate(mask_b.unsqueeze(0), size=(height, width), mode="bicubic")[0]
+    return (mask_a, mask_b)
+
+
+class SwarmOverMergeMasksForOverlapFix:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "mask_a": ("MASK",),
+                "mask_b": ("MASK",),
+            }
+        }
+
+    CATEGORY = "SwarmUI/masks"
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "mask_overmerge"
+    DESCRIPTION = "Merges two masks by simply adding them together, without any overlap handling. Intended for use with the Overlap nodes."
+
+    def mask_overmerge(self, mask_a, mask_b):
+        mask_a, mask_b = mask_size_match(mask_a, mask_b)
+        mask_sum = mask_a + mask_b
+        return (mask_sum,)
+
+
+class SwarmCleanOverlapMasks:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "mask_a": ("MASK",),
+                "mask_b": ("MASK",),
+            }
+        }
+
+    CATEGORY = "SwarmUI/masks"
+    RETURN_TYPES = ("MASK","MASK",)
+    FUNCTION = "mask_overlap"
+    DESCRIPTION = "Normalizes the overlap between two masks, such that where they overlap each mask will receive only partial strength that sums to no more than 1.0. This allows you to then add the masks together and the result will not exceed 1 at any point."
+
+    def mask_overlap(self, mask_a, mask_b):
+        mask_a, mask_b = mask_size_match(mask_a, mask_b)
+        mask_sum = mask_a + mask_b
+        mask_sum = mask_sum.clamp(1.0, 9999.0)
+        mask_a = mask_a / mask_sum
+        mask_b = mask_b / mask_sum
+        return (mask_a, mask_b)
+
+
+class SwarmCleanOverlapMasksExceptSelf:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "mask_self": ("MASK",),
+                "mask_merged": ("MASK",),
+            }
+        }
+
+    CATEGORY = "SwarmUI/masks"
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "mask_clean"
+    DESCRIPTION = "If masks have been overmerged, this takes a single mask and grabs just the CleanOverlap result for the one mask relative to the overmerge result."
+
+    def mask_clean(self, mask_self, mask_merged):
+        mask_self, mask_merged = mask_size_match(mask_self, mask_merged)
+        mask_sum = mask_merged.clamp(1.0, 9999.0)
+        mask_self = mask_self / mask_sum
+        return (mask_self,)
+
+
+class SwarmExcludeFromMask:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "main_mask": ("MASK",),
+                "exclude_mask": ("MASK",),
+            }
+        }
+
+    CATEGORY = "SwarmUI/masks"
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "mask_exclude"
+    DESCRIPTION = "Excludes the area of the exclude mask from the main mask, such that the main mask will be black in the area of the exclude mask. This is a simple subtract and clamp."
+
+    def mask_exclude(self, main_mask, exclude_mask):
+        main_mask, exclude_mask = mask_size_match(main_mask, exclude_mask)
+        main_mask = main_mask - exclude_mask
+        main_mask = main_mask.clamp(0.0, 1.0)
+        return (main_mask,)
+
+
+class SwarmMaskBounds:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "mask": ("MASK",),
+                "grow": ("INT", {"default": 0, "min": 0, "max": 1024, "tooltip": "Number of pixels to grow the mask by."}),
+            },
+            "optional": {
+                "aspect_x": ("INT", {"default": 0, "min": 0, "max": 4096, "tooltip": "An X width value, used to indicate a target aspect ratio. 0 to allow any aspect."}),
+                "aspect_y": ("INT", {"default": 0, "min": 0, "max": 4096, "tooltip": "A Y height value, used to indicate a target aspect ratio. 0 to allow any aspect."}),
+            }
+        }
+
+    CATEGORY = "SwarmUI/masks"
+    RETURN_TYPES = ("INT", "INT", "INT", "INT")
+    RETURN_NAMES = ("x", "y", "width", "height")
+    FUNCTION = "get_bounds"
+    DESCRIPTION = "Returns the bounding box of the mask (as pixel coordinates x,y,width,height), optionally grown by the number of pixels specified in 'grow'."
+
+    def get_bounds(self, mask, grow, aspect_x=0, aspect_y=0):
+        if len(mask.shape) == 3:
+            mask = mask[0]
+        sum_x = (torch.sum(mask, dim=0) != 0).to(dtype=torch.int)
+        sum_y = (torch.sum(mask, dim=1) != 0).to(dtype=torch.int)
+        def getval(arr, direction):
+            val = torch.argmax(arr).item()
+            val += grow * direction
+            val = max(0, min(val, arr.shape[0] - 1))
+            return val
+        x_start = getval(sum_x, -1)
+        x_end = mask.shape[1] - getval(sum_x.flip(0), -1)
+        y_start = getval(sum_y, -1)
+        y_end = mask.shape[0] - getval(sum_y.flip(0), -1)
+        if aspect_x > 0 and aspect_y > 0:
+            actual_aspect = aspect_x / aspect_y
+            width = x_end - x_start
+            height = y_end - y_start
+            found_aspect = width / height
+            if found_aspect > actual_aspect:
+                desired_height = width / actual_aspect
+                y_start = max(0, y_start - (desired_height - height) / 2)
+                y_end = min(mask.shape[0], y_start + desired_height)
+            else:
+                desired_width = height * actual_aspect
+                x_start = max(0, x_start - (desired_width - width) / 2)
+                x_end = min(mask.shape[1], x_start + desired_width)
+        return (int(x_start), int(y_start), int(x_end - x_start), int(y_end - y_start))
+
+
+class SwarmMaskGrow:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "mask": ("MASK",),
+                "grow": ("INT", {"default": 0, "min": 0, "max": 1024, "tooltip": "Number of pixels to grow the mask by."}),
+            }
+        }
+
+    CATEGORY = "SwarmUI/masks"
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "grow"
+    DESCRIPTION = "Expands the contents of the max, such that masked (white) areas grow and cover the unmasked (black) areas by the number of pixels specified in 'grow'."
+
+    def grow(self, mask, grow):
+        while mask.ndim < 4:
+            mask = mask.unsqueeze(0)
+        mask = mask.to(device=main_device)
+        # iterate rather than all at once - this avoids padding and runs much faster for large sizes
+        for _ in range((grow + 1) // 2):
+            mask = torch.nn.functional.max_pool2d(mask, kernel_size=3, stride=1, padding=1)
+        return (mask.to(device=intermediate_device),)
+
+
+# Blur code is copied out of ComfyUI's default ImageBlur
+def gaussian_kernel(kernel_size: int, sigma: float, device=None):
+    x, y = torch.meshgrid(torch.linspace(-1, 1, kernel_size, device=device), torch.linspace(-1, 1, kernel_size, device=device), indexing="ij")
+    d = torch.sqrt(x * x + y * y)
+    g = torch.exp(-(d * d) / (2.0 * sigma * sigma))
+    return g / g.sum()
+
+
+class SwarmMaskBlur:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "mask": ("MASK",),
+                "blur_radius": ("INT", { "default": 1, "min": 1, "max": 64, "step": 1, "tooltip": "The radius of the blur kernel." }),
+                "sigma": ("FLOAT", { "default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1, "tooltip": "The standard deviation of the Gaussian blur kernel." }),
+            },
+        }
+
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "blur"
+    CATEGORY = "SwarmUI/masks"
+    DESCRIPTION = "Blurs the contents of the mask."
+
+    def blur(self, mask, blur_radius, sigma):
+        if blur_radius == 0:
+            return (mask,)
+        mask = mask.to(device=main_device)
+        kernel_size = blur_radius * 2 + 1
+        kernel = gaussian_kernel(kernel_size, sigma, device=mask.device).repeat(1, 1, 1).unsqueeze(1)
+        while mask.ndim < 4:
+            mask = mask.unsqueeze(0)
+        padded_mask = torch.nn.functional.pad(mask, (blur_radius,blur_radius,blur_radius,blur_radius), 'reflect')
+        blurred = torch.nn.functional.conv2d(padded_mask, kernel, padding=kernel_size // 2, groups=1)[:,:,blur_radius:-blur_radius, blur_radius:-blur_radius]
+        blurred = blurred.squeeze(0).squeeze(0)
+        mask = mask.to(device=intermediate_device)
+        return (blurred.to(device=intermediate_device),)
+
+
+class SwarmMaskThreshold:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "mask": ("MASK",),
+                "min": ("FLOAT", { "default": 0.2, "min": 0, "max": 1, "step": 0.01, "tooltip": "The minimum value to threshold the mask to." }),
+                "max": ("FLOAT", { "default": 0.8, "min": 0, "max": 1, "step": 0.01, "tooltip": "The maximum value to threshold the mask to." }),
+            },
+        }
+
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "threshold"
+    CATEGORY = "SwarmUI/masks"
+    DESCRIPTION = "Thresholds the mask to the specified range, clamping any lower or higher values and rescaling the range to 0-1."
+
+    def threshold(self, mask, min, max):
+        mask = mask.clamp(min, max)
+        mask = mask - min
+        mask = mask / (max - min)
+        return (mask,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmSquareMaskFromPercent": SwarmSquareMaskFromPercent,
+    "SwarmCleanOverlapMasks": SwarmCleanOverlapMasks,
+    "SwarmCleanOverlapMasksExceptSelf": SwarmCleanOverlapMasksExceptSelf,
+    "SwarmExcludeFromMask": SwarmExcludeFromMask,
+    "SwarmOverMergeMasksForOverlapFix": SwarmOverMergeMasksForOverlapFix,
+    "SwarmMaskBounds": SwarmMaskBounds,
+    "SwarmMaskGrow": SwarmMaskGrow,
+    "SwarmMaskBlur": SwarmMaskBlur,
+    "SwarmMaskThreshold": SwarmMaskThreshold,
+}

SwarmMath.py

@@ -0,0 +1,23 @@
+
+class SwarmIntAdd:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "a": ("INT", {"default": 0, "min": -2147483647, "max": 2147483647}),
+                "b": ("INT", {"default": 0, "min": -2147483647, "max": 2147483647})
+            }
+        }
+
+    CATEGORY = "SwarmUI/math"
+    RETURN_TYPES = ("INT",)
+    FUNCTION = "add"
+    DESCRIPTION = "Adds two integers. Use a negative number to subtract."
+
+    def add(self, a, b):
+        return (a + b,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmIntAdd": SwarmIntAdd
+}

SwarmReference.py

@@ -0,0 +1,56 @@
+import torch
+
+# This code copied from https://github.com/comfyanonymous/ComfyUI_experiments/blob/master/reference_only.py
+# And modified to work better in Swarm generated workflows
+
+class SwarmReferenceOnly:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "reference": ("LATENT",),
+                "latent": ("LATENT",)
+            }
+        }
+
+    CATEGORY = "SwarmUI/sampling"
+    RETURN_TYPES = ("MODEL", "LATENT")
+    FUNCTION = "reference_only"
+    DESCRIPTION = "Applies 'reference only' image-prompting to the generation. Must forward the new model, and the new latent, to work properly."
+
+    def reference_only(self, model, reference, latent):
+        model_reference = model.clone()
+        reference["samples"] = torch.nn.functional.interpolate(reference["samples"], size=(latent["samples"].shape[2], latent["samples"].shape[3]), mode="bilinear")
+
+        batch = latent["samples"].shape[0] + reference["samples"].shape[0]
+        def reference_apply(q, k, v, extra_options):
+            k = k.clone().repeat(1, 2, 1)
+            offset = 0
+            if q.shape[0] > batch:
+                offset = batch
+
+            for o in range(0, q.shape[0], batch):
+                for x in range(1, batch):
+                    k[x + o, q.shape[1]:] = q[o,:]
+
+            return q, k, k
+
+        model_reference.set_model_attn1_patch(reference_apply)
+        out_latent = torch.cat((reference["samples"], latent["samples"]))
+        if "noise_mask" in latent:
+            mask = latent["noise_mask"]
+        else:
+            mask = torch.ones((64,64), dtype=torch.float32, device="cpu")
+
+        if len(mask.shape) < 3:
+            mask = mask.unsqueeze(0)
+        if mask.shape[0] < latent["samples"].shape[0]:
+            mask = mask.repeat(latent["samples"].shape[0], 1, 1)
+
+        out_mask = torch.zeros((1,mask.shape[1],mask.shape[2]), dtype=torch.float32, device="cpu")
+        return (model_reference, {"samples": out_latent, "noise_mask": torch.cat((out_mask, mask))})
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmReferenceOnly": SwarmReferenceOnly,
+}

SwarmSaveImageWS.py

@@ -0,0 +1,152 @@
+from PIL import Image
+import numpy as np
+import comfy.utils
+from server import PromptServer, BinaryEventTypes
+import time, io, struct
+
+SPECIAL_ID = 12345 # Tells swarm that the node is going to output final images
+VIDEO_ID = 12346
+TEXT_ID = 12347
+
+def send_image_to_server_raw(type_num: int, save_me: callable, id: int, event_type: int = BinaryEventTypes.PREVIEW_IMAGE):
+    out = io.BytesIO()
+    header = struct.pack(">I", type_num)
+    out.write(header)
+    save_me(out)
+    out.seek(0)
+    preview_bytes = out.getvalue()
+    server = PromptServer.instance
+    server.send_sync("progress", {"value": id, "max": id}, sid=server.client_id)
+    server.send_sync(event_type, preview_bytes, sid=server.client_id)
+
+class SwarmSaveImageWS:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE", ),
+            },
+            "optional": {
+                "bit_depth": (["8bit", "16bit", "raw"], {"default": "8bit"})
+            }
+        }
+
+    CATEGORY = "SwarmUI/images"
+    RETURN_TYPES = ()
+    FUNCTION = "save_images"
+    OUTPUT_NODE = True
+    DESCRIPTION = "Acts like a special version of 'SaveImage' that doesn't actual save to disk, instead it sends directly over websocket. This is intended so that SwarmUI can save the image itself rather than having Comfy's Core save it."
+
+    def save_images(self, images, bit_depth = "8bit"):
+        pbar = comfy.utils.ProgressBar(SPECIAL_ID)
+        step = 0
+        for image in images:
+            if bit_depth == "raw":
+                i = 255.0 * image.cpu().numpy()
+                img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+                def do_save(out):
+                    img.save(out, format='BMP')
+                send_image_to_server_raw(1, do_save, SPECIAL_ID, event_type=10)
+            elif bit_depth == "16bit":
+                i = 65535.0 * image.cpu().numpy()
+                img = self.convert_img_16bit(np.clip(i, 0, 65535).astype(np.uint16))
+                send_image_to_server_raw(2, lambda out: out.write(img), SPECIAL_ID)
+            else:
+                i = 255.0 * image.cpu().numpy()
+                img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+                pbar.update_absolute(step, SPECIAL_ID, ("PNG", img, None))
+            step += 1
+
+        return {}
+
+    def convert_img_16bit(self, img_np):
+        try:
+            import cv2
+            img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
+            success, img_encoded = cv2.imencode('.png', img_np)
+
+            if img_encoded is None or not success:
+                raise RuntimeError("OpenCV failed to encode image.")
+
+            return img_encoded.tobytes()
+        except Exception as e:
+            print(f"Error converting OpenCV image to PIL: {e}")
+            raise
+
+    @classmethod
+    def IS_CHANGED(s, images, bit_depth = "8bit"):
+        return time.time()
+
+
+class SwarmSaveAnimatedWebpWS:
+    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, "tooltip": "Frames per second, must match the actual generated speed or else you will get slow/fast motion."}),
+                "lossless": ("BOOLEAN", {"default": True, "tooltip": "If true, the image will be saved losslessly, otherwise it will be saved with the quality specified. Lossless is best quality, but takes more file space."}),
+                "quality": ("INT", {"default": 80, "min": 0, "max": 100, "tooltip": "Quality of the image as a percentage, only used if lossless is false. Smaller values save more space but look worse. 80 is a fine general value."}),
+                "method": (list(s.methods.keys()),),
+            },
+        }
+
+    CATEGORY = "SwarmUI/video"
+    RETURN_TYPES = ()
+    FUNCTION = "save_images"
+    OUTPUT_NODE = True
+    DESCRIPTION = "Acts like a special version of 'SaveAnimatedWEBP' that doesn't actual save to disk, instead it sends directly over websocket. This is intended so that SwarmUI can save the image itself rather than having Comfy's Core save it."
+
+    def save_images(self, images, fps, lossless, quality, method):
+        method = self.methods.get(method)
+        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)
+
+        def do_save(out):
+            pil_images[0].save(out, save_all=True, duration=int(1000.0/fps), append_images=pil_images[1 : len(pil_images)], lossless=lossless, quality=quality, method=method, format='WEBP')
+        send_image_to_server_raw(3, do_save, VIDEO_ID)
+
+        return { }
+
+    @classmethod
+    def IS_CHANGED(s, images, fps, lossless, quality, method):
+        return time.time()
+
+
+class SwarmAddSaveMetadataWS:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "key": ("STRING", {"tooltip": "The key to add to the metadata tracker. Must be simple A-Z plain text or underscores."}),
+                "value": ("STRING", {"tooltip": "The value to add to the metadata tracker."}),
+            }
+        }
+
+    CATEGORY = "SwarmUI/images"
+    RETURN_TYPES = ()
+    FUNCTION = "add_save_metadata"
+    OUTPUT_NODE = True
+    DESCRIPTION = "Adds a metadata key/value pair to SwarmUI's metadata tracker for this generation, which will be appended to any images saved after this node triggers. Note that keys overwrite, not add. Any key can have only one value."
+
+    def add_save_metadata(self, key, value):
+        full_text = f"{key}:{value}"
+        full_text_bytes = full_text.encode('utf-8')
+        send_image_to_server_raw(0, lambda out: out.write(full_text_bytes), TEXT_ID, event_type=BinaryEventTypes.TEXT)
+        return {}
+
+    @classmethod
+    def IS_CHANGED(s, key, value):
+        return time.time()
+
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmSaveImageWS": SwarmSaveImageWS,
+    "SwarmSaveAnimatedWebpWS": SwarmSaveAnimatedWebpWS,
+    "SwarmAddSaveMetadataWS": SwarmAddSaveMetadataWS,
+}

SwarmTextHandling.py

@@ -0,0 +1,211 @@
+import torch, comfy
+from nodes import MAX_RESOLUTION
+
+
+# LLaMA template for Hunyuan Image2Video.
+# This is actually a single-line monstrosity due to the way it's formatted.
+# This is probably an accident from the python devs misunderstanding how string lines work,
+# but, well, we're just matching what they did and that's what they did.
+PROMPT_TEMPLATE_ENCODE_VIDEO_I2V = (
+    "<|start_header_id|>system<|end_header_id|>\n\n<image>\nDescribe the video by detailing the following aspects according to the reference image: "
+    "1. The main content and theme of the video."
+    "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects."
+    "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects."
+    "4. background environment, light, style and atmosphere."
+    "5. camera angles, movements, and transitions used in the video:<|eot_id|>\n\n"
+    "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>"
+    "<|start_header_id|>assistant<|end_header_id|>\n\n"
+)
+# LLaMA template for Qwen Image Edit Plus.
+PROMPT_TEMPLATE_QWEN_IMAGE_EDIT_PLUS = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+
+class SwarmClipTextEncodeAdvanced:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "clip": ("CLIP", ),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "How many sampling steps will be ran - this is needed for per-step features (from-to/alternate/...) to work properly."}),
+                "prompt": ("STRING", {"multiline": True, "dynamicPrompts": True, "tooltip": "Your actual prompt text."} ),
+                "width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION, "tooltip": "Intended width of the image, used by some models (eg SDXL)."}),
+                "height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION, "tooltip": "Intended height of the image, used by some models (eg SDXL)."}),
+                "target_width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION, "tooltip": "Actual width of the image, used by some models (eg SDXL)."}),
+                "target_height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION, "tooltip": "Actual height of the image, used by some models (eg SDXL)."}),
+            },
+            "optional": {
+                "guidance": ("FLOAT", {"default": -1, "min": -1, "max": 100.0, "step": 0.1, "tooltip": "Guidance value to embed, used by some models (eg Flux)."}),
+                "llama_template": ("STRING", {"default": "", "multiline": True, "tooltip": "Template for the LLaMA model, if applicable."}),
+                "clip_vision_output": ("CLIP_VISION_OUTPUT", {"default": None, "tooltip": "Optional CLIP Vision Output to use for the LLaMA model, if applicable."}),
+                "images": ("IMAGE", {"default": None, "tooltip": "Optional images to use for a text-vision model, if applicable."}),
+            }
+        }
+
+    CATEGORY = "SwarmUI/clip"
+    RETURN_TYPES = ("CONDITIONING",)
+    FUNCTION = "encode"
+    DESCRIPTION = "Acts like the regular CLIPTextEncode, but supports more advanced special features like '<break>', '[from:to:when]', '[alter|nate]', ..."
+
+    def encode(self, clip, steps: int, prompt: str, width: int, height: int, target_width: int, target_height: int, guidance: float = -1, llama_template = None, clip_vision_output = None, images = None):
+        image_prompt = ""
+        if llama_template == "hunyuan_image":
+            llama_template = PROMPT_TEMPLATE_ENCODE_VIDEO_I2V
+        elif llama_template == "qwen_image_edit_plus":
+            llama_template = PROMPT_TEMPLATE_QWEN_IMAGE_EDIT_PLUS
+            if images is not None:
+                if len(images.shape) == 3:
+                    images = [images]
+                else:
+                    images = [i.unsqueeze(0) for i in images]
+                for i, image in enumerate(images):
+                    image_prompt += f"Picture {i + 1}: <|vision_start|><|image_pad|><|vision_end|>"
+
+        def tokenize(text: str):
+            if clip_vision_output is not None:
+                return clip.tokenize(text, llama_template=llama_template, image_embeds=clip_vision_output.mm_projected)
+            elif images is not None:
+                return clip.tokenize(image_prompt + text, llama_template=llama_template, images=images)
+            else:
+                return clip.tokenize(text)
+
+        encoding_cache = {}
+
+        def text_to_cond(text: str, start_percent: float, end_percent: float):
+            text = text.replace("\0\1", "[").replace("\0\2", "]").replace("\0\3", "embedding:")
+            if text in encoding_cache:
+                cond_arr = encoding_cache[text]
+            else:
+                cond_chunks = text.split("<break>")
+                tokens = tokenize(cond_chunks[0])
+                cond_arr = clip.encode_from_tokens_scheduled(tokens)
+                if len(cond_chunks) > 1:
+                    for chunk in cond_chunks[1:]:
+                        tokens = tokenize(chunk)
+                        cond_arr_chunk = clip.encode_from_tokens_scheduled(tokens)
+                        catted_cond = torch.cat([cond_arr[0][0], cond_arr_chunk[0][0]], dim=1)
+                        cond_arr[0] = [catted_cond, cond_arr[0][1]]
+                encoding_cache[text] = cond_arr
+            result = {"pooled_output": cond_arr[0][1]["pooled_output"], "width": width, "height": height, "crop_w": 0, "crop_h": 0, "target_width": target_width, "target_height": target_height, "start_percent": start_percent, "end_percent": end_percent}
+            if guidance >= 0:
+                result["guidance"] = guidance
+            out_cond_arr = [[cond_arr[0][0], result]]
+            out_cond_arr.extend(cond_arr[1:])
+            return out_cond_arr
+
+        prompt = prompt.replace("\\[", "\0\1").replace("\\]", "\0\2").replace("embedding:", "\0\3")
+
+        chunks = []
+        any = [False]
+        escapable = ["\\", "[", "]", ":", "|", "(", ")", "<", ">"]
+
+        def append_chunk(text: str, applies_to: list[int], can_subprocess: bool, limit_to: list[int]):
+            applies_to = [i for i in applies_to if i in limit_to]
+            fixed_text = ""
+            do_skip = False
+            for i in range(len(text)):
+                if text[i] == "\\" and not do_skip and i + 1 < len(text) and text[i + 1] in escapable:
+                    do_skip = True
+                else:
+                    do_skip = False
+                    fixed_text += text[i]
+            if can_subprocess and '[' in fixed_text:
+                get_chunks(fixed_text, applies_to)
+            else:
+                chunks.append({'text': text, 'applies_to': applies_to})
+
+        def get_chunks(remaining: str, limit_to: list[int] = [i for i in range(steps)]):
+            while True:
+                start = remaining.find("[")
+                if start == -1:
+                    append_chunk(remaining, [i for i in range(steps)], False, limit_to)
+                    break
+
+                end = -1
+                count = 0
+                do_skip = False
+                colon_indices = []
+                pipe_indices = []
+                for i in range(start + 1, len(remaining)):
+                    char = remaining[i]
+                    if char == "\\" and not do_skip and i + 1 < len(remaining) and remaining[i + 1] in escapable:
+                        do_skip = True
+                    elif do_skip:
+                        do_skip = False
+                    elif char == "[":
+                        count += 1
+                    elif char == "]":
+                        if count == 0:
+                            end = i
+                            break
+                        count -= 1
+                    elif char == ":" and count == 0 and len(pipe_indices) == 0:
+                        colon_indices.append(i)
+                    elif char == "|" and count == 0 and len(colon_indices) == 0:
+                        pipe_indices.append(i)
+
+                if end == -1:
+                    chunks[-1].text += remaining
+                    break
+                append_chunk(remaining[:start], [i for i in range(steps)], False, limit_to)
+                control = remaining[start + 1:end]
+
+                if len(pipe_indices) > 0:
+                    data = split_text_on(control, pipe_indices, start + 1)
+                    for i in range(len(data)):
+                        append_chunk(data[i], [step for step in range(steps) if step % len(data) == i], True, limit_to)
+                    any[0] = True
+                elif len(colon_indices) == 2:
+                    coloned = split_text_on(control, colon_indices, start + 1)
+                    when = float(coloned[2])
+                    if when < 1:
+                        when = when * steps
+                    append_chunk(coloned[0], [i for i in range(steps) if i < when], True, limit_to)
+                    append_chunk(coloned[1], [i for i in range(steps) if i >= when], True, limit_to)
+                    any[0] = True
+                elif len(colon_indices) == 1:
+                    coloned = split_text_on(control, colon_indices, start + 1)
+                    when = float(coloned[1])
+                    if when < 1:
+                        when = when * steps
+                    append_chunk(coloned[0], [i for i in range(steps) if i >= when], True, limit_to)
+                    any[0] = True
+                else:
+                    append_chunk(control, [i for i in range(steps)], False, limit_to)
+
+                remaining = remaining[end + 1:]
+
+        get_chunks(prompt)
+
+        if not any[0]:
+            return (text_to_cond(prompt, 0, 1), )
+
+        conds_out = []
+        last_text = ""
+        start_perc = 0
+        for i in range(steps):
+            perc = i / steps
+            text = ""
+            for chunk in chunks:
+                if i in chunk['applies_to']:
+                    text += chunk['text']
+            if text != last_text or i == 0:
+                if i != 0:
+                    conds_out.extend(text_to_cond(last_text, start_perc - 0.001, perc + 0.001))
+                last_text = text
+                start_perc = perc
+        conds_out.extend(text_to_cond(last_text, start_perc - 0.001, 1))
+        return (conds_out, )
+
+
+def split_text_on(text: str, indices: list[str], offset: int) -> list[str]:
+    indices = [i - offset for i in indices]
+    result = []
+    result.append(text[:indices[0]])
+    for i in range(len(indices) - 1):
+        result.append(text[indices[i] + 1:indices[i + 1]])
+    result.append(text[indices[-1] + 1:])
+    return result
+
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmClipTextEncodeAdvanced": SwarmClipTextEncodeAdvanced,
+}

SwarmTiling.py

@@ -0,0 +1,86 @@
+import torch, copy
+from torch.nn import functional as F
+
+def make_circular_assym(m, assym_mode):
+    def _conv_forward(self, input, weight, bias):
+        if self.padding_mode == "x_circular":
+            padded = F.pad(input, (self._reversed_padding_repeated_twice[0], self._reversed_padding_repeated_twice[1], 0, 0), "circular")
+            padded = F.pad(padded, (0, 0, self._reversed_padding_repeated_twice[2], self._reversed_padding_repeated_twice[3]), "constant", 0)
+            return F.conv2d(padded, weight, bias, self.stride, (0, 0), self.dilation, self.groups)
+        elif self.padding_mode == "y_circular":
+            padded = F.pad(input, (self._reversed_padding_repeated_twice[0], self._reversed_padding_repeated_twice[1], 0, 0), "constant", 0)
+            padded = F.pad(padded, (0, 0, self._reversed_padding_repeated_twice[2], self._reversed_padding_repeated_twice[3]), "circular")
+            return F.conv2d(padded, weight, bias, self.stride, (0, 0), self.dilation, self.groups)
+        elif self.padding_mode != "zeros":
+            padded = F.pad(input, self._reversed_padding_repeated_twice, mode=self.padding_mode)
+            return F.conv2d(padded, weight, bias, self.stride, (0, 0), self.dilation, self.groups)
+        else:
+            return F.conv2d(input, weight, bias, self.stride, self.padding, self.dilation, self.groups)
+    if isinstance(m, torch.nn.Conv2d):
+        m._conv_forward = _conv_forward.__get__(m, torch.nn.Conv2d)
+        m.padding_mode = assym_mode
+
+def make_circular(m):
+    if isinstance(m, torch.nn.Conv2d):
+        m.padding_mode = "circular"
+
+class SwarmModelTiling:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL", ),
+            },
+            "optional": {
+                "tile_axis": (["Both", "X", "Y"], )
+            }
+        }
+
+    CATEGORY = "SwarmUI/sampling"
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "adapt"
+    DESCRIPTION = "Adapts a model to use circular padding to enable tiled image results. Only works on UNet based models (eg SDv1, SDXL), not on DiT models (eg SD3, Flux). Use with SwarmTileableVAE."
+
+    def adapt(self, model, tile_axis=None):
+        m = copy.deepcopy(model)
+        if tile_axis is not None and tile_axis != "Both":
+            if tile_axis == "X":
+                m.model.apply(lambda x: make_circular_assym(x, "x_circular"))
+            elif tile_axis == "Y":
+                m.model.apply(lambda x: make_circular_assym(x, "y_circular"))
+        else:
+            m.model.apply(make_circular)
+        return (m,)
+
+class SwarmTileableVAE:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "vae": ("VAE", )
+            },
+            "optional": {
+                "tile_axis": (["Both", "X", "Y"], )
+            }
+        }
+
+    CATEGORY = "SwarmUI/sampling"
+    RETURN_TYPES = ("VAE",)
+    FUNCTION = "adapt"
+    DESCRIPTION = "Adapts a VAE to use circular padding to enable tiled image results. Use with SwarmModelTiling."
+
+    def adapt(self, vae, tile_axis=None):
+        vae = copy.deepcopy(vae)
+        if tile_axis is not None and tile_axis != "Both":
+            if tile_axis == "X":
+                vae.first_stage_model.apply(lambda x: make_circular_assym(x, "x_circular"))
+            elif tile_axis == "Y":
+                vae.first_stage_model.apply(lambda x: make_circular_assym(x, "y_circular"))
+        else:
+            vae.first_stage_model.apply(make_circular)
+        return (vae,)
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmModelTiling": SwarmModelTiling,
+    "SwarmTileableVAE": SwarmTileableVAE,
+}

SwarmUnsampler.py

@@ -0,0 +1,50 @@
+import torch, comfy
+from .SwarmKSampler import make_swarm_sampler_callback
+
+class SwarmUnsampler:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
+                "scheduler": (["turbo"] + comfy.samplers.KSampler.SCHEDULERS, ),
+                "positive": ("CONDITIONING", ),
+                "negative": ("CONDITIONING", ),
+                "latent_image": ("LATENT", ),
+                "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
+                "previews": (["default", "none", "one"], )
+            }
+        }
+
+    CATEGORY = "SwarmUI/sampling"
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "unsample"
+    DESCRIPTION = "Runs sampling in reverse. The function of this is to create noise that matches an image, such that you can the run forward sampling with an altered version of the unsampling prompt to get a closely altered image. May not work on all models, may not work perfectly. Input values should largely match your Sampler inputs."
+
+    def unsample(self, model, steps, sampler_name, scheduler, positive, negative, latent_image, start_at_step, previews):
+        device = comfy.model_management.get_torch_device()
+        latent_samples = latent_image["samples"].to(device)
+
+        noise = torch.zeros(latent_samples.size(), dtype=latent_samples.dtype, layout=latent_samples.layout, device=device)
+        noise_mask = None
+        if "noise_mask" in latent_image:
+            noise_mask = latent_image["noise_mask"]
+
+        sampler = comfy.samplers.KSampler(model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=1.0, model_options=model.model_options)
+        sigmas = sampler.sigmas.flip(0) + 0.0001
+
+        callback = make_swarm_sampler_callback(steps, device, model, previews)
+
+        samples = comfy.sample.sample(model, noise, steps, 1, sampler_name, scheduler, positive, negative, latent_samples,
+                                    denoise=1.0, disable_noise=False, start_step=0, last_step=steps - start_at_step,
+                                    force_full_denoise=False, noise_mask=noise_mask, sigmas=sigmas, callback=callback, seed=0)
+        out = latent_image.copy()
+        out["samples"] = samples
+        return (out, )
+
+
+NODE_CLASS_MAPPINGS = {
+    "SwarmUnsampler": SwarmUnsampler,
+}

__init__.py

@@ -0,0 +1,34 @@
+import os, folder_paths
+
+from . import SwarmBlending, SwarmClipSeg, SwarmImages, SwarmInternalUtil, SwarmKSampler, SwarmLoadImageB64, SwarmLoraLoader, SwarmMasks, SwarmSaveImageWS, SwarmTiling, SwarmExtractLora, SwarmUnsampler, SwarmLatents, SwarmInputNodes, SwarmTextHandling, SwarmReference, SwarmMath
+
+WEB_DIRECTORY = "./web"
+
+NODE_CLASS_MAPPINGS = (
+    SwarmBlending.NODE_CLASS_MAPPINGS
+    | SwarmClipSeg.NODE_CLASS_MAPPINGS
+    | SwarmImages.NODE_CLASS_MAPPINGS
+    | SwarmInternalUtil.NODE_CLASS_MAPPINGS
+    | SwarmKSampler.NODE_CLASS_MAPPINGS
+    | SwarmLoadImageB64.NODE_CLASS_MAPPINGS
+    | SwarmLoraLoader.NODE_CLASS_MAPPINGS
+    | SwarmMasks.NODE_CLASS_MAPPINGS
+    | SwarmSaveImageWS.NODE_CLASS_MAPPINGS
+    | SwarmTiling.NODE_CLASS_MAPPINGS
+    | SwarmExtractLora.NODE_CLASS_MAPPINGS
+    | SwarmUnsampler.NODE_CLASS_MAPPINGS
+    | SwarmLatents.NODE_CLASS_MAPPINGS
+    | SwarmInputNodes.NODE_CLASS_MAPPINGS
+    | SwarmTextHandling.NODE_CLASS_MAPPINGS
+    | SwarmReference.NODE_CLASS_MAPPINGS
+    | SwarmMath.NODE_CLASS_MAPPINGS
+)
+
+# TODO: Why is there no comfy core register method? 0.o
+def register_model_folder(name):
+    if name not in folder_paths.folder_names_and_paths:
+        folder_paths.folder_names_and_paths[name] = ([os.path.join(folder_paths.models_dir, name)], folder_paths.supported_pt_extensions)
+    else:
+        folder_paths.folder_names_and_paths[name] = (folder_paths.folder_names_and_paths[name][0], folder_paths.supported_pt_extensions)
+
+register_model_folder("yolov8")