import gradio as gr
from PIL import Image
from utils import outpaint_image

DEFAULT_MODEL_ID = "stabilityai/stable-diffusion-2-inpainting"

def generate(image, prompt, directions):
    try:
        yield gr.Textbox.update(value="正在生成扩图中，请稍候...", visible=True), None
        result = outpaint_image(image, prompt, directions, DEFAULT_MODEL_ID)
        yield gr.Textbox.update(visible=False), result
    except Exception as e:
        yield gr.Textbox.update(value=f"❌ 错误：{str(e)}", visible=True), None

with gr.Blocks(title="描述词扩图 App（手机专用）") as demo:
    gr.Markdown("""# 🎨 描述词驱动图像扩图（适配手机）
上传图像 + 描述词，点击按钮即可生成，使用轻量模型 `stabilityai/stable-diffusion-2-inpainting`
""")
    with gr.Row():
        image_input = gr.Image(label="上传原图", type="pil")
        prompt_input = gr.Textbox(label="请输入描述词 (prompt)")

    directions_input = gr.CheckboxGroup(
        choices=["上", "下", "左", "右"],
        label="选择扩展方向（可多选）",
        value=["上", "下", "左", "右"]
    )

    status_output = gr.Textbox(label="状态", visible=False)
    image_output = gr.Image(label="扩图结果")

    run_btn = gr.Button("开始生成")
    run_btn.click(
        fn=generate,
        inputs=[image_input, prompt_input, directions_input],
        outputs=[status_output, image_output]
    )

demo.launch()