upscaler: dynamic/static refactoring
This commit is contained in:
robertvoy
@@ -611,6 +611,79 @@ async def get_managed_workers_endpoint(request):
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/local-worker-status")
|
||||
async def get_local_worker_status_endpoint(request):
|
||||
"""Check status of all local workers (localhost/no host specified)."""
|
||||
try:
|
||||
config = load_config()
|
||||
worker_statuses = {}
|
||||
|
||||
for worker in config.get("workers", []):
|
||||
# Only check local workers
|
||||
if not worker.get("host") or worker.get("host") in ["localhost", "127.0.0.1"]:
|
||||
worker_id = worker["id"]
|
||||
port = worker["port"]
|
||||
|
||||
# Check if worker is enabled
|
||||
if not worker.get("enabled", False):
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": False,
|
||||
"processing": False,
|
||||
"queue_count": 0
|
||||
}
|
||||
continue
|
||||
|
||||
# Try to connect to worker
|
||||
try:
|
||||
session = await get_client_session()
|
||||
async with session.get(
|
||||
f"http://localhost:{port}/prompt",
|
||||
timeout=aiohttp.ClientTimeout(total=2)
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
data = await resp.json()
|
||||
queue_remaining = data.get("exec_info", {}).get("queue_remaining", 0)
|
||||
worker_statuses[worker_id] = {
|
||||
"online": True,
|
||||
"enabled": True,
|
||||
"processing": queue_remaining > 0,
|
||||
"queue_count": queue_remaining
|
||||
}
|
||||
else:
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": True,
|
||||
"processing": False,
|
||||
"queue_count": 0,
|
||||
"error": f"HTTP {resp.status}"
|
||||
}
|
||||
except asyncio.TimeoutError:
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": True,
|
||||
"processing": False,
|
||||
"queue_count": 0,
|
||||
"error": "Timeout"
|
||||
}
|
||||
except Exception as e:
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": True,
|
||||
"processing": False,
|
||||
"queue_count": 0,
|
||||
"error": str(e)
|
||||
}
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"worker_statuses": worker_statuses
|
||||
})
|
||||
except Exception as e:
|
||||
debug_log(f"Error checking local worker status: {e}")
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/worker_log/{worker_id}")
|
||||
async def get_worker_log_endpoint(request):
|
||||
"""Get log content for a specific worker."""
|
||||
|
||||
+774
-500
File diff suppressed because it is too large
Load Diff
+5
-1
@@ -1,6 +1,7 @@
|
||||
"""
|
||||
Shared constants for ComfyUI-Distributed.
|
||||
"""
|
||||
import os
|
||||
|
||||
# Timeouts (in seconds)
|
||||
WORKER_JOB_TIMEOUT = 30.0
|
||||
@@ -34,4 +35,7 @@ TILE_SEND_TIMEOUT = 60.0
|
||||
MEMORY_CLEAR_DELAY = 0.5
|
||||
|
||||
# Batch processing
|
||||
MAX_BATCH = 20 # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs)
|
||||
MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs)
|
||||
|
||||
# Heartbeat monitoring
|
||||
HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60')) # Worker heartbeat timeout in seconds
|
||||
@@ -0,0 +1,803 @@
|
||||
import asyncio
|
||||
import time
|
||||
import json
|
||||
import copy
|
||||
import os
|
||||
import io
|
||||
from aiohttp import web
|
||||
import server
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# Import from other utilities
|
||||
from .logging import debug_log, log
|
||||
from .network import handle_api_error, get_client_session
|
||||
from .image import tensor_to_pil
|
||||
|
||||
# Configure maximum payload size (50MB default, configurable via environment variable)
|
||||
MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))
|
||||
|
||||
# Import HEARTBEAT_TIMEOUT from constants
|
||||
from .constants import HEARTBEAT_TIMEOUT
|
||||
|
||||
|
||||
# Unified Job Data Structure Keys
|
||||
JOB_QUEUE = 'queue'
|
||||
JOB_MODE = 'mode'
|
||||
JOB_COMPLETED_TASKS = 'completed_tasks'
|
||||
JOB_WORKER_STATUS = 'worker_status'
|
||||
JOB_ASSIGNED_TO_WORKERS = 'assigned_to_workers'
|
||||
JOB_PENDING_TASKS = 'pending_tasks'
|
||||
JOB_BATCH_SIZE = 'batch_size' # For dynamic
|
||||
JOB_NUM_TILES_PER_IMAGE = 'num_tiles_per_image' # For static
|
||||
|
||||
# Task Types
|
||||
TASK_TYPE_TILE = 'tile'
|
||||
TASK_TYPE_IMAGE = 'image'
|
||||
|
||||
async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_image=None, all_indices=None, enabled_workers=None, task_assignments=None):
|
||||
"""Unified initialization for job queues in static and dynamic modes."""
|
||||
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:
|
||||
debug_log(f"Queue already exists for {multi_job_id}")
|
||||
return
|
||||
|
||||
job_data = {
|
||||
JOB_QUEUE: asyncio.Queue(),
|
||||
JOB_MODE: mode,
|
||||
JOB_COMPLETED_TASKS: {},
|
||||
JOB_WORKER_STATUS: {w: time.time() for w in enabled_workers or []},
|
||||
JOB_ASSIGNED_TO_WORKERS: {w: [] for w in enabled_workers or []},
|
||||
JOB_PENDING_TASKS: asyncio.Queue(),
|
||||
}
|
||||
|
||||
if mode == 'dynamic':
|
||||
job_data[JOB_BATCH_SIZE] = batch_size
|
||||
pending_queue = job_data[JOB_PENDING_TASKS]
|
||||
for i in all_indices or range(batch_size):
|
||||
await pending_queue.put(i)
|
||||
debug_log(f"Initialized dynamic queue with {batch_size} pending images")
|
||||
elif mode == 'static':
|
||||
job_data[JOB_NUM_TILES_PER_IMAGE] = num_tiles_per_image
|
||||
# For static, pending_tasks starts empty; requeues will be added later
|
||||
debug_log(f"Initialized static queue for {batch_size * num_tiles_per_image} tiles")
|
||||
|
||||
# If task assignments provided, populate JOB_ASSIGNED_TO_WORKERS
|
||||
if task_assignments and enabled_workers:
|
||||
# task_assignments[0] is for master, 1+ are for workers
|
||||
for i, worker_id in enumerate(enabled_workers):
|
||||
if i + 1 < len(task_assignments):
|
||||
job_data[JOB_ASSIGNED_TO_WORKERS][worker_id] = task_assignments[i + 1]
|
||||
debug_log(f"Worker {worker_id} assigned {len(task_assignments[i + 1])} tasks")
|
||||
else:
|
||||
raise ValueError(f"Unknown mode: {mode}")
|
||||
|
||||
prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
|
||||
|
||||
def _distribute_tasks(items: list, num_participants: int) -> list[list[any]]:
|
||||
"""Distribute a list of items among N participants (master + workers)."""
|
||||
if num_participants <= 1:
|
||||
return [items]
|
||||
|
||||
items_per_participant = len(items) // num_participants
|
||||
remainder = len(items) % num_participants
|
||||
assignments = []
|
||||
start_idx = 0
|
||||
for i in range(num_participants):
|
||||
count = items_per_participant + (1 if i < remainder else 0)
|
||||
end_idx = start_idx + count
|
||||
assignments.append(items[start_idx:end_idx])
|
||||
start_idx = end_idx
|
||||
return assignments
|
||||
|
||||
async def _get_next_task(multi_job_id):
|
||||
"""Get next task from pending queue (generalized for tiles/images)."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if not job_data or JOB_PENDING_TASKS not in job_data:
|
||||
return None
|
||||
try:
|
||||
task_id = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=1.0)
|
||||
return task_id
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
|
||||
async def _drain_results_queue(multi_job_id):
|
||||
"""Drain pending results from queue and update completed_tasks. Returns count drained."""
|
||||
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_QUEUE not in job_data or JOB_COMPLETED_TASKS not in job_data:
|
||||
return 0
|
||||
q = job_data[JOB_QUEUE]
|
||||
completed_tasks = job_data[JOB_COMPLETED_TASKS]
|
||||
|
||||
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:
|
||||
task_id = result['image_idx']
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = result['image']
|
||||
collected += 1
|
||||
elif 'tiles' in result:
|
||||
for tile_data in result['tiles']:
|
||||
task_id = tile_data.get('global_idx', tile_data['tile_idx'])
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = tile_data
|
||||
collected += 1
|
||||
elif 'tensor' in result and 'tile_idx' in result: # Single tile backward compat
|
||||
task_id = result.get('global_idx', result['tile_idx'])
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = {
|
||||
'tensor': result['tensor'],
|
||||
'tile_idx': result['tile_idx'],
|
||||
'x': result['x'],
|
||||
'y': result['y'],
|
||||
'extracted_width': result['extracted_width'],
|
||||
'extracted_height': result['extracted_height'],
|
||||
'padding': result['padding'],
|
||||
'batch_idx': result.get('batch_idx', 0),
|
||||
'global_idx': task_id
|
||||
}
|
||||
collected += 1
|
||||
|
||||
if is_last:
|
||||
# Track worker completion
|
||||
if worker_id in job_data[JOB_WORKER_STATUS]:
|
||||
del job_data[JOB_WORKER_STATUS][worker_id]
|
||||
except asyncio.TimeoutError:
|
||||
break
|
||||
|
||||
return collected
|
||||
|
||||
async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
|
||||
"""Check timed out workers and requeue their tasks. Returns requeued 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 not job_data:
|
||||
return 0
|
||||
|
||||
current_time = time.time()
|
||||
requeued_count = 0
|
||||
completed_tasks = job_data.get(JOB_COMPLETED_TASKS, {})
|
||||
|
||||
for worker, last_heartbeat in list(job_data.get(JOB_WORKER_STATUS, {}).items()):
|
||||
if current_time - last_heartbeat > HEARTBEAT_TIMEOUT:
|
||||
log(f"Worker {worker} timed out")
|
||||
for task_id in job_data.get(JOB_ASSIGNED_TO_WORKERS, {}).get(worker, []):
|
||||
if task_id not in completed_tasks:
|
||||
await job_data[JOB_PENDING_TASKS].put(task_id)
|
||||
requeued_count += 1
|
||||
if JOB_WORKER_STATUS in job_data:
|
||||
del job_data[JOB_WORKER_STATUS][worker]
|
||||
if JOB_ASSIGNED_TO_WORKERS in job_data:
|
||||
job_data[JOB_ASSIGNED_TO_WORKERS][worker] = []
|
||||
|
||||
return requeued_count
|
||||
|
||||
async def _get_completed_count(multi_job_id):
|
||||
"""Get count of completed tasks."""
|
||||
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 len(job_data[JOB_COMPLETED_TASKS])
|
||||
return 0
|
||||
|
||||
async def _mark_task_completed(multi_job_id, task_id, result):
|
||||
"""Mark a task as completed."""
|
||||
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:
|
||||
job_data[JOB_COMPLETED_TASKS][task_id] = result
|
||||
|
||||
async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id):
|
||||
"""Send heartbeat to master."""
|
||||
try:
|
||||
data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)}
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/heartbeat"
|
||||
async with session.post(url, json=data) as response:
|
||||
response.raise_for_status()
|
||||
except Exception as e:
|
||||
debug_log(f"Heartbeat failed: {e}")
|
||||
|
||||
async def _cleanup_job(multi_job_id):
|
||||
"""Cleanup the job data."""
|
||||
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:
|
||||
del prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
debug_log(f"Cleaned up job {multi_job_id}")
|
||||
|
||||
# API Endpoints (generalized)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/request_task")
|
||||
async def request_task_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
|
||||
if not worker_id or not multi_job_id:
|
||||
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
||||
|
||||
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]
|
||||
mode = job_data.get(JOB_MODE)
|
||||
pending_queue = job_data.get(JOB_PENDING_TASKS)
|
||||
if pending_queue:
|
||||
try:
|
||||
task_id = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
||||
if JOB_ASSIGNED_TO_WORKERS in job_data and worker_id in job_data[JOB_ASSIGNED_TO_WORKERS]:
|
||||
job_data[JOB_ASSIGNED_TO_WORKERS][worker_id].append(task_id)
|
||||
if JOB_WORKER_STATUS in job_data:
|
||||
job_data[JOB_WORKER_STATUS][worker_id] = time.time()
|
||||
remaining = pending_queue.qsize()
|
||||
debug_log(f"Assigned task {task_id} to worker {worker_id} in {mode} mode")
|
||||
return web.json_response({"task_id": task_id, "estimated_remaining": remaining, "mode": mode})
|
||||
except asyncio.TimeoutError:
|
||||
return web.json_response({"task_id": None})
|
||||
else:
|
||||
return await handle_api_error(request, "No pending tasks", 400)
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/heartbeat")
|
||||
async def heartbeat_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
|
||||
if not worker_id or not multi_job_id:
|
||||
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
||||
|
||||
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]
|
||||
if JOB_WORKER_STATUS in job_data:
|
||||
job_data[JOB_WORKER_STATUS][worker_id] = time.time()
|
||||
debug_log(f"Heartbeat from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Worker status tracking not available", 400)
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/submit_tiles")
|
||||
async def submit_tiles_endpoint(request):
|
||||
"""Endpoint for workers to submit processed tiles in static mode."""
|
||||
try:
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
batch_size = int(data.get('batch_size', 0))
|
||||
tiles = []
|
||||
|
||||
# Handle completion signal
|
||||
if batch_size == 0 and is_last:
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
|
||||
return await handle_api_error(request, "Mode mismatch: expected static mode", 400)
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': []
|
||||
})
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
# Handle batch tiles with metadata
|
||||
if batch_size > 0:
|
||||
padding = int(data.get('padding', 32))
|
||||
metadata_field = data.get('tiles_metadata')
|
||||
if metadata_field:
|
||||
if hasattr(metadata_field, 'file'):
|
||||
metadata_str = metadata_field.file.read().decode('utf-8')
|
||||
elif isinstance(metadata_field, (bytes, bytearray)):
|
||||
metadata_str = metadata_field.decode('utf-8')
|
||||
else:
|
||||
metadata_str = str(metadata_field)
|
||||
|
||||
metadata = json.loads(metadata_str)
|
||||
if len(metadata) != batch_size:
|
||||
return await handle_api_error(request, "Metadata length mismatch", 400)
|
||||
|
||||
tile_data_list = []
|
||||
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
continue
|
||||
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
if i < len(metadata):
|
||||
tile_meta = metadata[i]
|
||||
tile_idx = tile_meta.get('tile_idx', i)
|
||||
tile_info = {
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': tile_meta['x'],
|
||||
'y': tile_meta['y'],
|
||||
'extracted_width': tile_meta['extracted_width'],
|
||||
'extracted_height': tile_meta['extracted_height'],
|
||||
'padding': padding,
|
||||
'batch_idx': tile_meta.get('batch_idx', 0),
|
||||
'global_idx': tile_meta.get('global_idx', tile_idx)
|
||||
}
|
||||
tile_data_list.append(tile_info)
|
||||
|
||||
tile_data_list.sort(key=lambda x: x['tile_idx'])
|
||||
tiles.extend(tile_data_list)
|
||||
else:
|
||||
# Legacy format
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
continue
|
||||
|
||||
tile_idx = int(data.get(f'tile_{i}_idx', i))
|
||||
x = int(data.get(f'tile_{i}_x', 0))
|
||||
y = int(data.get(f'tile_{i}_y', 0))
|
||||
extracted_width = int(data.get(f'tile_{i}_width', 512))
|
||||
extracted_height = int(data.get(f'tile_{i}_height', 512))
|
||||
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles.append({
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding
|
||||
})
|
||||
else:
|
||||
# Single tile legacy
|
||||
image_file = data.get('image')
|
||||
if not image_file:
|
||||
return await handle_api_error(request, "Missing image data", 400)
|
||||
|
||||
tile_idx = int(data.get('tile_idx', 0))
|
||||
x = int(data.get('x', 0))
|
||||
y = int(data.get('y', 0))
|
||||
extracted_width = int(data.get('extracted_width', 512))
|
||||
extracted_height = int(data.get('extracted_height', 512))
|
||||
padding = int(data.get('padding', 32))
|
||||
|
||||
img_data = image_file.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles = [{
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding,
|
||||
'batch_idx': 0,
|
||||
'global_idx': tile_idx
|
||||
}]
|
||||
|
||||
# Submit tiles to queue
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
|
||||
return await handle_api_error(request, "Mode mismatch: expected static mode", 400)
|
||||
|
||||
q = job_data[JOB_QUEUE]
|
||||
if batch_size > 0 or len(tiles) > 0:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'tiles': tiles,
|
||||
'is_last': is_last
|
||||
})
|
||||
else:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': []
|
||||
})
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/submit_image")
|
||||
async def submit_image_endpoint(request):
|
||||
"""Endpoint for workers to submit processed images in dynamic mode."""
|
||||
try:
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
# Handle image submission
|
||||
if 'full_image' in data and 'image_idx' in data:
|
||||
image_idx = int(data.get('image_idx'))
|
||||
img_data = data['full_image'].file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
|
||||
debug_log(f"Received full image {image_idx} from worker {worker_id}")
|
||||
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
|
||||
return await handle_api_error(request, "Mode mismatch: expected dynamic mode", 400)
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
'image_idx': image_idx,
|
||||
'image': img,
|
||||
'is_last': is_last
|
||||
})
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
# Handle completion signal (no image data)
|
||||
elif is_last:
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
|
||||
return await handle_api_error(request, "Mode mismatch: expected dynamic mode", 400)
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': [] # For compatibility
|
||||
})
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Missing image data or invalid request", 400)
|
||||
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
# Keep legacy endpoint for backward compatibility
|
||||
@server.PromptServer.instance.routes.post("/distributed/tile_complete")
|
||||
async def tile_complete_endpoint(request):
|
||||
try:
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
if 'full_image' in data and 'image_idx' in data:
|
||||
image_idx = int(data.get('image_idx'))
|
||||
img_data = data['full_image'].file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
|
||||
debug_log(f"Received full image {image_idx} from worker {worker_id}")
|
||||
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
|
||||
return await handle_api_error(request, "Mode mismatch for image submission", 400)
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
'image_idx': image_idx,
|
||||
'image': img,
|
||||
'is_last': is_last
|
||||
})
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
batch_size = int(data.get('batch_size', 0))
|
||||
tiles = []
|
||||
|
||||
if batch_size == 0 and is_last:
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': []
|
||||
})
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
if batch_size > 0:
|
||||
padding = int(data.get('padding', 32))
|
||||
metadata_field = data.get('tiles_metadata')
|
||||
if metadata_field:
|
||||
if hasattr(metadata_field, 'file'):
|
||||
metadata_str = metadata_field.file.read().decode('utf-8')
|
||||
elif isinstance(metadata_field, (bytes, bytearray)):
|
||||
metadata_str = metadata_field.decode('utf-8')
|
||||
else:
|
||||
metadata_str = str(metadata_field)
|
||||
|
||||
metadata = json.loads(metadata_str)
|
||||
if len(metadata) != batch_size:
|
||||
return await handle_api_error(request, "Metadata length mismatch", 400)
|
||||
|
||||
tile_data_list = []
|
||||
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
continue
|
||||
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
if i < len(metadata):
|
||||
tile_meta = metadata[i]
|
||||
tile_idx = tile_meta.get('tile_idx', i)
|
||||
tile_info = {
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': tile_meta['x'],
|
||||
'y': tile_meta['y'],
|
||||
'extracted_width': tile_meta['extracted_width'],
|
||||
'extracted_height': tile_meta['extracted_height'],
|
||||
'padding': padding,
|
||||
'batch_idx': tile_meta.get('batch_idx', 0),
|
||||
'global_idx': tile_meta.get('global_idx', tile_idx)
|
||||
}
|
||||
tile_data_list.append(tile_info)
|
||||
|
||||
tile_data_list.sort(key=lambda x: x['tile_idx'])
|
||||
tiles.extend(tile_data_list)
|
||||
else:
|
||||
# Legacy
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
continue
|
||||
|
||||
tile_idx = int(data.get(f'tile_{i}_idx', i))
|
||||
x = int(data.get(f'tile_{i}_x', 0))
|
||||
y = int(data.get(f'tile_{i}_y', 0))
|
||||
extracted_width = int(data.get(f'tile_{i}_width', 512))
|
||||
extracted_height = int(data.get(f'tile_{i}_height', 512))
|
||||
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles.append({
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding
|
||||
})
|
||||
else:
|
||||
# Single tile legacy
|
||||
image_file = data.get('image')
|
||||
if not image_file:
|
||||
return await handle_api_error(request, "Missing image data", 400)
|
||||
|
||||
tile_idx = int(data.get('tile_idx', 0))
|
||||
x = int(data.get('x', 0))
|
||||
y = int(data.get('y', 0))
|
||||
extracted_width = int(data.get('extracted_width', 512))
|
||||
extracted_height = int(data.get('extracted_height', 512))
|
||||
padding = int(data.get('padding', 32))
|
||||
|
||||
img_data = image_file.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles = [{
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding,
|
||||
'batch_idx': 0,
|
||||
'global_idx': tile_idx
|
||||
}]
|
||||
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
|
||||
return await handle_api_error(request, "Mode mismatch for tile submission", 400)
|
||||
|
||||
q = job_data[JOB_QUEUE]
|
||||
if batch_size > 0 or len(tiles) > 0:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'tiles': tiles,
|
||||
'is_last': is_last
|
||||
})
|
||||
else:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': []
|
||||
})
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
|
||||
# Helper functions for shallow copying conditioning without duplicating models
|
||||
def clone_control_chain(control, clone_hint=True):
|
||||
"""Shallow copy the ControlNet chain, optionally cloning hints but sharing models."""
|
||||
if control is None:
|
||||
return None
|
||||
new_control = copy.copy(control) # Shallow copy (shares model)
|
||||
if clone_hint and hasattr(control, 'cond_hint_original'):
|
||||
new_control.cond_hint_original = control.cond_hint_original.clone()
|
||||
if hasattr(control, 'previous_controlnet'):
|
||||
new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint)
|
||||
return new_control
|
||||
|
||||
def clone_conditioning(cond_list, clone_hints=True):
|
||||
"""Clone conditioning without duplicating ControlNet models."""
|
||||
new_cond = []
|
||||
for emb, cond_dict in cond_list:
|
||||
new_emb = emb.clone() if emb is not None else None
|
||||
new_dict = cond_dict.copy()
|
||||
if 'control' in new_dict:
|
||||
new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints)
|
||||
if 'mask' in new_dict:
|
||||
new_dict['mask'] = new_dict['mask'].clone()
|
||||
# Handle other potential fields if needed
|
||||
if 'pooled_output' in new_dict:
|
||||
new_dict['pooled_output'] = new_dict['pooled_output'].clone()
|
||||
if 'area' in new_dict:
|
||||
new_dict['area'] = new_dict['area'][:] # Shallow copy list/tuple
|
||||
new_cond.append([new_emb, new_dict])
|
||||
return new_cond
|
||||
|
||||
# Helper function to ensure persistent state is initialized
|
||||
def ensure_tile_jobs_initialized():
|
||||
"""Ensure tile job storage is initialized on the server instance."""
|
||||
prompt_server = server.PromptServer.instance
|
||||
if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
|
||||
debug_log("Initializing persistent tile job queue on server instance.")
|
||||
prompt_server.distributed_pending_tile_jobs = {}
|
||||
prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
|
||||
else:
|
||||
# Clean up any legacy queue structures that don't have the 'mode' field
|
||||
# (Should be rare after fixes, but keep for safety)
|
||||
to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items()
|
||||
if not isinstance(job_data, dict) or 'mode' not in job_data]
|
||||
for job_id in to_remove:
|
||||
debug_log(f"Removing legacy queue structure for job {job_id}")
|
||||
del prompt_server.distributed_pending_tile_jobs[job_id]
|
||||
return prompt_server
|
||||
|
||||
# API Endpoint for tile completion
|
||||
@server.PromptServer.instance.routes.post("/distributed/request_image")
|
||||
async def request_image_endpoint(request):
|
||||
"""Endpoint for workers to request images in dynamic mode."""
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
|
||||
if not worker_id or not multi_job_id:
|
||||
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
||||
|
||||
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]
|
||||
if not isinstance(job_data, dict) or 'mode' not in job_data:
|
||||
return await handle_api_error(request, "Invalid job data structure", 500)
|
||||
if job_data['mode'] != 'dynamic':
|
||||
return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400)
|
||||
if 'pending_images' in job_data:
|
||||
pending_queue = job_data['pending_images']
|
||||
try:
|
||||
image_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
||||
# Track assigned image
|
||||
if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']:
|
||||
job_data['assigned_to_workers'][worker_id].append(image_idx)
|
||||
# Update worker heartbeat
|
||||
if 'worker_status' in job_data:
|
||||
job_data['worker_status'][worker_id] = time.time()
|
||||
# Get estimated remaining count
|
||||
remaining = pending_queue.qsize() # Approximate
|
||||
debug_log(f"UltimateSDUpscale API - Assigned image {image_idx} to worker {worker_id}")
|
||||
return web.json_response({"image_idx": image_idx, "estimated_remaining": remaining})
|
||||
except asyncio.TimeoutError:
|
||||
return web.json_response({"image_idx": None})
|
||||
else:
|
||||
return await handle_api_error(request, "Invalid dynamic mode configuration", 400)
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/job_status")
|
||||
async def job_status_endpoint(request):
|
||||
"""Endpoint to check if a job is ready."""
|
||||
multi_job_id = request.query.get('multi_job_id')
|
||||
if not multi_job_id:
|
||||
return web.json_response({"ready": False})
|
||||
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)
|
||||
ready = bool(job_data and isinstance(job_data, dict) and 'queue' in job_data)
|
||||
return web.json_response({"ready": ready})
|
||||
|
||||
|
||||
+1
-411
@@ -4,27 +4,13 @@ import torch
|
||||
import torch.nn.functional as F
|
||||
from torchvision.transforms import GaussianBlur
|
||||
import math
|
||||
import copy
|
||||
import asyncio
|
||||
import json
|
||||
import io
|
||||
import time
|
||||
import os
|
||||
import server
|
||||
from aiohttp import web
|
||||
|
||||
# Import other utilities from the same package
|
||||
from .logging import debug_log, log
|
||||
from .network import handle_api_error
|
||||
|
||||
if (not hasattr(Image, 'Resampling')): # For older versions of Pillow
|
||||
Image.Resampling = Image
|
||||
|
||||
BLUR_KERNEL_SIZE = 15
|
||||
|
||||
# Configure maximum payload size (50MB default, configurable via environment variable)
|
||||
MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))
|
||||
|
||||
|
||||
def tensor_to_pil(img_tensor, batch_index=0):
|
||||
# Takes a batch of images in the form of a tensor of shape [batch_size, height, width, channels]
|
||||
@@ -467,400 +453,4 @@ def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0)
|
||||
crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
cropped.append(n)
|
||||
return cropped
|
||||
|
||||
|
||||
# Helper functions for shallow copying conditioning without duplicating models
|
||||
def clone_control_chain(control, clone_hint=True):
|
||||
"""Shallow copy the ControlNet chain, optionally cloning hints but sharing models."""
|
||||
if control is None:
|
||||
return None
|
||||
new_control = copy.copy(control) # Shallow copy (shares model)
|
||||
if clone_hint and hasattr(control, 'cond_hint_original'):
|
||||
new_control.cond_hint_original = control.cond_hint_original.clone()
|
||||
if hasattr(control, 'previous_controlnet'):
|
||||
new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint)
|
||||
return new_control
|
||||
|
||||
def clone_conditioning(cond_list, clone_hints=True):
|
||||
"""Clone conditioning without duplicating ControlNet models."""
|
||||
new_cond = []
|
||||
for emb, cond_dict in cond_list:
|
||||
new_emb = emb.clone() if emb is not None else None
|
||||
new_dict = cond_dict.copy()
|
||||
if 'control' in new_dict:
|
||||
new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints)
|
||||
if 'mask' in new_dict:
|
||||
new_dict['mask'] = new_dict['mask'].clone()
|
||||
# Handle other potential fields if needed
|
||||
if 'pooled_output' in new_dict:
|
||||
new_dict['pooled_output'] = new_dict['pooled_output'].clone()
|
||||
if 'area' in new_dict:
|
||||
new_dict['area'] = new_dict['area'][:] # Shallow copy list/tuple
|
||||
new_cond.append([new_emb, new_dict])
|
||||
return new_cond
|
||||
|
||||
# Helper function to ensure persistent state is initialized
|
||||
def ensure_tile_jobs_initialized():
|
||||
"""Ensure tile job storage is initialized on the server instance."""
|
||||
prompt_server = server.PromptServer.instance
|
||||
if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
|
||||
debug_log("Initializing persistent tile job queue on server instance.")
|
||||
prompt_server.distributed_pending_tile_jobs = {}
|
||||
prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
|
||||
else:
|
||||
# Clean up any legacy queue structures that don't have the 'mode' field
|
||||
# (Should be rare after fixes, but keep for safety)
|
||||
to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items()
|
||||
if not isinstance(job_data, dict) or 'mode' not in job_data]
|
||||
for job_id in to_remove:
|
||||
debug_log(f"Removing legacy queue structure for job {job_id}")
|
||||
del prompt_server.distributed_pending_tile_jobs[job_id]
|
||||
return prompt_server
|
||||
|
||||
# API Endpoint for tile completion
|
||||
@server.PromptServer.instance.routes.post("/distributed/request_image")
|
||||
async def request_image_endpoint(request):
|
||||
"""Endpoint for workers to request images in dynamic mode."""
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
|
||||
if not worker_id or not multi_job_id:
|
||||
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
||||
|
||||
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]
|
||||
if not isinstance(job_data, dict) or 'mode' not in job_data:
|
||||
return await handle_api_error(request, "Invalid job data structure", 500)
|
||||
if job_data['mode'] != 'dynamic':
|
||||
return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400)
|
||||
if 'pending_images' in job_data:
|
||||
pending_queue = job_data['pending_images']
|
||||
try:
|
||||
image_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
||||
# Track assigned image
|
||||
if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']:
|
||||
job_data['assigned_to_workers'][worker_id].append(image_idx)
|
||||
# Update worker heartbeat
|
||||
if 'worker_status' in job_data:
|
||||
job_data['worker_status'][worker_id] = time.time()
|
||||
# Get estimated remaining count
|
||||
remaining = pending_queue.qsize() # Approximate
|
||||
debug_log(f"UltimateSDUpscale API - Assigned image {image_idx} to worker {worker_id}")
|
||||
return web.json_response({"image_idx": image_idx, "estimated_remaining": remaining})
|
||||
except asyncio.TimeoutError:
|
||||
return web.json_response({"image_idx": None})
|
||||
else:
|
||||
return await handle_api_error(request, "Invalid dynamic mode configuration", 400)
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/job_status")
|
||||
async def job_status_endpoint(request):
|
||||
"""Endpoint to check if a job is ready."""
|
||||
multi_job_id = request.query.get('multi_job_id')
|
||||
if not multi_job_id:
|
||||
return web.json_response({"ready": False})
|
||||
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)
|
||||
ready = bool(job_data and isinstance(job_data, dict) and 'queue' in job_data)
|
||||
return web.json_response({"ready": ready})
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/heartbeat")
|
||||
async def heartbeat_endpoint(request):
|
||||
"""Endpoint for workers to send heartbeat signals."""
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
|
||||
if not worker_id or not multi_job_id:
|
||||
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
||||
|
||||
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]
|
||||
if 'worker_status' in job_data:
|
||||
job_data['worker_status'][worker_id] = time.time()
|
||||
debug_log(f"UltimateSDUpscale API - Heartbeat from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Worker status tracking not available", 400)
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/tile_complete")
|
||||
async def tile_complete_endpoint(request):
|
||||
"""Endpoint for receiving completed tiles from workers."""
|
||||
try:
|
||||
# Check payload size before processing
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes (max: {MAX_PAYLOAD_SIZE})", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
# Check if this is a full image submission (dynamic mode)
|
||||
if 'full_image' in data and 'image_idx' in data:
|
||||
image_idx = int(data.get('image_idx'))
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
try:
|
||||
# Process full image
|
||||
img_data = data['full_image'].file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
|
||||
debug_log(f"UltimateSDUpscale API - Received full image {image_idx} from worker {worker_id}")
|
||||
|
||||
# Put into dynamic mode queue
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, dict) or 'mode' not in job_data:
|
||||
return await handle_api_error(request, "Invalid job data structure", 500)
|
||||
if job_data['mode'] != 'dynamic':
|
||||
return await handle_api_error(request, f"Mode mismatch: expected dynamic, got {job_data['mode']}", 400)
|
||||
# Check for duplicate image submission
|
||||
if 'completed_images' in job_data and image_idx in job_data['completed_images']:
|
||||
log(f"UltimateSDUpscale API - Duplicate image {image_idx} from {worker_id}, ignoring")
|
||||
return web.json_response({"status": "duplicate"})
|
||||
if 'queue' in job_data:
|
||||
await job_data['queue'].put({
|
||||
'worker_id': worker_id,
|
||||
'image_idx': image_idx,
|
||||
'image': img,
|
||||
'is_last': is_last
|
||||
})
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Invalid dynamic mode configuration", 400)
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
log(f"Error processing full image from worker {worker_id}: {e}")
|
||||
return await handle_api_error(request, f"Image processing error: {e}", 400)
|
||||
|
||||
# Check for batch mode
|
||||
batch_size = int(data.get('batch_size', 0))
|
||||
tiles = []
|
||||
|
||||
# Handle completion signals (batch_size=0 with is_last=True)
|
||||
if batch_size == 0 and is_last:
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if isinstance(job_data, dict) and 'queue' in job_data:
|
||||
await job_data['queue'].put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': []
|
||||
})
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Invalid job structure", 500)
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
|
||||
if batch_size > 0:
|
||||
# Batch mode: Extract multiple tiles
|
||||
padding = int(data.get('padding', 32))
|
||||
debug_log(f"UltimateSDUpscale - tile_complete batch - job_id: {multi_job_id}, worker: {worker_id}, batch_size: {batch_size}")
|
||||
|
||||
# Check for JSON metadata
|
||||
metadata_field = data.get('tiles_metadata')
|
||||
if metadata_field:
|
||||
# New JSON metadata format
|
||||
try:
|
||||
# Handle different types of metadata field
|
||||
if hasattr(metadata_field, 'file'):
|
||||
# File-like object
|
||||
metadata_str = metadata_field.file.read().decode('utf-8')
|
||||
elif isinstance(metadata_field, (bytes, bytearray)):
|
||||
# Direct bytes/bytearray
|
||||
metadata_str = metadata_field.decode('utf-8')
|
||||
else:
|
||||
# String
|
||||
metadata_str = str(metadata_field)
|
||||
|
||||
metadata = json.loads(metadata_str)
|
||||
if len(metadata) != batch_size:
|
||||
return await handle_api_error(request, "Metadata length mismatch", 400)
|
||||
|
||||
tile_data_list = [] # Temporary list to collect and sort tiles
|
||||
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
log(f"Missing tile_{i} from worker {worker_id}, skipping")
|
||||
continue
|
||||
|
||||
try:
|
||||
# Process image
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
# Get metadata for this tile
|
||||
if i < len(metadata):
|
||||
tile_data = metadata[i]
|
||||
tile_idx = tile_data.get('tile_idx', i)
|
||||
|
||||
tile_info = {
|
||||
'tensor': 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': padding,
|
||||
'batch_idx': tile_data.get('batch_idx', 0),
|
||||
'global_idx': tile_data.get('global_idx', tile_idx)
|
||||
}
|
||||
tile_data_list.append(tile_info)
|
||||
else:
|
||||
log(f"Missing metadata for tile {i} from worker {worker_id}, skipping")
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile")
|
||||
continue
|
||||
|
||||
# Sort tiles by tile_idx to ensure correct order
|
||||
tile_data_list.sort(key=lambda x: x['tile_idx'])
|
||||
tiles.extend(tile_data_list)
|
||||
|
||||
if tile_data_list:
|
||||
debug_log(f"Reordered {len(tile_data_list)} tiles based on tile_idx")
|
||||
except Exception as e:
|
||||
log(f"Error processing JSON metadata from worker {worker_id}: {e}")
|
||||
return await handle_api_error(request, f"Metadata processing error: {e}", 400)
|
||||
else:
|
||||
# Legacy format: individual fields per tile (backward compatibility)
|
||||
debug_log(f"WARNING: Worker {worker_id} using legacy field format. Please update to use JSON metadata.")
|
||||
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
log(f"Missing tile_{i} from worker {worker_id}, skipping")
|
||||
continue
|
||||
|
||||
# Get metadata for this tile
|
||||
tile_idx = int(data.get(f'tile_{i}_idx', i))
|
||||
x = int(data.get(f'tile_{i}_x', 0))
|
||||
y = int(data.get(f'tile_{i}_y', 0))
|
||||
extracted_width = int(data.get(f'tile_{i}_width', 512))
|
||||
extracted_height = int(data.get(f'tile_{i}_height', 512))
|
||||
|
||||
try:
|
||||
# Process image
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles.append({
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding
|
||||
})
|
||||
except Exception as e:
|
||||
log(f"Error processing tile {i} from worker {worker_id}: {e}, skipping this tile")
|
||||
continue
|
||||
else:
|
||||
# Single tile mode (backward compatibility)
|
||||
image_file = data.get('image')
|
||||
if not image_file:
|
||||
return await handle_api_error(request, "Missing image data", 400)
|
||||
|
||||
tile_idx = int(data.get('tile_idx', 0))
|
||||
x = int(data.get('x', 0))
|
||||
y = int(data.get('y', 0))
|
||||
extracted_width = int(data.get('extracted_width', 512))
|
||||
extracted_height = int(data.get('extracted_height', 512))
|
||||
padding = int(data.get('padding', 32))
|
||||
|
||||
debug_log(f"UltimateSDUpscale - tile_complete single - job_id: {multi_job_id}, worker: {worker_id}, tile: {tile_idx}")
|
||||
|
||||
try:
|
||||
# Process image
|
||||
img_data = image_file.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles = [{
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding,
|
||||
'batch_idx': 0, # Default for legacy single-tile
|
||||
'global_idx': tile_idx # Default for legacy single-tile
|
||||
}]
|
||||
except Exception as e:
|
||||
log(f"Error processing tile from worker {worker_id}: {e}")
|
||||
return await handle_api_error(request, f"Tile processing error: {e}", 400)
|
||||
|
||||
# Put tiles into queue
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, dict) or 'mode' not in job_data:
|
||||
return await handle_api_error(request, "Invalid job data structure", 500)
|
||||
if job_data['mode'] != 'static':
|
||||
return await handle_api_error(request, f"Mode mismatch: expected static, got {job_data['mode']}", 400)
|
||||
|
||||
q = job_data['queue']
|
||||
if batch_size > 0:
|
||||
# Put batch as single item
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'tiles': tiles,
|
||||
'is_last': is_last
|
||||
})
|
||||
else:
|
||||
# Put single tile (backward compat)
|
||||
tile_data = tiles[0]
|
||||
await q.put({
|
||||
'tensor': tile_data['tensor'],
|
||||
'worker_id': worker_id,
|
||||
'tile_idx': tile_data['tile_idx'],
|
||||
'x': tile_data['x'],
|
||||
'y': tile_data['y'],
|
||||
'extracted_width': tile_data['extracted_width'],
|
||||
'extracted_height': tile_data['extracted_height'],
|
||||
'padding': tile_data['padding'],
|
||||
'is_last': is_last
|
||||
})
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
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)
|
||||
return cropped
|
||||
Reference in New Issue
Block a user