import gradio as gr
from PIL import Image
import torch
from diffusers import AutoPipelineForImage2Image

# 1. Ładujemy model WindowSeat z Hugging Face
MODEL_ID = "huawei-bayerlab/windowseat-reflection-removal-v1-0"

pipe = AutoPipelineForImage2Image.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
).to("cuda" if torch.cuda.is_available() else "cpu")

# 2. Funkcja, która przyjmuje jedno zdjęcie i zwraca przetworzone
def remove_reflection(img: Image.Image) -> Image.Image:
    result = pipe(
        image=img,
        prompt="",          # nie potrzebujemy promptu, model wie co robić
        strength=1.0,
        num_inference_steps=4,
        guidance_scale=1.0,
    ).images[0]
    return result

# 3. Gradio – proste API: jedno wejście (image), jedno wyjście (image)
demo = gr.Interface(
    fn=remove_reflection,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="WindowSeat Reflection Removal - API",
    description="Upload a photo with window reflections – get cleaned version.",
)

if __name__ == "__main__":
    demo.launch()