280 lines
14 KiB
Python
280 lines
14 KiB
Python
import json
|
|
import math
|
|
from functools import wraps
|
|
|
|
import comfy.samplers
|
|
|
|
from .utils.logging import debug_log, log
|
|
from .utils.async_helpers import run_async_in_server_loop
|
|
from .utils.usdu_managment import ensure_tile_jobs_initialized
|
|
|
|
from .upscale.tile_ops import TileOpsMixin
|
|
from .upscale.result_collector import ResultCollectorMixin
|
|
from .upscale.worker_comms import WorkerCommsMixin
|
|
from .upscale.job_state import JobStateMixin
|
|
from .upscale.modes.single_gpu import SingleGpuModeMixin
|
|
from .upscale.modes.static import StaticModeMixin
|
|
from .upscale.modes.dynamic import DynamicModeMixin
|
|
|
|
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
|
|
|
|
def _parse_enabled_worker_ids(enabled_worker_ids):
|
|
"""Parse enabled worker IDs from either JSON or list input."""
|
|
if isinstance(enabled_worker_ids, list):
|
|
return [str(worker_id) for worker_id in enabled_worker_ids]
|
|
if not enabled_worker_ids:
|
|
return []
|
|
if isinstance(enabled_worker_ids, str):
|
|
try:
|
|
parsed = json.loads(enabled_worker_ids)
|
|
except json.JSONDecodeError:
|
|
log("USDU Dist: Invalid enabled_worker_ids JSON; defaulting to no workers.")
|
|
return []
|
|
if isinstance(parsed, list):
|
|
return [str(wid) for wid in parsed]
|
|
return []
|
|
|
|
class UltimateSDUpscaleDistributed(
|
|
DynamicModeMixin,
|
|
StaticModeMixin,
|
|
SingleGpuModeMixin,
|
|
ResultCollectorMixin,
|
|
WorkerCommsMixin,
|
|
JobStateMixin,
|
|
TileOpsMixin,
|
|
):
|
|
|
|
"""
|
|
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}),
|
|
},
|
|
"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,
|
|
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."""
|
|
# Strict WAN/FLOW batching: error if batch is not 4n+1 (except allow 1)
|
|
try:
|
|
batch_size = int(getattr(upscaled_image, 'shape', [1])[0])
|
|
except Exception:
|
|
batch_size = 1
|
|
# Enforce 4n+1 batches globally for any model when batch > 1 (master only)
|
|
if not is_worker and batch_size != 1 and (batch_size % 4 != 1):
|
|
raise ValueError(
|
|
f"Batch size {batch_size} is not of the form 4n+1. "
|
|
"This node requires batch sizes of 1 or 4n+1 (1, 5, 9, 13, ...). "
|
|
"Please adjust the batch size."
|
|
)
|
|
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)
|
|
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)
|
|
|
|
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):
|
|
"""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)
|
|
# Compute number of tiles for this image to decide if tile distribution makes sense
|
|
_, height, width, _ = upscaled_image.shape
|
|
all_tiles = self.calculate_tiles(width, height, self.round_to_multiple(tile_width), self.round_to_multiple(tile_height), force_uniform_tiles)
|
|
num_tiles_per_image = len(all_tiles)
|
|
|
|
mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
|
|
# For USDU-style processing, we want tile distribution whenever workers are available
|
|
# and there is more than one tile to process, even if batch == 1.
|
|
if num_workers > 0 and num_tiles_per_image > 1:
|
|
mode = "static"
|
|
|
|
debug_log(f"USDU Dist Worker - 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)
|
|
|
|
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):
|
|
"""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 and grid
|
|
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
|
num_tiles_per_image = len(all_tiles)
|
|
rows = math.ceil(height / tile_height)
|
|
cols = math.ceil(width / tile_width)
|
|
log(
|
|
f"USDU Dist: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({num_tiles_per_image} tiles/image) | Batch {batch_size}"
|
|
)
|
|
|
|
# Parse enabled workers
|
|
enabled_workers = json.loads(enabled_worker_ids)
|
|
num_workers = len(enabled_workers)
|
|
|
|
# Determine processing mode
|
|
mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
|
|
# Prefer tile-based static distribution when workers are available and there are multiple tiles,
|
|
# even for batch == 1, to spread tiles across GPUs like the legacy dynamic tile queue.
|
|
if num_workers > 0 and num_tiles_per_image > 1:
|
|
mode = "static"
|
|
|
|
log(f"USDU Dist: Workers {num_workers} | Mode {mode} | Threshold {dynamic_threshold}")
|
|
|
|
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)
|
|
|
|
def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str:
|
|
"""Determines processing mode per requested policy:
|
|
- any workers => prefer static (tile-based) for USDU
|
|
- no workers => single_gpu
|
|
"""
|
|
if num_workers == 0:
|
|
return "single_gpu"
|
|
# Default to static when distributed; master/worker may still override if special cases arise
|
|
return "static"
|
|
|
|
# Ensure initialization before registering routes
|
|
ensure_tile_jobs_initialized()
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
|
|
}
|