clean up
This commit is contained in:
robertvoy
+177
-223
@@ -33,13 +33,14 @@ from .utils.usdu_utils import (
|
||||
from .utils.usdu_managment import (
|
||||
clone_conditioning, ensure_tile_jobs_initialized,
|
||||
# Job management functions
|
||||
_init_job_queue, _drain_results_queue,
|
||||
_drain_results_queue,
|
||||
_check_and_requeue_timed_out_workers, _get_completed_count, _mark_task_completed,
|
||||
_send_heartbeat_to_master, _cleanup_job,
|
||||
# Constants
|
||||
JOB_COMPLETED_TASKS, JOB_WORKER_STATUS, JOB_PENDING_TASKS,
|
||||
MAX_PAYLOAD_SIZE
|
||||
)
|
||||
from .utils.usdu_managment import init_dynamic_job, init_static_job_batched
|
||||
|
||||
|
||||
# Note: MAX_BATCH and HEARTBEAT_TIMEOUT are imported from utils.constants
|
||||
@@ -172,13 +173,18 @@ class UltimateSDUpscaleDistributed:
|
||||
# Determine mode (must match master's logic)
|
||||
enabled_workers = json.loads(enabled_worker_ids)
|
||||
num_workers = len(enabled_workers)
|
||||
|
||||
# Compute number of tiles for this image to decide if tile distribution makes sense
|
||||
_, height, width, _ = upscaled_image.shape
|
||||
all_tiles = self.calculate_tiles(width, height, self.round_to_multiple(tile_width), self.round_to_multiple(tile_height), force_uniform_tiles)
|
||||
num_tiles_per_image = len(all_tiles)
|
||||
|
||||
mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
|
||||
if batch_size > 1 and num_workers > 0:
|
||||
debug_log("UltimateSDUpscale Worker - Forcing static mode for USDU-style batched-per-tile processing")
|
||||
# For USDU-style processing, we want tile distribution whenever workers are available
|
||||
# and there is more than one tile to process, even if batch == 1.
|
||||
if num_workers > 0 and num_tiles_per_image > 1:
|
||||
mode = "static"
|
||||
|
||||
debug_log(f"UltimateSDUpscale Worker - Mode: {mode}, batch_size: {batch_size}")
|
||||
debug_log(f"USDU Dist Worker - Batch size {batch_size}")
|
||||
|
||||
if mode == "dynamic":
|
||||
return self.process_worker_dynamic(upscaled_image, model, positive, negative, vae,
|
||||
@@ -207,11 +213,14 @@ class UltimateSDUpscaleDistributed:
|
||||
# Get image dimensions and batch size
|
||||
batch_size, height, width, _ = upscaled_image.shape
|
||||
|
||||
# Calculate all tiles
|
||||
# Calculate all tiles and grid
|
||||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||||
num_tiles_per_image = len(all_tiles)
|
||||
|
||||
debug_log(f"UltimateSDUpscale Master - Tiles per image: {num_tiles_per_image}")
|
||||
rows = math.ceil(height / tile_height)
|
||||
cols = math.ceil(width / tile_width)
|
||||
log(
|
||||
f"USDU Dist: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({num_tiles_per_image} tiles/image) | Batch {batch_size}"
|
||||
)
|
||||
|
||||
# Parse enabled workers
|
||||
enabled_workers = json.loads(enabled_worker_ids)
|
||||
@@ -219,11 +228,12 @@ class UltimateSDUpscaleDistributed:
|
||||
|
||||
# Determine processing mode
|
||||
mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
|
||||
if batch_size > 1 and num_workers > 0 and num_tiles_per_image > 1:
|
||||
debug_log("UltimateSDUpscale Master - Forcing static mode for USDU-style batched-per-tile processing")
|
||||
# Prefer tile-based static distribution when workers are available and there are multiple tiles,
|
||||
# even for batch == 1, to spread tiles across GPUs like the legacy dynamic tile queue.
|
||||
if num_workers > 0 and num_tiles_per_image > 1:
|
||||
mode = "static"
|
||||
|
||||
debug_log(f"UltimateSDUpscale Master - Mode: {mode}, batch_size: {batch_size}, workers: {num_workers}")
|
||||
log(f"USDU Dist: Workers {num_workers}")
|
||||
|
||||
if mode == "single_gpu":
|
||||
# No workers, process all tiles locally
|
||||
@@ -275,61 +285,6 @@ class UltimateSDUpscaleDistributed:
|
||||
|
||||
return result_image
|
||||
|
||||
async def _init_job_queue_unified(self, multi_job_id, mode='static', batch_size=None,
|
||||
total_tiles=None, enabled_workers=None, num_tiles_per_image=None,
|
||||
assigned_to_workers=None, all_indices=None, task_assignments=None,
|
||||
batched_static: bool = True):
|
||||
"""Unified job queue initialization for both static and dynamic modes."""
|
||||
if mode == 'dynamic':
|
||||
# Dynamic mode initialization
|
||||
await _init_job_queue(multi_job_id, 'dynamic', batch_size=batch_size,
|
||||
all_indices=all_indices,
|
||||
enabled_workers=list(assigned_to_workers.keys()) if assigned_to_workers else [])
|
||||
# Add additional fields for backward compatibility
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
job_data['completed_images'] = {}
|
||||
job_data['completed_tiles'] = {} # For uniformity
|
||||
# Fix: The endpoint expects 'pending_images', not 'pending_tasks'
|
||||
job_data['pending_images'] = job_data[JOB_PENDING_TASKS]
|
||||
debug_log(f"UltimateSDUpscale Master Dynamic - Initialized job queue with {batch_size} pending images for all participants")
|
||||
else:
|
||||
# Static mode initialization
|
||||
# If num_tiles_per_image not provided, calculate it
|
||||
if num_tiles_per_image is None and total_tiles is not None and batch_size is not None:
|
||||
num_tiles_per_image = total_tiles // batch_size if batch_size > 0 else total_tiles
|
||||
|
||||
await _init_job_queue(multi_job_id, 'static',
|
||||
batch_size=batch_size,
|
||||
num_tiles_per_image=num_tiles_per_image,
|
||||
all_indices=None,
|
||||
enabled_workers=enabled_workers,
|
||||
task_assignments=task_assignments,
|
||||
batched_static=batched_static)
|
||||
|
||||
# Add completed_tiles and total_items for backward compatibility
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
job_data['completed_tiles'] = {}
|
||||
if total_tiles is not None:
|
||||
job_data['total_items'] = total_tiles
|
||||
debug_log(f"UltimateSDUpscale Master - Initialized job queue for {multi_job_id}")
|
||||
|
||||
# Keep compatibility wrappers
|
||||
async def _init_job_queue_dynamic(self, multi_job_id, batch_size, assigned_to_workers=None, worker_status=None, all_indices=None):
|
||||
"""Initialize the job queue for dynamic mode with pending images."""
|
||||
await self._init_job_queue_unified(multi_job_id, mode='dynamic', batch_size=batch_size,
|
||||
assigned_to_workers=assigned_to_workers, all_indices=all_indices)
|
||||
|
||||
async def _init_job_queue(self, multi_job_id, total_tiles, enabled_workers, batch_size=1, num_tiles_per_image=None, task_assignments=None, batched_static: bool = True):
|
||||
"""Initialize the job queue for collecting tiles."""
|
||||
await self._init_job_queue_unified(multi_job_id, mode='static', batch_size=batch_size,
|
||||
total_tiles=total_tiles, enabled_workers=enabled_workers,
|
||||
num_tiles_per_image=num_tiles_per_image, task_assignments=task_assignments,
|
||||
batched_static=batched_static)
|
||||
|
||||
async def _async_collect_results(self, multi_job_id, num_workers, mode='static',
|
||||
remaining_to_collect=None, batch_size=None):
|
||||
@@ -355,9 +310,8 @@ class UltimateSDUpscaleDistributed:
|
||||
# Calculate expected count from job data for static mode
|
||||
expected_count = job_data.get('total_items', 0) - len(job_data.get('completed_items', set()))
|
||||
|
||||
mode_str = "Dynamic" if mode == 'dynamic' else "Static"
|
||||
item_type = "images" if mode == 'dynamic' else "tiles"
|
||||
debug_log(f"UltimateSDUpscale Master {mode_str} - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers")
|
||||
debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers")
|
||||
|
||||
collected_results = {}
|
||||
workers_done = set()
|
||||
@@ -426,11 +380,11 @@ class UltimateSDUpscaleDistributed:
|
||||
completed_images[image_idx] = image_pil
|
||||
collected_results[image_idx] = image_pil
|
||||
collected_count += 1
|
||||
debug_log(f"UltimateSDUpscale Master Dynamic - Received image {image_idx} from worker {worker_id}")
|
||||
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 {mode_str} - Worker {worker_id} completed")
|
||||
debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed")
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
if mode == 'dynamic':
|
||||
@@ -440,18 +394,18 @@ class UltimateSDUpscaleDistributed:
|
||||
# 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 Dynamic - Requeued {requeued} images from timed out workers")
|
||||
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 Dynamic - Overall timeout waiting for images")
|
||||
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 {mode_str} - Collection complete. Got {len(collected_results)} {item_type} from {len(workers_done)} workers")
|
||||
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:
|
||||
@@ -833,7 +787,7 @@ class UltimateSDUpscaleDistributed:
|
||||
sliced_conditioning_cache = {}
|
||||
|
||||
# Dynamic queue mode (static processing): process batched-per-tile
|
||||
debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} ready to process tiles dynamically (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
|
||||
@@ -844,7 +798,7 @@ class UltimateSDUpscaleDistributed:
|
||||
timeout=5.0
|
||||
)
|
||||
if ready:
|
||||
debug_log(f"Job {multi_job_id} ready after {attempt} attempts")
|
||||
debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts")
|
||||
break
|
||||
time.sleep(1.0)
|
||||
else:
|
||||
@@ -860,11 +814,11 @@ class UltimateSDUpscaleDistributed:
|
||||
)
|
||||
|
||||
if tile_idx is None:
|
||||
debug_log(f"UltimateSDUpscale Worker - No more tiles to process")
|
||||
debug_log(f"Worker[{worker_id[:8]}] - No more tiles to process")
|
||||
break
|
||||
|
||||
# Always batched-per-tile in static mode
|
||||
debug_log(f"UltimateSDUpscale Worker - Assigned tile_id {tile_idx} to worker {worker_id}")
|
||||
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]
|
||||
@@ -900,15 +854,15 @@ class UltimateSDUpscaleDistributed:
|
||||
timeout=5.0
|
||||
)
|
||||
except Exception as e:
|
||||
debug_log(f"Heartbeat failed: {e}")
|
||||
debug_log(f"Worker[{worker_id[:8]}] heartbeat failed: {e}")
|
||||
|
||||
# Send tiles in batches
|
||||
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 = []
|
||||
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:
|
||||
@@ -1004,13 +958,14 @@ class UltimateSDUpscaleDistributed:
|
||||
|
||||
sliced_conditioning_cache = {}
|
||||
# Initialize queue: pending queue holds tile ids (batched per tile)
|
||||
debug_log(f"UltimateSDUpscale Master - Using dynamic tile distribution")
|
||||
log("USDU Dist: Using tile queue distribution")
|
||||
run_async_in_server_loop(
|
||||
self._init_job_queue(multi_job_id, total_tiles, enabled_workers, batch_size, num_tiles_per_image, batched_static=True),
|
||||
init_static_job_batched(multi_job_id, batch_size, num_tiles_per_image, enabled_workers),
|
||||
timeout=10.0
|
||||
)
|
||||
debug_log(f"UltimateSDUpscale Master - Initialized batched static job queue with {num_tiles_per_image} tile ids for dynamic distribution")
|
||||
debug_log(f"UltimateSDUpscale Master - Processing tiles dynamically from queue")
|
||||
debug_log(
|
||||
f"Initialized tile-id queue with {num_tiles_per_image} ids for batch {batch_size}"
|
||||
)
|
||||
|
||||
processed_count = 0
|
||||
consecutive_no_tile = 0
|
||||
@@ -1047,6 +1002,7 @@ class UltimateSDUpscaleDistributed:
|
||||
_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:
|
||||
@@ -1111,7 +1067,7 @@ class UltimateSDUpscaleDistributed:
|
||||
local_tile_idx, all_tiles[local_tile_idx], upscaled_image[batch_idx:batch_idx+1], result_images[batch_idx],
|
||||
model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height,
|
||||
padding, mask_blur, width, height, tiled_decode, batch_idx=batch_idx
|
||||
padding, mask_blur, width, height, force_uniform_tiles, tiled_decode, batch_idx=batch_idx
|
||||
)
|
||||
|
||||
# Mark as completed and store in collected_tasks
|
||||
@@ -1222,7 +1178,7 @@ class UltimateSDUpscaleDistributed:
|
||||
return # Early exit if empty
|
||||
|
||||
total_tiles = len(processed_tiles)
|
||||
debug_log(f"UltimateSDUpscale Worker - Preparing to send {total_tiles} 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 = []
|
||||
@@ -1303,7 +1259,7 @@ class UltimateSDUpscaleDistributed:
|
||||
log(f"UltimateSDUpscale Worker - Failed to send chunk {chunk_index} after {max_retries} attempts: {e}")
|
||||
raise
|
||||
|
||||
debug_log(f"UltimateSDUpscale Worker - Sent chunk {chunk_index} ({chunk_size} tiles, ~{used/1e6:.2f} MB)")
|
||||
debug_log(f"Worker[{worker_id[:8]}] - Sent chunk {chunk_index} ({chunk_size} tiles, ~{used/1e6:.2f} MB)")
|
||||
chunk_index += 1
|
||||
i = j
|
||||
|
||||
@@ -1321,7 +1277,11 @@ class UltimateSDUpscaleDistributed:
|
||||
# Calculate all tiles
|
||||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||||
|
||||
debug_log(f"UltimateSDUpscale - Processing {len(all_tiles)} tiles locally for batch of {batch_size}, batched per tile")
|
||||
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 = []
|
||||
@@ -1368,27 +1328,24 @@ class UltimateSDUpscaleDistributed:
|
||||
batch_size, height, width, _ = upscaled_image.shape
|
||||
num_workers = len(enabled_workers)
|
||||
|
||||
debug_log(f"UltimateSDUpscale Master Dynamic - Processing batch of {batch_size} images with {num_workers} 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))
|
||||
|
||||
log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.")
|
||||
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:
|
||||
assigned_to_workers = {w: [] for w in enabled_workers}
|
||||
worker_status = {w: time.time() for w in enabled_workers}
|
||||
|
||||
run_async_in_server_loop(
|
||||
self._init_job_queue_dynamic(multi_job_id, batch_size, assigned_to_workers, worker_status, all_indices),
|
||||
init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices),
|
||||
timeout=2.0
|
||||
)
|
||||
except Exception as e:
|
||||
debug_log(f"UltimateSDUpscale Master Dynamic - Queue initialization error: {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
|
||||
@@ -1472,7 +1429,7 @@ class UltimateSDUpscaleDistributed:
|
||||
self._get_total_completed_count(multi_job_id),
|
||||
timeout=1.0
|
||||
)
|
||||
log(f"Progress: {completed_now}/{batch_size} images completed")
|
||||
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)
|
||||
@@ -1502,7 +1459,7 @@ class UltimateSDUpscaleDistributed:
|
||||
timeout=1.0
|
||||
)
|
||||
|
||||
log(f"Progress: {completed_now}/{batch_size} images completed")
|
||||
log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
|
||||
|
||||
if completed_now >= batch_size:
|
||||
break
|
||||
@@ -1559,7 +1516,7 @@ class UltimateSDUpscaleDistributed:
|
||||
if upscaled_image.is_cuda:
|
||||
result_tensor = result_tensor.cuda()
|
||||
|
||||
debug_log(f"UltimateSDUpscale Master Dynamic - Job {multi_job_id} complete")
|
||||
debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
|
||||
log(f"Completed processing all {batch_size} images")
|
||||
return (result_tensor,)
|
||||
|
||||
@@ -1665,128 +1622,127 @@ class UltimateSDUpscaleDistributed:
|
||||
|
||||
|
||||
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
|
||||
_, height, width, _ = upscaled_image.shape
|
||||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||||
|
||||
debug_log(f"UltimateSDUpscale Worker Dynamic - Worker {worker_id} ready to process images")
|
||||
|
||||
# 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),
|
||||
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
|
||||
)
|
||||
if ready:
|
||||
debug_log(f"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"UltimateSDUpscale Worker Dynamic - No more images to process")
|
||||
break
|
||||
|
||||
debug_log(f"UltimateSDUpscale Worker Dynamic - Mode: dynamic, assigned image {image_idx} to worker {worker_id}")
|
||||
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"UltimateSDUpscale Worker Dynamic - Error sending image {image_idx}: {e}")
|
||||
# Continue processing other images
|
||||
|
||||
# Send final is_last signal
|
||||
debug_log(f"UltimateSDUpscale Worker Dynamic - Worker {worker_id} processed {processed_count} images, sending completion signal")
|
||||
|
||||
# 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_worker_complete_signal(multi_job_id, master_url, worker_id),
|
||||
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"UltimateSDUpscale Worker Dynamic - Error sending completion signal: {e}")
|
||||
|
||||
return (upscaled_image,)
|
||||
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,)
|
||||
|
||||
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."""
|
||||
@@ -1957,15 +1913,13 @@ class UltimateSDUpscaleDistributed:
|
||||
|
||||
def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str:
|
||||
"""Determines processing mode per requested policy:
|
||||
- batch_size > 1 => static (tile-based)
|
||||
- batch_size == 1 => dynamic (per-image)
|
||||
- any workers => prefer static (tile-based) for USDU
|
||||
- no workers => single_gpu
|
||||
"""
|
||||
if num_workers == 0:
|
||||
return "single_gpu"
|
||||
return "static" if batch_size > 1 else "dynamic"
|
||||
|
||||
# Note: Using _distribute_tasks from usdu_managment.py instead of duplicate _distribute_items
|
||||
# Default to static when distributed; master/worker may still override if special cases arise
|
||||
return "static"
|
||||
|
||||
# Ensure initialization before registering routes
|
||||
ensure_tile_jobs_initialized()
|
||||
|
||||
+128
-315
@@ -7,13 +7,11 @@ import io
|
||||
from aiohttp import web
|
||||
import server
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# Import from other utilities
|
||||
from .logging import debug_log, log
|
||||
from .network import handle_api_error, get_client_session
|
||||
from .image import tensor_to_pil
|
||||
from .image import pil_to_tensor
|
||||
|
||||
# Configure maximum payload size (50MB default, configurable via environment variable)
|
||||
MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))
|
||||
@@ -22,6 +20,87 @@ MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024
|
||||
from .constants import HEARTBEAT_TIMEOUT
|
||||
|
||||
|
||||
def _parse_tiles_from_form(data):
|
||||
"""Parse tiles submitted via multipart/form-data into a list of tile dicts.
|
||||
|
||||
Expects the following fields in the aiohttp form data:
|
||||
- 'tiles_metadata': JSON list with per-tile metadata items containing at least
|
||||
'tile_idx', 'x', 'y', 'extracted_width', 'extracted_height'. Optional
|
||||
'batch_idx' and 'global_idx' are included when available.
|
||||
- 'tile_{i}': image bytes for each tile described in tiles_metadata (PNG).
|
||||
- 'padding': integer padding used during extraction (optional; defaults 0).
|
||||
|
||||
Returns: list of dicts with keys: 'tensor', 'tile_idx', 'x', 'y',
|
||||
'extracted_width', 'extracted_height', and optional 'batch_idx', 'global_idx',
|
||||
plus 'padding'.
|
||||
"""
|
||||
try:
|
||||
# Parse padding if present
|
||||
padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0
|
||||
except Exception:
|
||||
padding = 0
|
||||
|
||||
# Parse tiles metadata (JSON list)
|
||||
meta_raw = data.get('tiles_metadata')
|
||||
if meta_raw is None:
|
||||
raise ValueError("Missing tiles_metadata")
|
||||
|
||||
try:
|
||||
metadata = json.loads(meta_raw)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid tiles_metadata JSON: {e}")
|
||||
|
||||
if not isinstance(metadata, list):
|
||||
raise ValueError("tiles_metadata must be a list")
|
||||
|
||||
tiles = []
|
||||
# Iterate over metadata items and corresponding uploaded files tile_0, tile_1, ...
|
||||
for i, meta in enumerate(metadata):
|
||||
file_field = data.get(f'tile_{i}')
|
||||
if file_field is None or not hasattr(file_field, 'file'):
|
||||
raise ValueError(f"Missing tile data for index {i}")
|
||||
|
||||
# Read image bytes and decode to PIL
|
||||
raw = file_field.file.read()
|
||||
try:
|
||||
img = Image.open(io.BytesIO(raw)).convert("RGB")
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid image data for tile {i}: {e}")
|
||||
|
||||
# Convert to tensor [1, H, W, C]
|
||||
tile_tensor = pil_to_tensor(img)
|
||||
|
||||
# Build tile dictionary
|
||||
try:
|
||||
tile_info = {
|
||||
'tensor': tile_tensor,
|
||||
'tile_idx': int(meta.get('tile_idx', i)),
|
||||
'x': int(meta.get('x', 0)),
|
||||
'y': int(meta.get('y', 0)),
|
||||
'extracted_width': int(meta.get('extracted_width', img.width)),
|
||||
'extracted_height': int(meta.get('extracted_height', img.height)),
|
||||
'padding': int(padding),
|
||||
}
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid metadata values for tile {i}: {e}")
|
||||
|
||||
# Optional fields
|
||||
if 'batch_idx' in meta:
|
||||
try:
|
||||
tile_info['batch_idx'] = int(meta['batch_idx'])
|
||||
except Exception:
|
||||
pass
|
||||
if 'global_idx' in meta:
|
||||
try:
|
||||
tile_info['global_idx'] = int(meta['global_idx'])
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
tiles.append(tile_info)
|
||||
|
||||
return tiles
|
||||
|
||||
|
||||
# Unified Job Data Structure Keys
|
||||
JOB_QUEUE = 'queue'
|
||||
JOB_MODE = 'mode'
|
||||
@@ -36,6 +115,45 @@ JOB_NUM_TILES_PER_IMAGE = 'num_tiles_per_image' # For static
|
||||
TASK_TYPE_TILE = 'tile'
|
||||
TASK_TYPE_IMAGE = 'image'
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
async def init_dynamic_job(multi_job_id: str, batch_size: int, enabled_workers: List[str], all_indices: Optional[List[int]] = None):
|
||||
"""Initialize queue for dynamic mode (per-image), with collector fields.
|
||||
|
||||
- Creates JOB_PENDING_TASKS with image indices
|
||||
- Adds 'completed_images' dict and 'pending_images' alias used by collectors
|
||||
"""
|
||||
await _init_job_queue(
|
||||
multi_job_id,
|
||||
'dynamic',
|
||||
batch_size=batch_size,
|
||||
all_indices=all_indices or list(range(batch_size)),
|
||||
enabled_workers=enabled_workers,
|
||||
)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
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")
|
||||
|
||||
|
||||
async def init_static_job_batched(multi_job_id: str, batch_size: int, num_tiles_per_image: int, enabled_workers: List[str]):
|
||||
"""Initialize queue for static mode (batched-per-tile).
|
||||
|
||||
- Populates JOB_PENDING_TASKS with tile ids [0..num_tiles_per_image-1]
|
||||
"""
|
||||
await _init_job_queue(
|
||||
multi_job_id,
|
||||
'static',
|
||||
batch_size=batch_size,
|
||||
num_tiles_per_image=num_tiles_per_image,
|
||||
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}")
|
||||
|
||||
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."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
@@ -87,34 +205,7 @@ async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_ima
|
||||
|
||||
prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
|
||||
|
||||
def _distribute_tasks(items: list, num_participants: int) -> list[list[any]]:
|
||||
"""Distribute a list of items among N participants (master + workers)."""
|
||||
if num_participants <= 1:
|
||||
return [items]
|
||||
|
||||
items_per_participant = len(items) // num_participants
|
||||
remainder = len(items) % num_participants
|
||||
assignments = []
|
||||
start_idx = 0
|
||||
for i in range(num_participants):
|
||||
count = items_per_participant + (1 if i < remainder else 0)
|
||||
end_idx = start_idx + count
|
||||
assignments.append(items[start_idx:end_idx])
|
||||
start_idx = end_idx
|
||||
return assignments
|
||||
|
||||
async def _get_next_task(multi_job_id):
|
||||
"""Get next task from pending queue (generalized for tiles/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 JOB_PENDING_TASKS not in job_data:
|
||||
return None
|
||||
try:
|
||||
task_id = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=1.0)
|
||||
return task_id
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
# Note: legacy task distribution and queue pull helpers removed
|
||||
|
||||
async def _drain_results_queue(multi_job_id):
|
||||
"""Drain pending results from queue and update completed_tasks. Returns count drained."""
|
||||
@@ -331,108 +422,10 @@ async def submit_tiles_endpoint(request):
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
# Handle batch tiles with metadata
|
||||
if batch_size > 0:
|
||||
padding = int(data.get('padding', 32))
|
||||
metadata_field = data.get('tiles_metadata')
|
||||
if metadata_field:
|
||||
if hasattr(metadata_field, 'file'):
|
||||
metadata_str = metadata_field.file.read().decode('utf-8')
|
||||
elif isinstance(metadata_field, (bytes, bytearray)):
|
||||
metadata_str = metadata_field.decode('utf-8')
|
||||
else:
|
||||
metadata_str = str(metadata_field)
|
||||
|
||||
metadata = json.loads(metadata_str)
|
||||
if len(metadata) != batch_size:
|
||||
return await handle_api_error(request, "Metadata length mismatch", 400)
|
||||
|
||||
tile_data_list = []
|
||||
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
continue
|
||||
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
if i < len(metadata):
|
||||
tile_meta = metadata[i]
|
||||
tile_idx = tile_meta.get('tile_idx', i)
|
||||
tile_info = {
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': tile_meta['x'],
|
||||
'y': tile_meta['y'],
|
||||
'extracted_width': tile_meta['extracted_width'],
|
||||
'extracted_height': tile_meta['extracted_height'],
|
||||
'padding': padding,
|
||||
'batch_idx': tile_meta.get('batch_idx', 0),
|
||||
'global_idx': tile_meta.get('global_idx', tile_idx)
|
||||
}
|
||||
tile_data_list.append(tile_info)
|
||||
|
||||
tile_data_list.sort(key=lambda x: x['tile_idx'])
|
||||
tiles.extend(tile_data_list)
|
||||
else:
|
||||
# Legacy format
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
continue
|
||||
|
||||
tile_idx = int(data.get(f'tile_{i}_idx', i))
|
||||
x = int(data.get(f'tile_{i}_x', 0))
|
||||
y = int(data.get(f'tile_{i}_y', 0))
|
||||
extracted_width = int(data.get(f'tile_{i}_width', 512))
|
||||
extracted_height = int(data.get(f'tile_{i}_height', 512))
|
||||
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles.append({
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding
|
||||
})
|
||||
else:
|
||||
# Single tile legacy
|
||||
image_file = data.get('image')
|
||||
if not image_file:
|
||||
return await handle_api_error(request, "Missing image data", 400)
|
||||
|
||||
tile_idx = int(data.get('tile_idx', 0))
|
||||
x = int(data.get('x', 0))
|
||||
y = int(data.get('y', 0))
|
||||
extracted_width = int(data.get('extracted_width', 512))
|
||||
extracted_height = int(data.get('extracted_height', 512))
|
||||
padding = int(data.get('padding', 32))
|
||||
|
||||
img_data = image_file.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles = [{
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding,
|
||||
'batch_idx': 0,
|
||||
'global_idx': tile_idx
|
||||
}]
|
||||
try:
|
||||
tiles = _parse_tiles_from_form(data)
|
||||
except ValueError as e:
|
||||
return await handle_api_error(request, str(e), 400)
|
||||
|
||||
# Submit tiles to queue
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
@@ -448,6 +441,7 @@ async def submit_tiles_endpoint(request):
|
||||
'tiles': tiles,
|
||||
'is_last': is_last
|
||||
})
|
||||
debug_log(f"Received {len(tiles)} tiles from worker {worker_id} (is_last={is_last})")
|
||||
else:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
@@ -523,188 +517,7 @@ async def submit_image_endpoint(request):
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
# Keep legacy endpoint for backward compatibility
|
||||
@server.PromptServer.instance.routes.post("/distributed/tile_complete")
|
||||
async def tile_complete_endpoint(request):
|
||||
try:
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
if 'full_image' in data and 'image_idx' in data:
|
||||
image_idx = int(data.get('image_idx'))
|
||||
img_data = data['full_image'].file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
|
||||
debug_log(f"Received full image {image_idx} from worker {worker_id}")
|
||||
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
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 for image submission", 400)
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
'image_idx': image_idx,
|
||||
'image': img,
|
||||
'is_last': is_last
|
||||
})
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
batch_size = int(data.get('batch_size', 0))
|
||||
tiles = []
|
||||
|
||||
if batch_size == 0 and is_last:
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': []
|
||||
})
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
if batch_size > 0:
|
||||
padding = int(data.get('padding', 32))
|
||||
metadata_field = data.get('tiles_metadata')
|
||||
if metadata_field:
|
||||
if hasattr(metadata_field, 'file'):
|
||||
metadata_str = metadata_field.file.read().decode('utf-8')
|
||||
elif isinstance(metadata_field, (bytes, bytearray)):
|
||||
metadata_str = metadata_field.decode('utf-8')
|
||||
else:
|
||||
metadata_str = str(metadata_field)
|
||||
|
||||
metadata = json.loads(metadata_str)
|
||||
if len(metadata) != batch_size:
|
||||
return await handle_api_error(request, "Metadata length mismatch", 400)
|
||||
|
||||
tile_data_list = []
|
||||
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
continue
|
||||
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
if i < len(metadata):
|
||||
tile_meta = metadata[i]
|
||||
tile_idx = tile_meta.get('tile_idx', i)
|
||||
tile_info = {
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': tile_meta['x'],
|
||||
'y': tile_meta['y'],
|
||||
'extracted_width': tile_meta['extracted_width'],
|
||||
'extracted_height': tile_meta['extracted_height'],
|
||||
'padding': padding,
|
||||
'batch_idx': tile_meta.get('batch_idx', 0),
|
||||
'global_idx': tile_meta.get('global_idx', tile_idx)
|
||||
}
|
||||
tile_data_list.append(tile_info)
|
||||
|
||||
tile_data_list.sort(key=lambda x: x['tile_idx'])
|
||||
tiles.extend(tile_data_list)
|
||||
else:
|
||||
# Legacy
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
continue
|
||||
|
||||
tile_idx = int(data.get(f'tile_{i}_idx', i))
|
||||
x = int(data.get(f'tile_{i}_x', 0))
|
||||
y = int(data.get(f'tile_{i}_y', 0))
|
||||
extracted_width = int(data.get(f'tile_{i}_width', 512))
|
||||
extracted_height = int(data.get(f'tile_{i}_height', 512))
|
||||
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles.append({
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding
|
||||
})
|
||||
else:
|
||||
# Single tile legacy
|
||||
image_file = data.get('image')
|
||||
if not image_file:
|
||||
return await handle_api_error(request, "Missing image data", 400)
|
||||
|
||||
tile_idx = int(data.get('tile_idx', 0))
|
||||
x = int(data.get('x', 0))
|
||||
y = int(data.get('y', 0))
|
||||
extracted_width = int(data.get('extracted_width', 512))
|
||||
extracted_height = int(data.get('extracted_height', 512))
|
||||
padding = int(data.get('padding', 32))
|
||||
|
||||
img_data = image_file.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles = [{
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding,
|
||||
'batch_idx': 0,
|
||||
'global_idx': tile_idx
|
||||
}]
|
||||
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
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 for tile submission", 400)
|
||||
|
||||
q = job_data[JOB_QUEUE]
|
||||
if batch_size > 0 or len(tiles) > 0:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'tiles': tiles,
|
||||
'is_last': is_last
|
||||
})
|
||||
else:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': []
|
||||
})
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
# Note: Removed legacy /distributed/tile_complete endpoint. Use /distributed/submit_tiles.
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user