someone-in-the-world Claude Sonnet 5 commited on
Commit
0a2678b
·
1 Parent(s): c0f3a61

Free text encoder before VAE decode and enable VAE tiling to fix OOM

Browse files

Fast-path inference keeps the full ~35GB pipeline resident on the
47.4GB MIG slice, leaving too little headroom for the VAE decoder's
fp32-upcast memory spike on large (2048px) outputs. Offload the
~15GB text encoder to CPU once the last denoising step finishes
(skipped when accelerate hooks already own placement), and bound
decode memory with tiling regardless of output size.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Files changed (1) hide show
  1. app.py +8 -0
app.py CHANGED
@@ -190,6 +190,7 @@ pipe = QwenImageEditPlusPipeline.from_pretrained(
190
  torch_dtype=dtype,
191
  )
192
  _hb.set()
 
193
  print(f"[startup] pipeline loaded in {time.perf_counter()-_t1_load:.1f}s", flush=True)
194
 
195
  print("[startup] using default attention processor.", flush=True)
@@ -531,6 +532,13 @@ def _make_step_callback(steps, timer, t0):
531
  if step_idx == 0:
532
  timer.mark("first_step")
533
  timer.mark("last_step") # overwritten each step; final value = end of last step
 
 
 
 
 
 
 
534
  delta_ms = (now - (step_times[-2] if len(step_times) > 1 else t0)) * 1000
535
  tag = " ← includes cold-start (offload hook install + first weight transfer)" if step_idx == 0 else ""
536
  print(f"[infer] step {step_idx+1}/{steps} done — {delta_ms:.0f}ms{tag} | t={now-t0:.1f}s")
 
190
  torch_dtype=dtype,
191
  )
192
  _hb.set()
193
+ pipe.vae.enable_tiling()
194
  print(f"[startup] pipeline loaded in {time.perf_counter()-_t1_load:.1f}s", flush=True)
195
 
196
  print("[startup] using default attention processor.", flush=True)
 
532
  if step_idx == 0:
533
  timer.mark("first_step")
534
  timer.mark("last_step") # overwritten each step; final value = end of last step
535
+ # Text encoder is done after prompt encoding, before the denoising loop starts.
536
+ # Drop it (~15GB) ahead of VAE decode's fp32-upcast memory spike. Skipped when
537
+ # accelerate hooks are managing placement (offload-fallback path) to avoid
538
+ # fighting their own device bookkeeping — see the finally block below.
539
+ if step_idx == steps - 1 and getattr(pipeline.text_encoder, "_hf_hook", None) is None:
540
+ pipeline.text_encoder.to("cpu")
541
+ torch.cuda.empty_cache()
542
  delta_ms = (now - (step_times[-2] if len(step_times) > 1 else t0)) * 1000
543
  tag = " ← includes cold-start (offload hook install + first weight transfer)" if step_idx == 0 else ""
544
  print(f"[infer] step {step_idx+1}/{steps} done — {delta_ms:.0f}ms{tag} | t={now-t0:.1f}s")