diff --git a/README.md b/README.md index 919e99f..a1bade1 100644 --- a/README.md +++ b/README.md @@ -176,6 +176,17 @@ Download the following models: ## Advanced Techniques +### Tiled Fusion Sampler + +`LTXVTiledFusionSampler` (🅛🅣🅧 Tiled Fusion Sampler) denoises a canvas larger +than the model's trained window by running each denoise step on overlapping +spatial tiles that share one latent canvas and one noise field. It is meant for +IC-LoRA video-to-video where `LTXAddVideoICLoRAGuide` has already attached the +guide. Pair it with `LTXVGetTilingSizes` (🅛🅣🅧 Get Tiling Sizes) for legal +canvas / tile / `tile_frames` values. See +[tiled_fusion_sampler.md](./tiled_fusion_sampler.md). +Do not confuse it with `LTXVTiledSampler`, which samples tiles independently. + ### Low VRAM * For systems with low VRAM you can use the model loader nodes from [low_vram_loaders.py](./low_vram_loaders.py). Those nodes ensure the correct order of execution and perform the model offloading such that generation fits in 32 GB VRAM. * Use --reserve-vram ComfyUI parameter: `python -m main --reserve-vram 5` (or other number in GB). diff --git a/__init__.py b/__init__.py index 236493e..928ef28 100644 --- a/__init__.py +++ b/__init__.py @@ -60,6 +60,8 @@ from .stg import ( STGGuiderAdvancedNode, STGGuiderNode, ) +from .tiled_fusion_sampler import LTXVTiledFusionSampler +from .tiled_fusion_sizes import LTXVGetTilingSizes from .tiled_sampler import LTXVTiledSampler from .tiled_vae_decode import LTXVTiledVAEDecode from .tricks import NODE_CLASS_MAPPINGS as TRICKS_NODE_CLASS_MAPPINGS @@ -91,6 +93,8 @@ NODE_CLASS_MAPPINGS = { "LTXVSelectLatents": LTXVSelectLatents, "LTXVSetVideoLatentNoiseMasks": LTXVSetVideoLatentNoiseMasks, "LTXVTiledSampler": LTXVTiledSampler, + "LTXVTiledFusionSampler": LTXVTiledFusionSampler, + "LTXVGetTilingSizes": LTXVGetTilingSizes, "LTXVLoopingSampler": LTXVLoopingSampler, "LTXVTiledVAEDecode": LTXVTiledVAEDecode, "MultimodalGuider": MultimodalGuiderNode, diff --git a/iclora.py b/iclora.py index ca54698..cfc837f 100644 --- a/iclora.py +++ b/iclora.py @@ -12,6 +12,7 @@ from comfy_api.latest import io from .latents import LTXVDilateLatent from .nodes_registry import NODES_DISPLAY_NAME_PREFIX, comfy_node +from .tiled_fusion_plan import snap_frames FRAME_IDX_TOOLTIP = ( "RoPE position of the reference, in pixel frames.\n" @@ -37,7 +38,9 @@ class LTXAddVideoICLoRAGuide(io.ComfyNode): "Adds one or more conditioning frames starting at the specified frame index. " "Supports both single images and multi-frame videos. " "The latent_downscale_factor resizes input to a fraction of the target size " - "(1 = original, 2 = half, 3 = third, etc.) for IC-LoRA on small grids." + "(1 = original, 2 = half, 3 = third, etc.) for IC-LoRA on small grids. " + "For LTXVTiledFusionSampler temporal windows, turn on use_streaming and " + "set tile_frames to the same 8n+1 value as the sampler (97 for LTX)." ), inputs=[ io.Conditioning.Input("positive"), @@ -97,6 +100,31 @@ class LTXAddVideoICLoRAGuide(io.ComfyNode): step=16, tooltip="Overlap between tiles for tiled encoding. Only used when use_tiled_encode is enabled.", ), + io.Boolean.Input( + "use_streaming", + optional=True, + default=False, + tooltip=( + "Off (default): encode the whole clip as one IC-LoRA guide. " + "On: re-encode one fresh guide per LTXVTiledFusionSampler " + "temporal window and write the plan the sampler reads. " + "Set tile_frames on this node and the sampler to the same " + "value (97 for LTX)." + ), + ), + io.Int.Input( + "tile_frames", + optional=True, + default=97, + min=0, + max=1000, + step=1, + tooltip=( + "Window length in pixel frames, used only when use_streaming " + "is on. Snapped to 8n+1. Must match LTXVTiledFusionSampler. " + "Ignored when use_streaming is off." + ), + ), ], outputs=[ io.Conditioning.Output("positive"), @@ -177,6 +205,186 @@ class LTXAddVideoICLoRAGuide(io.ComfyNode): guide_latent = vae.encode(encode_pixels) return encode_pixels, guide_latent + @classmethod + def encode_windowed_guides( + cls, + positive, + negative, + vae, + latent_image, + noise_mask, + image, + latent_length, + latent_width, + latent_height, + scale_factors, + strength, + crop, + tile_frames, + attention_strength=1.0, + attention_mask=None, + frame_idx=0, + use_tiled_encode=False, + ): + """Re-encode each fusion window as its own clip, or return None. + + Mid-clip causal latents cannot be sliced into windows, so each window's + pixels are encoded as a clip starting at 0. The shared plan is written + onto the conditioning so LTXVTiledFusionSampler pairs window w with + guide block w. Returns None when the clip fits one window. + """ + from .iclora_attention import append_guide_attention_entry + from .tiled_fusion_plan import attach_temporal_plan, temporal_plan + + t_scale = scale_factors[0] + tf_lat = max(1, (tile_frames - 1) // t_scale + 1) + if not (0 < tf_lat < latent_length): + return None + if int(frame_idx) != 0: + raise ValueError( + "use_streaming encodes each window as a fresh clip at frame 0; " + f"frame_idx={frame_idx} is not supported. Use the whole-clip path " + "(use_streaming off) for stills or mid-clip placement." + ) + if use_tiled_encode: + raise ValueError( + "use_streaming cannot use tiled VAE encode; tiled guide encodes " + "imprint a grid into the conditioning. Leave use_tiled_encode off." + ) + keep = ((image.shape[0] - 1) // t_scale) * t_scale + 1 + image = image[:keep] + ts = temporal_plan(latent_length, tf_lat) + span_px = (tf_lat - 1) * t_scale + 1 + need = ts[-1] * t_scale + span_px + if image.shape[0] < need: + raise ValueError( + f"windowed guide needs {need} frames for {latent_length} " + f"content latents, found {image.shape[0]}" + ) + for w, t0 in enumerate(ts): + span = image[t0 * t_scale : t0 * t_scale + span_px] + _, guide_latent = cls.encode( + vae, + latent_width, + latent_height, + span, + scale_factors, + 1.0, + crop, + False, + 256, + 64, + ) + shape = list(guide_latent.shape[2:]) + toks = shape[0] * shape[1] * shape[2] + positive, negative, latent_image, noise_mask = ( + nodes_lt.LTXVAddGuide.append_keyframe( + positive, + negative, + 0, + latent_image, + noise_mask, + guide_latent, + strength, + scale_factors, + guide_mask=None, + latent_downscale_factor=1.0, + causal_fix=True, + ) + ) + # Pixel mask is (1, 1, F, H, W) after normalize_mask. F==1 broadcasts; + # otherwise it must match this window's encoded span, not the full clip. + window_mask = attention_mask + if window_mask is not None and window_mask.shape[2] != 1: + start = t0 * t_scale + window_mask = window_mask[:, :, start : start + span_px] + positive = append_guide_attention_entry( + positive, + toks, + shape, + attention_strength=attention_strength, + attention_mask=window_mask, + ) + negative = append_guide_attention_entry( + negative, + toks, + shape, + attention_strength=attention_strength, + attention_mask=window_mask, + ) + logging.info( + "window %d/%d: guide frames %d-%d -> %s", + w + 1, + len(ts), + t0 * t_scale, + t0 * t_scale + span_px - 1, + tuple(shape), + ) + plan = { + "tile_frames": int(tile_frames), + "tf_lat": int(tf_lat), + "n_content": int(latent_length), + "ts": [int(v) for v in ts], + } + positive = attach_temporal_plan(positive, plan) + negative = attach_temporal_plan(negative, plan) + logging.info("%d per-window guides (t0 = %s)", len(ts), plan["ts"]) + return io.NodeOutput( + positive, negative, {"samples": latent_image, "noise_mask": noise_mask} + ) + + @classmethod + def maybe_encode_windowed( + cls, + positive, + negative, + vae, + latent_image, + noise_mask, + image, + latent_length, + latent_width, + latent_height, + scale_factors, + strength, + crop, + use_streaming, + tile_frames, + latent_downscale_factor=1.0, + attention_strength=1.0, + attention_mask=None, + frame_idx=0, + use_tiled_encode=False, + ): + """Run the windowed encode only when use_streaming is on. Else None.""" + if not use_streaming: + return None + if latent_downscale_factor != 1: + raise ValueError( + "use_streaming requires latent_downscale_factor=1; " + "windowed fusion guides cannot be dilated." + ) + tf = snap_frames(int(tile_frames) if tile_frames and tile_frames > 0 else 97) + return cls.encode_windowed_guides( + positive, + negative, + vae, + latent_image, + noise_mask, + image, + latent_length, + latent_width, + latent_height, + scale_factors, + strength, + crop, + tf, + attention_strength=attention_strength, + attention_mask=attention_mask, + frame_idx=frame_idx, + use_tiled_encode=use_tiled_encode, + ) + @classmethod def execute( cls, @@ -192,6 +400,8 @@ class LTXAddVideoICLoRAGuide(io.ComfyNode): use_tiled_encode, tile_size, tile_overlap, + use_streaming=False, + tile_frames=0, ) -> io.NodeOutput: scale_factors = vae.downscale_index_formula latent_image = latent["samples"] @@ -199,6 +409,28 @@ class LTXAddVideoICLoRAGuide(io.ComfyNode): _, _, latent_length, latent_height, latent_width = latent_image.shape + windowed = cls.maybe_encode_windowed( + positive, + negative, + vae, + latent_image, + noise_mask, + image, + latent_length, + latent_width, + latent_height, + scale_factors, + strength, + crop, + use_streaming, + tile_frames, + latent_downscale_factor=latent_downscale_factor, + frame_idx=frame_idx, + use_tiled_encode=use_tiled_encode, + ) + if windowed is not None: + return windowed + time_scale_factor = scale_factors[0] num_frames_to_keep = ( (image.shape[0] - 1) // time_scale_factor @@ -306,7 +538,8 @@ class LTXAddVideoICLoRAGuideAdvanced(LTXAddVideoICLoRAGuide): description=( "Adds IC-LoRA guide conditioning with per-guide attention strength control. " "Same as LTXAddVideoICLoRAGuide, but allows controlling how strongly this " - "guide influences generation via self-attention, optionally with a spatial mask." + "guide influences generation via self-attention, optionally with a spatial mask. " + "use_streaming + matching tile_frames re-encodes one guide per Tiled Fusion window." ), inputs=[ io.Conditioning.Input("positive"), @@ -362,6 +595,27 @@ class LTXAddVideoICLoRAGuideAdvanced(LTXAddVideoICLoRAGuide): "Multiplied by attention_strength." ), ), + io.Boolean.Input( + "use_streaming", + optional=True, + default=False, + tooltip=( + "Off (default): whole-clip guide. On: one fresh encode per " + "fusion temporal window. Must match the sampler's tile_frames." + ), + ), + io.Int.Input( + "tile_frames", + optional=True, + default=97, + min=0, + max=1000, + step=1, + tooltip=( + "Used only when use_streaming is on. Snapped to 8n+1. " + "97 for LTX. Must match the sampler." + ), + ), ], outputs=[ io.Conditioning.Output("positive"), @@ -387,6 +641,8 @@ class LTXAddVideoICLoRAGuideAdvanced(LTXAddVideoICLoRAGuide): tile_overlap, attention_strength=1.0, attention_mask=None, + use_streaming=False, + tile_frames=0, ) -> io.NodeOutput: from .iclora_attention import normalize_mask from .latents import LTXVDilateLatent @@ -397,6 +653,30 @@ class LTXAddVideoICLoRAGuideAdvanced(LTXAddVideoICLoRAGuide): _, _, latent_length, latent_height, latent_width = latent_image.shape + windowed = cls.maybe_encode_windowed( + positive, + negative, + vae, + latent_image, + noise_mask, + image, + latent_length, + latent_width, + latent_height, + scale_factors, + strength, + crop, + use_streaming, + tile_frames, + latent_downscale_factor=latent_downscale_factor, + attention_strength=attention_strength, + attention_mask=normalize_mask(attention_mask), + frame_idx=frame_idx, + use_tiled_encode=use_tiled_encode, + ) + if windowed is not None: + return windowed + time_scale_factor = scale_factors[0] num_frames_to_keep = ( (image.shape[0] - 1) // time_scale_factor diff --git a/tiled_fusion_plan.py b/tiled_fusion_plan.py new file mode 100644 index 0000000..1bc7c05 --- /dev/null +++ b/tiled_fusion_plan.py @@ -0,0 +1,80 @@ +"""Shared temporal-window plan for tiled fusion. + +The guide node encodes one fresh IC-LoRA clip per window and writes the plan +onto conditioning. The sampler reads it so window w always pairs with guide +block w. Both sides must use the same origins; slicing a whole-clip causal +guide cannot be repaired by rewriting coordinates. +""" + +from __future__ import annotations + +import os + +# Minimum temporal overlap fraction. Keep in sync with the sampler env knobs. +_T_OV = float(os.environ.get("LTXV_TILED_FUSION_T_OVERLAP", "0.5")) + +TEMPORAL_PLAN_KEY = "ltxv_tiled_fusion_temporal_plan" +GEAR_TEMPORAL_PLAN_KEY = "gear_temporal_plan" +FRAME_STEP = 8 + + +def snap_frames(n: int) -> int: + """Nearest 8n+1 length. Ties keep the shorter clip (no pad).""" + n = max(1, int(n)) + below = ((n - 1) // FRAME_STEP) * FRAME_STEP + 1 + above = below + FRAME_STEP + if above - n < n - below: + return above + return below + + +def plan_grid(canvas_lat, tile_lat, min_overlap): + """Latent-aligned tile origins covering [0, canvas_lat) with >= min_overlap.""" + if canvas_lat <= tile_lat: + return [0] + tile_lat = max(1, int(tile_lat)) + canvas_lat = int(canvas_lat) + # Overlap cannot reach tile_lat, so an unclamped request would never exit. + min_overlap = min(max(0, min_overlap), tile_lat - 1) + n = 2 + while n < canvas_lat and (n * tile_lat - canvas_lat) / (n - 1) < min_overlap: + n += 1 + step = (canvas_lat - tile_lat) / (n - 1) + return [round(i * step) for i in range(n)] + + +def shifted(base, canvas_lat, tile_lat, off): + """Shift interior origins by `off`; endpoints stay pinned so coverage holds.""" + out = [] + for i, v in enumerate(base): + if i in (0, len(base) - 1): + out.append(v) + else: + out.append(max(1, min(canvas_lat - tile_lat - 1, v + off))) + return sorted(set(out)) + + +def temporal_plan(n_lat, tf_lat): + """Shared window plan so a guide node and the sampler always agree.""" + return plan_grid(n_lat, tf_lat, max(0, int(tf_lat * _T_OV))) + + +def attach_temporal_plan(cond, plan): + """Write the plan onto every conditioning entry.""" + out = [] + for emb, d in cond: + d = dict(d) + d[TEMPORAL_PLAN_KEY] = plan + out.append([emb, d]) + return out + + +def read_temporal_plan(positive): + """Return the plan from conditioning, including the Gear-era key.""" + for _, d in positive: + if not isinstance(d, dict): + continue + plan = d.get(TEMPORAL_PLAN_KEY) or d.get(GEAR_TEMPORAL_PLAN_KEY) + if plan: + return plan + return None diff --git a/tiled_fusion_sampler.md b/tiled_fusion_sampler.md new file mode 100644 index 0000000..d7c3be5 --- /dev/null +++ b/tiled_fusion_sampler.md @@ -0,0 +1,70 @@ +# LTXVTiledFusionSampler + +Per-step tiled latent fusion for IC-LoRA (and other LTX conditioning) at +resolutions larger than the model's trained window. + +This is **not** `LTXVTiledSampler`. That node independently samples a grid of +tiles. This node keeps **one** latent canvas and **one** noise field, runs the +model on overlapping crops at **every** denoise step, and Gaussian-blends the +stepped tiles back onto the canvas. Tiles share a trajectory, so seams do not +form. Peak VRAM tracks **spatial tile size** and the **temporal window** +actually stepped (one content window plus one guide block when streaming), +not the full canvas. Cached tile latents/noise/masks follow `canvas_device`; +processed conds stay on the GPU. Streaming also does one full-res VAE encode +per temporal window. + +## Graph shape + +``` +UNET + IC-LoRA +Gemma CLIP → CLIPTextEncode (pos / neg) +guide frames → LTX Add Video IC-LoRA Guide (positive, negative, latents) +KSamplerSelect (euler / heun / …) → SAMPLER +sigmas (Manual Sigmas or any scheduler, descending, ending at 0) + │ + ▼ +LTXV Tiled Fusion Sampler → LATENT → VAE Decode (tiled) +``` + +The latent **must** come from the IC-LoRA guide node. It carries appended guide +frames and a noise mask. This node crops those extra frames itself, so +`LTXVCropGuides` is not needed. + +Leave `use_streaming` **off** on the guide (the default) for a whole-clip encode. +For clips longer than the trained window, turn `use_streaming` **on** and set the +same `tile_frames` (typically 97) on both the guide and this sampler. That +re-encodes each temporal window as its own clip and writes +`ltxv_tiled_fusion_temporal_plan` onto the conditioning. The sampler disables +windowing if the plan is missing. Do not slice a whole-clip guide: mid-clip +causal latents are not fresh-clip-start latents. + +## Sampler input + +The `SAMPLER` socket is the same object `SamplerCustomAdvanced` takes. Fusion +owns the sigma loop and calls `sampler.sampler_function` once per tile with a +two-sigma slice. Use discrete step methods (`euler`, `heun`, `dpm_2`). History +methods (`dpmpp_2m`, `lms`) lose their multi-step memory. Adaptive samplers +cannot fuse. Ancestral methods add independent noise per tile and can seam. + +## Settings that matter + +| Input | Notes | +|---|---| +| `tile_width` / `tile_height` | The IC-LoRA's trained spatial window in pixels (multiples of 32). Smaller tiles see content at the wrong scale. | +| `overlap_frac` | Stay at **0.5** or above. Below that a periodic grid appears on structured content. | +| `blend_var` | Gaussian overlap variance. Default 0.05. | +| `grid_cycle` | 1 = fixed grid (cheapest), 4 = shift the grid across steps. Spatial only. | +| `tile_frames` | Temporal window in pixel frames (97 for LTX). 0 = one extent. >0 needs `use_streaming` on `LTXAddVideoICLoRAGuide` with the same value. | +| `vae` | Optional. Supplies temporal/spatial downscale (otherwise read from the diffusion model). | +| `use_tiled_encode` on the guide node | Must stay **false**. Spatially tiled guide encodes imprint a grid into the conditioning. | +| `use_streaming` on the guide node | Off by default. On = per-window fresh encode + plan. | + +`overlap_frac`, `blend_var`, and `grid_cycle` are spatial fusion knobs. They do +not replace an untiled composition pass when the LoRA is generating a new look +at a canvas far above the trained window. + +Take `SAMPLER` from `KSamplerSelect` (`euler`) and sigmas from +`Manual Sigmas` (descending, ending at 0). Keep `use_tiled_encode` **false** on +every `LTXAddVideoICLoRAGuide`. Turn `use_streaming` **on** and set the same +`tile_frames` on the guide and the sampler when the clip is longer than one +window. diff --git a/tiled_fusion_sampler.py b/tiled_fusion_sampler.py new file mode 100644 index 0000000..4176d37 --- /dev/null +++ b/tiled_fusion_sampler.py @@ -0,0 +1,832 @@ +"""Per-step tiled latent fusion for LTX video models. + +Keeps one latent canvas and one noise field, runs the model once per overlapping +spatial tile at every denoise step (each tile conditioned on its own crop of the +IC-LoRA guide), and merges the stepped tiles back onto the canvas with Gaussian +weights. Tiles follow a single coupled trajectory, so seams do not form. Peak +VRAM is set by the spatial tile and by the temporal window actually stepped +(content window plus one guide block when streaming), not the full canvas +spatial size. Cached tile latents/noise/masks follow canvas_device; processed +conds stay on the GPU for the run. + +This is distinct from LTXVTiledSampler, which independently samples a grid of +tiles. Fusion has to happen between denoise steps, so this node calls +``sampler.sampler_function`` once per tile with a two-sigma slice rather than +``sampler.sample()`` on the whole trajectory. + +The latent input must be the one returned by an IC-LoRA guide node: it carries +guide frames appended along the frame axis plus the noise mask, and the +conditioning carries ``keyframe_idxs`` / ``guide_attention_entries``, all of +which this node crops per tile. +""" + +import copy +import logging +import os + +import comfy.model_management +import comfy.sampler_helpers +import comfy.samplers +import comfy.utils +import torch + +from .nodes_registry import comfy_node +from .tiled_fusion_plan import ( + plan_grid, + read_temporal_plan, + shifted, + snap_frames, + temporal_plan, +) + +logger = logging.getLogger(__name__) + +# Per-step/per-window mean logging for drift forensics (set env before launch). +_DEBUG_STATS = os.environ.get("LTXV_TILED_FUSION_DEBUG", "") not in ("", "0") + +# Temporal-window experiment knobs (env, not node inputs, while under study): +# LTXV_TILED_FUSION_T_OVERLAP minimum temporal overlap fraction (default 0.5) +# LTXV_TILED_FUSION_SKIP_W0_LAT 0 = also merge the first latent of window w>0 +# (that slot is a fresh-clip-start / causal frame +# and poisons a mid-clip canvas position). +# Default on: skip that latent. +_SKIP_FIRST_LAT = os.environ.get("LTXV_TILED_FUSION_SKIP_W0_LAT", "1") not in ( + "0", + "", +) + +_NODE = "LTXV Tiled Fusion Sampler" + + +def _scale_factors(model, vae=None): + """Return (temporal, height, width) pixel-per-latent downscale factors.""" + if vae is not None: + formula = getattr(vae, "downscale_index_formula", None) + if formula is not None: + st, sw, sh = formula + return int(st), int(sh), int(sw) + try: + sf = model.model.latent_format + sh = int( + getattr(sf, "spatial_downscale", None) + or getattr(sf, "spacial_downscale_ratio", None) + or 32 + ) + sw = sh + st = int( + getattr(sf, "temporal_downscale", None) + or getattr(sf, "temporal_downscale_ratio", None) + or 8 + ) + return st, sh, sw + except Exception: + return 8, 32, 32 + + +def gaussian_weight(tile_w, tile_h, var, device): + """Separable Gaussian tile weight. Pure Gaussian, deliberately no floor.""" + + def axis(n): + x = torch.arange(n, dtype=torch.float32, device=device) + mid = (n - 1) / 2 + return torch.exp(-((x - mid) ** 2) / (n * n * 2 * var)) + + return axis(tile_h)[:, None] * axis(tile_w)[None, :] + + +def _crop_guide_metadata( + conds, y0, x0, th, tw, canvas_h, canvas_w, sh, sw, entry_idx=None +): + """Crop keyframe_idxs + guide_attention_entries to a spatial tile. + + y0/x0/th/tw are in latent units; sh/sw are pixels per latent step. Token + order inside each guide's block is the patchifier's row-major (t, h, w), + which is what makes the reshape below valid. Spatial positions are + normalised to the tile so each tile is the same problem the LoRA was + trained on. + + entry_idx selects a SINGLE guide entry: temporal windowing pairs window w + with guide block w, each encoded FRESH from that window's pixel span, so + its coords already describe a clip starting at 0 and are never rewritten + here. None keeps every entry (whole-clip guide). + Slicing a whole-clip guide's latents into windows does NOT work and no + coordinate rewrite can save it: a mid-clip causal latent (each frame + encodes 8 pixels given prior context) is not a fresh-clip latent (frame 0 + encodes exactly 1 pixel frame). + """ + out = [] + py0, px0 = y0 * sh, x0 * sw + for emb, d in conds: + d = dict(d) + entries = d.get("guide_attention_entries") + kf = d.get("keyframe_idxs") + if kf is not None and entries: + pieces = [] + new_entries = [] + tok0 = 0 + for j, e in enumerate(entries): + f, gh, gw = (int(v) for v in e["latent_shape"]) + n_tok = f * gh * gw + blk = kf[:, :, tok0 : tok0 + n_tok, :] + tok0 += n_tok + if entry_idx is not None and j != entry_idx: + continue + if (gh, gw) != (canvas_h, canvas_w): + raise ValueError( + f"{_NODE}: guide latent grid {gh}x{gw} does not match " + f"the canvas {canvas_h}x{canvas_w}. Downscaled " + f"(factor > 1) IC-LoRA guides are not supported; " + f"re-attach the guide at factor 1." + ) + b = blk.shape[0] + blk = blk.reshape(b, 3, f, gh, gw, 2) + blk = blk[:, :, :, y0 : y0 + th, x0 : x0 + tw, :].clone() + blk[:, 1, ...] -= py0 + blk[:, 2, ...] -= px0 + pieces.append(blk.reshape(b, 3, f * th * tw, 2)) + ne = dict(e) + ne["pre_filter_count"] = f * th * tw + ne["latent_shape"] = [f, th, tw] + pm = e.get("pixel_mask") + if pm is not None: + ne["pixel_mask"] = pm[ + :, :, :, py0 : py0 + th * sh, px0 : px0 + tw * sw + ] + new_entries.append(ne) + d["keyframe_idxs"] = torch.cat(pieces, dim=2) + d["guide_attention_entries"] = new_entries + out.append([emb, d]) + return out + + +def _crop_spatial(t, y0, x0, th, tw): + """Crop the last two dims when they are real spatial dims (>1).""" + if t is None: + return None + if t.shape[-2] > 1 or t.shape[-1] > 1: + return t[..., y0 : y0 + th, x0 : x0 + tw] + return t + + +class _TileGuider: + """CFGGuider look-alike so KSamplerX0Inpaint / k-diffusion can call us. + + `inner_model` is the object `prepare_sampling` returned (has + `scale_latent_inpaint`). `__call__` is the transformer+CFG pass. + """ + + def __init__(self, inner_model, conds, cfg): + self.inner_model = inner_model + self.conds = conds + self.cfg = cfg + + def __call__(self, x, timestep, model_options={}, seed=None): + return comfy.samplers.sampling_function( + self.inner_model, + x, + timestep, + self.conds.get("negative"), + self.conds.get("positive"), + self.cfg, + model_options=model_options, + seed=seed, + ) + + +def _one_sampler_step( + sampler, + inner_model, + conds, + cfg, + x, + sig, + sig_next, + noise, + latent_image, + denoise_mask, + model_options, + seed, + full_sigmas, +): + """Run one k-diffusion step (sig -> sig_next) on a tile. + + Uses `sampler.sampler_function` (what KSamplerSelect builds), not + `sampler.sample()`, which would re-scale noise and walk the whole + schedule — fusion has to happen between steps. + """ + fn = getattr(sampler, "sampler_function", None) + if fn is None: + raise ValueError( + f"{_NODE} needs a SAMPLER from KSamplerSelect (or another node " + "that exposes sampler_function). Adaptive samplers that run the " + "whole trajectory in one go cannot fuse tiles between steps." + ) + wrap = _TileGuider(inner_model, conds, cfg) + model_k = comfy.samplers.KSamplerX0Inpaint(wrap, full_sigmas) + model_k.latent_image = latent_image + inpaint_opts = getattr(sampler, "inpaint_options", None) or {} + if inpaint_opts.get("random", False): + generator = torch.manual_seed(seed + 1) + model_k.noise = torch.randn(noise.shape, generator=generator, device="cpu").to( + device=noise.device, dtype=noise.dtype + ) + else: + model_k.noise = noise + extra_args = { + "model_options": model_options, + "seed": seed, + "denoise_mask": denoise_mask, + } + sigmas_pair = torch.tensor([sig, sig_next], device=x.device, dtype=torch.float32) + extra_opts = dict(getattr(sampler, "extra_options", None) or {}) + return fn( + model_k, + x, + sigmas_pair, + extra_args=extra_args, + callback=None, + disable=True, + **extra_opts, + ) + + +@comfy_node( + name="LTXVTiledFusionSampler", + description="Tiled Fusion Sampler", +) +class LTXVTiledFusionSampler: + CATEGORY = "sampling" + RETURN_TYPES = ("LATENT",) + RETURN_NAMES = ("output",) + FUNCTION = "sample" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ( + "MODEL", + {"tooltip": "The diffusion model to use."}, + ), + "positive": ( + "CONDITIONING", + { + "tooltip": "Positive conditioning from the IC-LoRA guide node. " + "Must carry keyframe_idxs / guide_attention_entries so each " + "tile can be cropped to the LoRA's trained window." + }, + ), + "negative": ( + "CONDITIONING", + {"tooltip": "Negative conditioning from the IC-LoRA guide node."}, + ), + "latents": ( + "LATENT", + { + "tooltip": "The latent RETURNED BY the IC-LoRA guide node " + "(it carries appended guide frames plus the noise mask). " + "A bare empty video latent will not work." + }, + ), + "sigmas": ( + "SIGMAS", + { + "tooltip": "Denoise schedule. Must be descending and end at 0. " + "Fusion walks this list one step at a time." + }, + ), + "sampler": ( + "SAMPLER", + { + "tooltip": "From KSamplerSelect (same type as SamplerCustomAdvanced). " + "Discrete step samplers only: euler, euler_ancestral, heun, dpm_2, … " + "History methods (DPM++ 2M, LMS) lose their multi-step memory because " + "fusion owns the loop. Adaptive samplers cannot fuse. Ancestral methods " + "add independent noise per tile and can seam." + }, + ), + "seed": ( + "INT", + { + "default": 42, + "min": 0, + "max": 0xFFFFFFFFFFFFFFFF, + "control_after_generate": True, + "tooltip": "Seed for the single shared noise field across all tiles.", + }, + ), + "cfg": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 20.0, + "step": 0.1, + "tooltip": "Classifier-free guidance scale. Distilled models typically use 1.0.", + }, + ), + "tile_width": ( + "INT", + { + "default": 1024, + "min": 256, + "max": 4096, + "step": 32, + "tooltip": "Tile width in pixels, multiple of 32. Set this to the " + "IC-LoRA's trained spatial window. Tiles smaller than the trained " + "window make every tile see content at the wrong scale.", + }, + ), + "tile_height": ( + "INT", + { + "default": 576, + "min": 256, + "max": 4096, + "step": 32, + "tooltip": "Tile height in pixels, multiple of 32. Match tile_width to " + "the IC-LoRA's trained spatial window.", + }, + ), + "overlap_frac": ( + "FLOAT", + { + "default": 0.5, + "min": 0.25, + "max": 0.75, + "step": 0.05, + "tooltip": "Fraction of each tile overlapping its neighbor. Stay at " + "0.5 or above: below that a periodic grid appears on structured content.", + }, + ), + "blend_var": ( + "FLOAT", + { + "default": 0.05, + "min": 0.01, + "max": 0.2, + "step": 0.01, + "tooltip": "Variance of the Gaussian blend weights in the overlaps. " + "Leave at 0.05 unless you are measuring seams.", + }, + ), + "grid_cycle": ( + "INT", + { + "default": 1, + "min": 1, + "max": 4, + "tooltip": "Cycle shifted tile grids across steps. 1 = fixed grid " + "(cheapest), 4 = thorough. Spatial only: cycling the temporal grid " + "slides content in time.", + }, + ), + "tile_frames": ( + "INT", + { + "default": 0, + "min": 0, + "max": 1000, + "step": 1, + "tooltip": "Temporal window in pixel frames (97 for LTX). 0 = one " + "extent. Other values are snapped to 8n+1. >0 needs " + "LTXAddVideoICLoRAGuide with use_streaming on and the SAME " + "tile_frames: that re-encodes a fresh guide per window and " + "writes the plan this sampler reads. A whole-clip guide " + "cannot be sliced, and windowing is disabled if the plan " + "is missing.", + }, + ), + }, + "optional": { + "vae": ( + "VAE", + { + "tooltip": "Used to read the model's temporal and spatial downscale " + "factors. If omitted, they are taken from the diffusion model " + "(LTX defaults: 8 frames and 32 pixels per latent cell)." + }, + ), + "canvas_device": ( + ["auto", "gpu", "cpu"], + { + "default": "auto", + "tooltip": "Where to keep the full-canvas tensors and cached tile " + "latents/noise/masks. auto = GPU if there is enough free " + "memory, otherwise CPU. Processed conds stay on the GPU.", + }, + ), + }, + } + + @torch.inference_mode() + def sample( + self, + model, + positive, + negative, + latents, + sigmas, + sampler, + seed, + cfg, + tile_width, + tile_height, + overlap_frac, + blend_var, + grid_cycle, + tile_frames=0, + vae=None, + canvas_device="auto", + ): + samples = latents["samples"] + noise_mask = latents.get("noise_mask") + if samples.ndim != 5: + raise ValueError(f"{_NODE} needs a 5D video latent.") + b, C, F_total, H, W = samples.shape + if b != 1: + raise ValueError(f"{_NODE} supports batch size 1.") + if noise_mask is None: + raise ValueError( + f"{_NODE} requires a noise_mask; the latent must be the one " + "returned by the IC-LoRA guide node (it carries noise_mask)." + ) + + device = comfy.model_management.get_torch_device() + st, sh, sw = _scale_factors(model, vae) + + tl_w, tl_h = tile_width // sw, tile_height // sh + tl_w, tl_h = min(tl_w, W), min(tl_h, H) + + min_ov_x = max(2, int(tl_w * overlap_frac)) + min_ov_y = max(2, int(tl_h * overlap_frac)) + xs = plan_grid(W, tl_w, min_ov_x) + ys = plan_grid(H, tl_h, min_ov_y) + + ov_x = tl_w - (xs[1] - xs[0]) if len(xs) > 1 else tl_w + ov_y = tl_h - (ys[1] - ys[0]) if len(ys) > 1 else tl_h + # Interior shift must leave at least min_ov overlap on cycled grids. + max_shift = max(0, min(ov_x - min_ov_x, ov_y - min_ov_y)) + off = min(4, max_shift) + if grid_cycle > 1 and off < 1: + logger.warning( + "grid_cycle=%d ignored: no interior shift can keep overlap_frac=%.2f", + grid_cycle, + overlap_frac, + ) + offsets = [(0, 0)] + else: + offsets = [(0, 0), (off, off), (off, 0), (0, off)][: max(1, grid_cycle)] + grids = [ + [(x, y) for y in shifted(ys, H, tl_h, oy) for x in shifted(xs, W, tl_w, ox)] + for ox, oy in offsets + ] + all_origins = sorted({o for g in grids for o in g}) + logger.info( + "tiled fusion: canvas %dx%d lat, tile %dx%d, grid %dx%d, " + "cycle=%d, %d unique origins", + W, + H, + tl_w, + tl_h, + len(xs), + len(ys), + len(grids), + len(all_origins), + ) + + if len(all_origins) == 1: + logger.info("canvas fits one tile; falling back to plain sampling") + + n_content = F_total + if noise_mask is not None: + fm = ( + noise_mask.to(dtype=torch.float32) + .reshape(1, 1, F_total, -1) + .amax(dim=-1) + ) + for f in range(F_total - 1, -1, -1): + if float(fm[0, 0, f]) < 0.999: + n_content = f + else: + break + n_guide = F_total - n_content + + plan = read_temporal_plan(positive) + + tile_frames = int(tile_frames) + if tile_frames > 0: + snapped = snap_frames(tile_frames) + if snapped != tile_frames: + logger.info("tile_frames %d snapped to %d (8n+1)", tile_frames, snapped) + tile_frames = snapped + + tf_lat = 0 + if tile_frames > 0: + tf_lat = max(1, (tile_frames - 1) // st + 1) + use_temporal = 0 < tf_lat < n_content + + if plan is not None and int(plan["tile_frames"]) != int(tile_frames): + raise ValueError( + f"{_NODE}: the guide was encoded for tile_frames=" + f"{plan['tile_frames']} but the sampler got tile_frames=" + f"{tile_frames}; set both to the same value." + ) + + per_window_guides = False + if use_temporal: + if plan is not None: + ts = [int(v) for v in plan["ts"]] + if int(plan["n_content"]) != n_content or n_guide != len(ts) * tf_lat: + raise ValueError( + f"{_NODE}: temporal plan does not match the latent " + f"(plan: n_content={plan['n_content']}, " + f"{len(ts)} windows of {tf_lat}; latent: " + f"content={n_content}, guide={n_guide})." + ) + per_window_guides = True + elif n_guide == 0: + ts = temporal_plan(n_content, tf_lat) + else: + logger.warning( + "tile_frames=%d with a whole-clip guide: temporal windows " + "need a fresh IC-LoRA encode per window. Turn on " + "use_streaming on LTXAddVideoICLoRAGuide with the same " + "tile_frames; windowing DISABLED for this run.", + tile_frames, + ) + use_temporal = False + + if use_temporal: + logger.info( + "temporal windows: %d x %d latent frames (%d content, " "%d guide%s)", + len(ts), + tf_lat, + n_content, + n_guide, + ", per-window guides" if per_window_guides else "", + ) + window_len = min(tf_lat, n_content) + if per_window_guides: + window_len += tf_lat + else: + ts = [0] + tf_lat = F_total + window_len = F_total + if tile_frames > 0: + logger.info( + "temporal windowing not active (clip fits one window, " + "or disabled above)" + ) + + if canvas_device == "cpu": + cdev = torch.device("cpu") + elif canvas_device == "gpu": + cdev = device + else: + free = comfy.model_management.get_free_memory(device) + need = samples.numel() * 4 * 4 + cdev = device if need < free * 0.25 else torch.device("cpu") + logger.info("canvas tensors on %s", cdev) + + g = torch.Generator(device="cpu").manual_seed(seed) + noise = torch.randn(samples.shape, generator=g, dtype=torch.float32).to(cdev) + + tile_shape = [1, C, window_len, tl_h, tl_w] + conds0 = { + "positive": comfy.sampler_helpers.convert_cond(copy.deepcopy(positive)), + "negative": comfy.sampler_helpers.convert_cond(copy.deepcopy(negative)), + } + inner_model, conds0, loaded = comfy.sampler_helpers.prepare_sampling( + model, tile_shape, conds0, model.model_options + ) + try: + model_sampling = ( + getattr(inner_model, "model_sampling", None) + or model.model.model_sampling + ) + + model.pre_run() + + latent_in = ( + model.model.process_latent_in(samples).to(torch.float32).to(cdev) + ) + sigmas = sigmas.to(torch.float32) + s0 = sigmas[0].item() + + canvas = model_sampling.noise_scaling( + torch.tensor(s0, device=cdev), noise, latent_in, False + ) + + if noise_mask is not None: + mask_full = noise_mask.to(torch.float32).to(cdev) + else: + mask_full = None + + def _window_slice(w, t0, tensor): + """Take content[t0:t0+tf] plus window w's own guide block.""" + if not use_temporal: + return tensor + parts = [tensor[:, :, t0 : t0 + tf_lat]] + if per_window_guides: + g0 = n_content + w * tf_lat + parts.append(tensor[:, :, g0 : g0 + tf_lat]) + return torch.cat(parts, dim=2) if len(parts) > 1 else parts[0] + + tile_data = {} + for w, t0_ in enumerate(ts): + eidx = w if (use_temporal and per_window_guides) else None + for x0_, y0_ in all_origins: + pos_t = _crop_guide_metadata( + copy.deepcopy(positive), + y0_, + x0_, + tl_h, + tl_w, + H, + W, + sh, + sw, + entry_idx=eidx, + ) + neg_t = _crop_guide_metadata( + copy.deepcopy(negative), + y0_, + x0_, + tl_h, + tl_w, + H, + W, + sh, + sw, + entry_idx=eidx, + ) + lat_t = _window_slice( + w, + t0_, + latent_in[:, :, :, y0_ : y0_ + tl_h, x0_ : x0_ + tl_w], + ) + noise_t = _window_slice( + w, + t0_, + noise[:, :, :, y0_ : y0_ + tl_h, x0_ : x0_ + tl_w], + ) + mask_t = None + if mask_full is not None: + mask_t = _crop_spatial(mask_full, y0_, x0_, tl_h, tl_w) + if mask_t.shape[2] == F_total: + mask_t = _window_slice(w, t0_, mask_t) + if mask_t.shape[2] == lat_t.shape[2]: + mask_t = mask_t.expand(1, 1, lat_t.shape[2], tl_h, tl_w) + # process_conds needs GPU tensors; lat/noise/mask stay on cdev + # until each step so a CPU canvas does not pin every tile. + conds_t = comfy.samplers.process_conds( + inner_model, + noise_t.to(device), + { + "positive": comfy.sampler_helpers.convert_cond(pos_t), + "negative": comfy.sampler_helpers.convert_cond(neg_t), + }, + device, + latent_image=lat_t.to(device), + denoise_mask=None if mask_t is None else mask_t.to(device), + seed=seed, + latent_shapes=[lat_t.shape], + ) + tile_data[(w, x0_, y0_)] = (conds_t, lat_t, noise_t, mask_t) + + w_tile = gaussian_weight(tl_w, tl_h, blend_var, cdev) + + w_len = tf_lat if use_temporal else n_content + if use_temporal and len(ts) > 1: + xs_t = torch.arange(w_len, dtype=torch.float32, device=cdev) + mid_t = (w_len - 1) / 2 + w_time = torch.exp( + -((xs_t - mid_t) ** 2) / (w_len * w_len * 2 * blend_var) + ) + else: + w_time = torch.ones(w_len, dtype=torch.float32, device=cdev) + + def _w_skip(w): + return int( + _SKIP_FIRST_LAT + and use_temporal + and w > 0 + and ts[w] < ts[w - 1] + tf_lat + ) + + wsums = [] + for grid in grids: + ws = torch.zeros((1, 1, n_content, H, W), device=cdev) + for w, t0_ in enumerate(ts): + sk = _w_skip(w) + tw_ = w_time[sk : min(tf_lat, n_content - t0_)].reshape( + 1, 1, -1, 1, 1 + ) + for x0_, y0_ in grid: + ws[ + :, + :, + t0_ + sk : t0_ + sk + tw_.shape[2], + y0_ : y0_ + tl_h, + x0_ : x0_ + tl_w, + ] += ( + w_tile * tw_ + ) + wsums.append(ws.clamp_min(1e-12)) + + model_options = model.model_options + + pbar = comfy.utils.ProgressBar(len(sigmas) - 1) + for step in range(len(sigmas) - 1): + sig = sigmas[step].item() + sig_next = sigmas[step + 1].item() + grid = grids[step % len(grids)] + acc = torch.zeros( + (1, C, n_content, H, W), device=cdev, dtype=canvas.dtype + ) + + for w, t0_ in enumerate(ts): + for x0_, y0_ in grid: + conds_t, lat_t, noise_t, mask_t = tile_data[(w, x0_, y0_)] + x = _window_slice( + w, + t0_, + canvas[:, :, :, y0_ : y0_ + tl_h, x0_ : x0_ + tl_w], + ).to(device, torch.float32) + + x_next = _one_sampler_step( + sampler, + inner_model, + conds_t, + cfg, + x, + sig, + sig_next, + noise_t.to(device, torch.float32), + lat_t.to(device, torch.float32), + None if mask_t is None else mask_t.to(device), + model_options, + seed, + sigmas, + ) + + nc = min(tf_lat, n_content - t0_) if use_temporal else n_content + sk = _w_skip(w) + tw_ = w_time[sk:nc].reshape(1, 1, -1, 1, 1) + acc[ + :, + :, + t0_ + sk : t0_ + nc, + y0_ : y0_ + tl_h, + x0_ : x0_ + tl_w, + ] += ( + x_next[:, :, sk:nc].to(cdev) * w_tile * tw_ + ) + + if _DEBUG_STATS: + logger.info( + "dbg step %d w%d (t0=%d xy=%d,%d): x %.4f " + "x_next %.4f", + step, + w, + t0_, + x0_, + y0_, + float(x[:, :, :nc].mean()), + float(x_next[:, :, :nc].mean()), + ) + + blended = acc / wsums[step % len(grids)] + if _DEBUG_STATS: + logger.info( + "dbg step %d blended: mean %.4f std %.4f", + step, + float(blended.mean()), + float(blended.std()), + ) + if n_guide: + canvas = torch.cat([blended, canvas[:, :, n_content:]], dim=2) + else: + canvas = blended + pbar.update(1) + + out = canvas + if mask_full is not None: + frame_mask = mask_full.reshape(1, 1, F_total, -1).amax(dim=-1) + keep = F_total + for f in range(F_total - 1, -1, -1): + if frame_mask[0, 0, f] < 0.999: + keep = f + else: + break + out = canvas[:, :, :keep] + if keep < F_total: + logger.info("dropped %d appended guide frame(s)", F_total - keep) + + out = model.model.process_latent_out(out.to(torch.float32).cpu()) + finally: + model.cleanup() + comfy.sampler_helpers.cleanup_models(conds0, loaded) + + return ({"samples": out},) diff --git a/tiled_fusion_sizes.py b/tiled_fusion_sizes.py new file mode 100644 index 0000000..6cc608c --- /dev/null +++ b/tiled_fusion_sizes.py @@ -0,0 +1,445 @@ +"""Canvas sizes and temporal window for tiled-fusion graphs. + +Stage 1 aims at full HD in the source aspect without spatial tiling. Stage 2 +is 2x that (~4K). 8K is another 2x. HD is the 720p step of the same ladder; +qHD is the 960 long-edge step. + +Dimensions stay multiples of 32. ``tile_frames`` stays 8n+1, targeting 97, +stretching up to 121 only when that lets a slightly longer clip skip chunking. +""" + +from __future__ import annotations + +import math +from typing import NamedTuple, Optional + +from .nodes_registry import comfy_node +from .tiled_fusion_plan import plan_grid, snap_frames, temporal_plan + +SPATIAL_ALIGN = 32 +FRAME_STEP = 8 +SLACK_PX = 200 +QHD_LONG_EDGE = 960 +HD_LONG_EDGE = 1280 +FULLHD_LONG_EDGE = 1920 +TILE_FRAMES_TARGET = 97 +TILE_FRAMES_MAX = 121 +OVERLAP_FRAC = 0.5 + +TILE_CHOICES = ("qHD", "HD", "FullHD") +OUTPUT_CHOICES = ("FullHD", "4K", "8K") + +_QHD_KEYS = frozenset({"qhd", "quarterhd", "quarter", "960"}) +_HD_KEYS = frozenset({"hd", "720p", "720"}) +_FULLHD_KEYS = frozenset({"fullhd", "fhd", "1080p", "1080"}) +_4K_KEYS = frozenset({"4k", "uhd", "2160p", "2160"}) +_8K_KEYS = frozenset({"8k", "eightk", "4320p", "4320"}) + + +class TilingSizes(NamedTuple): + qhd_width: int + qhd_height: int + hd_width: int + hd_height: int + fullhd_width: int + fullhd_height: int + overlap_frac: float + tile_frames: int + width_4k: int + height_4k: int + width_8k: int + height_8k: int + + +class TilingPlan(NamedTuple): + initial_canvas_width: int + initial_canvas_height: int + tile_width: int + tile_height: int + output_width: int + output_height: int + overlap_frac: float + tile_frames: int + spatial_tiles_x: int + spatial_tiles_y: int + output_tiles_x: int + output_tiles_y: int + temporal_chunks: int + canvas_label: str + tile_label: str + output_label: str + tile_clamped: bool + is_8k: bool + summary: str + + +def _norm_label(value: str) -> str: + return str(value or "").lower().replace(" ", "").replace("_", "").replace("-", "") + + +def snap_spatial(value: float, align: int = SPATIAL_ALIGN) -> int: + """Nearest multiple of ``align``, ties round up, never below ``align``.""" + value = max(float(align), float(value)) + return max(align, int(math.floor(value / align + 0.5) * align)) + + +def tile_frames_for(frame_count: int) -> int: + """Temporal window: 97 or the clip itself, stretching to 121 to skip chunks.""" + n = max(1, int(frame_count)) + if n <= TILE_FRAMES_TARGET: + tf = snap_frames(n) + return min(TILE_FRAMES_TARGET, tf) + if n <= TILE_FRAMES_MAX: + covering = ((n - 1 + FRAME_STEP - 1) // FRAME_STEP) * FRAME_STEP + 1 + return min(TILE_FRAMES_MAX, covering) + return TILE_FRAMES_TARGET + + +def _fit_long_edge(width: float, height: float, long_edge: int) -> tuple[float, float]: + long = max(width, height) + scale = long_edge / long + return width * scale, height * scale + + +def _keep_source_if_close( + target_w: float, target_h: float, src_w: int, src_h: int, slack: int = SLACK_PX +) -> tuple[float, float]: + if abs(src_w - target_w) <= slack and abs(src_h - target_h) <= slack: + return float(src_w), float(src_h) + return target_w, target_h + + +def _aligned_size(width: float, height: float) -> tuple[int, int]: + return snap_spatial(width), snap_spatial(height) + + +def _ladder_step(src_w: int, src_h: int, long_edge: int) -> tuple[int, int]: + tw, th = _fit_long_edge(src_w, src_h, long_edge) + tw, th = _keep_source_if_close(tw, th, src_w, src_h) + return _aligned_size(tw, th) + + +def compute_tiling_sizes(width: int, height: int, frame_count: int) -> TilingSizes: + """Return the qHD / HD / full-HD / 4K / 8K ladder plus overlap and tile_frames.""" + src_w = max(1, int(width)) + src_h = max(1, int(height)) + + qhd_w, qhd_h = _ladder_step(src_w, src_h, QHD_LONG_EDGE) + hd_w, hd_h = _ladder_step(src_w, src_h, HD_LONG_EDGE) + fhd_w, fhd_h = _ladder_step(src_w, src_h, FULLHD_LONG_EDGE) + + k4_w, k4_h = _keep_source_if_close(fhd_w * 2, fhd_h * 2, src_w, src_h) + k4_w, k4_h = _aligned_size(k4_w, k4_h) + + k8_w, k8_h = _keep_source_if_close(k4_w * 2, k4_h * 2, src_w, src_h) + k8_w, k8_h = _aligned_size(k8_w, k8_h) + + return TilingSizes( + qhd_width=qhd_w, + qhd_height=qhd_h, + hd_width=hd_w, + hd_height=hd_h, + fullhd_width=fhd_w, + fullhd_height=fhd_h, + overlap_frac=OVERLAP_FRAC, + tile_frames=tile_frames_for(frame_count), + width_4k=k4_w, + height_4k=k4_h, + width_8k=k8_w, + height_8k=k8_h, + ) + + +def _pick_tile_or_canvas(sizes: TilingSizes, choice: str) -> tuple[int, int, str]: + key = _norm_label(choice) + if key in _QHD_KEYS: + return sizes.qhd_width, sizes.qhd_height, "qHD" + if key in _HD_KEYS: + return sizes.hd_width, sizes.hd_height, "HD" + if key in _FULLHD_KEYS: + return sizes.fullhd_width, sizes.fullhd_height, "FullHD" + raise ValueError( + f"Unknown tile/canvas size {choice!r}; expected qHD, HD, or FullHD." + ) + + +def _pick_output(sizes: TilingSizes, choice: str) -> tuple[int, int, str]: + key = _norm_label(choice) + if key in _8K_KEYS: + return sizes.width_8k, sizes.height_8k, "8K" + if key in _4K_KEYS: + return sizes.width_4k, sizes.height_4k, "4K" + if key in _FULLHD_KEYS: + return sizes.fullhd_width, sizes.fullhd_height, "FullHD" + raise ValueError(f"Unknown output size {choice!r}; expected FullHD, 4K, or 8K.") + + +def _latent_wh(width: int, height: int) -> tuple[int, int]: + return max(1, int(width) // SPATIAL_ALIGN), max(1, int(height) // SPATIAL_ALIGN) + + +def spatial_tile_grid( + canvas_w: int, + canvas_h: int, + tile_w: int, + tile_h: int, + overlap_frac: float = OVERLAP_FRAC, +) -> tuple[int, int]: + """Number of spatial tiles covering the canvas, matching the sampler grid.""" + cw, ch = _latent_wh(canvas_w, canvas_h) + tw, th = _latent_wh(tile_w, tile_h) + tw, th = min(tw, cw), min(th, ch) + min_ov_x = max(2, int(tw * overlap_frac)) + min_ov_y = max(2, int(th * overlap_frac)) + return ( + len(plan_grid(cw, tw, min_ov_x)), + len(plan_grid(ch, th, min_ov_y)), + ) + + +def temporal_chunk_count(frame_count: int, tile_frames: int) -> int: + """Number of temporal windows, matching the sampler / guide handshake.""" + n = max(1, int(frame_count)) + tf = max(1, int(tile_frames)) + n_lat = (n - 1) // FRAME_STEP + 1 + tf_lat = min(n_lat, (tf - 1) // FRAME_STEP + 1) + return len(temporal_plan(n_lat, tf_lat)) + + +def format_summary( + *, + src_w: int, + src_h: int, + frame_count: int, + plan_fields: dict, +) -> str: + sx, sy = plan_fields["spatial_tiles_x"], plan_fields["spatial_tiles_y"] + ox, oy = plan_fields["output_tiles_x"], plan_fields["output_tiles_y"] + chunks = plan_fields["temporal_chunks"] + tile_note = " (clamped to canvas)" if plan_fields["tile_clamped"] else "" + if (sx, sy) == (ox, oy): + spatial = f"spatial {sx}×{sy} tiles" + else: + spatial = f"spatial {sx}×{sy} on canvas, {ox}×{oy} on output" + chunk_word = "chunk" if chunks == 1 else "chunks" + return "\n".join( + ( + f"source {src_w}×{src_h} × {frame_count}f", + f"canvas {plan_fields['initial_canvas_width']}×{plan_fields['initial_canvas_height']}" + f" {plan_fields['canvas_label']}", + f"tile {plan_fields['tile_width']}×{plan_fields['tile_height']}" + f" {plan_fields['tile_label']}{tile_note}", + f"output {plan_fields['output_width']}×{plan_fields['output_height']}" + f" {plan_fields['output_label']}", + f"overlap {plan_fields['overlap_frac']:g}", + f"tile_frames {plan_fields['tile_frames']}", + spatial, + f"temporal {chunks} {chunk_word}", + ) + ) + + +def resolve_tiling_plan( + width: int, + height: int, + frame_count: int, + tile_size: str = "FullHD", + initial_canvas_size: str = "FullHD", + output_size: str = "FullHD", +) -> TilingPlan: + """Pick canvas / tile / output from the ladder and count tiles and chunks.""" + src_w = max(1, int(width)) + src_h = max(1, int(height)) + frames = max(1, int(frame_count)) + sizes = compute_tiling_sizes(src_w, src_h, frames) + + canvas_w, canvas_h, canvas_label = _pick_tile_or_canvas(sizes, initial_canvas_size) + tile_w, tile_h, tile_label = _pick_tile_or_canvas(sizes, tile_size) + out_w, out_h, out_label = _pick_output(sizes, output_size) + + clamped_w, clamped_h = min(tile_w, canvas_w), min(tile_h, canvas_h) + tile_clamped = (clamped_w, clamped_h) != (tile_w, tile_h) + tile_w, tile_h = clamped_w, clamped_h + + sx, sy = spatial_tile_grid(canvas_w, canvas_h, tile_w, tile_h, sizes.overlap_frac) + ox, oy = spatial_tile_grid(out_w, out_h, tile_w, tile_h, sizes.overlap_frac) + chunks = temporal_chunk_count(frames, sizes.tile_frames) + + fields = dict( + initial_canvas_width=canvas_w, + initial_canvas_height=canvas_h, + tile_width=tile_w, + tile_height=tile_h, + output_width=out_w, + output_height=out_h, + overlap_frac=sizes.overlap_frac, + tile_frames=sizes.tile_frames, + spatial_tiles_x=sx, + spatial_tiles_y=sy, + output_tiles_x=ox, + output_tiles_y=oy, + temporal_chunks=chunks, + canvas_label=canvas_label, + tile_label=tile_label, + output_label=out_label, + tile_clamped=tile_clamped, + is_8k=out_label == "8K", + ) + summary = format_summary( + src_w=src_w, src_h=src_h, frame_count=frames, plan_fields=fields + ) + return TilingPlan(summary=summary, **fields) + + +def _dims_from_image(image) -> tuple[int, int, int]: + """Comfy IMAGE is [frames, height, width, channels].""" + frames, height, width = ( + int(image.shape[0]), + int(image.shape[1]), + int(image.shape[2]), + ) + return width, height, frames + + +@comfy_node( + name="LTXVGetTilingSizes", + description="Get Tiling Sizes", +) +class LTXVGetTilingSizes: + """Size calculator in the spirit of Comfy's Get Image Size, for tiled fusion.""" + + CATEGORY = "ltxvideo" + FUNCTION = "get_sizes" + OUTPUT_NODE = True + RETURN_TYPES = ( + "INT", + "INT", + "INT", + "INT", + "INT", + "INT", + "FLOAT", + "INT", + "BOOLEAN", + "STRING", + ) + RETURN_NAMES = ( + "initial_canvas_width", + "initial_canvas_height", + "tile_width", + "tile_height", + "output_width", + "output_height", + "overlap_frac", + "tile_frames", + "is_8k", + "summary", + ) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "width": ( + "INT", + { + "default": 1920, + "min": 1, + "max": 16384, + "tooltip": "Source width in pixels (from Load Video / Get Image Size).", + }, + ), + "height": ( + "INT", + { + "default": 1080, + "min": 1, + "max": 16384, + "tooltip": "Source height in pixels.", + }, + ), + "frame_count": ( + "INT", + { + "default": 97, + "min": 1, + "max": 10000, + "tooltip": "Source frame count. tile_frames is 8n+1, 97 or less, " + "or up to 121 when that avoids temporal chunking.", + }, + ), + "tile_size": ( + list(TILE_CHOICES), + { + "default": "FullHD", + "tooltip": "Spatial tile long-edge. qHD=960, HD=1280, FullHD=1920. " + "Clamped to the initial canvas so a FullHD tile on a smaller " + "canvas does not force tiling.", + }, + ), + "initial_canvas_size": ( + list(TILE_CHOICES), + { + "default": "FullHD", + "tooltip": "Stage-1 canvas long-edge (same choices as tile size). " + "Source within 200px of FullHD is kept instead of rescaled.", + }, + ), + "output_size": ( + list(OUTPUT_CHOICES), + { + "default": "FullHD", + "tooltip": "Final canvas. FullHD is the 1920 long-edge step; " + "4K is 2× that; 8K is 4×. Single-stage graphs use this as the " + "canvas. Two-stage graphs use initial canvas for stage 1 and " + "this for stage 2.", + }, + ), + }, + "optional": { + "image": ( + "IMAGE", + { + "tooltip": "Optional video/image batch. When connected, width, " + "height, and frame_count are taken from it (Get Image Size style)." + }, + ), + }, + } + + def get_sizes( + self, + width, + height, + frame_count, + tile_size="FullHD", + initial_canvas_size="FullHD", + output_size="FullHD", + image: Optional[object] = None, + **_kwargs, + ): + if image is not None: + width, height, frame_count = _dims_from_image(image) + plan = resolve_tiling_plan( + width, + height, + frame_count, + tile_size=tile_size, + initial_canvas_size=initial_canvas_size, + output_size=output_size, + ) + return { + "ui": {"text": [plan.summary]}, + "result": ( + plan.initial_canvas_width, + plan.initial_canvas_height, + plan.tile_width, + plan.tile_height, + plan.output_width, + plan.output_height, + plan.overlap_frac, + plan.tile_frames, + plan.is_8k, + plan.summary, + ), + } diff --git a/web/js/tiling_sizes_preview.js b/web/js/tiling_sizes_preview.js new file mode 100644 index 0000000..d5f5d0e --- /dev/null +++ b/web/js/tiling_sizes_preview.js @@ -0,0 +1,94 @@ +import { app } from "../../../scripts/app.js"; + +const NODE_NAME = "LTXVGetTilingSizes"; + +function linesOf(text) { + return String(text || "").split("\n").length; +} + +function ensurePreview(node, text) { + const value = Array.isArray(text) ? text.join("\n") : String(text ?? ""); + if (node._tilingPreviewEl) { + node._tilingPreviewEl.textContent = value; + const widget = (node.widgets || []).find((w) => w.name === "tiling_preview"); + if (widget?.computeSize) { + const lines = linesOf(value); + widget.computedHeight = Math.max(88, 14 + lines * 16); + } + return; + } + + const el = document.createElement("pre"); + el.className = "ltx-tiling-sizes-preview"; + el.textContent = value || "Queue prompt to preview tiling sizes."; + Object.assign(el.style, { + margin: "0", + padding: "6px 8px", + fontFamily: "ui-monospace, SFMono-Regular, Menlo, Consolas, monospace", + fontSize: "11px", + lineHeight: "1.4", + whiteSpace: "pre-wrap", + color: "var(--fg-color, #ddd)", + background: "var(--bg-color, transparent)", + border: "1px solid var(--border-color, rgba(255,255,255,0.12))", + borderRadius: "6px", + width: "100%", + boxSizing: "border-box", + pointerEvents: "none", + }); + + const widget = node.addDOMWidget("tiling_preview", "preview", el, { + serialize: false, + hideOnZoom: false, + getValue() { + return el.textContent; + }, + setValue(v) { + el.textContent = v ?? ""; + }, + }); + widget.computeSize = function (width) { + const lines = linesOf(el.textContent); + const height = Math.max(88, 14 + lines * 16); + this.computedHeight = height; + return [width, height]; + }; + + node._tilingPreviewEl = el; +} + +app.registerExtension({ + name: "LTXVideo.GetTilingSizesPreview", + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData?.name !== NODE_NAME) { + return; + } + + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + onNodeCreated?.apply(this, arguments); + ensurePreview(this, "Queue prompt to preview tiling sizes."); + }; + + const onExecuted = nodeType.prototype.onExecuted; + nodeType.prototype.onExecuted = function (message) { + onExecuted?.apply(this, arguments); + const text = message?.text; + if (text == null) { + return; + } + ensurePreview(this, text); + requestAnimationFrame(() => { + const size = this.computeSize?.() || this.size; + if (size && this.onResize) { + const next = [ + Math.max(size[0], this.size?.[0] || 0), + Math.max(size[1], this.size?.[1] || 0), + ]; + this.onResize(next); + } + app.graph?.setDirtyCanvas?.(true, false); + }); + }; + }, +});