From 71f6bcbfa35b6710b10cc29071266e553906b3d7 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Fri, 18 Jul 2025 09:13:14 +1000 Subject: [PATCH] remove exessive logging, fixed tiled decode, reduce wait fo requeue --- distributed.py | 32 --------------- distributed_upscale.py | 93 ++++++++++++------------------------------ 2 files changed, 25 insertions(+), 100 deletions(-) diff --git a/distributed.py b/distributed.py index 7b53c41..46df518 100644 --- a/distributed.py +++ b/distributed.py @@ -1777,14 +1777,12 @@ class DistributedCollectorNode: if batch_size == 0: return - debug_log(f"Worker - Sending {batch_size} images in chunks of max {MAX_BATCH}") for start in range(0, batch_size, MAX_BATCH): chunk = image_batch[start:start + MAX_BATCH] chunk_size = chunk.shape[0] is_chunk_last = (start + chunk_size == batch_size) # True only for final chunk - debug_log(f"Worker - Processing chunk {start//MAX_BATCH + 1}, images {start} to {start + chunk_size - 1}") data = aiohttp.FormData() data.add_field('multi_job_id', multi_job_id) @@ -1806,21 +1804,11 @@ class DistributedCollectorNode: data.add_field(f'image_{j}', byte_io, filename=f'image_{j}.png', content_type='image/png') try: - # Estimate payload size for logging - payload_size = 0 - for field in data._fields: - if hasattr(field[2], 'read'): - field[2].seek(0) - payload_size += len(field[2].read()) - field[2].seek(0) - debug_log(f"Worker - Chunk payload size: {payload_size:,} bytes") session = await get_client_session() url = f"{master_url}/distributed/job_complete" - debug_log(f"Worker - Sending chunk of {chunk_size} images to URL: {url}") async with session.post(url, data=data) as response: response.raise_for_status() - debug_log(f"Worker - Successfully sent chunk of {chunk_size} images") except Exception as e: log(f"Worker - Failed to send chunk to master: {e}") debug_log(f"Worker - Full error details: URL={url}") @@ -1847,8 +1835,6 @@ class DistributedCollectorNode: # Ensure master images are contiguous images_on_cpu = ensure_contiguous(images_on_cpu) - # Debug: Check master tensor properties - debug_log(f"Master tensor - shape: {images_on_cpu.shape}, dtype: {images_on_cpu.dtype}, device: {images_on_cpu.device}") # Initialize storage for collected images worker_images = {} # Dict to store images by worker_id and index @@ -1869,13 +1855,11 @@ class DistributedCollectorNode: # Use a reasonable timeout for the first image timeout = WORKER_JOB_TIMEOUT - debug_log(f"Master - Starting collection loop, expecting {num_workers} workers") # Get queue size before starting async with prompt_server.distributed_jobs_lock: q = prompt_server.distributed_pending_jobs[multi_job_id] initial_size = q.qsize() - debug_log(f"Master - Queue size before collection: {initial_size}") # NEW: Initialize progress bar for workers (total = num_workers) p = ProgressBar(num_workers) @@ -1887,7 +1871,6 @@ class DistributedCollectorNode: q = prompt_server.distributed_pending_jobs[multi_job_id] current_size = q.qsize() - debug_log(f"Master - Waiting for queue item, timeout={timeout}s, queue size={current_size}") result = await asyncio.wait_for(q.get(), timeout=timeout) worker_id = result['worker_id'] @@ -1906,9 +1889,6 @@ class DistributedCollectorNode: for idx, tensor in enumerate(tensors): worker_images[worker_id][idx] = tensor - # Debug: Check worker tensor properties - if collected_count == 0: # Only print once - debug_log(f"Master - Worker tensor - shape: {tensor.shape}, dtype: {tensor.dtype}, device: {tensor.device}") collected_count += len(tensors) else: @@ -1922,9 +1902,6 @@ class DistributedCollectorNode: worker_images[worker_id] = {} worker_images[worker_id][image_index] = tensor - # Debug: Check worker tensor properties - if collected_count == 0: # Only print once - debug_log(f"Master - Worker tensor - shape: {tensor.shape}, dtype: {tensor.dtype}, device: {tensor.device}") collected_count += 1 @@ -1934,9 +1911,7 @@ class DistributedCollectorNode: if is_last: workers_done.add(worker_id) p.update(1) # +1 per completed worker - debug_log(f"Master - Worker {worker_id} done. Collected {len(worker_images[worker_id])} images") else: - debug_log(f"Master - Collected image {image_index + 1} from worker {worker_id}") except asyncio.TimeoutError: missing_workers = set(str(w) for w in enabled_workers) - workers_done @@ -1947,7 +1922,6 @@ class DistributedCollectorNode: if multi_job_id in prompt_server.distributed_pending_jobs: final_q = prompt_server.distributed_pending_jobs[multi_job_id] final_size = final_q.qsize() - debug_log(f"Master - Queue size at timeout: {final_size}") # Try to drain any remaining items remaining_items = [] @@ -1959,7 +1933,6 @@ class DistributedCollectorNode: break if remaining_items: - debug_log(f"Master - Found {len(remaining_items)} items in queue after timeout!") # Process them for item in remaining_items: worker_id = item['worker_id'] @@ -1990,15 +1963,11 @@ class DistributedCollectorNode: if is_last: workers_done.add(worker_id) p.update(1) # +1 here too - debug_log(f"Master - Worker {worker_id} done (found in timeout drain)") else: log(f"Master - Queue {multi_job_id} no longer exists!") break total_collected = sum(len(imgs) for imgs in worker_images.values()) - debug_log(f"Master - Collection complete. Received {total_collected} images from {len(workers_done)} workers") - debug_log(f"Master - Workers done: {workers_done}, Enabled workers: {enabled_workers}") - debug_log(f"Master - Worker images keys: {list(worker_images.keys())}") # Clean up job queue async with prompt_server.distributed_jobs_lock: @@ -2038,7 +2007,6 @@ class DistributedCollectorNode: return (combined,) except Exception as e: log(f"Master - Error combining images: {e}") - debug_log(f"Master - Tensor shapes: {[t.shape for t in cpu_tensors]}") # Return just the master images as fallback return (images,) diff --git a/distributed_upscale.py b/distributed_upscale.py index dc3356c..fefa0cd 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -572,8 +572,6 @@ class UltimateSDUpscaleDistributed: """Initialize the job queue for collecting tiles.""" prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: - debug_log(f"UltimateSDUpscale Master - _init_job_queue: Checking if {multi_job_id} exists in distributed_pending_tile_jobs") - debug_log(f"UltimateSDUpscale Master - _init_job_queue: Current jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}") if multi_job_id not in prompt_server.distributed_pending_tile_jobs: prompt_server.distributed_pending_tile_jobs[multi_job_id] = { 'queue': asyncio.Queue(), @@ -581,9 +579,6 @@ class UltimateSDUpscaleDistributed: 'completed_tiles': {} # For static mode data } debug_log(f"UltimateSDUpscale Master - Initialized job queue for {multi_job_id}") - debug_log(f"UltimateSDUpscale Master - _init_job_queue: Created new queue for {multi_job_id}") - else: - debug_log(f"UltimateSDUpscale Master - _init_job_queue: Queue already exists for {multi_job_id}") async def _async_collect_worker_tiles(self, multi_job_id, num_workers): """Async helper to collect tiles from workers.""" @@ -606,7 +601,6 @@ class UltimateSDUpscaleDistributed: debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {num_workers} workers to complete") while len(workers_done) < num_workers: - debug_log(f"UltimateSDUpscale Master - Loop status: {len(workers_done)}/{num_workers} workers done, {len(collected_tiles)} tiles collected, workers_done set: {workers_done}") try: result = await asyncio.wait_for(q.get(), timeout=timeout) worker_id = result['worker_id'] @@ -651,12 +645,9 @@ class UltimateSDUpscaleDistributed: if is_last: workers_done.add(worker_id) debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed") - debug_log(f"UltimateSDUpscale Master - Worker '{worker_id}' marked as done. Total workers done: {len(workers_done)}/{num_workers}") except asyncio.TimeoutError: log(f"UltimateSDUpscale Master - Timeout waiting for tiles") - debug_log(f"UltimateSDUpscale Master - Final status: {len(workers_done)}/{num_workers} workers done, {len(collected_tiles)} tiles collected") - debug_log(f"UltimateSDUpscale Master - Workers that completed: {workers_done}") break debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_tiles)} tiles from {len(workers_done)} workers") @@ -724,7 +715,7 @@ class UltimateSDUpscaleDistributed: job_data_check = prompt_server.distributed_pending_tile_jobs[multi_job_id] # Check for timed out workers (60 second timeout) for worker, last_heartbeat in list(job_data_check.get('worker_status', {}).items()): - if current_time - last_heartbeat > 60: + if current_time - last_heartbeat > 30: log(f"UltimateSDUpscale Master Dynamic - Worker {worker} timed out") # Requeue assigned images from this worker for idx in job_data_check.get('assigned_to_workers', {}).get(worker, []): @@ -875,22 +866,6 @@ class UltimateSDUpscaleDistributed: 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}") - # Log conditioning details - debug_log(f"[process_tile] Positive conditioning items: {len(positive)}") - debug_log(f"[process_tile] Negative conditioning items: {len(negative)}") - - # Check for ControlNet in conditioning - for i, (emb, cond_dict) in enumerate(positive): - if 'control' in cond_dict: - control = cond_dict['control'] - debug_log(f"[process_tile] Found ControlNet in positive[{i}]") - debug_log(f"[process_tile] Control id: {id(control)}") - if hasattr(control, 'control_model'): - debug_log(f"[process_tile] control_model id: {id(control.control_model)}") - if hasattr(control, 'cond_hint_original'): - debug_log(f"[process_tile] cond_hint_original shape: {control.cond_hint_original.shape}") - debug_log(f"[process_tile] Expected to use hint at index {batch_idx} from batch") - debug_log(f"[process_tile] Embedding shape: {emb.shape}") # Import here to avoid circular dependencies from nodes import common_ksampler, VAEEncode, VAEDecode @@ -930,26 +905,16 @@ class UltimateSDUpscaleDistributed: positive_cropped = positive_tile negative_cropped = negative_tile - # Encode to latent - if tiled_decode and tiled_vae_available: - latent = VAEEncodeTiled().encode(vae, clean_tensor)[0] - else: - latent = VAEEncode().encode(vae, clean_tensor)[0] - - # Log latent shape - if 'samples' in latent: - debug_log(f"[process_tile] Latent shape: {latent['samples'].shape}") + # Encode to latent (always non-tiled, matching original node) + latent = VAEEncode().encode(vae, clean_tensor)[0] # Sample - debug_log(f"[process_tile] About to call common_ksampler with batch_idx={batch_idx}") - debug_log(f"[process_tile] Model id: {id(model)}") samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive_cropped, negative_cropped, latent, denoise=denoise)[0] - debug_log(f"[process_tile] common_ksampler completed") # Decode back to image if tiled_decode and tiled_vae_available: - image = VAEDecodeTiled().decode(vae, samples)[0] + image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0] else: image = VAEDecode().decode(vae, samples)[0] @@ -1032,8 +997,6 @@ class UltimateSDUpscaleDistributed: chunk_size = len(chunk) is_chunk_last = (start + chunk_size == total_tiles) # True only for final chunk - debug_log(f"UltimateSDUpscale Worker - Processing chunk {start//MAX_BATCH + 1}, tiles {start} to {start + chunk_size - 1}") - data = aiohttp.FormData() data.add_field('multi_job_id', multi_job_id) data.add_field('worker_id', str(worker_id)) @@ -1072,14 +1035,6 @@ class UltimateSDUpscaleDistributed: # Add image field data.add_field(f'tile_{j}', byte_io, filename=f'tile_{j}.png', content_type='image/png') - # Estimate payload size for logging - payload_size = 0 - for field in data._fields: - if hasattr(field[2], 'read'): - field[2].seek(0) - payload_size += len(field[2].read()) - field[2].seek(0) - debug_log(f"UltimateSDUpscale Worker - Chunk payload size: {payload_size:,} bytes") # Retry logic with exponential backoff max_retries = 5 @@ -1090,16 +1045,12 @@ class UltimateSDUpscaleDistributed: session = await get_client_session() url = f"{master_url}/distributed/tile_complete" - debug_log(f"UltimateSDUpscale Worker - Sending chunk of {chunk_size} tiles to {url}, attempt {attempt + 1}") - async with session.post(url, data=data) as response: response.raise_for_status() - debug_log(f"UltimateSDUpscale Worker - Successfully sent chunk of {chunk_size} tiles") break # Success, move to next chunk except Exception as e: if attempt < max_retries - 1: - debug_log(f"UltimateSDUpscale Worker - Retry {attempt + 1}/{max_retries} after error: {e}") await asyncio.sleep(retry_delay) retry_delay *= 2 # Exponential backoff else: @@ -1260,6 +1211,7 @@ class UltimateSDUpscaleDistributed: sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, width, height, tiled_decode, batch_idx=image_idx ) + # Note: No per-tile drain here – that's what makes this "per-image" result_images[image_idx] = local_image @@ -1268,6 +1220,22 @@ class UltimateSDUpscaleDistributed: 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"Progress: {completed_now}/{batch_size} images completed") else: # Queue empty: collect any queued worker results to update progress drained_count = run_async_in_server_loop( @@ -1291,11 +1259,9 @@ class UltimateSDUpscaleDistributed: timeout=1.0 ) + log(f"Progress: {completed_now}/{batch_size} images completed") - debug_log(f"Queue empty, completed count: {completed_now}/{batch_size}") - if completed_now >= batch_size: - debug_log("All images completed, exiting master loop") break # Check if there are pending images in the queue (could be requeued) @@ -1305,7 +1271,6 @@ class UltimateSDUpscaleDistributed: ) if pending_count > 0: - debug_log(f"Found {pending_count} pending images (possibly requeued), continuing to process") consecutive_retries = 0 # Reset retries since there's work to do continue @@ -1314,7 +1279,7 @@ class UltimateSDUpscaleDistributed: log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.") break # Force exit to collection phase - log("Queue empty temporarily, waiting 2s before retry") + debug_log("Waiting for workers") time.sleep(2) debug_log(f"Master processed {processed_count} images locally") @@ -1455,7 +1420,7 @@ class UltimateSDUpscaleDistributed: # Check for timed out workers (60 second timeout) for worker, last_heartbeat in list(job_data.get('worker_status', {}).items()): - if current_time - last_heartbeat > 60: + if current_time - last_heartbeat > 30: log(f"Worker {worker} timed out") # Requeue assigned images from this worker for idx in job_data.get('assigned_to_workers', {}).get(worker, []): @@ -1701,9 +1666,8 @@ class UltimateSDUpscaleDistributed: async with session.post(url, json=data) as response: response.raise_for_status() - debug_log(f"Worker {worker_id} sent heartbeat") except Exception as e: - debug_log(f"Worker {worker_id} heartbeat failed: {e}") + pass async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id): """Send completion signal to master in dynamic mode.""" @@ -2062,9 +2026,6 @@ async def tile_complete_endpoint(request): # Put tiles into queue async with prompt_server.distributed_tile_jobs_lock: - debug_log(f"UltimateSDUpscale - tile_complete: Checking distributed_pending_tile_jobs for job {multi_job_id}") - debug_log(f"UltimateSDUpscale - tile_complete: Current jobs in distributed_pending_tile_jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}") - if multi_job_id in prompt_server.distributed_pending_tile_jobs: job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] if not isinstance(job_data, dict) or 'mode' not in job_data: @@ -2080,7 +2041,6 @@ async def tile_complete_endpoint(request): 'tiles': tiles, 'is_last': is_last }) - debug_log(f"UltimateSDUpscale - Received batch of {len(tiles)} tiles for job {multi_job_id} from worker {worker_id}") else: # Put single tile (backward compat) tile_data = tiles[0] @@ -2095,12 +2055,9 @@ async def tile_complete_endpoint(request): 'padding': tile_data['padding'], 'is_last': is_last }) - debug_log(f"UltimateSDUpscale - Received single tile {tile_data['tile_idx']} for job {multi_job_id} from worker {worker_id}") return web.json_response({"status": "success"}) else: - debug_log(f"UltimateSDUpscale - API Error: Job {multi_job_id} not found in distributed_pending_tile_jobs") - debug_log(f"UltimateSDUpscale - API Error: Available jobs: {list(prompt_server.distributed_pending_tile_jobs.keys())}") return await handle_api_error(request, "Job not found or already complete", 404) except Exception as e: return await handle_api_error(request, e, 500)