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.
This commit is contained in:
parent
a8a4d4fe0e
commit
613afa9697
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue