Update app.py

app.py

@@ -2,17 +2,18 @@ import gradio as gr
 from random import randint
 from all_models import models
 
-models_load = {}
-def load_models():
+def load_fn(models):
     global models_load
+    models_load = {}
     for model in models:
-        try:
-            models_load[model] = gr.Interface.load(f'models/{model}')
-        except Exception as error:
-            print(f"Failed to load model {model}: {error}")
-            models_load[model] = gr.Interface(lambda txt: None, ['text'], ['image'])
+        if model not in models_load.keys():
+            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_models()
+load_fn(models)
 
 num_models = len(models)
 default_models = models[:num_models]
@@ -27,31 +28,25 @@ def update_imgbox(choices):
 def gen_fn(model_str, prompt):
     if model_str == 'NA':
         return None
-    noise = str(randint(0, 99999999))
-    return models_load[model_str](f'{prompt} {noise}').outputs[0]
+    noise = str(randint(0, 9999999))
+    return models_load[model_str](f'{prompt} {noise}')
 
 def make_me():
     with gr.Row():
-        txt_input = gr.Textbox(label='Your prompt:', lines=3, container=False, elem_id="custom_textbox", placeholder="Prompt")
-        gen_button = gr.Button('Generate images', elem_id="custom_gen_button")
-        stop_button = gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
-        
-        def on_generate_click():
-            nonlocal gen_button, stop_button
-            gen_button.interactive = True
-            stop_button.interactive = True
-            gen_button.label = 'Generating...'
-            stop_button.label = 'Stop'
-        
-        def on_stop_click():
-            nonlocal gen_button, stop_button
-            gen_button.interactive = False
-            stop_button.interactive = False
-            gen_button.label = 'Generate images'
-            stop_button.label = 'Stopped'
-        
-        gen_button.click(on_generate_click, inputs=None, outputs=None)
-        stop_button.click(on_stop_click, inputs=None, outputs=None)
+        with gr.Column(scale=1):
+            txt_input = gr.Textbox(label='Your prompt:', lines=3, container=False, elem_id="custom_textbox", placeholder="Prompt")
+            with gr.Row():
+                gen_button = gr.Button('Generate images', elem_id="custom_gen_button")
+                stop_button = gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
+                
+                def on_generate_click():
+                    return gr.Button('Generate images', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=True, elem_id="custom_stop_button")
+                
+                def on_stop_click():
+                    return gr.Button('Generate images', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
+                
+                gen_button.click(on_generate_click, inputs=None, outputs=[gen_button, stop_button])
+                stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button])
     
     with gr.Row():
         output = [gr.Image(label=m, min_width=250, height=250, elem_id="custom_image") for m in default_models]
@@ -74,7 +69,7 @@ custom_css = """
 }
 body {
     background-color: var(--body-background-fill) !important;
-    color: #c5c6c7;
+    color: #2d3d4f;
     margin: 0;
     padding: 0;
     font-family: Arial, sans-serif;
@@ -223,4 +218,4 @@ with gr.Blocks(css=custom_css) as demo:
     make_me()
 
 demo.queue(concurrency_count=100)
-demo.launch()
+demo.launch()