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

# Hugging Face API 資訊
API_URL = "https://api-inference.huggingface.co/models/KappaNeuro/ukiyo-e-art"
headers = {"Authorization": "Bearer hf_MySpaceToken"}

# 呼叫模型的函數，並增加日誌以排查錯誤
def query(payload):
    print("呼叫模型中...")
    response = requests.post(API_URL, headers=headers, json=payload)

    if response.status_code != 200:
        print(f"Error {response.status_code}: {response.text}")
        return None  # 回傳 None 表示呼叫失敗

    print("模型呼叫成功，正在解析回應...")
    return response.content

# 定義 Gradio 的生成函數
def generate_image(prompt):
    result = query({"inputs": prompt})

    if result is None:
        return "模型回應錯誤，請檢查 API Token 或模型狀態。"

    try:
        image = Image.open(BytesIO(result))  # 將回傳的二進位資料轉為圖片
        return image
    except Exception as e:
        print(f"解析圖片失敗：{e}")
        return "無法生成圖片，請稍後再試。"

# 建立 Gradio 介面
interface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title="浮世繪生成器",
    description="輸入一句話，生成浮世繪風格的藝術圖片。",
)

# 啟動應用程式，設定 share=True 以方便測試
interface.launch(share=True)