# app.py
import gradio as gr
from huggingface_hub import InferenceClient
import os

HF_TOKEN = os.environ["HF_TOKEN"]
client = InferenceClient(token=HF_TOKEN)

def add_curtains(editor_data, style="classic", color="beige"):
    if editor_data is None:
        return "❌ Пожалуйста, загрузите фото и нарисуйте область окна."
    
    image = editor_data["image"]      # PIL.Image
    mask = editor_data["mask"]        # PIL.Image (чёрно-белая маска)

    if image is None or mask is None:
        return "❌ Нарисуйте маску на окне."

    prompt = f"{style} {color} curtains on the window, photorealistic, natural lighting"

    try:
        result = client.inpainting(
            image=image,
            mask_image=mask,
            prompt=prompt,
            model="diffusers/stable-diffusion-xl-1.0-inpainting-0.1"
        )
        return result
    except Exception as e:
        return f"❌ Ошибка: {str(e)}"

demo = gr.Interface(
    fn=add_curtains,
    inputs=[
        gr.ImageEditor(
            label="1. Загрузите фото окна\n2. Нажмите 'Edit' → выберите кисть → закрасьте область окна белым",
            type="pil"
        ),
        gr.Radio(["classic", "modern", "roman"], value="classic", label="Стиль штор"),
        gr.Textbox(value="beige", label="Цвет (например: white, navy blue)")
    ],
    outputs=gr.Image(label="Результат со шторами"),
    title="🪟 Виртуальная примерка штор",
    description="Нарисуйте маску на окне — ИИ добавит туда шторы!"
)

demo.launch()