import gradio as gr
from functools import partial
from all_models import models

NUM_MODELS = 6
DEFAULT_MODELS = models[:NUM_MODELS]

def load_model(model_name):
    try:
        return gr.load(f'models/{model_name}')
    except Exception:
        return gr.Interface(lambda txt: None, ['text'], ['image'])

def load_models(model_list):
    return {model: load_model(model) for model in model_list}

MODELS_LOADED = load_models(models)

def extend_choices(choices):
    return choices + ['NA'] * (NUM_MODELS - len(choices))

def update_imgbox(choices):
    choices_extended = extend_choices(choices)
    return [gr.Image(None, label=m, visible=(m != 'NA')) for m in choices_extended]

def gen_fn(model_str, prompt, negative_prompt=None, image_style="Default"):
    if model_str == 'NA':
        return None
    modified_prompt = f"{prompt}, {image_style}" if image_style != "Default" else prompt
    if negative_prompt:
        modified_prompt += f", not {negative_prompt}"
    return MODELS_LOADED[model_str](modified_prompt)

def create_interface():
    with gr.Blocks() as demo:
        with gr.Tab('The Dream'):
            txt_input = gr.Textbox(label='Your prompt:', lines=4)
            with gr.Accordion("Advanced Settings", open=False):
                neg_prompt = gr.Textbox(label='Negative prompt (Optional):', placeholder='Enter undesirable attributes here', lines=2)
                image_style = gr.Dropdown(label='Select Style', choices=["Default", "Realistic", "Portrait", "Anime"], value="Default")
            
            gen_button = gr.Button('Generate up to 6 images in up to 2 minutes total')
            stop_button = gr.Button('Stop', variant='secondary', interactive=False)
            
            gen_button.click(lambda: gr.update(interactive=True), None, stop_button, concurrency_limit=10)
            
            gr.HTML(
                """
                <div style="text-align: center; max-width: 1200px; margin: 0 auto;">
                    <p style="margin-bottom: 10px; color: var(--body-text-color);">Scroll down to see more images and select models.</p>
                </div>
                """
            )
            
            with gr.Row():
                output = [gr.Image(label=m, min_width=460) 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(partial(gen_fn, prompt=txt_input, negative_prompt=neg_prompt, image_style=image_style), 
                                                 inputs=[m], outputs=[o], concurrency_limit=10)
                    stop_button.click(lambda: gr.update(interactive=False), None, stop_button, cancels=[gen_event])
            
            with gr.Accordion('Model selection'):
                model_choice = gr.CheckboxGroup(models, label=f'Choose up to {NUM_MODELS} different models from the best available!', 
                                                value=DEFAULT_MODELS, interactive=True)
                model_choice.change(update_imgbox, model_choice, output)
                model_choice.change(extend_choices, model_choice, current_models)
            
            gr.HTML(
                """
                <div class="footer">
                    <p>Have fun!</p>
                </div>
                """
            )
    
    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch(max_threads=200)