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)