diff --git a/dumas_h3_longvideos.py b/dumas_h3_longvideos.py index 998e70a..bc0494c 100644 --- a/dumas_h3_longvideos.py +++ b/dumas_h3_longvideos.py @@ -30,9 +30,11 @@ video; each later paragraph = a scene beat), a shot length, and a resolution fro 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 -the joint audio/video schedule. +Requirements: H3 is CFG-free (cfg 1) and needs no negative prompt -- the node +makes an empty one internally. The main pass keeps denoise fixed at 1.0: a +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" @@ -4002,16 +4006,47 @@ def _silent_audio_latent(audio_vae, frame_count, fps): return None # never fail a render for a nicety -def _decode_audio(audio_vae, out_latent): - latent = out_latent["samples"] - if latent.is_nested: - latent = latent.unbind()[-1] - audio = audio_vae.decode(latent).movedim(-1, 1) - std = torch.std(audio, dim=[1, 2], keepdim=True) * 5.0 - std[std < 1.0] = 1.0 - audio = audio / std - sr = getattr(audio_vae, "audio_sample_rate_output", getattr(audio_vae, "audio_sample_rate", 44100)) - return {"waveform": audio, "sample_rate": sr} +def _decode_audio(audio_vae, out_latent): + latent = out_latent["samples"] + if latent.is_nested: + latent = latent.unbind()[-1] + audio = audio_vae.decode(latent).movedim(-1, 1) + std = torch.std(audio, dim=[1, 2], keepdim=True) * 5.0 + std[std < 1.0] = 1.0 + audio = audio / std + sr = getattr(audio_vae, "audio_sample_rate_output", getattr(audio_vae, "audio_sample_rate", 44100)) + 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 ---------------------------------------- @@ -6084,11 +6119,11 @@ class H3LongVideos: "when your scene contains distress sounds that H3 would otherwise " "suppress. Keep auto_silence_nonspeech ON for shots that should be " "truly silent."}), - "character_memory": ("STRING", {"multiline": True, "forceInput": True, "default": "", - "tooltip": "Optional dedicated wardrobe channel (same role as a 'wardrobe:' line in " - "the first paragraph -- use whichever you prefer; this field wins if both " - "are set). Re-stamped into every shot so clothing holds even when the " - "camera crops it out. IMPORTANT: this is the ONLY place clothing should " + "character_memory": ("STRING", {"multiline": True, "forceInput": True, "default": "", + "tooltip": "Optional dedicated wardrobe channel (same role as a 'wardrobe:' line in " + "the first paragraph -- use whichever you prefer; this field wins if both " + "are set). Re-stamped into every shot so clothing holds even when the " + "camera crops it out. IMPORTANT: this is the ONLY place clothing should " "live -- keep it out of the anchor prose, or a removal won't stick because " "the immutable anchor keeps re-adding it. To change/remove an item " "mid-chain, put 'wardrobe: ' inside the beat where it " @@ -6097,11 +6132,24 @@ class H3LongVideos: "'a woman with silver hair'. A noun phrase renders as 'She (a woman with...)', " "i.e. two subjects in one clause, which causes character duplication. The node " "strips them automatically, but writing attributes directly is cleaner. " - "ONE-TOKEN EDITS (no restating the outfit): 'wardrobe: -= jacket' removes " - "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."}), - }, + "ONE-TOKEN EDITS (no restating the outfit): 'wardrobe: -= jacket' removes " + "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 # place it survives. Named 'graph'/'node_id' rather than the usual @@ -6122,48 +6170,53 @@ class H3LongVideos: opt[name] = opt.pop(name) # re-insert at the end, value unchanged return schema - 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): - 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, - audio_vae=audio_vae, silent=silent) - seed, steps, cfg, sn, sch, denoise = sa + 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, + 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, + audio_vae=audio_vae, silent=silent) + seed, steps, cfg, sn, sch, denoise = sa # Conditioning is built, so the text encoder and VAEs are dead weight for the # whole sampling loop -- evict them and keep only the DiT on the card. _evict_all_but(model) - try: - (out,) = nodes.common_ksampler(model, seed, steps, cfg, sn, sch, positive, negative, - latent, denoise=denoise) - except Exception as e: + try: + (out,) = nodes.common_ksampler(model, seed, steps, cfg, sn, sch, positive, negative, + latent, denoise=denoise) + except Exception as e: # Mark WHERE this failed. `tiled` only affects the DECODE, so the caller's # OOM retry cannot help an OOM raised here -- it just re-runs the whole # sampling pass and fails the same way, which on a 362-frame shot is four # more minutes for nothing. - if _is_oom(e): - e._h3_stage = "sampling" - raise - # 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 = None - 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 - _deep_cleanup() - return video, audio, shot_latent + 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. + 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) + audio = _decode_audio(audio_vae, out) + del out, refined_out, positive, latent + _deep_cleanup() + return video, audio, shot_latent def run(self, model, clip, vae, audio_vae, prompt, resolution, steps, cfg, sampler_name, scheduler, seed, @@ -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. @@ -6250,11 +6305,17 @@ class H3LongVideos: # Patch the dual video/audio schedule onto the model here, so a missing # upstream ModelSamplingMiniMaxH3 can't silently produce gibberish audio. # Shifts come from the widgets (12/3 base default; MXFP8/turbo differ). - ms_note = "" - if apply_model_sampling: - model, ms_note = apply_h3_model_sampling(model, shift_video, shift_audio) - - paras = split_paragraphs(prompt, "##") + 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(): anchor, beat_paras = anchor_override.strip(), paras elif paras: @@ -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 "") diff --git a/tests/test_dumas_h3_longvideos.py b/tests/test_dumas_h3_longvideos.py index fddb3a3..c8a93d8 100644 --- a/tests/test_dumas_h3_longvideos.py +++ b/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( "",