import gradio as gr
import json
import requests
import io
import os
import base64
from PIL import Image, PngImagePlugin

url = os.environ['url']
status_code = requests.get(url).status_code


with gr.Blocks() as block:
    with gr.Row():
        txt_prompt = gr.Text(label="Prompt",value="1girl",lines=2)
    with gr.Row():
        if status_code == requests.codes.ok:
            btn_generate = gr.Button("Generate",variant="primary")
        else:
            btn_generate = gr.Button("Sorry, not working",variant="secondary")
    with gr.Row():
        image = gr.Image()

    def generate(prompt):
        payload = {
            "seed": -1,
            "sampler_name": "Euler",
            "batch_size": 1,
            "n_iter": 1,
            "steps": 30,
            "cfg_scale": 7,
            "width": 512,
            "height": 768,
            "prompt":prompt,
            "negative_prompt": "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username",
        }
        
        response = requests.post(url=f'{url}/sdapi/v1/txt2img', json=payload)
        
        r = response.json()
    
        for i in r['images']:
            image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))
            return gr.Image.update(image)

    btn_generate.click(
        fn=generate,
        inputs=[txt_prompt],
        outputs=[image]
    )

block.queue().launch()