Harden H3 detail-pass cleanup

2026-09-03 07:18:13 +00:00
parent ab6ba86d76
commit 769257513b
2 changed files with 133 additions and 27 deletions
@@ -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