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

# Hugging Face API 配置
API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-inpainting"
HEADERS = {"Authorization": "Bearer YOUR_HF_hf_kluxSzzsiuuNUVcNvmvjnNbFnMMSwQMQNM"}  # 需要替换为你的 Hugging Face API Key

# 处理函数
def inpaint_image(image, mask, prompt):
    # 将图片转换为字节流
    image_bytes = io.BytesIO()
    mask_bytes = io.BytesIO()
    image.save(image_bytes, format="PNG")
    mask.save(mask_bytes, format="PNG")

    # 发送请求到 Hugging Face API
    response = requests.post(API_URL, headers=HEADERS, files={
        "image": ("image.png", image_bytes.getvalue(), "image/png"),
        "mask": ("mask.png", mask_bytes.getvalue(), "image/png"),
    }, data={"prompt": prompt})

    # 解析返回的图片
    if response.status_code == 200:
        return Image.open(io.BytesIO(response.content))
    else:
        return "Error: " + response.text

# 创建 Gradio 界面
iface = gr.Interface(
    fn=inpaint_image,
    inputs=[
        gr.Image(type="pil", label="原始图片"),
        gr.Image(type="pil", label="掩码图片 (白色区域为编辑区)"),
        gr.Textbox(label="修改描述 (Prompt)")
    ],
    outputs=gr.Image(type="pil", label="结果图片"),
    title="AI 图片编辑 (Stable Diffusion Inpainting)",
    description="上传图片、绘制或上传掩码（白色部分为要修改的区域），并输入修改描述（英文），AI 将根据提示修改图片。",
    # allow_flagging="never" # 暂时注释掉这一行
)
iface.launch()