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
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,