import gradio as gr
import requests
import os

# API ссылка
url = "https://stablediffusionapi.com/api/v4/dreambooth"

# API ключ
api_key = os.getenv("KEY")

model_list = ["realistic-vision-51", "anything-v5"]


# Функция для отправки запроса
def render(model, prompt, negative_prompt, width, height, guidance, seed, upscale):
    data = {
        "key":  api_key,  
        "model_id": model,  
        "prompt": prompt, 
        "negative_prompt": negative_prompt,  
        "width": width,  
        "height": height,  
        "samples": "1",  
        "num_inference_steps": "40",  
        "safety_checker": "no",  
        "enhance_prompt": "yes",  
        "seed": seed,  
        "guidance_scale": guidance,  
        "multi_lingual": "no",  
        "panorama": "no",  
        "self_attention": "no",  
        "upscale": "no",  
        "embeddings": "embeddings_model_id",  
        "lora": "lora_model_id",  
        "webhook": None,  
        "track_id": None,
    }
    response = requests.post(url, json=data)
    if response.status_code == 200:
        return response.content
    else:
        return None

# UI
gr.Interface(
    render,
    inputs=[
        gr.Dropdown(model_list, label="Модель", value="realistic-vision-51"),
        gr.Textbox(label="Описание изображения:", placeholder="Введите описание изображения", lines=3),
        gr.Textbox(label="Negative Prompt:", placeholder="Введите Negative Prompt", value="[deformed | disfigured], poorly drawn, [bad : wrong] anatomy, [extra | missing | floating | disconnected] limb, (mutated hands and fingers), blurry"),
        gr.Slider(show_label=True, minimum=256, maximum=2048, label="Ширина:", value="512", step=1),
        gr.Slider(show_label=True, minimum=256, maximum=2048, label="Высота:", value="512", step=1),
        gr.Slider(show_label=True, minimum=1, maximum=20, label="CFG Scale:", value="7.5", step=0.5),
        gr.Slider(show_label=True, label="Seed:", minimum=-1, maximum=1000000, value=-1, step=1),
        gr.Dropdown(["yes", "no"], label="Upscale", value="no"),
    ],
    outputs=gr.Image()
).launch()