From d3f9460967c269eb44ceb286217230502c86c785 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Thu, 11 Jun 2026 12:07:09 +0000 Subject: [PATCH] feat: integrate distributed LTX tiled sampler --- api/__init__.py | 1 + api/ltx_routes.py | 76 ++++ nodes/__init__.py | 6 + nodes/ltx_tiled_sampler.py | 607 ++++++++++++++++++++++++++++++++ pyproject.toml | 2 +- tests/test_ltx_tiled_sampler.py | 129 +++++++ 6 files changed, 820 insertions(+), 1 deletion(-) create mode 100644 api/ltx_routes.py create mode 100644 nodes/ltx_tiled_sampler.py create mode 100644 tests/test_ltx_tiled_sampler.py diff --git a/api/__init__.py b/api/__init__.py index d20f665..e2dd3f5 100644 --- a/api/__init__.py +++ b/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 diff --git a/api/ltx_routes.py b/api/ltx_routes.py new file mode 100644 index 0000000..796af38 --- /dev/null +++ b/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) diff --git a/nodes/__init__.py b/nodes/__init__.py index d7327b3..5e37139 100644 --- a/nodes/__init__.py +++ b/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, } diff --git a/nodes/ltx_tiled_sampler.py b/nodes/ltx_tiled_sampler.py new file mode 100644 index 0000000..16d1e74 --- /dev/null +++ b/nodes/ltx_tiled_sampler.py @@ -0,0 +1,607 @@ +"""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 + + +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) + + 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)"} diff --git a/pyproject.toml b/pyproject.toml index 8994bc8..d743b49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ name = "ComfyUI-Distributed" description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud" version = "1.4.6" license = {file = "LICENSE"} -dependencies = [] +dependencies = ["safetensors"] [project.urls] Repository = "https://github.com/robertvoy/ComfyUI-Distributed" diff --git a/tests/test_ltx_tiled_sampler.py b/tests/test_ltx_tiled_sampler.py new file mode 100644 index 0000000..f29763c --- /dev/null +++ b/tests/test_ltx_tiled_sampler.py @@ -0,0 +1,129 @@ +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) + + +if __name__ == "__main__": + unittest.main()