fix: keep local LTX tiling in execution context

Fixes #85

.github/FUNDING.yml

@@ -1,4 +1,3 @@
 # These are supported funding model platforms
 
 github: robertvoy
-buy_me_a_coffee: robertvoy

api/__init__.py

@@ -3,3 +3,4 @@ from . import tunnel_routes  # noqa: F401
 from . import worker_routes  # noqa: F401
 from . import job_routes  # noqa: F401
 from . import usdu_routes  # noqa: F401
+from . import ltx_routes  # noqa: F401

api/ltx_routes.py

@@ -0,0 +1,76 @@
+"""HTTP routes for distributed LTX latent tile results."""
+
+from __future__ import annotations
+
+import time
+
+from aiohttp import web
+import server
+
+from ..upscale.job_models import TileJobState
+from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized
+from ..utils.logging import debug_log
+from ..utils.network import handle_api_error
+
+
+@server.PromptServer.instance.routes.post("/distributed/ltx/submit_tile")
+async def submit_ltx_tile_endpoint(request):
+    """Accept one processed LTX latent tile from a worker.
+
+    Body is JSON:
+    - multi_job_id: distributed job id
+    - worker_id: worker id
+    - tile_index: spatial tile id, or -1 for completion-only messages
+    - payload: base64 safetensors payload produced by nodes.ltx_tiled_sampler
+    - is_last: optional worker completion flag
+    """
+    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.json()
+        multi_job_id = data.get("multi_job_id")
+        worker_id = str(data.get("worker_id") or "")
+        is_last = bool(data.get("is_last", False))
+        payload = data.get("payload")
+        tile_index = data.get("tile_index")
+
+        if not multi_job_id or not worker_id:
+            return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
+
+        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 isinstance(job_data, TileJobState):
+                return await handle_api_error(request, "Job not configured for tile submissions", 404)
+
+            job_data.worker_status[worker_id] = time.time()
+
+            if is_last and (payload is None or int(tile_index or -1) < 0):
+                await job_data.queue.put({
+                    "worker_id": worker_id,
+                    "is_last": True,
+                    "tiles": [],
+                })
+                debug_log(f"Received LTX completion signal from worker {worker_id}")
+                return web.json_response({"status": "success"})
+
+            if payload is None or tile_index is None:
+                return await handle_api_error(request, "Missing tile_index or payload", 400)
+
+            tile_index = int(tile_index)
+            await job_data.queue.put({
+                "worker_id": worker_id,
+                "is_last": is_last,
+                "tiles": [{
+                    "tile_idx": tile_index,
+                    "global_idx": tile_index,
+                    "payload": payload,
+                    "worker_id": worker_id,
+                }],
+            })
+            debug_log(f"Received LTX tile {tile_index} from worker {worker_id} (is_last={is_last})")
+            return web.json_response({"status": "success"})
+    except Exception as exc:
+        return await handle_api_error(request, exc, 500)

api/orchestration/prompt_transform.py

@@ -125,6 +125,209 @@ def _find_upstream_nodes(prompt_obj, start_ids):
     return connected
 
 
+_DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES = {
+    "PrimitiveBoolean",
+    "PrimitiveFloat",
+    "PrimitiveInt",
+    "PrimitiveNode",
+    "PrimitiveString",
+}
+
+_DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES = {
+    "LoadImage",
+}
+
+_DELEGATE_MASTER_SAFE_SCALAR_TYPES = {"BOOLEAN", "FLOAT", "INT", "STRING"}
+_DELEGATE_MASTER_SAFE_LIST_TYPES = {"LIST"}
+
+# ComfyUI 0.23 exposes CreateList via the newer schema API rather than the
+# legacy RETURN_TYPES/INPUT_TYPES attributes. Treat it as a safe config utility
+# only after its connected inputs recursively prove safe.
+_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES = {"CreateList"}
+_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES = {
+    "CreateList": ("inputs.",),
+}
+
+# Test hook. At runtime this stays None and the ComfyUI node registry is loaded lazily.
+_DELEGATE_MASTER_NODE_CLASS_MAPPINGS = None
+
+
+def _get_delegate_master_node_class_mappings():
+    """Return ComfyUI node-class mappings when available."""
+    if _DELEGATE_MASTER_NODE_CLASS_MAPPINGS is not None:
+        return _DELEGATE_MASTER_NODE_CLASS_MAPPINGS
+    try:
+        import nodes as comfy_nodes  # type: ignore
+    except Exception:  # pragma: no cover - depends on ComfyUI runtime imports
+        return {}
+    return getattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {}) or {}
+
+
+def _get_delegate_master_node_class(class_type):
+    mappings = _get_delegate_master_node_class_mappings()
+    return mappings.get(class_type) if isinstance(mappings, dict) else None
+
+
+def _normalize_delegate_master_return_type(return_type):
+    if return_type is None:
+        return ""
+    return str(return_type).strip().upper()
+
+
+def _delegate_master_type_is_safe_scalar(type_name):
+    return type_name in _DELEGATE_MASTER_SAFE_SCALAR_TYPES
+
+
+def _delegate_master_type_is_safe_config(type_name):
+    return _delegate_master_type_is_safe_scalar(type_name) or type_name in _DELEGATE_MASTER_SAFE_LIST_TYPES
+
+
+def _delegate_master_output_is_safe_scalar(class_type, output_index):
+    """Return True when a registered node output is lightweight config data."""
+    if class_type in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES:
+        return True
+    node_class = _get_delegate_master_node_class(class_type)
+    return_types = getattr(node_class, "RETURN_TYPES", ()) if node_class is not None else ()
+    try:
+        output_type = return_types[int(output_index)]
+    except (IndexError, TypeError, ValueError):
+        return False
+    return _delegate_master_type_is_safe_config(_normalize_delegate_master_return_type(output_type))
+
+
+def _get_delegate_master_input_types(class_type):
+    node_class = _get_delegate_master_node_class(class_type)
+    input_types = getattr(node_class, "INPUT_TYPES", None) if node_class is not None else None
+    if callable(input_types):
+        try:
+            input_types = input_types()
+        except TypeError:
+            return {}
+    return input_types if isinstance(input_types, dict) else {}
+
+
+def _normalize_delegate_master_input_type(input_spec):
+    if isinstance(input_spec, (list, tuple)) and input_spec:
+        return _normalize_delegate_master_return_type(input_spec[0])
+    return _normalize_delegate_master_return_type(input_spec)
+
+
+def _delegate_master_input_is_safe_scalar(class_type, input_name):
+    """Return True when a registered downstream input expects config data."""
+    for prefix in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES.get(class_type, ()):
+        if input_name.startswith(prefix):
+            return True
+    input_types = _get_delegate_master_input_types(class_type)
+    for section_name in ("required", "optional"):
+        section = input_types.get(section_name, {})
+        if isinstance(section, dict) and input_name in section:
+            input_type = _normalize_delegate_master_input_type(section[input_name])
+            return _delegate_master_type_is_safe_config(input_type)
+    return False
+
+
+def _is_delegate_master_always_retained_upstream_node(node):
+    if not isinstance(node, dict):
+        return False
+    class_type = node.get("class_type")
+    return isinstance(class_type, str) and class_type in _DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES
+
+
+def _is_delegate_master_retained_upstream_node(node, output_index=0):
+    """Return True for lightweight upstream nodes safe to keep on the master."""
+    if not isinstance(node, dict):
+        return False
+    class_type = node.get("class_type")
+    if not isinstance(class_type, str):
+        return False
+    return (
+        class_type in _DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES
+        or class_type.startswith("Primitive")
+        or _delegate_master_output_is_safe_scalar(class_type, output_index)
+    )
+
+
+def _collect_delegate_master_retained_upstream_branch(
+    prompt_obj,
+    node_id,
+    output_index,
+    memo,
+    visiting,
+):
+    """Return safe retained branch nodes, or None when the branch is not safe."""
+    node_id = str(node_id)
+    cache_key = (node_id, output_index)
+    if cache_key in memo:
+        cached = memo[cache_key]
+        return None if cached is None else set(cached)
+    if cache_key in visiting:
+        memo[cache_key] = None
+        return None
+
+    node = prompt_obj.get(node_id)
+    if not _is_delegate_master_retained_upstream_node(node, output_index):
+        memo[cache_key] = None
+        return None
+
+    visiting.add(cache_key)
+    retained = {node_id}
+    inputs = node.get("inputs", {}) if isinstance(node, dict) else {}
+    class_type = node.get("class_type") if isinstance(node, dict) else None
+    for input_name, value in inputs.items():
+        if not (isinstance(value, list) and len(value) == 2):
+            continue
+        if not _delegate_master_input_is_safe_scalar(class_type, input_name):
+            visiting.remove(cache_key)
+            memo[cache_key] = None
+            return None
+        source_id = str(value[0])
+        branch = _collect_delegate_master_retained_upstream_branch(
+            prompt_obj,
+            source_id,
+            value[1],
+            memo,
+            visiting,
+        )
+        if branch is None:
+            visiting.remove(cache_key)
+            memo[cache_key] = None
+            return None
+        retained.update(branch)
+
+    visiting.remove(cache_key)
+    memo[cache_key] = frozenset(retained)
+    return retained
+
+
+def _find_delegate_master_retained_upstream_nodes(prompt_obj, start_ids):
+    """Return lightweight upstream nodes needed by kept delegate-master nodes."""
+    connected = set()
+    memo = {}
+    for node_id in start_ids:
+        node = prompt_obj.get(str(node_id)) or {}
+        inputs = node.get("inputs", {})
+        class_type = node.get("class_type") if isinstance(node, dict) else None
+        for input_name, value in inputs.items():
+            if not (isinstance(value, list) and len(value) == 2):
+                continue
+            source_node = prompt_obj.get(str(value[0]))
+            if _is_delegate_master_always_retained_upstream_node(source_node):
+                connected.add(str(value[0]))
+                continue
+            if not _delegate_master_input_is_safe_scalar(class_type, input_name):
+                continue
+            branch = _collect_delegate_master_retained_upstream_branch(
+                prompt_obj,
+                value[0],
+                value[1],
+                memo,
+                set(),
+            )
+            if branch is not None:
+                connected.update(branch)
+    return connected
+
+
 def prune_prompt_for_worker(prompt_obj):
     """Prune worker prompt to distributed nodes and their upstream dependencies."""
     collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
