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

Refactor distributed upscale node into mixin modules

Browse Source
This commit is contained in:
Robert Wojciechowski
2026-02-23 02:12:13 +00:00
parent 30e55dc7df
commit d8fc6751f2
10 changed files with 1749 additions and 1714 deletions
Download Patch File Download Diff File Expand all files Collapse all files
distributed_upscale.py
+22 -1714
View File
File diff suppressed because it is too large Load Diff
upscale/__init__.py
View File
upscale/job_state.py
+113
View File
@@ -0,0 +1,113 @@
import asyncio, time
import server
from ..utils.logging import debug_log
from ..utils.usdu_managment import (
ensure_tile_jobs_initialized,
_check_and_requeue_timed_out_workers as _requeue_usdu,
JOB_COMPLETED_TASKS, JOB_PENDING_TASKS,
)
class JobStateMixin:
async def _get_all_completed_tasks(self, multi_job_id):
"""Helper to retrieve all completed tasks from the job data."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if job_data and JOB_COMPLETED_TASKS in job_data:
return dict(job_data[JOB_COMPLETED_TASKS]) # Return a copy
return {}
async def _get_next_image_index(self, multi_job_id):
"""Get next image index from pending queue for master."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if not job_data or 'pending_images' not in job_data:
return None
try:
image_idx = await asyncio.wait_for(job_data['pending_images'].get(), timeout=1.0)
return image_idx
except asyncio.TimeoutError:
return None
async def _get_next_tile_index(self, multi_job_id):
"""Get next tile index from pending queue for master in static mode."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if not job_data or JOB_PENDING_TASKS not in job_data:
return None
try:
tile_idx = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=0.1)
return tile_idx
except asyncio.TimeoutError:
return None
async def _get_total_completed_count(self, multi_job_id):
"""Get total count of all completed images (master + workers)."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if job_data and 'completed_images' in job_data:
return len(job_data['completed_images'])
return 0
async def _get_all_completed_images(self, multi_job_id):
"""Get all completed images."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if job_data and 'completed_images' in job_data:
return job_data['completed_images'].copy()
return {}
async def _get_pending_count(self, multi_job_id):
"""Get count of pending images in the queue."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if job_data and 'pending_images' in job_data:
return job_data['pending_images'].qsize()
return 0
async def _drain_worker_results_queue(self, multi_job_id):
"""Drain pending worker results from queue and update completed_images. Returns count of drained images."""
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if not job_data or 'queue' not in job_data or 'completed_images' not in job_data:
return 0
q = job_data['queue']
completed_images = job_data['completed_images']
collected = 0
while not q.empty():
try:
result = await asyncio.wait_for(q.get(), timeout=0.1)
worker_id = result['worker_id']
is_last = result.get('is_last', False)
if 'image_idx' in result and 'image' in result:
image_idx = result['image_idx']
image_pil = result['image']
if image_idx not in completed_images:
completed_images[image_idx] = image_pil
collected += 1
debug_log(f"Drained image {image_idx} from worker {worker_id}")
if is_last:
# Optional: track worker completion if needed
pass
except asyncio.TimeoutError:
break # No more immediately available
if collected > 0:
debug_log(f"Drained {collected} worker images during retry")
return collected
async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
"""Check for timed out workers and requeue their assigned images. Returns count of requeued images."""
# Use the original function from usdu_managment.py
return await _requeue_usdu(multi_job_id, batch_size)
upscale/modes/__init__.py
View File
upscale/modes/dynamic.py
+340
View File
@@ -0,0 +1,340 @@
import asyncio, torch
from PIL import Image
import comfy.model_management
from ...utils.logging import debug_log, log
from ...utils.image import tensor_to_pil, pil_to_tensor
from ...utils.async_helpers import run_async_in_server_loop
from ...utils.config import get_worker_timeout_seconds
from ...utils.constants import TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT
from ...utils.usdu_managment import (
ensure_tile_jobs_initialized, init_dynamic_job,
clone_conditioning, _send_heartbeat_to_master,
)
class DynamicModeMixin:
def process_master_dynamic(self, upscaled_image, model, positive, negative, vae,
seed, steps, cfg, sampler_name, scheduler, denoise,
tile_width, tile_height, padding, mask_blur,
force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers):
"""Dynamic mode for large batches - assigns whole images to workers dynamically, including master."""
# Get batch size and dimensions
batch_size, height, width, _ = upscaled_image.shape
num_workers = len(enabled_workers)
log(f"USDU Dist: Image queue distribution | Batch {batch_size} | Workers {num_workers} | Canvas {width}x{height} | Tile {tile_width}x{tile_height}")
# No fixed share - all images are dynamic
all_indices = list(range(batch_size))
debug_log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.")
# Calculate tiles for processing
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
# Initialize job queue for communication
try:
run_async_in_server_loop(
init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices),
timeout=2.0
)
except Exception as e:
debug_log(f"UltimateSDUpscale Master - Queue initialization error: {e}")
raise RuntimeError(f"Failed to initialize dynamic mode queue: {e}")
# Convert batch to PIL list
result_images = [tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB').copy() for b in range(batch_size)]
# Process images dynamically with master participating
prompt_server = ensure_tile_jobs_initialized()
processed_count = 0
consecutive_retries = 0
max_consecutive_retries = 10
# Process loop - master pulls from queue and processes synchronously
while processed_count < batch_size:
# Try to get an image to process
image_idx = run_async_in_server_loop(
self._get_next_image_index(multi_job_id),
timeout=5.0 # Short timeout to allow frequent checks
)
if image_idx is not None:
# Reset retry counter and process locally
consecutive_retries = 0
debug_log(f"Master processing image {image_idx} dynamically")
processed_count += 1
# Process locally
single_tensor = upscaled_image[image_idx:image_idx+1]
local_image = result_images[image_idx]
image_seed = seed + image_idx * 1000
# Pre-slice conditioning once per image (not per tile)
positive_sliced = clone_conditioning(positive)
negative_sliced = clone_conditioning(negative)
for cond_list in [positive_sliced, negative_sliced]:
for i in range(len(cond_list)):
emb, cond_dict = cond_list[i]
if emb.shape[0] > 1:
cond_list[i][0] = emb[image_idx:image_idx+1]
if 'control' in cond_dict:
control = cond_dict['control']
while control is not None:
hint = control.cond_hint_original
if hint.shape[0] > 1:
control.cond_hint_original = hint[image_idx:image_idx+1]
control = control.previous_controlnet
if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1:
cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1]
for tile_idx, pos in enumerate(all_tiles):
local_image = self._process_and_blend_tile(
tile_idx, pos, single_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,
tiled_decode, batch_idx=image_idx
)
# Yield after each tile to minimize worker downtime
run_async_in_server_loop(self._async_yield(), timeout=0.1)
# Note: No per-tile drain here – that's what makes this "per-image"
result_images[image_idx] = local_image
# Mark as completed
run_async_in_server_loop(
self._mark_image_completed(multi_job_id, image_idx, local_image),
timeout=5.0
)
# NEW: Drain after the full image is marked complete (catches workers who finished during master's processing)
drained_count = run_async_in_server_loop(
self._drain_worker_results_queue(multi_job_id),
timeout=5.0
)
if drained_count > 0:
debug_log(f"Drained {drained_count} worker images after master's image {image_idx}")
# NEW: Log overall progress (includes master's image + any drained workers)
completed_now = run_async_in_server_loop(
self._get_total_completed_count(multi_job_id),
timeout=1.0
)
log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
# Yield to allow workers to get new images after completing one
run_async_in_server_loop(self._async_yield(), timeout=0.1)
else:
# Queue empty: collect any queued worker results to update progress
drained_count = run_async_in_server_loop(
self._drain_worker_results_queue(multi_job_id),
timeout=5.0
)
run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after drain
# Check for timed out workers and requeue their images
requeued_count = run_async_in_server_loop(
self._check_and_requeue_timed_out_workers(multi_job_id, batch_size),
timeout=5.0
)
run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after requeue
if requeued_count > 0:
log(f"Requeued {requeued_count} images from timed out workers")
consecutive_retries = 0 # Reset since we have work to do
continue
# Now check total completed (includes newly collected)
completed_now = run_async_in_server_loop(
self._get_total_completed_count(multi_job_id),
timeout=1.0
)
log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
if completed_now >= batch_size:
break
run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield before pending check
# Check if there are pending images in the queue (could be requeued)
pending_count = run_async_in_server_loop(
self._get_pending_count(multi_job_id),
timeout=1.0
)
if pending_count > 0:
consecutive_retries = 0 # Reset retries since there's work to do
continue
consecutive_retries += 1
if consecutive_retries >= max_consecutive_retries:
log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.")
break # Force exit to collection phase
debug_log("Waiting for workers")
# Use async sleep to allow event loop to process worker requests
run_async_in_server_loop(asyncio.sleep(2), timeout=3.0)
debug_log(f"Master processed {processed_count} images locally")
# Get all completed images to check what needs to be collected
all_completed = run_async_in_server_loop(
self._get_all_completed_images(multi_job_id),
timeout=5.0
)
# Calculate how many we still need to collect
remaining_to_collect = batch_size - len(all_completed)
if remaining_to_collect > 0:
debug_log(f"Waiting for {remaining_to_collect} more images from workers")
# Use the unified worker timeout for the collection phase
collection_timeout = float(get_worker_timeout_seconds())
collected_images = run_async_in_server_loop(
self._async_collect_dynamic_images(multi_job_id, remaining_to_collect, num_workers, batch_size, processed_count),
timeout=collection_timeout
)
# Merge collected with already completed
all_completed.update(collected_images)
# Update result images with all completed images
for idx, processed_img in all_completed.items():
if idx < batch_size:
result_images[idx] = processed_img
# Convert back to tensor
result_tensor = torch.cat([pil_to_tensor(img) for img in result_images], dim=0) if batch_size > 1 else pil_to_tensor(result_images[0])
if upscaled_image.is_cuda:
result_tensor = result_tensor.cuda()
debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
log(f"Completed processing all {batch_size} images")
return (result_tensor,)
def process_worker_dynamic(self, upscaled_image, model, positive, negative, vae,
seed, steps, cfg, sampler_name, scheduler, denoise,
tile_width, tile_height, padding, mask_blur,
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
worker_id, enabled_worker_ids, dynamic_threshold):
"""Worker processing in dynamic mode - processes whole images."""
# Round tile dimensions
tile_width = self.round_to_multiple(tile_width)
tile_height = self.round_to_multiple(tile_height)
# Get dimensions and tile grid
batch_size, height, width, _ = upscaled_image.shape
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
log(f"USDU Dist Worker[{worker_id[:8]}]: Processing image queue | Batch {batch_size}")
# Keep track of processed images for is_last detection
processed_count = 0
# Poll for job readiness to avoid races during master init
max_poll_attempts = 20 # ~20s at 1s sleep
for attempt in range(max_poll_attempts):
ready = run_async_in_server_loop(
self._check_job_status(multi_job_id, master_url),
timeout=5.0
)
if ready:
debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts")
break
time.sleep(1.0) # Poll every 1s
else:
log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
return (upscaled_image,)
# Loop to request and process images
while True:
# Request an image to process
image_idx, estimated_remaining = run_async_in_server_loop(
self._request_image_from_master(multi_job_id, master_url, worker_id),
timeout=TILE_WAIT_TIMEOUT
)
if image_idx is None:
debug_log(f"USDU Dist Worker - No more images to process")
break
debug_log(f"Worker[{worker_id[:8]}] - Assigned image {image_idx}")
processed_count += 1
# Determine if this should be marked as last for this worker
is_last_for_worker = (estimated_remaining == 0)
# Extract single image tensor
single_tensor = upscaled_image[image_idx:image_idx+1]
# Convert to PIL for processing
local_image = tensor_to_pil(single_tensor, 0).copy()
# Process all tiles for this image
image_seed = seed + image_idx * 1000
# Pre-slice conditioning once per image (not per tile)
positive_sliced = clone_conditioning(positive)
negative_sliced = clone_conditioning(negative)
for cond_list in [positive_sliced, negative_sliced]:
for i in range(len(cond_list)):
emb, cond_dict = cond_list[i]
if emb.shape[0] > 1:
cond_list[i][0] = emb[image_idx:image_idx+1]
if 'control' in cond_dict:
control = cond_dict['control']
while control is not None:
hint = control.cond_hint_original
if hint.shape[0] > 1:
control.cond_hint_original = hint[image_idx:image_idx+1]
control = control.previous_controlnet
if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1:
cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1]
for tile_idx, pos in enumerate(all_tiles):
local_image = self._process_and_blend_tile(
tile_idx, pos, single_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,
tiled_decode, batch_idx=image_idx
)
run_async_in_server_loop(
_send_heartbeat_to_master(multi_job_id, master_url, worker_id),
timeout=5.0
)
# Send processed image back to master
try:
# Use the estimated remaining to determine if this is the last image
is_last = is_last_for_worker
run_async_in_server_loop(
self._send_full_image_to_master(local_image, image_idx, multi_job_id,
master_url, worker_id, is_last),
timeout=TILE_SEND_TIMEOUT
)
# Send heartbeat after processing
run_async_in_server_loop(
_send_heartbeat_to_master(multi_job_id, master_url, worker_id),
timeout=5.0
)
if is_last:
break
except Exception as e:
log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending image {image_idx}: {e}")
# Continue processing other images
# Send final is_last signal
debug_log(f"Worker[{worker_id[:8]}] processed {processed_count} images, sending completion signal")
try:
run_async_in_server_loop(
self._send_worker_complete_signal(multi_job_id, master_url, worker_id),
timeout=TILE_SEND_TIMEOUT
)
except Exception as e:
log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending completion signal: {e}")
return (upscaled_image,)
upscale/modes/single_gpu.py
+67
View File
@@ -0,0 +1,67 @@
import math, torch
from PIL import Image
from ...utils.logging import debug_log, log
from ...utils.image import tensor_to_pil, pil_to_tensor
class SingleGpuModeMixin:
def process_single_gpu(self, upscaled_image, model, positive, negative, vae,
seed, steps, cfg, sampler_name, scheduler, denoise,
tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode):
"""Process all tiles on a single GPU (no distribution), batching per tile like USDU."""
# Round tile dimensions
tile_width = self.round_to_multiple(tile_width)
tile_height = self.round_to_multiple(tile_height)
# Get image dimensions and batch size
batch_size, height, width, _ = upscaled_image.shape
# Calculate all tiles
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
rows = math.ceil(height / tile_height)
cols = math.ceil(width / tile_width)
log(
f"USDU Dist: Single GPU | Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({len(all_tiles)} tiles/image) | Batch {batch_size}"
)
# Prepare result images list
result_images = []
for b in range(batch_size):
image_pil = tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB')
result_images.append(image_pil.copy())
# Precompute tile masks once
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))
# Process tiles batched across images
for tile_idx, (tx, ty) in enumerate(all_tiles):
# 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
)
# Process batch
region = (x1, y1, x1 + ew, y1 + eh)
processed_batch = self.process_tiles_batch(tile_batch, model, positive, negative, vae,
seed, steps, cfg, sampler_name, scheduler, denoise,
tiled_decode, region, (width, height))
# Blend results back into each image using cached mask
tile_mask = tile_masks[tile_idx]
for b in range(batch_size):
tile_pil = tensor_to_pil(processed_batch, b)
# Resize back to extracted size
if tile_pil.size != (ew, eh):
tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding)
# Convert back to tensor
result_tensors = [pil_to_tensor(img) for img in result_images]
result_tensor = torch.cat(result_tensors, dim=0)
if upscaled_image.is_cuda:
result_tensor = result_tensor.cuda()
return (result_tensor,)
upscale/modes/static.py
+367
View File
@@ -0,0 +1,367 @@
import asyncio, time, torch
from PIL import Image
import comfy.model_management
from ...utils.logging import debug_log, log
from ...utils.image import tensor_to_pil, pil_to_tensor
from ...utils.async_helpers import run_async_in_server_loop
from ...utils.config import get_worker_timeout_seconds
from ...utils.constants import TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT, MAX_BATCH
from ...utils.usdu_managment import (
ensure_tile_jobs_initialized, init_static_job_batched,
_mark_task_completed, _cleanup_job, _drain_results_queue,
_send_heartbeat_to_master,
JOB_PENDING_TASKS, JOB_WORKER_STATUS,
)
class StaticModeMixin:
def _process_worker_static_sync(self, upscaled_image, model, positive, negative, vae,
seed, steps, cfg, sampler_name, scheduler, denoise,
tile_width, tile_height, padding, mask_blur,
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
worker_id, enabled_workers):
"""Worker static mode processing with optional dynamic queue pulling."""
# Round tile dimensions
tile_width = self.round_to_multiple(tile_width)
tile_height = self.round_to_multiple(tile_height)
# Get dimensions and calculate tiles
_, height, width, _ = upscaled_image.shape
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
num_tiles_per_image = len(all_tiles)
batch_size = upscaled_image.shape[0]
total_tiles = batch_size * num_tiles_per_image
processed_tiles = []
sliced_conditioning_cache = {}
# 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}")
processed_count = 0
# Poll for job readiness
max_poll_attempts = 20
for attempt in range(max_poll_attempts):
ready = run_async_in_server_loop(
self._check_job_status(multi_job_id, master_url),
timeout=5.0
)
if ready:
debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts")
break
time.sleep(1.0)
else:
log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
return (upscaled_image,)
# Main processing loop - pull tile ids from queue
while True:
# Request a tile to process
tile_idx, estimated_remaining, batched_static = run_async_in_server_loop(
self._request_tile_from_master(multi_job_id, master_url, worker_id),
timeout=TILE_WAIT_TIMEOUT
)
if tile_idx is None:
debug_log(f"Worker[{worker_id[:8]}] - No more tiles to process")
break
# Always batched-per-tile in static mode
debug_log(f"Worker[{worker_id[:8]}] - Assigned tile_id {tile_idx}")
processed_count += batch_size
tile_id = tile_idx
tx, ty = all_tiles[tile_id]
# Extract tile for entire batch
tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding(
upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles
)
# Process batch
region = (x1, y1, x1 + ew, y1 + eh)
processed_batch = self.process_tiles_batch(
tile_batch, model, positive, negative, vae,
seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode,
region, (width, height)
)
# Queue results
for b in range(batch_size):
processed_tiles.append({
'tile': processed_batch[b:b+1],
'tile_idx': tile_id,
'x': x1,
'y': y1,
'extracted_width': ew,
'extracted_height': eh,
'padding': padding,
'batch_idx': b,
'global_idx': b * num_tiles_per_image + tile_id
})
# Send heartbeat
try:
run_async_in_server_loop(
_send_heartbeat_to_master(multi_job_id, master_url, worker_id),
timeout=5.0
)
except Exception as e:
debug_log(f"Worker[{worker_id[:8]}] heartbeat failed: {e}")
# Send tiles in batches within loop
if len(processed_tiles) >= MAX_BATCH:
run_async_in_server_loop(
self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id),
timeout=TILE_SEND_TIMEOUT
)
processed_tiles = []
# Send any remaining tiles
if processed_tiles:
run_async_in_server_loop(
self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id),
timeout=TILE_SEND_TIMEOUT
)
debug_log(f"Worker {worker_id} completed all assigned and requeued tiles")
return (upscaled_image,)
async def _async_collect_and_monitor_static(self, multi_job_id, total_tiles, expected_total):
"""Async helper for collection and monitoring in static mode.
Returns collected tasks dict. Caller should check if all tasks are complete."""
last_progress_log = time.time()
progress_interval = 5.0
last_heartbeat_check = time.time()
last_completed_count = 0
while True:
# Check for user interruption
if comfy.model_management.processing_interrupted():
log("Processing interrupted by user")
raise comfy.model_management.InterruptProcessingException()
# Drain any pending results
collected_count = await _drain_results_queue(multi_job_id)
# Check and requeue timed-out workers periodically
current_time = time.time()
if current_time - last_heartbeat_check >= 10.0:
requeued_count = await self._check_and_requeue_timed_out_workers(multi_job_id, expected_total)
if requeued_count > 0:
log(f"Requeued {requeued_count} tasks from timed-out workers")
last_heartbeat_check = current_time
# Get current completion count
completed_count = await _get_completed_count(multi_job_id)
# Progress logging
if current_time - last_progress_log >= progress_interval:
log(f"Progress: {completed_count}/{expected_total} tasks completed")
last_progress_log = current_time
# Check if all tasks are completed
if completed_count >= expected_total:
debug_log(f"All {expected_total} tasks completed")
break
# If no active workers remain and there are pending tasks, return for local processing
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if job_data:
pending_queue = job_data.get(JOB_PENDING_TASKS)
active_workers = list(job_data.get(JOB_WORKER_STATUS, {}).keys())
if pending_queue and not pending_queue.empty() and len(active_workers) == 0:
log(f"No active workers remaining with {expected_total - completed_count} tasks pending. Returning for local processing.")
break
# Wait a bit before next check
await asyncio.sleep(0.1)
# Get all completed tasks for return
return await self._get_all_completed_tasks(multi_job_id)
def _process_master_static_sync(self, upscaled_image, model, positive, negative, vae,
seed, steps, cfg, sampler_name, scheduler, denoise,
tile_width, tile_height, padding, mask_blur,
force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers,
all_tiles, num_tiles_per_image):
"""Static mode master processing with optional dynamic queue pulling."""
batch_size = upscaled_image.shape[0]
_, height, width, _ = upscaled_image.shape
total_tiles = batch_size * num_tiles_per_image
# Convert batch to PIL list for processing
result_images = []
for b in range(batch_size):
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(
init_static_job_batched(multi_job_id, batch_size, num_tiles_per_image, enabled_workers),
timeout=10.0
)
debug_log(
f"Initialized tile-id queue with {num_tiles_per_image} ids for batch {batch_size}"
)
# Precompute masks for all tile positions to avoid repeated Gaussian blur work during blending
tile_masks = []
for idx, (tx, ty) in enumerate(all_tiles):
tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
processed_count = 0
consecutive_no_tile = 0
max_consecutive_no_tile = 2
while processed_count < total_tiles:
comfy.model_management.throw_exception_if_processing_interrupted()
tile_idx = run_async_in_server_loop(
self._get_next_tile_index(multi_job_id),
timeout=5.0
)
if tile_idx is not None:
consecutive_no_tile = 0
processed_count += batch_size
tile_id = tile_idx
tx, ty = all_tiles[tile_id]
tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding(
upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles
)
region = (x1, y1, x1 + ew, y1 + eh)
processed_batch = self.process_tiles_batch(
tile_batch, model, positive, negative, vae,
seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode,
region, (width, height)
)
tile_mask = tile_masks[tile_id]
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)
result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding)
global_idx = b * num_tiles_per_image + tile_id
run_async_in_server_loop(
_mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}),
timeout=5.0
)
log(f"USDU Dist: Tiles progress {processed_count}/{total_tiles} (tile {tile_id})")
else:
consecutive_no_tile += 1
if consecutive_no_tile >= max_consecutive_no_tile:
debug_log(f"Master processed {processed_count} tiles, moving to collection phase")
break
time.sleep(0.1)
master_processed_count = processed_count
# Continue processing any remaining tiles while collecting worker results
remaining_tiles = total_tiles - master_processed_count
if remaining_tiles > 0:
debug_log(f"Master waiting for {remaining_tiles} tiles from workers")
# Collect worker results using async operations
try:
# Wait until either all tasks are collected or there are no active workers left
collected_tasks = run_async_in_server_loop(
self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles),
timeout=None
)
except comfy.model_management.InterruptProcessingException:
# Clean up job on interruption
run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0)
raise
# Check if we need to process any remaining tasks locally after collection
completed_count = len(collected_tasks)
if completed_count < total_tiles:
log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures")
# Process any remaining pending tasks (batched-per-tile)
while True:
# Check for user interruption
comfy.model_management.throw_exception_if_processing_interrupted()
# Get next tile_id from pending queue
tile_id = run_async_in_server_loop(
self._get_next_tile_index(multi_job_id),
timeout=5.0
)
if tile_id is None:
break
# Extract batched tile and process across available batch
tx, ty = all_tiles[tile_id]
tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding(
upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles
)
region = (x1, y1, x1 + ew, y1 + eh)
processed_batch = self.process_tiles_batch(
tile_batch, model, positive, negative, vae,
seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode,
region, (width, height)
)
tile_mask = tile_masks[tile_id]
out_bs = processed_batch.shape[0] if hasattr(processed_batch, 'shape') else batch_size
for b in range(min(batch_size, out_bs)):
tile_pil = tensor_to_pil(processed_batch, b)
if tile_pil.size != (ew, eh):
tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding)
global_idx = b * num_tiles_per_image + tile_id
# Mark as completed so the collector state is consistent
run_async_in_server_loop(
_mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}),
timeout=5.0
)
else:
# Master processed all tiles
collected_tasks = run_async_in_server_loop(
self._get_all_completed_tasks(multi_job_id),
timeout=5.0
)
# Blend worker tiles synchronously
for global_idx, tile_data in collected_tasks.items():
# Skip tiles that don't have tensor data (already processed)
if 'tensor' not in tile_data and 'image' not in tile_data:
continue
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)
if batch_idx >= batch_size:
continue
# Blend tile synchronously
x = tile_data.get('x', 0)
y = tile_data.get('y', 0)
# Prefer PIL image if present to avoid reconversion
if 'image' in tile_data:
tile_pil = tile_data['image']
else:
tile_tensor = tile_data['tensor']
tile_pil = tensor_to_pil(tile_tensor, 0)
orig_x, orig_y = all_tiles[tile_idx]
tile_mask = tile_masks[tile_idx]
extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding)
extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding)
result_images[batch_idx] = self.blend_tile(result_images[batch_idx], tile_pil,
x, y, (extracted_width, extracted_height), tile_mask, padding)
try:
# Convert back to tensor
if batch_size == 1:
result_tensor = pil_to_tensor(result_images[0])
else:
result_tensors = [pil_to_tensor(img) for img in result_images]
result_tensor = torch.cat(result_tensors, dim=0)
if upscaled_image.is_cuda:
result_tensor = result_tensor.cuda()
log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
return (result_tensor,)
finally:
# Cleanup (async operation) - always execute
run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0)
upscale/result_collector.py
+165
View File
@@ -0,0 +1,165 @@
import asyncio, time
import comfy.model_management
import server
from ..utils.logging import debug_log, log
from ..utils.config import get_worker_timeout_seconds
from ..utils.usdu_managment import (
ensure_tile_jobs_initialized, _mark_task_completed,
_check_and_requeue_timed_out_workers,
JOB_COMPLETED_TASKS, JOB_WORKER_STATUS, JOB_PENDING_TASKS,
)
class ResultCollectorMixin:
async def _async_collect_results(self, multi_job_id, num_workers, mode='static',
remaining_to_collect=None, batch_size=None):
"""Unified async helper to collect results from workers (tiles or images)."""
# Get the already initialized queue
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id not in prompt_server.distributed_pending_tile_jobs:
raise RuntimeError(f"Job queue not initialized for {multi_job_id}")
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if not isinstance(job_data, dict) or 'mode' not in job_data:
raise RuntimeError("Invalid job data structure")
if job_data['mode'] != mode:
raise RuntimeError(f"Mode mismatch: expected {mode}, got {job_data['mode']}")
q = job_data['queue']
# For dynamic mode, get reference to completed_images
if mode == 'dynamic':
completed_images = job_data['completed_images']
# Calculate expected count for logging
expected_count = remaining_to_collect or batch_size
else:
# Calculate expected count from job data for static mode
expected_count = job_data.get('total_items', 0) - len(job_data.get('completed_items', set()))
item_type = "images" if mode == 'dynamic' else "tiles"
debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers")
collected_results = {}
workers_done = set()
# Unify collector/upscaler wait behavior with the UI worker timeout
timeout = float(get_worker_timeout_seconds())
last_heartbeat_check = time.time()
collected_count = 0
while len(workers_done) < num_workers:
# Check for user interruption
if comfy.model_management.processing_interrupted():
log("Processing interrupted by user")
raise comfy.model_management.InterruptProcessingException()
# For dynamic mode with remaining_to_collect, check if we've collected enough
if mode == 'dynamic' and remaining_to_collect and collected_count >= remaining_to_collect:
break
try:
# Shorter poll for dynamic mode, but never exceed the configured timeout
wait_timeout = (min(10.0, timeout) if mode == 'dynamic' else timeout)
result = await asyncio.wait_for(q.get(), timeout=wait_timeout)
worker_id = result['worker_id']
is_last = result.get('is_last', False)
if mode == 'static':
# Handle tiles
tiles = result.get('tiles', [])
if tiles:
# Batch mode
debug_log(f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker '{worker_id}' (is_last={is_last})")
for tile_data in tiles:
# Validate required fields
if 'batch_idx' not in tile_data:
log(f"UltimateSDUpscale Master - Missing batch_idx in tile data, skipping")
continue
tile_idx = tile_data['tile_idx']
# Use global_idx as key if available (for batch processing)
key = tile_data.get('global_idx', tile_idx)
# Store the full tile data including metadata; prefer PIL image if present
entry = {
'tile_idx': tile_idx,
'x': tile_data['x'],
'y': tile_data['y'],
'extracted_width': tile_data['extracted_width'],
'extracted_height': tile_data['extracted_height'],
'padding': tile_data['padding'],
'worker_id': worker_id,
'batch_idx': tile_data.get('batch_idx', 0),
'global_idx': tile_data.get('global_idx', tile_idx)
}
if 'image' in tile_data:
entry['image'] = tile_data['image']
elif 'tensor' in tile_data:
entry['tensor'] = tile_data['tensor']
collected_results[key] = entry
else:
# Single tile mode (backward compat)
tile_idx = result['tile_idx']
collected_results[tile_idx] = result
debug_log(f"UltimateSDUpscale Master - Received single tile {tile_idx} from worker '{worker_id}' (is_last={is_last})")
elif mode == 'dynamic':
# Handle full images
if 'image_idx' in result and 'image' in result:
image_idx = result['image_idx']
image_pil = result['image']
completed_images[image_idx] = image_pil
collected_results[image_idx] = image_pil
collected_count += 1
debug_log(f"UltimateSDUpscale Master - Received image {image_idx} from worker {worker_id}")
if is_last:
workers_done.add(worker_id)
debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed")
except asyncio.TimeoutError:
if mode == 'dynamic':
# Check for worker timeouts periodically
current_time = time.time()
if current_time - last_heartbeat_check >= 10.0:
# Use the class method to check and requeue
requeued = await self._check_and_requeue_timed_out_workers(multi_job_id, batch_size)
if requeued > 0:
log(f"UltimateSDUpscale Master - Requeued {requeued} images from timed out workers")
last_heartbeat_check = current_time
# Check if we've been waiting too long overall
if current_time - last_heartbeat_check > timeout:
log(f"UltimateSDUpscale Master - Overall timeout waiting for images")
break
else:
log(f"UltimateSDUpscale Master - Timeout waiting for {item_type}")
break
debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_results)} {item_type} from {len(workers_done)} workers")
# Clean up job queue
async with prompt_server.distributed_tile_jobs_lock:
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
del prompt_server.distributed_pending_tile_jobs[multi_job_id]
return collected_results if mode == 'static' else completed_images
async def _async_collect_worker_tiles(self, multi_job_id, num_workers):
"""Async helper to collect tiles from workers."""
return await self._async_collect_results(multi_job_id, num_workers, mode='static')
async def _mark_image_completed(self, multi_job_id, image_idx, image_pil):
"""Mark an image as completed in the job data."""
# Mark the image as completed with the image data
await _mark_task_completed(multi_job_id, image_idx, {'image': image_pil})
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if job_data and 'completed_images' in job_data:
job_data['completed_images'][image_idx] = image_pil
async def _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count):
"""Collect remaining processed images from workers."""
return await self._async_collect_results(multi_job_id, num_workers, mode='dynamic',
remaining_to_collect=remaining_to_collect,
batch_size=batch_size)
upscale/tile_ops.py
+404
View File
@@ -0,0 +1,404 @@
import math, torch
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.usdu_managment import clone_conditioning
class TileOpsMixin:
def round_to_multiple(self, value: int, multiple: int = 8) -> int:
"""Round value to nearest multiple."""
return round(value / multiple) * multiple
def calculate_tiles(self, image_width: int, image_height: int,
tile_width: int, tile_height: int, force_uniform_tiles: bool = True) -> List[Tuple[int, int]]:
"""Calculate tile positions to match Ultimate SD Upscale.
Positions are a simple grid starting at (0,0) with steps of
`tile_width` and `tile_height`, using ceil(rows/cols) to cover edges.
Uniform vs non-uniform affects only crop/resize, not positions.
"""
rows = math.ceil(image_height / tile_height)
cols = math.ceil(image_width / tile_width)
tiles: List[Tuple[int, int]] = []
for yi in range(rows):
for xi in range(cols):
tiles.append((xi * tile_width, yi * tile_height))
return tiles
def extract_tile_with_padding(self, image: torch.Tensor, x: int, y: int,
tile_width: int, tile_height: int, padding: int,
force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]:
"""Extract a tile region and resize to match USDU cropping logic.
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
- 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.
"""
_, h, w, _ = image.shape
# Create mask and compute initial padded crop region
mask = Image.new('L', (w, h), 0)
draw = ImageDraw.Draw(mask)
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)
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)
else:
crop_w = x2 - x1
crop_h = y2 - y1
target_w = max(8, math.ceil(crop_w / 8) * 8)
target_h = max(8, math.ceil(crop_h / 8) * 8)
(x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
# Actual extracted size before resizing (for blending)
extracted_width = x2 - x1
extracted_height = y2 - y1
# Extract tile and resize to processing size
tile = image[:, y1:y2, x1:x2, :]
tile_pil = tensor_to_pil(tile, 0)
if tile_pil.size != (target_w, target_h):
tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS)
tile_tensor = pil_to_tensor(tile_pil)
if image.is_cuda:
tile_tensor = tile_tensor.cuda()
return tile_tensor, x1, y1, extracted_width, extracted_height
def extract_batch_tile_with_padding(self, images: torch.Tensor, x: int, y: int,
tile_width: int, tile_height: int, padding: int,
force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]:
"""Extract a tile region for the entire batch and resize to USDU logic.
- Computes a single crop region from a mask at (x,y,w,h) with padding
- force_uniform_tiles controls target processing size logic
- Returns a batched tensor [B,H',W',C] and crop origin/size for blending
"""
batch, h, w, _ = images.shape
# Create mask and compute initial padded crop region (same for all images)
mask = Image.new('L', (w, h), 0)
draw = ImageDraw.Draw(mask)
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
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)
else:
crop_w = x2 - x1
crop_h = y2 - y1
target_w = max(8, math.ceil(crop_w / 8) * 8)
target_h = max(8, math.ceil(crop_h / 8) * 8)
(x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
extracted_width = x2 - x1
extracted_height = y2 - y1
# Slice batch region
tiles = images[:, y1:y2, x1:x2, :]
# Resize each tile to target size
resized_tiles = []
for i in range(batch):
tile_pil = tensor_to_pil(tiles, i)
if tile_pil.size != (target_w, target_h):
tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS)
resized_tiles.append(pil_to_tensor(tile_pil))
tile_batch = torch.cat(resized_tiles, dim=0)
if images.is_cuda:
tile_batch = tile_batch.cuda()
return tile_batch, x1, y1, extracted_width, extracted_height
def process_tile(self, tile_tensor: torch.Tensor, model, positive, negative, vae,
seed: int, steps: int, cfg: float, sampler_name: str,
scheduler: str, denoise: float, tiled_decode: bool = False,
batch_idx: int = 0, region: Tuple[int, int, int, int] = None,
image_size: Tuple[int, int] = None) -> torch.Tensor:
"""Process a single tile through SD sampling.
Note: positive and negative should already be pre-sliced for the current batch_idx."""
debug_log(f"[process_tile] Processing tile for batch_idx={batch_idx}, seed={seed}, region={region}")
# Import here to avoid circular dependencies
from nodes import common_ksampler, VAEEncode, VAEDecode
# Try to import tiled VAE nodes if available
try:
from nodes import VAEEncodeTiled, VAEDecodeTiled
tiled_vae_available = True
except ImportError:
tiled_vae_available = False
if tiled_decode:
debug_log("Tiled VAE nodes not available, falling back to standard VAE")
# Convert to PIL and back to ensure clean tensor without gradient tracking
tile_pil = tensor_to_pil(tile_tensor, 0)
clean_tensor = pil_to_tensor(tile_pil)
# Ensure tensor is detached and doesn't require gradients
clean_tensor = clean_tensor.detach()
if hasattr(clean_tensor, 'requires_grad_'):
clean_tensor.requires_grad_(False)
# Move to correct device
if tile_tensor.is_cuda:
clean_tensor = clean_tensor.cuda()
clean_tensor = clean_tensor.detach() # Detach again after device transfer
# Clone conditioning per tile (shares models, clones hints for cropping)
positive_tile = clone_conditioning(positive, clone_hints=True)
negative_tile = clone_conditioning(negative, clone_hints=True)
# Crop conditioning to tile region if provided (assumes hints at image resolution)
if region is not None and image_size is not None:
init_size = image_size # (width, height) of full image
canvas_size = image_size
tile_size = (tile_tensor.shape[2], tile_tensor.shape[1]) # (width, height)
w_pad = 0 # No extra pad needed; region already includes padding
h_pad = 0
positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
else:
# No region cropping needed, use cloned conditioning as-is
positive_cropped = positive_tile
negative_cropped = negative_tile
# 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]
# Decode back to image
if tiled_decode and tiled_vae_available:
image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0]
else:
image = VAEDecode().decode(vae, samples)[0]
return image
def process_tiles_batch(self, tile_batch: torch.Tensor, model, positive, negative, vae,
seed: int, steps: int, cfg: float, sampler_name: str,
scheduler: str, denoise: float, tiled_decode: bool,
region: Tuple[int, int, int, int], image_size: Tuple[int, int]) -> torch.Tensor:
"""Process a batch of tiles together (USDU behavior).
tile_batch: [B, H, W, C]
Returns image batch tensor [B, H, W, C]
"""
# Import locally to avoid circular deps
from nodes import common_ksampler, VAEEncode, VAEDecode
try:
from nodes import VAEEncodeTiled, VAEDecodeTiled
tiled_vae_available = True
except ImportError:
tiled_vae_available = False
# Detach and move device
clean = tile_batch.detach()
if hasattr(clean, 'requires_grad_'):
clean.requires_grad_(False)
if tile_batch.is_cuda:
clean = clean.cuda().detach()
# Clone/crop conditioning once for the region
positive_tile = clone_conditioning(positive, clone_hints=True)
negative_tile = clone_conditioning(negative, clone_hints=True)
init_size = image_size
canvas_size = image_size
tile_size = (clean.shape[2], clean.shape[1]) # (W,H)
w_pad = 0
h_pad = 0
positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
# 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]
if tiled_decode and tiled_vae_available:
image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0]
else:
image = VAEDecode().decode(vae, samples)[0]
return image
def create_tile_mask(self, image_width: int, image_height: int,
x: int, y: int, tile_width: int, tile_height: int,
mask_blur: int) -> Image.Image:
"""Create a mask for blending tiles - matches Ultimate SD Upscale approach.
Creates a black image with a white rectangle at the tile position,
then applies blur to create soft edges.
"""
# Create a full-size mask matching the image dimensions
mask = Image.new('L', (image_width, image_height), 0) # Black background
# Draw white rectangle at tile position
draw = ImageDraw.Draw(mask)
draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255)
# Apply blur to soften edges
if mask_blur > 0:
mask = mask.filter(ImageFilter.GaussianBlur(mask_blur))
return mask
def blend_tile(self, base_image: Image.Image, tile_image: Image.Image,
x: int, y: int, extracted_size: Tuple[int, int],
mask: Image.Image, padding: int) -> Image.Image:
"""Blend a processed tile back into the base image using Ultimate SD Upscale's exact approach.
This follows the exact method from ComfyUI_UltimateSDUpscale/modules/processing.py
"""
extracted_width, extracted_height = extracted_size
# Debug logging (uncomment if needed)
# debug_log(f"[Blend] Placing tile at ({x}, {y}), size: {extracted_width}x{extracted_height}")
# Calculate the crop region that was used for extraction
crop_region = (x, y, x + extracted_width, y + extracted_height)
# The mask is already full-size, no need to crop
# Resize the processed tile back to the extracted size
if tile_image.size != (extracted_width, extracted_height):
tile_resized = tile_image.resize((extracted_width, extracted_height), Image.LANCZOS)
else:
tile_resized = tile_image
# Follow Ultimate SD Upscale blending approach:
# Put the tile into position
image_tile_only = Image.new('RGBA', base_image.size)
image_tile_only.paste(tile_resized, crop_region[:2])
# Add the mask as an alpha channel
# Must make a copy due to the possibility of an edge becoming black
temp = image_tile_only.copy()
temp.putalpha(mask) # Use the full image mask
image_tile_only.paste(temp, image_tile_only)
# Add back the tile to the initial image according to the mask in the alpha channel
result = base_image.convert('RGBA')
result.alpha_composite(image_tile_only)
# Convert back to RGB
return result.convert('RGB')
def _slice_conditioning(self, positive, negative, batch_idx):
"""Helper to slice conditioning for a specific batch index."""
# Clone and slice conditioning properly, including ControlNet hints
positive_sliced = clone_conditioning(positive)
negative_sliced = clone_conditioning(negative)
for cond_list in [positive_sliced, negative_sliced]:
for i in range(len(cond_list)):
emb, cond_dict = cond_list[i]
if emb.shape[0] > 1:
cond_list[i][0] = emb[batch_idx:batch_idx+1]
if 'control' in cond_dict:
control = cond_dict['control']
while control is not None:
hint = control.cond_hint_original
if hint.shape[0] > 1:
control.cond_hint_original = hint[batch_idx:batch_idx+1]
control = control.previous_controlnet
if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1:
cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1]
return positive_sliced, negative_sliced
def _process_and_blend_tile(self, tile_idx, tile_pos, upscaled_image, result_image,
model, positive, negative, vae, seed, steps, cfg,
sampler_name, scheduler, denoise, tile_width, tile_height,
padding, mask_blur, image_width, image_height, force_uniform_tiles,
tiled_decode, batch_idx: int = 0):
"""Process a single tile and blend it into the result image."""
x, y = tile_pos
# Extract and process tile
tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding(
upscaled_image, x, y, tile_width, tile_height, padding, force_uniform_tiles
)
processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae,
seed, steps, cfg, sampler_name,
scheduler, denoise, tiled_decode, batch_idx=batch_idx,
region=(x1, y1, x1 + ew, y1 + eh), image_size=(image_width, image_height))
# Convert and blend
processed_pil = tensor_to_pil(processed_tile, 0)
# Create mask for this specific tile (no cache here; only used in single-tile path)
tile_mask = self.create_tile_mask(image_width, image_height, x, y, tile_width, tile_height, mask_blur)
# Use extraction position and size for blending
result_image = self.blend_tile(result_image, processed_pil,
x1, y1, (ew, eh), tile_mask, padding)
return result_image
def _process_single_tile(self, global_idx, num_tiles_per_image, upscaled_image, all_tiles,
model, positive, negative, vae, seed, steps, cfg, sampler_name,
scheduler, denoise, tiled_decode, tile_width, tile_height, padding,
width, height, force_uniform_tiles, sliced_conditioning_cache):
"""Process a single tile."""
# Calculate which image and tile this corresponds to
batch_idx = global_idx // num_tiles_per_image
tile_idx = global_idx % num_tiles_per_image
# Skip if batch_idx is out of range
if batch_idx >= upscaled_image.shape[0]:
debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds batch size {upscaled_image.shape[0]}")
return None
# Get or create sliced conditioning for this batch index
if batch_idx not in sliced_conditioning_cache:
positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx)
sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced)
else:
positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx]
x, y = all_tiles[tile_idx]
# Extract tile from the specific image in the batch
tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding(
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
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,
region=(x1, y1, x1 + ew, y1 + eh), image_size=(width, height))
return {
'tile': processed_tile,
'global_idx': global_idx,
'batch_idx': batch_idx,
'tile_idx': tile_idx,
'x': x1,
'y': y1,
'extracted_width': ew,
'extracted_height': eh
}
upscale/worker_comms.py
+271
View File
@@ -0,0 +1,271 @@
import asyncio, io, json, time
import aiohttp
from PIL import Image
from ..utils.logging import debug_log, log
from ..utils.network import get_client_session
from ..utils.constants import TILE_SEND_TIMEOUT
from ..utils.usdu_managment import MAX_PAYLOAD_SIZE, _send_heartbeat_to_master
from ..utils.image import tensor_to_pil
class WorkerCommsMixin:
async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url,
padding, worker_id):
"""Send all processed tiles to master, chunked if large."""
if not processed_tiles:
return # Early exit if empty
total_tiles = len(processed_tiles)
debug_log(f"Worker[{worker_id[:8]}] - Preparing to send {total_tiles} tiles (size-aware chunks)")
# Prepare encoded images and sizes to enable size-aware chunking
encoded = []
for idx, tile_data in enumerate(processed_tiles):
img = tensor_to_pil(tile_data['tile'], 0)
bio = io.BytesIO()
# Keep compression low to balance speed and size; adjust if needed
img.save(bio, format='PNG', compress_level=0)
raw = bio.getvalue()
encoded.append({
'bytes': raw,
'meta': {
'tile_idx': tile_data['tile_idx'],
'x': tile_data['x'],
'y': tile_data['y'],
'extracted_width': tile_data['extracted_width'],
'extracted_height': tile_data['extracted_height'],
**({'batch_idx': tile_data['batch_idx']} if 'batch_idx' in tile_data else {}),
**({'global_idx': tile_data['global_idx']} if 'global_idx' in tile_data else {}),
}
})
# Size-aware chunking
max_bytes = int(MAX_PAYLOAD_SIZE) - (1024 * 1024) # 1MB headroom
i = 0
chunk_index = 0
while i < total_tiles:
data = aiohttp.FormData()
data.add_field('multi_job_id', multi_job_id)
data.add_field('worker_id', str(worker_id))
data.add_field('padding', str(padding))
metadata = []
used = 0
j = i
while j < total_tiles:
img_bytes = encoded[j]['bytes']
meta = encoded[j]['meta']
# Rough overhead for fields + JSON
overhead = 1024
if used + len(img_bytes) + overhead > max_bytes and j > i:
break
# Accept this tile in this chunk
metadata.append(meta)
data.add_field(f'tile_{j - i}', io.BytesIO(img_bytes), filename=f'tile_{j}.png', content_type='image/png')
used += len(img_bytes) + overhead
j += 1
# Ensure at least one tile per chunk
if j == i:
# Single oversized tile, send anyway
meta = encoded[j]['meta']
metadata.append(meta)
data.add_field('tile_0', io.BytesIO(encoded[j]['bytes']), filename=f'tile_{j}.png', content_type='image/png')
j += 1
chunk_size = j - i
is_chunk_last = (j >= total_tiles)
data.add_field('is_last', str(is_chunk_last))
data.add_field('batch_size', str(chunk_size))
data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json')
# Retry logic with exponential backoff
max_retries = 5
retry_delay = 0.5
for attempt in range(max_retries):
try:
session = await get_client_session()
url = f"{master_url}/distributed/submit_tiles"
async with session.post(url, data=data) as response:
response.raise_for_status()
break
except Exception as e:
if attempt < max_retries - 1:
await asyncio.sleep(retry_delay)
retry_delay = min(retry_delay * 2, 5.0)
else:
log(f"UltimateSDUpscale Worker - Failed to send chunk {chunk_index} after {max_retries} attempts: {e}")
raise
debug_log(f"Worker[{worker_id[:8]}] - Sent chunk {chunk_index} ({chunk_size} tiles, ~{used/1e6:.2f} MB)")
chunk_index += 1
i = j
async def _request_image_from_master(self, multi_job_id, master_url, worker_id):
"""Request an image index to process from master in dynamic mode."""
# Enhanced retries with 404-specific delay, backoff cap, and total timeout to handle init races
max_retries = 10
retry_delay = 0.5
start_time = time.time()
for attempt in range(max_retries):
# Check total timeout
if time.time() - start_time > 30:
log(f"Total request timeout after 30s for worker {worker_id}")
return None, 0
try:
session = await get_client_session()
url = f"{master_url}/distributed/request_image"
async with session.post(url, json={
'worker_id': str(worker_id),
'multi_job_id': multi_job_id
}) as response:
if response.status == 200:
data = await response.json()
image_idx = data.get('image_idx')
estimated_remaining = data.get('estimated_remaining', 0)
return image_idx, estimated_remaining
elif response.status == 404:
# Special handling for 404 - job not found yet
text = await response.text()
debug_log(f"Job not found (404), will retry: {text}")
await asyncio.sleep(1.0)
else:
text = await response.text()
debug_log(f"Request image failed: {response.status} - {text}")
except Exception as e:
if attempt < max_retries - 1:
debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
await asyncio.sleep(retry_delay)
retry_delay *= 2
retry_delay = min(retry_delay, 5.0) # Cap backoff at 5s
else:
log(f"Failed to request image after {max_retries} attempts: {e}")
raise
return None, 0
async def _request_tile_from_master(self, multi_job_id, master_url, worker_id):
"""Request a tile index to process from master in static mode (reusing dynamic infrastructure)."""
# Reuse the same retry logic as dynamic mode
max_retries = 10
retry_delay = 0.5
start_time = time.time()
for attempt in range(max_retries):
# Check total timeout
if time.time() - start_time > 30:
log(f"Total request timeout after 30s for worker {worker_id}")
return None, 0
try:
session = await get_client_session()
url = f"{master_url}/distributed/request_image" # Same endpoint
async with session.post(url, json={
'worker_id': str(worker_id),
'multi_job_id': multi_job_id
}) as response:
if response.status == 200:
data = await response.json()
tile_idx = data.get('tile_idx')
estimated_remaining = data.get('estimated_remaining', 0)
batched_static = data.get('batched_static', False)
return tile_idx, estimated_remaining, batched_static
elif response.status == 404:
# Special handling for 404 - job not found yet
text = await response.text()
debug_log(f"Job not found (404), will retry: {text}")
await asyncio.sleep(1.0)
else:
text = await response.text()
debug_log(f"Request tile failed: {response.status} - {text}")
except Exception as e:
if attempt < max_retries - 1:
debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
await asyncio.sleep(retry_delay)
retry_delay *= 2
retry_delay = min(retry_delay, 5.0) # Cap backoff at 5s
else:
log(f"Failed to request tile after {max_retries} attempts: {e}")
raise
return None, 0, False
async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id,
master_url, worker_id, is_last):
"""Send a processed full image back to master in dynamic mode."""
# Serialize image to PNG
byte_io = io.BytesIO()
image_pil.save(byte_io, format='PNG', compress_level=0)
byte_io.seek(0)
# Prepare form data
data = aiohttp.FormData()
data.add_field('multi_job_id', multi_job_id)
data.add_field('worker_id', str(worker_id))
data.add_field('image_idx', str(image_idx))
data.add_field('is_last', str(is_last))
data.add_field('full_image', byte_io, filename=f'image_{image_idx}.png',
content_type='image/png')
# Retry logic
max_retries = 5
retry_delay = 0.5
for attempt in range(max_retries):
try:
session = await get_client_session()
url = f"{master_url}/distributed/submit_image"
async with session.post(url, data=data) as response:
response.raise_for_status()
debug_log(f"Successfully sent image {image_idx} to master")
return
except Exception as e:
if attempt < max_retries - 1:
debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
await asyncio.sleep(retry_delay)
retry_delay *= 2
else:
log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}")
raise
async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id):
"""Send completion signal to master in dynamic mode."""
# Send a dummy request with is_last=True
data = aiohttp.FormData()
data.add_field('multi_job_id', multi_job_id)
data.add_field('worker_id', str(worker_id))
data.add_field('is_last', 'true')
# No image data - just completion signal
session = await get_client_session()
url = f"{master_url}/distributed/submit_image"
async with session.post(url, data=data) as response:
response.raise_for_status()
debug_log(f"Worker {worker_id} sent completion signal")
async def _check_job_status(self, multi_job_id, master_url):
"""Check if job is ready on the master."""
try:
session = await get_client_session()
url = f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}"
async with session.get(url) as response:
if response.status == 200:
data = await response.json()
return data.get('ready', False)
return False
except Exception as e:
debug_log(f"Job status check failed: {e}")
return False
async def _async_yield(self):
"""Simple async yield to allow event loop processing."""
await asyncio.sleep(0)