Add spatial batching to latent upscale

H3_LONG_VIDEOS_GUIDE.md

@@ -1163,6 +1163,7 @@ What this really means:
 - the sampled latent is upscaled in latent space to the target size
 - the conditioning is rebuilt at that target size
 - the node then runs a short refinement pass over the upscaled latent with the sampler, scheduler, step count, denoise, and megapixel target you picked on the latent-upscale params node
+- if the target is larger than the spatial tile size, that refinement pass is processed in spatial batches using the same tile defaults as the upstream latent-split node
 
 Good starting point:
 
@@ -1170,6 +1171,7 @@ Good starting point:
 - use the interpolation mode when you want a cheaper resize-only path
 - start with `euler_ancestral`, `simple`, `2` steps, and `0.2` denoise
 - leave width and height at `0` unless you want an exact override; otherwise `megapixels` drives the target size
+- leave the spatial tile inputs at their defaults first: `512x512` tiles, `64` overlap, `0` fade width, `earlier` overlap mode
 
 The important part is that this stage is still a latent pass, not a pixel-space resize:
 

README.md

@@ -47,9 +47,9 @@
   - Per-shot directives now support `continuity:`, `ref_mode:`, `ref_noise_aug:`, `anchor_add:`, `soundscape:`, and `music:` in addition to the existing timing and wardrobe directives.
 
 - `Dumas H3 Latent Upscale Params`
-  - Inputs: `mode`, `model_name`, `method`, `width`, `height`, `device`, `precision`, `sampler_name`, `scheduler`, `steps`, `denoise`, `megapixels`
+  - Inputs: `mode`, `model_name`, `method`, `width`, `height`, `device`, `precision`, `sampler_name`, `scheduler`, `steps`, `denoise`, `megapixels`, `tile_width`, `tile_height`, `overlap`, `fade_width`, `overlap_mode`
   - Output: `latent_upscale_param`
-  - Bundles the optional latent-space upscaler settings used by `Dumas H3 Long Videos` before decode, so the main node can rebuild conditioning at the target size and run a short refinement pass with your chosen sampler, scheduler, step count, and denoise.
+  - Bundles the optional latent-space upscaler settings used by `Dumas H3 Long Videos` before decode, so the main node can rebuild conditioning at the target size and run a short refinement pass with your chosen sampler, scheduler, step count, denoise, and optional spatial batching.
 
 - `Dumas H3 Beat Prompt`
   - Inputs: authored through the custom front-end beat editor

dumas_h3_latent_upscale.py

@@ -463,13 +463,27 @@ class H3LatentUpscaleParams:
                     "tooltip": "How much the refinement pass may rewrite the upscaled latent. Lower = safer, higher = freer."}),
                 "megapixels": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 4.0, "step": 0.01,
                     "tooltip": "Primary target size for the latent refinement stage. If width and height are both set, they win; otherwise the node scales the current shot to this pixel budget while preserving aspect ratio. 0 keeps the incoming latent size."}),
+                "tile_width": ("INT", {"default": 512, "min": 32, "max": 4096, "step": 32,
+                    "tooltip": "Spatial tile width for the refinement stage in pixels. 512 matches the upstream latent-split default."}),
+                "tile_height": ("INT", {"default": 512, "min": 32, "max": 4096, "step": 32,
+                    "tooltip": "Spatial tile height for the refinement stage in pixels. 512 matches the upstream latent-split default."}),
+                "overlap": ("INT", {"default": 64, "min": 0, "max": 4096, "step": 32,
+                    "tooltip": "Pixel overlap between neighbouring spatial tiles. 64 matches the upstream latent-split default."}),
+                "fade_width": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 32,
+                    "tooltip": "Width in pixels of the freeze-to-free transition inside each overlap strip. 0 freezes the whole strip, matching the upstream default."}),
+                "overlap_mode": (["earlier", "later"], {"default": "earlier",
+                    "tooltip": "Which tile wins the overlap band when stitching the spatial batches back together."}),
             }
         }
 
-    def build(self, mode, model_name, method, width, height, device, precision, sampler_name, scheduler, steps, denoise, megapixels):
+    def build(self, mode, model_name, method, width, height, device, precision, sampler_name, scheduler, steps, denoise, megapixels, tile_width, tile_height, overlap, fade_width, overlap_mode):
         width = int(width)
         height = int(height)
         steps = int(steps)
+        tile_width = int(tile_width)
+        tile_height = int(tile_height)
+        overlap = int(overlap)
+        fade_width = int(fade_width)
         if mode == "off":
             return ({
                 "mode": "off",
@@ -481,6 +495,11 @@ class H3LatentUpscaleParams:
                 "denoise": float(denoise),
                 "refine_denoise": float(denoise),
                 "megapixels": float(megapixels),
+                "tile_width": tile_width,
+                "tile_height": tile_height,
+                "overlap": overlap,
+                "fade_width": fade_width,
+                "overlap_mode": overlap_mode,
             },)
         if width > 0:
             width = int(round(width / 32.0)) * 32