@@ -167,6 +370,9 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
     downstream = _find_downstream_nodes(prompt_obj, collector_ids)
     nodes_to_keep = set(collector_ids)
     nodes_to_keep.update(downstream)
+    nodes_to_keep.update(
+        _find_delegate_master_retained_upstream_nodes(prompt_obj, nodes_to_keep)
+    )
 
     pruned_prompt = {}
     for node_id in nodes_to_keep:

nodes/__init__.py

@@ -10,6 +10,10 @@ from .utilities import (
     any_type,
 )
 from .collector import DistributedCollectorNode
+from .ltx_tiled_sampler import (
+    NODE_CLASS_MAPPINGS as LTX_CLASS_MAPPINGS,
+    NODE_DISPLAY_NAME_MAPPINGS as LTX_DISPLAY_NAME_MAPPINGS,
+)
 
 NODE_CLASS_MAPPINGS = {
     "DistributedCollector": DistributedCollectorNode,
@@ -19,6 +23,7 @@ NODE_CLASS_MAPPINGS = {
     "ImageBatchDivider": ImageBatchDivider,
     "AudioBatchDivider": AudioBatchDivider,
     "DistributedEmptyImage": DistributedEmptyImage,
+    **LTX_CLASS_MAPPINGS,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "DistributedCollector": "Distributed Collector",
@@ -28,4 +33,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "ImageBatchDivider": "Image Batch Divider",
     "AudioBatchDivider": "Audio Batch Divider",
     "DistributedEmptyImage": "Distributed Empty Image",
+    **LTX_DISPLAY_NAME_MAPPINGS,
 }

nodes/collector.py

@@ -22,6 +22,7 @@ prompt_server = _server.PromptServer.instance
 
 
 class DistributedCollectorNode:
+    INPUT_IS_LIST = True
     EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
 
     @classmethod
@@ -55,7 +56,65 @@ class DistributedCollectorNode:
     FUNCTION = "run"
     CATEGORY = "image"
     
+    @staticmethod
+    def _unwrap_list_input(value):
+        """Unwrap scalar inputs when ComfyUI passes them via INPUT_IS_LIST."""
+        if isinstance(value, (list, tuple)) and len(value) == 1:
+            return value[0]
+        return value
+
+    def _normalize_images_input(self, images):
+        """Collapse ComfyUI list IMAGE inputs into a normal batched IMAGE tensor."""
+        if isinstance(images, (list, tuple)):
+            if not images:
+                raise ValueError("Collector received an empty image list")
+            if not all(isinstance(image, torch.Tensor) for image in images):
+                raise TypeError("Collector expected IMAGE list items to be torch.Tensor instances")
+            if len(images) == 1:
+                return ensure_contiguous(images[0])
+            return ensure_contiguous(torch.cat([ensure_contiguous(image) for image in images], dim=0))
+        return ensure_contiguous(images)
+
+    def _normalize_audio_input(self, audio):
+        """Collapse ComfyUI list AUDIO inputs into a single AUDIO payload when present."""
+        if not isinstance(audio, (list, tuple)):
+            return audio
+
+        audio_items = [item for item in audio if item is not None]
+        if not audio_items:
+            return None
+        if len(audio_items) == 1:
+            return audio_items[0]
+
+        waveforms = []
+        sample_rate = 44100
+        for item in audio_items:
+            if not isinstance(item, dict):
+                raise TypeError("Collector expected AUDIO list items to be dictionaries")
+            waveform = item.get("waveform")
+            if waveform is None or waveform.numel() == 0:
+                continue
+            waveforms.append(waveform)
+            if sample_rate == 44100:
+                sample_rate = item.get("sample_rate", 44100)
+
+        if not waveforms:
+            return None
+        return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
+
     def run(self, images, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
+        images = self._normalize_images_input(images)
+        audio = self._normalize_audio_input(audio)
+        load_balance = self._unwrap_list_input(load_balance)
+        multi_job_id = self._unwrap_list_input(multi_job_id)
+        is_worker = self._unwrap_list_input(is_worker)
+        master_url = self._unwrap_list_input(master_url)
+        enabled_worker_ids = self._unwrap_list_input(enabled_worker_ids)
+        worker_batch_size = self._unwrap_list_input(worker_batch_size)
+        worker_id = self._unwrap_list_input(worker_id)
+        pass_through = self._unwrap_list_input(pass_through)
+        delegate_only = self._unwrap_list_input(delegate_only)
+
         # Create empty audio if not provided
         empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
 

nodes/ltx_tiled_sampler.py

@@ -0,0 +1,668 @@
+"""Distributed LTX tiled sampler node.
+
+This node adapts the upstream LTX tiled sampler interface to the existing
+ComfyUI-Distributed job queue. It distributes spatial latent tiles across the
+same worker/master infrastructure used by the distributed upscaler instead of
+creating a parallel queue implementation.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import copy
+import math
+from typing import Callable, Dict, Iterable, List, Optional, Tuple
+
+import aiohttp
+import torch
+from safetensors.torch import load as safetensors_load
+from safetensors.torch import save as safetensors_save
+
+from ..upscale.job_models import TileJobState
+from ..upscale.job_store import (
+    _cleanup_job,
+    _drain_results_queue,
+    _get_completed_count,
+    _mark_task_completed,
+    ensure_tile_jobs_initialized,
+    init_static_job_batched,
+)
+from ..upscale.job_timeout import _check_and_requeue_timed_out_workers
+from ..utils.async_helpers import run_async_in_server_loop
+from ..utils.constants import HEARTBEAT_INTERVAL, TILE_WAIT_TIMEOUT
+from ..utils.logging import debug_log, log
+from ..utils.network import get_client_session
+from ..utils.usdu_managment import _send_heartbeat_to_master
+
+try:  # Optional external LTX tiled sampler dependency.
+    from latent_tiled_sampler import LTXTiledSampler  # type: ignore
+except Exception:  # pragma: no cover - depends on optional custom node install.
+    LTXTiledSampler = None  # type: ignore
+
+
+def _default_ltx_inputs() -> dict:
+    """Fallback schema used when the source LTX sampler is not importable."""
+    return {
+        "required": {
+            "noise": ("NOISE",),
+            "guider": ("GUIDER",),
+            "sampler": ("SAMPLER",),
+            "sigmas": ("SIGMAS",),
+            "latent_image": ("LATENT",),
+            "h_tiles": ("INT", {"default": 1, "min": 1, "max": 64}),
+            "w_tiles": ("INT", {"default": 1, "min": 1, "max": 64}),
+            "overlap": ("INT", {"default": 0, "min": 0, "max": 4096}),
+        },
+        "optional": {
+            "denoise_mask": ("MASK",),
+        },
+    }
+
+
+def _copy_base_inputs() -> dict:
+    if LTXTiledSampler is None:
+        return _default_ltx_inputs()
+    try:
+        return copy.deepcopy(LTXTiledSampler.INPUT_TYPES())
+    except Exception:
+        return _default_ltx_inputs()
+
+
+def _compute_tile_starts(total_size: int, n_tiles: int, overlap: int) -> Tuple[List[int], int]:
+    """Compute evenly spaced tile starts and tile size for one dimension."""
+    total_size = int(total_size)
+    n_tiles = max(1, int(n_tiles))
+    overlap = max(0, int(overlap))
+    if n_tiles <= 1 or total_size <= 1:
+        return [0], total_size
+    tile_size = math.ceil((total_size + (n_tiles - 1) * overlap) / n_tiles)
+    tile_size = min(tile_size, total_size)
+    if n_tiles == 2:
+        return [0, max(0, total_size - tile_size)], tile_size
+    stride = (total_size - tile_size) / (n_tiles - 1)
+    return [int(round(i * stride)) for i in range(n_tiles)], tile_size
+
+
+def _make_window_1d(size: int, fade_left: int, fade_right: int, dtype, device):
+    """Create a one-dimensional cosine edge window for tile blending."""
+    win = torch.ones(int(size), dtype=dtype, device=device)
+    if fade_left > 0:
+        fl = min(int(fade_left), int(size))
+        i = torch.arange(fl, dtype=dtype, device=device)
+        # Use fl rather than fl - 1 to preserve a nonzero center plateau for small fades.
+        win[:fl] = 0.5 * (1.0 - torch.cos(math.pi * i / max(1, fl)))
+    if fade_right > 0:
+        fr = min(int(fade_right), int(size))
+        i = torch.arange(fr, dtype=dtype, device=device)
+        win[int(size) - fr:] = 0.5 * (1.0 + torch.cos(math.pi * i / max(1, fr)))
+    return win
+
+
+def _pack_tensors(tensors: Dict[str, torch.Tensor]) -> str:
+    """Serialize tensors to a base64 safetensors payload for JSON transport."""
+    cpu_tensors = {
+        name: tensor.detach().contiguous().to("cpu")
+        for name, tensor in tensors.items()
+        if tensor is not None
+    }
+    payload = safetensors_save(cpu_tensors)
+    return base64.b64encode(payload).decode("ascii")
+
+
+def _unpack_tensors(payload_b64: str, device) -> Dict[str, torch.Tensor]:
+    payload = base64.b64decode(payload_b64.encode("ascii"))
+    return {name: tensor.to(device) for name, tensor in safetensors_load(payload).items()}
+
+
+def _extract_tensor_and_rebuilder(value) -> Tuple[torch.Tensor, Callable[[torch.Tensor], object]]:
+    """Extract the tensor from common Comfy latent containers and preserve shape."""
+    if isinstance(value, dict) and "samples" in value:
+        original = value
+
+        def rebuild(tensor):
+            updated = dict(original)
+            updated["samples"] = tensor
+            return updated
+
+        return value["samples"], rebuild
+
+    if isinstance(value, tuple) and value and torch.is_tensor(value[0]):
+        original = value
+
+        def rebuild(tensor):
+            return (tensor, *original[1:])
+
+        return value[0], rebuild
+
+    if isinstance(value, list) and value and torch.is_tensor(value[0]):
+        original = value
+
+        def rebuild(tensor):
+            updated = list(original)
+            updated[0] = tensor
+            return updated
+
+        return value[0], rebuild
+
+    if torch.is_tensor(value):
+        return value, lambda tensor: tensor
+
+    raise TypeError("Expected latent_image to be a tensor or latent container with a tensor samples field")
+
+
+def _extract_optional_mask(mask):
+    if mask is None:
+        return None
+    if isinstance(mask, dict) and "samples" in mask:
+        return mask["samples"]
+    return mask
+
+
+def _tile_latent(value: torch.Tensor, tile: dict) -> torch.Tensor:
+    return value[:, :, :, tile["h_start"]:tile["h_end"], tile["w_start"]:tile["w_end"]].contiguous()
+
+
+def _tile_mask(mask, tile: dict):
+    if mask is None:
+        return None
+    if torch.is_tensor(mask) and mask.ndim >= 5:
+        return mask[:, :, :, tile["h_start"]:tile["h_end"], tile["w_start"]:tile["w_end"]].contiguous()
+    if torch.is_tensor(mask) and mask.ndim >= 4:
+        return mask[:, :, tile["h_start"]:tile["h_end"], tile["w_start"]:tile["w_end"]].contiguous()
+    return mask
+
+
+def _sample_tile(guider, tile_noise, tile_latent, sampler, sigmas, tile_mask):
+    result = guider.sample(
+        tile_noise,
+        tile_latent,
+        sampler,
+        sigmas,
+        denoise_mask=tile_mask,
+        disable_pbar=True,
+    )
+    if isinstance(result, tuple):
+        tile_samples = result[0]
+        tile_denoised = result[1] if len(result) > 1 else None
+    else:
+        tile_samples = result
+        tile_denoised = None
+    return tile_samples, tile_denoised
+
+
+def _build_tile_queue(latent_shape: Iterable[int], n_h_tiles: int, n_w_tiles: int, overlap: int) -> List[dict]:
+    _batch, _channels, _frames, height, width = [int(v) for v in latent_shape]
+    starts_y, tile_h = _compute_tile_starts(height, n_h_tiles, overlap)
+    starts_x, tile_w = _compute_tile_starts(width, n_w_tiles, overlap)
+    queue = []
+    tile_idx = 0
+    for grid_y, y_start in enumerate(starts_y):
+        y_end = min(y_start + tile_h, height)
+        for grid_x, x_start in enumerate(starts_x):
+            x_end = min(x_start + tile_w, width)
+            queue.append(
+                {
+                    "tile_index": tile_idx,
+                    "h_start": y_start,
+                    "h_end": y_end,
+                    "w_start": x_start,
+                    "w_end": x_end,
+                    "grid_y": grid_y,
+                    "grid_x": grid_x,
+                }
+            )
+            tile_idx += 1
+    return queue
+
+
+def _blend_tile_into_sums(
+    tile: dict,
+    samples: torch.Tensor,
+    denoised: Optional[torch.Tensor],
+    output_sum: torch.Tensor,
+    weight_sum: torch.Tensor,
+    denoised_sum: Optional[torch.Tensor],
+    n_h_tiles: int,
+    n_w_tiles: int,
+    overlap: int,
+):
+    y_start = tile["h_start"]
+    y_end = tile["h_end"]
+    x_start = tile["w_start"]
+    x_end = tile["w_end"]
+    grid_y = tile["grid_y"]
+    grid_x = tile["grid_x"]
+    fade_top = overlap if grid_y > 0 else 0
+    fade_bottom = overlap if grid_y < n_h_tiles - 1 else 0
+    fade_left = overlap if grid_x > 0 else 0
+    fade_right = overlap if grid_x < n_w_tiles - 1 else 0
+    dtype = output_sum.dtype
+    device = output_sum.device
+    wy = _make_window_1d(y_end - y_start, fade_top, fade_bottom, dtype, device)
+    wx = _make_window_1d(x_end - x_start, fade_left, fade_right, dtype, device)
+    weights = wy.view(1, 1, 1, y_end - y_start, 1) * wx.view(1, 1, 1, 1, x_end - x_start)
+    output_sum[:, :, :, y_start:y_end, x_start:x_end] += samples.to(device=device, dtype=dtype) * weights
+    weight_sum[:, :, :, y_start:y_end, x_start:x_end] += weights
+    if denoised_sum is not None and denoised is not None:
+        denoised_sum[:, :, :, y_start:y_end, x_start:x_end] += denoised.to(device=device, dtype=dtype) * weights
+
+
+async def _get_next_ltx_tile_index(multi_job_id: str) -> Optional[int]:
+    prompt_server = ensure_tile_jobs_initialized()
+    pending_queue = None
+    async with prompt_server.distributed_tile_jobs_lock:
+        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        if isinstance(job_data, TileJobState):
+            pending_queue = job_data.pending_tasks
+    if pending_queue is None:
+        return None
+    try:
+        return await asyncio.wait_for(pending_queue.get(), timeout=0.1)
+    except asyncio.TimeoutError:
+        return None
+
+
+async def _get_ltx_completed_tasks(multi_job_id: str) -> dict:
+    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 isinstance(job_data, TileJobState):
+            return dict(job_data.completed_tasks)
+    return {}
+
+
+async def _check_ltx_job_status(multi_job_id: str, master_url: str) -> bool:
+    session = await get_client_session()
+    async with session.get(f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}") as response:
+        if response.status != 200:
+            return False
+        data = await response.json()
+        return bool(data.get("ready", False))
+
+
+async def _request_tile_from_master(multi_job_id: str, master_url: str, worker_id: str):
+    session = await get_client_session()
+    async with session.post(
+        f"{master_url}/distributed/request_image",
+        json={"multi_job_id": multi_job_id, "worker_id": str(worker_id)},
+        timeout=aiohttp.ClientTimeout(total=TILE_WAIT_TIMEOUT),
+    ) as response:
+        if response.status == 200:
+            return await response.json()
+        if response.status == 404:
+            return None
+        response.raise_for_status()
+
+
+async def _submit_ltx_tile_to_master(
+    multi_job_id: str,
+    master_url: str,
+    worker_id: str,
+    tile_index: int,
+    payload_b64: Optional[str] = None,
+    is_last: bool = False,
+):
+    session = await get_client_session()
+    body = {
+        "multi_job_id": multi_job_id,
+        "worker_id": str(worker_id),
+        "tile_index": int(tile_index) if tile_index is not None else None,
+        "payload": payload_b64,
+        "is_last": bool(is_last),
+    }
+    async with session.post(f"{master_url}/distributed/ltx/submit_tile", json=body) as response:
+        response.raise_for_status()
+        return await response.json()
+
+
+async def _ltx_worker_process(
+    node,
+    noise,
+    guider,
+    sampler,
+    sigmas,
+    latent_tensor,
+    denoise_mask,
+    n_h_tiles,
+    n_w_tiles,
+    tile_overlap,
+):
+    job_id = node.multi_job_id
+    master_url = node.master_url
+    worker_id = node.worker_id or "ltx-worker"
+    tiles_by_index = {
+        tile["tile_index"]: tile
+        for tile in _build_tile_queue(latent_tensor.shape, n_h_tiles, n_w_tiles, tile_overlap)
+    }
+    for attempt in range(120):
+        if await _check_ltx_job_status(job_id, master_url):
+            break
+        await asyncio.sleep(1.0)
+    else:
+        log(f"LTX Dist Worker[{worker_id[:8]}]: job {job_id} was not ready; aborting")
+        return
+
+    full_noise = noise.generate_noise({"samples": latent_tensor})
+
+    last_heartbeat = 0.0
+    while True:
+        data = await _request_tile_from_master(job_id, master_url, worker_id)
+        if not data:
+            log(f"LTX Dist Worker[{worker_id[:8]}]: master job unavailable; stopping")
+            break
+        tile_index = data.get("tile_idx")
+        if tile_index is None:
+            await _submit_ltx_tile_to_master(job_id, master_url, worker_id, -1, is_last=True)
+            break
+        tile = tiles_by_index.get(int(tile_index))
+        if tile is None:
+            raise RuntimeError(f"Master assigned unknown LTX tile index {tile_index}")
+
+        tile_latent = _tile_latent(latent_tensor, tile)
+        tile_noise = _tile_latent(full_noise, tile)
+        tile_mask = _tile_mask(denoise_mask, tile)
+        tile_samples, tile_denoised = _sample_tile(guider, tile_noise, tile_latent, sampler, sigmas, tile_mask)
+        tensors = {"samples": tile_samples}
+        if tile_denoised is not None:
+            tensors["denoised"] = tile_denoised
+        await _submit_ltx_tile_to_master(
+            job_id,
+            master_url,
+            worker_id,
+            int(tile_index),
+            _pack_tensors(tensors),
+        )
+        now = asyncio.get_running_loop().time()
+        if now - last_heartbeat >= HEARTBEAT_INTERVAL:
+            await _send_heartbeat_to_master(job_id, master_url, worker_id)
+            last_heartbeat = now
+
+
+async def _ltx_master_process(
+    node,
+    noise,
+    guider,
+    sampler,
+    sigmas,
+    latent_tensor,
+    denoise_mask,
+    n_h_tiles,
+    n_w_tiles,
+    tile_overlap,
+):
+    job_id = node.multi_job_id
+    enabled_workers = node.enabled_worker_ids
+    tiles = _build_tile_queue(latent_tensor.shape, n_h_tiles, n_w_tiles, tile_overlap)
+    tiles_by_index = {tile["tile_index"]: tile for tile in tiles}
+    total_tiles = len(tiles)
+    await init_static_job_batched(job_id, batch_size=1, num_tiles_per_image=total_tiles, enabled_workers=enabled_workers)
+
+    full_noise = noise.generate_noise({"samples": latent_tensor})
+    output_sum = torch.zeros_like(latent_tensor)
+    weight_sum = torch.zeros_like(latent_tensor)
+    denoised_sum = torch.zeros_like(latent_tensor) if denoise_mask is not None else None
+    integrated = set()
+    last_timeout_check = 0.0
+
+    try:
+        while len(integrated) < total_tiles:
+            await _drain_results_queue(job_id)
+            completed = await _get_ltx_completed_tasks(job_id)
+            for tile_index, result in list(completed.items()):
+                if tile_index in integrated:
+                    continue
+                if not isinstance(result, dict) or "payload" not in result:
+                    # Locally completed tiles are blended immediately and then marked complete.
+                    continue
+                tile = tiles_by_index.get(int(tile_index))
+                if tile is None:
+                    continue
+                tensors = _unpack_tensors(result["payload"], latent_tensor.device)
+                _blend_tile_into_sums(
+                    tile,
+                    tensors["samples"],
+                    tensors.get("denoised"),
+                    output_sum,
+                    weight_sum,
+                    denoised_sum,
+                    n_h_tiles,
+                    n_w_tiles,
+                    tile_overlap,
+                )
+                integrated.add(int(tile_index))
+
+            if len(integrated) >= total_tiles:
+                break
+
+            tile_index = await _get_next_ltx_tile_index(job_id)
+            if tile_index is not None and int(tile_index) not in integrated:
+                tile = tiles_by_index[int(tile_index)]
+                tile_latent = _tile_latent(latent_tensor, tile)
+                tile_noise = _tile_latent(full_noise, tile)
+                tile_mask = _tile_mask(denoise_mask, tile)
+                tile_samples, tile_denoised = _sample_tile(guider, tile_noise, tile_latent, sampler, sigmas, tile_mask)
+                _blend_tile_into_sums(
+                    tile,
+                    tile_samples,
+                    tile_denoised,
+                    output_sum,
+                    weight_sum,
+                    denoised_sum,
+                    n_h_tiles,
+                    n_w_tiles,
+                    tile_overlap,
+                )
+                integrated.add(int(tile_index))
+                await _mark_task_completed(job_id, int(tile_index), {"tile_idx": int(tile_index), "source": "master"})
+                continue
+
+            now = asyncio.get_running_loop().time()
+            if now - last_timeout_check >= 10.0:
+                requeued = await _check_and_requeue_timed_out_workers(job_id, total_tiles)
+                if requeued:
+                    log(f"LTX Dist: Requeued {requeued} timed-out tile(s)")
+                last_timeout_check = now
+
+            completed_count = await _get_completed_count(job_id)
+            debug_log(f"LTX Dist: Progress {completed_count}/{total_tiles} completed, {len(integrated)} integrated")
+            await asyncio.sleep(0.1)
+    finally:
+        await _cleanup_job(job_id)
+
+    weight_sum = torch.where(weight_sum == 0, torch.ones_like(weight_sum), weight_sum)
+    samples = output_sum / weight_sum
+    denoised = denoised_sum / weight_sum if denoised_sum is not None else None
+    return samples, denoised
+
+
+def _ltx_local_tiled_process(
+    noise,
+    guider,
+    sampler,
+    sigmas,
+    latent_tensor,
+    denoise_mask,
+    n_h_tiles,
+    n_w_tiles,
+    tile_overlap,
+):
+    """Process all LTX latent tiles locally in the current ComfyUI execution context.
+
+    This path is used when no workers are enabled. Keeping it synchronous avoids
+    crossing ComfyUI's execution/inference context boundary while still using the
+    same tile extraction and blend logic as the distributed master/worker path.
+    """
+    tiles = _build_tile_queue(latent_tensor.shape, n_h_tiles, n_w_tiles, tile_overlap)
+    full_noise = noise.generate_noise({"samples": latent_tensor})
+    output_sum = torch.zeros_like(latent_tensor)
+    weight_sum = torch.zeros_like(latent_tensor)
+    denoised_sum = None
+
+    for tile in tiles:
+        tile_latent = _tile_latent(latent_tensor, tile)
+        tile_noise = _tile_latent(full_noise, tile)
+        tile_mask = _tile_mask(denoise_mask, tile)
+        tile_samples, tile_denoised = _sample_tile(guider, tile_noise, tile_latent, sampler, sigmas, tile_mask)
+        if denoised_sum is None and tile_denoised is not None:
+            denoised_sum = torch.zeros_like(latent_tensor)
+        _blend_tile_into_sums(
+            tile,
+            tile_samples,
+            tile_denoised,
+            output_sum,
+            weight_sum,
+            denoised_sum,
+            n_h_tiles,
+            n_w_tiles,
+            tile_overlap,
+        )
+
+    weight_sum = torch.where(weight_sum == 0, torch.ones_like(weight_sum), weight_sum)
+    samples = output_sum / weight_sum
+    denoised = denoised_sum / weight_sum if denoised_sum is not None else None
+    return samples, denoised
+
+
+class LTXTiledSamplerDistributed:
+    """Distributed wrapper around the LTX tiled sampler."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        base_inputs = _copy_base_inputs()
+        base_inputs.setdefault("hidden", {})
+        base_inputs["hidden"].update(
+            {
+                "multi_job_id": ("STRING", {"default": ""}),
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "master_url": ("STRING", {"default": ""}),
+                "worker_id": ("STRING", {"default": ""}),
+                "enabled_worker_ids": ("STRING", {"default": "[]"}),
+            }
+        )
+        return base_inputs
+
+    RETURN_TYPES = getattr(LTXTiledSampler, "RETURN_TYPES", ("LATENT", "LATENT", "LATENT"))
+    RETURN_NAMES = getattr(LTXTiledSampler, "RETURN_NAMES", ("samples", "denoised", "preview"))
+    FUNCTION = "sample_tiled_distributed"
+    CATEGORY = "latent/diffusion"
+
+    @classmethod
+    def IS_CHANGED(cls, **kwargs):
+        return float("nan")
+
+    def __init__(self):
+        self.base = LTXTiledSampler() if LTXTiledSampler is not None else None
+
+    def _run_base(self, noise, guider, sampler, sigmas, latent_image, denoise_mask, h_tiles, w_tiles, overlap):
+        if self.base is None:
+            raise RuntimeError(
+                "LTXTiledSamplerDistributed requires the base LTXTiledSampler node for local fallback. "
+                "Install the LTX tiled sampler dependency or run with distributed multi_job_id."
+            )
+        return self.base.sample_tiled(
+            noise,
+            guider,
+            sampler,
+            sigmas,
+            latent_image,
+            denoise_mask,
+            h_tiles,
+            w_tiles,
+            overlap,
+        )
+
+    def sample_tiled_distributed(
+        self,
+        noise,
+        guider,
+        sampler,
+        sigmas,
+        latent_image,
+        denoise_mask=None,
+        h_tiles: int = 1,
+        w_tiles: int = 1,
+        overlap: int = 0,
+        multi_job_id: str = "",
+        is_worker: bool = False,
+        master_url: str = "",
+        worker_id: str = "",
+        enabled_worker_ids: str = "[]",
+    ):
+        if not multi_job_id:
+            return self._run_base(noise, guider, sampler, sigmas, latent_image, denoise_mask, h_tiles, w_tiles, overlap)
+
+        latent_tensor, rebuild_latent = _extract_tensor_and_rebuilder(latent_image)
+        if latent_tensor.ndim != 5:
+            raise ValueError(
+                f"LTXTiledSamplerDistributed expects a 5D latent tensor [B,C,F,H,W]; got shape {tuple(latent_tensor.shape)}"
+            )
+        mask_tensor = _extract_optional_mask(denoise_mask)
+        self.multi_job_id = str(multi_job_id)
+        self.master_url = str(master_url or "").rstrip("/")
+        self.worker_id = str(worker_id or "")
+        if isinstance(enabled_worker_ids, str):
+            try:
+                import json
+
+                parsed_workers = json.loads(enabled_worker_ids) if enabled_worker_ids else []
+            except Exception:
+                parsed_workers = []
+        else:
+            parsed_workers = enabled_worker_ids or []
+        self.enabled_worker_ids = [str(worker) for worker in parsed_workers]
+
+        if is_worker:
+            run_async_in_server_loop(
+                _ltx_worker_process(
+                    self,
+                    noise,
+                    guider,
+                    sampler,
+                    sigmas,
+                    latent_tensor,
+                    mask_tensor,
+                    int(h_tiles),
+                    int(w_tiles),
+                    int(overlap),
+                ),
+                timeout=None,
+            )
+            placeholder = rebuild_latent(latent_tensor)
+            return (placeholder, placeholder, placeholder)
+
+        if not self.enabled_worker_ids:
+            samples, denoised = _ltx_local_tiled_process(
+                noise,
+                guider,
+                sampler,
+                sigmas,
+                latent_tensor,
+                mask_tensor,
+                int(h_tiles),
+                int(w_tiles),
+                int(overlap),
+            )
+        else:
+            samples, denoised = run_async_in_server_loop(
+                _ltx_master_process(
+                    self,
+                    noise,
+                    guider,
+                    sampler,
+                    sigmas,
+                    latent_tensor,
+                    mask_tensor,
+                    int(h_tiles),
+                    int(w_tiles),
+                    int(overlap),
+                ),
+                timeout=None,
+            )
+        samples_latent = rebuild_latent(samples)
+        denoised_latent = rebuild_latent(denoised if denoised is not None else samples)
+        return (samples_latent, denoised_latent, samples_latent)
+
+
+NODE_CLASS_MAPPINGS = {"LTXTiledSamplerDistributed": LTXTiledSamplerDistributed}
+NODE_DISPLAY_NAME_MAPPINGS = {"LTXTiledSamplerDistributed": "LTX Tiled Sampler (Distributed)"}

pyproject.toml

@@ -1,9 +1,9 @@
 [project]
 name = "ComfyUI-Distributed"
 description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
-version = "1.4.3"
+version = "1.4.6"
 license = {file = "LICENSE"}
-dependencies  = []
+dependencies  = ["safetensors"]
 
 [project.urls]
 Repository = "https://github.com/robertvoy/ComfyUI-Distributed"

tests/test_collector_list_inputs.py

@@ -0,0 +1,205 @@
+import importlib.util
+import sys
+import types
+import asyncio
+from pathlib import Path
+
+import torch
+
+
+def _load_collector_module():
+    module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
+    package_name = "dist_collector_list_testpkg"
+
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    nodes_pkg = types.ModuleType(f"{package_name}.nodes")
+    nodes_pkg.__path__ = []
+    sys.modules[f"{package_name}.nodes"] = nodes_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    class _Routes:
+        def post(self, _path):
+            return lambda fn: fn
+
+        def get(self, _path):
+            return lambda fn: fn
+
+    prompt_server = types.SimpleNamespace(
+        routes=_Routes(),
+        distributed_jobs_lock=None,
+        distributed_pending_jobs={},
+    )
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
+    sys.modules["server"] = server_module
+
+    comfy_module = types.ModuleType("comfy")
+    model_management = types.ModuleType("comfy.model_management")
+
+    class InterruptProcessingException(Exception):
+        pass
+
+    model_management.InterruptProcessingException = InterruptProcessingException
+    model_management.throw_exception_if_processing_interrupted = lambda: None
+    comfy_module.model_management = model_management
+
+    comfy_utils = types.ModuleType("comfy.utils")
+
+    class ProgressBar:
+        def __init__(self, _total):
+            self.total = _total
+            self.updates = []
+
+        def update(self, value):
+            self.updates.append(value)
+
+    comfy_utils.ProgressBar = ProgressBar
+    comfy_module.utils = comfy_utils
+    sys.modules["comfy"] = comfy_module
+    sys.modules["comfy.model_management"] = model_management
+    sys.modules["comfy.utils"] = comfy_utils
+
+    aiohttp_module = types.ModuleType("aiohttp")
+    aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
+    sys.modules["aiohttp"] = aiohttp_module
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.get_worker_timeout_seconds = lambda: 0.1
+    config_module.load_config = lambda: {"workers": []}
+    config_module.is_master_delegate_only = lambda: False
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.HEARTBEAT_INTERVAL = 1.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+    def _ensure_contiguous(tensor):
+        return tensor.contiguous() if hasattr(tensor, "contiguous") else tensor
+
+    image_module.ensure_contiguous = _ensure_contiguous
+    image_module.tensor_to_pil = lambda *_args, **_kwargs: None
+    image_module.pil_to_tensor = lambda value: value
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.build_worker_url = lambda worker: "http://worker"
+    network_module.get_client_session = lambda: None
+    network_module.probe_worker = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
+    audio_payload_module.encode_audio_payload = lambda audio: audio
+    sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.collector", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_collector_opts_into_comfyui_list_inputs():
+    collector = _load_collector_module().DistributedCollectorNode
+
+    assert collector.INPUT_IS_LIST is True
+
+
+def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
+    collector = _load_collector_module().DistributedCollectorNode()
+    first = torch.zeros(1, 2, 2, 3)
+    second = torch.ones(1, 2, 2, 3)
+
+    images, audio = collector.run(
+        images=[first, second],
+        load_balance=[False],
+        audio=[None],
+        multi_job_id=[""],
+        is_worker=[False],
+        master_url=[""],
+        enabled_worker_ids=["[]"],
+        worker_batch_size=[1],
+        worker_id=[""],
+        pass_through=[False],
+        delegate_only=[False],
+    )
+
+    assert tuple(images.shape) == (2, 2, 2, 3)
+    assert torch.equal(images[0:1], first)
+    assert torch.equal(images[1:2], second)
+    assert tuple(audio["waveform"].shape) == (1, 2, 1)
+
+
+def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final_item():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    first = torch.zeros(1, 2, 2, 3)
+    second = torch.ones(1, 2, 2, 3)
+    posted_payloads = []
+
+    class _FakeImage:
+        def save(self, fp, format=None, compress_level=None):
+            fp.write(b"png-bytes")
+
+    class _FakeResponse:
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc_type, exc, tb):
+            return False
+
+        def raise_for_status(self):
+            return None
+
+    class _FakeSession:
+        def post(self, url, json, timeout):
+            posted_payloads.append(json)
+            return _FakeResponse()
+
+    async def _fake_get_client_session():
+        return _FakeSession()
+
+    module.tensor_to_pil = lambda *_args, **_kwargs: _FakeImage()
+    module.get_client_session = _fake_get_client_session
+    module.encode_audio_payload = lambda _audio: None
+
+    images, audio = collector.run(
+        images=[first, second],
+        load_balance=[False],
+        audio=[None],
+        multi_job_id=["job-list-1"],
+        is_worker=[True],
+        master_url=["http://master"],
+        enabled_worker_ids=["[]"],
+        worker_batch_size=[1],
+        worker_id=["worker-a"],
+        pass_through=[False],
+        delegate_only=[False],
+    )
+
+    assert tuple(images.shape) == (2, 2, 2, 3)
+    assert tuple(audio["waveform"].shape) == (1, 2, 1)
+    assert len(posted_payloads) == 2
+    assert [payload["batch_idx"] for payload in posted_payloads] == [0, 1]
+    assert [payload["is_last"] for payload in posted_payloads] == [False, True]
+    assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
+    assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}

