2207 lines
110 KiB
Python
2207 lines
110 KiB
Python
# Import statements (original + new)
|
||
import torch
|
||
import numpy as np
|
||
from PIL import Image, ImageFilter, ImageDraw
|
||
import json
|
||
import asyncio
|
||
import aiohttp
|
||
from aiohttp import web
|
||
import io
|
||
import server
|
||
import math
|
||
import os
|
||
import time
|
||
from typing import List, Tuple
|
||
from functools import wraps
|
||
|
||
# Import ComfyUI modules
|
||
import comfy.samplers
|
||
import comfy.model_management
|
||
|
||
# Import shared utilities
|
||
from .utils.logging import debug_log, log
|
||
from .utils.config import CONFIG_FILE
|
||
from .utils.image import tensor_to_pil, pil_to_tensor
|
||
from .utils.network import get_server_port, get_server_loop, get_client_session, handle_api_error
|
||
from .utils.async_helpers import run_async_in_server_loop
|
||
from .utils.constants import (
|
||
TILE_COLLECTION_TIMEOUT, TILE_WAIT_TIMEOUT, TILE_TRANSFER_TIMEOUT,
|
||
QUEUE_INIT_TIMEOUT, TILE_SEND_TIMEOUT, MAX_BATCH, HEARTBEAT_TIMEOUT
|
||
)
|
||
|
||
# Import for controller support
|
||
from .utils.usdu_utils import crop_cond
|
||
from .utils.usdu_managment import (
|
||
clone_conditioning, ensure_tile_jobs_initialized,
|
||
# Job management functions
|
||
_init_job_queue, _distribute_tasks, _get_next_task, _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
|
||
)
|
||
|
||
|
||
# Note: MAX_BATCH and HEARTBEAT_TIMEOUT are imported from utils.constants
|
||
# They can be overridden via environment variables:
|
||
# - COMFYUI_MAX_BATCH (default: 20)
|
||
# - COMFYUI_HEARTBEAT_TIMEOUT (default: 120)
|
||
|
||
# Sync wrapper decorator for async methods
|
||
def sync_wrapper(async_func):
|
||
"""Decorator to wrap async methods for synchronous execution."""
|
||
@wraps(async_func)
|
||
def sync_func(self, *args, **kwargs):
|
||
# Use run_async_in_server_loop for ComfyUI compatibility
|
||
return run_async_in_server_loop(
|
||
async_func(self, *args, **kwargs),
|
||
timeout=600.0 # 10 minute timeout for long operations
|
||
)
|
||
return sync_func
|
||
|
||
# Note: tensor_to_pil and pil_to_tensor are imported from utils.image
|
||
|
||
class UltimateSDUpscaleDistributed:
|
||
"""
|
||
Distributed version of Ultimate SD Upscale (No Upscale).
|
||
|
||
Supports three processing modes:
|
||
1. Single GPU: No workers available, process everything locally
|
||
2. Static Mode: Small batches, distributes tiles across workers (flattened)
|
||
3. Dynamic Mode: Large batches, assigns whole images to workers dynamically
|
||
|
||
Features:
|
||
- Multi-mode batch handling for efficient video/image upscaling
|
||
- Tiled VAE support for memory efficiency
|
||
- Dynamic load balancing for large batches
|
||
- Backward compatible with single-image workflows
|
||
|
||
Environment Variables:
|
||
- COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
|
||
- COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB)
|
||
|
||
Threshold: dynamic_threshold input controls mode switch (default 8)
|
||
"""
|
||
|
||
def __init__(self):
|
||
"""Initialize the node and ensure persistent storage exists."""
|
||
# Pre-initialize the persistent storage on node creation
|
||
ensure_tile_jobs_initialized()
|
||
debug_log("UltimateSDUpscaleDistributed - Node initialized")
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"upscaled_image": ("IMAGE",),
|
||
"model": ("MODEL",),
|
||
"positive": ("CONDITIONING",),
|
||
"negative": ("CONDITIONING",),
|
||
"vae": ("VAE",),
|
||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
||
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
|
||
"sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
|
||
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
|
||
"denoise": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"tile_width": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}),
|
||
"tile_height": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}),
|
||
"padding": ("INT", {"default": 32, "min": 0, "max": 256, "step": 8}),
|
||
"mask_blur": ("INT", {"default": 8, "min": 0, "max": 256}),
|
||
"force_uniform_tiles": ("BOOLEAN", {"default": True}),
|
||
"tiled_decode": ("BOOLEAN", {"default": False}),
|
||
"static_distribution": ("BOOLEAN", {"default": False, "label": "Use Static Distribution (legacy)"}),
|
||
},
|
||
"hidden": {
|
||
"multi_job_id": ("STRING", {"default": ""}),
|
||
"is_worker": ("BOOLEAN", {"default": False}),
|
||
"master_url": ("STRING", {"default": ""}),
|
||
"enabled_worker_ids": ("STRING", {"default": "[]"}),
|
||
"worker_id": ("STRING", {"default": ""}),
|
||
"tile_indices": ("STRING", {"default": ""}), # Unused - kept for compatibility
|
||
"dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
FUNCTION = "run"
|
||
CATEGORY = "image/upscaling"
|
||
|
||
@classmethod
|
||
def IS_CHANGED(cls, **kwargs):
|
||
"""Force re-execution."""
|
||
return float("nan") # Always re-execute
|
||
|
||
|
||
def run(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, static_distribution=False,
|
||
multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]",
|
||
worker_id="", tile_indices="", dynamic_threshold=8):
|
||
"""Entry point - runs SYNCHRONOUSLY like Ultimate SD Upscaler."""
|
||
if not multi_job_id:
|
||
# No distributed processing, run single GPU version
|
||
return self.process_single_gpu(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)
|
||
|
||
if is_worker:
|
||
# Worker mode: process tiles synchronously
|
||
return self.process_worker(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, static_distribution)
|
||
else:
|
||
# Master mode: distribute and collect synchronously
|
||
return self.process_master(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_worker_ids,
|
||
dynamic_threshold, static_distribution)
|
||
|
||
def process_worker(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, static_distribution=False):
|
||
"""Unified worker processing - handles both static and dynamic modes."""
|
||
# Get batch size to determine mode
|
||
batch_size = upscaled_image.shape[0]
|
||
|
||
# Ensure mode consistency across master/workers via shared threshold
|
||
# Determine mode (must match master's logic)
|
||
enabled_workers = json.loads(enabled_worker_ids)
|
||
num_workers = len(enabled_workers)
|
||
|
||
mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
|
||
|
||
debug_log(f"UltimateSDUpscale Worker - Mode: {mode}, batch_size: {batch_size}")
|
||
|
||
if mode == "dynamic":
|
||
return self.process_worker_dynamic(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)
|
||
|
||
# Static mode - enhanced with health monitoring and retry logic
|
||
return self._process_worker_static_sync(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, static_distribution)
|
||
|
||
def process_master(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_worker_ids,
|
||
dynamic_threshold, static_distribution=False):
|
||
"""Unified master processing with enhanced monitoring and failure handling."""
|
||
# 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)
|
||
num_tiles_per_image = len(all_tiles)
|
||
|
||
debug_log(f"UltimateSDUpscale Master - Tiles per image: {num_tiles_per_image}")
|
||
|
||
# Parse enabled workers
|
||
enabled_workers = json.loads(enabled_worker_ids)
|
||
num_workers = len(enabled_workers)
|
||
|
||
# Determine processing mode
|
||
mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
|
||
|
||
debug_log(f"UltimateSDUpscale Master - Mode: {mode}, batch_size: {batch_size}, workers: {num_workers}")
|
||
|
||
if mode == "single_gpu":
|
||
# No workers, process all tiles locally
|
||
return self.process_single_gpu(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)
|
||
|
||
elif mode == "dynamic":
|
||
# Dynamic mode for large batches
|
||
return self.process_master_dynamic(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)
|
||
|
||
# Static mode - enhanced with unified job management
|
||
return self._process_master_static_sync(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_distribution)
|
||
|
||
def _get_worker_global_indices(self, total_tiles, enabled_workers, worker_id):
|
||
"""Calculate which global tile indices are assigned to a specific worker in flattened mode."""
|
||
num_workers = len(enabled_workers)
|
||
|
||
debug_log(f"UltimateSDUpscale Worker - Worker ID: {worker_id}, Enabled workers: {enabled_workers}")
|
||
|
||
try:
|
||
worker_index = enabled_workers.index(worker_id)
|
||
except ValueError:
|
||
log(f"UltimateSDUpscale Worker - Worker {worker_id} not found in enabled list {enabled_workers}")
|
||
return []
|
||
|
||
global_indices = list(range(total_tiles))
|
||
all_assignments = _distribute_tasks(global_indices, num_workers + 1)
|
||
return all_assignments[worker_index + 1] # +1 because 0 is master
|
||
|
||
def _get_master_global_indices(self, total_tiles, num_workers):
|
||
"""Calculate which global tile indices are assigned to the master in flattened mode."""
|
||
global_indices = list(range(total_tiles))
|
||
all_assignments = _distribute_tasks(global_indices, num_workers + 1)
|
||
return all_assignments[0]
|
||
|
||
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, 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
|
||
)
|
||
|
||
processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae,
|
||
seed + tile_idx, 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
|
||
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
|
||
|
||
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):
|
||
"""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)
|
||
|
||
# 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):
|
||
"""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)
|
||
|
||
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()))
|
||
|
||
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")
|
||
|
||
collected_results = {}
|
||
workers_done = set()
|
||
timeout = 120.0 if mode == 'dynamic' else 60.0 # Longer timeout for full images
|
||
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 timeout for regular checks in dynamic mode
|
||
wait_timeout = 10.0 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
|
||
collected_results[key] = {
|
||
'tensor': tile_data['tensor'],
|
||
'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)
|
||
}
|
||
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 Dynamic - 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")
|
||
|
||
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 Dynamic - 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")
|
||
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")
|
||
|
||
# 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
|
||
|
||
# Keep compatibility wrappers for existing code
|
||
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)
|
||
|
||
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 for the image.
|
||
|
||
Tiles are placed at grid positions without overlap in their placement.
|
||
The overlap happens during extraction with padding.
|
||
|
||
If force_uniform_tiles is True, tiles are distributed evenly across the image."""
|
||
if force_uniform_tiles:
|
||
# Calculate number of tiles needed
|
||
num_tiles_x = math.ceil(image_width / tile_width)
|
||
num_tiles_y = math.ceil(image_height / tile_height)
|
||
|
||
# Calculate uniform spacing
|
||
tiles = []
|
||
|
||
# Handle X dimension
|
||
if num_tiles_x == 1:
|
||
tiles_x = [0]
|
||
else:
|
||
# Calculate step to distribute tiles evenly
|
||
step_x = (image_width - tile_width) / (num_tiles_x - 1)
|
||
# Check for excessive overlap
|
||
if step_x < tile_width / 2:
|
||
debug_log(f"Warning: Uniform tiles would have excessive overlap (step_x={step_x}, tile_width={tile_width}). Falling back to non-uniform mode.")
|
||
force_uniform_tiles = False
|
||
else:
|
||
tiles_x = [round(i * step_x) for i in range(num_tiles_x)]
|
||
# Ensure tiles are aligned to 8-pixel boundaries
|
||
tiles_x = [self.round_to_multiple(x, 8) for x in tiles_x]
|
||
|
||
# Handle Y dimension (if still using uniform tiles)
|
||
if force_uniform_tiles:
|
||
if num_tiles_y == 1:
|
||
tiles_y = [0]
|
||
else:
|
||
# Calculate step to distribute tiles evenly
|
||
step_y = (image_height - tile_height) / (num_tiles_y - 1)
|
||
# Check for excessive overlap
|
||
if step_y < tile_height / 2:
|
||
debug_log(f"Warning: Uniform tiles would have excessive overlap (step_y={step_y}, tile_height={tile_height}). Falling back to non-uniform mode.")
|
||
force_uniform_tiles = False
|
||
else:
|
||
tiles_y = [round(i * step_y) for i in range(num_tiles_y)]
|
||
# Ensure tiles are aligned to 8-pixel boundaries
|
||
tiles_y = [self.round_to_multiple(y, 8) for y in tiles_y]
|
||
|
||
# If still using uniform tiles, generate positions
|
||
if force_uniform_tiles:
|
||
# Generate all tile positions
|
||
for y in tiles_y:
|
||
for x in tiles_x:
|
||
tiles.append((x, y))
|
||
|
||
return tiles
|
||
|
||
# Fall through to non-uniform mode if needed
|
||
if not force_uniform_tiles:
|
||
# Original non-uniform tile calculation
|
||
tiles = []
|
||
for y in range(0, image_height, tile_height):
|
||
for x in range(0, image_width, tile_width):
|
||
tiles.append((x, y))
|
||
return tiles
|
||
|
||
def extract_tile_with_padding(self, image: torch.Tensor, x: int, y: int,
|
||
tile_width: int, tile_height: int, padding: int) -> Tuple[torch.Tensor, int, int, int, int]:
|
||
"""Extract a tile from the image with padding."""
|
||
_, h, w, _ = image.shape
|
||
|
||
# Calculate extraction bounds with padding
|
||
x1 = max(0, x - padding)
|
||
y1 = max(0, y - padding)
|
||
x2 = min(w, x + tile_width + padding)
|
||
y2 = min(h, y + tile_height + padding)
|
||
|
||
# Store the actual extracted size before resizing
|
||
extracted_width = x2 - x1
|
||
extracted_height = y2 - y1
|
||
|
||
# Debug logging (uncomment if needed)
|
||
# debug_log(f"[Extract] Tile at ({x}, {y}) -> Extract from ({x1}, {y1}) to ({x2}, {y2}), size: {extracted_width}x{extracted_height}")
|
||
|
||
# Extract tile
|
||
tile = image[:, y1:y2, x1:x2, :]
|
||
|
||
# Convert to PIL using utility function
|
||
tile_pil = tensor_to_pil(tile, 0)
|
||
|
||
# Resize to target dimensions for processing
|
||
tile_pil = tile_pil.resize((tile_width, tile_height), Image.LANCZOS)
|
||
|
||
# Convert back to tensor using utility function
|
||
tile_tensor = pil_to_tensor(tile_pil)
|
||
|
||
# Move to same device as original image
|
||
if image.is_cuda:
|
||
tile_tensor = tile_tensor.cuda()
|
||
|
||
return tile_tensor, 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 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
|
||
|
||
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 _process_worker_static_async(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):
|
||
"""Enhanced static mode worker with health monitoring and retry logic."""
|
||
# Round tile dimensions
|
||
tile_width = self.round_to_multiple(tile_width)
|
||
tile_height = self.round_to_multiple(tile_height)
|
||
|
||
# Calculate all tile positions
|
||
batch_size = upscaled_image.shape[0]
|
||
_, height, width, _ = upscaled_image.shape
|
||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||
|
||
# Calculate total tiles across all images
|
||
num_tiles_per_image = len(all_tiles)
|
||
total_tiles = batch_size * num_tiles_per_image
|
||
|
||
# Get initial assigned global tile indices for this worker
|
||
assigned_global_indices = self._get_worker_global_indices(total_tiles, enabled_workers, worker_id)
|
||
|
||
debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} initially assigned {len(assigned_global_indices)} tiles")
|
||
|
||
# Process initially assigned tiles
|
||
processed_tiles = []
|
||
sliced_conditioning_cache = {}
|
||
|
||
# Process assigned tiles with heartbeat after each
|
||
for global_idx in assigned_global_indices:
|
||
# Process the tile
|
||
processed_tile = await self._process_single_tile(
|
||
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, sliced_conditioning_cache
|
||
)
|
||
|
||
if processed_tile:
|
||
processed_tiles.append(processed_tile)
|
||
|
||
# Send heartbeat after each tile
|
||
try:
|
||
await _send_heartbeat_to_master(multi_job_id, master_url, worker_id)
|
||
except Exception as e:
|
||
debug_log(f"Heartbeat failed: {e}")
|
||
|
||
# Send initial batch of processed tiles
|
||
if processed_tiles:
|
||
await self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id)
|
||
processed_tiles = []
|
||
|
||
# Check for requeued tiles from failed workers
|
||
debug_log(f"Worker {worker_id} checking for requeued tiles...")
|
||
while True:
|
||
# Request next task from pending queue
|
||
task_id = await _get_next_task(multi_job_id)
|
||
if task_id is None:
|
||
break
|
||
|
||
debug_log(f"Worker {worker_id} processing requeued tile {task_id}")
|
||
|
||
# Process the requeued tile
|
||
processed_tile = await self._process_single_tile(
|
||
task_id, 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, sliced_conditioning_cache
|
||
)
|
||
|
||
if processed_tile:
|
||
processed_tiles.append(processed_tile)
|
||
|
||
# Send heartbeat
|
||
try:
|
||
await _send_heartbeat_to_master(multi_job_id, master_url, worker_id)
|
||
except Exception as e:
|
||
debug_log(f"Heartbeat failed: {e}")
|
||
|
||
# Send tiles in batches
|
||
if len(processed_tiles) >= MAX_BATCH:
|
||
await self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id)
|
||
processed_tiles = []
|
||
|
||
# Send any remaining tiles
|
||
if processed_tiles:
|
||
await self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id)
|
||
|
||
debug_log(f"Worker {worker_id} completed all assigned and requeued tiles")
|
||
return (upscaled_image,)
|
||
|
||
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, static_distribution=False):
|
||
"""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 = {}
|
||
|
||
if static_distribution:
|
||
# Legacy mode: Use pre-assigned tiles
|
||
debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} using static tile assignment")
|
||
worker_global_indices = self._get_worker_global_indices(total_tiles, enabled_workers, worker_id)
|
||
debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} processing {len(worker_global_indices)} pre-assigned tiles")
|
||
|
||
# Process assigned tiles synchronously
|
||
for global_idx in worker_global_indices:
|
||
batch_idx = global_idx // num_tiles_per_image
|
||
tile_idx = global_idx % num_tiles_per_image
|
||
|
||
if batch_idx >= batch_size:
|
||
continue
|
||
|
||
# Get or create sliced conditioning
|
||
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]
|
||
|
||
# Process tile synchronously
|
||
processed_tile = self._process_single_tile(
|
||
global_idx, num_tiles_per_image, upscaled_image, all_tiles,
|
||
model, positive_sliced, negative_sliced, vae, seed, steps, cfg, sampler_name,
|
||
scheduler, denoise, tiled_decode, tile_width, tile_height, padding,
|
||
width, height, sliced_conditioning_cache
|
||
)
|
||
|
||
if processed_tile:
|
||
processed_tiles.append(processed_tile)
|
||
|
||
# Send heartbeat (async operation)
|
||
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"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 = []
|
||
|
||
# Send 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
|
||
)
|
||
processed_tiles = []
|
||
|
||
# Check for requeued tiles from failed workers
|
||
debug_log(f"Worker {worker_id} checking for requeued tiles...")
|
||
while True:
|
||
# Request next task from pending queue (async operation)
|
||
task_id = run_async_in_server_loop(
|
||
_get_next_task(multi_job_id),
|
||
timeout=5.0
|
||
)
|
||
if task_id is None:
|
||
break
|
||
|
||
debug_log(f"Worker {worker_id} processing requeued tile {task_id}")
|
||
|
||
# Process the requeued tile synchronously
|
||
batch_idx = task_id // num_tiles_per_image
|
||
tile_idx = task_id % num_tiles_per_image
|
||
|
||
if batch_idx >= batch_size:
|
||
continue
|
||
|
||
# Get or create sliced conditioning
|
||
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]
|
||
|
||
processed_tile = self._process_single_tile(
|
||
task_id, num_tiles_per_image, upscaled_image, all_tiles,
|
||
model, positive_sliced, negative_sliced, vae, seed, steps, cfg, sampler_name,
|
||
scheduler, denoise, tiled_decode, tile_width, tile_height, padding,
|
||
width, height, sliced_conditioning_cache
|
||
)
|
||
|
||
if processed_tile:
|
||
processed_tiles.append(processed_tile)
|
||
|
||
# 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"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 = []
|
||
|
||
else:
|
||
# Dynamic queue mode
|
||
debug_log(f"UltimateSDUpscale Worker - Worker {worker_id} ready to process tiles dynamically")
|
||
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"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 tiles from queue
|
||
while True:
|
||
# Request a tile to process
|
||
tile_idx, estimated_remaining = 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"UltimateSDUpscale Worker - No more tiles to process")
|
||
break
|
||
|
||
debug_log(f"UltimateSDUpscale Worker - Assigned tile {tile_idx} to worker {worker_id}")
|
||
processed_count += 1
|
||
|
||
# Calculate batch and local tile indices
|
||
batch_idx = tile_idx // num_tiles_per_image
|
||
local_tile_idx = tile_idx % num_tiles_per_image
|
||
|
||
if batch_idx >= batch_size:
|
||
continue
|
||
|
||
# Get or create sliced conditioning
|
||
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]
|
||
|
||
# Process tile synchronously
|
||
processed_tile = self._process_single_tile(
|
||
tile_idx, num_tiles_per_image, upscaled_image, all_tiles,
|
||
model, positive_sliced, negative_sliced, vae, seed, steps, cfg, sampler_name,
|
||
scheduler, denoise, tiled_decode, tile_width, tile_height, padding,
|
||
width, height, sliced_conditioning_cache
|
||
)
|
||
|
||
if processed_tile:
|
||
processed_tiles.append(processed_tile)
|
||
|
||
# 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"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 = []
|
||
|
||
# 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()
|
||
stall_count = 0
|
||
max_stall_count = 10 # Allow 10 iterations with no progress before returning
|
||
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")
|
||
stall_count = 0 # Reset stall count when we requeue
|
||
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
|
||
|
||
# Check if we're making progress
|
||
if completed_count > last_completed_count:
|
||
stall_count = 0 # Reset stall count on progress
|
||
last_completed_count = completed_count
|
||
else:
|
||
stall_count += 1
|
||
|
||
# If no progress for too long, check if we need to process locally
|
||
if stall_count >= max_stall_count:
|
||
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)
|
||
# If there are pending tasks but no active workers, or we're stalled
|
||
if pending_queue and not pending_queue.empty():
|
||
log(f"Collection stalled with {expected_total - completed_count} tasks remaining. Returning for local processing.")
|
||
break
|
||
|
||
# Wait a bit before next check
|
||
await asyncio.sleep(0.5)
|
||
|
||
# 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_distribution=False):
|
||
"""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 = {}
|
||
master_indices = set() # Track which tiles master processed
|
||
master_processed_count = 0 # Track how many tiles master processed
|
||
|
||
if static_distribution:
|
||
# Legacy mode: Pre-assign tiles to master and workers
|
||
debug_log(f"UltimateSDUpscale Master - Using static tile distribution")
|
||
|
||
# Distribute tasks across master and workers
|
||
all_indices = list(range(total_tiles))
|
||
num_participants = len(enabled_workers) + 1
|
||
task_assignments = _distribute_tasks(all_indices, num_participants)
|
||
master_indices = set(task_assignments[0]) # Convert to set for efficient lookup
|
||
|
||
# Initialize job queue with task assignments (don't populate pending queue)
|
||
run_async_in_server_loop(
|
||
self._init_job_queue(multi_job_id, total_tiles, enabled_workers, batch_size, num_tiles_per_image, task_assignments),
|
||
timeout=10.0
|
||
)
|
||
|
||
debug_log(f"UltimateSDUpscale Master - Initialized static job queue for {total_tiles} tiles")
|
||
debug_log(f"UltimateSDUpscale Master - Processing {len(master_indices)} tiles locally")
|
||
master_processed_count = len(master_indices)
|
||
|
||
# Process master's tiles synchronously
|
||
for global_idx in master_indices:
|
||
# Check for user interruption
|
||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||
batch_idx = global_idx // num_tiles_per_image
|
||
tile_idx = global_idx % num_tiles_per_image
|
||
|
||
if batch_idx >= batch_size:
|
||
continue
|
||
|
||
# Get or create sliced conditioning
|
||
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]
|
||
|
||
# Use unique seed per image
|
||
image_seed = seed + batch_idx * 1000
|
||
|
||
# Process tile synchronously - no async context for tensor operations
|
||
result_images[batch_idx] = self._process_and_blend_tile(
|
||
tile_idx, all_tiles[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
|
||
)
|
||
|
||
# Mark as completed (async operation)
|
||
run_async_in_server_loop(
|
||
_mark_task_completed(multi_job_id, global_idx, {'batch_idx': batch_idx, 'tile_idx': tile_idx}),
|
||
timeout=5.0
|
||
)
|
||
else:
|
||
# Dynamic queue mode: All tiles go to pending queue
|
||
debug_log(f"UltimateSDUpscale Master - Using dynamic tile distribution")
|
||
|
||
# Initialize job queue with all tiles in pending queue for dynamic distribution
|
||
run_async_in_server_loop(
|
||
self._init_job_queue(multi_job_id, total_tiles, enabled_workers, batch_size, num_tiles_per_image),
|
||
timeout=10.0
|
||
)
|
||
|
||
debug_log(f"UltimateSDUpscale Master - Initialized static job queue for {total_tiles} tiles with dynamic distribution")
|
||
debug_log(f"UltimateSDUpscale Master - Processing tiles dynamically from queue")
|
||
|
||
# Process tiles by pulling from queue (similar to dynamic mode)
|
||
processed_count = 0
|
||
consecutive_no_tile = 0
|
||
max_consecutive_no_tile = 5
|
||
|
||
while processed_count < total_tiles:
|
||
# Check for user interruption
|
||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||
|
||
# Try to get a tile to process
|
||
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 += 1
|
||
|
||
# Calculate batch and local tile indices
|
||
batch_idx = tile_idx // num_tiles_per_image
|
||
local_tile_idx = tile_idx % num_tiles_per_image
|
||
|
||
if batch_idx >= batch_size:
|
||
continue
|
||
|
||
debug_log(f"Master processing tile {tile_idx} (batch {batch_idx}, tile {local_tile_idx})")
|
||
|
||
# Get or create sliced conditioning
|
||
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]
|
||
|
||
# Use unique seed per image
|
||
image_seed = seed + batch_idx * 1000
|
||
|
||
# Process tile synchronously
|
||
result_images[batch_idx] = self._process_and_blend_tile(
|
||
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
|
||
)
|
||
|
||
# Mark as completed
|
||
run_async_in_server_loop(
|
||
_mark_task_completed(multi_job_id, tile_idx, {'batch_idx': batch_idx, 'tile_idx': local_tile_idx}),
|
||
timeout=5.0
|
||
)
|
||
else:
|
||
consecutive_no_tile += 1
|
||
if consecutive_no_tile >= max_consecutive_no_tile:
|
||
# No more tiles available, break to start collecting
|
||
debug_log(f"Master processed {processed_count} tiles, moving to collection phase")
|
||
break
|
||
# Small wait before retry
|
||
time.sleep(0.1)
|
||
|
||
# Set master_processed_count for the collection phase
|
||
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:
|
||
collected_tasks = run_async_in_server_loop(
|
||
self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles),
|
||
timeout=300.0 # 5 minutes to allow for retries and requeuing
|
||
)
|
||
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
|
||
while True:
|
||
# Check for user interruption
|
||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||
|
||
# Get next task from pending queue
|
||
task_id = run_async_in_server_loop(
|
||
self._get_next_tile_index(multi_job_id),
|
||
timeout=5.0
|
||
)
|
||
|
||
if task_id is None:
|
||
break
|
||
|
||
batch_idx = task_id // num_tiles_per_image
|
||
local_tile_idx = task_id % num_tiles_per_image
|
||
|
||
if batch_idx >= batch_size:
|
||
continue
|
||
|
||
# Get or create sliced conditioning
|
||
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]
|
||
|
||
# Use unique seed per image
|
||
image_seed = seed + batch_idx * 1000
|
||
|
||
# Process tile synchronously
|
||
result_images[batch_idx] = self._process_and_blend_tile(
|
||
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
|
||
)
|
||
|
||
# Mark as completed and store in collected_tasks
|
||
run_async_in_server_loop(
|
||
_mark_task_completed(multi_job_id, task_id, {'batch_idx': batch_idx, 'tile_idx': local_tile_idx}),
|
||
timeout=5.0
|
||
)
|
||
# Add to collected_tasks so we skip it in the blending loop
|
||
collected_tasks[task_id] = {'batch_idx': batch_idx, 'tile_idx': local_tile_idx}
|
||
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 already processed by master
|
||
if static_distribution and global_idx in master_indices:
|
||
continue # Skip master's tiles in static mode
|
||
|
||
# Skip tiles that don't have tensor data (already processed)
|
||
if 'tensor' 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)
|
||
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)
|
||
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)
|
||
|
||
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, 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
|
||
)
|
||
|
||
# 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 + tile_idx, 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
|
||
}
|
||
|
||
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"UltimateSDUpscale Worker - Sending {total_tiles} tiles in chunks of max {MAX_BATCH}")
|
||
|
||
for start in range(0, total_tiles, MAX_BATCH):
|
||
chunk = processed_tiles[start:start + MAX_BATCH]
|
||
chunk_size = len(chunk)
|
||
is_chunk_last = (start + chunk_size == total_tiles) # True only for final chunk
|
||
|
||
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', str(is_chunk_last))
|
||
data.add_field('batch_size', str(chunk_size))
|
||
data.add_field('padding', str(padding))
|
||
|
||
# Chunk metadata: Keep absolute tile_idx and batch_idx if present
|
||
metadata = []
|
||
for tile_data in chunk:
|
||
meta = {
|
||
'tile_idx': tile_data['tile_idx'], # Original/absolute idx
|
||
'x': tile_data['x'],
|
||
'y': tile_data['y'],
|
||
'extracted_width': tile_data['extracted_width'],
|
||
'extracted_height': tile_data['extracted_height']
|
||
}
|
||
# Add batch_idx and global_idx if present (for batch processing)
|
||
if 'batch_idx' in tile_data:
|
||
meta['batch_idx'] = tile_data['batch_idx']
|
||
if 'global_idx' in tile_data:
|
||
meta['global_idx'] = tile_data['global_idx']
|
||
metadata.append(meta)
|
||
|
||
# Add JSON metadata
|
||
data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json')
|
||
|
||
# Add chunk images as 'tile_{j}' where j=0 to chunk_size-1
|
||
for j, tile_data in enumerate(chunk):
|
||
# Convert tensor to PIL
|
||
img = tensor_to_pil(tile_data['tile'], 0)
|
||
byte_io = io.BytesIO()
|
||
img.save(byte_io, format='PNG', compress_level=0)
|
||
byte_io.seek(0)
|
||
|
||
# Add image field
|
||
data.add_field(f'tile_{j}', byte_io, filename=f'tile_{j}.png', content_type='image/png')
|
||
|
||
|
||
# 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 # Success, move to next chunk
|
||
|
||
except Exception as e:
|
||
if attempt < max_retries - 1:
|
||
await asyncio.sleep(retry_delay)
|
||
retry_delay *= 2 # Exponential backoff
|
||
else:
|
||
log(f"UltimateSDUpscale Worker - Failed to send chunk after {max_retries} attempts: {e}")
|
||
raise
|
||
|
||
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)."""
|
||
# 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)
|
||
|
||
debug_log(f"UltimateSDUpscale - Processing {len(all_tiles)} tiles locally for batch of {batch_size}")
|
||
|
||
# Convert batch to PIL list
|
||
result_images = []
|
||
for b in range(batch_size):
|
||
image_pil = tensor_to_pil(upscaled_image[b:b+1], 0)
|
||
# Ensure RGB mode for consistency
|
||
if image_pil.mode != 'RGB':
|
||
image_pil = image_pil.convert('RGB')
|
||
result_images.append(image_pil.copy())
|
||
|
||
# Process each image in the batch
|
||
for batch_idx in range(batch_size):
|
||
debug_log(f"[process_single_gpu] Processing batch_idx {batch_idx}")
|
||
# 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[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]
|
||
|
||
# Process each tile for this image with pre-sliced conditioning
|
||
for tile_idx, tile_pos in enumerate(all_tiles):
|
||
# Use unique seed per image
|
||
image_seed = seed + batch_idx * 1000
|
||
result_images[batch_idx] = self._process_and_blend_tile(
|
||
tile_idx, tile_pos, 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
|
||
)
|
||
|
||
# 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()
|
||
|
||
return (result_tensor,)
|
||
|
||
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)
|
||
|
||
debug_log(f"UltimateSDUpscale Master Dynamic - Processing batch of {batch_size} images with {num_workers} workers")
|
||
|
||
# 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.")
|
||
|
||
# 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),
|
||
timeout=2.0
|
||
)
|
||
except Exception as e:
|
||
debug_log(f"UltimateSDUpscale Master Dynamic - 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, 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"Progress: {completed_now}/{batch_size} images completed")
|
||
|
||
# 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"Progress: {completed_now}/{batch_size} images completed")
|
||
|
||
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")
|
||
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=TILE_COLLECTION_TIMEOUT * 2
|
||
)
|
||
|
||
# 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 Dynamic - Job {multi_job_id} complete")
|
||
log(f"Completed processing all {batch_size} images")
|
||
return (result_tensor,)
|
||
|
||
|
||
async def _async_yield(self):
|
||
"""Simple async yield to allow event loop processing."""
|
||
await asyncio.sleep(0)
|
||
|
||
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=1.0)
|
||
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 _check_and_requeue_timed_out_workers(multi_job_id, batch_size)
|
||
|
||
|
||
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),
|
||
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, 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")
|
||
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"UltimateSDUpscale Worker Dynamic - 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."""
|
||
# 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)
|
||
return tile_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 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
|
||
|
||
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 _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
|
||
|
||
# Note: Using _send_heartbeat_to_master from usdu_managment.py instead of duplicate
|
||
|
||
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")
|
||
|
||
def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str:
|
||
"""Determines the processing mode based on batch size and worker count."""
|
||
if num_workers == 0:
|
||
return "single_gpu"
|
||
if batch_size <= dynamic_threshold:
|
||
return "static"
|
||
return "dynamic"
|
||
|
||
# Note: Using _distribute_tasks from usdu_managment.py instead of duplicate _distribute_items
|
||
|
||
# Ensure initialization before registering routes
|
||
ensure_tile_jobs_initialized()
|
||
|
||
# Node registration
|
||
NODE_CLASS_MAPPINGS = {
|
||
"UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
|
||
} |