nodes/robertvoy-ComfyUI-Distributed
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

USDU parity: progressive tiles, deterministic blend, crop/seed/model fixes

Browse Source
This commit is contained in:
Robert Wojciechowski
2026-02-28 04:53:20 +00:00
parent 487f5b9c85
commit e4c025e06e
7 changed files with 240 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_static_mode.py
+7
View File
@@ -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)
upscale/modes/dynamic.py
+10 -4
View File
@@ -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,
upscale/modes/single_gpu.py
+6 -1
View File
@@ -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
upscale/modes/static.py
+35 -6
View File
@@ -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
upscale/tile_ops.py
+59 -17
View File
@@ -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,
utils/crop_model_patch.py
+114
View File
@@ -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
utils/usdu_utils.py
+9 -3
View File
@@ -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
return cropped