commit 2564377aa161a8d62860b08b56645c2e0c5d3842 Author: Ethanfel Date: Wed Feb 25 15:52:32 2026 +0100 feat: initial ComfyUI_UltimateSGUpscale Recreates UltimateSDUpscale features using built-in ComfyUI nodes (SplitImageToTileList, ImageMergeTileList) plus a small GenerateSeamMask helper node. - GenerateSeamMask: creates white-band mask at tile seam positions - Workflow JSON: 22-node tiled upscale pipeline with: - Pass 1: model upscale + tiled img2img redraw - Pass 2: targeted seam fix using SetLatentNoiseMask - Unit tests for seam mask generation Co-Authored-By: Claude Opus 4.6 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7a60b85 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +__pycache__/ +*.pyc diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..03fbe96 --- /dev/null +++ b/__init__.py @@ -0,0 +1,11 @@ +from .seam_mask_node import GenerateSeamMask + +NODE_CLASS_MAPPINGS = { + "GenerateSeamMask": GenerateSeamMask, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "GenerateSeamMask": "Generate Seam Mask", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/example_workflows/tiled-upscale-builtin-nodes.json b/example_workflows/tiled-upscale-builtin-nodes.json new file mode 100644 index 0000000..c4640ec --- /dev/null +++ b/example_workflows/tiled-upscale-builtin-nodes.json @@ -0,0 +1,494 @@ +{ + "last_node_id": 22, + "last_link_id": 36, + "nodes": [ + { + "id": 1, + "type": "CheckpointLoaderSimple", + "pos": [50, 200], + "size": [315, 98], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + {"name": "MODEL", "type": "MODEL", "slot_index": 0, "links": [1, 2]}, + {"name": "CLIP", "type": "CLIP", "slot_index": 1, "links": [3, 4]}, + {"name": "VAE", "type": "VAE", "slot_index": 2, "links": [5, 6, 7, 8]} + ], + "properties": {"Node name for S&R": "CheckpointLoaderSimple"}, + "widgets_values": ["v1-5-pruned-emaonly-fp16.safetensors"] + }, + { + "id": 2, + "type": "CLIPTextEncode", + "pos": [50, 400], + "size": [400, 150], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + {"name": "clip", "type": "CLIP", "link": 3} + ], + "outputs": [ + {"name": "CONDITIONING", "type": "CONDITIONING", "slot_index": 0, "links": [9, 10]} + ], + "properties": {"Node name for S&R": "CLIPTextEncode"}, + "widgets_values": ["high quality, detailed, sharp"], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 3, + "type": "CLIPTextEncode", + "pos": [50, 600], + "size": [400, 150], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + {"name": "clip", "type": "CLIP", "link": 4} + ], + "outputs": [ + {"name": "CONDITIONING", "type": "CONDITIONING", "slot_index": 0, "links": [11, 12]} + ], + "properties": {"Node name for S&R": "CLIPTextEncode"}, + "widgets_values": ["blurry, low quality, artifacts, watermark"], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 4, + "type": "LoadImage", + "pos": [500, 50], + "size": [315, 314], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + {"name": "IMAGE", "type": "IMAGE", "slot_index": 0, "links": [13]}, + {"name": "MASK", "type": "MASK", "slot_index": 1, "links": []} + ], + "properties": {"Node name for S&R": "LoadImage"}, + "widgets_values": ["example.png", "image"] + }, + { + "id": 5, + "type": "UpscaleModelLoader", + "pos": [500, 400], + "size": [315, 58], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + {"name": "UPSCALE_MODEL", "type": "UPSCALE_MODEL", "slot_index": 0, "links": [14]} + ], + "properties": {"Node name for S&R": "UpscaleModelLoader"}, + "widgets_values": ["4x-UltraSharp.pth"] + }, + { + "id": 6, + "type": "ImageUpscaleWithModel", + "pos": [870, 200], + "size": [241, 46], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + {"name": "upscale_model", "type": "UPSCALE_MODEL", "link": 14}, + {"name": "image", "type": "IMAGE", "link": 13} + ], + "outputs": [ + {"name": "IMAGE", "type": "IMAGE", "slot_index": 0, "links": [15, 16]} + ], + "properties": {"Node name for S&R": "ImageUpscaleWithModel"}, + "widgets_values": [] + }, + { + "id": 7, + "type": "GetImageSize", + "pos": [870, 300], + "size": [200, 66], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + {"name": "image", "type": "IMAGE", "link": 15} + ], + "outputs": [ + {"name": "width", "type": "INT", "slot_index": 0, "links": [17, 18, 35]}, + {"name": "height", "type": "INT", "slot_index": 1, "links": [19, 20, 36]}, + {"name": "batch_size", "type": "INT", "slot_index": 2, "links": []} + ], + "properties": {"Node name for S&R": "GetImageSize"}, + "widgets_values": [] + }, + { + "id": 8, + "type": "SplitImageToTileList", + "pos": [1170, 200], + "size": [250, 106], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + {"name": "image", "type": "IMAGE", "link": 16} + ], + "outputs": [ + {"name": "image", "type": "IMAGE", "slot_index": 0, "links": [21]} + ], + "properties": {"Node name for S&R": "SplitImageToTileList"}, + "widgets_values": [1024, 1024, 128] + }, + { + "id": 9, + "type": "VAEEncode", + "pos": [1470, 200], + "size": [170, 46], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + {"name": "pixels", "type": "IMAGE", "link": 21}, + {"name": "vae", "type": "VAE", "link": 5} + ], + "outputs": [ + {"name": "LATENT", "type": "LATENT", "slot_index": 0, "links": [22]} + ], + "properties": {"Node name for S&R": "VAEEncode"}, + "widgets_values": [] + }, + { + "id": 10, + "type": "KSampler", + "pos": [1470, 300], + "size": [300, 474], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + {"name": "model", "type": "MODEL", "link": 1}, + {"name": "positive", "type": "CONDITIONING", "link": 9}, + {"name": "negative", "type": "CONDITIONING", "link": 11}, + {"name": "latent_image", "type": "LATENT", "link": 22} + ], + "outputs": [ + {"name": "LATENT", "type": "LATENT", "slot_index": 0, "links": [23]} + ], + "properties": {"Node name for S&R": "KSampler"}, + "widgets_values": [0, "fixed", 20, 7.0, "euler", "normal", 0.35] + }, + { + "id": 11, + "type": "VAEDecode", + "pos": [1820, 200], + "size": [170, 46], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + {"name": "samples", "type": "LATENT", "link": 23}, + {"name": "vae", "type": "VAE", "link": 6} + ], + "outputs": [ + {"name": "IMAGE", "type": "IMAGE", "slot_index": 0, "links": [24]} + ], + "properties": {"Node name for S&R": "VAEDecode"}, + "widgets_values": [] + }, + { + "id": 12, + "type": "ImageMergeTileList", + "pos": [2040, 200], + "size": [250, 106], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + {"name": "image_list", "type": "IMAGE", "link": 24}, + {"name": "final_width", "type": "INT", "link": 17, "widget": {"name": "final_width"}}, + {"name": "final_height", "type": "INT", "link": 19, "widget": {"name": "final_height"}} + ], + "outputs": [ + {"name": "IMAGE", "type": "IMAGE", "slot_index": 0, "links": [25]} + ], + "properties": {"Node name for S&R": "ImageMergeTileList"}, + "widgets_values": [2048, 2048, 128] + }, + { + "id": 13, + "type": "GenerateSeamMask", + "pos": [2040, 500], + "size": [250, 170], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + {"name": "image_width", "type": "INT", "link": 18, "widget": {"name": "image_width"}}, + {"name": "image_height", "type": "INT", "link": 20, "widget": {"name": "image_height"}} + ], + "outputs": [ + {"name": "IMAGE", "type": "IMAGE", "slot_index": 0, "links": [27]} + ], + "properties": {"Node name for S&R": "GenerateSeamMask"}, + "widgets_values": [2048, 2048, 1024, 1024, 128, 64] + }, + { + "id": 14, + "type": "SplitImageToTileList", + "pos": [2370, 200], + "size": [250, 106], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + {"name": "image", "type": "IMAGE", "link": 25} + ], + "outputs": [ + {"name": "image", "type": "IMAGE", "slot_index": 0, "links": [28]} + ], + "properties": {"Node name for S&R": "SplitImageToTileList"}, + "widgets_values": [768, 768, 128] + }, + { + "id": 15, + "type": "SplitImageToTileList", + "pos": [2370, 500], + "size": [250, 106], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + {"name": "image", "type": "IMAGE", "link": 27} + ], + "outputs": [ + {"name": "image", "type": "IMAGE", "slot_index": 0, "links": [29]} + ], + "properties": {"Node name for S&R": "SplitImageToTileList"}, + "widgets_values": [768, 768, 128] + }, + { + "id": 16, + "type": "ImageToMask", + "pos": [2670, 500], + "size": [200, 58], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + {"name": "image", "type": "IMAGE", "link": 29} + ], + "outputs": [ + {"name": "MASK", "type": "MASK", "slot_index": 0, "links": [30]} + ], + "properties": {"Node name for S&R": "ImageToMask"}, + "widgets_values": ["red"] + }, + { + "id": 17, + "type": "VAEEncode", + "pos": [2670, 200], + "size": [170, 46], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + {"name": "pixels", "type": "IMAGE", "link": 28}, + {"name": "vae", "type": "VAE", "link": 7} + ], + "outputs": [ + {"name": "LATENT", "type": "LATENT", "slot_index": 0, "links": [31]} + ], + "properties": {"Node name for S&R": "VAEEncode"}, + "widgets_values": [] + }, + { + "id": 18, + "type": "SetLatentNoiseMask", + "pos": [2670, 350], + "size": [250, 46], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + {"name": "samples", "type": "LATENT", "link": 31}, + {"name": "mask", "type": "MASK", "link": 30} + ], + "outputs": [ + {"name": "LATENT", "type": "LATENT", "slot_index": 0, "links": [32]} + ], + "properties": {"Node name for S&R": "SetLatentNoiseMask"}, + "widgets_values": [] + }, + { + "id": 19, + "type": "KSampler", + "pos": [2970, 200], + "size": [300, 474], + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + {"name": "model", "type": "MODEL", "link": 2}, + {"name": "positive", "type": "CONDITIONING", "link": 10}, + {"name": "negative", "type": "CONDITIONING", "link": 12}, + {"name": "latent_image", "type": "LATENT", "link": 32} + ], + "outputs": [ + {"name": "LATENT", "type": "LATENT", "slot_index": 0, "links": [33]} + ], + "properties": {"Node name for S&R": "KSampler"}, + "widgets_values": [0, "fixed", 20, 7.0, "euler", "normal", 0.35] + }, + { + "id": 20, + "type": "VAEDecode", + "pos": [3320, 200], + "size": [170, 46], + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + {"name": "samples", "type": "LATENT", "link": 33}, + {"name": "vae", "type": "VAE", "link": 8} + ], + "outputs": [ + {"name": "IMAGE", "type": "IMAGE", "slot_index": 0, "links": [34]} + ], + "properties": {"Node name for S&R": "VAEDecode"}, + "widgets_values": [] + }, + { + "id": 21, + "type": "ImageMergeTileList", + "pos": [3540, 200], + "size": [250, 106], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + {"name": "image_list", "type": "IMAGE", "link": 34}, + {"name": "final_width", "type": "INT", "link": 35, "widget": {"name": "final_width"}}, + {"name": "final_height", "type": "INT", "link": 36, "widget": {"name": "final_height"}} + ], + "outputs": [ + {"name": "IMAGE", "type": "IMAGE", "slot_index": 0, "links": [26]} + ], + "properties": {"Node name for S&R": "ImageMergeTileList"}, + "widgets_values": [2048, 2048, 128] + }, + { + "id": 22, + "type": "SaveImage", + "pos": [3840, 200], + "size": [400, 400], + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + {"name": "images", "type": "IMAGE", "link": 26} + ], + "outputs": [], + "properties": {"Node name for S&R": "SaveImage"}, + "widgets_values": ["UltimateSG/upscale"] + } + ], + "links": [ + [1, 1, 0, 10, 0, "MODEL"], + [2, 1, 0, 19, 0, "MODEL"], + [3, 1, 1, 2, 0, "CLIP"], + [4, 1, 1, 3, 0, "CLIP"], + [5, 1, 2, 9, 1, "VAE"], + [6, 1, 2, 11, 1, "VAE"], + [7, 1, 2, 17, 1, "VAE"], + [8, 1, 2, 20, 1, "VAE"], + [9, 2, 0, 10, 1, "CONDITIONING"], + [10, 2, 0, 19, 1, "CONDITIONING"], + [11, 3, 0, 10, 2, "CONDITIONING"], + [12, 3, 0, 19, 2, "CONDITIONING"], + [13, 4, 0, 6, 1, "IMAGE"], + [14, 5, 0, 6, 0, "UPSCALE_MODEL"], + [15, 6, 0, 7, 0, "IMAGE"], + [16, 6, 0, 8, 0, "IMAGE"], + [17, 7, 0, 12, 1, "INT"], + [18, 7, 0, 13, 0, "INT"], + [19, 7, 1, 12, 2, "INT"], + [20, 7, 1, 13, 1, "INT"], + [21, 8, 0, 9, 0, "IMAGE"], + [22, 9, 0, 10, 3, "LATENT"], + [23, 10, 0, 11, 0, "LATENT"], + [24, 11, 0, 12, 0, "IMAGE"], + [25, 12, 0, 14, 0, "IMAGE"], + [26, 21, 0, 22, 0, "IMAGE"], + [27, 13, 0, 15, 0, "IMAGE"], + [28, 14, 0, 17, 0, "IMAGE"], + [29, 15, 0, 16, 0, "IMAGE"], + [30, 16, 0, 18, 1, "MASK"], + [31, 17, 0, 18, 0, "LATENT"], + [32, 18, 0, 19, 3, "LATENT"], + [33, 19, 0, 20, 0, "LATENT"], + [34, 20, 0, 21, 0, "IMAGE"], + [35, 7, 0, 21, 1, "INT"], + [36, 7, 1, 21, 2, "INT"] + ], + "groups": [ + { + "id": 1, + "title": "Load Models", + "bounding": [40, 130, 345, 400], + "color": "#3f789e", + "font_size": 24, + "flags": {} + }, + { + "id": 2, + "title": "Prompts", + "bounding": [40, 330, 420, 440], + "color": "#3f789e", + "font_size": 24, + "flags": {} + }, + { + "id": 3, + "title": "Upscale", + "bounding": [490, 30, 500, 400], + "color": "#3f789e", + "font_size": 24, + "flags": {} + }, + { + "id": 4, + "title": "Pass 1: Tiled Redraw", + "bounding": [1150, 130, 970, 730], + "color": "#3a7e3a", + "font_size": 24, + "flags": {} + }, + { + "id": 5, + "title": "Pass 2: Seam Fix (bypass/mute to skip)", + "bounding": [2030, 130, 1800, 600], + "color": "#7e3a3a", + "font_size": 24, + "flags": {} + }, + { + "id": 6, + "title": "Output", + "bounding": [3830, 130, 420, 500], + "color": "#3f789e", + "font_size": 24, + "flags": {} + } + ], + "config": {}, + "extra": { + "ds": { + "scale": 0.4, + "offset": [50, 50] + } + }, + "version": 0.4 +} diff --git a/seam_mask_node.py b/seam_mask_node.py new file mode 100644 index 0000000..43423f5 --- /dev/null +++ b/seam_mask_node.py @@ -0,0 +1,52 @@ +import torch + + +class GenerateSeamMask: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_width": ("INT", {"default": 2048, "min": 64, "max": 16384, "step": 1, + "tooltip": "Width of the image (from GetImageSize)."}), + "image_height": ("INT", {"default": 2048, "min": 64, "max": 16384, "step": 1, + "tooltip": "Height of the image (from GetImageSize)."}), + "tile_width": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8, + "tooltip": "Tile width used in the main tiled redraw pass."}), + "tile_height": ("INT", {"default": 1024, "min": 64, "max": 8192, "step": 8, + "tooltip": "Tile height used in the main tiled redraw pass."}), + "overlap": ("INT", {"default": 128, "min": 0, "max": 4096, "step": 1, + "tooltip": "Overlap used in the main tiled redraw pass."}), + "seam_width": ("INT", {"default": 64, "min": 8, "max": 512, "step": 8, + "tooltip": "Width of the seam bands to fix (in pixels)."}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "generate" + CATEGORY = "image/upscaling" + DESCRIPTION = "Generates a mask image with white bands at tile seam positions. Used for targeted seam fix denoising." + + def generate(self, image_width, image_height, tile_width, tile_height, overlap, seam_width): + mask = torch.zeros(1, image_height, image_width, 3) + + stride_x = max(1, tile_width - overlap) + stride_y = max(1, tile_height - overlap) + half_w = seam_width // 2 + + # Vertical seam bands + x = stride_x + while x < image_width: + x_start = max(0, x - half_w) + x_end = min(image_width, x + half_w) + mask[:, :, x_start:x_end, :] = 1.0 + x += stride_x + + # Horizontal seam bands + y = stride_y + while y < image_height: + y_start = max(0, y - half_w) + y_end = min(image_height, y + half_w) + mask[:, y_start:y_end, :, :] = 1.0 + y += stride_y + + return (mask,) diff --git a/tests/test_seam_mask.py b/tests/test_seam_mask.py new file mode 100644 index 0000000..51585de --- /dev/null +++ b/tests/test_seam_mask.py @@ -0,0 +1,67 @@ +import sys +import os + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from seam_mask_node import GenerateSeamMask + + +def test_output_shape(): + node = GenerateSeamMask() + result = node.generate(image_width=2048, image_height=2048, + tile_width=1024, tile_height=1024, + overlap=128, seam_width=64) + mask = result[0] + assert mask.shape == (1, 2048, 2048, 3), f"Expected (1, 2048, 2048, 3), got {mask.shape}" + + +def test_seam_positions(): + node = GenerateSeamMask() + result = node.generate(image_width=2048, image_height=2048, + tile_width=1024, tile_height=1024, + overlap=128, seam_width=64) + mask = result[0] + # Stride = 1024 - 128 = 896 + # Seams at x=896, 1792 and y=896, 1792 + assert mask[0, 0, 896, 0].item() == 1.0, "Center of vertical seam should be white" + assert mask[0, 896, 0, 0].item() == 1.0, "Center of horizontal seam should be white" + assert mask[0, 0, 400, 0].item() == 0.0, "Far from any seam should be black" + + +def test_no_seams_single_tile(): + """If image fits in one tile, no seams should exist.""" + node = GenerateSeamMask() + result = node.generate(image_width=512, image_height=512, + tile_width=1024, tile_height=1024, + overlap=128, seam_width=64) + mask = result[0] + assert mask.sum().item() == 0.0, "Single tile image should have no seams" + + +def test_seam_band_width(): + node = GenerateSeamMask() + result = node.generate(image_width=2048, image_height=1024, + tile_width=1024, tile_height=1024, + overlap=0, seam_width=64) + mask = result[0] + # Stride = 1024, seam at x=1024, band from 992 to 1056 + assert mask[0, 0, 1023, 0].item() == 1.0, "Inside band should be white" + assert mask[0, 0, 991, 0].item() == 0.0, "Outside band should be black" + + +def test_values_are_binary(): + node = GenerateSeamMask() + result = node.generate(image_width=2048, image_height=2048, + tile_width=1024, tile_height=1024, + overlap=128, seam_width=64) + mask = result[0] + unique = mask.unique() + assert len(unique) <= 2, f"Mask should only contain 0.0 and 1.0, got {unique}" + + +if __name__ == "__main__": + test_output_shape() + test_seam_positions() + test_no_seams_single_tile() + test_seam_band_width() + test_values_are_binary() + print("All tests passed!")