ausboss
5d720be96f
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.
218 lines
11 KiB
Python
218 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
import sys
|
|
import types
|
|
import unittest
|
|
|
|
import torch
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
sys.path.insert(0, str(ROOT))
|
|
if "nodes" in sys.modules and not hasattr(sys.modules["nodes"], "__path__"):
|
|
del sys.modules["nodes"]
|
|
|
|
from nodes import node_tiled_upscale as module
|
|
from nodes._resize_helpers import resample_batch
|
|
from nodes.node_tiled_upscale import AusBossTiledUpscaleStitch, AusBossTiledUpscale
|
|
|
|
|
|
def picture(batch: int, height: int, width: int, seed: int = 0, channels: int = 3) -> torch.Tensor:
|
|
"""A smooth picture with a little texture: sharp noise would not survive a resample."""
|
|
generator = torch.Generator().manual_seed(seed)
|
|
ys = torch.linspace(0, 1, height)[None, :, None, None]
|
|
xs = torch.linspace(0, 1, width)[None, None, :, None]
|
|
base = torch.cat([ys.expand(1, height, width, 1), xs.expand(1, height, width, 1), (ys * xs).expand(1, height, width, 1)], dim=-1)
|
|
noise = torch.rand((batch, height, width, 3), generator=generator) * 0.05
|
|
out = (base * 0.9 + noise).clamp(0, 1)
|
|
if channels == 4:
|
|
out = torch.cat([out, torch.ones((batch, height, width, 1))], dim=-1)
|
|
return out
|
|
|
|
|
|
def stitch(tiles, stitcher, keep_colors=1.0):
|
|
# the node is INPUT_IS_LIST: every input arrives as a list
|
|
pictures = AusBossTiledUpscaleStitch().stitch(list(tiles), [stitcher], [keep_colors])[0]
|
|
assert len(pictures) == 1
|
|
return pictures[0]
|
|
|
|
|
|
class SplitTests(unittest.TestCase):
|
|
def test_outputs(self):
|
|
tiles, stitcher, before, width, height, report = AusBossTiledUpscale().split(picture(1, 327, 184), 2.0)
|
|
self.assertEqual((width, height), (1056, 1888))
|
|
self.assertEqual(len(tiles), 1)
|
|
self.assertEqual(tuple(tiles[0].shape), (1, 1888, 1056, 3))
|
|
self.assertEqual(tuple(before.shape), (1, 1888, 1056, 3))
|
|
self.assertEqual(stitcher["kind"], module.STITCHER_KIND)
|
|
self.assertIn("one tile", report)
|
|
|
|
def test_tiles_are_cut_from_the_enlarged_picture(self):
|
|
source = picture(1, 654, 368)
|
|
tiles, stitcher, before, width, height, _ = AusBossTiledUpscale().split(source, 8.0, max_growth=64.0)
|
|
plan = stitcher["plan"]
|
|
self.assertEqual(len(tiles), 4)
|
|
base = resample_batch(source, *plan["redraw"], "lanczos")
|
|
for tile, (left, top, right, bottom) in zip(tiles, plan["boxes"]):
|
|
self.assertTrue(torch.equal(tile, base[:, top:bottom, left:right, :]))
|
|
self.assertEqual((before.shape[2], before.shape[1]), (width, height))
|
|
|
|
def test_alpha_and_single_channel_pictures(self):
|
|
tiles, *_ = AusBossTiledUpscale().split(picture(1, 96, 64, channels=4), 0.5)
|
|
self.assertEqual(tiles[0].shape[-1], 3)
|
|
gray = picture(1, 96, 64)[..., :1]
|
|
tiles, *_ = AusBossTiledUpscale().split(gray, 0.5)
|
|
self.assertEqual(tiles[0].shape[-1], 3)
|
|
|
|
def test_rejects_what_is_not_a_picture(self):
|
|
with self.assertRaisesRegex(ValueError, "BHWC"):
|
|
AusBossTiledUpscale().split(torch.zeros((8, 8, 3)), 2.0)
|
|
|
|
|
|
class StitchTests(unittest.TestCase):
|
|
def test_untouched_tiles_give_the_enlarged_picture_back(self):
|
|
source = picture(1, 654, 368, seed=3)
|
|
for asked, count in ((2.0, 1), (4.0, 2), (8.0, 4)):
|
|
with self.subTest(megapixels=asked):
|
|
tiles, stitcher, before, width, height, _ = AusBossTiledUpscale().split(source, asked, max_growth=64.0)
|
|
self.assertEqual(len(tiles), count)
|
|
out = stitch(tiles, stitcher, 0.0)
|
|
self.assertEqual(tuple(out.shape), (1, height, width, 3))
|
|
# blended joins of identical content: the picture itself, to float precision
|
|
self.assertLess(float((out - before).abs().max()), 1e-5)
|
|
matched = stitch(tiles, stitcher, 1.0)
|
|
self.assertLess(float((matched - before).abs().max()), 2e-3)
|
|
|
|
def test_a_small_picture_ends_at_the_size_asked(self):
|
|
tiles, stitcher, before, width, height, report = AusBossTiledUpscale().split(picture(1, 256, 147), 8.0)
|
|
self.assertEqual(len(tiles), 1)
|
|
self.assertIn("enlarged", report)
|
|
out = stitch(tiles, stitcher)
|
|
self.assertEqual(tuple(out.shape), (1, height, width, 3))
|
|
self.assertGreater(width * height, tiles[0].shape[1] * tiles[0].shape[2] * 3)
|
|
|
|
def test_tiles_of_another_size_or_with_alpha_are_fitted(self):
|
|
tiles, stitcher, before, width, height, _ = AusBossTiledUpscale().split(picture(1, 654, 368, seed=5), 4.0, max_growth=64.0)
|
|
changed = []
|
|
for tile in tiles:
|
|
half = resample_batch(tile, tile.shape[2] // 2, tile.shape[1] // 2, "bilinear")
|
|
changed.append(torch.cat([half, torch.ones_like(half[..., :1])], dim=-1))
|
|
out = stitch(changed, stitcher, 0.0)
|
|
self.assertEqual(tuple(out.shape), (1, height, width, 3))
|
|
self.assertLess(float((out - before).abs().mean()), 0.02)
|
|
|
|
def test_keep_colors_pulls_a_tinted_tile_back(self):
|
|
tiles, stitcher, before, *_ = AusBossTiledUpscale().split(picture(1, 654, 368, seed=7), 4.0, max_growth=64.0)
|
|
self.assertEqual(len(tiles), 2)
|
|
tinted = [(tile * 0.8 + 0.1).clamp(0, 1) if index == 0 else tile for index, tile in enumerate(tiles)]
|
|
loose = stitch(tinted, stitcher, 0.0)
|
|
kept = stitch(tinted, stitcher, 1.0)
|
|
self.assertGreater(float((loose - before).abs().mean()), 3 * float((kept - before).abs().mean()))
|
|
half = stitch(tinted, stitcher, 0.5)
|
|
self.assertLess(float((kept - before).abs().mean()), float((half - before).abs().mean()))
|
|
|
|
def test_keep_colors_is_not_fooled_by_a_noisy_source(self):
|
|
# A pale picture with coarse color noise, the way an old phone photo looks. The redraw is the clean picture.
|
|
clean = picture(1, 480, 320, seed=11) * 0.35 + 0.4
|
|
generator = torch.Generator().manual_seed(12)
|
|
blotches = torch.nn.functional.interpolate(torch.randn((1, 3, 60, 40), generator=generator), size=(480, 320), mode="bilinear")
|
|
noisy = (clean + blotches.movedim(1, -1) * 0.06).clamp(0, 1)
|
|
kept = module._keep_colors(clean, noisy, 1.0)
|
|
spread = lambda image: float(module._averaged_lab(image)[..., 1:].std(dim=(1, 2)).norm())
|
|
self.assertLess(abs(spread(kept) / spread(clean) - 1.0), 0.05)
|
|
# the plain spread match takes the noise for color: this is what the fit is for
|
|
from nodes._color_helpers import match_colors
|
|
self.assertGreater(spread(match_colors(clean, noisy, 1.0, method="lab")) / spread(clean), 1.15)
|
|
|
|
def test_keep_colors_brings_a_faded_source_back_faded(self):
|
|
fresh = picture(1, 480, 320, seed=13)
|
|
faded = fresh * 0.5 + 0.3
|
|
kept = module._keep_colors(fresh, faded, 1.0)
|
|
# a fade that is a straight line in RGB is a slight curve in LAB: close, not exact
|
|
self.assertLess(float((kept - faded).abs().mean()), 0.02)
|
|
self.assertGreater(float((fresh - faded).abs().mean()), 5 * float((kept - faded).abs().mean()))
|
|
|
|
def test_keep_colors_leaves_a_gray_picture_gray(self):
|
|
gray = picture(1, 240, 160, seed=15).mean(dim=-1, keepdim=True).expand(-1, -1, -1, 3)
|
|
tinted = (gray * torch.tensor([1.03, 1.0, 0.97])).clamp(0, 1)
|
|
kept = module._keep_colors(tinted, gray, 1.0)
|
|
chroma = lambda image: float(module._averaged_lab(image)[..., 1:].abs().mean())
|
|
self.assertLess(chroma(kept), 0.3 * chroma(tinted))
|
|
|
|
def test_a_batch_of_pictures_and_one_batched_tensor(self):
|
|
source = picture(2, 654, 368, seed=9)
|
|
tiles, stitcher, before, width, height, _ = AusBossTiledUpscale().split(source, 4.0, max_growth=64.0)
|
|
self.assertEqual((len(tiles), tiles[0].shape[0]), (2, 2))
|
|
out = stitch(tiles, stitcher, 0.0)
|
|
self.assertEqual(tuple(out.shape), (2, height, width, 3))
|
|
self.assertLess(float((out - before).abs().max()), 1e-5)
|
|
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"):
|
|
stitch(tiles[:3], stitcher)
|
|
with self.assertRaisesRegex(ValueError, "stitcher from Tiled Upscale"):
|
|
stitch(tiles, {"kind": "something else"})
|
|
|
|
|
|
class UpscaleModelTests(unittest.TestCase):
|
|
"""The upscale model runs through ComfyUI's own node; here a stand-in doubles the picture."""
|
|
|
|
def setUp(self):
|
|
self.saved = sys.modules.get("nodes")
|
|
calls = self.calls = []
|
|
|
|
class Doubler:
|
|
@classmethod
|
|
def execute(cls, upscale_model, image):
|
|
calls.append(tuple(image.shape))
|
|
big = resample_batch(image, image.shape[2] * 2, image.shape[1] * 2, "bilinear")
|
|
return types.SimpleNamespace(result=(big,))
|
|
|
|
self.core = {"ImageUpscaleWithModel": Doubler}
|
|
self.patch = module._core_nodes
|
|
module._core_nodes = lambda: self.core
|
|
|
|
def tearDown(self):
|
|
module._core_nodes = self.patch
|
|
|
|
def test_model_runs_until_the_picture_covers_the_redraw(self):
|
|
tiles, stitcher, before, width, height, _ = AusBossTiledUpscale().split(picture(1, 654, 368), 8.0, upscale_model=object(), max_growth=64.0)
|
|
self.assertEqual(len(self.calls), 3) # 368 wide, doubled three times, covers 2112
|
|
self.assertEqual(len(tiles), 4)
|
|
self.calls.clear()
|
|
AusBossTiledUpscale().split(picture(1, 654, 368), 8.0, upscale_model=object())
|
|
self.assertEqual(len(self.calls), 2) # a small picture is redrawn at 2 MP: two passes cover 1056
|
|
self.assertEqual((before.shape[2], before.shape[1]), (width, height))
|
|
|
|
def test_model_is_left_alone_when_nothing_grows(self):
|
|
AusBossTiledUpscale().split(picture(1, 654, 368), 0, upscale_model=object())
|
|
self.assertEqual(self.calls, [])
|
|
|
|
def test_missing_core_node_says_what_to_do(self):
|
|
self.core.clear()
|
|
with self.assertRaisesRegex(RuntimeError, "Upscale Image"):
|
|
AusBossTiledUpscale().split(picture(1, 654, 368), 8.0, upscale_model=object())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|