@@ -500,6 +519,11 @@ class H3LatentUpscaleParams:
             "denoise": float(denoise),
             "refine_denoise": float(denoise),
             "megapixels": float(megapixels),
+            "tile_width": tile_width,
+            "tile_height": tile_height,
+            "overlap": overlap,
+            "fade_width": fade_width,
+            "overlap_mode": overlap_mode,
         },)
 
 

dumas_h3_longvideos.py

@@ -4084,6 +4084,32 @@ def _latent_upscale_target_size(base_w, base_h, param):
     return int(base_w), int(base_h)
 
 
+def _latent_spatial_grid(h, w, th, tw, ol_h, ol_w):
+    if th <= 0 or tw <= 0:
+        raise ValueError("tile dimensions must be positive")
+    if ol_h >= th or ol_w >= tw:
+        raise ValueError("overlap must be smaller than the tile size")
+    sh = th - ol_h
+    sw = tw - ol_w
+    nrows = 1 if h <= th else math.ceil((h - ol_h) / sh)
+    if (nrows - 1) * sh + th < h:
+        nrows += 1
+    ncols = 1 if w <= tw else math.ceil((w - ol_w) / sw)
+    if (ncols - 1) * sw + tw < w:
+        ncols += 1
+    rows = [i * sh for i in range(nrows)]
+    cols = [j * sw for j in range(ncols)]
+    trows = [min(th, h - r) for r in rows]
+    tcols = [min(tw, w - c) for c in cols]
+    return rows, cols, trows, tcols
+
+
+def _latent_spatial_blend_weights(t, overlap_mode):
+    if overlap_mode == "later":
+        return 1.0 - t
+    return t
+
+
 def _nested_tensor_parts(samples):
     if samples is None:
         return ()
