822 lines
38 KiB
Python
822 lines
38 KiB
Python
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 with dynamic distribution, populate all tile indices in pending queue
|
|
pending_queue = job_data[JOB_PENDING_TASKS]
|
|
total_tiles = batch_size * num_tiles_per_image
|
|
for i in range(total_tiles):
|
|
await pending_queue.put(i)
|
|
debug_log(f"Initialized static queue with {total_tiles} pending tiles for dynamic distribution")
|
|
|
|
# Keep backward compatibility - if task assignments provided, still track them
|
|
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} pre-assigned {len(task_assignments[i + 1])} tasks (legacy mode)")
|
|
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 tasks (images in dynamic mode, tiles in static 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)
|
|
|
|
mode = job_data['mode']
|
|
|
|
# Handle both dynamic and static modes
|
|
if mode == 'dynamic' and 'pending_images' in job_data:
|
|
pending_queue = job_data['pending_images']
|
|
elif mode == 'static' and JOB_PENDING_TASKS in job_data:
|
|
pending_queue = job_data[JOB_PENDING_TASKS]
|
|
else:
|
|
return await handle_api_error(request, f"Invalid {mode} mode configuration", 400)
|
|
|
|
try:
|
|
task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
|
# Track assigned task
|
|
if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']:
|
|
job_data['assigned_to_workers'][worker_id].append(task_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
|
|
|
|
# Return appropriate response based on mode
|
|
if mode == 'dynamic':
|
|
debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}")
|
|
return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining})
|
|
else: # static
|
|
debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
|
|
return web.json_response({"tile_idx": task_idx, "estimated_remaining": remaining})
|
|
except asyncio.TimeoutError:
|
|
if mode == 'dynamic':
|
|
return web.json_response({"image_idx": None})
|
|
else:
|
|
return web.json_response({"tile_idx": None})
|
|
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})
|
|
|
|
|