import gradio as gr
import requests
import os

# Замените "YOUR_MODEL_ENDPOINT_HERE" и "YOUR_BEARER_TOKEN_HERE" на фактический адрес вашей модели Hugging Face и ваш Bearer токен
MODEL_ENDPOINT = os.getenv("BASE_URL")
BEARER_TOKEN = os.getenv("API_KEY")

def generate_image(prompt, negative_prompt, sampling_steps, cfg_scale, seed):
    # Создаем заголовок авторизации Bearer
    headers = {"Authorization": f"Bearer {BEARER_TOKEN}"}

    # Отправляем запрос к API модели Hugging Face с использованием введенных параметров и заголовка авторизации
    response = requests.post(
        MODEL_ENDPOINT,
        json={
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "sampling_steps": sampling_steps,
            "cfg_scale": cfg_scale,
            "seed": seed
        },
        headers=headers
    )

    # Получаем сгенерированное изображение из ответа API
    generated_image = response.content

    return generated_image

iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox("text", label="Prompt"),
        gr.Textbox("text", label="Negative Prompt"),
    ],
    outputs=gr.Image(label="Generated Image"),
    live=True,
    title="Hugging Face Image Generator",
    theme="compact",
    layout="vertical",
    sidebar=[
        gr.Slider("number", label="Sampling Steps", min=1, max=30, default=10),
        gr.Slider("number", label="CFG Scale", min=0.1, max=2.0, default=1.0),
        gr.Textbox("number", label="Seed (default -1)", default=-1)
    ]
)

iface.launch()