import gradio as gr
from random import randint

# Список моделей (пример)


# Загрузка моделей
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:
                m = gr.Interface(lambda txt: None, ['text'], ['image'])
            models_load[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):
    if model_str == 'NA':
        return None
    noise = str(randint(0, 9999999))
    return models_load[model_str](f'{prompt} {noise}')

# Создание интерфейса
def make_me():
    with gr.Row():
        with gr.Column(scale=1):
            txt_input = gr.Textbox(
                label='Your prompt:', 
                lines=3, 
                elem_id="custom_textbox", 
                placeholder="Enter your prompt"
            )
            with gr.Row():
                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 gen_button, stop_button.interactive(True)

                def on_stop_click():
                    return gen_button, stop_button.interactive(False)

                gen_button.click(on_generate_click, None, [gen_button, stop_button])
                stop_button.click(on_stop_click, None, [gen_button, stop_button])

    with gr.Row():
        output = [gr.Image(label=m, elem_id="custom_image") for m in default_models]  # Убрали tool="editor"
        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, None, [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} models selected', 
            value=default_models, 
            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("")

# Кастомный CSS для интерфейса
custom_css = """
body, html {
    overflow-y: auto;
    height: 100%;
    margin: 0;
    padding: 0;
}

.gradio-container {
    overflow-y: auto !important;
    height: 100vh;
    display: flex;
    flex-direction: column;
}

/* Стили для текста */
.custom_textbox {
    background-color: #2d343f;
    border: 1px solid #3b4252;
    color: #7f8184;
    padding: 10px;
    border-radius: 4px;
    width: 100%;
    box-sizing: border-box;
}

/* Стили для кнопок */
.custom_gen_button, .custom_stop_button {
    padding: 15px;
    border-radius: 4px;
    transition: all 0.2s;
    cursor: pointer;
}

.custom_gen_button:hover, .custom_stop_button:hover {
    transform: scale(1.05);
}

/* Стили для изображений */
.custom_image {
    max-width: 100%;
    height: auto;
    object-fit: contain;
    overflow: auto;
    border: 1px solid #3b4252;
    background-color: #2d343f;
    border-radius: 4px;
    cursor: grab;
    margin: 10px;
}

/* Стили для выбора моделей */
.custom_accordion {
    margin-top: 20px;
    border: 1px solid #3b4252;
    border-radius: 4px;
}
"""

# Создание интерфейса
with gr.Blocks(css=custom_css) as demo:
    make_me()

# Очередь и запуск интерфейса с параметрами
demo.queue(default_concurrency_limit=240, max_size=240)
demo.launch(max_threads=400, ssr_mode=True)