tests/test_ltx_tiled_sampler.py

@@ -0,0 +1,157 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import torch
+
+
+def _load_ltx_module():
+    module_path = Path(__file__).resolve().parents[1] / "nodes" / "ltx_tiled_sampler.py"
+    package_name = "dist_ltx_testpkg"
+
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    for subpkg in ["nodes", "upscale", "utils"]:
+        pkg = types.ModuleType(f"{package_name}.{subpkg}")
+        pkg.__path__ = []
+        sys.modules[f"{package_name}.{subpkg}"] = pkg
+
+    job_models = types.ModuleType(f"{package_name}.upscale.job_models")
+
+    class TileJobState:
+        pass
+
+    job_models.TileJobState = TileJobState
+    sys.modules[f"{package_name}.upscale.job_models"] = job_models
+
+    async def _noop(*_args, **_kwargs):
+        return None
+
+    job_store = types.ModuleType(f"{package_name}.upscale.job_store")
+    job_store._cleanup_job = _noop
+    job_store._drain_results_queue = _noop
+    job_store._get_completed_count = _noop
+    job_store._mark_task_completed = _noop
+    job_store.ensure_tile_jobs_initialized = lambda: types.SimpleNamespace(
+        distributed_tile_jobs_lock=asyncio.Lock(),
+        distributed_pending_tile_jobs={},
+    )
+    job_store.init_static_job_batched = _noop
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store
+
+    job_timeout = types.ModuleType(f"{package_name}.upscale.job_timeout")
+    job_timeout._check_and_requeue_timed_out_workers = _noop
+    sys.modules[f"{package_name}.upscale.job_timeout"] = job_timeout
+
+    async_helpers = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers.run_async_in_server_loop = lambda coro, timeout=None: asyncio.run(coro)
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers
+
+    constants = types.ModuleType(f"{package_name}.utils.constants")
+    constants.HEARTBEAT_INTERVAL = 10.0
+    constants.TILE_WAIT_TIMEOUT = 1.0
+    sys.modules[f"{package_name}.utils.constants"] = constants
+
+    logging = types.ModuleType(f"{package_name}.utils.logging")
+    logging.debug_log = lambda *_args, **_kwargs: None
+    logging.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging
+
+    network = types.ModuleType(f"{package_name}.utils.network")
+
+    async def _missing_session():
+        raise RuntimeError("network not available in unit test")
+
+    network.get_client_session = _missing_session
+    sys.modules[f"{package_name}.utils.network"] = network
+
+    usdu = types.ModuleType(f"{package_name}.utils.usdu_managment")
+    usdu._send_heartbeat_to_master = _noop
+    sys.modules[f"{package_name}.utils.usdu_managment"] = usdu
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.ltx_tiled_sampler", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+ltx = _load_ltx_module()
+
+
+class LtxTiledSamplerTests(unittest.TestCase):
+    def test_compute_tile_starts_with_overlap_covers_edges(self):
+        starts, tile_size = ltx._compute_tile_starts(total_size=10, n_tiles=3, overlap=2)
+
+        self.assertEqual(starts[0], 0)
+        self.assertEqual(starts[-1] + tile_size, 10)
+        self.assertGreaterEqual(tile_size, 4)
+
+    def test_pack_unpack_tensors_round_trips_with_safetensors(self):
+        tensor = torch.arange(12, dtype=torch.float32).reshape(1, 1, 3, 2, 2)
+
+        payload = ltx._pack_tensors({"samples": tensor})
+        restored = ltx._unpack_tensors(payload, device="cpu")
+
+        self.assertTrue(torch.equal(restored["samples"], tensor))
+
+    def test_latent_dict_rebuilder_preserves_metadata(self):
+        samples = torch.zeros((1, 2, 3, 4, 5))
+        latent = {"samples": samples, "noise_mask": "keep-me"}
+
+        extracted, rebuild = ltx._extract_tensor_and_rebuilder(latent)
+        rebuilt = rebuild(torch.ones_like(extracted))
+
+        self.assertIs(extracted, samples)
+        self.assertEqual(rebuilt["noise_mask"], "keep-me")
+        self.assertTrue(torch.equal(rebuilt["samples"], torch.ones_like(samples)))
+
+    def test_node_schema_includes_distributed_hidden_inputs(self):
+        schema = ltx.LTXTiledSamplerDistributed.INPUT_TYPES()
+
+        self.assertIn("multi_job_id", schema["hidden"])
+        self.assertIn("is_worker", schema["hidden"])
+        self.assertIn("master_url", schema["hidden"])
+        self.assertIn("enabled_worker_ids", schema["hidden"])
+        self.assertIn("LTXTiledSamplerDistributed", ltx.NODE_CLASS_MAPPINGS)
+
+
+    def test_local_tiled_process_runs_without_async_master_path(self):
+        class _Noise:
+            def generate_noise(self, latent):
+                return torch.ones_like(latent['samples']) * 0.25
+
+        class _Guider:
+            def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, disable_pbar=True):
+                return latent_image + noise
+
+        latent = torch.zeros((1, 1, 2, 4, 4), dtype=torch.float32)
+
+        samples, denoised = ltx._ltx_local_tiled_process(
+            _Noise(),
+            _Guider(),
+            sampler=object(),
+            sigmas=torch.tensor([1.0, 0.0]),
+            latent_tensor=latent,
+            denoise_mask=None,
+            n_h_tiles=2,
+            n_w_tiles=2,
+            tile_overlap=0,
+        )
+
+        self.assertIsNone(denoised)
+        self.assertEqual(tuple(samples.shape), tuple(latent.shape))
+        self.assertTrue(torch.allclose(samples, torch.full_like(latent, 0.25)))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_prompt_transform.py

