From 6ddaca2b4eacb2e884d08c5fd142f4c2fbd797b7 Mon Sep 17 00:00:00 2001 From: mo-akkakk Date: Sun, 1 Nov 2026 17:07:12 +0100 Subject: [PATCH] Add short final window option --- README.md | 2 ++ __init__.py | 11 +++++++-- h3_windows.py | 20 +++++++++++----- pyproject.toml | 2 +- tests/test_windows.py | 53 +++++++++++++++++++++++++++++++++++++++++-- 5 files changed, 77 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 642f06e..075b959 100644 --- a/README.md +++ b/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. | diff --git a/__init__.py b/__init__.py index dd99ae8..e31ad99 100644 --- a/__init__.py +++ b/__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) diff --git a/h3_windows.py b/h3_windows.py index 7de1e85..45c9259 100644 --- a/h3_windows.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index e55b65a..73e9a2b 100644 --- a/pyproject.toml +++ b/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 = [] diff --git a/tests/test_windows.py b/tests/test_windows.py index 192e118..de44252 100644 --- a/tests/test_windows.py +++ b/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)