Retry long videos decode with tiling fallback

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

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