@@ -6119,7 +6145,8 @@ class H3LongVideos:
                 "latent_upscale_param": ("DUMAS_H3_LATENT_UPSCALE_PARAM", {
                     "tooltip": "Output of 'Dumas H3 Latent Upscale Params'. When connected, the first-pass "
                                "latent is upscaled and run through a short refinement pass before decode, "
-                               "using the sampler, scheduler, steps, denoise, and megapixel target from that node. "
+                               "using the sampler, scheduler, steps, denoise, megapixel target, and optional "
+                               "spatial batching from that node. "
                                "Leave unconnected to skip latent upscaling entirely."}),
                 "upscale": (["off", "rtx", "model", "lanczos"], {"default": "off",
                     "tooltip": "Optional post-pass on the finished frames. 'rtx' = NVIDIA RTX Video Super "
@@ -6478,9 +6505,68 @@ class H3LongVideos:
                 refine_scheduler = latent_upscale_param.get("scheduler", sch)
                 refine_denoise_value = latent_upscale_param.get("denoise", latent_upscale_param.get("refine_denoise", 0.2))
                 refine_denoise = 0.2 if refine_denoise_value is None else float(refine_denoise_value)
-                (refined_out,) = nodes.common_ksampler(
-                    model, seed, refine_steps, cfg, refine_sampler, refine_scheduler, upscale_cond, negative, upscale_latent,
-                    denoise=refine_denoise)
+                tile_w_px = int(latent_upscale_param.get("tile_width", 512) or 512)
+                tile_h_px = int(latent_upscale_param.get("tile_height", 512) or 512)
+                overlap_px = max(0, int(latent_upscale_param.get("overlap", 64) or 64))
+                fade_px = max(0, int(latent_upscale_param.get("fade_width", 0) or 0))
+                overlap_mode = str(latent_upscale_param.get("overlap_mode", "earlier"))
+                tile_tw = max(1, min(int(up_w), max(1, tile_w_px // 16)))
+                tile_th = max(1, min(int(up_h), max(1, tile_h_px // 16)))
+                ol_tw = max(0, min(tile_tw - 1, overlap_px // 16))
+                ol_th = max(0, min(tile_th - 1, overlap_px // 16))
+                fw_tw = max(0, min(ol_tw, fade_px // 16))
+                fw_th = max(0, min(ol_th, fade_px // 16))
+                rows, cols, trows, tcols = _latent_spatial_grid(int(up_h), int(up_w), tile_th, tile_tw, ol_th, ol_tw)
+                if len(rows) == 1 and len(cols) == 1:
+                    (refined_out,) = nodes.common_ksampler(
+                        model, seed, refine_steps, cfg, refine_sampler, refine_scheduler, upscale_cond, negative, upscale_latent,
+                        denoise=refine_denoise)
+                else:
+                    refined_video = upscaled_video.clone()
+                    full_audio = parts[1]
+                    for row_index, r0 in enumerate(rows):
+                        tr = trows[row_index]
+                        for col_index, c0 in enumerate(cols):
+                            tc = tcols[col_index]
+                            tile_target_w = int(tc) * 16
+                            tile_target_h = int(tr) * 16
+                            tile_cond, tile_latent = _build_shot_conditioning(
+                                clip, vae, prompt, tile_target_w, tile_target_h, ln, fps, handoff,
+                                ref_images=refs, ref_image_size=ref_image_size,
+                                ref_noise_aug=ref_noise_aug, audio_vae=audio_vae, silent=silent)
+                            tile_video = upscaled_video[:, :, :, r0:r0 + tr, c0:c0 + tc].contiguous()
+                            tile_latent["samples"] = comfy.nested_tensor.NestedTensor((tile_video, full_audio))
+                            tile_out, = nodes.common_ksampler(
+                                model, seed, refine_steps, cfg, refine_sampler, refine_scheduler, tile_cond, negative, tile_latent,
+                                denoise=refine_denoise)
+                            tile_out = _video_only_refined_latent(
+                                {"samples": comfy.nested_tensor.NestedTensor((tile_video, full_audio))},
+                                tile_out)
+                            tile_video_out = tile_out["samples"].tensors[0]
+                            region = refined_video[:, :, :, r0:r0 + tr, c0:c0 + tc]
+                            base_region = region.clone()
+                            region.copy_(tile_video_out)
+                            if col_index > 0 and ol_tw > 0:
+                                t = torch.linspace(0.0, 1.0, ol_tw, device=region.device, dtype=region.dtype)
+                                w = _latent_spatial_blend_weights(t, overlap_mode)
+                                if fw_tw > 0:
+                                    w = w.clone()
+                                    w[:fw_tw] = 0.0
+                                region[:, :, :, :, :ol_tw] = (
+                                    base_region[:, :, :, :, :ol_tw] * (1.0 - w[None, None, None, None, :]) +
+                                    tile_video_out[:, :, :, :, :ol_tw] * w[None, None, None, None, :]
+                                )
+                            if row_index > 0 and ol_th > 0:
+                                t = torch.linspace(0.0, 1.0, ol_th, device=region.device, dtype=region.dtype)
+                                w = _latent_spatial_blend_weights(t, overlap_mode)
+                                if fw_th > 0:
+                                    w = w.clone()
+                                    w[:fw_th] = 0.0
+                                region[:, :, :, :ol_th, :] = (
+                                    base_region[:, :, :, :ol_th, :] * (1.0 - w[None, None, None, :, None]) +
+                                    tile_video_out[:, :, :, :ol_th, :] * w[None, None, None, :, None]
+                                )
+                    refined_out = {"samples": comfy.nested_tensor.NestedTensor((refined_video, full_audio))}
                 timing["latent_upscale_sample"] += time.perf_counter() - latent_start
                 refined_out = _video_only_refined_latent(out, refined_out)
             except Exception as e:
@@ -6617,10 +6703,16 @@ class H3LongVideos:
             denoise = 0.2 if denoise_value is None else float(denoise_value)
             refine_sampler = latent_upscale_param.get("sampler_name", "euler_ancestral")
             refine_scheduler = latent_upscale_param.get("scheduler", "simple")
+            batch_note = ""
+            tile_w_px = int(latent_upscale_param.get("tile_width", 512) or 512)
+            tile_h_px = int(latent_upscale_param.get("tile_height", 512) or 512)
+            overlap_px = max(0, int(latent_upscale_param.get("overlap", 64) or 64))
+            if tile_w_px > 0 and tile_h_px > 0 and (tile_w_px < target_w or tile_h_px < target_h):
+                batch_note = f"; spatial batches {tile_w_px}x{tile_h_px}px overlap {overlap_px}px"
             latent_upscale_note = (
                 f" latent upscale: target {target_w}x{target_h}px{detail}; "
                 f"{int(latent_upscale_param.get('steps', 2) or 2)}-step refinement "
-                f"{refine_sampler}/{refine_scheduler} denoise {denoise:.2f}"
+                f"{refine_sampler}/{refine_scheduler} denoise {denoise:.2f}{batch_note}"
             )
 
         paras = split_paragraphs(prompt, "##")

tests/test_dumas_h3_longvideos.py

@@ -1127,6 +1127,11 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
         self.assertEqual(required["steps"][1]["default"], 2)
         self.assertEqual(required["denoise"][1]["default"], 0.2)
         self.assertEqual(required["megapixels"][1]["default"], 1.0)
+        self.assertEqual(required["tile_width"][1]["default"], 512)
+        self.assertEqual(required["tile_height"][1]["default"], 512)
+        self.assertEqual(required["overlap"][1]["default"], 64)
+        self.assertEqual(required["fade_width"][1]["default"], 0)
+        self.assertEqual(required["overlap_mode"][1]["default"], "earlier")
 
     def test_compose_persistent_does_not_expand_ambiguous_plural_to_full_cast(self):
         active = self.module.parse_wardrobe(