app.py

import os
import gradio as gr
from random import randint
from all_models import models

css_code = os.getenv("DazDinGo_CSS", "No CSS found")

def load_fn(models):
    global models_load
    models_load = {}
    for model in models:
        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_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, negative_prompt, cfg_scale, sampler, seed):
    if model_str == 'NA':
        return None
    # Use the provided seed or generate a random one if seed is -1
    noise = seed if seed != -1 else randint(0, 99999999)
    
    # Prepare the full prompt with negative prompt and CFG scale
    full_prompt = f'{prompt} --negative_prompt "{negative_prompt}" --cfg_scale {cfg_scale} --sampler {sampler} --seed {noise}'
    
    return models_load[model_str](full_prompt)

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)
            negative_prompt = gr.Textbox(label='Negative Prompt:', lines=2, placeholder="What to avoid in the image", elem_id="custom_negative_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 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=None, height=None, 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]
        
        # Additional inputs for CFG, sampler, seed, and negative prompt
        cfg_scale = gr.Slider(minimum=1, maximum=20, step=0.1, value=7, label="CFG Scale", elem_id="custom_cfg_scale")
        sampler = gr.Dropdown(choices=["DDIM", "PLMS", "Euler", "Heun"], value="DDIM", label="Sampler", elem_id="custom_sampler")
        seed = gr.Number(value=-1, label="Seed (-1 for random)", elem_id="custom_seed")
        
        for m, o in zip(current_models, output):
            gen_event = gen_button.click(
                gen_fn, 
                [m, txt_input, negative_prompt, cfg_scale, sampler, seed], 
                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("")

with gr.Blocks(css=css_code) as demo: 
    make_me()

demo.queue(concurrency_count=50)
demo.launch()