Tiled Upscale Stitch: work with a list of pictures

A list of pictures runs Tiled Upscale once per picture, so Stitch gets
every picture's tiles in one list and one stitcher per picture. It only
read the first stitcher and stopped with a tile count error (#205).
Each stitcher now takes its own tiles off the list and Stitch returns a
list, one finished picture each. One picture comes out as before.

CHANGELOG.md

@@ -4,6 +4,11 @@ All notable changes to ComfyUI-AusBoss are documented here.
 
 ## Unreleased
 
+- **Tiled Upscale Stitch 🆎 works with a list of pictures (#205).** Give
+  Tiled Upscale a list of pictures, such as a folder of images, and Stitch
+  now gives back one finished picture for each, in the same order. Before,
+  it stopped with "got 12 tiles and the stitcher has 4". No loop is needed:
+  the nodes after Stitch run once per picture.
 - **LoRA Loader 🆎: show the rows on the face of a subgraph box.** When the
   loader sits inside a subgraph, its gear has a new first item, **Rows on
   the box**. Press **Show** and the rows appear on the box itself, where you

js/docs/AUSBOSS_NODES_TiledUpscaleStitch.md

@@ -7,6 +7,13 @@ you the finished picture, at the size you asked for in
 Wire the redrawn tiles into `tiles` and the `stitcher` from Tiled
 Upscale into `stitcher`. That is all it needs.
 
+## Several pictures at once
+
+Give Tiled Upscale a list of pictures, such as a folder of images, and
+Stitch gives back a list too: one finished picture for each, in the same
+order. The nodes after it, such as Save Image, run once per picture. No
+loop is needed.
+
 ## Keep colors
 
 **Keep colors** (`keep_colors`) matches each tile to the same tile of your
@@ -30,6 +37,10 @@ photo comes back faded.
 - The tiles can arrive one after another or as one batch in the same
   order. A different number of tiles than the stitcher expects stops the
   run and says how many it got.
+- With a list of pictures, Tiled Upscale runs once per picture, so the
+  tiles arrive as one list, picture after picture, with one stitcher per
+  picture. Each stitcher takes its own number of tiles off the front of
+  the list.
 - When the picture was redrawn smaller than the size you asked for (a
   small picture, see Redraw limit on Tiled Upscale), the joined redraw
   is enlarged the plain way to that size.

nodes/node_tiled_upscale.py

@@ -331,57 +331,77 @@ class AusBossTiledUpscaleStitch:
 
     RETURN_TYPES = ("IMAGE",)
     RETURN_NAMES = ("image",)
-    OUTPUT_TOOLTIPS = ("The finished picture at the size you asked for.",)
+    OUTPUT_IS_LIST = (True,)
+    OUTPUT_TOOLTIPS = (
+        "The finished picture at the size you asked for. A list of pictures "
+        "comes out as a list, one finished picture each.",
+    )
     FUNCTION = "stitch"
 
     def stitch(self, tiles, stitcher, keep_colors=None):
-        stitcher = stitcher[0] if isinstance(stitcher, (list, tuple)) else stitcher
-        if not isinstance(stitcher, dict) or stitcher.get("kind") != STITCHER_KIND:
+        # A list of pictures runs Tiled Upscale once per picture: the tiles
+        # arrive as one list, picture after picture, with one stitcher each.
+        stitchers = list(stitcher) if isinstance(stitcher, (list, tuple)) else [stitcher]
+        if not stitchers or any(
+            not isinstance(item, dict) or item.get("kind") != STITCHER_KIND for item in stitchers
+        ):
             raise ValueError(
                 "Tiled Upscale Stitch needs the stitcher from Tiled Upscale 🆎."
             )
         strength = keep_colors[0] if isinstance(keep_colors, (list, tuple)) and keep_colors else keep_colors
         strength = 1.0 if strength is None else max(0.0, min(1.0, float(strength)))
-        plan = stitcher["plan"]
-        boxes = plan["boxes"]
-        sources = stitcher["tiles"]
+        counts = [len(item["plan"]["boxes"]) for item in stitchers]
+        expected = sum(counts)
 
         redrawn = [_rgb(tile, "Tiled Upscale Stitch") for tile in (tiles if isinstance(tiles, (list, tuple)) else [tiles])]
-        if len(redrawn) == 1 and len(boxes) > 1 and redrawn[0].shape[0] % len(boxes) == 0:
+        if len(redrawn) == 1 and expected > 1 and redrawn[0].shape[0] % expected == 0:
             # the tiles came back as one batch: split it the way it was made
-            per_tile = redrawn[0].shape[0] // len(boxes)
+            per_tile = redrawn[0].shape[0] // expected
             redrawn = list(redrawn[0].split(per_tile, dim=0))
-        if len(redrawn) != len(boxes):
+        if len(redrawn) != expected:
+            has = f"{len(stitchers)} stitchers have" if len(stitchers) > 1 else "the stitcher has"
             raise ValueError(
-                f"Tiled Upscale Stitch got {len(redrawn)} tiles and the stitcher has {len(boxes)}. "
+                f"Tiled Upscale Stitch got {len(redrawn)} tiles and {has} {expected}. "
                 "Send every tile from Tiled Upscale 🆎 through the same nodes, in order."
             )
 
-        redraw_width, redraw_height = plan["redraw"]
-        result_width, result_height = plan["result"]
-        tile_width, tile_height = plan["tile"]
-        lefts = sorted({box[0] for box in boxes})
-        tops = sorted({box[1] for box in boxes})
-        batch = max(tile.shape[0] for tile in redrawn)
-        total = torch.zeros((batch, redraw_height, redraw_width, 3), dtype=torch.float32)
-        weight = torch.zeros((redraw_height, redraw_width), dtype=torch.float32)
+        pictures = []
+        start = 0
+        for item, count in zip(stitchers, counts):
+            pictures.append(_join(redrawn[start : start + count], item, strength))
+            start += count
+        return (pictures,)
 
-        for tile, source, (left, top, right, bottom) in zip(redrawn, sources, boxes):
-            tile = resample_batch(tile.cpu(), tile_width, tile_height, "lanczos")
-            if strength > 0.0:
-                reference = source.cpu()
-                if reference.shape[0] not in (1, tile.shape[0]):
-                    reference = reference[:1]
-                tile = _keep_colors(tile, reference, strength)
-            share_left, share_right = _shared(lefts, tile_width, left)
-            share_top, share_bottom = _shared(tops, tile_height, top)
-            mask = _ramp(tile_height, share_top, share_bottom)[:, None] * _ramp(tile_width, share_left, share_right)[None, :]
-            total[:, top:bottom, left:right, :] += tile * mask[None, :, :, None]
-            weight[top:bottom, left:right] += mask
 
-        merged = total / weight.clamp_min(1e-6)[None, :, :, None]
-        merged = resample_batch(merged.clamp_(0.0, 1.0), result_width, result_height, "lanczos")
-        return (merged,)
+def _join(redrawn: list[torch.Tensor], stitcher: dict, strength: float) -> torch.Tensor:
+    """One picture from its redrawn tiles, blended at the joins."""
+    plan = stitcher["plan"]
+    boxes = plan["boxes"]
+    sources = stitcher["tiles"]
+    redraw_width, redraw_height = plan["redraw"]
+    result_width, result_height = plan["result"]
+    tile_width, tile_height = plan["tile"]
+    lefts = sorted({box[0] for box in boxes})
+    tops = sorted({box[1] for box in boxes})
+    batch = max(tile.shape[0] for tile in redrawn)
+    total = torch.zeros((batch, redraw_height, redraw_width, 3), dtype=torch.float32)
+    weight = torch.zeros((redraw_height, redraw_width), dtype=torch.float32)
+
+    for tile, source, (left, top, right, bottom) in zip(redrawn, sources, boxes):
+        tile = resample_batch(tile.cpu(), tile_width, tile_height, "lanczos")
+        if strength > 0.0:
+            reference = source.cpu()
+            if reference.shape[0] not in (1, tile.shape[0]):
+                reference = reference[:1]
+            tile = _keep_colors(tile, reference, strength)
+        share_left, share_right = _shared(lefts, tile_width, left)
+        share_top, share_bottom = _shared(tops, tile_height, top)
+        mask = _ramp(tile_height, share_top, share_bottom)[:, None] * _ramp(tile_width, share_left, share_right)[None, :]
+        total[:, top:bottom, left:right, :] += tile * mask[None, :, :, None]
+        weight[top:bottom, left:right] += mask
+
+    merged = total / weight.clamp_min(1e-6)[None, :, :, None]
+    return resample_batch(merged.clamp_(0.0, 1.0), result_width, result_height, "lanczos")
 
 
 NODE_CLASS_MAPPINGS = {

tests/test_tiled_upscale.py

@@ -32,7 +32,9 @@ def picture(batch: int, height: int, width: int, seed: int = 0, channels: int =
 
 def stitch(tiles, stitcher, keep_colors=1.0):
     # the node is INPUT_IS_LIST: every input arrives as a list
-    return AusBossTiledUpscaleStitch().stitch(list(tiles), [stitcher], [keep_colors])[0]
+    pictures = AusBossTiledUpscaleStitch().stitch(list(tiles), [stitcher], [keep_colors])[0]
+    assert len(pictures) == 1
+    return pictures[0]
 
 
 class SplitTests(unittest.TestCase):
@@ -147,6 +149,22 @@ class StitchTests(unittest.TestCase):
         together = stitch([torch.cat(tiles, dim=0)], stitcher, 0.0)
         self.assertTrue(torch.allclose(together, out))
 
+    def test_a_list_of_pictures_comes_back_as_a_list(self):
+        # ComfyUI runs Tiled Upscale once per picture in a list: the tiles
+        # arrive as one list, picture after picture, with one stitcher each.
+        sources = [picture(1, 654, 368, seed=1), picture(1, 512, 512, seed=2), picture(1, 300, 700, seed=3)]
+        runs = [AusBossTiledUpscale().split(source, 8.0, max_growth=64.0) for source in sources]
+        tiles = [tile for run in runs for tile in run[0]]
+        stitchers = [run[1] for run in runs]
+        self.assertGreater(len(tiles), len(runs))
+        pictures = AusBossTiledUpscaleStitch().stitch(tiles, stitchers, [0.0])[0]
+        self.assertEqual(len(pictures), len(runs))
+        for out, run in zip(pictures, runs):
+            self.assertEqual(tuple(out.shape), (1, run[4], run[3], 3))
+            self.assertLess(float((out - run[2]).abs().max()), 1e-5)
+        with self.assertRaisesRegex(ValueError, f"3 stitchers have {len(tiles)}"):
+            AusBossTiledUpscaleStitch().stitch(tiles[:-1], stitchers, [0.0])
+
     def test_wrong_wiring_is_a_clear_error(self):
         tiles, stitcher, *_ = AusBossTiledUpscale().split(picture(1, 654, 368), 8.0, max_growth=64.0)
         with self.assertRaisesRegex(ValueError, "got 3 tiles and the stitcher has 4"):