Update app.py

app.py

@@ -7,6 +7,8 @@ import random
 import uuid
 from typing import Tuple
 import numpy as np
+import time
+import zipfile
 
 DESCRIPTION = """## 
 """
@@ -82,6 +84,8 @@ def generate(
     randomize_seed: bool = False,
     style_name: str = DEFAULT_STYLE_NAME,
     num_inference_steps: int = 30,
+    num_images: int = 1,
+    zip_images: bool = False,
     progress=gr.Progress(track_tqdm=True),
 ):
     positive_prompt, style_negative_prompt = apply_style(style_name, prompt)
@@ -99,6 +103,8 @@ def generate(
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device="cuda").manual_seed(seed)
     
+    start_time = time.time()
+    
     images = pipe(
         prompt=positive_prompt,
         negative_prompt=final_negative_prompt if final_negative_prompt else None,
@@ -106,12 +112,25 @@ def generate(
         height=height,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
-        num_images_per_prompt=1,
+        num_images_per_prompt=num_images,
         generator=generator,
         output_type="pil",
     ).images
+    
+    end_time = time.time()
+    duration = end_time - start_time
+    
     image_paths = [save_image(img) for img in images]
-    return image_paths, seed
+    
+    zip_path = None
+    if zip_images:
+        zip_name = str(uuid.uuid4()) + ".zip"
+        with zipfile.ZipFile(zip_name, 'w') as zipf:
+            for i, img_path in enumerate(image_paths):
+                zipf.write(img_path, arcname=f"Img_{i}.png")
+        zip_path = zip_name
+    
+    return image_paths, seed, f"{duration:.2f}", zip_path
 
 # Example prompts
 examples = [
@@ -148,6 +167,9 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         )
         run_button = gr.Button("Run", scale=0, variant="primary")
     result = gr.Gallery(label="Result", columns=1, show_label=False, preview=True)
+    seed_display = gr.Textbox(label="Seed used", interactive=False)
+    generation_time = gr.Textbox(label="Generation time (seconds)", interactive=False)
+    zip_file = gr.File(label="Download ZIP")
     
     with gr.Accordion("Advanced options", open=False):
         style_selection = gr.Dropdown(
@@ -200,11 +222,19 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             step=1,
             value=30,
         )
+        num_images = gr.Slider(
+            label="Number of images",
+            minimum=1,
+            maximum=101,
+            step=1,
+            value=1,
+        )
+        zip_images = gr.Checkbox(label="Zip generated images", value=False)
     
     gr.Examples(
         examples=examples,
         inputs=prompt,
-        outputs=[result, seed],
+        outputs=[result, seed_display, generation_time, zip_file],
         fn=generate,
         cache_examples=False,
     )
@@ -235,8 +265,10 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             randomize_seed,
             style_selection,
             num_inference_steps,
+            num_images,
+            zip_images,
         ],
-        outputs=[result, seed],
+        outputs=[result, seed_display, generation_time, zip_file],
         api_name="run",
     )