@@ -296,6 +296,501 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
         empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
         self.assertEqual(len(empty_nodes), 2)
 
+    def test_preserves_primitive_string_for_downstream_required_input(self):
+        """Delegate-only master keeps primitive inputs needed by SaveImage."""
+        save_image = type(
+            "SaveImage",
+            (),
+            {
+                "INPUT_TYPES": classmethod(
+                    lambda cls: {
+                        "required": {
+                            "images": ("IMAGE",),
+                            "filename_prefix": ("STRING",),
+                        }
+                    }
+                )
+            },
+        )
+        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
+        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {"SaveImage": save_image}
+        try:
+            prompt = {
+                "8": {"class_type": "KSampler", "inputs": {}},
+                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
+                "15": {"class_type": "PrimitiveString", "inputs": {"value": "test/input_bug"}},
+                "9": {
+                    "class_type": "SaveImage",
+                    "inputs": {
+                        "images": ["11", 0],
+                        "filename_prefix": ["15", 0],
+                    },
+                },
+            }
+
+            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
+        finally:
+            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
+
+        self.assertIn("15", result)
+        self.assertEqual(result["15"], prompt["15"])
+        self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["15", 0])
+
+    def test_does_not_preserve_non_primitive_upstream_for_collector(self):
+        prompt = _delegate_prompt()
+        result = pt.prepare_delegate_master_prompt(prompt, ["3"])
+
+        self.assertNotIn("2", result)
+        self.assertNotEqual(result["3"]["inputs"]["images"], ["2", 0])
+
+    def test_preserves_load_image_for_switch_alternate_required_input(self):
+        """Delegate-only master keeps LoadImage inputs needed by switches."""
+        prompt = {
+            "848": {"class_type": "LoadImage", "inputs": {"image": "input_bug2_00003_.png"}},
+            "862": {"class_type": "VAEDecode", "inputs": {}},
+            "854": {"class_type": "DistributedCollector", "inputs": {"images": ["862", 0]}},
+            "850": {
+                "class_type": "ComfySwitchNode",
+                "inputs": {
+                    "on_false": ["848", 0],
+                    "on_true": ["854", 0],
+                },
+            },
+            "851": {"class_type": "PreviewImage", "inputs": {"images": ["850", 0]}},
+        }
+
+        result = pt.prepare_delegate_master_prompt(prompt, ["854"])
+
+        self.assertIn("848", result)
+        self.assertEqual(result["850"]["inputs"]["on_false"], ["848", 0])
+        self.assertEqual(result["850"]["inputs"]["on_true"], ["854", 0])
+        self.assertNotIn("862", result)
+
+    def test_preserves_registered_string_utility_subgraph_for_downstream_required_input(self):
+        """Delegate-only master keeps scalar utility chains used by SaveImage."""
+        string_concat = type(
+            "StringConcatenate",
+            (),
+            {
+                "RETURN_TYPES": ("STRING",),
+                "INPUT_TYPES": classmethod(
+                    lambda cls: {
+                        "required": {
+                            "string_a": ("STRING",),
+                            "string_b": ("STRING",),
+                        }
+                    }
+                ),
+            },
+        )
+        save_image = type(
+            "SaveImage",
+            (),
+            {
+                "INPUT_TYPES": classmethod(
+                    lambda cls: {
+                        "required": {
+                            "images": ("IMAGE",),
+                            "filename_prefix": ("STRING",),
+                        }
+                    }
+                )
+            },
+        )
+        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
+        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
+            "SaveImage": save_image,
+            "StringConcatenate": string_concat,
+        }
+        try:
+            prompt = {
+                "8": {"class_type": "KSampler", "inputs": {}},
+                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
+                "15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
+                "16": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug3"}},
+                "17": {
+                    "class_type": "StringConcatenate",
+                    "inputs": {
+                        "string_a": ["15", 0],
+                        "string_b": ["16", 0],
+                    },
+                },
+                "9": {
+                    "class_type": "SaveImage",
+                    "inputs": {
+                        "images": ["11", 0],
+                        "filename_prefix": ["17", 0],
+                    },
+                },
+            }
+
+            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
+        finally:
+            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
+
+        self.assertIn("15", result)
+        self.assertIn("16", result)
+        self.assertIn("17", result)
+        self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["17", 0])
+        self.assertEqual(result["17"]["inputs"]["string_a"], ["15", 0])
+        self.assertEqual(result["17"]["inputs"]["string_b"], ["16", 0])
+
+    def test_preserves_registered_multi_string_join_subgraph_for_downstream_required_input(self):
+        """Delegate-only master keeps multi-input scalar utility chains."""
+        join_string_multi = type(
+            "JoinStringMulti",
+            (),
+            {
+                "RETURN_TYPES": ("STRING",),
+                "INPUT_TYPES": classmethod(
+                    lambda cls: {
+                        "required": {"string_1": ("STRING",)},
+                        "optional": {
+                            "string_2": ("STRING",),
+                            "string_3": ("STRING",),
+                        },
+                    }
+                ),
+            },
+        )
+        save_image = type(
+            "SaveImage",
+            (),
+            {
+                "INPUT_TYPES": classmethod(
+                    lambda cls: {
+                        "required": {
+                            "images": ("IMAGE",),
+                            "filename_prefix": ("STRING",),
+                        }
+                    }
+                )
+            },
+        )
+        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
+        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
+            "JoinStringMulti": join_string_multi,
+            "SaveImage": save_image,
+        }
+        try:
+            prompt = {
+                "8": {"class_type": "KSampler", "inputs": {}},
+                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
+                "15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
+                "16": {"class_type": "PrimitiveString", "inputs": {"value": "input"}},
+                "17": {"class_type": "PrimitiveString", "inputs": {"value": "bug4"}},
+                "18": {
+                    "class_type": "JoinStringMulti",
+                    "inputs": {
+                        "string_1": ["15", 0],
+                        "string_2": ["16", 0],
+                        "string_3": ["17", 0],
+                    },
+                },
+                "9": {
+                    "class_type": "SaveImage",
+                    "inputs": {
+                        "images": ["11", 0],
+                        "filename_prefix": ["18", 0],
+                    },
+                },
+            }
+
+            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
+        finally:
+            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
+
+        self.assertIn("15", result)
+        self.assertIn("16", result)
+        self.assertIn("17", result)
+        self.assertIn("18", result)
+        self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["18", 0])
+
+    def test_does_not_preserve_scalar_utility_with_heavy_upstream_dependency(self):
+        """Scalar utility nodes are retained only when their full input branch is safe."""
+        string_concat = type(
+            "StringConcatenate",
+            (),
+            {
+                "RETURN_TYPES": ("STRING",),
+                "INPUT_TYPES": classmethod(
+                    lambda cls: {
+                        "required": {
+                            "string_a": ("STRING",),
+                            "string_b": ("STRING",),
+                        }
+                    }
+                ),
+            },
+        )
+        save_image = type(
+            "SaveImage",
+            (),
+            {
+                "INPUT_TYPES": classmethod(
+                    lambda cls: {
+                        "required": {
+                            "images": ("IMAGE",),
+                            "filename_prefix": ("STRING",),
+                        }
+                    }
+                )
+            },
+        )
+        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
+        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
+            "SaveImage": save_image,
+            "StringConcatenate": string_concat,
+        }
+        try:
+            prompt = {
+                "8": {"class_type": "KSampler", "inputs": {}},
+                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
+                "15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
+                "17": {
+                    "class_type": "StringConcatenate",
+                    "inputs": {
+                        "string_a": ["15", 0],
+                        "string_b": ["8", 0],
+                    },
+                },
+                "9": {
+                    "class_type": "SaveImage",
+                    "inputs": {
+                        "images": ["11", 0],
+                        "filename_prefix": ["17", 0],
+                    },
+                },
+            }
+
+            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
+        finally:
+            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
+
+        self.assertNotIn("8", result)
+        self.assertNotIn("17", result)
+        self.assertNotIn("filename_prefix", result["9"]["inputs"])
+
+    def test_does_not_preserve_scalar_output_for_non_scalar_downstream_input(self):
+        """Scalar outputs are retained only for scalar/config downstream inputs."""
+        string_provider = type("StringProvider", (), {"RETURN_TYPES": ("STRING",)})
+        image_consumer = type(
+            "ImageConsumer",
+            (),
+            {
+                "INPUT_TYPES": classmethod(
+                    lambda cls: {
+                        "required": {
+                            "images": ("IMAGE",),
+                            "mask": ("IMAGE",),
+                        }
+                    }
+                )
+            },
+        )
+        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
+        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
+            "ImageConsumer": image_consumer,
+            "StringProvider": string_provider,
+        }
+        try:
+            prompt = {
+                "8": {"class_type": "KSampler", "inputs": {}},
+                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
+                "17": {"class_type": "StringProvider", "inputs": {}},
+                "9": {
+                    "class_type": "ImageConsumer",
+                    "inputs": {
+                        "images": ["11", 0],
+                        "mask": ["17", 0],
+                    },
+                },
+            }
+
+            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
+        finally:
+            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
+
+        self.assertNotIn("17", result)
+        self.assertNotIn("mask", result["9"]["inputs"])
+
+    def test_preserves_scalar_list_join_subgraph_for_downstream_required_input(self):
+        """Delegate-only master keeps list-shaped scalar config chains."""
+        create_list = type(
+            "CreateList",
+            (),
+            {
+                "RETURN_TYPES": ("LIST",),
+                "INPUT_TYPES": classmethod(
+                    lambda cls: {
+                        "required": {"inputs.input0": ("STRING",)},
+                        "optional": {
+                            "inputs.input1": ("STRING",),
+                            "inputs.input2": ("STRING",),
+                        },
+                    }
+                ),
+            },
+        )
+        string_data_list_join = type(
+            "StringDataListJoin",
+            (),
+            {
+                "RETURN_TYPES": ("STRING",),
+                "INPUT_TYPES": classmethod(
+                    lambda cls: {
+                        "required": {
+                            "strings": ("LIST",),
+                            "delimiter": ("STRING",),
+                        }
+                    }
+                ),
+            },
+        )
+        save_image = type(
+            "SaveImage",
+            (),
+            {
+                "INPUT_TYPES": classmethod(
+                    lambda cls: {
+                        "required": {
+                            "images": ("IMAGE",),
+                            "filename_prefix": ("STRING",),
+                        }
+                    }
+                )
+            },
+        )
+        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
+        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
+            "CreateList": create_list,
+            "SaveImage": save_image,
+            "StringDataListJoin": string_data_list_join,
+        }
+        try:
+            prompt = {
+                "8": {"class_type": "KSampler", "inputs": {}},
+                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
+                "15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
+                "17": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug"}},
+                "29": {"class_type": "PrimitiveString", "inputs": {"value": "new"}},
+                "28": {
+                    "class_type": "CreateList",
+                    "inputs": {
+                        "inputs.input0": ["15", 0],
+                        "inputs.input1": ["17", 0],
+                        "inputs.input2": ["29", 0],
+                    },
+                },
+                "32": {
+                    "class_type": "StringDataListJoin",
+                    "inputs": {
+                        "strings": ["28", 0],
+                        "delimiter": "/",
+                    },
+                },
+                "9": {
+                    "class_type": "SaveImage",
+                    "inputs": {
+                        "images": ["11", 0],
+                        "filename_prefix": ["32", 0],
+                    },
+                },
+            }
+
+            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
+        finally:
+            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
+
+        self.assertIn("15", result)
+        self.assertIn("17", result)
+        self.assertIn("28", result)
+        self.assertIn("29", result)
+        self.assertIn("32", result)
+        self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["32", 0])
+        self.assertEqual(result["32"]["inputs"]["strings"], ["28", 0])
+        self.assertEqual(result["28"]["inputs"]["inputs.input0"], ["15", 0])
+        self.assertNotIn("8", result)
+
+    def test_preserves_builtin_create_list_string_data_list_join_subgraph(self):
+        """Delegate-only master handles ComfyUI 0.23 CreateList data-list output."""
+        string_data_list_join = type(
+            "StringDataListJoin",
+            (),
+            {
+                "RETURN_TYPES": ("STRING",),
+                "INPUT_IS_LIST": True,
+                "INPUT_TYPES": classmethod(
+                    lambda cls: {
+                        "required": {
+                            "strings": ("STRING", {"forceInput": True}),
+                            "sep": ("STRING", {"default": " "}),
+                        }
+                    }
+                ),
+            },
+        )
+        save_image = type(
+            "SaveImage",
+            (),
+            {
+                "INPUT_TYPES": classmethod(
+                    lambda cls: {
+                        "required": {
+                            "images": ("IMAGE",),
+                            "filename_prefix": ("STRING",),
+                        }
+                    }
+                )
+            },
+        )
+        previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
+        pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
+            "Basic data handling: StringDataListJoin": string_data_list_join,
+            "SaveImage": save_image,
+        }
+        try:
+            prompt = {
+                "8": {"class_type": "KSampler", "inputs": {}},
+                "11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
+                "15": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug"}},
+                "17": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
+                "29": {"class_type": "PrimitiveString", "inputs": {"value": "new"}},
+                "28": {
+                    "class_type": "CreateList",
+                    "inputs": {
+                        "inputs.input0": ["17", 0],
+                        "inputs.input1": ["15", 0],
+                        "inputs.input2": ["29", 0],
+                    },
+                },
+                "32": {
+                    "class_type": "Basic data handling: StringDataListJoin",
+                    "inputs": {
+                        "strings": ["28", 0],
+                        "sep": "/",
+                    },
+                },
+                "9": {
+                    "class_type": "SaveImage",
+                    "inputs": {
+                        "images": ["11", 0],
+                        "filename_prefix": ["32", 0],
+                    },
+                },
+            }
+
+            result = pt.prepare_delegate_master_prompt(prompt, ["11"])
+        finally:
+            pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
+
+        for node_id in ("15", "17", "28", "29", "32"):
+            self.assertIn(node_id, result)
+        self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["32", 0])
+        self.assertEqual(result["32"]["inputs"]["strings"], ["28", 0])
+        self.assertEqual(result["28"]["inputs"]["inputs.input0"], ["17", 0])
+        self.assertNotIn("8", result)
+
     def test_result_is_independent_copy(self):
         prompt = _delegate_prompt()
         result = pt.prepare_delegate_master_prompt(prompt, ["3"])