gheinrich commited on
Commit
cd66dcf
·
1 Parent(s): a794ce5

Add generation progress feedback

Browse files

Surface three phases in the UI via gr.Progress: model loading (first run),
live denoising-step bar (track_tqdm with disable_tqdm off), and safety filter.

Files changed (1) hide show
  1. app.py +12 -1
app.py CHANGED
@@ -343,7 +343,10 @@ class T2IEngine:
343
  micro_cond: str,
344
  return_animation: bool,
345
  enable_image_guard: bool,
 
346
  ):
 
 
347
  self._lazy_load()
348
 
349
  gen_cfg = dict(DEFAULT_GENERATION_CONFIG)
@@ -367,6 +370,9 @@ class T2IEngine:
367
 
368
  n_tokens = n_tokens_from_resolution(int(image_resolution))
369
 
 
 
 
370
  t0 = time.time()
371
  with torch.no_grad():
372
  with torch.inference_mode():
@@ -376,7 +382,7 @@ class T2IEngine:
376
  **gen_cfg,
377
  image_resolution=int(image_resolution),
378
  n_tokens=n_tokens,
379
- disable_tqdm=True,
380
  return_intermediate_steps=return_animation,
381
  )
382
  latency = time.time() - t0
@@ -391,6 +397,9 @@ class T2IEngine:
391
  f"edit_threshold={edit_threshold:.3f} | gen_time={latency:.2f}s"
392
  )
393
 
 
 
 
394
  guard_ok, guard_meta = self._moderate_output(
395
  result, return_animation, enable_image_guard
396
  )
@@ -425,6 +434,7 @@ def generate(
425
  micro_cond: str,
426
  return_animation: bool,
427
  enable_image_guard: bool,
 
428
  ):
429
  return engine.generate(
430
  prompt,
@@ -443,6 +453,7 @@ def generate(
443
  micro_cond,
444
  return_animation,
445
  enable_image_guard,
 
446
  )
447
 
448
 
 
343
  micro_cond: str,
344
  return_animation: bool,
345
  enable_image_guard: bool,
346
+ progress: gr.Progress | None = None,
347
  ):
348
+ if progress is not None and self._model is None:
349
+ progress(0.0, desc="Loading model (first run, this can take 1-2 min)…")
350
  self._lazy_load()
351
 
352
  gen_cfg = dict(DEFAULT_GENERATION_CONFIG)
 
370
 
371
  n_tokens = n_tokens_from_resolution(int(image_resolution))
372
 
373
+ if progress is not None:
374
+ progress(0.0, desc="Generating…")
375
+
376
  t0 = time.time()
377
  with torch.no_grad():
378
  with torch.inference_mode():
 
382
  **gen_cfg,
383
  image_resolution=int(image_resolution),
384
  n_tokens=n_tokens,
385
+ disable_tqdm=progress is None,
386
  return_intermediate_steps=return_animation,
387
  )
388
  latency = time.time() - t0
 
397
  f"edit_threshold={edit_threshold:.3f} | gen_time={latency:.2f}s"
398
  )
399
 
400
+ if progress is not None and enable_image_guard:
401
+ progress(1.0, desc="Running safety filter…")
402
+
403
  guard_ok, guard_meta = self._moderate_output(
404
  result, return_animation, enable_image_guard
405
  )
 
434
  micro_cond: str,
435
  return_animation: bool,
436
  enable_image_guard: bool,
437
+ progress: gr.Progress = gr.Progress(track_tqdm=True),
438
  ):
439
  return engine.generate(
440
  prompt,
 
453
  micro_cond,
454
  return_animation,
455
  enable_image_guard,
456
+ progress=progress,
457
  )
458
 
459