import gradio as gr
import requests
import io
import random
import os
from PIL import Image
from deep_translator import GoogleTranslator

# Словарь URL для различных моделей
MODEL_URLS = {
    "DALL-E 3 XL": "https://api-inference.huggingface.co/models/openskyml/dalle-3-xl",
    "Playground 2": "https://api-inference.huggingface.co/models/playgroundai/playground-v2-1024px-aesthetic",
    "Openjourney 4": "https://api-inference.huggingface.co/models/prompthero/openjourney-v4",
    "AbsoluteReality 1.8.1": "https://api-inference.huggingface.co/models/digiplay/AbsoluteReality_v1.8.1",
    # ...
}

# Функция для получения изображения от API модели
def get_image_from_api(prompt, model, headers, payload):
    response = requests.post(MODEL_URLS[model], headers=headers, json=payload)
    if response.status_code == 200:
        image_bytes = response.content
        image = Image.open(io.BytesIO(image_bytes))
        return image
    else:
        raise gr.Error(f"Ошибка {response.status_code}: {response.reason}")

# Функция для запроса к GPT (вы должны добавить свой код здесь)
def enhance_prompt_with_gpt(prompt):
    # Добавьте ваш код для запроса к GPT здесь...
    pass

# Функция для перевода текста через Google Translator
def translate_prompt(prompt, source_lang="ru", target_lang="en"):
    return GoogleTranslator(source=source_lang, target=target_lang).translate(prompt)

# Функция для генерации изображения
def generate_image(prompt, model, negative_prompt, steps, cfg_scale, sampler, seed, strength, use_gpt):
    if not prompt:
        raise gr.Error("Prompt не может быть пустым")

    if use_gpt:
        prompt = enhance_prompt_with_gpt(prompt)

    translated_prompt = translate_prompt(prompt)

    # Добавьте к prompt дополнительные параметры если нужно
    enhanced_prompt = f"{translated_prompt} | ultra detail, ultra elaboration, ultra quality, perfect."

    payload = {
        "inputs": enhanced_prompt,
        "is_negative": negative_prompt,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "sampler": sampler,
        "seed": seed if seed != -1 else random.randint(1, 1000000000),
        "strength": strength
    }

    # Получаем токен из переменных окружения
    api_token = os.getenv("HF_API_TOKEN")
    headers = {"Authorization": f"Bearer {api_token}"}

    try:
        image = get_image_from_api(enhanced_prompt, model, headers, payload)
        return image
    except Exception as e:
        raise gr.Error(f"Ошибка при генерации изображения: {e}")

css = """
* {}
footer {visibility: hidden !important;}
"""

# Создание интерфейса Gradio
with gr.Blocks(css=css) as dalle_interface:
    with gr.Tab("Базовые настройки"):
        with gr.Row():
            text_prompt = gr.Textbox(label="Prompt", placeholder="Описание изображения", lines=3)
            model = gr.Radio(label="Модель", value="DALL-E 3 XL", choices=list(MODEL_URLS.keys()))

    with gr.Tab("Расширенные настройки"):
        negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Чего не должно быть на изображении", lines=3)
        steps = gr.Slider(label="Sampling steps", value=35, minimum=1, maximum=100, step=1)
        cfg = gr.Slider(label="CFG Scale", value=7, minimum=1, maximum=20, step=1)
        method = gr.Radio(label="Sampling method", value="DPM++ 2M Karras", choices=["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"])
        strength = gr.Slider(label="Strength", value=0.7, minimum=0, maximum=1, step=0.001)
        seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=1000000000, step=1)
        use_gpt = gr.Checkbox(label="Use GPT to enhance prompt", value=False)

    with gr.Tab("Информация"):
        gr.Textbox(label="Prompt template", value="{prompt} | ultra detail, ultra elaboration, ultra quality, perfect.")

    with gr.Row():
        generate_button = gr.Button("Generate Image")

    with gr.Row():
        image_output = gr.Image(type="pil", label="Generated Image")

    generate_button.click(
        generate_image,
        inputs=[text_prompt, model, negative_prompt, steps, cfg, method, seed, strength, use_gpt],
        outputs=image_output
    )

# Запуск интерфейса
dalle_interface.launch(show_api=False, share=False)