From 613afa9697a506b84901e54233ba3cd64c659098 Mon Sep 17 00:00:00 2001 From: chelsealong Date: Mon, 10 Aug 2026 07:24:39 +0000 Subject: [PATCH] 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)