feat(offload): on-demand vae under group offload

Vae-class components never take group hooks, so group mode kept them
resident on the gpu; a MiniMax-class video vae holds about 10GB that
way while running only seconds per generation. Components above 1GB
now rest in system memory: the apply_forward_hook bridge on encode and
decode fires an on-demand hook that moves the whole module to the
device, so tiled calls find every weight already loaded, and the
processing seams return it to cpu once outputs are materialized. Small
vaes stay resident since the transfer would cost more than it frees.

- placement is decided per component by measured size and requires the
  entry bridge; components without it stay resident
- move_model no longer forces on-demand vaes to the gpu for
  non-txt2img tasks, and full_vae_encode onloads before binding the
  input, which otherwise lands on the resting device
- mode switches clear the stamp and hook in both directions
This commit is contained in:
CalamitousFelicitousness
2026-08-07 23:44:18 +01:00
parent b84b782ba4
commit abfb5ac3ed
5 changed files with 96 additions and 9 deletions
+4
View File
@@ -189,6 +189,8 @@ def full_vae_encode(image, model):
sd_models.move_model(model.unet, devices.cpu)
if shared.opts.diffusers_offload_mode != "sequential" and hasattr(model, 'vae'):
sd_models.move_model(model.vae, devices.device)
if getattr(model.vae, 'sdnext_ondemand', False):
model.vae.to(devices.device) # the image placement below derives from vae.device, and the entry bridge would onload the weights only after the input is already bound
vae_name = sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "default"
log_debug(f'Encode vae="{vae_name}" dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)}')
@@ -369,6 +371,7 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he
if shared.cmd_opts.profile or debug:
t1 = time.time()
log.debug(f'Profile: VAE decode: {t1-t0:.2f}')
sd_models.offload_ondemand(model)
devices.torch_gc()
shared.state.end(jobid)
return images
@@ -393,6 +396,7 @@ def vae_encode(image, model, vae_type='Full'): # pylint: disable=unused-variable
else:
log.error('VAE not found in model')
latents = []
sd_models.offload_ondemand(model)
devices.torch_gc()
shared.state.end(jobid)
return latents