Add short final window option

README.md

@@ -41,6 +41,8 @@ Everything else is the stock H3 workflow: Reference to Video, Add Guide, Sampler
 | `temporal_windows` | Split the shot into overlapping windows. Off = the whole length runs as one pass. |
 | `window_frames` | Frames per window, on H3's 17k+5 grid (124 = 5.2 s, 175 = 7.3 s). Keep it within what H3 handles accurately on your setup. |
 | `overlap_frames` | Frames shared by neighbouring windows (5, 22, 39…). They are carried into the next window as its frame-0 guide, then crossfaded. The last window is pulled back to end on the last frame, so its overlap can be larger. |
+| `short_final_window` | Off (default): the last window is pulled back so it is a full `window_frames` long, which re-generates some frames at the end. On: the last window keeps its natural, shorter length, which saves those frames — most useful with large windows. |
+| `min_final_window_frames` | With `short_final_window` on: the shortest the last window may be (17k+5: 5, 22, 39, 56…). A shorter tail is pulled back only far enough to reach it. |
 | `carry_overlap` | Carry the finished overlap (video and audio) into the next window. |
 | `spatial_tiles` | Split each window into overlapping tiles (stage 1 or stage 2). Needs a plate: without a reference video, a full-denoise pass in tiles repeats the subject in every tile. |
 | `tile_width` / `tile_height` / `tile_overlap` | Tile size in pixels. Use a size H3 handles well on your setup, with ≥ 192 px overlap. |

__init__.py

@@ -26,6 +26,12 @@ class MiniMaxH3TemporalWindow(io.ComfyNode):
                              tooltip="Frames per window (17k+5). Keep it inside what works natively for you, e.g. <= 192."),
                 io.Int.Input("overlap_frames", default=22, min=5, max=1000, step=17,
                              tooltip="Frames shared by neighbouring windows (5, 22, 39...). Carried into the next window as its frame-0 guide, then crossfaded."),
+                io.Boolean.Input("short_final_window", default=False,
+                                 tooltip="Let the last window be shorter than window_frames instead of pulling it back to full length. "
+                                         "Saves the frames that would be generated twice at the end, which matters most with large windows. Off = previous behaviour."),
+                io.Int.Input("min_final_window_frames", default=22, min=5, max=3600, step=17,
+                             tooltip="With short_final_window on: the last window is never shorter than this (17k+5: 5, 22, 39, 56...). "
+                                     "If the natural tail is shorter, the window is pulled back only far enough to reach it."),
                 io.Boolean.Input("carry_overlap", default=True,
                                  tooltip="Give each window the previous window's finished overlap frames (video and audio) as its frame-0 guide."),
                 io.Boolean.Input("spatial_tiles", default=False,
@@ -40,11 +46,12 @@ class MiniMaxH3TemporalWindow(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model, temporal_windows, window_frames, overlap_frames, carry_overlap, spatial_tiles, tile_width, tile_height,
+    def execute(cls, model, temporal_windows, window_frames, overlap_frames, short_final_window, min_final_window_frames, carry_overlap, spatial_tiles, tile_width, tile_height,
                 tile_overlap, crop_ref_images_per_tile, clip=None) -> io.NodeOutput:
         wrapper = WindowsWrapper(window_frames if temporal_windows else 0, overlap_frames, carry_overlap,
                                  tile_height // 16 if spatial_tiles else 0, tile_width // 16 if spatial_tiles else 0,
-                                 tile_overlap // 16, crop_ref_images_per_tile)
+                                 tile_overlap // 16, crop_ref_images_per_tile,
+                                 min_final_window_frames if short_final_window else 0)
         return io.NodeOutput(apply(model, wrapper), recording_clip(clip) if clip is not None else None)
 
 

h3_windows.py

@@ -45,17 +45,23 @@ def audio_ticks(frames):
     return round(frames * FRAME_RESCALE)
 
 
-def plan_windows(latent_t, window_frames, overlap_frames):
-    """[(start, length)] in latent tokens. Clip-aligned; the last window is pulled back to end on the last frame."""
+def plan_windows(latent_t, window_frames, overlap_frames, min_final_frames=0):
+    """[(start, length)] in latent tokens, clip-aligned.
+
+    The last window ends on the last frame. By default it is pulled back to the full window
+    length. With min_final_frames > 0 it keeps its natural (shorter) length when that is at
+    least min_final_frames, and is otherwise pulled back only far enough to reach that minimum.
+    """
     win = tokens_for_frames(window_frames)
     if window_frames <= 0 or win >= latent_t:
         return [(0, latent_t)]
     ov = min(tokens_for_frames(overlap_frames), win - CLIP_TOKENS)
+    min_tok = min(win, tokens_for_frames(min_final_frames)) if min_final_frames > 0 else win
     windows = [(0, win)]
     while windows[-1][0] + win < latent_t:
         start = windows[-1][0] + win - ov
-        if start + win > latent_t:
-            start = (latent_t - win) // CLIP_TOKENS * CLIP_TOKENS
+        if start + win > latent_t and latent_t - start < min_tok:
+            start = (latent_t - min_tok) // CLIP_TOKENS * CLIP_TOKENS
         windows.append((start, min(win, latent_t - start)))
     return windows
 
@@ -247,9 +253,11 @@ def piece_conds(conds, k0, k1, ys, xs, lat_h, lat_w, crop_images, carry, warned,
 
 
 class WindowsWrapper:
-    def __init__(self, window_frames, overlap_frames, carry_overlap, tile_h, tile_w, tile_overlap, crop_images):
+    def __init__(self, window_frames, overlap_frames, carry_overlap, tile_h, tile_w, tile_overlap, crop_images,
+                 min_final_frames=0):
         self.window_frames = window_frames
         self.overlap_frames = overlap_frames
+        self.min_final_frames = min_final_frames
         self.carry_overlap = carry_overlap
         self.tile_h, self.tile_w, self.tile_overlap = tile_h, tile_w, tile_overlap
         self.crop_images = crop_images
@@ -261,7 +269,7 @@ class WindowsWrapper:
             return executor(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes=latent_shapes)
         vshape, ashape = latent_shapes[0], latent_shapes[1]
         latent_t, lat_h, lat_w, audio_t = vshape[2], vshape[3], vshape[4], ashape[-1]
-        windows = plan_windows(latent_t, self.window_frames, self.overlap_frames)
+        windows = plan_windows(latent_t, self.window_frames, self.overlap_frames, self.min_final_frames)
         ys = plan_tiles(lat_h, self.tile_h, self.tile_overlap)
         xs = plan_tiles(lat_w, self.tile_w, self.tile_overlap)
         if len(windows) == 1 and len(ys) == 1 and len(xs) == 1:

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-minimax-h3-temporal-window-spatial-tiles"
 description = "Minimax H3 Temporal Window & Spatial Tiles: long video-to-video for the open-weights MiniMax H3 as overlapping temporal windows sampled one after another (stage 1 and 2), plus spatial tiles that follow their crop of the plate. One model patch, works with native ComfyUI H3 nodes."
-version = "1.0.2"
+version = "1.0.3"
 license = { file = "LICENSE" }
 dependencies = []
 

tests/test_windows.py

@@ -103,8 +103,8 @@ def run(model, pos, lat, sigmas, seed=3):
     return SamplerCustomAdvanced.execute(Noise_RandomNoise(seed), g, comfy.samplers.sampler_object("euler"), sigmas, lat).result
 
 
-def patch(model, clip, frames=124, overlap=22, tiles=False, tw=960, th=544, tov=256):
-    return pack.MiniMaxH3TemporalWindow.execute(model, frames > 0, max(frames, 22), overlap, True, tiles, tw, th, tov, True, clip=clip).result
+def patch(model, clip, frames=124, overlap=22, tiles=False, tw=960, th=544, tov=256, short_final=False, min_final=22):
+    return pack.MiniMaxH3TemporalWindow.execute(model, frames > 0, max(frames, 22), overlap, short_final, min_final, True, tiles, tw, th, tov, True, clip=clip).result
 
 
 def test_plans():
@@ -129,6 +129,53 @@ def test_switch_off_is_native(model):
     check("temporal_windows off: whole length in one native pass", torch.equal(native[0], mine[0]) and torch.equal(native[1], mine[1]))
 
 
+def test_short_final_window_plan():
+    T = h3.temporal_shape(362)[1]
+    F = hw.frames_for_tokens
+    full = hw.plan_windows(T, 243, 22)
+    short = hw.plan_windows(T, 243, 22, 22)
+    check("default: final window pulled back to full length {}".format([(F(s), F(n)) for s, n in full]), F(full[-1][1]) == 243)
+    check("short final window keeps its natural length {}".format([(F(s), F(n)) for s, n in short]),
+          F(short[-1][1]) == 141 and short[-1][0] + short[-1][1] == T)
+    T2 = h3.temporal_shape(719)[1]
+    w = hw.plan_windows(T2, 192, 22, 124)
+    check("short tail below the minimum is pulled back only to the minimum {}".format([(F(s), F(n)) for s, n in w]),
+          F(w[-1][1]) == 124 and w[-1][0] + w[-1][1] == T2)
+    ok = True
+    for frames in (124, 209, 362, 400, 719):
+        Tn = h3.temporal_shape(frames)[1]
+        for mn in (5, 22, 39, 56):
+            ws = hw.plan_windows(Tn, 124, 22, mn)
+            ok &= ws[0][0] == 0 and ws[-1][0] + ws[-1][1] == Tn and all(s % 5 == 0 and (n - 2) % 5 == 0 for s, n in ws)
+            ok &= all(ws[i][0] + ws[i][1] > ws[i + 1][0] for i in range(len(ws) - 1))
+            ok &= F(ws[-1][1]) >= min(mn, 124)
+    check("every plan with a short final window is valid, overlapping and covers the shot", bool(ok))
+
+
+def test_short_final_window_sampling(model):
+    W, H, N = 384, 256, 260
+    plate, ref = torch.rand(N, H, W, 3), torch.rand(1, H, W, 3)
+    m, c = patch(model, FakeClip(), frames=124, overlap=22, short_final=True, min_final=5)
+    pos, lat = conditioning(c, plate, ref, W, H, N)
+    T = lat["samples"].unbind()[0].shape[2]
+    wins = hw.plan_windows(T, 124, 22, 5)
+    lengths = []
+    orig = comfy.samplers.CFGGuider.inner_sample
+
+    def rec(self, *a, **k):
+        lengths.append(k["latent_shapes"][0][2])
+        return orig(self, *a, **k)
+    comfy.samplers.CFGGuider.inner_sample = rec
+    try:
+        sigmas = comfy.samplers.calculate_sigmas(model.get_model_object("model_sampling"), "simple", 1)
+        out = run(m, pos, lat, sigmas)[0]["samples"].unbind()
+    finally:
+        comfy.samplers.CFGGuider.inner_sample = orig
+    check("sampling with a short final window ({} frames) covers the shot".format(hw.frames_for_tokens(wins[-1][1])),
+          lengths == [n for _, n in wins] and wins[-1][1] < hw.tokens_for_frames(124)
+          and out[0].shape == lat["samples"].unbind()[0].shape and bool(torch.isfinite(out[0]).all()))
+
+
 def test_one_window_is_native(model):
     W, H, N = 384, 256, 56
     plate, ref = torch.rand(N, H, W, 3), torch.rand(1, H, W, 3)
@@ -237,10 +284,12 @@ def test_stage2_tiles(model):
 
 if __name__ == "__main__":
     test_plans()
+    test_short_final_window_plan()
     with tempfile.TemporaryDirectory() as d:
         model = tiny_model(os.path.join(d, "tiny_h3.safetensors"))
         test_one_window_is_native(model)
         test_switch_off_is_native(model)
+        test_short_final_window_sampling(model)
         test_windows(model)
         test_ksampler(model)
         test_stage2_tiles(model)