import gradio as gr
import numpy as np
import PIL.Image as Image
import random
import time

# Список моделей с отображаемыми названиями и внутренними именами
models = {
    "FLUX.1-dev": "black-forest-labs/FLUX.1-dev",
    "flux-RealismLora": "XLabs-AI/flux-RealismLora",
    "Midjourney": "Jovie/Midjourney",
    "Stable Diffusion v1-5": "stable-diffusion-v1-5/stable-diffusion-v1-5",
    "Pixel Art XL": "nerijs/pixel-art-xl",
    "Knitted Character Flux LoRA": "prithivMLmods/Knitted-Character-Flux-LoRA",
    "Flux Ghibsky Illustration": "aleksa-codes/flux-ghibsky-illustration",
    "Flux Super Realism LoRA": "strangerzonehf/Flux-Super-Realism-LoRA",
    "Flux Game Assets LoRA v2": "gokaygokay/Flux-Game-Assets-LoRA-v2",
    "Flux Animex v2 LoRA": "strangerzonehf/Flux-Animex-v2-LoRA",
    "Flux Qwen Capybara": "cfahlgren1/flux-qwen-capybara",
    "Softserve Anime": "alvdansen/softserve_anime",
    "FLUX Aesthetic Anime": "dataautogpt3/FLUX-AestheticAnime",
    "FLUX Animeo v1 LoRA": "strangerzonehf/Flux-Animeo-v1-LoRA",
    "Fractured Line Flare": "prithivMLmods/Fractured-Line-Flare",
    "Castor 3D Portrait Flux LoRA": "prithivMLmods/Castor-3D-Portrait-Flux-LoRA"
}

# Очередь на генерации
queue = []

def generate_image(prompt, model_name, width, height, seed, steps, guidance_scale, random_seed):
    # Проверка очереди
    if len(queue) >= 150:
        yield None, "Очередь переполнена. Пожалуйста, подождите.", None
        return
    
    # Добавляем задачу в очередь
    task_id = random.randint(100000, 999999)
    queue.append(task_id)
    
    # Генерация сида, если выбрана галочка "Случайный сид"
    if random_seed:
        seed = random.randint(0, 1000000)
    
    # Формирование текста параметров
    params_text = (
        f"Промпт: {prompt}\n"
        f"Модель: {model_name}\n"
        f"Размер: {width}x{height}\n"
        f"Сид: {seed}\n"
        f"Шаги: {steps}\n"
        f"Гуидансе скейл: {guidance_scale}"
    )
    
    # Отправляем текст параметров
    yield None, params_text, seed
    
    for step in range(steps):
        # Имитация генерации изображения
        time.sleep(random.uniform(0.1, 0.5))  # Таймаут до 0.5 секунды
        
        # Имитация изображения (шум)
        image_array = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
        image = Image.fromarray(image_array)
        
        # Обновление изображения
        yield image, None, None
    
    # Удаляем задачу из очереди
    queue.remove(task_id)
    
    # Возвращаем финальное изображение и параметры
    yield image, params_text, seed

# Создание интерфейса
with gr.Blocks() as demo:
    gr.Markdown("# Генератор изображений")
    
    with gr.Row():
        prompt_textbox = gr.Textbox(label="Промпт", placeholder="Введите ваш промпт здесь", lines=2)
    
    with gr.Row():
        model_dropdown = gr.Dropdown(list(models.keys()), label="Модель", value="FLUX.1-dev")
    
    with gr.Row():
        with gr.Column():
            width_slider = gr.Slider(128, 1024, value=512, step=64, label="Ширина")
        with gr.Column():
            height_slider = gr.Slider(128, 1024, value=512, step=64, label="Высота")
    
    with gr.Row():
        seed_slider = gr.Slider(0, 1000000, value=random.randint(0, 1000000), step=1, label="Сид")
        random_seed_checkbox = gr.Checkbox(label="Случайный сид")
    
    with gr.Row():
        steps_slider = gr.Slider(1, 100, value=50, step=1, label="Шаги")
        guidance_scale_slider = gr.Slider(0.1, 20.0, value=7.5, step=0.1, label="Гуидансе скейл")
    
    generate_button = gr.Button("Сгенерировать изображение")
    output_image = gr.Image(label="Сгенерированное изображение")
    params_textbox = gr.Textbox(label="Параметры генерации", interactive=False)
    
    generate_button.click(
        fn=generate_image,
        inputs=[
            prompt_textbox,
            model_dropdown,
            width_slider,
            height_slider,
            seed_slider,
            steps_slider,
            guidance_scale_slider,
            random_seed_checkbox
        ],
        outputs=[output_image, params_textbox, seed_slider]
    )
    
    demo.queue(max_size=150, api_open=False)
    demo.launch(share=False, server_port=7860)

# Скрытие футера
hide_footer = """
<style>
    .gradio-container .footer {
        display: none !important;
    }
</style>
"""
demo = demo.add_css(hide_footer)

# Установка таймаута
demo.queue(api_name="/generate").set_config(max_size=150, api_open=False, max_threads=1, timeout=125)

demo.launch()