Reuse conditioning for latent upscale refine

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)