From a8a4d4fe0ef840070253ed4b939f5343b852b55b Mon Sep 17 00:00:00 2001 From: chelsealong Date: Mon, 10 Aug 2026 01:10:11 +0000 Subject: [PATCH] Free other models' memory before retrying handles_tiling VAE decode For VAEs that set handles_tiling (MiniMax H3 video, SeedVR2), the tiled decode fallback ends up calling the exact same decode path as the regular decode that just OOM'd, so retrying without changing anything about available memory just fails identically. Free memory held by other loaded models (including dynamically streamed weights) before the retry, keeping only the VAE itself loaded, so the retry actually has a chance to succeed. --- comfy/sd.py | 6 ++++ .../comfy_test/test_seedvr2_vae_tiled.py | 29 +++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/comfy/sd.py b/comfy/sd.py index 5fed4ca9a..52be28137 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -1215,6 +1215,9 @@ class VAE: if self.handles_tiling: tile = 256 // self.spacial_compression_decode() overlap = tile // 4 + # decode_tiled is identical to decode for these VAEs, so freeing the + # memory other models hold onto is the only thing that can make the retry succeed. + model_management.free_memory(1e30, self.device, keep_loaded=[model_management.LoadedModel(self.patcher)]) pixel_samples = self._decode_tiled_owned(samples_in, tile_x=tile, tile_y=tile, overlap=overlap) else: pixel_samples = self.decode_tiled_(samples_in) @@ -1222,6 +1225,9 @@ class VAE: tile = 256 // self.spacial_compression_decode() overlap = tile // 4 if self.handles_tiling: + # decode_tiled is identical to decode for these VAEs, so freeing the + # memory other models hold onto is the only thing that can make the retry succeed. + model_management.free_memory(1e30, self.device, keep_loaded=[model_management.LoadedModel(self.patcher)]) pixel_samples = self._decode_tiled_owned(samples_in, tile_x=tile, tile_y=tile, overlap=overlap) else: pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) diff --git a/tests-unit/comfy_test/test_seedvr2_vae_tiled.py b/tests-unit/comfy_test/test_seedvr2_vae_tiled.py index d64f51918..26141c644 100644 --- a/tests-unit/comfy_test/test_seedvr2_vae_tiled.py +++ b/tests-unit/comfy_test/test_seedvr2_vae_tiled.py @@ -392,6 +392,35 @@ def test_seedvr2_3d_routes_to_owned_encode_tiled_on_oom(): ) +def test_handles_tiling_decode_oom_frees_memory_before_retry(): + wrapper = seedvr_vae_mod.VideoAutoencoderKLWrapper.__new__( + seedvr_vae_mod.VideoAutoencoderKLWrapper) + vae = _make_vae(wrapper, latent_channels=_LATENT_CHANNELS, latent_dim=3) + + call_order = [] + free_memory_call = MagicMock(side_effect=lambda *a, **k: call_order.append("free_memory")) + seedvr2_call = MagicMock( + side_effect=lambda *a, **k: call_order.append("decode_tiled_owned") or torch.zeros(1, 3, 9, 64, 64)) + + mm = sd_mod.model_management + with ExitStack() as stack: + stack.enter_context(patch.object(mm, "raise_non_oom", lambda e: None)) + stack.enter_context(patch.object(mm, "load_models_gpu", lambda *a, **k: None)) + stack.enter_context(patch.object(mm, "soft_empty_cache", lambda: None)) + stack.enter_context(patch.object(mm, "free_memory", free_memory_call)) + stack.enter_context(patch.object(sd_mod.VAE, "_decode_tiled_owned", seedvr2_call)) + stack.enter_context(patch.object( + seedvr_vae_mod.VideoAutoencoderKLWrapper, "decode", + side_effect=_force_oom)) + vae.decode(torch.zeros(1, _LATENT_CHANNELS * 3, 8, 8)) + + assert call_order == ["free_memory", "decode_tiled_owned"], ( + "decode_tiled retry for a handles_tiling VAE must free other models' " + f"memory before retrying; call order was {call_order}" + ) + assert free_memory_call.call_args.kwargs["keep_loaded"][0].model is vae.patcher + + def test_non_seedvr2_encode_tiled_3d_default_overlap_is_concrete(): vae = _make_non_seedvr2_vae_fallback() vae.downscale_ratio = (lambda a: max(1, a // 4), 8, 8)