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

# Hugging Face APIトークンを環境変数から取得
api_token = os.getenv("HF_API_KEY")
headers = {"Authorization": f"Bearer {api_token}"}

# 背景を削除する関数
def remove_background(image):
    # Hugging Face APIエンドポイントにリクエストを送信
    image_bytes = BytesIO()
    image.save(image_bytes, format="PNG")
    response = requests.post(
        "https://api-inference.huggingface.co/models/chwshuang/Stable_diffusion_remove_background_model", 
        headers=headers, 
        files={"file": image_bytes.getvalue()}
    )
    
    if response.status_code == 200:
        # 応答を画像に変換して返す
        result_image = Image.open(BytesIO(response.content))
        return result_image
    else:
        # エラー時の処理
        return f"Error: {response.status_code}, {response.text}"

# Gradioインターフェースを作成
interface = gr.Interface(
    fn=remove_background,
    inputs=gr.Image(type="pil"),
    outputs="image",
    title="Background Remover"
)

interface.launch()