robertvoy-ComfyUI-Distributed/upscale/modes/dynamic.py

from __future__ import annotations

import asyncio, time, torch
from typing import Any
from PIL import Image
from ...utils.logging import debug_log, log
from ...utils.image import tensor_to_pil, pil_to_tensor
from ...utils.async_helpers import run_async_in_server_loop
from ...utils.config import get_worker_timeout_seconds
from ...utils.constants import (
    JOB_POLL_INTERVAL,
    JOB_POLL_MAX_ATTEMPTS,
    TILE_WAIT_TIMEOUT,
    TILE_SEND_TIMEOUT,
)
from ..job_store import ensure_tile_jobs_initialized, init_dynamic_job
from ..mode_contexts import (
    DynamicModeContext,
    JobStateCollaborator,
    ResultCollectorCollaborator,
    TileOpsCollaborator,
    WorkerCommsCollaborator,
)
from ..processing_args import UpscaleCoreArgs


class DynamicModeMixin:
    """
    Dynamic (per-image queue) USDU mode behaviors for master and worker roles.

    Expected co-mixins on `self`:
    - TileOpsMixin (`calculate_tiles`, `slice_conditioning`, `process_and_blend_tile`).
    - JobStateMixin (image queue/task completion helpers).
    - WorkerCommsMixin (`request_assignment`, `send_full_image`, `send_heartbeat`).
    """

    def _build_dynamic_mode_context(self) -> DynamicModeContext:
        """Build explicit collaborators for dynamic mode execution."""
        return DynamicModeContext(
            tile_ops=TileOpsCollaborator(self),
            job_state=JobStateCollaborator(self),
            worker_comms=WorkerCommsCollaborator(self),
            result_collector=ResultCollectorCollaborator(self),
        )

    def _poll_job_ready_dynamic(
        self,
        multi_job_id: str,
        master_url: str,
        worker_id: str | None = None,
        max_attempts: int = JOB_POLL_MAX_ATTEMPTS,
        mode_context: DynamicModeContext | None = None,
    ) -> bool:
        """Poll master for job readiness using explicit worker-comms collaborator."""
        context = mode_context or self._build_dynamic_mode_context()
        for attempt in range(max_attempts):
            ready = run_async_in_server_loop(
                context.worker_comms.check_job_status(multi_job_id, master_url),
                timeout=5.0,
            )
            if ready:
                if worker_id:
                    debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts")
                else:
                    debug_log(f"Job {multi_job_id} ready after {attempt} attempts")
                return True
            time.sleep(JOB_POLL_INTERVAL)
        return False

    def _process_image_tiles(
        self,
        image_idx: int,
        local_image: Image.Image,
        single_tensor: torch.Tensor,
        core_args: UpscaleCoreArgs,
        all_tiles: list,
        tile_width: int,
        tile_height: int,
        padding: int,
        mask_blur: int,
        width: int,
        height: int,
        force_uniform_tiles: bool,
        context: DynamicModeContext,
        per_tile_callback: Any | None = None,
    ) -> Image.Image:
        """Process all tiles for a single image. Returns the processed PIL image."""
        positive_sliced, negative_sliced = context.tile_ops.slice_conditioning(
            core_args.positive, core_args.negative, image_idx,
        )
        for tile_idx, pos in enumerate(all_tiles):
            source_tensor = pil_to_tensor(local_image)
            if single_tensor.is_cuda:
                source_tensor = source_tensor.cuda()
            local_image = context.tile_ops.process_and_blend_tile(
                tile_idx, pos, source_tensor, local_image,
                core_args.model,
                positive_sliced,
                negative_sliced,
                core_args.vae,
                core_args.seed,
                core_args.steps,
                core_args.cfg,
                core_args.sampler_name,
                core_args.scheduler,
                core_args.denoise,
                tile_width,
                tile_height,
                padding, mask_blur, width, height, force_uniform_tiles,
                core_args.tiled_decode,
                batch_idx=image_idx,
            )
            if per_tile_callback:
                per_tile_callback()
        return local_image

    def _handle_master_dynamic_idle_state(
        self,
        multi_job_id,
        batch_size,
        consecutive_retries,
        max_consecutive_retries,
        mode_context: DynamicModeContext | None = None,
    ):
        """Handle queue-empty state while master waits for worker progress."""
        context = mode_context or self._build_dynamic_mode_context()
        drained_count = run_async_in_server_loop(
            context.job_state.drain_worker_results_queue(multi_job_id),
            timeout=5.0,
        )
        run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)

        requeued_count = run_async_in_server_loop(
            context.job_state.check_and_requeue_timed_out_workers(multi_job_id, batch_size),
            timeout=5.0,
        )
        run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)

        if requeued_count > 0:
            log(f"Requeued {requeued_count} images from timed out workers")
            return False, True, 0

        completed_now = run_async_in_server_loop(
            context.job_state.total_completed_count(multi_job_id),
            timeout=1.0,
        )
        log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
        if completed_now >= batch_size:
            return True, False, consecutive_retries

        run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)

        pending_count = run_async_in_server_loop(
            context.job_state.pending_count(multi_job_id),
            timeout=1.0,
        )
        if pending_count > 0:
            return False, True, 0

        consecutive_retries += 1
        if consecutive_retries >= max_consecutive_retries:
            log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.")
            return True, False, consecutive_retries

        if drained_count > 0:
            debug_log(f"Master idle drain picked up {drained_count} images while waiting for workers")
        debug_log("Waiting for workers")
        run_async_in_server_loop(asyncio.sleep(2), timeout=3.0)
        return False, True, consecutive_retries

    def _finalize_master_dynamic_results(
        self,
        multi_job_id,
        batch_size,
        num_workers,
        processed_count,
        result_images,
        upscaled_image,
        mode_context: DynamicModeContext | None = None,
    ):
        """Collect remaining worker outputs and convert final images back to tensor."""
        context = mode_context or self._build_dynamic_mode_context()
        all_completed = run_async_in_server_loop(
            context.job_state.all_completed_images(multi_job_id),
            timeout=5.0,
        )
        remaining_to_collect = batch_size - len(all_completed)
        if remaining_to_collect > 0:
            debug_log(f"Waiting for {remaining_to_collect} more images from workers")
            collection_timeout = float(get_worker_timeout_seconds())
            collected_images = run_async_in_server_loop(
                context.result_collector.collect_dynamic_images(
                    multi_job_id,
                    remaining_to_collect,
                    num_workers,
                    batch_size,
                    processed_count,
                ),
                timeout=collection_timeout,
            )
            all_completed.update(collected_images)

        for idx, processed_img in all_completed.items():
            if idx < batch_size:
                result_images[idx] = processed_img

        result_tensor = (
            torch.cat([pil_to_tensor(img) for img in result_images], dim=0)
            if batch_size > 1
            else pil_to_tensor(result_images[0])
        )
        if upscaled_image.is_cuda:
            result_tensor = result_tensor.cuda()
        return result_tensor

    def process_master_dynamic(
        self,
        upscaled_image: torch.Tensor,
        core_args: UpscaleCoreArgs,
        tile_width: int,
        tile_height: int,
        padding: int,
        mask_blur: int,
        force_uniform_tiles: bool,
        multi_job_id: str,
        enabled_workers: list[str],
        mode_context: DynamicModeContext | None = None,
    ) -> tuple[torch.Tensor]:
        """Dynamic mode for large batches - assigns whole images to workers dynamically, including master."""
        context = mode_context or self._build_dynamic_mode_context()
        # Get batch size and dimensions
        batch_size, height, width, _ = upscaled_image.shape
        num_workers = len(enabled_workers)
        
        log(f"USDU Dist: Image queue distribution | Batch {batch_size} | Workers {num_workers} | Canvas {width}x{height} | Tile {tile_width}x{tile_height}")

        # No fixed share - all images are dynamic
        all_indices = list(range(batch_size))
        
        debug_log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.")
        
        # Calculate tiles for processing
        all_tiles = context.tile_ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
        
        # Initialize job queue for communication
        try:
            run_async_in_server_loop(
                init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices),
                timeout=2.0
            )
        except Exception as e:
            debug_log(f"UltimateSDUpscale Master - Queue initialization error: {e}")
            raise RuntimeError(f"Failed to initialize dynamic mode queue: {e}")
        
        # Convert batch to PIL list
        result_images = [tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB').copy() for b in range(batch_size)]
        
        # Process images dynamically with master participating
        prompt_server = ensure_tile_jobs_initialized()
        processed_count = 0
        consecutive_retries = 0
        max_consecutive_retries = 10
        
        # Process loop - master pulls from queue and processes synchronously
        while processed_count < batch_size:
            # Try to get an image to process
            image_idx = run_async_in_server_loop(
                context.job_state.next_image_index(multi_job_id),
                timeout=5.0  # Short timeout to allow frequent checks
            )

            if image_idx is not None:
                # Reset retry counter and process locally
                consecutive_retries = 0
                debug_log(f"Master processing image {image_idx} dynamically")
                processed_count += 1

                # Process locally
                single_tensor = upscaled_image[image_idx:image_idx+1]
                local_image = result_images[image_idx]

                local_image = self._process_image_tiles(
                    image_idx, local_image, single_tensor, core_args, all_tiles,
                    tile_width, tile_height, padding, mask_blur, width, height,
                    force_uniform_tiles, context,
                    per_tile_callback=lambda: run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1),
                )

                result_images[image_idx] = local_image
                
                # Mark as completed
                run_async_in_server_loop(
                    context.result_collector.mark_image_completed(multi_job_id, image_idx, local_image),
                    timeout=5.0
                )
                
                # NEW: Drain after the full image is marked complete (catches workers who finished during master's processing)
                drained_count = run_async_in_server_loop(
                    context.job_state.drain_worker_results_queue(multi_job_id),
                    timeout=5.0
                )
                
                if drained_count > 0:
                    debug_log(f"Drained {drained_count} worker images after master's image {image_idx}")
                
                # NEW: Log overall progress (includes master's image + any drained workers)
                completed_now = run_async_in_server_loop(
                    context.job_state.total_completed_count(multi_job_id),
                    timeout=1.0
                )
                log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
                
                # Yield to allow workers to get new images after completing one
                run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)
            else:
                should_break, should_continue, consecutive_retries = self._handle_master_dynamic_idle_state(
                    multi_job_id=multi_job_id,
                    batch_size=batch_size,
                    consecutive_retries=consecutive_retries,
                    max_consecutive_retries=max_consecutive_retries,
                    mode_context=context,
                )
                if should_break:
                    break
                if should_continue:
                    continue
        
        debug_log(f"Master processed {processed_count} images locally")
        result_tensor = self._finalize_master_dynamic_results(
            multi_job_id=multi_job_id,
            batch_size=batch_size,
            num_workers=num_workers,
            processed_count=processed_count,
            result_images=result_images,
            upscaled_image=upscaled_image,
            mode_context=context,
        )
        
        debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
        log(f"Completed processing all {batch_size} images")
        return (result_tensor,)

    def process_worker_dynamic(
        self,
        upscaled_image: torch.Tensor,
        core_args: UpscaleCoreArgs,
        tile_width: int,
        tile_height: int,
        padding: int,
        mask_blur: int,
        force_uniform_tiles: bool,
        multi_job_id: str,
        master_url: str,
        worker_id: str,
        enabled_worker_ids: list[str] | str,
        dynamic_threshold: int,
        mode_context: DynamicModeContext | None = None,
    ) -> tuple[torch.Tensor]:
        """Worker processing in dynamic mode - processes whole images."""
        context = mode_context or self._build_dynamic_mode_context()
        # Round tile dimensions
        tile_width = context.tile_ops.round_to_multiple(tile_width)
        tile_height = context.tile_ops.round_to_multiple(tile_height)

        # Get dimensions and tile grid
        batch_size, height, width, _ = upscaled_image.shape
        all_tiles = context.tile_ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
        log(f"USDU Dist Worker[{worker_id[:8]}]: Processing image queue | Batch {batch_size}")

        # Keep track of processed images for is_last detection
        processed_count = 0

        # Ensure the dynamic job queue exists on the master.  In delegate-only
        # mode the master's USDU node is replaced with a lightweight collector
        # that cannot initialise the queue itself, so workers do it instead.
        # The call is idempotent — if the master already created the queue this
        # is a no-op.
        from ...utils.worker_ids import coerce_enabled_worker_ids as _coerce
        _enabled = _coerce(enabled_worker_ids)
        run_async_in_server_loop(
            context.worker_comms.init_dynamic_job_on_master(
                multi_job_id, master_url, batch_size, _enabled,
            ),
            timeout=10.0,
        )

        # Poll for job readiness to avoid races during master init
        max_poll_attempts = JOB_POLL_MAX_ATTEMPTS
        if not self._poll_job_ready_dynamic(
            multi_job_id,
            master_url,
            worker_id=worker_id,
            max_attempts=max_poll_attempts,
            mode_context=context,
        ):
            log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
            return (upscaled_image,)

        # Loop to request and process images
        while True:
            # Request an image to process
            assignment = run_async_in_server_loop(
                context.worker_comms.request_assignment(multi_job_id, master_url, worker_id),
                timeout=TILE_WAIT_TIMEOUT,
            )
            image_idx = assignment.task_idx if assignment.kind == "image" else None
            estimated_remaining = assignment.estimated_remaining

            if image_idx is None:
                debug_log(f"USDU Dist Worker - No more images to process")
                break

            debug_log(f"Worker[{worker_id[:8]}] - Assigned image {image_idx}")
            processed_count += 1

            # Determine if this should be marked as last for this worker
            is_last_for_worker = (estimated_remaining == 0)

            # Extract single image tensor
            single_tensor = upscaled_image[image_idx:image_idx+1]

            # Convert to PIL for processing
            local_image = tensor_to_pil(single_tensor, 0).copy()

            # Process all tiles for this image
            local_image = self._process_image_tiles(
                image_idx, local_image, single_tensor, core_args, all_tiles,
                tile_width, tile_height, padding, mask_blur, width, height,
                force_uniform_tiles, context,
                per_tile_callback=lambda: run_async_in_server_loop(
                    context.worker_comms.send_heartbeat(multi_job_id, master_url, worker_id),
                    timeout=5.0,
                ),
            )

            # Send processed image back to master
            try:
                # Use the estimated remaining to determine if this is the last image
                is_last = is_last_for_worker
                run_async_in_server_loop(
                    context.worker_comms.send_full_image(
                        local_image,
                        image_idx,
                        multi_job_id,
                        master_url,
                        worker_id,
                        is_last,
                    ),
                    timeout=TILE_SEND_TIMEOUT
                )
                # Send heartbeat after processing
                run_async_in_server_loop(
                    context.worker_comms.send_heartbeat(multi_job_id, master_url, worker_id),
                    timeout=5.0
                )
                if is_last:
                    break
            except Exception as e:
                log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending image {image_idx}: {e}")
                # Continue processing other images

        # Send final is_last signal
        debug_log(f"Worker[{worker_id[:8]}] processed {processed_count} images, sending completion signal")
        try:
            run_async_in_server_loop(
                context.worker_comms.send_worker_complete_signal(multi_job_id, master_url, worker_id),
                timeout=TILE_SEND_TIMEOUT
            )
        except Exception as e:
            log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending completion signal: {e}")

        return (upscaled_image,)