Update

app.py

@@ -3,8 +3,10 @@
 from __future__ import annotations
 
 import os
+import random
 
 import gradio as gr
+import numpy as np
 import PIL.Image
 import torch
 from diffusers import StableDiffusionAttendAndExcitePipeline, StableDiffusionPipeline
@@ -29,6 +31,16 @@ if torch.cuda.is_available():
     sd_pipe.to(device)
 
 
+MAX_INFERENCE_STEPS = 100
+MAX_SEED = np.iinfo(np.int32).max
+
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+
 def get_token_table(prompt: str) -> list[tuple[int, str]]:
     tokens = [ax_pipe.tokenizer.decode(t) for t in ax_pipe.tokenizer(prompt)["input_ids"]]
     tokens = tokens[1:-1]
@@ -40,7 +52,7 @@ def run(
     indices_to_alter_str: str,
     seed: int = 0,
     apply_attend_and_excite: bool = True,
-    num_steps: int = 50,
+    num_inference_steps: int = 50,
     guidance_scale: float = 7.5,
     scale_factor: int = 20,
     thresholds: dict[int, float] = {
@@ -48,9 +60,12 @@ def run(
         20: 0.8,
     },
     max_iter_to_alter: int = 25,
+    progress=gr.Progress(track_tqdm=True),
 ) -> PIL.Image.Image:
-    generator = torch.Generator(device=device).manual_seed(seed)
+    if num_inference_steps > MAX_INFERENCE_STEPS:
+        raise gr.Error(f"Number of steps cannot exceed {MAX_INFERENCE_STEPS}.")
 
+    generator = torch.Generator(device=device).manual_seed(seed)
     if apply_attend_and_excite:
         try:
             token_indices = list(map(int, indices_to_alter_str.split(",")))
@@ -61,7 +76,7 @@ def run(
             token_indices=token_indices,
             guidance_scale=guidance_scale,
             generator=generator,
-            num_inference_steps=num_steps,
+            num_inference_steps=num_inference_steps,
             max_iter_to_alter=max_iter_to_alter,
             thresholds=thresholds,
             scale_factor=scale_factor,
@@ -71,7 +86,7 @@ def run(
             prompt=prompt,
             guidance_scale=guidance_scale,
             generator=generator,
-            num_inference_steps=num_steps,
+            num_inference_steps=num_inference_steps,
         )
     return out.images[0]
 
@@ -115,23 +130,24 @@ with gr.Blocks(css="style.css") as demo:
                 max_lines=1,
                 placeholder="4,16",
             )
+            apply_attend_and_excite = gr.Checkbox(label="Apply Attend-and-Excite", value=True)
             seed = gr.Slider(
                 label="Seed",
                 minimum=0,
-                maximum=100000,
+                maximum=MAX_SEED,
                 step=1,
                 value=0,
             )
-            apply_attend_and_excite = gr.Checkbox(label="Apply Attend-and-Excite", value=True)
-            num_steps = gr.Slider(
-                label="Number of steps",
-                minimum=0,
-                maximum=100,
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            num_inference_steps = gr.Slider(
+                label="Number of inference steps",
+                minimum=1,
+                maximum=MAX_INFERENCE_STEPS,
                 step=1,
                 value=50,
             )
             guidance_scale = gr.Slider(
-                label="CFG scale",
+                label="Guidance scale",
                 minimum=0,
                 maximum=50,
                 step=0.1,
@@ -246,10 +262,16 @@ with gr.Blocks(css="style.css") as demo:
         token_indices_str,
         seed,
         apply_attend_and_excite,
-        num_steps,
+        num_inference_steps,
         guidance_scale,
     ]
     prompt.submit(
+        fn=randomize_seed_fn,
+        inputs=[seed, randomize_seed],
+        outputs=seed,
+        queue=False,
+        api_name=False,
+    ).then(
         fn=get_token_table,
         inputs=prompt,
         outputs=token_indices_table,
@@ -262,6 +284,12 @@ with gr.Blocks(css="style.css") as demo:
         api_name=False,
     )
     token_indices_str.submit(
+        fn=randomize_seed_fn,
+        inputs=[seed, randomize_seed],
+        outputs=seed,
+        queue=False,
+        api_name=False,
+    ).then(
         fn=get_token_table,
         inputs=prompt,
         outputs=token_indices_table,
@@ -274,6 +302,12 @@ with gr.Blocks(css="style.css") as demo:
         api_name=False,
     )
     run_button.click(
+        fn=randomize_seed_fn,
+        inputs=[seed, randomize_seed],
+        outputs=seed,
+        queue=False,
+        api_name=False,
+    ).then(
         fn=get_token_table,
         inputs=prompt,
         outputs=token_indices_table,