Touch ups

app.py

@@ -1,15 +1,12 @@
 """
 """
-from datetime import datetime
-t0 = datetime.now()
-
 # Upgrade PyTorch
 import os
 os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')
-print('torch upgrade', -(t0 - (t0 := datetime.now())))
 
 # Actual app.py
 import os
+from datetime import datetime
 
 import gradio as gr
 import spaces
@@ -22,7 +19,6 @@ from zerogpu import aoti_compile
 
 
 pipeline = FluxPipeline.from_pretrained('black-forest-labs/FLUX.1-schnell', torch_dtype=torch.bfloat16).to('cuda')
-print('FluxPipeline.from_pretrained', -(t0 - (t0 := datetime.now())))
 
 
 @spaces.GPU(duration=1500)
@@ -70,17 +66,16 @@ transformer_config = pipeline.transformer.config
 pipeline.transformer = compile_transformer()
 pipeline.transformer.config = transformer_config
 
+
 @spaces.GPU
-def _generate_image(prompt: str, t0: datetime):
-    print('@spaces.GPU', -(t0 - (t0 := datetime.now())))
+def generate_image(prompt: str, progress=gr.Progress(track_tqdm=True)):
+    t0 = datetime.now()
     images = []
     for _ in range(4):
-        images += pipeline(prompt, num_inference_steps=4).images
-        print('pipeline', -(t0 - (t0 := datetime.now())))
+        image = pipeline(prompt, num_inference_steps=4).images[0]
+        elapsed = -(t0 - (t0 := datetime.now()))
+        images += [(image, f'{elapsed.total_seconds():.2f}s')]
     return images
 
 
-def generate_image(prompt: str, progress=gr.Progress(track_tqdm=True)):
-    return _generate_image(prompt, datetime.now())
-
-gr.Interface(generate_image, gr.Text(), gr.Gallery()).launch(show_error=True)
+gr.Interface(generate_image, gr.Text(), gr.Gallery(rows=3, columns=3, height='60vh')).launch()