import os
import requests
import numpy as np
from PIL import Image
import gradio as gr
from io import BytesIO


def gen_image(desc: str, style: str):
    """generate the image from the wukong huahua model of ascend server in Wuhan AICC
    Args:
        desc(str): the input description text
    """
    if not desc:
        return
    access_token = os.environ['token']
    headers = {'content-type': "application/json", 'X-Subject-Token': access_token}

    # url = "https://a2f051d4cabf45f885d7b0108edc9b9c.infer.ovaijisuan.com/v1/infers/dce9ad51-7cde-4eeb-8291-ae29f267ed2c/wukong_hf"
    url = "https://a2f051d4cabf45f885d7b0108edc9b9c.infer.ovaijisuan.com/v1/infers/975eedfd-6e15-4571-8ca9-b945da0da24b/wukong_hf"
    body = {
        "user_name": 'huggingface',
        "desc": desc,
        "style": style
    }

    resp_data = requests.post(url, json=body, headers=headers)
    print(resp_data['status'])
    if resp_data['status'] != 200:
        return
    # img_np = np.array(resp_data['output_image_list'][0])
    # image = Image.fromarray(np.uint8(img_np))
    img_rep = requests.get(resp_data['output_image_url'][0])
    image = Image.open(BytesIO(img_rep.content))
    image_np = np.asarray(image)

    return image_np


demo = gr.Interface(
    fn=gen_image,
    inputs=["text", "text"],
    outputs=gr.outputs.Image(type='numpy'),
)
demo.launch()