import gradio as gr
import requests
from PIL import Image
from io import BytesIO
import base64

api_url = "https://5cb20b40-572c-426f-9466-995256f9b6eb.id.repl.co/generate_image"

def generate_image(model="DucHaiten", prompt="", seed=0, negative_prompt="", sampler="k_dpmpp_2s_a", steps=50, guidance_scale=1.0):
    data = "?model=" + model + "&prompt=" + prompt + "&seed=" + str(seed) + "&negative_prompt=" + negative_prompt + "&sampler=" + sampler + "&steps=" + str(steps) + "&guidance_scale=" + str(guidance_scale)
    response = requests.post(api_url + data, json={"guidance_scale": guidance_scale}, timeout=400)
    if response.status_code == 200:
        img_base64 = response.json()["url"]
        img_bytes = base64.b64decode(img_base64)
        img = Image.open(BytesIO(img_bytes))
        return img
    else:
        return None

model_dropdown = gr.components.Dropdown(['3DKX',  'DucHaiten', 'iCoMix', 'Mega Merge Diffusion', 'Midjourney Diffusion', 'mo-di-diffusion',  'NeverEnding Dream',  'Perfect World',  'PFG',  'Poison', 'Project Unreal Engine 5', 'PVC', 'RCNZ Dumb Monkey', 'RealBiter',  'Redshift Diffusion', 'Rev Animated', 'Unstable Ink Dream'], label="Model", default="DucHaiten")
prompt_textbox = gr.components.Textbox(label="Prompt")
seed_number = gr.components.Number(label="Seed", default=0)
negative_prompt_textbox = gr.components.Textbox(label="Negative Prompt", default="")
sampler_dropdown = gr.components.Dropdown(["k_lms", "k_heun", "k_euler", "k_euler_a", "k_dpm_2", "k_dpm_2_a", "DDIM", "k_dpm_fast", "k_dpm_adaptive", "k_dpmpp_2m", "k_dpmpp_2s_a", "k_dpmpp_sde"], label="Sampler", default="k_dpmpp_2s_a")
steps_number = gr.components.Number(label="Steps", default=50)
guidance_scale_number = gr.components.Number(label="Guidance Scale", default=1.0)

def image_fn(model, prompt, seed, negative_prompt, sampler, steps, guidance_scale):
    return generate_image(model, prompt, seed, negative_prompt, sampler, steps, guidance_scale)

output_image = gr.outputs.Image(label="Generated Image", type="pil")

interface = gr.Interface(fn=image_fn, inputs=[model_dropdown, prompt_textbox, seed_number, negative_prompt_textbox, sampler_dropdown, steps_number, guidance_scale_number], outputs=output_image, title="Text to Image Generator", description="Generate images from text prompts")

interface.launch()