import streamlit as st
from transformers import pipeline
from PIL import Image

# pipeline = pipeline(task="text-to-image", model="nf-model/GuoFeng4_XL")


# import requests

# API_URL = "http://ai.gitee.com/api/inference-api/models/hf-models/GuoFeng3"
# headers = {
# 	"Authorization": "Bearer eyJpc3MiOiJodHRwczovL2FpLmdpdGVlLmNvbSIsInN1YiI6IjM0MjQ5In0.gbBUBqG7MBy1YuRba-MH5HoLzz7-QpXSpU2xp2ixMU0pg1b1r1TDq9SUmk5v2EIct1Le44llSNqmpWRX2qjNCA",
# 	"X-Compute-Id": "iluvatar/prod-il-sh-if/inference/zk1.small",
# 	"Content-Type": "application/json"
# }

# def query(payload):
# 	response = requests.post(API_URL, headers=headers, json=payload)
# 	return response.content

# image_bytes = query({
#   "inputs": "Astronaut riding a horse"
# })
# print(image_bytes)
# 你可以使用 PIL.Image 访问图像，例如：
# import io
# from PIL import Image
# image = Image.open(io.BytesIO(image_bytes))


from diffusers import StableDiffusionPipeline
import torch

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

# prompt = "a photo of an astronaut riding a horse on mars"

    
# image.save("astronaut_rides_horse.png")


st.title("输入一篇古诗词，我将会为你生成一张壁纸。")

# 创建输入框让用户输入古诗词
poem = st.text_input("请输入古诗词：")

if st.button("生成壁纸"):
    if poem:
        # 使用模型生成图片
        image = pipe(poem).images[0]  
        # result = pipeline(poem)
        # image = Image.fromarray(result[0]['image'])
        
        # 显示生成的图片
        st.image(image, caption='生成的壁纸', use_column_width=True)
    else:
        st.warning("请输入古诗词内容！")

# file_name = st.file_uploader("上传一张图片，将会输出是热狗的概率")

# if file_name is not None:
#     col1, col2 = st.columns(2)

#     image = Image.open(file_name)
#     col1.image(image, use_column_width=True)
#     predictions = pipeline(image)

#     col2.header("Probabilities")
#     for p in predictions:
#         col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")