Free first-pass latent before refinement

dumas_h3_longvideos.py

@@ -6648,6 +6648,13 @@ class H3LongVideos:
                 if not getattr(out_samples, "is_nested", False) or len(parts) < 2:
                     raise RuntimeError("latent upscale expects a nested AV latent")
                 upscaled_video, up_h, up_w = _upscale_latent_video(parts[0], latent_upscale_param)
+                full_audio = parts[1]
+                # Drop the first-pass sampling state before we start the refinement
+                # 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
+                mm.soft_empty_cache()
                 target_w = int(up_w) * 16
                 target_h = int(up_h) * 16
                 if target_w <= 0 or target_h <= 0:
@@ -6660,7 +6667,7 @@ class H3LongVideos:
                     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, parts[1]))
+                    (upscaled_video, full_audio))
                 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)
@@ -6733,7 +6740,6 @@ class H3LongVideos:
                         denoise=refine_denoise)
                 else:
                     refined_video = upscaled_video.clone()
-                    full_audio = parts[1]
                     for row_index, r0 in enumerate(rows):
                         tr = trows[row_index]
                         ovh = row_ovl[row_index]
@@ -6811,7 +6817,7 @@ class H3LongVideos:
                                 )
                     refined_out = {"samples": comfy.nested_tensor.NestedTensor((refined_video, full_audio))}
                 timing["latent_upscale_sample"] += time.perf_counter() - latent_start
-                refined_out = _video_only_refined_latent(out, refined_out)
+                refined_out = _video_only_refined_latent(upscale_latent, refined_out)
             except Exception as e:
                 if _is_oom(e):
                     e._h3_stage = "sampling"
@@ -6821,19 +6827,16 @@ class H3LongVideos:
         # 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.
-        decode_video_start = time.perf_counter()
-        shot_latent = _copy_sample_latent(refined_out)
         decode_audio_start = time.perf_counter()
         audio = _decode_audio(audio_vae, out)
         timing["decode_audio"] += time.perf_counter() - decode_audio_start
-        # Audio is much smaller than the video decode. Drop the first-pass
-        # conditioning before the VAE work so the optional latent upscale pass does
-        # not keep both sampled latents resident across the heaviest allocation.
-        del out, positive, latent
+        decode_video_start = time.perf_counter()
+        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)
         timing["decode_video"] += time.perf_counter() - decode_video_start
         cleanup_start = time.perf_counter()
+        del out
         del refined_out
         _deep_cleanup()
         timing["cleanup"] += time.perf_counter() - cleanup_start