From 0a4e94ee7d022f45efdf5ceac00127d9c01cdfc7 Mon Sep 17 00:00:00 2001 From: larsupb Date: Sun, 7 Dec 2025 21:08:10 +0100 Subject: [PATCH 01/11] Add UltimateSDUpscaleTiler node, add batch processing for UltimateSDUpscaleNoUpscale --- modules/processing.py | 2 + nodes.py | 154 ++++++++++++- usdu_patch.py | 518 ++++++++++++++++++++++++++++++++++++++---- 3 files changed, 625 insertions(+), 49 deletions(-) diff --git a/modules/processing.py b/modules/processing.py index c2a652f..ba53306 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -50,6 +50,7 @@ class StableDiffusionProcessing: seam_fix_mode, custom_sampler=None, custom_sigmas=None, + batch_size=1, ): # Variables used by the USDU script self.init_images = [init_img] @@ -85,6 +86,7 @@ class StableDiffusionProcessing: self.upscale_by = upscale_by self.uniform_tile_mode = uniform_tile_mode self.tiled_decode = tiled_decode + self.batch_size = batch_size self.vae_decoder = VAEDecode() self.vae_encoder = VAEEncode() self.vae_decoder_tiled = VAEDecodeTiled() diff --git a/nodes.py b/nodes.py index 5753474..accf4d0 100644 --- a/nodes.py +++ b/nodes.py @@ -4,7 +4,7 @@ import logging import torch import comfy from usdu_patch import usdu -from utils import tensor_to_pil, pil_to_tensor +from utils import tensor_to_pil, pil_to_tensor, pad_image2 from modules.processing import StableDiffusionProcessing import modules.shared as shared from modules.upscaler import UpscalerData @@ -133,13 +133,18 @@ class UltimateSDUpscale: shared.batch = [tensor_to_pil(image, i) for i in range(len(image))] shared.batch_as_tensor = image + # Get batch_size from instance if available (for UltimateSDUpscaleNoUpscale) + batch_size = getattr(self, 'batch_size', 1) + print(f"[USDU Batch Debug] UltimateSDUpscale.upscale() using batch_size={batch_size}") + # Processing sdprocessing = StableDiffusionProcessing( shared.batch[0], model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, upscale_by, force_uniform_tiles, tiled_decode, tile_width, tile_height, MODES[self.mode_type], SEAM_FIX_MODES[self.seam_fix_mode], - custom_sampler, custom_sigmas, + custom_sampler, custom_sigmas, batch_size, ) + print(f"[USDU Batch Debug] StableDiffusionProcessing created with batch_size={sdprocessing.batch_size}") # Disable logging logger = logging.getLogger() @@ -173,6 +178,7 @@ class UltimateSDUpscaleNoUpscale(UltimateSDUpscale): remove_input(required, "upscale_model") remove_input(required, "upscale_by") rename_input(required, "image", "upscaled_image") + required.append(("batch_size", ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}))) return prepare_inputs(required, optional) RETURN_TYPES = ("IMAGE",) @@ -183,8 +189,13 @@ class UltimateSDUpscaleNoUpscale(UltimateSDUpscale): steps, cfg, sampler_name, scheduler, denoise, mode_type, tile_width, tile_height, mask_blur, tile_padding, seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur, - seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode): + seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, batch_size): upscale_by = 1.0 + + # Store batch_size for use in processing + self.batch_size = batch_size + print(f"[USDU Batch Debug] UltimateSDUpscaleNoUpscale.upscale() received batch_size={batch_size}") + return super().upscale(upscaled_image, model, positive, negative, vae, upscale_by, seed, steps, cfg, sampler_name, scheduler, denoise, None, mode_type, tile_width, tile_height, mask_blur, tile_padding, @@ -200,7 +211,7 @@ class UltimateSDUpscaleCustomSample(UltimateSDUpscale): optional.append(("custom_sampler", ("SAMPLER",))) optional.append(("custom_sigmas", ("SIGMAS",))) return prepare_inputs(required, optional) - + RETURN_TYPES = ("IMAGE",) FUNCTION = "upscale" CATEGORY = "image/upscaling" @@ -220,17 +231,148 @@ class UltimateSDUpscaleCustomSample(UltimateSDUpscale): custom_sampler, custom_sigmas) +class UltimateSDUpscaleTiler: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "tile_width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}), + "tile_height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}), + "tile_padding": ("INT", {"default": 32, "min": 0, "max": MAX_RESOLUTION, "step": 8}), + "mode_type": (list(MODES.keys()),), + "force_uniform_tiles": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("IMAGE", "INT", "INT", "INT") + RETURN_NAMES = ("tiles", "rows", "cols", "tile_count") + FUNCTION = "tile_image" + CATEGORY = "image/upscaling" + + def calc_rectangle(self, xi, yi, tile_width, tile_height): + """Calculate tile rectangle coordinates""" + x1 = xi * tile_width + y1 = yi * tile_height + x2 = xi * tile_width + tile_width + y2 = yi * tile_height + tile_height + return x1, y1, x2, y2 + + def tile_image(self, image, tile_width, tile_height, tile_padding, mode_type, force_uniform_tiles): + from PIL import Image + import math + + # Get the image dimensions (batch, height, width, channels) + batch_size = len(image) + img_height = image.shape[1] + img_width = image.shape[2] + + # Calculate grid dimensions + rows = math.ceil(img_height / tile_height) + cols = math.ceil(img_width / tile_width) + + mode = MODES[mode_type] + + # Process each image in the batch + all_tiles = [] + + for batch_idx in range(batch_size): + # Convert tensor to PIL for easier cropping + pil_image = tensor_to_pil(image, batch_idx) + + # If force_uniform_tiles, resize the image to fit the grid exactly + if force_uniform_tiles: + target_width = cols * tile_width + target_height = rows * tile_height + if pil_image.width != target_width or pil_image.height != target_height: + pil_image = pil_image.resize((target_width, target_height), Image.Resampling.LANCZOS) + + batch_tiles = [] + + if mode == usdu.USDUMode.LINEAR: + # Linear mode: process tiles row by row + for yi in range(rows): + for xi in range(cols): + x1, y1, x2, y2 = self.calc_rectangle(xi, yi, tile_width, tile_height) + # Crop the tile + tile = pil_image.crop((x1, y1, min(x2, pil_image.width), min(y2, pil_image.height))) + + # Add padding if specified + if tile_padding > 0: + tile = pad_image2(tile, tile_padding, tile_padding, tile_padding, tile_padding, fill=True, blur=False) + + batch_tiles.append(pil_to_tensor(tile)) + + elif mode == usdu.USDUMode.CHESS: + # Chess mode: process tiles in checkerboard pattern + # First, determine tile colors + tiles_map = [] + for yi in range(rows): + tiles_map.append([]) + for xi in range(cols): + color = xi % 2 == 0 + if yi > 0 and yi % 2 != 0: + color = not color + tiles_map[yi].append(color) + + # Process white tiles first + for yi in range(rows): + for xi in range(cols): + if tiles_map[yi][xi]: # White tiles + x1, y1, x2, y2 = self.calc_rectangle(xi, yi, tile_width, tile_height) + tile = pil_image.crop((x1, y1, min(x2, pil_image.width), min(y2, pil_image.height))) + + if tile_padding > 0: + tile = pad_image2(tile, tile_padding, tile_padding, tile_padding, tile_padding, fill=True, blur=False) + + batch_tiles.append(pil_to_tensor(tile)) + + # Then process black tiles + for yi in range(rows): + for xi in range(cols): + if not tiles_map[yi][xi]: # Black tiles + x1, y1, x2, y2 = self.calc_rectangle(xi, yi, tile_width, tile_height) + tile = pil_image.crop((x1, y1, min(x2, pil_image.width), min(y2, pil_image.height))) + + if tile_padding > 0: + tile = pad_image2(tile, tile_padding, tile_padding, tile_padding, tile_padding, fill=True, blur=False) + + batch_tiles.append(pil_to_tensor(tile)) + + else: # USDUMode.NONE + # None mode: return the entire image as a single tile + tile = pil_image + if tile_padding > 0: + tile = pad_image2(tile, tile_padding, tile_padding, tile_padding, tile_padding, fill=True, blur=False) + batch_tiles.append(pil_to_tensor(tile)) + + all_tiles.extend(batch_tiles) + + # Stack all tiles into a single tensor + tiles_tensor = torch.cat(all_tiles, dim=0) + + # Calculate total tile count + if mode == usdu.USDUMode.NONE: + tile_count = batch_size + else: + tile_count = rows * cols * batch_size + + return (tiles_tensor, rows, cols, tile_count) + + # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { "UltimateSDUpscale": UltimateSDUpscale, "UltimateSDUpscaleNoUpscale": UltimateSDUpscaleNoUpscale, - "UltimateSDUpscaleCustomSample": UltimateSDUpscaleCustomSample + "UltimateSDUpscaleCustomSample": UltimateSDUpscaleCustomSample, + "UltimateSDUpscaleTiler": UltimateSDUpscaleTiler } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { "UltimateSDUpscale": "Ultimate SD Upscale", "UltimateSDUpscaleNoUpscale": "Ultimate SD Upscale (No Upscale)", - "UltimateSDUpscaleCustomSample": "Ultimate SD Upscale (Custom Sample)" + "UltimateSDUpscaleCustomSample": "Ultimate SD Upscale (Custom Sample)", + "UltimateSDUpscaleTiler": "Ultimate SD Upscale Tiler" } diff --git a/usdu_patch.py b/usdu_patch.py index f96f07e..d90a8b5 100644 --- a/usdu_patch.py +++ b/usdu_patch.py @@ -1,71 +1,503 @@ -# Make some patches to the script -from repositories import ultimate_upscale as usdu -import modules.shared as shared +""" +Refactored USD Upscaler batch processing patch. + +Preserves original behavior but: +- Organizes imports and helpers +- Replaces prints with logging +- Factors duplicated logic (tile preparation, batching, decoding) +- Uses functools.wraps when monkey-patching methods +- Adds type hints and docstrings for clarity +""" + +from __future__ import annotations + +import logging import math -from PIL import Image +import numpy as np +import torch + +from functools import wraps +from typing import Tuple, List, Iterable +from PIL import Image, ImageFilter, ImageDraw +from comfy_extras.nodes_custom_sampler import SamplerCustom + +import modules.shared as shared +from nodes import common_ksampler, VAEEncode, VAEDecode, VAEDecodeTiled +from repositories import ultimate_upscale as usdu + +logger = logging.getLogger(__name__) +logger.addHandler(logging.StreamHandler()) +logger.setLevel(logging.INFO) -if (not hasattr(Image, 'Resampling')): # For older versions of Pillow - Image.Resampling = Image - -# -# Instead of using multiples of 64, use multiples of 8 -# +# Compatibility for older Pillow versions +try: + Image.Resampling # type: ignore +except Exception: + Image.Resampling = Image # type: ignore -def round_length(length, multiple=8): +# ------------------------- +# Utility helpers +# ------------------------- +def round_length(length: int, multiple: int = 8) -> int: + """Round length to nearest multiple (default 8).""" return round(length / multiple) * multiple -# Upscaler -old_init = usdu.USDUpscaler.__init__ +# PIL <-> tensor helpers (adapted from your inline utilities) +def _pil_to_tensor(image: Image.Image) -> torch.Tensor: + """Convert PIL image to CHW-like float tensor in [0,1], with batch dim omitted.""" + arr = np.array(image).astype(np.float32) / 255.0 + t = torch.from_numpy(arr) + # Ensure a channel dimension: HxW -> HxWx1, or HxWxC + if t.ndim == 2: + t = t.unsqueeze(-1) + # Move channel last to channel-first if needed by your VAE? You used unsqueeze(0) previously, + # so preserve the previous behavior: add batch dim at dim=0 but do not permute channels. + t = t.unsqueeze(0) + return t -def new_init(self, p, image, upscaler_index, save_redraw, save_seams_fix, tile_width, tile_height): - p.width = round_length(image.width * p.upscale_by) - p.height = round_length(image.height * p.upscale_by) - old_init(self, p, image, upscaler_index, save_redraw, save_seams_fix, tile_width, tile_height) +def _tensor_to_pil(img_tensor: torch.Tensor, batch_index: int = 0) -> Image.Image: + """Convert tensor (with batch) to PIL image for a specific batch index.""" + safe = torch.nan_to_num(img_tensor[batch_index]) + arr = (255 * safe.cpu().numpy()).astype(np.uint8) + return Image.fromarray(arr) -usdu.USDUpscaler.__init__ = new_init - -# Redraw -old_setup_redraw = usdu.USDURedraw.init_draw +def _fix_crop_region(region: Tuple[int, int, int, int], image_size: Tuple[int, int]) -> Tuple[int, int, int, int]: + """Adjust crop region to remove trailing pixel if not touching border.""" + image_width, image_height = image_size + x1, y1, x2, y2 = region + if x2 < image_width: + x2 -= 1 + if y2 < image_height: + y2 -= 1 + return x1, y1, x2, y2 -def new_setup_redraw(self, p, width, height): - mask, draw = old_setup_redraw(self, p, width, height) - p.width = round_length(self.tile_width + self.padding) - p.height = round_length(self.tile_height + self.padding) - return mask, draw +def _get_crop_region(mask: Image.Image, pad: int = 0) -> Tuple[int, int, int, int]: + """Get the bounding box of the white region in a mask and pad it.""" + coords = mask.getbbox() + if coords is not None: + x1, y1, x2, y2 = coords + else: + # empty bbox => use inverted (no area) + x1, y1, x2, y2 = mask.width, mask.height, 0, 0 + x1 = max(x1 - pad, 0) + y1 = max(y1 - pad, 0) + x2 = min(x2 + pad, mask.width) + y2 = min(y2 + pad, mask.height) + return _fix_crop_region((x1, y1, x2, y2), (mask.width, mask.height)) -usdu.USDURedraw.init_draw = new_setup_redraw +def _expand_crop(region: Tuple[int, int, int, int], width: int, height: int, target_width: int, target_height: int) -> Tuple[Tuple[int, int, int, int], Tuple[int, int]]: + """Expand a crop region to target size while keeping it inside image.""" + x1, y1, x2, y2 = region + actual_w = x2 - x1 + actual_h = y2 - y1 -# Seams fix -old_setup_seams_fix = usdu.USDUSeamsFix.init_draw + # Expand horizontally + w_diff = target_width - actual_w + x2 = min(x2 + w_diff // 2, width) + w_diff = target_width - (x2 - x1) + x1 = max(x1 - w_diff, 0) + w_diff = target_width - (x2 - x1) + x2 = min(x2 + w_diff, width) + + # Expand vertically + h_diff = target_height - actual_h + y2 = min(y2 + h_diff // 2, height) + h_diff = target_height - (y2 - y1) + y1 = max(y1 - h_diff, 0) + h_diff = target_height - (y2 - y1) + y2 = min(y2 + h_diff, height) + + return (x1, y1, x2, y2), (target_width, target_height) -def new_setup_seams_fix(self, p): - old_setup_seams_fix(self, p) - p.width = round_length(self.tile_width + self.padding) - p.height = round_length(self.tile_height + self.padding) +def _crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad: int = 0, h_pad: int = 0): + """Placeholder simplified crop conditioning for batch processing (keeps original behavior).""" + # This intentionally mirrors your simplified version: returns same conditioning. + return cond -usdu.USDUSeamsFix.init_draw = new_setup_seams_fix +def _sample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise, custom_sampler, custom_sigmas): + """Sampling wrapper that supports a custom sampler or falls back to common_ksampler.""" + if custom_sampler is not None and custom_sigmas is not None: + kwargs = dict( + model=model, + add_noise=True, + noise_seed=seed, + cfg=cfg, + positive=positive, + negative=negative, + sampler=custom_sampler, + sigmas=custom_sigmas, + latent_image=latent + ) + if hasattr(SamplerCustom, "execute"): + (samples, _) = SamplerCustom.execute(**kwargs) + else: + custom_sample = SamplerCustom() + (samples, _) = getattr(custom_sample, custom_sample.FUNCTION)(**kwargs) + return samples + + (samples,) = common_ksampler(model, seed, steps, cfg, sampler_name, + scheduler, positive, negative, latent, denoise=denoise) + return samples -# -# Make the script upscale on a batch of images instead of one image -# +# ------------------------- +# Monkey patches for USDUpscaler sizing / redraw / seams fix +# ------------------------- +def patch_usdu_upscaler_init(): + """Patch USDUpscaler.__init__ to round upscaler p.width/p.height to multiples.""" + old_init = usdu.USDUpscaler.__init__ -old_upscale = usdu.USDUpscaler.upscale + @wraps(old_init) + def new_init(self, p, image, upscaler_index, save_redraw, save_seams_fix, tile_width, tile_height): + p.width = round_length(image.width * p.upscale_by) + p.height = round_length(image.height * p.upscale_by) + return old_init(self, p, image, upscaler_index, save_redraw, save_seams_fix, tile_width, tile_height) + + usdu.USDUpscaler.__init__ = new_init -def new_upscale(self): - old_upscale(self) - shared.batch = [self.image] + \ - [img.resize((self.p.width, self.p.height), resample=Image.LANCZOS) for img in shared.batch[1:]] +def patch_usdu_redraw_init(): + """Patch USDURedraw.init_draw to round tile size used for redraw.""" + old_init_draw = usdu.USDURedraw.init_draw + + @wraps(old_init_draw) + def new_init_draw(self, p, width, height): + mask, draw = old_init_draw(self, p, width, height) + p.width = round_length(self.tile_width + self.padding) + p.height = round_length(self.tile_height + self.padding) + return mask, draw + + usdu.USDURedraw.init_draw = new_init_draw -usdu.USDUpscaler.upscale = new_upscale +def patch_usdu_seams_fix_init(): + old_init = usdu.USDUSeamsFix.init_draw + + @wraps(old_init) + def new_init(self, p): + old_init(self, p) + p.width = round_length(self.tile_width + self.padding) + p.height = round_length(self.tile_height + self.padding) + + usdu.USDUSeamsFix.init_draw = new_init + + +def patch_usdu_upscale_method(): + """Patch USDUpscaler.upscale to keep shared.batch resized to p.width/p.height.""" + old_upscale = usdu.USDUpscaler.upscale + + @wraps(old_upscale) + def new_upscale(self): + old_upscale(self) + # Keep shared.batch consistent with the upscaling width/height for subsequent processing. + shared.batch = [self.image] + [ + img.resize((self.p.width, self.p.height), resample=Image.LANCZOS) + for img in shared.batch[1:] + ] + + usdu.USDUpscaler.upscale = new_upscale + + +# Apply patches +patch_usdu_upscaler_init() +patch_usdu_redraw_init() +patch_usdu_seams_fix_init() +patch_usdu_upscale_method() + + +# ------------------------- +# Patched script.run replacement +# ------------------------- +def patched_script_run(self, p, _, tile_width, tile_height, mask_blur, padding, seams_fix_width, seams_fix_denoise, seams_fix_padding, + upscaler_index, save_upscaled_image, redraw_mode, save_seams_fix_image, seams_fix_mask_blur, + seams_fix_type, target_size_type, custom_width, custom_height, custom_scale): + """ + Replacement for usdu.Script.run that preserves the original batch_size + and delegates to the (patched) USDUpscaler and redraw pipeline. + """ + preserved_batch_size = getattr(p, 'batch_size', 1) + logger.info("[USDU Batch Debug] Patched script.run() preserving batch_size=%s", preserved_batch_size) + + # Init (matching original code) + usdu.processing.fix_seed(p) + usdu.devices.torch_gc() + + # Keep original file-saving flags as in original code + p.do_not_save_grid = True + p.do_not_save_samples = True + p.inpaint_full_res = False + + p.inpainting_fill = 1 + p.n_iter = 1 + p.batch_size = preserved_batch_size + + seed = p.seed + + # Init image + init_img = p.init_images[0] + if init_img is None: + return usdu.processing.Processed(p, [], seed, "Empty image") + init_img = usdu.images.flatten(init_img, usdu.shared.opts.img2img_background_color) + + # Override size by user choice + if target_size_type == 1: + p.width = custom_width + p.height = custom_height + elif target_size_type == 2: + p.width = math.ceil((init_img.width * custom_scale) / 64) * 64 + p.height = math.ceil((init_img.height * custom_scale) / 64) * 64 + + # Create and run upscaler + upscaler = usdu.USDUpscaler(p, init_img, upscaler_index, save_upscaled_image, save_seams_fix_image, tile_width, tile_height) + upscaler.upscale() + + # Drawing & seams fix setup + upscaler.setup_redraw(redraw_mode, padding, mask_blur) + upscaler.setup_seams_fix(seams_fix_padding, seams_fix_denoise, seams_fix_mask_blur, seams_fix_width, seams_fix_type) + upscaler.print_info() + upscaler.add_extra_info() + upscaler.process() + result_images = upscaler.result_images + + logger.info("[USDU Batch Debug] Patched script.run() complete, batch_size=%s", p.batch_size) + return usdu.processing.Processed(p, result_images, seed, upscaler.initial_info or "") + + +# Replace the original script.run with patched version +usdu.Script.run = patched_script_run + + +# ------------------------- +# Batch processing helpers shared between linear and chess modes +# ------------------------- +def _prepare_tile_for_batch(calc_rectangle_fn, current_image: Image.Image, tx: int, ty: int, p) -> Tuple[Image.Image, Tuple[int, int, int, int], Image.Image, Tuple[int, int]]: + """ + Prepare cropped/resized tile, mask, crop-region and tile-size for encoding. + Returns: (cropped_tile, initial_tile_size, tile_mask, tile_size) + """ + tile_mask = Image.new("L", (current_image.width, current_image.height), "black") + tile_draw = ImageDraw.Draw(tile_mask) + tile_draw.rectangle(calc_rectangle_fn(tx, ty), fill="white") + + crop_region = _get_crop_region(tile_mask, p.inpaint_full_res_padding) + + if p.uniform_tile_mode: + x1, y1, x2, y2 = crop_region + crop_w = x2 - x1 + crop_h = y2 - y1 + crop_ratio = crop_w / crop_h if crop_h != 0 else 1.0 + p_ratio = p.width / p.height if p.height != 0 else 1.0 + if crop_ratio > p_ratio: + target_w = crop_w + target_h = round(crop_w / p_ratio) + else: + target_w = round(crop_h * p_ratio) + target_h = crop_h + crop_region, _ = _expand_crop(crop_region, tile_mask.width, tile_mask.height, target_w, target_h) + tile_size = (p.width, p.height) + else: + x1, y1, x2, y2 = crop_region + crop_w = x2 - x1 + crop_h = y2 - y1 + target_w = math.ceil(crop_w / 8) * 8 + target_h = math.ceil(crop_h / 8) * 8 + crop_region, tile_size = _expand_crop(crop_region, tile_mask.width, tile_mask.height, target_w, target_h) + + # Optional blur + if getattr(p, "mask_blur", 0) > 0: + tile_mask = tile_mask.filter(ImageFilter.GaussianBlur(p.mask_blur)) + + cropped_tile = current_image.crop(crop_region) + initial_tile_size = cropped_tile.size + if cropped_tile.size != tile_size: + cropped_tile = cropped_tile.resize(tile_size, Image.Resampling.LANCZOS) + + return cropped_tile, initial_tile_size, tile_mask, crop_region, tile_size + + +def _process_batch_tiles(p, + tiles_coords: Iterable[Tuple[int, int]], + current_image: Image.Image, + calc_rectangle_fn, + vae_encoder: VAEEncode, + vae_decoder: VAEDecode, + vae_decoder_tiled: VAEDecodeTiled) -> Image.Image: + """Encode, sample and decode a batch of tiles and composite them into current_image.""" + if not tiles_coords: + return current_image + + batch_tiles = [] + batch_masks = [] + batch_crop_regions = [] + batch_tile_sizes = [] + + for tx, ty in tiles_coords: + cropped_tile, initial_tile_size, tile_mask, crop_region, tile_size = _prepare_tile_for_batch(calc_rectangle_fn, current_image, tx, ty, p) + batch_tiles.append((cropped_tile, initial_tile_size)) + batch_masks.append(tile_mask) + batch_crop_regions.append(crop_region) + batch_tile_sizes.append(tile_size) + + # Encode tiles -> latent + batched_tensors = torch.cat([_pil_to_tensor(tile) for tile, _ in batch_tiles], dim=0) + (latent,) = vae_encoder.encode(p.vae, batched_tensors) + + # Condition from first tile (assume same) + first_crop_region = batch_crop_regions[0] + first_tile_size = batch_tile_sizes[0] + positive_cropped = _crop_cond(p.positive, first_crop_region, p.init_size, current_image.size, first_tile_size) + negative_cropped = _crop_cond(p.negative, first_crop_region, p.init_size, current_image.size, first_tile_size) + + # Sampling + samples = _sample(p.model, p.seed, p.steps, p.cfg, p.sampler_name, p.scheduler, + positive_cropped, negative_cropped, latent, p.denoise, + p.custom_sampler, p.custom_sigmas) + + # Update progress bar if present + if getattr(p, "progress_bar_enabled", False) and getattr(p, "pbar", None) is not None: + p.pbar.update(len(list(tiles_coords))) + + # Decode + if not getattr(p, "tiled_decode", False): + (decoded,) = vae_decoder.decode(p.vae, samples) + else: + (decoded,) = vae_decoder_tiled.decode(p.vae, samples, 512) + + # Composite tiles back + result_img = current_image + for idx, (tx, ty) in enumerate(tiles_coords): + tile_sampled = _tensor_to_pil(decoded, idx) + initial_tile_size = batch_tiles[idx][1] + crop_region = batch_crop_regions[idx] + tile_mask = batch_masks[idx] + + if tile_sampled.size != initial_tile_size: + tile_sampled = tile_sampled.resize(initial_tile_size, Image.Resampling.LANCZOS) + + image_tile_only = Image.new('RGBA', result_img.size) + image_tile_only.paste(tile_sampled, crop_region[:2]) + + # Add mask as alpha and composite + temp = image_tile_only.copy() + temp.putalpha(tile_mask) + image_tile_only.paste(temp, image_tile_only) + + result = result_img.convert('RGBA') + result.alpha_composite(image_tile_only) + result_img = result.convert('RGB') + + return result_img + + +# ------------------------- +# Replace USDURedraw.linear_process and chess_process with batched variants +# ------------------------- +def patch_usdu_linear_and_chess_process(): + old_linear = usdu.USDURedraw.linear_process + old_chess = usdu.USDURedraw.chess_process + + @wraps(old_linear) + def new_linear_process(self, p, image, rows, cols): + batch_size = getattr(p, 'batch_size', 1) + logger.info("[USDU Batch Debug] linear_process called batch_size=%s rows=%s cols=%s total_tiles=%s", batch_size, rows, cols, rows * cols) + + if batch_size <= 1: + logger.info("[USDU Batch Debug] Using original single-tile processing (batch_size=%s)", batch_size) + return old_linear(self, p, image, rows, cols) + + # Batch mode + vae_encoder = VAEEncode() + vae_decoder = VAEDecode() + vae_decoder_tiled = VAEDecodeTiled() + + mask_template, draw_template = self.init_draw(p, image.width, image.height) + tiles_to_process: List[Tuple[int, int]] = [] + batch_count = 0 + + for yi in range(rows): + for xi in range(cols): + if shared.state.interrupted: + break + + tiles_to_process.append((xi, yi)) + + if len(tiles_to_process) >= batch_size or (yi == rows - 1 and xi == cols - 1): + batch_count += 1 + logger.info("[USDU Batch Debug] Processing batch #%s with %s tiles: %s", batch_count, len(tiles_to_process), tiles_to_process) + image = _process_batch_tiles(p, tiles_to_process, image, self.calc_rectangle, vae_encoder, vae_decoder, vae_decoder_tiled) + tiles_to_process = [] + + logger.info("[USDU Batch Debug] Linear processing complete. Processed %s batches total.", batch_count) + + # Update shared.batch[0] with the processed image so it can be retrieved later + shared.batch[0] = image + + p.width = image.width + p.height = image.height + return image + + @wraps(old_chess) + def new_chess_process(self, p, image, rows, cols): + batch_size = getattr(p, 'batch_size', 1) + if batch_size <= 1: + return old_chess(self, p, image, rows, cols) + + vae_encoder = VAEEncode() + vae_decoder = VAEDecode() + vae_decoder_tiled = VAEDecodeTiled() + + mask_template, draw_template = self.init_draw(p, image.width, image.height) + + # Determine tile "white/black" order + tile_colors = [] + for yi in range(rows): + row_colors = [] + for xi in range(cols): + color = xi % 2 == 0 + if yi > 0 and yi % 2 != 0: + color = not color + row_colors.append(color) + tile_colors.append(row_colors) + + # Helper to iterate tiles in chess order: white first, then black + def chess_order_iter(white: bool): + for yi in range(rows): + for xi in range(cols): + if tile_colors[yi][xi] == white: + yield (xi, yi) + + # Process white tiles then black tiles + for color in (True, False): + tiles_to_process: List[Tuple[int, int]] = [] + for tx, ty in chess_order_iter(color): + if shared.state.interrupted: + break + tiles_to_process.append((tx, ty)) + if len(tiles_to_process) >= batch_size: + image = _process_batch_tiles(p, tiles_to_process, image, self.calc_rectangle, vae_encoder, vae_decoder, vae_decoder_tiled) + tiles_to_process = [] + if tiles_to_process: + image = _process_batch_tiles(p, tiles_to_process, image, self.calc_rectangle, vae_encoder, vae_decoder, vae_decoder_tiled) + + # Update shared.batch[0] with the processed image so it can be retrieved later + shared.batch[0] = image + + p.width = image.width + p.height = image.height + return image + + usdu.USDURedraw.linear_process = new_linear_process + usdu.USDURedraw.chess_process = new_chess_process + + +patch_usdu_linear_and_chess_process() +logger.info("USDU batch patches applied successfully.") From 5e02909c4136a3e1a029d9ddf454ed451ef19c82 Mon Sep 17 00:00:00 2001 From: larsupb Date: Sun, 7 Dec 2025 21:27:57 +0100 Subject: [PATCH 02/11] add batch_size to other USDU nodes --- nodes.py | 14 ++++++++------ usdu_patch.py | 1 - 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/nodes.py b/nodes.py index accf4d0..5a503d7 100644 --- a/nodes.py +++ b/nodes.py @@ -94,6 +94,7 @@ class UltimateSDUpscale: @classmethod def INPUT_TYPES(s): required, optional = USDU_base_inputs() + required.append(("batch_size", ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}))) return prepare_inputs(required, optional) RETURN_TYPES = ("IMAGE",) @@ -104,7 +105,7 @@ class UltimateSDUpscale: steps, cfg, sampler_name, scheduler, denoise, upscale_model, mode_type, tile_width, tile_height, mask_blur, tile_padding, seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur, - seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, + seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, batch_size, custom_sampler=None, custom_sigmas=None): # Store params self.tile_width = tile_width @@ -133,8 +134,8 @@ class UltimateSDUpscale: shared.batch = [tensor_to_pil(image, i) for i in range(len(image))] shared.batch_as_tensor = image - # Get batch_size from instance if available (for UltimateSDUpscaleNoUpscale) - batch_size = getattr(self, 'batch_size', 1) + # Store batch_size for use in processing + self.batch_size = batch_size print(f"[USDU Batch Debug] UltimateSDUpscale.upscale() using batch_size={batch_size}") # Processing @@ -200,13 +201,14 @@ class UltimateSDUpscaleNoUpscale(UltimateSDUpscale): steps, cfg, sampler_name, scheduler, denoise, None, mode_type, tile_width, tile_height, mask_blur, tile_padding, seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur, - seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode) + seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, batch_size) class UltimateSDUpscaleCustomSample(UltimateSDUpscale): @classmethod def INPUT_TYPES(s): required, optional = USDU_base_inputs() remove_input(required, "upscale_model") + required.append(("batch_size", ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}))) optional.append(("upscale_model", ("UPSCALE_MODEL",))) optional.append(("custom_sampler", ("SAMPLER",))) optional.append(("custom_sigmas", ("SIGMAS",))) @@ -220,14 +222,14 @@ class UltimateSDUpscaleCustomSample(UltimateSDUpscale): steps, cfg, sampler_name, scheduler, denoise, mode_type, tile_width, tile_height, mask_blur, tile_padding, seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur, - seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, + seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, batch_size, upscale_model=None, custom_sampler=None, custom_sigmas=None): return super().upscale(image, model, positive, negative, vae, upscale_by, seed, steps, cfg, sampler_name, scheduler, denoise, upscale_model, mode_type, tile_width, tile_height, mask_blur, tile_padding, seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur, - seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, + seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, batch_size, custom_sampler, custom_sigmas) diff --git a/usdu_patch.py b/usdu_patch.py index d90a8b5..f90a79d 100644 --- a/usdu_patch.py +++ b/usdu_patch.py @@ -45,7 +45,6 @@ def round_length(length: int, multiple: int = 8) -> int: return round(length / multiple) * multiple -# PIL <-> tensor helpers (adapted from your inline utilities) def _pil_to_tensor(image: Image.Image) -> torch.Tensor: """Convert PIL image to CHW-like float tensor in [0,1], with batch dim omitted.""" arr = np.array(image).astype(np.float32) / 255.0 From 1ff35eb196b246fc634a40a98f83eaa36e4b26f8 Mon Sep 17 00:00:00 2001 From: ssitu Date: Sat, 7 Feb 2026 12:35:08 -0500 Subject: [PATCH 03/11] Remove tiler --- nodes.py | 134 +------------------------------------------------------ 1 file changed, 1 insertion(+), 133 deletions(-) diff --git a/nodes.py b/nodes.py index 5a503d7..130039b 100644 --- a/nodes.py +++ b/nodes.py @@ -4,7 +4,7 @@ import logging import torch import comfy from usdu_patch import usdu -from utils import tensor_to_pil, pil_to_tensor, pad_image2 +from utils import tensor_to_pil, pil_to_tensor from modules.processing import StableDiffusionProcessing import modules.shared as shared from modules.upscaler import UpscalerData @@ -232,143 +232,12 @@ class UltimateSDUpscaleCustomSample(UltimateSDUpscale): seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, batch_size, custom_sampler, custom_sigmas) - -class UltimateSDUpscaleTiler: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "tile_width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}), - "tile_height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}), - "tile_padding": ("INT", {"default": 32, "min": 0, "max": MAX_RESOLUTION, "step": 8}), - "mode_type": (list(MODES.keys()),), - "force_uniform_tiles": ("BOOLEAN", {"default": True}), - } - } - - RETURN_TYPES = ("IMAGE", "INT", "INT", "INT") - RETURN_NAMES = ("tiles", "rows", "cols", "tile_count") - FUNCTION = "tile_image" - CATEGORY = "image/upscaling" - - def calc_rectangle(self, xi, yi, tile_width, tile_height): - """Calculate tile rectangle coordinates""" - x1 = xi * tile_width - y1 = yi * tile_height - x2 = xi * tile_width + tile_width - y2 = yi * tile_height + tile_height - return x1, y1, x2, y2 - - def tile_image(self, image, tile_width, tile_height, tile_padding, mode_type, force_uniform_tiles): - from PIL import Image - import math - - # Get the image dimensions (batch, height, width, channels) - batch_size = len(image) - img_height = image.shape[1] - img_width = image.shape[2] - - # Calculate grid dimensions - rows = math.ceil(img_height / tile_height) - cols = math.ceil(img_width / tile_width) - - mode = MODES[mode_type] - - # Process each image in the batch - all_tiles = [] - - for batch_idx in range(batch_size): - # Convert tensor to PIL for easier cropping - pil_image = tensor_to_pil(image, batch_idx) - - # If force_uniform_tiles, resize the image to fit the grid exactly - if force_uniform_tiles: - target_width = cols * tile_width - target_height = rows * tile_height - if pil_image.width != target_width or pil_image.height != target_height: - pil_image = pil_image.resize((target_width, target_height), Image.Resampling.LANCZOS) - - batch_tiles = [] - - if mode == usdu.USDUMode.LINEAR: - # Linear mode: process tiles row by row - for yi in range(rows): - for xi in range(cols): - x1, y1, x2, y2 = self.calc_rectangle(xi, yi, tile_width, tile_height) - # Crop the tile - tile = pil_image.crop((x1, y1, min(x2, pil_image.width), min(y2, pil_image.height))) - - # Add padding if specified - if tile_padding > 0: - tile = pad_image2(tile, tile_padding, tile_padding, tile_padding, tile_padding, fill=True, blur=False) - - batch_tiles.append(pil_to_tensor(tile)) - - elif mode == usdu.USDUMode.CHESS: - # Chess mode: process tiles in checkerboard pattern - # First, determine tile colors - tiles_map = [] - for yi in range(rows): - tiles_map.append([]) - for xi in range(cols): - color = xi % 2 == 0 - if yi > 0 and yi % 2 != 0: - color = not color - tiles_map[yi].append(color) - - # Process white tiles first - for yi in range(rows): - for xi in range(cols): - if tiles_map[yi][xi]: # White tiles - x1, y1, x2, y2 = self.calc_rectangle(xi, yi, tile_width, tile_height) - tile = pil_image.crop((x1, y1, min(x2, pil_image.width), min(y2, pil_image.height))) - - if tile_padding > 0: - tile = pad_image2(tile, tile_padding, tile_padding, tile_padding, tile_padding, fill=True, blur=False) - - batch_tiles.append(pil_to_tensor(tile)) - - # Then process black tiles - for yi in range(rows): - for xi in range(cols): - if not tiles_map[yi][xi]: # Black tiles - x1, y1, x2, y2 = self.calc_rectangle(xi, yi, tile_width, tile_height) - tile = pil_image.crop((x1, y1, min(x2, pil_image.width), min(y2, pil_image.height))) - - if tile_padding > 0: - tile = pad_image2(tile, tile_padding, tile_padding, tile_padding, tile_padding, fill=True, blur=False) - - batch_tiles.append(pil_to_tensor(tile)) - - else: # USDUMode.NONE - # None mode: return the entire image as a single tile - tile = pil_image - if tile_padding > 0: - tile = pad_image2(tile, tile_padding, tile_padding, tile_padding, tile_padding, fill=True, blur=False) - batch_tiles.append(pil_to_tensor(tile)) - - all_tiles.extend(batch_tiles) - - # Stack all tiles into a single tensor - tiles_tensor = torch.cat(all_tiles, dim=0) - - # Calculate total tile count - if mode == usdu.USDUMode.NONE: - tile_count = batch_size - else: - tile_count = rows * cols * batch_size - - return (tiles_tensor, rows, cols, tile_count) - - # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { "UltimateSDUpscale": UltimateSDUpscale, "UltimateSDUpscaleNoUpscale": UltimateSDUpscaleNoUpscale, "UltimateSDUpscaleCustomSample": UltimateSDUpscaleCustomSample, - "UltimateSDUpscaleTiler": UltimateSDUpscaleTiler } # A dictionary that contains the friendly/humanly readable titles for the nodes @@ -376,5 +245,4 @@ NODE_DISPLAY_NAME_MAPPINGS = { "UltimateSDUpscale": "Ultimate SD Upscale", "UltimateSDUpscaleNoUpscale": "Ultimate SD Upscale (No Upscale)", "UltimateSDUpscaleCustomSample": "Ultimate SD Upscale (Custom Sample)", - "UltimateSDUpscaleTiler": "Ultimate SD Upscale Tiler" } From fe0196319f19b71248269b73f0b0287f33cc2e47 Mon Sep 17 00:00:00 2001 From: ssitu Date: Sat, 7 Feb 2026 22:56:58 -0500 Subject: [PATCH 04/11] refactor: move batch_size param to one place --- usdu_nodes.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/usdu_nodes.py b/usdu_nodes.py index 51e3728..b277ab3 100644 --- a/usdu_nodes.py +++ b/usdu_nodes.py @@ -56,6 +56,7 @@ def USDU_base_inputs(): # Misc ("force_uniform_tiles", ("BOOLEAN", {"default": True, "tooltip": "Force all tiles to be the same as the set tile size, even when tiles could be smaller. This can help prevent the model from working with irregular tile sizes."})), ("tiled_decode", ("BOOLEAN", {"default": False, "tooltip": "Whether to use tiled decoding when decoding tiles."})), + ("batch_size", ("INT", {"default": 1, "min": 1, "max": 4096, "step": 1, "tooltip": "The number of tiles to process in a batch. Higher values can reduce processing time but use more VRAM."})), ] optional = [] @@ -94,7 +95,6 @@ class UltimateSDUpscale: @classmethod def INPUT_TYPES(s): required, optional = USDU_base_inputs() - required.append(("batch_size", ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}))) return prepare_inputs(required, optional) RETURN_TYPES = ("IMAGE",) @@ -108,7 +108,7 @@ class UltimateSDUpscale: steps, cfg, sampler_name, scheduler, denoise, upscale_model, mode_type, tile_width, tile_height, mask_blur, tile_padding, seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur, - seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, batch_size, + seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, batch_size=1, custom_sampler=None, custom_sigmas=None): # Store params self.tile_width = tile_width @@ -182,7 +182,6 @@ class UltimateSDUpscaleNoUpscale(UltimateSDUpscale): remove_input(required, "upscale_model") remove_input(required, "upscale_by") rename_input(required, "image", "upscaled_image") - required.append(("batch_size", ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}))) return prepare_inputs(required, optional) RETURN_TYPES = ("IMAGE",) @@ -195,7 +194,7 @@ class UltimateSDUpscaleNoUpscale(UltimateSDUpscale): steps, cfg, sampler_name, scheduler, denoise, mode_type, tile_width, tile_height, mask_blur, tile_padding, seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur, - seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, batch_size): + seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, batch_size=1): upscale_by = 1.0 # Store batch_size for use in processing @@ -213,7 +212,6 @@ class UltimateSDUpscaleCustomSample(UltimateSDUpscale): def INPUT_TYPES(s): required, optional = USDU_base_inputs() remove_input(required, "upscale_model") - required.append(("batch_size", ("INT", {"default": 1, "min": 1, "max": 16, "step": 1}))) optional.append(("upscale_model", ("UPSCALE_MODEL", {"tooltip": "The model to use for upscaling the image. If not provided, a simple Lanczos scaling will be used instead."}))) optional.append(("custom_sampler", ("SAMPLER", {"tooltip": "A custom sampler to use instead of the built-in ComfyUI sampler specified by sampler_name. Only used if both custom_sampler and custom_sigmas are provided."}))) optional.append(("custom_sigmas", ("SIGMAS", {"tooltip": "A custom noise schedule to use during sampling. Only used if both custom_sampler and custom_sigmas are provided."}))) @@ -229,7 +227,7 @@ class UltimateSDUpscaleCustomSample(UltimateSDUpscale): steps, cfg, sampler_name, scheduler, denoise, mode_type, tile_width, tile_height, mask_blur, tile_padding, seam_fix_mode, seam_fix_denoise, seam_fix_mask_blur, - seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, batch_size, + seam_fix_width, seam_fix_padding, force_uniform_tiles, tiled_decode, batch_size=1, upscale_model=None, custom_sampler=None, custom_sigmas=None): return super().upscale(image, model, positive, negative, vae, upscale_by, seed, From 4eed7a7d76765cf8b23e5286f8b016880b72df8e Mon Sep 17 00:00:00 2001 From: ssitu Date: Mon, 9 Feb 2026 17:40:11 -0500 Subject: [PATCH 05/11] fix: assert for batch size conflicts with non uniform tiles --- usdu_nodes.py | 1 + 1 file changed, 1 insertion(+) diff --git a/usdu_nodes.py b/usdu_nodes.py index b277ab3..59b8723 100644 --- a/usdu_nodes.py +++ b/usdu_nodes.py @@ -140,6 +140,7 @@ class UltimateSDUpscale: # Store batch_size for use in processing self.batch_size = batch_size print(f"[USDU Batch Debug] UltimateSDUpscale.upscale() using batch_size={batch_size}") + assert batch_size == 1 or force_uniform_tiles, "batch_size greater than 1 requires force_uniform_tiles to be True; all tiles in the batch must be the same size." # Processing sdprocessing = StableDiffusionProcessing( From 5d3887fad7142b4688b5eeb0d6c7c6e0593f4b48 Mon Sep 17 00:00:00 2001 From: ssitu Date: Mon, 9 Feb 2026 18:36:27 -0500 Subject: [PATCH 06/11] test: add batch size = 2 to tests, change params to better identify expected behavior - Higher denoise - Larger band pass size - Make image MAE thresholds more strict: from 0.05 -> 0.01 - Use a helper function for building test image names --- test/io_utils.py | 6 +++ test/test_base_images.py | 4 +- test/test_controlnet.py | 30 ++++++------- test/test_main_workflow.py | 88 ++++++++++++++++++++------------------ test/test_settings.py | 81 ++++++++++++++++++++--------------- test/test_tiling_modes.py | 30 +++++++------ 6 files changed, 131 insertions(+), 108 deletions(-) diff --git a/test/io_utils.py b/test/io_utils.py index 8f4dabf..f099ed9 100644 --- a/test/io_utils.py +++ b/test/io_utils.py @@ -18,3 +18,9 @@ def save_image(tensor, path: pathlib.Path): def load_image(path: pathlib.Path, device=None): """Load an image from disk and convert it to a tensor.""" return usdu_utils.pil_to_tensor(Image.open(path)).to(device=device) + + +def image_name_format(name: str, extension: str, batch_size: int = 1) -> str: + """Helper for building image names for tests.""" + batch_suffix = f"_batch{batch_size}" if batch_size > 1 else "" + return f"{name}{batch_suffix}{extension}" diff --git a/test/test_base_images.py b/test/test_base_images.py index 0f3642b..e06b45b 100644 --- a/test/test_base_images.py +++ b/test/test_base_images.py @@ -27,5 +27,5 @@ def test_base_image_matches_reference(base_image, test_dirs: DirectoryConfig): diff1 = img_tensor_mae(blur(im1), blur(test_im1)) diff2 = img_tensor_mae(blur(im2), blur(test_im2)) logger.info(f"Base Image Diff1: {diff1}, Diff2: {diff2}") - assert diff1 < 0.05, "Image 1 does not match its test image." - assert diff2 < 0.05, "Image 2 does not match its test image." + assert diff1 < 0.01, "Image 1 does not match its test image." + assert diff2 < 0.01, "Image 2 does not match its test image." diff --git a/test/test_controlnet.py b/test/test_controlnet.py index acd1afc..11e7101 100644 --- a/test/test_controlnet.py +++ b/test/test_controlnet.py @@ -9,27 +9,26 @@ import torch from setup_utils import execute from tensor_utils import img_tensor_mae, blur -from io_utils import save_image, load_image +from io_utils import save_image, load_image, image_name_format from configs import DirectoryConfig from fixtures_images import EXT CATEGORY = pathlib.Path(pathlib.Path(__file__).stem.removeprefix("test_")) -CONTROLNET_TILE_OUTPUT_IMAGE = "controlnet_tile" + EXT TEST_CONTROLNET_TILE_MODEL = "control_v11f1e_sd15_tile.pth" +@pytest.mark.parametrize("batch_size", [1, 2]) class TestControlNet: """Integration tests for the upscaling workflow with ControlNet.""" - @pytest.fixture(scope="class") - def controlnet_upscaled_image( + def test_controlnet_tile( self, base_image, loaded_checkpoint, - upscale_model, node_classes, seed, - test_dirs, + batch_size, + test_dirs: DirectoryConfig, ): """Generate upscaled images using ControlNet.""" image, positive, negative = base_image @@ -72,25 +71,20 @@ class TestControlNet: seam_fix_padding=16, force_uniform_tiles=True, tiled_decode=False, + batch_size=batch_size, ) # Save and reload sample image sample_dir = test_dirs.sample_images - filename = CATEGORY / CONTROLNET_TILE_OUTPUT_IMAGE + filename = CATEGORY / image_name_format("controlnet_tile", EXT, batch_size) save_image(upscaled[0], sample_dir / filename) upscaled = load_image(sample_dir / filename) - return upscaled - def test_controlnet_upscaled_image_matches_reference( - self, controlnet_upscaled_image, test_dirs: DirectoryConfig - ): - """ - Verify ControlNet upscaled images match reference images. - """ - logger = logging.getLogger("test_controlnet_upscaled_image_matches_reference") + # Verify against reference image + logger = logging.getLogger("test_controlnet_tile") test_img_dir = test_dirs.test_images - test_img = load_image(test_img_dir / CATEGORY / CONTROLNET_TILE_OUTPUT_IMAGE) + test_img = load_image(test_img_dir / filename) # Reduce high-frequency noise differences with gaussian blur - diff = img_tensor_mae(blur(controlnet_upscaled_image), blur(test_img)) + diff = img_tensor_mae(blur(upscaled), blur(test_img)) logger.info(f"ControlNet Upscaled Image Diff: {diff}") - assert diff < 0.05, "ControlNet upscaled image does not match its test image." + assert diff < 0.01, "ControlNet upscaled image does not match its test image." diff --git a/test/test_main_workflow.py b/test/test_main_workflow.py index 275655f..ee183d0 100644 --- a/test/test_main_workflow.py +++ b/test/test_main_workflow.py @@ -4,24 +4,20 @@ Tests a common workflow for UltimateSDUpscale. import logging import pathlib +import pytest import torch from setup_utils import execute from tensor_utils import img_tensor_mae, blur -from io_utils import save_image, load_image +from io_utils import save_image, load_image, image_name_format from configs import DirectoryConfig from fixtures_images import EXT # Image file names CATEGORY = pathlib.Path(pathlib.Path(__file__).stem.removeprefix("test_")) -IMAGE_1 = CATEGORY / ("main1_sd15_upscaled" + EXT) -IMAGE_2 = CATEGORY / ("main2_sd15_upscaled" + EXT) -NO_UPSCALE_IMAGE_1 = CATEGORY / ("main1_sd15_upscaled_no_upscale" + EXT) -NO_UPSCALE_IMAGE_2 = CATEGORY / ("main2_sd15_upscaled_no_upscale" + EXT) -CUSTOM_SAMPLER_IMAGE_1 = CATEGORY / ("main1_sd15_upscaled_custom_sampler" + EXT) -CUSTOM_SAMPLER_IMAGE_2 = CATEGORY / ("main2_sd15_upscaled_custom_sampler" + EXT) +@pytest.mark.parametrize("batch_size", [1, 2]) class TestMainWorkflow: """Integration tests for the main upscaling workflow.""" @@ -32,6 +28,7 @@ class TestMainWorkflow: upscale_model, node_classes, seed, + batch_size, test_dirs: DirectoryConfig, ): """Generate upscaled images using standard workflow.""" @@ -48,11 +45,11 @@ class TestMainWorkflow: vae=vae, upscale_by=2.00000004, # Test small float difference doesn't add extra tiles seed=seed, - steps=10, + steps=5, cfg=8, sampler_name="euler", scheduler="normal", - denoise=0.2, + denoise=0.6, upscale_model=upscale_model, mode_type="Chess", tile_width=512, @@ -66,11 +63,14 @@ class TestMainWorkflow: seam_fix_padding=16, force_uniform_tiles=True, tiled_decode=False, + batch_size=batch_size, ) # Save images + im1_filename = image_name_format("upscaled_image1", EXT, batch_size) + im2_filename = image_name_format("upscaled_image2", EXT, batch_size) sample_dir = test_dirs.sample_images - upscaled_img1_path = sample_dir / IMAGE_1 - upscaled_img2_path = sample_dir / IMAGE_2 + upscaled_img1_path = sample_dir / CATEGORY / im1_filename + upscaled_img2_path = sample_dir / CATEGORY / im2_filename save_image(upscaled[0], upscaled_img1_path) save_image(upscaled[1], upscaled_img2_path) # Load to account for compression @@ -83,16 +83,15 @@ class TestMainWorkflow: im1_upscaled = upscaled[0] im2_upscaled = upscaled[1] - test_im1_upscaled = load_image(test_image_dir / IMAGE_1) - test_im2_upscaled = load_image(test_image_dir / IMAGE_2) - - diff1 = img_tensor_mae(blur(im1_upscaled), blur(test_im1_upscaled)) - diff2 = img_tensor_mae(blur(im2_upscaled), blur(test_im2_upscaled)) + test_im1 = load_image(test_image_dir / CATEGORY / im1_filename) + test_im2 = load_image(test_image_dir / CATEGORY / im2_filename) + diff1 = img_tensor_mae(blur(im1_upscaled), blur(test_im1)) + diff2 = img_tensor_mae(blur(im2_upscaled), blur(test_im2)) # This tolerance is enough to handle both cpu and gpu as the device, as well as jpg compression differences. logger.info(f"Diff1: {diff1}, Diff2: {diff2}") - assert diff1 < 0.05, "Upscaled Image 1 doesn't match its test image." - assert diff2 < 0.05, "Upscaled Image 2 doesn't match its test image." + assert diff1 < 0.01, "Upscaled Image 1 doesn't match its test image." + assert diff2 < 0.01, "Upscaled Image 2 doesn't match its test image." def test_upscale_no_upscale( self, @@ -100,6 +99,7 @@ class TestMainWorkflow: loaded_checkpoint, node_classes, seed, + batch_size, test_dirs: DirectoryConfig, ): """Generate upscaled images using standard workflow using the no upscale node.""" @@ -121,11 +121,11 @@ class TestMainWorkflow: negative=negative, vae=vae, seed=seed, - steps=10, + steps=5, cfg=8, sampler_name="euler", scheduler="normal", - denoise=0.2, + denoise=0.6, mode_type="Chess", tile_width=512, tile_height=512, @@ -138,11 +138,14 @@ class TestMainWorkflow: seam_fix_padding=16, force_uniform_tiles=True, tiled_decode=False, + batch_size=batch_size, ) # Save images + im1_filename = image_name_format("no_upscale_image_1", EXT, batch_size) + im2_filename = image_name_format("no_upscale_image_2", EXT, batch_size) sample_dir = test_dirs.sample_images - upscaled_img1_path = sample_dir / NO_UPSCALE_IMAGE_1 - upscaled_img2_path = sample_dir / NO_UPSCALE_IMAGE_2 + upscaled_img1_path = sample_dir / CATEGORY / im1_filename + upscaled_img2_path = sample_dir / CATEGORY / im2_filename save_image(upscaled[0], upscaled_img1_path) save_image(upscaled[1], upscaled_img2_path) # Load to account for compression @@ -155,15 +158,15 @@ class TestMainWorkflow: im1_upscaled = upscaled[0] im2_upscaled = upscaled[1] - test_im1_upscaled = load_image(test_image_dir / NO_UPSCALE_IMAGE_1) - test_im2_upscaled = load_image(test_image_dir / NO_UPSCALE_IMAGE_2) + test_im1 = load_image(test_image_dir / CATEGORY / im1_filename) + test_im2 = load_image(test_image_dir / CATEGORY / im2_filename) - diff1 = img_tensor_mae(blur(im1_upscaled), blur(test_im1_upscaled)) - diff2 = img_tensor_mae(blur(im2_upscaled), blur(test_im2_upscaled)) + diff1 = img_tensor_mae(blur(im1_upscaled), blur(test_im1)) + diff2 = img_tensor_mae(blur(im2_upscaled), blur(test_im2)) # This tolerance is enough to handle both cpu and gpu as the device, as well as jpg compression differences. logger.info(f"Diff1: {diff1}, Diff2: {diff2}") - assert diff1 < 0.05, "No Upscale Image 1 doesn't match its test image." - assert diff2 < 0.05, "No Upscale Image 2 doesn't match its test image." + assert diff1 < 0.01, f"{im1_filename} doesn't match its test image." + assert diff2 < 0.01, f"{im2_filename} doesn't match its test image." def test_upscale_with_custom_sampler( self, @@ -172,6 +175,7 @@ class TestMainWorkflow: upscale_model, node_classes, seed, + batch_size, test_dirs: DirectoryConfig, ): """Generate upscaled images using standard workflow using the custom sampler node.""" @@ -181,8 +185,8 @@ class TestMainWorkflow: with torch.inference_mode(): # Setup custom scheduler and sampler custom_scheduler = node_classes["KarrasScheduler"] - (sigmas,) = execute(custom_scheduler, 20, 14.614642, 0.0291675, 7.0) - (_, sigmas) = execute(node_classes["SplitSigmasDenoise"], sigmas, 0.15) + (sigmas,) = execute(custom_scheduler, 9, 14.614642, 0.0291675, 7.0) + (_, sigmas) = execute(node_classes["SplitSigmasDenoise"], sigmas, 0.6) custom_sampler = node_classes["KSamplerSelect"] (sampler,) = execute(custom_sampler, "dpmpp_2m") @@ -201,7 +205,7 @@ class TestMainWorkflow: cfg=8, sampler_name="euler", scheduler="normal", - denoise=0.2, + denoise=1.0, upscale_model=upscale_model, mode_type="Chess", tile_width=512, @@ -209,19 +213,22 @@ class TestMainWorkflow: mask_blur=8, tile_padding=32, seam_fix_mode="None", - seam_fix_denoise=1.0, + seam_fix_denoise=0.5, seam_fix_width=64, seam_fix_mask_blur=8, seam_fix_padding=16, force_uniform_tiles=True, tiled_decode=False, + batch_size=batch_size, custom_sampler=sampler, custom_sigmas=sigmas, ) # Save images + im1_filename = image_name_format("custom_sampler1", EXT, batch_size) + im2_filename = image_name_format("custom_sampler2", EXT, batch_size) sample_dir = test_dirs.sample_images - upscaled_img1_path = sample_dir / CUSTOM_SAMPLER_IMAGE_1 - upscaled_img2_path = sample_dir / CUSTOM_SAMPLER_IMAGE_2 + upscaled_img1_path = sample_dir / CATEGORY / im1_filename + upscaled_img2_path = sample_dir / CATEGORY / im2_filename save_image(upscaled[0], upscaled_img1_path) save_image(upscaled[1], upscaled_img2_path) # Load to account for compression @@ -233,14 +240,13 @@ class TestMainWorkflow: test_image_dir = test_dirs.test_images im1_upscaled = upscaled[0] im2_upscaled = upscaled[1] + test_im1 = load_image(test_image_dir / CATEGORY / im1_filename) + test_im2 = load_image(test_image_dir / CATEGORY / im2_filename) - test_im1_upscaled = load_image(test_image_dir / CUSTOM_SAMPLER_IMAGE_1) - test_im2_upscaled = load_image(test_image_dir / CUSTOM_SAMPLER_IMAGE_2) - - diff1 = img_tensor_mae(blur(im1_upscaled), blur(test_im1_upscaled)) - diff2 = img_tensor_mae(blur(im2_upscaled), blur(test_im2_upscaled)) + diff1 = img_tensor_mae(blur(im1_upscaled), blur(test_im1)) + diff2 = img_tensor_mae(blur(im2_upscaled), blur(test_im2)) # This tolerance is enough to handle both cpu and gpu as the device, as well as jpg compression differences. logger.info(f"Diff1: {diff1}, Diff2: {diff2}") - assert diff1 < 0.05, "Upscaled Image 1 doesn't match its test image." - assert diff2 < 0.05, "Upscaled Image 2 doesn't match its test image." + assert diff1 < 0.01, f"{im1_filename} doesn't match its test image." + assert diff2 < 0.01, f"{im2_filename} doesn't match its test image." diff --git a/test/test_settings.py b/test/test_settings.py index 84641ac..3141e4a 100644 --- a/test/test_settings.py +++ b/test/test_settings.py @@ -6,9 +6,10 @@ import logging import pathlib import pytest import torch +from contextlib import nullcontext from tensor_utils import img_tensor_mae, blur -from io_utils import save_image, load_image +from io_utils import save_image, load_image, image_name_format from configs import DirectoryConfig from fixtures_images import EXT @@ -16,47 +17,59 @@ from fixtures_images import EXT CATEGORY = pathlib.Path(pathlib.Path(__file__).stem.removeprefix("test_")) +@pytest.mark.parametrize("batch_size", [1, 2]) def test_minimal_tile_sizes( - base_image, loaded_checkpoint, node_classes, seed, test_dirs: DirectoryConfig + base_image, + loaded_checkpoint, + node_classes, + seed, + batch_size, + test_dirs: DirectoryConfig, ): """Test upscaling with minimal tile sizes.""" - filename = "non_uniform_tiles" image, positive, negative = base_image + image = image[0:1] # 1 image for simplicity model, clip, vae = loaded_checkpoint with torch.inference_mode(): - usdu = node_classes["UltimateSDUpscale"] - (upscaled,) = usdu().upscale( - image=image[0:1], - model=model, - positive=positive, - negative=negative, - vae=vae, - upscale_by=1.5, - seed=seed, - steps=5, - cfg=8, - sampler_name="euler", - scheduler="normal", - denoise=0.15, - upscale_model=None, - mode_type="Chess", - tile_width=512, - tile_height=512, - mask_blur=8, - tile_padding=8, - seam_fix_mode="None", - seam_fix_denoise=1.0, - seam_fix_width=16, - seam_fix_mask_blur=8, - seam_fix_padding=4, - force_uniform_tiles=False, - tiled_decode=False, - ) + with pytest.raises(AssertionError) if batch_size > 1 else nullcontext(): + usdu = node_classes["UltimateSDUpscale"] + (upscaled,) = usdu().upscale( + image=image, + model=model, + positive=positive, + negative=negative, + vae=vae, + upscale_by=1.5, + seed=seed, + steps=5, + cfg=8, + sampler_name="euler", + scheduler="normal", + denoise=0.6, + upscale_model=None, + mode_type="Chess", + tile_width=512, + tile_height=512, + mask_blur=8, + tile_padding=8, + seam_fix_mode="None", + seam_fix_denoise=1.0, + seam_fix_width=16, + seam_fix_mask_blur=8, + seam_fix_padding=4, + force_uniform_tiles=False, # This should trigger the assertion for batch_size > 1 + tiled_decode=False, + batch_size=batch_size, + ) + + if batch_size > 1: + return # Test passed if assertion was raised # Save and reload sample image sample_dir = test_dirs.sample_images - filename_path = CATEGORY / (filename + EXT) + filename = image_name_format("non_uniform_tiles", EXT, batch_size) + filename_path = CATEGORY / filename save_image(upscaled[0], sample_dir / filename_path) upscaled = load_image(sample_dir / filename_path) @@ -65,5 +78,5 @@ def test_minimal_tile_sizes( test_image = load_image(test_image_dir / filename_path) diff = img_tensor_mae(blur(upscaled), blur(test_image)) logger = logging.getLogger(__name__) - logger.info(f"{filename} MAE: {diff}") - assert diff < 0.05, f"{filename} output doesn't match reference" + logger.info(f"{filename_path} MAE: {diff}") + assert diff < 0.01, f"{filename_path} output doesn't match reference" diff --git a/test/test_tiling_modes.py b/test/test_tiling_modes.py index b25437c..e276b6e 100644 --- a/test/test_tiling_modes.py +++ b/test/test_tiling_modes.py @@ -8,7 +8,7 @@ import pytest import torch from tensor_utils import img_tensor_mae, blur -from io_utils import save_image, load_image +from io_utils import save_image, load_image, image_name_format from configs import DirectoryConfig from fixtures_images import EXT @@ -16,11 +16,7 @@ from fixtures_images import EXT CATEGORY = pathlib.Path(pathlib.Path(__file__).stem.removeprefix("test_")) -def image_name_format(prefix: str, mode: str) -> str: - """Helper for the image name format for the tests below.""" - return f"{prefix}_{mode.lower().replace(' ', '_')}{EXT}" - - +@pytest.mark.parametrize("batch_size", [1, 2]) class TestTilingModes: def _test_upscale_variant( self, @@ -33,6 +29,7 @@ class TestTilingModes: seam_fix_mode, seam_fix_denoise, filename_prefix, + batch_size, ): """Helper method to test upscale variants with different parameters.""" logger = logging.getLogger(f"test_{filename_prefix}") @@ -53,7 +50,7 @@ class TestTilingModes: cfg=8, sampler_name="euler", scheduler="normal", - denoise=0.2, + denoise=0.9, upscale_model=None, mode_type=mode_type, tile_width=512, @@ -62,11 +59,12 @@ class TestTilingModes: tile_padding=32, seam_fix_mode=seam_fix_mode, seam_fix_denoise=seam_fix_denoise, - seam_fix_width=64, + seam_fix_width=256, seam_fix_mask_blur=8, seam_fix_padding=16, force_uniform_tiles=True, tiled_decode=False, + batch_size=batch_size, ) # Save and reload sample image @@ -79,8 +77,8 @@ class TestTilingModes: test_image_dir = test_dirs.test_images test_image = load_image(test_image_dir / filename) diff = img_tensor_mae(blur(upscaled), blur(test_image)) - logger.info(f"{filename_prefix} MAE: {diff}") - assert diff < 0.05, f"{filename_prefix} output doesn't match reference" + logger.info(f"{filename} MAE: {diff}") + assert diff < 0.01, f"{filename} output doesn't match reference" # "Chess" is tested in the main workflow test @pytest.mark.parametrize("mode_type", ["Linear", "None"]) @@ -91,10 +89,11 @@ class TestTilingModes: node_classes, seed, mode_type, + batch_size, test_dirs: DirectoryConfig, ): """Test different tiling mode types.""" - filename = image_name_format("mode", mode_type) + filename = image_name_format("mode_" + mode_type.lower(), EXT, batch_size) self._test_upscale_variant( base_image, loaded_checkpoint, @@ -105,6 +104,7 @@ class TestTilingModes: seam_fix_mode="None", seam_fix_denoise=1.0, filename_prefix=filename, + batch_size=batch_size, ) @pytest.mark.parametrize( @@ -117,10 +117,13 @@ class TestTilingModes: node_classes, seed, seam_fix_mode, + batch_size, test_dirs: DirectoryConfig, ): """Test different seam fix modes.""" - filename = image_name_format("seamfix", seam_fix_mode) + filename = image_name_format( + "seamfix_" + seam_fix_mode.lower().replace(" ", "_"), EXT, batch_size + ) self._test_upscale_variant( base_image, loaded_checkpoint, @@ -129,6 +132,7 @@ class TestTilingModes: test_dirs, mode_type="None", seam_fix_mode=seam_fix_mode, - seam_fix_denoise=0.5, + seam_fix_denoise=0.6, filename_prefix=filename, + batch_size=batch_size, ) From 6a4594da4c52043d388599efc20893ef8d07b5a0 Mon Sep 17 00:00:00 2001 From: ssitu Date: Wed, 11 Feb 2026 19:02:29 -0500 Subject: [PATCH 07/11] refactor: reuse functions in patch for batched tiles --- test/test_main_workflow.py | 8 ++-- test/test_settings.py | 15 +++---- usdu_patch.py | 90 ++++---------------------------------- 3 files changed, 19 insertions(+), 94 deletions(-) diff --git a/test/test_main_workflow.py b/test/test_main_workflow.py index ee183d0..78b2e7b 100644 --- a/test/test_main_workflow.py +++ b/test/test_main_workflow.py @@ -49,7 +49,7 @@ class TestMainWorkflow: cfg=8, sampler_name="euler", scheduler="normal", - denoise=0.6, + denoise=0.7, upscale_model=upscale_model, mode_type="Chess", tile_width=512, @@ -125,7 +125,7 @@ class TestMainWorkflow: cfg=8, sampler_name="euler", scheduler="normal", - denoise=0.6, + denoise=0.7, mode_type="Chess", tile_width=512, tile_height=512, @@ -185,8 +185,8 @@ class TestMainWorkflow: with torch.inference_mode(): # Setup custom scheduler and sampler custom_scheduler = node_classes["KarrasScheduler"] - (sigmas,) = execute(custom_scheduler, 9, 14.614642, 0.0291675, 7.0) - (_, sigmas) = execute(node_classes["SplitSigmasDenoise"], sigmas, 0.6) + (sigmas,) = execute(custom_scheduler, 10, 14.614642, 0.0291675, 7.0) + (_, sigmas) = execute(node_classes["SplitSigmasDenoise"], sigmas, 0.7) custom_sampler = node_classes["KSamplerSelect"] (sampler,) = execute(custom_sampler, "dpmpp_2m") diff --git a/test/test_settings.py b/test/test_settings.py index 3141e4a..f97fbd5 100644 --- a/test/test_settings.py +++ b/test/test_settings.py @@ -68,15 +68,14 @@ def test_minimal_tile_sizes( # Save and reload sample image sample_dir = test_dirs.sample_images - filename = image_name_format("non_uniform_tiles", EXT, batch_size) - filename_path = CATEGORY / filename - save_image(upscaled[0], sample_dir / filename_path) - upscaled = load_image(sample_dir / filename_path) + filename = CATEGORY / image_name_format("non_uniform_tiles", EXT, batch_size) + save_image(upscaled[0], sample_dir / filename) + upscaled = load_image(sample_dir / filename) # Compare with reference test_image_dir = test_dirs.test_images - test_image = load_image(test_image_dir / filename_path) - diff = img_tensor_mae(blur(upscaled), blur(test_image)) + test_image = load_image(test_image_dir / filename) logger = logging.getLogger(__name__) - logger.info(f"{filename_path} MAE: {diff}") - assert diff < 0.01, f"{filename_path} output doesn't match reference" + diff = img_tensor_mae(blur(upscaled), blur(test_image)) + logger.info(f"{filename} MAE: {diff}") + assert diff < 0.02, f"{filename} does not match reference (MAE {diff})" diff --git a/usdu_patch.py b/usdu_patch.py index f90a79d..8d9444a 100644 --- a/usdu_patch.py +++ b/usdu_patch.py @@ -24,6 +24,7 @@ from comfy_extras.nodes_custom_sampler import SamplerCustom import modules.shared as shared from nodes import common_ksampler, VAEEncode, VAEDecode, VAEDecodeTiled from repositories import ultimate_upscale as usdu +import usdu_utils logger = logging.getLogger(__name__) logger.addHandler(logging.StreamHandler()) @@ -45,81 +46,6 @@ def round_length(length: int, multiple: int = 8) -> int: return round(length / multiple) * multiple -def _pil_to_tensor(image: Image.Image) -> torch.Tensor: - """Convert PIL image to CHW-like float tensor in [0,1], with batch dim omitted.""" - arr = np.array(image).astype(np.float32) / 255.0 - t = torch.from_numpy(arr) - # Ensure a channel dimension: HxW -> HxWx1, or HxWxC - if t.ndim == 2: - t = t.unsqueeze(-1) - # Move channel last to channel-first if needed by your VAE? You used unsqueeze(0) previously, - # so preserve the previous behavior: add batch dim at dim=0 but do not permute channels. - t = t.unsqueeze(0) - return t - - -def _tensor_to_pil(img_tensor: torch.Tensor, batch_index: int = 0) -> Image.Image: - """Convert tensor (with batch) to PIL image for a specific batch index.""" - safe = torch.nan_to_num(img_tensor[batch_index]) - arr = (255 * safe.cpu().numpy()).astype(np.uint8) - return Image.fromarray(arr) - - -def _fix_crop_region(region: Tuple[int, int, int, int], image_size: Tuple[int, int]) -> Tuple[int, int, int, int]: - """Adjust crop region to remove trailing pixel if not touching border.""" - image_width, image_height = image_size - x1, y1, x2, y2 = region - if x2 < image_width: - x2 -= 1 - if y2 < image_height: - y2 -= 1 - return x1, y1, x2, y2 - - -def _get_crop_region(mask: Image.Image, pad: int = 0) -> Tuple[int, int, int, int]: - """Get the bounding box of the white region in a mask and pad it.""" - coords = mask.getbbox() - if coords is not None: - x1, y1, x2, y2 = coords - else: - # empty bbox => use inverted (no area) - x1, y1, x2, y2 = mask.width, mask.height, 0, 0 - x1 = max(x1 - pad, 0) - y1 = max(y1 - pad, 0) - x2 = min(x2 + pad, mask.width) - y2 = min(y2 + pad, mask.height) - return _fix_crop_region((x1, y1, x2, y2), (mask.width, mask.height)) - - -def _expand_crop(region: Tuple[int, int, int, int], width: int, height: int, target_width: int, target_height: int) -> Tuple[Tuple[int, int, int, int], Tuple[int, int]]: - """Expand a crop region to target size while keeping it inside image.""" - x1, y1, x2, y2 = region - actual_w = x2 - x1 - actual_h = y2 - y1 - - # Expand horizontally - w_diff = target_width - actual_w - x2 = min(x2 + w_diff // 2, width) - w_diff = target_width - (x2 - x1) - x1 = max(x1 - w_diff, 0) - w_diff = target_width - (x2 - x1) - x2 = min(x2 + w_diff, width) - - # Expand vertically - h_diff = target_height - actual_h - y2 = min(y2 + h_diff // 2, height) - h_diff = target_height - (y2 - y1) - y1 = max(y1 - h_diff, 0) - h_diff = target_height - (y2 - y1) - y2 = min(y2 + h_diff, height) - - return (x1, y1, x2, y2), (target_width, target_height) - - -def _crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad: int = 0, h_pad: int = 0): - """Placeholder simplified crop conditioning for batch processing (keeps original behavior).""" - # This intentionally mirrors your simplified version: returns same conditioning. - return cond def _sample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise, custom_sampler, custom_sigmas): @@ -287,7 +213,7 @@ def _prepare_tile_for_batch(calc_rectangle_fn, current_image: Image.Image, tx: i tile_draw = ImageDraw.Draw(tile_mask) tile_draw.rectangle(calc_rectangle_fn(tx, ty), fill="white") - crop_region = _get_crop_region(tile_mask, p.inpaint_full_res_padding) + crop_region = usdu_utils.get_crop_region(tile_mask, p.inpaint_full_res_padding) if p.uniform_tile_mode: x1, y1, x2, y2 = crop_region @@ -301,7 +227,7 @@ def _prepare_tile_for_batch(calc_rectangle_fn, current_image: Image.Image, tx: i else: target_w = round(crop_h * p_ratio) target_h = crop_h - crop_region, _ = _expand_crop(crop_region, tile_mask.width, tile_mask.height, target_w, target_h) + crop_region, _ = usdu_utils.expand_crop(crop_region, tile_mask.width, tile_mask.height, target_w, target_h) tile_size = (p.width, p.height) else: x1, y1, x2, y2 = crop_region @@ -309,7 +235,7 @@ def _prepare_tile_for_batch(calc_rectangle_fn, current_image: Image.Image, tx: i crop_h = y2 - y1 target_w = math.ceil(crop_w / 8) * 8 target_h = math.ceil(crop_h / 8) * 8 - crop_region, tile_size = _expand_crop(crop_region, tile_mask.width, tile_mask.height, target_w, target_h) + crop_region, tile_size = usdu_utils.expand_crop(crop_region, tile_mask.width, tile_mask.height, target_w, target_h) # Optional blur if getattr(p, "mask_blur", 0) > 0: @@ -347,14 +273,14 @@ def _process_batch_tiles(p, batch_tile_sizes.append(tile_size) # Encode tiles -> latent - batched_tensors = torch.cat([_pil_to_tensor(tile) for tile, _ in batch_tiles], dim=0) + batched_tensors = torch.cat([usdu_utils.pil_to_tensor(tile) for tile, _ in batch_tiles], dim=0) (latent,) = vae_encoder.encode(p.vae, batched_tensors) # Condition from first tile (assume same) first_crop_region = batch_crop_regions[0] first_tile_size = batch_tile_sizes[0] - positive_cropped = _crop_cond(p.positive, first_crop_region, p.init_size, current_image.size, first_tile_size) - negative_cropped = _crop_cond(p.negative, first_crop_region, p.init_size, current_image.size, first_tile_size) + positive_cropped = usdu_utils.crop_cond(p.positive, first_crop_region, p.init_size, current_image.size, first_tile_size) + negative_cropped = usdu_utils.crop_cond(p.negative, first_crop_region, p.init_size, current_image.size, first_tile_size) # Sampling samples = _sample(p.model, p.seed, p.steps, p.cfg, p.sampler_name, p.scheduler, @@ -374,7 +300,7 @@ def _process_batch_tiles(p, # Composite tiles back result_img = current_image for idx, (tx, ty) in enumerate(tiles_coords): - tile_sampled = _tensor_to_pil(decoded, idx) + tile_sampled = usdu_utils.tensor_to_pil(decoded, idx) initial_tile_size = batch_tiles[idx][1] crop_region = batch_crop_regions[idx] tile_mask = batch_masks[idx] From bb8e9a3b38b56837fe07a4b70b4bf80c94f57bd3 Mon Sep 17 00:00:00 2001 From: ssitu Date: Wed, 11 Feb 2026 19:26:16 -0500 Subject: [PATCH 08/11] fix: individual controlnet hints for batched tiles --- usdu_patch.py | 5 ++-- usdu_utils.py | 82 +++++++++++++++++++++++++++++++++++++++++---------- 2 files changed, 69 insertions(+), 18 deletions(-) diff --git a/usdu_patch.py b/usdu_patch.py index 8d9444a..36f8d80 100644 --- a/usdu_patch.py +++ b/usdu_patch.py @@ -277,10 +277,9 @@ def _process_batch_tiles(p, (latent,) = vae_encoder.encode(p.vae, batched_tensors) # Condition from first tile (assume same) - first_crop_region = batch_crop_regions[0] first_tile_size = batch_tile_sizes[0] - positive_cropped = usdu_utils.crop_cond(p.positive, first_crop_region, p.init_size, current_image.size, first_tile_size) - negative_cropped = usdu_utils.crop_cond(p.negative, first_crop_region, p.init_size, current_image.size, first_tile_size) + positive_cropped = usdu_utils.crop_cond(p.positive, batch_crop_regions, p.init_size, current_image.size, first_tile_size) + negative_cropped = usdu_utils.crop_cond(p.negative, batch_crop_regions, p.init_size, current_image.size, first_tile_size) # Sampling samples = _sample(p.model, p.seed, p.steps, p.cfg, p.sampler_name, p.scheduler, diff --git a/usdu_utils.py b/usdu_utils.py index 479dd3d..2e015f7 100644 --- a/usdu_utils.py +++ b/usdu_utils.py @@ -295,19 +295,40 @@ def resize_and_pad_tensor(tensor, width, height, fill=False, blur=False): return result -def crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): +def crop_controlnet(cond_dict, regions, init_size, canvas_size, tile_size, w_pad, h_pad): + """ + Crop controlnet hints to the given region and resize them to the tile size. + If there are multiple regions, the hints will be cropped and resized for each region + and concatenated together in the batch dimension. + + Supports multiple regions. + + :param cond_dict: dict that contains the conditioning. + :param regions: A tuple or list of tuples of the form (x1, y1, x2, y2) denoting the + upper left and the lower right points of the rectangular region. + :param init_size: The original size of the image that the controlnet hints were generated for. + :param canvas_size: The size of the image that the controlnet hints will be resized to before cropping. + :param tile_size: The size to which each cropped hint will be resized. + :param w_pad: The horizontal padding added to each cropped hint. + :param h_pad: The vertical padding added to each cropped hint. + """ if "control" not in cond_dict: return + if not isinstance(regions, list): + regions = [regions] c = cond_dict["control"] controlnet = c.copy() cond_dict["control"] = controlnet while c is not None: # hint is shape (B, C, H, W) hint = controlnet.cond_hint_original - resized_crop = resize_region(region, canvas_size, hint.shape[:-3:-1]) - hint = crop_tensor(hint.movedim(1, -1), resized_crop).movedim(-1, 1) - hint = resize_tensor(hint, tile_size[::-1]) - controlnet.cond_hint_original = hint + tiled_hints = [] + for region in regions: + resized_crop = resize_region(region, canvas_size, hint.shape[:-3:-1]) + tiled_hint = crop_tensor(hint.movedim(1, -1), resized_crop).movedim(-1, 1) + tiled_hint = resize_tensor(tiled_hint, tile_size[::-1]) + tiled_hints.append(tiled_hint) + controlnet.cond_hint_original = torch.cat(tiled_hints, dim=0) c = c.previous_controlnet controlnet.set_previous_controlnet(c.copy() if c is not None else None) controlnet = controlnet.previous_controlnet @@ -333,9 +354,18 @@ def region_intersection(region1, region2): return (x1, y1, x2, y2) -def crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): +def crop_gligen(cond_dict, regions, init_size, canvas_size, tile_size, w_pad, h_pad): + """ + Crop gligen position conditioning to the given region. + + Does not support multiple regions. + """ if "gligen" not in cond_dict: return + + # Only use first region if multiple regions are given + region = regions if isinstance(regions, tuple) else regions[0] + type, model, cond = cond_dict["gligen"] if type != "position": from warnings import warn @@ -379,10 +409,18 @@ def crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_p cond_dict["gligen"] = (type, model, cropped) -def crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): +def crop_area(cond_dict, regions, init_size, canvas_size, tile_size, w_pad, h_pad): + """ + Crop area conditioning to the given region. + + Does not support multiple regions. + """ if "area" not in cond_dict: return + # Only use first region if multiple regions are given + region = regions if isinstance(regions, tuple) else regions[0] + # Resize the area conditioning to the canvas size and confine it to the tile region h, w, y, x = cond_dict["area"] w, h, x, y = 8 * w, 8 * h, 8 * x, 8 * y @@ -413,9 +451,18 @@ def crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad cond_dict["area"] = (h, w, y, x) -def crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): +def crop_mask(cond_dict, regions, init_size, canvas_size, tile_size, w_pad, h_pad): + """ + Crop the mask conditioning to the given region + + Does not support multiple regions. + """ if "mask" not in cond_dict: return + + # Only use first region if multiple regions are given + region = regions if isinstance(regions, tuple) else regions[0] + mask_tensor = cond_dict["mask"] # (B, H, W) masks = [] for i in range(mask_tensor.shape[0]): @@ -443,18 +490,23 @@ def crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad cond_dict["mask"] = torch.cat(masks, dim=0) # (B, H, W) # Added Flux-Kontext Support crop_reference_latents by TBG ETUR -def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad): +def crop_reference_latents(cond_dict, regions, init_size, canvas_size, tile_size, w_pad, h_pad): """ 1. Resize each latent to `canvas_size` in latent units. 2. Crop the rectangle `region` (pixel coordinates). 3. Down-sample the crop to latent-space `tile_size`. Expects a list of BCHW tensors under "reference_latents". + + Does not support multiple regions. """ latents = cond_dict.get("reference_latents") if not isinstance(latents, list): return # nothing to do + # Only use first region if multiple regions are given + region = regions if isinstance(regions, tuple) else regions[0] + k = 8 # down-sample factor from pixel space → latent space (SD-type models) W_can_px, H_can_px = canvas_size @@ -503,15 +555,15 @@ def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, -def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0): +def crop_cond(cond, regions, init_size, canvas_size, tile_size, w_pad=0, h_pad=0): cropped = [] for emb, x in cond: cond_dict = x.copy() n = [emb, cond_dict] - crop_controlnet(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) - crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) - crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) - crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) - crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad) + crop_controlnet(cond_dict, regions, init_size, canvas_size, tile_size, w_pad, h_pad) + crop_gligen(cond_dict, regions, init_size, canvas_size, tile_size, w_pad, h_pad) + crop_area(cond_dict, regions, init_size, canvas_size, tile_size, w_pad, h_pad) + crop_mask(cond_dict, regions, init_size, canvas_size, tile_size, w_pad, h_pad) + crop_reference_latents(cond_dict, regions, init_size, canvas_size, tile_size, w_pad, h_pad) cropped.append(n) return cropped From 3054e22c6902c52361df3b675b85e19864046079 Mon Sep 17 00:00:00 2001 From: ssitu Date: Thu, 12 Feb 2026 17:11:43 -0500 Subject: [PATCH 09/11] fix: sample whole batch instead of just the first image --- usdu_patch.py | 80 +++++++++++++++++++++++++-------------------------- 1 file changed, 39 insertions(+), 41 deletions(-) diff --git a/usdu_patch.py b/usdu_patch.py index 36f8d80..3a79537 100644 --- a/usdu_patch.py +++ b/usdu_patch.py @@ -250,27 +250,28 @@ def _prepare_tile_for_batch(calc_rectangle_fn, current_image: Image.Image, tx: i def _process_batch_tiles(p, - tiles_coords: Iterable[Tuple[int, int]], - current_image: Image.Image, + tiles_coords: List[Tuple[int, int]], + images: List[Image.Image], calc_rectangle_fn, vae_encoder: VAEEncode, vae_decoder: VAEDecode, - vae_decoder_tiled: VAEDecodeTiled) -> Image.Image: - """Encode, sample and decode a batch of tiles and composite them into current_image.""" - if not tiles_coords: - return current_image + vae_decoder_tiled: VAEDecodeTiled) -> List[Image.Image]: + """Encode, sample and decode a batch of tiles and composite them into the given images.""" + if not tiles_coords or not images: + return images batch_tiles = [] batch_masks = [] batch_crop_regions = [] batch_tile_sizes = [] - for tx, ty in tiles_coords: - cropped_tile, initial_tile_size, tile_mask, crop_region, tile_size = _prepare_tile_for_batch(calc_rectangle_fn, current_image, tx, ty, p) - batch_tiles.append((cropped_tile, initial_tile_size)) - batch_masks.append(tile_mask) - batch_crop_regions.append(crop_region) - batch_tile_sizes.append(tile_size) + for image in images: + for tx, ty in tiles_coords: + cropped_tile, initial_tile_size, tile_mask, crop_region, tile_size = _prepare_tile_for_batch(calc_rectangle_fn, image, tx, ty, p) + batch_tiles.append((cropped_tile, initial_tile_size)) + batch_masks.append(tile_mask) + batch_crop_regions.append(crop_region) + batch_tile_sizes.append(tile_size) # Encode tiles -> latent batched_tensors = torch.cat([usdu_utils.pil_to_tensor(tile) for tile, _ in batch_tiles], dim=0) @@ -278,8 +279,8 @@ def _process_batch_tiles(p, # Condition from first tile (assume same) first_tile_size = batch_tile_sizes[0] - positive_cropped = usdu_utils.crop_cond(p.positive, batch_crop_regions, p.init_size, current_image.size, first_tile_size) - negative_cropped = usdu_utils.crop_cond(p.negative, batch_crop_regions, p.init_size, current_image.size, first_tile_size) + positive_cropped = usdu_utils.crop_cond(p.positive, batch_crop_regions, p.init_size, images[0].size, first_tile_size) + negative_cropped = usdu_utils.crop_cond(p.negative, batch_crop_regions, p.init_size, images[0].size, first_tile_size) # Sampling samples = _sample(p.model, p.seed, p.steps, p.cfg, p.sampler_name, p.scheduler, @@ -297,29 +298,32 @@ def _process_batch_tiles(p, (decoded,) = vae_decoder_tiled.decode(p.vae, samples, 512) # Composite tiles back - result_img = current_image - for idx, (tx, ty) in enumerate(tiles_coords): - tile_sampled = usdu_utils.tensor_to_pil(decoded, idx) - initial_tile_size = batch_tiles[idx][1] - crop_region = batch_crop_regions[idx] - tile_mask = batch_masks[idx] + result_imgs = images + for i, result_img in enumerate(result_imgs): + for j, (tx, ty) in enumerate(tiles_coords): + idx = i * len(tiles_coords) + j + tile_sampled = usdu_utils.tensor_to_pil(decoded, idx) + initial_tile_size = batch_tiles[idx][1] + crop_region = batch_crop_regions[idx] + tile_mask = batch_masks[idx] - if tile_sampled.size != initial_tile_size: - tile_sampled = tile_sampled.resize(initial_tile_size, Image.Resampling.LANCZOS) + if tile_sampled.size != initial_tile_size: + tile_sampled = tile_sampled.resize(initial_tile_size, Image.Resampling.LANCZOS) - image_tile_only = Image.new('RGBA', result_img.size) - image_tile_only.paste(tile_sampled, crop_region[:2]) + image_tile_only = Image.new('RGBA', result_img.size) + image_tile_only.paste(tile_sampled, crop_region[:2]) - # Add mask as alpha and composite - temp = image_tile_only.copy() - temp.putalpha(tile_mask) - image_tile_only.paste(temp, image_tile_only) + # Add mask as alpha and composite + temp = image_tile_only.copy() + temp.putalpha(tile_mask) + image_tile_only.paste(temp, image_tile_only) - result = result_img.convert('RGBA') - result.alpha_composite(image_tile_only) - result_img = result.convert('RGB') + result = result_img.convert('RGBA') + result.alpha_composite(image_tile_only) + result_img = result.convert('RGB') + result_imgs[i] = result_img - return result_img + return result_imgs # ------------------------- @@ -357,14 +361,11 @@ def patch_usdu_linear_and_chess_process(): if len(tiles_to_process) >= batch_size or (yi == rows - 1 and xi == cols - 1): batch_count += 1 logger.info("[USDU Batch Debug] Processing batch #%s with %s tiles: %s", batch_count, len(tiles_to_process), tiles_to_process) - image = _process_batch_tiles(p, tiles_to_process, image, self.calc_rectangle, vae_encoder, vae_decoder, vae_decoder_tiled) + shared.batch = _process_batch_tiles(p, tiles_to_process, shared.batch, self.calc_rectangle, vae_encoder, vae_decoder, vae_decoder_tiled) tiles_to_process = [] logger.info("[USDU Batch Debug] Linear processing complete. Processed %s batches total.", batch_count) - # Update shared.batch[0] with the processed image so it can be retrieved later - shared.batch[0] = image - p.width = image.width p.height = image.height return image @@ -407,13 +408,10 @@ def patch_usdu_linear_and_chess_process(): break tiles_to_process.append((tx, ty)) if len(tiles_to_process) >= batch_size: - image = _process_batch_tiles(p, tiles_to_process, image, self.calc_rectangle, vae_encoder, vae_decoder, vae_decoder_tiled) + shared.batch = _process_batch_tiles(p, tiles_to_process, shared.batch, self.calc_rectangle, vae_encoder, vae_decoder, vae_decoder_tiled) tiles_to_process = [] if tiles_to_process: - image = _process_batch_tiles(p, tiles_to_process, image, self.calc_rectangle, vae_encoder, vae_decoder, vae_decoder_tiled) - - # Update shared.batch[0] with the processed image so it can be retrieved later - shared.batch[0] = image + shared.batch = _process_batch_tiles(p, tiles_to_process, shared.batch, self.calc_rectangle, vae_encoder, vae_decoder, vae_decoder_tiled) p.width = image.width p.height = image.height From 44d83685c706d1bd34d597c2a71cd3e3a42cddae Mon Sep 17 00:00:00 2001 From: ssitu Date: Thu, 12 Feb 2026 23:35:15 -0500 Subject: [PATCH 10/11] test: relax custom sampling MAE threshold --- test/test_main_workflow.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_main_workflow.py b/test/test_main_workflow.py index 78b2e7b..5880610 100644 --- a/test/test_main_workflow.py +++ b/test/test_main_workflow.py @@ -248,5 +248,5 @@ class TestMainWorkflow: # This tolerance is enough to handle both cpu and gpu as the device, as well as jpg compression differences. logger.info(f"Diff1: {diff1}, Diff2: {diff2}") - assert diff1 < 0.01, f"{im1_filename} doesn't match its test image." - assert diff2 < 0.01, f"{im2_filename} doesn't match its test image." + assert diff1 < 0.011, f"{im1_filename} doesn't match its test image." + assert diff2 < 0.011, f"{im2_filename} doesn't match its test image." From 6b210655dbdbc928bc5554d891af50010072800b Mon Sep 17 00:00:00 2001 From: ssitu Date: Thu, 12 Feb 2026 23:53:04 -0500 Subject: [PATCH 11/11] docs: more descriptive batch size tooltip --- usdu_nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/usdu_nodes.py b/usdu_nodes.py index 59b8723..4a476ea 100644 --- a/usdu_nodes.py +++ b/usdu_nodes.py @@ -56,7 +56,7 @@ def USDU_base_inputs(): # Misc ("force_uniform_tiles", ("BOOLEAN", {"default": True, "tooltip": "Force all tiles to be the same as the set tile size, even when tiles could be smaller. This can help prevent the model from working with irregular tile sizes."})), ("tiled_decode", ("BOOLEAN", {"default": False, "tooltip": "Whether to use tiled decoding when decoding tiles."})), - ("batch_size", ("INT", {"default": 1, "min": 1, "max": 4096, "step": 1, "tooltip": "The number of tiles to process in a batch. Higher values can reduce processing time but use more VRAM."})), + ("batch_size", ("INT", {"default": 1, "min": 1, "max": 4096, "step": 1, "tooltip": "The number of tiles to process in a batch. Higher values can reduce processing time but use more VRAM. Yields different results than individual tiles. Only affects the main redraw step, not the seam fix step."})), ] optional = []