# -*- coding: utf-8 -*-
# @时间       : 2024/5/15 18:03
# @作者       : caishilong
# @文件名      : main.py
# @项目名      : ai-platform
# @Software   : PyCharm
import io
import gradio as gr
import numpy as np
import requests
from PIL import Image
from dotenv import load_dotenv
import os
load_dotenv()
SECRET_TOKEN = os.getenv("SECRET_TOKEN")
account = os.getenv('ACCOUNT')
API_BASE_URL = f"https://api.cloudflare.com/client/v4/accounts/{account}/ai/run/"
headers = {"Authorization": f"Bearer {SECRET_TOKEN}"}


def run(model, inputs):
    input = inputs
    response = requests.post(f"{API_BASE_URL}{model}", headers=headers, json=input)
    # 返回的是"content-type": "image/png"
    try:
        print(response.json())
    except:
        pass
    image = Image.open(io.BytesIO(response.content))
    image_array = np.array(image)
    return image_array


# output = run("@cf/stabilityai/stable-diffusion-xl-base-1.0", inputs)
#
# with open("output.png", "wb") as f:
#     f.write(output)



def fake_diffusion(model, prompt, guidance, image: np.ndarray):
    # image = Image.fromarray(image)
    #
    # image = np.array(image)
    # print(image)
    # plt.plot(121)
    # plt.imshow(image)
    # plt.show()
    data = {"prompt": prompt,
            "guidance": guidance,
            # "image": image.tolist()
            }
    output = run(model, data)
    yield output


demo = gr.Interface(fake_diffusion,
                    [
                        gr.Dropdown(["@cf/stabilityai/stable-diffusion-xl-base-1.0",
                                     "@cf/lykon/dreamshaper-8-lcm",
                                     '@cf/bytedance/stable-diffusion-xl-lightning',
                                     '@cf/runwayml/stable-diffusion-v1-5-img2img'],
                                    label="Model", value='@cf/bytedance/stable-diffusion-xl-lightning'),
                        gr.Textbox(label='prompt', value='human playing with a dog'),
                        gr.Slider(1, 20, 20, label='guidance'),
                        # gr.Image()
                    ],
                    outputs="image")

if __name__ == "__main__":
    demo.launch(server_name='0.0.0.0',auth= ("admin", r'qwe[]\123'),)