Harden H3 detail-pass cleanup

dumas_h3_longvideos.py

@@ -6474,6 +6474,7 @@ class H3LongVideos:
                     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
@@ -6481,14 +6482,18 @@ class H3LongVideos:
         # 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
+        # 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 out, refined_out, positive, latent
+        del refined_out
         _deep_cleanup()
         timing["cleanup"] += time.perf_counter() - cleanup_start
         if timing_sink is not None:

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