From a8a4d4fe0ef840070253ed4b939f5343b852b55b Mon Sep 17 00:00:00 2001 From: chelsealong Date: Mon, 10 Aug 2026 01:10:11 +0000 Subject: [PATCH 1/2] 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) From 613afa9697a506b84901e54233ba3cd64c659098 Mon Sep 17 00:00:00 2001 From: chelsealong Date: Mon, 10 Aug 2026 07:24:39 +0000 Subject: [PATCH 2/2] Free other models' memory before retrying handles_tiling VAE encode Mirrors the decode fix: encode_tiled is identical to encode for handles_tiling VAEs, so the OOM-retry path needs to free other models' memory the same way decode's retry path already does. --- comfy/sd.py | 3 ++ .../comfy_test/test_seedvr2_vae_tiled.py | 28 +++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/comfy/sd.py b/comfy/sd.py index 52be28137..3bed1b844 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -1313,6 +1313,9 @@ class VAE: tile = 256 overlap = tile // 4 if self.handles_tiling: + # encode_tiled is identical to encode 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)]) samples = self._encode_tiled_owned(pixel_samples, tile_x=tile, tile_y=tile, overlap=overlap) else: samples = self.encode_tiled_3d(pixel_samples, 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 26141c644..d261bc912 100644 --- a/tests-unit/comfy_test/test_seedvr2_vae_tiled.py +++ b/tests-unit/comfy_test/test_seedvr2_vae_tiled.py @@ -421,6 +421,34 @@ def test_handles_tiling_decode_oom_frees_memory_before_retry(): assert free_memory_call.call_args.kwargs["keep_loaded"][0].model is vae.patcher +def test_handles_tiling_encode_oom_frees_memory_before_retry(): + vae = _make_seedvr2_vae_fallback() + pixel_samples = torch.zeros((1, 8, 64, 64, 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("encode_tiled_owned") or torch.zeros(1, _LATENT_CHANNELS, 2, 8, 8)) + + 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, "_encode_tiled_owned", seedvr2_call)) + stack.enter_context(patch.object( + seedvr_vae_mod.VideoAutoencoderKLWrapper, "encode", + side_effect=_force_regular_encode_oom)) + vae.encode(pixel_samples) + + assert call_order == ["free_memory", "encode_tiled_owned"], ( + "encode_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)