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

def load_models(model_names):
    loaded_models = {}
    for model_name in model_names:
        try:
            model = gr.load(f'models/{model_name}')
        except Exception:
            model = gr.Interface(lambda txt: "Model load error", inputs=["text"], outputs=["image"])
        loaded_models[model_name] = model
    return loaded_models

models_load = load_models(models)

num_models = 6
default_models = models[:num_models]

def extend_choices(choices, num_models=6):
    return choices + ['NA'] * (num_models - len(choices))

def update_image_boxes(choices):
    extended_choices = extend_choices(choices)
    return [gr.Image(label=model, visible=(model != 'NA')) for model in extended_choices]

def generate_image(model_name, prompt):
    if model_name == 'NA':
        return None
    # If you decide to add randomness to the prompt, uncomment the next line
    # prompt += f" {randint(0, 999999999)}"
    return models_load[model_name](prompt)

with gr.Blocks() as demo:
    with gr.Tab("Zexxxi"):
        txt_input = gr.Textbox(label="Your prompt:", lines=4).style(container=False, min_width=1200)
        gen_button = gr.Button("Generate up to 6 images")
        stop_button = gr.Button("Stop", variant="secondary")
        gen_button.click(lambda _: gr.update(interactive=True), None, stop_button)
        
        with gr.Row():
            output_images = [gr.Image(label=model, min_width=480) for model in default_models]
            model_textboxes = [gr.Textbox(model, visible=False) for model in default_models]
            
            for model_textbox, output_image in zip(model_textboxes, output_images):
                gen_event = gen_button.click(generate_image, [model_textbox, txt_input], output_image)
                stop_button.click(lambda _: gr.update(interactive=False), None, stop_button, cancels=[gen_event])
        
        with gr.Accordion("Model selection"):
            model_choice = gr.CheckboxGroup(choices=models, label=f"Choose up to {num_models} models from the 686 available!", value=default_models, multiselect=True, max_choices=num_models, interactive=True)
            model_choice.change(update_image_boxes, model_choice, output_images)
            model_choice.change(lambda choices: extend_choices(choices, num_models), model_choice, model_textboxes)
    
        with gr.Row():
            gr.HTML("""
                <div class="footer">
                    <p>zexxxi nsfw model</p>
                </div>
            """)

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