import gradio as gr
import requests
from PIL import Image
from io import BytesIO
from dotenv import load_dotenv
import os

# Загрузка токена из .env файла
load_dotenv()
API_TOKEN = os.getenv("HF_API_TOKEN")

def generate_image(prompt, model_name):
    API_URL = f"https://api-inference.huggingface.co/models/{model_name}"
    headers = {"Authorization": f"Bearer {API_TOKEN}"}

    response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
    if response.status_code == 200:
        return Image.open(BytesIO(response.content))
    else:
        return f"Ошибка: {response.status_code}, {response.text}"

models = {
    "Stable Diffusion v1.5": "Yntec/stable-diffusion-v1-5",
    "DreamBooth v2": "Yntec/stable-diffusion-v1-5",
    "DALL-E Mini": "Yntec/stable-diffusion-v1-5",
    "DeepAI Generator": "Jovie/Midjourney",
    "Realistic Vision": "Jovie/Midjourney",
    "Artistic AI": "Jovie/Midjourney"
}

def handle_input(prompt):
    outputs = {}
    for name, model in models.items():
        outputs[name] = generate_image(prompt, model)
    return list(outputs.values())


with gr.Blocks() as demo:
    gr.Markdown("## Демо-генерация изображений с помощью различных моделей нейросетей")
    with gr.Row():
        user_input = gr.Textbox(label="Введите описание изображения",
                                placeholder="Например, 'Астронавт верхом на лошади'")
    with gr.Row():
        outputs = [
            gr.Image(label=name, interactive=False)
            for name in models.keys()
        ]
    generate_button = gr.Button("Сгенерировать")

    # Обработка клика
    generate_button.click(
        fn=handle_input,
        inputs=[user_input],
        outputs=outputs
    )

demo.launch()