Update app.py

app.py

@@ -27,6 +27,7 @@ MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
 OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs")
+HISTORY_SECRET = os.getenv("HISTORY_SECRET", "default_secret")
 
 MODEL = os.getenv(
     "MODEL",
@@ -37,7 +38,7 @@ DESCRIPTION = '''
 <div>
 <h1 style="text-align: center;">High Definition Pony Diffusion</h1>
 <p>Gradio demo for PonyDiffusion v6 with image gallery, json prompt support, advanced options and more.</p>
-<p>❤️ Thanks for ✨4000 visits! Heart this space if you like it!</p>
+<p>❤️ Thanks for ✨5000 visits! Heart this space if you like it!</p>
 <p>🔎 For more details about me, take a look at <a href="https://sergidev.me">My website</a>.</p>
 <p>🌚 For dark mode compatibility, click <a href="https://sergidev.me/hdiffusion">here</a>.</p>
 </div>
@@ -114,6 +115,7 @@ def generate(
     upscaler_strength: float = 0.55,
     upscale_by: float = 1.5,
     json_params: str = "",
+    batch_size: int = 1,
     progress=gr.Progress(track_tqdm=True),
 ) -> Image:
     if json_params:
@@ -154,6 +156,7 @@ def generate(
         "num_inference_steps": num_inference_steps,
         "seed": seed,
         "sampler": sampler,
+        "batch_size": batch_size,
     }
 
     if use_upscaler:
@@ -170,46 +173,50 @@ def generate(
     logger.info(json.dumps(metadata, indent=4))
 
     try:
-        if use_upscaler:
-            latents = pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                width=width,
-                height=height,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                generator=generator,
-                output_type="latent",
-            ).images
-            upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
-            images = upscaler_pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                image=upscaled_latents,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                strength=upscaler_strength,
-                generator=generator,
-                output_type="pil",
-            ).images
-        else:
-            images = pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                width=width,
-                height=height,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                generator=generator,
-                output_type="pil",
-            ).images
-
-        if images and IS_COLAB:
-            for image in images:
+        all_images = []
+        for _ in range(batch_size):
+            batch_generator = utils.seed_everything(random.randint(0, utils.MAX_SEED))
+            if use_upscaler:
+                latents = pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    width=width,
+                    height=height,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    generator=batch_generator,
+                    output_type="latent",
+                ).images
+                upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
+                images = upscaler_pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    image=upscaled_latents,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    strength=upscaler_strength,
+                    generator=batch_generator,
+                    output_type="pil",
+                ).images
+            else:
+                images = pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    width=width,
+                    height=height,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    generator=batch_generator,
+                    output_type="pil",
+                ).images
+            all_images.extend(images)
+
+        if all_images and IS_COLAB:
+            for image in all_images:
                 filepath = utils.save_image(image, metadata, OUTPUT_DIR)
                 logger.info(f"Image saved as {filepath} with metadata")
 
-        return images, metadata
+        return all_images, metadata
     except Exception as e:
         logger.exception(f"An error occurred: {e}")
         raise
@@ -229,16 +236,18 @@ def handle_image_click(evt: gr.SelectData):
     return selected["image"], json.dumps(selected["metadata"], indent=2)
 
 def generate_and_update_history(*args, **kwargs):
+    global generation_history
     images, metadata = generate(*args, **kwargs)
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    generation_history.insert(0, {
-        "prompt": metadata["prompt"],
-        "timestamp": timestamp,
-        "image": images[0],
-        "metadata": metadata
-    })
+    for image in images:
+        generation_history.insert(0, {
+            "prompt": metadata["prompt"],
+            "timestamp": timestamp,
+            "image": image,
+            "metadata": metadata
+        })
     if len(generation_history) > 20:
-        generation_history.pop()
+        generation_history = generation_history[:20]
     return images[0], json.dumps(metadata, indent=2), update_history_list()
 
 with open('characterfull.txt', 'r') as f:
@@ -262,6 +271,12 @@ if torch.cuda.is_available():
 else:
     pipe = None
 
+def check_history_password(password):
+    if password == HISTORY_SECRET:
+        return gr.update(visible=True)
+    else:
+        return gr.update(visible=False)
+
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown(DESCRIPTION)
 
@@ -365,6 +380,13 @@ with gr.Blocks(css="style.css") as demo:
                     step=1,
                     value=28,
                 )
+        batch_size = gr.Slider(
+            label="Batch Size",
+            minimum=1,
+            maximum=4,
+            step=1,
+            value=1,
+        )
                 
     with gr.Accordion(label="Generation Parameters", open=False):
         gr_metadata = gr.JSON(label="Metadata", show_label=False)
@@ -372,17 +394,25 @@ with gr.Blocks(css="style.css") as demo:
         generate_from_json = gr.Button("Generate from JSON")
 
     with gr.Accordion("Generation History", open=False) as history_accordion:
-        history_gallery = gr.Gallery(
-            label="History",
-            show_label=False,
-            elem_id="history_gallery",
-            columns=5,
-            rows=2,
-            height="auto"
+        history_password = gr.Textbox(
+            label="Global generation history",
+            type="password",
+            placeholder="Enter secret for generation history"
         )
-        with gr.Row():
-            selected_image = gr.Image(label="Selected Image", interactive=False)
-            selected_metadata = gr.JSON(label="Selected Metadata", show_label=False)
+        history_submit = gr.Button("Submit")
+        
+        with gr.Group(visible=False) as history_content:
+            history_gallery = gr.Gallery(
+                label="History",
+                show_label=False,
+                elem_id="history_gallery",
+                columns=5,
+                rows=2,
+                height="auto"
+            )
+            with gr.Row():
+                selected_image = gr.Image(label="Selected Image", interactive=False)
+                selected_metadata = gr.JSON(label="Selected Metadata", show_label=False)
 
     gr.Examples(
         examples=config.examples,
@@ -422,6 +452,7 @@ with gr.Blocks(css="style.css") as demo:
         upscaler_strength,
         upscale_by,
         json_input,
+        batch_size,
     ]
 
     prompt.submit(
@@ -483,5 +514,11 @@ with gr.Blocks(css="style.css") as demo:
         inputs=[],
         outputs=[selected_image, selected_metadata]
     )
+
+    history_submit.click(
+        fn=check_history_password,
+        inputs=[history_password],
+        outputs=[history_content],
+    )
     
 demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)