Spaces:
Running on Zero
Running on Zero
Commit ·
0a2678b
1
Parent(s): c0f3a61
Free text encoder before VAE decode and enable VAE tiling to fix OOM
Browse filesFast-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>
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")
|