From 769257513bc2a568a7f46a0b50a392098a8db5a5 Mon Sep 17 00:00:00 2001 From: Chris Dumas Date: Thu, 3 Sep 2026 07:18:13 +0000 Subject: [PATCH] Harden H3 detail-pass cleanup --- dumas_h3_longvideos.py | 59 +++++++++-------- tests/test_dumas_h3_longvideos.py | 101 ++++++++++++++++++++++++++++++ 2 files changed, 133 insertions(+), 27 deletions(-) diff --git a/dumas_h3_longvideos.py b/dumas_h3_longvideos.py index 2bfb882..476a59f 100644 --- a/dumas_h3_longvideos.py +++ b/dumas_h3_longvideos.py @@ -6461,36 +6461,41 @@ class H3LongVideos: raise refined_out = out detail_pass = _coerce_bool_flag(detail_pass) - if detail_pass: - detail_latent = _latent_with_replaced_samples(latent, out) - try: - detail_start = time.perf_counter() - (refined_out,) = nodes.common_ksampler( + if detail_pass: + detail_latent = _latent_with_replaced_samples(latent, out) + try: + detail_start = time.perf_counter() + (refined_out,) = nodes.common_ksampler( model, seed, int(detail_steps), cfg, detail_sampler_name, detail_scheduler, positive, negative, detail_latent, denoise=float(detail_denoise)) timing["detail_sample"] += time.perf_counter() - detail_start - except Exception as e: - if _is_oom(e): - e._h3_stage = "sampling" - raise - refined_out = _video_only_refined_latent(out, refined_out) - # Keep a CPU copy of the sampled latent BEFORE decoding, for the `latent` - # output. Latents are ~1000x smaller than the frames they decode to (a - # 1344x768 124f shot is ~1.5MB against ~1.5GB), so carrying one per shot for - # the whole chain is free. Detached and moved off the card immediately, for - # the same reason the decoded frames are. - decode_video_start = time.perf_counter() - shot_latent = _copy_sample_latent(refined_out) - video = _decode_video(vae, refined_out, tiled, free_first=model, - tile_t=decode_tile_frames, tile_xy=decode_tile_size) - timing["decode_video"] += time.perf_counter() - decode_video_start - decode_audio_start = time.perf_counter() - audio = _decode_audio(audio_vae, out) - timing["decode_audio"] += time.perf_counter() - decode_audio_start - cleanup_start = time.perf_counter() - del out, refined_out, positive, latent - _deep_cleanup() - timing["cleanup"] += time.perf_counter() - cleanup_start + except Exception as e: + if _is_oom(e): + e._h3_stage = "sampling" + raise + refined_out = _video_only_refined_latent(out, refined_out) + del detail_latent + # Keep a CPU copy of the sampled latent BEFORE decoding, for the `latent` + # output. Latents are ~1000x smaller than the frames they decode to (a + # 1344x768 124f shot is ~1.5MB against ~1.5GB), so carrying one per shot for + # the whole chain is free. Detached and moved off the card immediately, for + # the same reason the decoded frames are. + decode_video_start = time.perf_counter() + shot_latent = _copy_sample_latent(refined_out) + decode_audio_start = time.perf_counter() + audio = _decode_audio(audio_vae, out) + timing["decode_audio"] += time.perf_counter() - decode_audio_start + # Audio is much smaller than the video decode. Drop the first-pass + # conditioning before the VAE work so the optional detail pass does not + # keep both sampled latents resident across the heaviest allocation. + del out, positive, latent + video = _decode_video(vae, refined_out, tiled, free_first=model, + tile_t=decode_tile_frames, tile_xy=decode_tile_size) + timing["decode_video"] += time.perf_counter() - decode_video_start + cleanup_start = time.perf_counter() + del refined_out + _deep_cleanup() + timing["cleanup"] += time.perf_counter() - cleanup_start if timing_sink is not None: timing["total"] = sum(timing.values()) timing_sink.append(timing) diff --git a/tests/test_dumas_h3_longvideos.py b/tests/test_dumas_h3_longvideos.py index 832c5ed..4fc65e4 100644 --- a/tests/test_dumas_h3_longvideos.py +++ b/tests/test_dumas_h3_longvideos.py @@ -328,6 +328,107 @@ class DumasH3LongVideosHelperTests(unittest.TestCase): else: self.module.comfy.nested_tensor.NestedTensor = original_nested + def test_detail_pass_decodes_audio_before_video_and_cleans_up(self): + class FakeTensor: + def __init__(self, name): + self.name = name + + def detach(self): + return self + + def to(self, *args, **kwargs): + return self + + class FakeNestedTensor: + def __init__(self, parts): + self._parts = tuple(parts) + self.is_nested = True + + def unbind(self): + return self._parts + + order = [] + first_out = {"samples": FakeNestedTensor((FakeTensor("v1"), FakeTensor("a1")))} + second_out = {"samples": FakeNestedTensor((FakeTensor("v2"), FakeTensor("a2")))} + + original_common_ksampler = self.module.nodes.common_ksampler + original_build = self.module._build_shot_conditioning + original_evict = self.module._evict_all_but + original_decode_video = self.module._decode_video + original_decode_audio = self.module._decode_audio + original_cleanup = self.module._deep_cleanup + original_nested = getattr(self.module.comfy.nested_tensor, "NestedTensor", None) + try: + self.module.comfy.nested_tensor.NestedTensor = FakeNestedTensor + + def common_ksampler(*args, **kwargs): + order.append("detail_sample" if len(order) else "sample") + return (first_out if len([x for x in order if x.endswith("sample")]) == 1 else second_out,) + + def decode_audio(_vae, out_latent): + order.append("audio") + self.assertIs(out_latent, first_out) + return out_latent + + def decode_video(_vae, out_latent, *_args, **_kwargs): + order.append("video") + self.assertIsNot(out_latent, first_out) + self.assertIs(out_latent["samples"].unbind()[0], second_out["samples"].unbind()[0]) + return out_latent + + def cleanup(): + order.append("cleanup") + + self.module.nodes.common_ksampler = common_ksampler + self.module._build_shot_conditioning = lambda *_args, **_kwargs: ( + "cond", + {"samples": FakeNestedTensor((FakeTensor("basev"), FakeTensor("basea")))}, + ) + self.module._evict_all_but = lambda *_args, **_kwargs: None + self.module._decode_video = decode_video + self.module._decode_audio = decode_audio + self.module._deep_cleanup = cleanup + + self.module.H3LongVideos()._render( + model=object(), + clip=types.SimpleNamespace( + tokenize=lambda text, **kwargs: text, + encode_from_tokens_scheduled=lambda tokens: tokens, + ), + vae=object(), + audio_vae=object(), + negative="negative", + prompt="beat", + w=128, + h=64, + ln=24, + fps=24, + tiled=False, + sa=(123, 20, 1.0, "res_multistep", "simple", 1.0), + handoff=None, + detail_pass=True, + detail_sampler_name="euler", + detail_scheduler="beta", + detail_steps=5, + detail_denoise=0.4, + ) + + self.assertEqual(order[0], "sample") + self.assertEqual(order[1], "detail_sample") + self.assertLess(order.index("audio"), order.index("video")) + self.assertEqual(order[-1], "cleanup") + finally: + self.module.nodes.common_ksampler = original_common_ksampler + self.module._build_shot_conditioning = original_build + self.module._evict_all_but = original_evict + self.module._decode_video = original_decode_video + self.module._decode_audio = original_decode_audio + self.module._deep_cleanup = original_cleanup + if original_nested is None: + delattr(self.module.comfy.nested_tensor, "NestedTensor") + else: + self.module.comfy.nested_tensor.NestedTensor = original_nested + def test_detail_pass_treats_falsey_strings_as_disabled(self): calls = [] original_common_ksampler = self.module.nodes.common_ksampler