import gradio as gr
import random

# Загрузка модели
model = gr.load(
    "models/fffiloni/cute-comic-800",
    provider="hf-inference",
)

# Создание интерфейса с детским стилем
with gr.Blocks() as demo:
    # Входное текстовое поле
    input_text = gr.Textbox(label="Напиши, что хочешь увидеть на картинке в стиле Minecraft!")

    # Параметры для настройки генерации
    environment_checkbox = gr.CheckboxGroup(
        label="Добавь объекты на картинку",
        choices=["Деревья", "Горы", "Водопад", "Домик", "Пещера", "Озеро"],
        value=[]
    )

    weather_checkbox = gr.CheckboxGroup(
        label="Добавь погоду",
        choices=["Солнечно", "Дождь", "Снег", "Туман"],
        value=[]
    )

    effects_checkbox = gr.CheckboxGroup(
        label="Добавь эффекты",
        choices=["Тень", "Свет", "Размытость"],
        value=[]
    )

    # Выбор стиля
    style_dropdown = gr.Dropdown(
        label="Выбери стиль",
        choices=["Пиксель-арт", "Комикс", "Мультфильм"],
        value="Пиксель-арт"
    )

    # Выбор цветовой палитры
    color_palette_dropdown = gr.Dropdown(
        label="Выбери цветовую палитру",
        choices=["Яркая", "Пастельная", "Темная"],
        value="Яркая"
    )

    # Кнопка для запуска модели
    submit_button = gr.Button("Создать картинку")

    # Выходное изображение
    output_image = gr.Image(label="Вот твоя картинка!", type="filepath", height=500)

    # Кнопка для скачивания изображения
    download_button = gr.DownloadButton("Скачать картинку")

    # Обработчик события нажатия кнопки
    def generate_and_display(text, environment, weather, effects, style, color_palette):
        # Перевод параметров на английский
        environment_mapping = {
            "Деревья": "trees",
            "Горы": "mountains",
            "Водопад": "waterfall",
            "Домик": "house",
            "Пещера": "cave",
            "Озеро": "lake"
        }

        weather_mapping = {
            "Солнечно": "sunny",
            "Дождь": "rain",
            "Снег": "snow",
            "Туман": "fog"
        }

        effects_mapping = {
            "Тень": "shadow",
            "Свет": "light",
            "Размытость": "blur"
        }

        style_mapping = {
            "Пиксель-арт": "pixel art",
            "Комикс": "comic",
            "Мультфильм": "cartoon"
        }

        color_palette_mapping = {
            "Яркая": "bright",
            "Пастельная": "pastel",
            "Темная": "dark"
        }

        # Формирование промта с выбранными параметрами
        prompt_parts = [text, f"in {style_mapping[style]} style", f"with a {color_palette_mapping[color_palette]} color palette"]

        if environment:
            prompt_parts.append(f"Add environment: {', '.join([environment_mapping[env] for env in environment])}")

        if weather:
            prompt_parts.append(f"Add weather: {', '.join([weather_mapping[w] for w in weather])}")

        if effects:
            prompt_parts.append(f"Add effects: {', '.join([effects_mapping[e] for e in effects])}")

        prompt_parts.append("in Minecraft style")

        seed = random.randint(0, 10000)
        prompt_parts.append(f"Seed: {seed}")

        prompt = " ".join(prompt_parts)

        # Вывод промта в лог
        print("Сгенерированный промт:", prompt)

        # Генерация изображения
        image_path = model(prompt)
        return image_path, image_path

    submit_button.click(generate_and_display, inputs=[input_text, environment_checkbox, weather_checkbox, effects_checkbox, style_dropdown, color_palette_dropdown], outputs=[output_image, download_button])

# Запуск интерфейса
demo.launch()