Improve generate timing

app.py

@@ -1,5 +1,3 @@
-import time
-
 import gradio as gr
 
 from generate import generate
@@ -23,24 +21,14 @@ def read_file(path: str) -> str:
 
 
 # don't request a GPU if input is bad
-def generate_btn_click(*args, **kwargs):
-    start = time.perf_counter()
-
-    if "prompt" in kwargs:
-        prompt = kwargs.get("prompt")
-    elif len(args) > 0:
+def generate_btn_click(*args):
+    if len(args) > 0:
         prompt = args[0]
     else:
         prompt = None
-
     if prompt is None or prompt.strip() == "":
         raise gr.Error("You must enter a prompt")
-
-    images = generate(*args, **kwargs, Error=gr.Error)
-    end = time.perf_counter()
-    diff = end - start
-    gr.Info(f"Generated {len(images)} images in {diff:.2f}s")
-    return images
+    return generate(*args)
 
 
 with gr.Blocks(

generate.py

@@ -1,4 +1,5 @@
 import re
+import time
 from contextlib import contextmanager
 from datetime import datetime
 from itertools import product
@@ -6,6 +7,7 @@ from os import environ
 from types import MethodType
 from warnings import filterwarnings
 
+import gradio as gr
 import spaces
 import tomesd
 import torch
@@ -242,10 +244,10 @@ def generate(
     deepcache_interval=1,
     tgate_step=0,
     tome_ratio=0,
-    Error=Exception,
+    progress=gr.Progress(track_tqdm=True),
 ):
     if not torch.cuda.is_available():
-        raise Error("CUDA not available")
+        raise gr.Error("CUDA not available")
 
     if seed is None:
         seed = int(datetime.now().timestamp())
@@ -263,6 +265,7 @@ def generate(
     )
 
     with torch.inference_mode():
+        start = time.perf_counter()
         loader = Loader()
         pipe = loader.load(model, scheduler, karras, taesd, deepcache_interval, TORCH_DTYPE)
 
@@ -319,4 +322,7 @@ def generate(
             # spaces always start fresh
             loader.pipe = None
 
+        end = time.perf_counter()
+        diff = end - start
+        gr.Info(f"Generated {len(images)} image{'s' if len(images) > 1 else ''} in {diff:.2f}s")
         return images