Update app.py

app.py

@@ -2,19 +2,9 @@ import os
 import gradio as gr
 from random import randint
 from all_models import models
-import uuid
-from fastapi.staticfiles import StaticFiles
-from PIL import Image
-import requests
-from io import BytesIO
-import tempfile
 
 css_code = os.getenv("DazDinGo_CSS")
 
-# Create a dedicated directory for images
-IMAGE_DIR = "/tmp/flux_images"
-os.makedirs(IMAGE_DIR, exist_ok=True)
-
 def load_fn(models):
     global models_load
     models_load = {}
@@ -28,108 +18,58 @@ def load_fn(models):
 
 load_fn(models)
 
-def process_image_output(image_data):
-    """Handle various image output formats and save to temp directory"""
-    if image_data is None:
-        return None
-        
-    # If it's already a file path that exists, use it
-    if isinstance(image_data, str) and os.path.isfile(image_data):
-        return image_data
-        
-    # If it's a URL, download it
-    if isinstance(image_data, str) and image_data.startswith(('http://', 'https://')):
-        try:
-            response = requests.get(image_data)
-            img = Image.open(BytesIO(response.content))
-        except:
-            return None
-    else:
-        # Assume it's a PIL Image or numpy array
-        img = image_data
-    
-    # Save to temp file
-    filename = f"{uuid.uuid4()}.jpg"
-    filepath = os.path.join(IMAGE_DIR, filename)
-    
-    try:
-        if isinstance(img, Image.Image):
-            img.save(filepath)
-        else:
-            # Handle numpy arrays or other formats
-            Image.fromarray(img).save(filepath)
-        return filepath
-    except Exception as e:
-        print(f"Error saving image: {e}")
-        return None
+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
-        
     noise = str(randint(0, 4294967296))
     klir = '| ultra detail, ultra elaboration, ultra quality, perfect'
-    result = models_load[model_str](f'{prompt} {klir} {noise}')
-    
-    if isinstance(result, list):
-        return [process_image_output(img) for img in result if img is not None]
-    return process_image_output(result)
+    return models_load[model_str](f'{prompt} {klir} {noise}')
 
-def create_interface():
-    with gr.Blocks(css=css_code) as demo:
-        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"
-                )
-            
-            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 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)
         
-        with gr.Row():
-            output = gr.Gallery(
-                label="Generated Images",
-                columns=3,
-                height="auto",
-                object_fit="contain",
-                show_download_button=True,
-                elem_id="custom_gallery"
-            )
+        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")
             
-            model_selector = gr.Dropdown(
-                models,
-                label="Select LoRA Model",
-                value=models[0] if models else None,
-                multiselect=False,
-                interactive=True,
-                elem_id="custom_dropdown"
-            )
-        
-        gen_event = gen_button.click(
-            fn=gen_fn,
-            inputs=[model_selector, txt_input],
-            outputs=output
-        )
-        
-        stop_button.click(
-            fn=lambda: [gr.Button(interactive=True), gr.Button(interactive=False)],
-            inputs=None,
-            outputs=[gen_button, stop_button],
-            cancels=[gen_event]
-        )
-        
-        return demo
-
-demo = create_interface()
+            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, 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.Row():
+        gr.HTML("")
 
-# Configure static files serving
-app = demo.app
-app.mount("/file", StaticFiles(directory=IMAGE_DIR), name="static")
+with gr.Blocks(css=css_code) as demo: 
+    make_me()
 
-if __name__ == "__main__":
-    demo.queue(concurrency_count=50).launch()
+demo.queue(concurrency_count=50)
+demo.launch()