diff --git a/tests/test_static_mode.py b/tests/test_static_mode.py index f69a865..a824493 100644 --- a/tests/test_static_mode.py +++ b/tests/test_static_mode.py @@ -176,6 +176,13 @@ class _FakeStaticWorker(static_mode.StaticModeMixin): processed_batch = torch.zeros((batch_size, 2, 2, 3), dtype=torch.float32) return processed_batch, 0, 0, 2, 2 + def create_tile_mask(self, *_args, **_kwargs): + from PIL import Image + return Image.new("L", (4, 4), 255) + + def blend_tile(self, base_image, *_args, **_kwargs): + return base_image + def _call_worker_static(fake_worker): image = torch.zeros((1, 4, 4, 3), dtype=torch.float32) diff --git a/upscale/modes/dynamic.py b/upscale/modes/dynamic.py index d7d9a96..f22292f 100644 --- a/upscale/modes/dynamic.py +++ b/upscale/modes/dynamic.py @@ -74,14 +74,17 @@ class DynamicModeMixin: # Process locally single_tensor = upscaled_image[image_idx:image_idx+1] local_image = result_images[image_idx] - image_seed = seed + image_idx * 1000 + image_seed = seed # Pre-slice conditioning once per image (not per tile) positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, image_idx) for tile_idx, pos in enumerate(all_tiles): + source_tensor = pil_to_tensor(local_image) + if single_tensor.is_cuda: + source_tensor = source_tensor.cuda() local_image = self._process_and_blend_tile( - tile_idx, pos, single_tensor, local_image, + tile_idx, pos, source_tensor, local_image, model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, width, height, force_uniform_tiles, @@ -256,14 +259,17 @@ class DynamicModeMixin: local_image = tensor_to_pil(single_tensor, 0).copy() # Process all tiles for this image - image_seed = seed + image_idx * 1000 + image_seed = seed # Pre-slice conditioning once per image (not per tile) positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, image_idx) for tile_idx, pos in enumerate(all_tiles): + source_tensor = pil_to_tensor(local_image) + if single_tensor.is_cuda: + source_tensor = source_tensor.cuda() local_image = self._process_and_blend_tile( - tile_idx, pos, single_tensor, local_image, + tile_idx, pos, source_tensor, local_image, model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, width, height, force_uniform_tiles, diff --git a/upscale/modes/single_gpu.py b/upscale/modes/single_gpu.py index 9b96257..5b86595 100644 --- a/upscale/modes/single_gpu.py +++ b/upscale/modes/single_gpu.py @@ -38,9 +38,14 @@ class SingleGpuModeMixin: # Process tiles batched across images for tile_idx, (tx, ty) in enumerate(all_tiles): + # Progressive state parity: extract each tile from the current updated image batch. + source_batch = torch.cat([pil_to_tensor(img) for img in result_images], dim=0) + if upscaled_image.is_cuda: + source_batch = source_batch.cuda() + # Extract batched tile tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding( - upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles + source_batch, tx, ty, tile_width, tile_height, padding, force_uniform_tiles ) # Process batch diff --git a/upscale/modes/static.py b/upscale/modes/static.py index 2c61eed..abd4d67 100644 --- a/upscale/modes/static.py +++ b/upscale/modes/static.py @@ -148,8 +148,11 @@ class StaticModeMixin: height, ): """Process one tile_id across the batch and blend into result_images.""" + source_batch = torch.cat([pil_to_tensor(img) for img in result_images], dim=0) + if upscaled_image.is_cuda: + source_batch = source_batch.cuda() processed_batch, x1, y1, ew, eh = self._extract_and_process_tile( - upscaled_image, + source_batch, tile_id, all_tiles, tile_width, @@ -203,7 +206,13 @@ class StaticModeMixin: total_tiles = batch_size * num_tiles_per_image processed_tiles = [] - sliced_conditioning_cache = {} + working_images = [] + for b in range(batch_size): + image_pil = tensor_to_pil(upscaled_image[b:b+1], 0) + working_images.append(image_pil.copy()) + tile_masks = [] + for tx, ty in all_tiles: + tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur)) # Dynamic queue mode (static processing): process batched-per-tile log(f"USDU Dist Worker[{worker_id[:8]}]: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Tiles/image {num_tiles_per_image} | Batch {batch_size}") @@ -230,8 +239,11 @@ class StaticModeMixin: debug_log(f"Worker[{worker_id[:8]}] - Assigned tile_id {tile_idx}") processed_count += batch_size tile_id = tile_idx + source_batch = torch.cat([pil_to_tensor(img) for img in working_images], dim=0) + if upscaled_image.is_cuda: + source_batch = source_batch.cuda() processed_batch, x1, y1, ew, eh = self._extract_and_process_tile( - upscaled_image, + source_batch, tile_id, all_tiles, tile_width, @@ -254,6 +266,18 @@ class StaticModeMixin: ) # Queue results for b in range(batch_size): + tile_pil = tensor_to_pil(processed_batch, b) + if tile_pil.size != (ew, eh): + tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS) + working_images[b] = self.blend_tile( + working_images[b], + tile_pil, + x1, + y1, + (ew, eh), + tile_masks[tile_id], + padding, + ) processed_tiles.append({ 'tile': processed_batch[b:b+1], 'tile_idx': tile_id, @@ -360,7 +384,6 @@ class StaticModeMixin: image_pil = tensor_to_pil(upscaled_image[b:b+1], 0) result_images.append(image_pil.copy()) - sliced_conditioning_cache = {} # Initialize queue: pending queue holds tile ids (batched per tile) log("USDU Dist: Using tile queue distribution") run_async_in_server_loop( @@ -495,8 +518,14 @@ class StaticModeMixin: timeout=5.0 ) - # Blend worker tiles synchronously - for global_idx, tile_data in collected_tasks.items(): + # Blend worker tiles synchronously in deterministic tile order. + def _sort_key(item): + global_idx, tile_data = item + batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image) + tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image) + return (tile_idx, batch_idx, global_idx) + + for global_idx, tile_data in sorted(collected_tasks.items(), key=_sort_key): # Skip tiles that don't have tensor data (already processed) if 'tensor' not in tile_data and 'image' not in tile_data: continue diff --git a/upscale/tile_ops.py b/upscale/tile_ops.py index 3a51f6b..f8b235f 100644 --- a/upscale/tile_ops.py +++ b/upscale/tile_ops.py @@ -1,10 +1,12 @@ import math, torch +from contextlib import nullcontext from PIL import Image, ImageFilter, ImageDraw from typing import List, Tuple import comfy.samplers, comfy.model_management from ..utils.logging import debug_log, log from ..utils.image import tensor_to_pil, pil_to_tensor from ..utils.usdu_utils import crop_cond, get_crop_region, expand_crop +from ..utils.crop_model_patch import crop_model_cond from .conditioning import clone_conditioning @@ -37,8 +39,8 @@ class TileOpsMixin: Mirrors ComfyUI_UltimateSDUpscale processing: - Build a mask with a white rectangle at the tile rect - Compute crop_region via get_crop_region(mask, padding) - - If force_uniform_tiles: shrink/expand to target size of - round_to_multiple(tile + padding) for each dimension + - If force_uniform_tiles: expand by crop/aspect ratio, then resize to + fixed processing size of round_to_multiple(tile + padding) - Else: target is ceil(crop_size/8)*8 per dimension - Extract the crop and resize to target tile_size Returns the resized tensor and crop origin/size for blending. @@ -51,11 +53,23 @@ class TileOpsMixin: draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255) x1, y1, x2, y2 = get_crop_region(mask, padding) - # Determine target tile size (processing size) + # Determine crop + processing size if force_uniform_tiles: - target_w = self.round_to_multiple(tile_width + padding, 8) - target_h = self.round_to_multiple(tile_height + padding, 8) - (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h) + process_w = self.round_to_multiple(tile_width + padding, 8) + process_h = self.round_to_multiple(tile_height + padding, 8) + crop_w = x2 - x1 + crop_h = y2 - y1 + crop_ratio = crop_w / crop_h if crop_h != 0 else 1.0 + process_ratio = process_w / process_h if process_h != 0 else 1.0 + if crop_ratio > process_ratio: + target_w = crop_w + target_h = round(crop_w / process_ratio) if process_ratio != 0 else crop_h + else: + target_w = round(crop_h * process_ratio) + target_h = crop_h + (x1, y1, x2, y2), _ = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h) + target_w = process_w + target_h = process_h else: crop_w = x2 - x1 crop_h = y2 - y1 @@ -96,11 +110,23 @@ class TileOpsMixin: draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255) x1, y1, x2, y2 = get_crop_region(mask, padding) - # Determine target processing size + # Determine crop + processing size if force_uniform_tiles: - target_w = self.round_to_multiple(tile_width + padding, 8) - target_h = self.round_to_multiple(tile_height + padding, 8) - (x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h) + process_w = self.round_to_multiple(tile_width + padding, 8) + process_h = self.round_to_multiple(tile_height + padding, 8) + crop_w = x2 - x1 + crop_h = y2 - y1 + crop_ratio = crop_w / crop_h if crop_h != 0 else 1.0 + process_ratio = process_w / process_h if process_h != 0 else 1.0 + if crop_ratio > process_ratio: + target_w = crop_w + target_h = round(crop_w / process_ratio) if process_ratio != 0 else crop_h + else: + target_w = round(crop_h * process_ratio) + target_h = crop_h + (x1, y1, x2, y2), _ = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h) + target_w = process_w + target_h = process_h else: crop_w = x2 - x1 crop_h = y2 - y1 @@ -185,9 +211,22 @@ class TileOpsMixin: # Encode to latent (always non-tiled, matching original node) latent = VAEEncode().encode(vae, clean_tensor)[0] - # Sample - samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, - positive_cropped, negative_cropped, latent, denoise=denoise)[0] + # Sample with model patch cropping parity (ControlNet patch hints) + if region is not None and image_size is not None: + model_ctx = crop_model_cond( + model, + region, + image_size, + image_size, + (clean_tensor.shape[2], clean_tensor.shape[1]), + ) + else: + model_ctx = nullcontext(model) + with model_ctx as model_for_sampling: + samples = common_ksampler( + model_for_sampling, seed, steps, cfg, sampler_name, scheduler, + positive_cropped, negative_cropped, latent, denoise=denoise + )[0] # Decode back to image if tiled_decode and tiled_vae_available: @@ -235,8 +274,11 @@ class TileOpsMixin: # Encode -> Sample -> Decode latent = VAEEncode().encode(vae, clean)[0] - samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, - positive_cropped, negative_cropped, latent, denoise=denoise)[0] + with crop_model_cond(model, region, image_size, image_size, tile_size) as model_for_sampling: + samples = common_ksampler( + model_for_sampling, seed, steps, cfg, sampler_name, scheduler, + positive_cropped, negative_cropped, latent, denoise=denoise + )[0] if tiled_decode and tiled_vae_available: image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0] else: @@ -385,8 +427,8 @@ class TileOpsMixin: upscaled_image[batch_idx:batch_idx+1], x, y, tile_width, tile_height, padding, force_uniform_tiles ) - # Process tile through SD with unique seed - image_seed = seed + batch_idx * 1000 + # Process tile through SD with the exact seed (USDU parity) + image_seed = seed processed_tile = self.process_tile(tile_tensor, model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode, batch_idx=batch_idx, diff --git a/utils/crop_model_patch.py b/utils/crop_model_patch.py new file mode 100644 index 0000000..17ba2a4 --- /dev/null +++ b/utils/crop_model_patch.py @@ -0,0 +1,114 @@ +from contextlib import contextmanager + +import torch + +from .logging import debug_log +from .usdu_utils import resize_region + + +@contextmanager +def crop_model_cond(model, crop_regions, init_size, canvas_size, tile_size, latent_crop=False): + """Clone model and crop compatible model patches for tile-local sampling.""" + try: + patched_model = model.clone() + except Exception: + # Fallback to original model when clone/patch access is unavailable. + yield model + return + + patches = ( + patched_model + .model_options + .get("transformer_options", {}) + .get("patches", {}) + ) + applied_croppers = {} + for _module, module_patches in patches.items(): + for patch in module_patches: + if id(patch) in applied_croppers: + continue + if type(patch).__name__ not in ("DiffSynthCnetPatch", "ZImageControlPatch"): + continue + try: + cropper = ModelPatchCropper(patch).crop(crop_regions, canvas_size, latent_crop) + applied_croppers[id(patch)] = cropper + except Exception as exc: + debug_log(f"crop_model_cond: patch crop skipped for {type(patch).__name__}: {exc}") + try: + yield patched_model + finally: + for cropper in applied_croppers.values(): + del cropper + + +class ModelPatchCropper: + """Stateful crop helper that restores model patch tensors on cleanup.""" + + def __init__(self, patch): + self.patch = patch + self.original_state = { + "image": patch.image.clone() if isinstance(patch.image, torch.Tensor) else patch.image, + "encoded_image": patch.encoded_image.clone() if isinstance(patch.encoded_image, torch.Tensor) else patch.encoded_image, + "encoded_image_size": patch.encoded_image_size, + } + self.patch_class = type(patch).__name__ + required_attrs = ( + "image", + "model_patch", + "vae", + "strength", + "encoded_image", + "encoded_image_size", + ) + missing = [attr for attr in required_attrs if not hasattr(patch, attr)] + if missing: + raise AttributeError( + f"{self.patch_class} missing required attrs: {', '.join(missing)}" + ) + + def __del__(self): + self.patch.image = self.original_state["image"] + self.patch.encoded_image = self.original_state["encoded_image"] + self.patch.encoded_image_size = self.original_state["encoded_image_size"] + + def crop(self, crop_regions, canvas_size, latent_crop=True): + patch = self.patch + + if not isinstance(crop_regions, list): + crop_regions = [crop_regions] + + image_size = (patch.image.shape[2], patch.image.shape[1]) # (W,H) + + cropped_images = [] + for crop_region in crop_regions: + resized_crop = resize_region(crop_region, canvas_size, image_size) + x1, y1, x2, y2 = resized_crop + cropped_image = patch.image[:, y1:y2, x1:x2, :] + cropped_images.append(cropped_image) + + concatenated_image = torch.cat(cropped_images, dim=0) + patch.image = concatenated_image + patch.encoded_image_size = ( + concatenated_image.shape[1], + concatenated_image.shape[2], + ) + + if latent_crop: + downscale_ratio = patch.vae.spacial_compression_encode() + cropped_latents = [] + for crop_region in crop_regions: + resized_crop = resize_region(crop_region, canvas_size, image_size) + x1, y1, x2, y2 = tuple(x // downscale_ratio for x in resized_crop) + cropped_latent = patch.encoded_image[:, :, y1:y2, x1:x2] + cropped_latents.append(cropped_latent) + patch.encoded_image = torch.cat(cropped_latents, dim=0) + else: + patch.__init__( + patch.model_patch, + patch.vae, + concatenated_image, + patch.strength, + inpaint_image=patch.inpaint_image, + mask=patch.mask, + ) + return self diff --git a/utils/usdu_utils.py b/utils/usdu_utils.py index 492033f..39b1cd3 100644 --- a/utils/usdu_utils.py +++ b/utils/usdu_utils.py @@ -466,9 +466,13 @@ def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, x1_px, y1_px, x2_px, y2_px = region new_latents = [] - for t in latents: # (B,C,H_lat_in,W_lat_in) + for t in latents: # (B,C,H_lat_in,W_lat_in) or (B,C,1,H_lat_in,W_lat_in) + has_5d = False + if t.ndim == 5: + has_5d = True + t = t.squeeze(2) if t.ndim != 4: - raise ValueError(f"expected BCHW, got {t.shape}") + raise ValueError(f"expected BCHW or BC1HW, got {t.shape}") # 1. Resize to canvas resolution in latent units only if needed if t.shape[-2:] != (H_can_lat, W_can_lat): @@ -491,6 +495,8 @@ def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, mode="bilinear", align_corners=False) + if has_5d: + cropped = cropped.unsqueeze(2) new_latents.append(cropped) cond_dict["reference_latents"] = new_latents @@ -508,4 +514,4 @@ def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0) 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) cropped.append(n) - return cropped \ No newline at end of file + return cropped