Spaces:
Running on Zero
Running on Zero
Add generation progress feedback
Browse filesSurface 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.
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=
|
| 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 |
|