180 lines
6.5 KiB
Python
180 lines
6.5 KiB
Python
import copy
|
|
import importlib.util
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from contextlib import nullcontext
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
|
|
def _bootstrap_package(package_name):
|
|
for mod_name in list(sys.modules):
|
|
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
|
del sys.modules[mod_name]
|
|
|
|
root_pkg = types.ModuleType(package_name)
|
|
root_pkg.__path__ = []
|
|
sys.modules[package_name] = root_pkg
|
|
|
|
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
|
utils_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.utils"] = utils_pkg
|
|
|
|
upscale_pkg = types.ModuleType(f"{package_name}.upscale")
|
|
upscale_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.upscale"] = upscale_pkg
|
|
|
|
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
|
logging_module.debug_log = lambda *_args, **_kwargs: None
|
|
logging_module.log = lambda *_args, **_kwargs: None
|
|
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
|
|
|
image_module = types.ModuleType(f"{package_name}.utils.image")
|
|
|
|
def tensor_to_pil(image_batch, index):
|
|
tensor = image_batch[index].detach().cpu().clamp(0, 1)
|
|
array = (tensor.numpy() * 255.0).astype("uint8")
|
|
return Image.fromarray(array)
|
|
|
|
def pil_to_tensor(image):
|
|
array = np.array(image).astype(np.float32) / 255.0
|
|
return torch.from_numpy(array).unsqueeze(0)
|
|
|
|
image_module.tensor_to_pil = tensor_to_pil
|
|
image_module.pil_to_tensor = pil_to_tensor
|
|
sys.modules[f"{package_name}.utils.image"] = image_module
|
|
|
|
usdu_utils_module = types.ModuleType(f"{package_name}.utils.usdu_utils")
|
|
|
|
def crop_cond(cond, *_args, **_kwargs):
|
|
return cond
|
|
|
|
def get_crop_region(mask, padding):
|
|
bbox = mask.getbbox()
|
|
if bbox is None:
|
|
return (0, 0, mask.size[0], mask.size[1])
|
|
x1, y1, x2, y2 = bbox
|
|
return (
|
|
max(0, x1 - int(padding)),
|
|
max(0, y1 - int(padding)),
|
|
min(mask.size[0], x2 + int(padding)),
|
|
min(mask.size[1], y2 + int(padding)),
|
|
)
|
|
|
|
def expand_crop(region, width, height, target_w, target_h):
|
|
x1, y1, x2, y2 = region
|
|
crop_w = x2 - x1
|
|
crop_h = y2 - y1
|
|
if crop_w >= target_w and crop_h >= target_h:
|
|
return (region, (crop_w, crop_h))
|
|
new_x2 = min(width, x1 + max(crop_w, target_w))
|
|
new_y2 = min(height, y1 + max(crop_h, target_h))
|
|
return ((x1, y1, new_x2, new_y2), (new_x2 - x1, new_y2 - y1))
|
|
|
|
usdu_utils_module.crop_cond = crop_cond
|
|
usdu_utils_module.get_crop_region = get_crop_region
|
|
usdu_utils_module.expand_crop = expand_crop
|
|
sys.modules[f"{package_name}.utils.usdu_utils"] = usdu_utils_module
|
|
|
|
crop_model_patch_module = types.ModuleType(f"{package_name}.utils.crop_model_patch")
|
|
crop_model_patch_module.crop_model_cond = lambda model, *_args, **_kwargs: nullcontext(model)
|
|
sys.modules[f"{package_name}.utils.crop_model_patch"] = crop_model_patch_module
|
|
|
|
conditioning_module = types.ModuleType(f"{package_name}.upscale.conditioning")
|
|
conditioning_module.clone_conditioning = lambda cond, clone_hints=False: copy.deepcopy(cond)
|
|
sys.modules[f"{package_name}.upscale.conditioning"] = conditioning_module
|
|
|
|
comfy_module = types.ModuleType("comfy")
|
|
comfy_samplers = types.ModuleType("comfy.samplers")
|
|
comfy_model_management = types.ModuleType("comfy.model_management")
|
|
comfy_module.samplers = comfy_samplers
|
|
comfy_module.model_management = comfy_model_management
|
|
sys.modules["comfy"] = comfy_module
|
|
sys.modules["comfy.samplers"] = comfy_samplers
|
|
sys.modules["comfy.model_management"] = comfy_model_management
|
|
|
|
|
|
def _load_module(package_name, module_rel_path, module_name):
|
|
module_path = Path(__file__).resolve().parents[3] / module_rel_path
|
|
spec = importlib.util.spec_from_file_location(
|
|
f"{package_name}.{module_name}",
|
|
module_path,
|
|
)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec is not None and spec.loader is not None
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def _load_tile_ops_module():
|
|
package_name = "dist_tile_ops_testpkg"
|
|
_bootstrap_package(package_name)
|
|
return _load_module(package_name, "upscale/tile_ops.py", "upscale.tile_ops")
|
|
|
|
|
|
tile_ops_module = _load_tile_ops_module()
|
|
|
|
|
|
class _DummyTileOps(tile_ops_module.TileOpsMixin):
|
|
pass
|
|
|
|
|
|
class _FakeControl:
|
|
def __init__(self, cond_hint_original, previous_controlnet=None):
|
|
self.cond_hint_original = cond_hint_original
|
|
self.previous_controlnet = previous_controlnet
|
|
|
|
|
|
class TileOpsTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.node = _DummyTileOps()
|
|
|
|
def test_round_and_calculate_tiles(self):
|
|
self.assertEqual(self.node.round_to_multiple(13, 8), 16)
|
|
tiles = self.node.calculate_tiles(10, 9, 4, 5, force_uniform_tiles=True)
|
|
self.assertEqual(len(tiles), 6)
|
|
self.assertEqual(tiles[0], (0, 0))
|
|
self.assertEqual(tiles[-1], (8, 5))
|
|
|
|
def test_create_tile_mask_and_blend_tile(self):
|
|
base = Image.new("RGB", (8, 8), (0, 0, 0))
|
|
tile = Image.new("RGB", (4, 4), (255, 255, 255))
|
|
mask = self.node.create_tile_mask(8, 8, 2, 2, 4, 4, mask_blur=0)
|
|
|
|
self.assertEqual(mask.getpixel((0, 0)), 0)
|
|
self.assertEqual(mask.getpixel((3, 3)), 255)
|
|
|
|
blended = self.node.blend_tile(
|
|
base_image=base,
|
|
tile_image=tile,
|
|
x=2,
|
|
y=2,
|
|
extracted_size=(4, 4),
|
|
mask=mask,
|
|
padding=0,
|
|
)
|
|
self.assertEqual(blended.size, (8, 8))
|
|
self.assertEqual(blended.getpixel((3, 3)), (255, 255, 255))
|
|
|
|
def test_slice_conditioning_selects_batch_index(self):
|
|
control = _FakeControl(cond_hint_original=torch.ones((2, 1, 1)))
|
|
positive = [[torch.arange(6, dtype=torch.float32).reshape(2, 3), {"control": control, "mask": torch.ones((2, 1, 1))}]]
|
|
negative = [[torch.arange(6, dtype=torch.float32).reshape(2, 3), {"mask": torch.zeros((2, 1, 1))}]]
|
|
|
|
pos_sliced, neg_sliced = self.node._slice_conditioning(positive, negative, batch_idx=1)
|
|
|
|
self.assertEqual(tuple(pos_sliced[0][0].shape), (1, 3))
|
|
self.assertEqual(tuple(neg_sliced[0][0].shape), (1, 3))
|
|
self.assertEqual(tuple(pos_sliced[0][1]["mask"].shape), (1, 1, 1))
|
|
self.assertEqual(tuple(neg_sliced[0][1]["mask"].shape), (1, 1, 1))
|
|
self.assertEqual(tuple(pos_sliced[0][1]["control"].cond_hint_original.shape), (1, 1, 1))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|