Update app.py

app.py

@@ -2,62 +2,70 @@ import gradio as gr
 from random import randint
 from all_models import models
 
-def load_fn(models):
-    global models_load
+# Load models
+def load_models(models):
     models_load = {}
-    
     for model in models:
-        if model not in models_load.keys():
+        if model not in models_load:
             try:
                 m = gr.load(f'models/{model}')
             except Exception as error:
                 m = gr.Interface(lambda txt: None, ['text'], ['image'])
-            models_load.update({model: m})
-
-
-load_fn(models)
+            models_load[model] = m
+    return models_load
 
+models_load = load_models(models)
 
 num_models = 6
 default_models = models[:num_models]
 
-
+# Extend choices to a fixed number of models
 def extend_choices(choices):
-    return choices + (num_models - len(choices)) * ['NA']
-
+    return choices + ['NA'] * (num_models - len(choices))
 
+# Dynamically update image boxes based on number of choices
 def update_imgbox(choices):
-    choices_plus = extend_choices(choices)
-    return [gr.Image(None, label = m, visible = (m != 'NA')) for m in choices_plus]
+    extended_choices = extend_choices(choices)
+    return [gr.Image(None, label=m, visible=(m != 'NA')) for m in extended_choices]
 
-    
-def gen_fn(model_str, prompt):
+# Generate function with noise added to prompt
+def generate_image(model_str, prompt):
     if model_str == 'NA':
         return None
     noise = str(randint(0, 99999999999))
     return models_load[model_str](f'{prompt} {noise}')
 
-
+# Gradio interface setup
 with gr.Blocks() as demo:
-    model_choice2 = gr.Dropdown(models, label = 'Choose model', value = models[0], filterable = False)
-    txt_input2 = gr.Textbox(label = 'Prompt text')
-        
+    model_dropdown = gr.Dropdown(models, label='Choose model', value=models[0], filterable=False)
+    text_input = gr.Textbox(label='Prompt text')
+
     max_images = 6
-    num_images = gr.Slider(1, max_images, value = max_images, step = 1, label = 'Number of images')
-        
-    gen_button2 = gr.Button('Generate')
-    stop_button2 = gr.Button('Stop', variant = 'secondary', interactive = False)
-    gen_button2.click(lambda s: gr.update(interactive = True), None, stop_button2)
-        
+    num_images_slider = gr.Slider(1, max_images, value=max_images, step=1, label='Number of images')
+
+    generate_button = gr.Button('Generate')
+    stop_button = gr.Button('Stop', variant='secondary', interactive=False)
+
+    # Enable the stop button when generation starts
+    generate_button.click(lambda: gr.update(interactive=True), None, stop_button)
+
     with gr.Row():
-        output2 = [gr.Image(label = '') for _ in range(max_images)]
+        output_images = [gr.Image(label='') for _ in range(max_images)]
+
+    for i, output in enumerate(output_images):
+        img_index = gr.Number(i, visible=False)
+        num_images_slider.change(
+            lambda idx, n: gr.update(visible=(idx < n)), 
+            [img_index, num_images_slider], output
+        )
+        generate_event = generate_button.click(
+            lambda idx, n, model, prompt: generate_image(model, prompt) if idx < n else None, 
+            [img_index, num_images_slider, model_dropdown, text_input], output
+        )
 
-    for i, o in enumerate(output2):
-        img_i = gr.Number(i, visible = False)
-        num_images.change(lambda i, n: gr.update(visible = (i < n)), [img_i, num_images], o)
-        gen_event2 = gen_button2.click(lambda i, n, m, t: gen_fn(m, t) if (i < n) else None, [img_i, num_images, model_choice2, txt_input2], o)
-        stop_button2.click(lambda s: gr.update(interactive = False), None, stop_button2, cancels = [gen_event2])
+        # Stop button functionality to cancel image generation
+        stop_button.click(lambda: gr.update(interactive=False), None, stop_button, cancels=[generate_event])
 
-    
-demo.queue(concurrency_count = 36)                        
-demo.launch()
+# Set queue concurrency and launch the demo
+demo.queue(concurrency_count=36)
+demo.launch()