import gradio as gr
import requests

# Базовые настройки
prompt = gr.Textbox(label="Prompt")
models = gr.Radio(
    ["Bard", "DALL-E 2", "VQGAN+CLIP", "VQGAN", "CLIP", "InceptionV3", "VQGAN+CLIP-ViT-B16", "VQGAN+CLIP-ViT-B32", "VQGAN+CLIP-ViT-L", "VQGAN+CLIP-ViT-H", "VQGAN+CLIP-ViT-XL"],
    label="Model",
)

# Расширенные настройки
negative_prompt = gr.Textbox(label="Negative Prompt")
sampling_method = gr.Radio(
    ["random", "greedy", "nucleus", "top_k", "top_p"], label="Sampling Method"
)
sampling_steps = gr.Number(
    label="Sampling Steps", value=30, minimum=1, maximum=100
)
cfg_scale = gr.Number(label="CFG Scale", minimum=0.1, maximum=10.0, value=1.0)
seed = gr.Number(label="Seed", minimum=0, maximum=2**31, value=0)

# Улучшение качества
algorithm = gr.Radio(
    ["nearest", "bilinear", "bicubic", "lanczos", "cubic", "mitchell", "bicubic_nn", "bicubic_nn_diff", "bicubic_nn_diff_v2"], label="Algorithm"
)

# Функция генерации изображения
def generate_image(prompt, model, negative_prompt, sampling_method, sampling_steps, cfg_scale, seed):
    url = "https://api.huggingface.co/models/text-to-image/v1/generate"
    data = {
        "prompt": prompt,
        "model": model,
        "negative_prompt": negative_prompt,
        "sampling_method": sampling_method,
        "sampling_steps": sampling_steps,
        "cfg_scale": cfg_scale,
        "seed": seed,
    }
    response = requests.post(url, json=data)
    image = response.json()["image"]
    return image

# Функция улучшения качества изображения
def improve_quality(image, algorithm):
    url = "https://api.huggingface.co/models/text-to-image/v1/improve-quality"
    data = {
        "image": image,
        "algorithm": algorithm,
    }
    response = requests.post(url, json=data)
    image = response.json()["image"]
    return image

# Основная функция
def main():
    interface = gr.Interface(
        generate_image,
        [
            gr.Row(
                gr.Column(prompt, "Базовые настройки"),
                gr.Column(models, "Базовые настройки"),
            ),
            gr.Row(
                gr.Column(negative_prompt, "Расширенные настройки"),
                gr.Column(sampling_method, "Расширенные настройки"),
                gr.Column(sampling_steps, "Расширенные настройки"),
                gr.Column(cfg_scale, "Расширенные настройки"),
                gr.Column(seed, "Расширенные настройки"),
            ),
            gr.Row(
                algorithm,  # Assuming 'algorithm' is a single component
            ),
        ],
        outputs=gr.Image(),
        title="Gradio Image Generator",
    )

    interface.launch()

if __name__ == "__main__":
    main()