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

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

# 呼叫模型的函數，並處理回應格式
def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)

    if response.status_code != 200:
        return f"Error {response.status_code}: {response.text}"  # 顯示 API 回傳的錯誤訊息

    # 檢查回應格式，如果是 JSON 說明模型未回傳圖片
    if response.headers["Content-Type"] == "application/json":
        response_data = response.json()
        return f"API 回應錯誤: {response_data.get('error', '未知錯誤')}"

    return response.content  # 如果回傳的是圖片，回傳二進位資料

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

    # 如果回應內容是字串（錯誤訊息），直接回傳
    if isinstance(result, str):
        return result

    try:
        image = Image.open(BytesIO(result))  # 將二進位資料轉為圖片
        return image
    except Exception as e:
        return f"圖片解析失敗: {e}"

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

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