app.py

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-from io import BytesIO
 from random import randint
 from all_models import models
 
@@ -30,37 +29,25 @@ def gen_fn(model_str, prompt):
     if model_str == 'NA':
         return None
     noise = str(randint(0, 99999999))
-    result = models_load[model_str](f'{prompt} {noise}')
-    
-    if isinstance(result, str):
-        response = requests.get(result)
-        img = Image.open(BytesIO(response.content))
-    elif isinstance(result, bytes):
-        img = Image.open(BytesIO(result))
-    else:
-        return None
-    
-    buffered = BytesIO()
-    img.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
-    return f"data:image/png;base64,{img_str}"
+    return models_load[model_str](f'{prompt} {noise}')
 
 def make_me():
     with gr.Row():
-        with gr.Column(scale=1):
+        with gr.Column(scale=4):
             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.Column(scale=1):
+            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", show_label=True, interactive=False, show_share_button=False) for m in default_models]