Update app.py

app.py

@@ -1,7 +1,10 @@
 import os
 import gradio as gr
+import base64
 from random import randint
 from all_models import models
+from io import BytesIO
+from PIL import Image
 
 css_code = os.getenv("DazDinGo_CSS")
 
@@ -18,52 +21,63 @@ def load_fn(models):
 
 load_fn(models)
 
-num_models = len(models)
-default_models = models[:num_models]
-
-def extend_choices(choices):
-    return choices + (num_models - len(choices)) * ['NA']
-
-def update_imgbox(choices):
-    choices_plus = extend_choices(choices)
-    return [gr.Image(None, label=m, visible=(m != 'NA'), elem_id="custom_image") for m in choices_plus]
-
 def gen_fn(model_str, prompt):
     if model_str == 'NA':
-        return None
+        return None, ""
     noise = str(randint(0, 4294967296))
     klir = '| ultra detail, ultra elaboration, ultra quality, perfect'
-    return models_load[model_str](f'{prompt} {klir} {noise}')
+    image = models_load[model_str](f'{prompt} {klir} {noise}')
+    
+    # Convert to base64
+    buffered = BytesIO()
+    if isinstance(image, str):  # if it's a file path
+        img = Image.open(image)
+        img.save(buffered, format="JPEG")
+    else:  # if it's a PIL Image
+        image.save(buffered, format="JPEG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    base64_code = f"data:image/jpeg;base64,{img_str}"
+    
+    return image, base64_code
 
 def make_me():
     with gr.Row():
         with gr.Column(scale=4):
-            txt_input = gr.Textbox(label='Your prompt:', lines=4, container=False, elem_id="custom_textbox", placeholder="Prompt", height=250)
+            txt_input = gr.Textbox(label='Your prompt:', lines=4, container=False, 
+                                 elem_id="custom_textbox", placeholder="Prompt", height=250)
         
         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")
+            gen_button = gr.Button('Generate image', 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")
+                return gr.Button('Generate image', 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")
+                return gr.Button('Generate image', 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, width=512, max_height=768, elem_id="custom_image", show_label=True, interactive=False, show_share_button=False) for m in default_models]
-        current_models = [gr.Textbox(m, visible=False) for m in default_models]
-        for m, o in zip(current_models, output):
-            gen_event = gen_button.click(gen_fn, [m, txt_input], o)
-            stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button], cancels=[gen_event])
-    
-    with gr.Accordion('Model selection', elem_id="custom_accordion"):
-        model_choice = gr.CheckboxGroup(models, label=f'{num_models} different models selected', value=default_models, interactive=True, elem_id="custom_checkbox_group")
-        model_choice.change(update_imgbox, model_choice, output)
-        model_choice.change(extend_choices, model_choice, current_models)
+        with gr.Column():
+            model_dropdown = gr.Dropdown(models, label="Select Model", 
+                                        value=models[0] if models else None)
+            output_image = gr.Image(label="Generated Image", width=512, 
+                                  max_height=768, elem_id="custom_image", 
+                                  show_label=True, interactive=False, 
+                                  show_share_button=False)
+            
+            # Base64 output
+            base64_output = gr.Textbox(label="Base64 Code", interactive=True, 
+                                     placeholder="Base64 code will appear here after generation",
+                                     lines=4)
+            
+            gen_event = gen_button.click(gen_fn, [model_dropdown, txt_input], 
+                                       [output_image, base64_output])
+            stop_button.click(on_stop_click, inputs=None, 
+                            outputs=[gen_button, stop_button], cancels=[gen_event])
     
     with gr.Row():
         gr.HTML("")