diff --git a/dumas_h3_longvideos.py b/dumas_h3_longvideos.py index 5cc4d23..ef874ee 100644 --- a/dumas_h3_longvideos.py +++ b/dumas_h3_longvideos.py @@ -6372,16 +6372,29 @@ class H3LongVideos: 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 + # 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) + try: + video = _decode_video(vae, refined_out, tiled, free_first=model, + tile_t=decode_tile_frames, tile_xy=decode_tile_size) + except Exception as e: + # Decode is the biggest allocation in the run. If the straight path + # fails, retry once with tiled decode and without the aggressive unload + # so a marginal card can still finish the render. + if not _is_oom(e) and "decode" not in str(e).lower(): + raise + mm.soft_empty_cache(True) + retry_tiled = True + retry_tile_t = decode_tile_frames or 16 + retry_tile_xy = decode_tile_size or 256 + video = _decode_video(vae, refined_out, retry_tiled, free_first=None, + tile_t=retry_tile_t, tile_xy=retry_tile_xy) + 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 diff --git a/tests/test_dumas_h3_longvideos.py b/tests/test_dumas_h3_longvideos.py index 97e77e8..eca4397 100644 --- a/tests/test_dumas_h3_longvideos.py +++ b/tests/test_dumas_h3_longvideos.py @@ -226,6 +226,7 @@ class DumasH3LongVideosHelperTests(unittest.TestCase): original_cleanup = self.module._deep_cleanup original_nested = getattr(self.module.comfy.nested_tensor, "NestedTensor", None) try: + sentinel_model = object() self.module.comfy.nested_tensor.NestedTensor = FakeNestedTensor def common_ksampler(*args, **kwargs): @@ -243,7 +244,7 @@ class DumasH3LongVideosHelperTests(unittest.TestCase): self.module._deep_cleanup = lambda: None result = self.module.H3LongVideos()._render( - model=object(), + model=sentinel_model, clip=types.SimpleNamespace( tokenize=lambda text, **kwargs: text, encode_from_tokens_scheduled=lambda tokens: tokens, @@ -289,6 +290,95 @@ class DumasH3LongVideosHelperTests(unittest.TestCase): else: self.module.comfy.nested_tensor.NestedTensor = original_nested + def test_render_retries_decode_with_tiling_after_decode_oom(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 + + decode_calls = [] + cleanup_calls = [] + first_out = {"samples": FakeNestedTensor((FakeTensor("v1"), FakeTensor("a1")))} + + 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_soft_empty_cache = self.module.mm.soft_empty_cache + original_nested = getattr(self.module.comfy.nested_tensor, "NestedTensor", None) + try: + sentinel_model = object() + self.module.comfy.nested_tensor.NestedTensor = FakeNestedTensor + self.module.nodes.common_ksampler = lambda *args, **kwargs: (first_out,) + self.module._build_shot_conditioning = lambda *_args, **_kwargs: ( + "cond", + {"samples": FakeNestedTensor((FakeTensor("basev"), FakeTensor("basea")))}, + ) + self.module._evict_all_but = lambda *_args, **_kwargs: None + + def decode_video(_vae, out_latent, tiled, free_first=None, tile_t=None, tile_xy=None): + decode_calls.append((tiled, free_first, tile_t, tile_xy)) + if len(decode_calls) == 1: + raise RuntimeError("CUDA out of memory during decode") + return out_latent + + self.module._decode_video = decode_video + self.module._decode_audio = lambda _vae, out_latent: out_latent + self.module.mm.soft_empty_cache = lambda *args, **kwargs: cleanup_calls.append((args, kwargs)) + self.module._deep_cleanup = lambda: None + + result = self.module.H3LongVideos()._render( + model=sentinel_model, + 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, + ) + + self.assertEqual(len(decode_calls), 2) + self.assertEqual(decode_calls[0], (False, sentinel_model, 0, 0)) + self.assertEqual(decode_calls[1], (True, None, 16, 256)) + self.assertTrue(cleanup_calls) + self.assertIs(result[1], first_out) + 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 + self.module.mm.soft_empty_cache = original_soft_empty_cache + 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