Add optional video-only detail pass

dumas_h3_longvideos.py

@@ -31,8 +31,10 @@ the VRAM-appropriate list. It splits the beats into shots that fit H3's ceiling
 and your VRAM, chains them, and returns the finished video + audio.
 
 Requirements: H3 is CFG-free (cfg 1) and needs no negative prompt -- the node
-makes an empty one internally. denoise is fixed at 1.0: a partial denoise desyncs
+makes an empty one internally. The main pass keeps denoise fixed at 1.0: a
-the joint audio/video schedule.
+partial denoise desyncs the joint audio/video schedule. An optional refinement
+pass can use its own denoise later, before any upscale, while keeping the output
+video-only.
 
 Verified against ComfyUI core (comfy_extras/nodes_minimax_h3.py, model_base.py,
 ldm/minimax/model.py, text_encoders/minimax.py, sd.py).
@@ -276,6 +278,8 @@ ADDED_WIDGETS = (
     "motion_guard", "contact_guard",
     "auto_soundscape", "allow_nonspeech_vocals",
     "ref_5", "ref_6", "ref_7", "ref_8", "ref_9",
+    "detail_pass", "detail_sampler_name", "detail_scheduler",
+    "detail_steps", "detail_denoise",
 )
 
 NL = "\n"
@@ -4014,6 +4018,37 @@ def _decode_audio(audio_vae, out_latent):
     return {"waveform": audio, "sample_rate": sr}
 
 
+def _copy_sample_latent(out_latent):
+    """Detach a sampled latent to CPU without changing its layout."""
+    raw = out_latent.get("samples") if isinstance(out_latent, dict) else None
+    if raw is None:
+        return None
+    try:
+        parts = raw.unbind() if hasattr(raw, "unbind") else None
+        return ([t.detach().to("cpu", copy=True) for t in parts]
+                if parts else raw.detach().to("cpu", copy=True))
+    except Exception:
+        return None
+
+
+def _video_only_refined_latent(base_latent, refined_latent):
+    """Keep the refined video latent, but preserve the original audio latent."""
+    base = base_latent.get("samples") if isinstance(base_latent, dict) else None
+    refined = refined_latent.get("samples") if isinstance(refined_latent, dict) else None
+    if base is None or refined is None:
+        return refined_latent
+    if not getattr(base, "is_nested", False) or not getattr(refined, "is_nested", False):
+        return refined_latent
+    try:
+        base_parts = base.unbind()
+        refined_parts = refined.unbind()
+        if len(base_parts) >= 2 and len(refined_parts) >= 1:
+            return {"samples": comfy.nested_tensor.NestedTensor((refined_parts[0], base_parts[-1]))}
+    except Exception:
+        return refined_latent
+    return refined_latent
+
+
 # --- ref2va reference conditioning ----------------------------------------
 # H3's reference pipeline encodes a reference image at up to a 2048 short edge.
 # Reference rows ride through EVERY sampling step, so this is also the setting
@@ -6101,6 +6136,19 @@ class H3LongVideos:
                                "the jacket, 'wardrobe: += sunglasses' adds one. TWO+ PEOPLE: name them -- "
                                "'Maya = grey shorts, red jacket; Jon = navy overalls', then edit one at a "
                                "time: 'wardrobe: Maya -= jacket' leaves Jon untouched."}),
+                "detail_pass": ("BOOLEAN", {"default": False,
+                    "tooltip": "Run a second refinement sampler on each beat BEFORE any upscale. "
+                               "It reuses the same conditioning and keeps the output video-only by "
+                               "preserving the first pass's audio latent. Good for extra detail without "
+                               "building a separate graph."}),
+                "detail_sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler",
+                    "tooltip": "Sampler used for the optional refinement pass."}),
+                "detail_scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "karras",
+                    "tooltip": "Scheduler used for the optional refinement pass."}),
+                "detail_steps": ("INT", {"default": 8, "min": 1, "max": 200,
+                    "tooltip": "Steps for the optional refinement pass."}),
+                "detail_denoise": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01,
+                    "tooltip": "How hard the refinement pass is allowed to rewrite the beat latent."}),
             },
             # Read-only graph access, for SLA-LoRA detection: a LoRA's filename is
             # the only thing that identifies an SLA build, and the graph is the only
