someone-in-the-world Claude Sonnet 4.6 commited on
Commit
a1fde2a
·
1 Parent(s): 2f7a843

Add VAE decode logging: pre-decode mem snapshot, tiling prediction, peak reset

Browse files

Adds three log points to support issue #6 investigation:
- [startup] VAE tile threshold confirmation after enable_tiling()
- [infer] whether tiling will activate for this image size
- [infer] pre-VAE-decode memory snapshot + reset_peak_memory_stats() so
the existing "VAE decode + postprocess done" peak reflects only VAE decode

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

Files changed (1) hide show
  1. app.py +12 -2
app.py CHANGED
@@ -192,6 +192,7 @@ pipe = QwenImageEditPlusPipeline.from_pretrained(
192
  )
193
  _hb.set()
194
  pipe.vae.enable_tiling(tile_sample_min_height=Mode.HIGH_DETAIL.max_dim, tile_sample_min_width=Mode.HIGH_DETAIL.max_dim)
 
195
  print(f"[startup] pipeline loaded in {time.perf_counter()-_t1_load:.1f}s", flush=True)
196
 
197
  print("[startup] setting cuDNN SDPA attention processor...", flush=True)
@@ -565,7 +566,7 @@ def _instrument_first_touch(modules_with_names, t0):
565
  handle_box["h"] = module.register_forward_pre_hook(_make_hook(name, handle_box))
566
 
567
 
568
- def _make_step_callback(steps, timer, t0, mode: Mode):
569
  """Build the diffusers step callback that logs per-step timing and marks timer checkpoints."""
570
  step_times = []
571
  def _step_cb(pipeline, step_idx, timestep, cb_kwargs):
@@ -592,6 +593,9 @@ def _make_step_callback(steps, timer, t0, mode: Mode):
592
  print("[infer] skipping text_encoder offload (int8, .to() unsupported / smaller footprint)")
593
  else:
594
  print(f"[infer] skipping text_encoder offload for mode={mode.value} (ample headroom at this resolution)")
 
 
 
595
  delta_ms = (now - (step_times[-2] if len(step_times) > 1 else t0)) * 1000
596
  tag = " ← includes cold-start (offload hook install + first weight transfer)" if step_idx == 0 else ""
597
  print(f"[infer] step {step_idx+1}/{steps} done — {delta_ms:.0f}ms{tag} | t={now-t0:.1f}s")
@@ -629,9 +633,15 @@ def _infer_gpu(pil_images, prompt, seed, guidance_scale, steps, width, height, m
629
  )
630
 
631
  print(f"[infer] {len(pil_images)} image(s) pre-decoded, output={width}x{height}, seed={seed}")
 
 
 
 
 
 
632
 
633
  generator = torch.Generator(device=device).manual_seed(seed)
634
- step_cb = _make_step_callback(steps, timer, t0, mode)
635
 
636
  timer.mark("pipe_start")
637
  print(f"[infer] calling pipe... t={time.perf_counter()-t0:.1f}s")
 
192
  )
193
  _hb.set()
194
  pipe.vae.enable_tiling(tile_sample_min_height=Mode.HIGH_DETAIL.max_dim, tile_sample_min_width=Mode.HIGH_DETAIL.max_dim)
195
+ print(f"[startup] VAE tiling: threshold={pipe.vae.tile_sample_min_height}x{pipe.vae.tile_sample_min_width}px use_tiling={pipe.vae.use_tiling}", flush=True)
196
  print(f"[startup] pipeline loaded in {time.perf_counter()-_t1_load:.1f}s", flush=True)
197
 
198
  print("[startup] setting cuDNN SDPA attention processor...", flush=True)
 
566
  handle_box["h"] = module.register_forward_pre_hook(_make_hook(name, handle_box))
567
 
568
 
569
+ def _make_step_callback(steps, timer, t0, mode: Mode, cuda_ok: bool = False):
570
  """Build the diffusers step callback that logs per-step timing and marks timer checkpoints."""
571
  step_times = []
572
  def _step_cb(pipeline, step_idx, timestep, cb_kwargs):
 
593
  print("[infer] skipping text_encoder offload (int8, .to() unsupported / smaller footprint)")
594
  else:
595
  print(f"[infer] skipping text_encoder offload for mode={mode.value} (ample headroom at this resolution)")
596
+ if step_idx == steps - 1 and cuda_ok:
597
+ print(f"[infer] pre-VAE-decode — {_gpu_mem_str(True, sync=True)} | t={time.perf_counter()-t0:.1f}s")
598
+ torch.cuda.reset_peak_memory_stats()
599
  delta_ms = (now - (step_times[-2] if len(step_times) > 1 else t0)) * 1000
600
  tag = " ← includes cold-start (offload hook install + first weight transfer)" if step_idx == 0 else ""
601
  print(f"[infer] step {step_idx+1}/{steps} done — {delta_ms:.0f}ms{tag} | t={now-t0:.1f}s")
 
633
  )
634
 
635
  print(f"[infer] {len(pil_images)} image(s) pre-decoded, output={width}x{height}, seed={seed}")
636
+ if _cuda_ok:
637
+ _will_tile = pipe.vae.use_tiling and (
638
+ width > pipe.vae.tile_sample_min_width or height > pipe.vae.tile_sample_min_height
639
+ )
640
+ print(f"[infer] VAE tiling will {'activate' if _will_tile else 'NOT activate'} "
641
+ f"(threshold={pipe.vae.tile_sample_min_height}x{pipe.vae.tile_sample_min_width}px)")
642
 
643
  generator = torch.Generator(device=device).manual_seed(seed)
644
+ step_cb = _make_step_callback(steps, timer, t0, mode, _cuda_ok)
645
 
646
  timer.mark("pipe_start")
647
  print(f"[infer] calling pipe... t={time.perf_counter()-t0:.1f}s")