Reuse conditioning for latent upscale refine

dumas_h3_longvideos.py

@@ -4040,6 +4040,42 @@ def _copy_sample_latent(out_latent):
         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):
@@ -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")
+                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))
+                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)

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
 
-            self.module.nodes.common_ksampler = common_ksampler
-            self.module._build_shot_conditioning = lambda *_args, **_kwargs: (
-                "cond",
+            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 = 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