@@ -6124,7 +6172,9 @@ class H3LongVideos:
 
     def _render(self, model, clip, vae, audio_vae, negative, prompt, w, h, ln, fps, tiled, sa,
                 handoff, decode_tile_frames=0, decode_tile_size=0,
-                refs=None, ref_image_size="match", ref_noise_aug=None, silent=False):
+                refs=None, ref_image_size="match", ref_noise_aug=None, silent=False,
+                detail_pass=False, detail_sampler_name="euler", detail_scheduler="karras",
+                detail_steps=8, detail_denoise=0.4):
         positive, latent = _build_shot_conditioning(clip, vae, prompt, w, h, ln, fps, handoff,
                                                     ref_images=refs, ref_image_size=ref_image_size,
                                                     ref_noise_aug=ref_noise_aug,
@@ -6144,24 +6194,27 @@ class H3LongVideos:
             if _is_oom(e):
                 e._h3_stage = "sampling"
             raise
+        refined_out = out
+        if detail_pass:
+            try:
+                (refined_out,) = nodes.common_ksampler(
+                    model, seed, int(detail_steps), cfg, detail_sampler_name, detail_scheduler,
+                    positive, negative, out, denoise=float(detail_denoise))
+            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.
-        raw = out.get("samples") if isinstance(out, dict) else None
+        shot_latent = _copy_sample_latent(refined_out)
-        shot_latent = None
+        video = _decode_video(vae, refined_out, tiled, free_first=model,
-        if raw is not None:
-            try:
-                parts = raw.unbind() if hasattr(raw, "unbind") else None
-                shot_latent = ([t.detach().to("cpu", copy=True) for t in parts]
-                               if parts else raw.detach().to("cpu", copy=True))
-            except Exception:
-                shot_latent = None      # never fail a render for the sake of an output
-        video = _decode_video(vae, out, tiled, free_first=model,
                               tile_t=decode_tile_frames, tile_xy=decode_tile_size)
         audio = _decode_audio(audio_vae, out)
-        del out, positive, latent
+        del out, refined_out, positive, latent
         _deep_cleanup()
         return video, audio, shot_latent
 
@@ -6193,6 +6246,8 @@ class H3LongVideos:
             ref_5=None, ref_6=None, ref_7=None, ref_8=None,
             ref_9=None,
             ref_mode="where tagged", ref_image_size="match", ref_noise_aug=0.999,
+            detail_pass=False, detail_sampler_name="euler", detail_scheduler="karras",
+            detail_steps=8, detail_denoise=0.4,
             graph=None, node_id=None):
 
         # FIRST: detect a checkpoint swap since the previous execution and hard-flush.
@@ -6253,6 +6308,12 @@ class H3LongVideos:
         ms_note = ""
         if apply_model_sampling:
             model, ms_note = apply_h3_model_sampling(model, shift_video, shift_audio)
+        detail_note = ""
+        if detail_pass:
+            detail_note = (f" detail pass: {int(detail_steps)} step(s) via "
+                           f"{detail_sampler_name}/{detail_scheduler} at denoise "
+                           f"{float(detail_denoise):.2f}; video-only refinement keeps "
+                           f"audio from the first pass")
 
         paras = split_paragraphs(prompt, "##")
         if anchor_override.strip():
@@ -6686,7 +6747,9 @@ class H3LongVideos:
                 while True:
                     try:
                         frames, audio, shot_latent = self._render(model, clip, vae, audio_vae, negative, gen_prompt, w, h, ln_i, fps, tiled, sa, shot_handoff, decode_tile_frames, decode_tile_size,
-                                                     shot_refs, ref_image_size, shot_aug, shot_silent)
+                                                     shot_refs, ref_image_size, shot_aug, shot_silent,
+                                                     detail_pass, detail_sampler_name, detail_scheduler,
+                                                     detail_steps, detail_denoise)
                         break
                     except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
                         if not _is_oom(e):
@@ -6702,7 +6765,9 @@ class H3LongVideos:
             else:
                 try:
                     frames, audio, shot_latent = self._render(model, clip, vae, audio_vae, negative, gen_prompt, w, h, ln_i, fps, tiled, sa, shot_handoff, decode_tile_frames, decode_tile_size,
-                                                     shot_refs, ref_image_size, shot_aug, shot_silent)
+                                                     shot_refs, ref_image_size, shot_aug, shot_silent,
+                                                     detail_pass, detail_sampler_name, detail_scheduler,
+                                                     detail_steps, detail_denoise)
                 except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
                     if _is_oom(e) and getattr(e, "_h3_stage", "") == "sampling":
                         # Retrying with tiles would re-run the whole sampling pass and
@@ -6715,7 +6780,9 @@ class H3LongVideos:
                         raise
                     mm.soft_empty_cache(True); tiled = True; backoff.append(f"shot {i+1}: tiled")
                     frames, audio, shot_latent = self._render(model, clip, vae, audio_vae, negative, gen_prompt, w, h, ln_i, fps, tiled, sa, shot_handoff, decode_tile_frames, decode_tile_size,
-                                                     shot_refs, ref_image_size, shot_aug, shot_silent)
+                                                     shot_refs, ref_image_size, shot_aug, shot_silent,
+                                                     detail_pass, detail_sampler_name, detail_scheduler,
+                                                     detail_steps, detail_denoise)
 
             if shot_latent is not None:
                 latent_chunks.append(shot_latent)
@@ -6970,6 +7037,7 @@ class H3LongVideos:
                    f"shot before them ended on dialogue." if mouth_settled else "")
                 + (f"{anatomy_note}." if anatomy_note else "")
                 + (f"{latent_note}." if latent_note else "")
+                + (f"{detail_note}." if detail_note else "")
                 + (f" SLA LoRA '{os.path.basename(str(sla_name))}' paired with sparse attention."
                    if sla_name and sparse_on else "")
                 + (f" {beats_note}." if beats_note else "")

tests/test_dumas_h3_longvideos.py

@@ -159,6 +159,99 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
             "keyframe carry",
         )
 
+    def test_detail_pass_refines_video_but_preserves_audio(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
+
+        calls = []
+        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):
+                calls.append((args, kwargs))
+                return (first_out if len(calls) == 1 else second_out,)
+
+            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 = lambda _vae, out_latent, *_args, **_kwargs: out_latent
+            self.module._decode_audio = lambda _vae, out_latent: out_latent
+            self.module._deep_cleanup = lambda: None
+
+            result = 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="karras",
+                detail_steps=5,
+                detail_denoise=0.4,
+            )
+
+            self.assertEqual(len(calls), 2)
+            self.assertIs(calls[1][0][8], first_out)
+            self.assertEqual(calls[1][0][4], "euler")
+            self.assertEqual(calls[1][0][5], "karras")
+            self.assertAlmostEqual(calls[1][1]["denoise"], 0.4)
+            self.assertEqual(result[1], first_out)
+            self.assertEqual(result[2][0].name, "v2")
+            self.assertEqual(result[2][1].name, "a1")
+            self.assertEqual(result[0]["samples"].unbind()[0].name, "v2")
+            self.assertEqual(result[0]["samples"].unbind()[-1].name, "a1")
+        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_distribute_generations_canonicalizes_per_shot_audio_and_anchor_directives(self):
         generations = self.module.distribute_generations(
             "",