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

speed optimisation

Browse Source
This commit is contained in:
robertvoy
2025-09-10 10:15:35 +10:00
parent 58f26a106f
commit 287e4c05fa
2 changed files with 38 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
distributed_upscale.py
+25 -12
View File
@@ -277,7 +277,7 @@ class UltimateSDUpscaleDistributed:
# Convert and blend
processed_pil = tensor_to_pil(processed_tile, 0)
# Create mask for this specific tile
# 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,
@@ -353,9 +353,8 @@ class UltimateSDUpscaleDistributed:
# 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
collected_results[key] = {
'tensor': tile_data['tensor'],
# 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'],
@@ -366,6 +365,11 @@ class UltimateSDUpscaleDistributed:
'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']
@@ -935,7 +939,7 @@ class UltimateSDUpscaleDistributed:
break
# Wait a bit before next check
await asyncio.sleep(0.5)
await asyncio.sleep(0.1)
# Get all completed tasks for return
return await self._get_all_completed_tasks(multi_job_id)
@@ -967,9 +971,14 @@ class UltimateSDUpscaleDistributed:
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 = 5
max_consecutive_no_tile = 2
while processed_count < total_tiles:
comfy.model_management.throw_exception_if_processing_interrupted()
@@ -991,7 +1000,7 @@ class UltimateSDUpscaleDistributed:
seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode,
region, (width, height)
)
tile_mask = self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur)
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):
@@ -1087,7 +1096,7 @@ class UltimateSDUpscaleDistributed:
# 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:
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)
@@ -1099,10 +1108,14 @@ class UltimateSDUpscaleDistributed:
# Blend tile synchronously
x = tile_data.get('x', 0)
y = tile_data.get('y', 0)
tile_tensor = tile_data['tensor']
tile_pil = tensor_to_pil(tile_tensor, 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 = self.create_tile_mask(width, height, orig_x, orig_y, tile_width, tile_height, mask_blur)
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,
@@ -1546,7 +1559,7 @@ class UltimateSDUpscaleDistributed:
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=1.0)
tile_idx = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=0.1)
return tile_idx
except asyncio.TimeoutError:
return None
utils/usdu_managment.py
+13 -12
View File
@@ -67,13 +67,14 @@ def _parse_tiles_from_form(data):
except Exception as e:
raise ValueError(f"Invalid image data for tile {i}: {e}")
# Convert to tensor [1, H, W, C]
# Convert to tensor [1, H, W, C] (keep PIL too for faster blending later)
tile_tensor = pil_to_tensor(img)
# Build tile dictionary
try:
tile_info = {
'tensor': tile_tensor,
'image': img, # keep PIL to avoid reconversion later
'tile_idx': int(meta.get('tile_idx', i)),
'x': int(meta.get('x', 0)),
'y': int(meta.get('y', 0)),
@@ -136,7 +137,7 @@ async def init_dynamic_job(multi_job_id: str, batch_size: int, enabled_workers:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
job_data['completed_images'] = {}
job_data['pending_images'] = job_data[JOB_PENDING_TASKS]
debug_log(f"Dynamic job {multi_job_id} initialized with {batch_size} images")
debug_log(f"Job {multi_job_id} initialized with {batch_size} images")
async def init_static_job_batched(multi_job_id: str, batch_size: int, num_tiles_per_image: int, enabled_workers: List[str]):
@@ -152,7 +153,7 @@ async def init_static_job_batched(multi_job_id: str, batch_size: int, num_tiles_
enabled_workers=enabled_workers,
batched_static=True,
)
debug_log(f"Static job {multi_job_id} initialized with {num_tiles_per_image} tile ids for batch {batch_size}")
debug_log(f"Job {multi_job_id} initialized with {num_tiles_per_image} tile ids for batch {batch_size}")
async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_image=None, all_indices=None, enabled_workers=None, task_assignments=None, batched_static: bool = False):
"""Unified initialization for job queues in static and dynamic modes."""
@@ -174,9 +175,9 @@ async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_ima
if mode == 'dynamic':
job_data[JOB_BATCH_SIZE] = batch_size
pending_queue = job_data[JOB_PENDING_TASKS]
for i in all_indices or range(batch_size):
for i in (all_indices or range(batch_size)):
await pending_queue.put(i)
debug_log(f"Initialized dynamic queue with {batch_size} pending images")
debug_log(f"Initialized image queue with {batch_size} pending items")
elif mode == 'static':
job_data[JOB_NUM_TILES_PER_IMAGE] = num_tiles_per_image
job_data[JOB_BATCH_SIZE] = batch_size
@@ -186,12 +187,12 @@ async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_ima
if batched_static and num_tiles_per_image is not None:
for i in range(num_tiles_per_image):
await pending_queue.put(i)
debug_log(f"Initialized static queue (batched) with {num_tiles_per_image} tile ids for dynamic distribution across batch of {batch_size}")
debug_log(f"Initialized tile-id queue with {num_tiles_per_image} ids for batch {batch_size}")
else:
total_tiles = batch_size * num_tiles_per_image
for i in range(total_tiles):
await pending_queue.put(i)
debug_log(f"Initialized static queue with {total_tiles} pending tiles for dynamic distribution")
debug_log(f"Initialized tile queue with {total_tiles} pending tasks")
# Keep backward compatibility - if task assignments provided, still track them
if task_assignments and enabled_workers:
@@ -412,7 +413,7 @@ async def submit_tiles_endpoint(request):
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
return await handle_api_error(request, "Mode mismatch: expected static mode", 400)
return await handle_api_error(request, "Job not configured for tile submissions", 400)
if JOB_QUEUE in job_data:
await job_data[JOB_QUEUE].put({
'worker_id': worker_id,
@@ -432,7 +433,7 @@ async def submit_tiles_endpoint(request):
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
return await handle_api_error(request, "Mode mismatch: expected static mode", 400)
return await handle_api_error(request, "Job not configured for tile submissions", 400)
q = job_data[JOB_QUEUE]
if batch_size > 0 or len(tiles) > 0:
@@ -485,7 +486,7 @@ async def submit_image_endpoint(request):
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
return await handle_api_error(request, "Mode mismatch: expected dynamic mode", 400)
return await handle_api_error(request, "Job not configured for image submissions", 400)
if JOB_QUEUE in job_data:
await job_data[JOB_QUEUE].put({
'worker_id': worker_id,
@@ -501,7 +502,7 @@ async def submit_image_endpoint(request):
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
return await handle_api_error(request, "Mode mismatch: expected dynamic mode", 400)
return await handle_api_error(request, "Job not configured for image submissions", 400)
if JOB_QUEUE in job_data:
await job_data[JOB_QUEUE].put({
'worker_id': worker_id,
@@ -599,7 +600,7 @@ async def request_image_endpoint(request):
elif mode == 'static' and JOB_PENDING_TASKS in job_data:
pending_queue = job_data[JOB_PENDING_TASKS]
else:
return await handle_api_error(request, f"Invalid {mode} mode configuration", 400)
return await handle_api_error(request, "Invalid job configuration", 400)
try:
task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)