From 7bb0b0abea75d948835747b4b430a1db873868c3 Mon Sep 17 00:00:00 2001 From: Chris Dumas Date: Fri, 4 Sep 2026 11:49:21 +0000 Subject: [PATCH] Reuse conditioning for latent upscale refine --- dumas_h3_longvideos.py | 80 +++++++++++++++++++++++-------- tests/test_dumas_h3_longvideos.py | 14 ++++-- 2 files changed, 70 insertions(+), 24 deletions(-) diff --git a/dumas_h3_longvideos.py b/dumas_h3_longvideos.py index 1aad51e..58d79b8 100644 --- a/dumas_h3_longvideos.py +++ b/dumas_h3_longvideos.py @@ -4027,23 +4027,59 @@ 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 +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 _latent_with_replaced_samples(template_latent, sampled_latent): - """Reuse the original latent payload, but swap in freshly sampled tensors.""" - if not isinstance(template_latent, dict): - return sampled_latent + except Exception: + return None + + +def _retarget_conditioning_spatial(cond, latent_h, latent_w): + """Resize H3 keyframe latents in existing conditioning to a new latent grid.""" + latent_h = int(latent_h) + latent_w = int(latent_w) + if latent_h <= 0 or latent_w <= 0: + raise RuntimeError("conditioning target latent size must be positive") + out = [] + for item in cond: + try: + tensor, data = item + except Exception: + out.append(item) + continue + nd = dict(data) + keyframes = nd.get("minimax_keyframes") + if keyframes: + resized_keyframes = [] + for keyframe in keyframes: + nkf = dict(keyframe) + latent_value = nkf.get("latent") + if latent_value is not None and len(getattr(latent_value, "shape", ())) >= 5: + if latent_value.shape[3] != latent_h or latent_value.shape[4] != latent_w: + b, c, t, h, w = latent_value.shape + resized = torch.nn.functional.interpolate( + latent_value.to(torch.float32).reshape(b * t, c, h, w), + size=(latent_h, latent_w), + mode="bilinear", + align_corners=False, + ).reshape(b, c, t, latent_h, latent_w) + nkf["latent"] = resized.to(device=latent_value.device, dtype=latent_value.dtype) + resized_keyframes.append(nkf) + nd["minimax_keyframes"] = resized_keyframes + out.append([tensor, nd]) + return out + + +def _latent_with_replaced_samples(template_latent, sampled_latent): + """Reuse the original latent payload, but swap in freshly sampled tensors.""" + if not isinstance(template_latent, dict): + return sampled_latent out = dict(template_latent) if isinstance(sampled_latent, dict): for key, value in sampled_latent.items(): @@ -6693,18 +6729,22 @@ class H3LongVideos: # pass; otherwise the 12-step base latent and the upscale latent sit # in memory together and can trigger a retry loop. out["samples"] = comfy.nested_tensor.NestedTensor((upscaled_video, full_audio)) - del out_samples, positive, latent, parts + del out_samples, parts mm.soft_empty_cache() target_w = int(up_w) * 16 target_h = int(up_h) * 16 if target_w <= 0 or target_h <= 0: raise RuntimeError("latent upscale target size must be positive") - upscale_cond, upscale_latent = _build_shot_conditioning( - clip, vae, prompt, target_w, target_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) - upscale_latent["samples"] = comfy.nested_tensor.NestedTensor( - (upscaled_video, full_audio)) + try: + upscale_cond = _retarget_conditioning_spatial(positive, int(up_h), int(up_w)) + upscale_latent = dict(latent) if isinstance(latent, dict) else {} + except Exception: + upscale_cond, upscale_latent = _build_shot_conditioning( + clip, vae, prompt, target_w, target_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) + upscale_latent["samples"] = comfy.nested_tensor.NestedTensor((upscaled_video, full_audio)) + del positive, latent refine_steps = int(latent_upscale_param.get("steps", 2) or 2) refine_sampler = latent_upscale_param.get("sampler_name", sn) refine_scheduler = latent_upscale_param.get("scheduler", sch) diff --git a/tests/test_dumas_h3_longvideos.py b/tests/test_dumas_h3_longvideos.py index 7522278..2293e91 100644 --- a/tests/test_dumas_h3_longvideos.py +++ b/tests/test_dumas_h3_longvideos.py @@ -362,6 +362,7 @@ class DumasH3LongVideosHelperTests(unittest.TestCase): return self._parts order = [] + build_calls = [] first_out = {"samples": FakeNestedTensor((FakeTensor("v1"), FakeTensor("a1")))} second_out = {"samples": FakeNestedTensor((FakeTensor("v2"), FakeTensor("a2")))} @@ -408,11 +409,15 @@ class DumasH3LongVideosHelperTests(unittest.TestCase): order.append("upscale") return FakeTensor("upv"), 8, 16 + def build_conditioning(*_args, **_kwargs): + build_calls.append(True) + return ( + [["cond", {}]], + {"samples": FakeNestedTensor((FakeTensor("basev"), FakeTensor("basea")))}, + ) + self.module.nodes.common_ksampler = common_ksampler - self.module._build_shot_conditioning = lambda *_args, **_kwargs: ( - "cond", - {"samples": FakeNestedTensor((FakeTensor("basev"), FakeTensor("basea")))}, - ) + self.module._build_shot_conditioning = build_conditioning self.module._evict_all_but = lambda *_args, **_kwargs: None self.module.mm.unload_model_and_clones = unload_model_and_clones self.module.mm.unload_all_models = unload_all_models @@ -457,6 +462,7 @@ class DumasH3LongVideosHelperTests(unittest.TestCase): self.assertLess(order.index("upscale"), order.index("latent_upscale_sample")) self.assertLess(order.index("audio"), order.index("video")) self.assertEqual(order[-1], "cleanup") + self.assertEqual(len(build_calls), 1) finally: self.module.nodes.common_ksampler = original_common_ksampler self.module._build_shot_conditioning = original_build