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
+1
View File
@@ -477,6 +477,7 @@ def process_decode(p: processing.StableDiffusionProcessing, output):
log.debug(f'Generated: frames={len(output.frames[0])}')
output.images = output.frames[0]
if output.images is not None and len(output.images) > 0 and isinstance(output.images[0], Image.Image):
sd_models.offload_ondemand(shared.sd_model) # in-pipe decode paths return materialized frames; the vae seam in processing_vae never runs
return attach_audio(output.images, audio)
model = shared.sd_model if not is_refiner_enabled(p) else shared.sd_refiner
if not hasattr(model, 'vae'):