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

# ------------------------------------------------------------------------------
# 1) PIPELINE LADEN
# ------------------------------------------------------------------------------
MODEL_ID = "runwayml/stable-diffusion-inpainting"

print("Lade Inpainting-Pipeline... (kann einen Moment dauern)")
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32
)
# CPU-Modus (gratis, aber langsam); wenn GPU verfügbar, pipe.to("cuda") für schnellere Ergebnisse
pipe.to("cpu")

# Wir reduzieren z. B. die Inference-Steps, damit es schneller geht
DEFAULT_STEPS = 20
GUIDANCE_SCALE = 7.5
IMG_SIZE = 512  # Wir verkleinern Bilder auf 512x512, um CPU-Zeit zu sparen

# ------------------------------------------------------------------------------
# 2) INPAINT-FUNKTION mit manuellem Masking
# ------------------------------------------------------------------------------
def inpaint_masked(img, mask, prompt):
    """
    Nimmt ein Originalbild (img) + eine vom User gezeichnete Maske (mask).
    Nur in maskierten Bereichen wird etwas verändert, z. B. Möbel eingefügt.
    Der Rest bleibt unverändert.
    """
    if img is None or mask is None:
        return None

    # 1) Auf 512x512 verkleinern (Bild & Maske gleich groß halten!)
    img = img.resize((IMG_SIZE, IMG_SIZE), resample=Image.LANCZOS)
    mask = mask.resize((IMG_SIZE, IMG_SIZE), resample=Image.LANCZOS)

    # 2) Pipeline-Aufruf
    result = pipe(
        prompt=prompt,
        image=img,
        mask_image=mask,
        num_inference_steps=DEFAULT_STEPS,
        guidance_scale=GUIDANCE_SCALE
    ).images[0]

    return result

# ------------------------------------------------------------------------------
# 3) GRADIO INTERFACE
# ------------------------------------------------------------------------------
def build_app():
    with gr.Blocks() as demo:
        gr.Markdown("# 🏠 Virtuelles Home Staging mit manueller Maske")
        gr.Markdown(
            "1. Lade dein Zimmerfoto hoch.\n"
            "2. **Male in der Maske** die Bereiche aus, wo Möbel erscheinen sollen.\n"
            "   - Weiß = Hier darf die KI etwas ändern.\n"
            "   - Schwarz = Bleibt unverändert.\n"
            "3. Prompt eingeben, auf **Inpaint** klicken."
        )

        with gr.Row():
            with gr.Column():
                input_image = gr.Image(
                    label="(1) Zimmerfoto hochladen",
                    type="pil"
                )
                mask_image = gr.Image(
                    label="(2) Maske hier zeichnen (weiß = veränderbar)",
                    tool="sketch",  # erlaubt freies Zeichnen
                    type="pil"
                )
                prompt_text = gr.Textbox(
                    label="(3) Prompt: Was soll eingefügt werden?",
                    value="Add modern furniture, keep walls the same",
                    lines=2
                )
                run_button = gr.Button("Inpaint")

            with gr.Column():
                result_image = gr.Image(label="Ergebnis")

        # Button-Event: Klick -> inpaint_masked()
        run_button.click(
            fn=inpaint_masked,
            inputs=[input_image, mask_image, prompt_text],
            outputs=[result_image]
        )

        gr.Markdown(
            "**Tipps:**\n"
            "- Standard-Inpainting erfordert weiße Stellen für Bereiche, die verändert werden dürfen.\n"
            "- Du kannst z. B. nur den Fußboden weiß maskieren, damit dort Möbel entstehen.\n"
            "- Für bessere Qualität könntest du `DEFAULT_STEPS` oder `IMG_SIZE` anpassen."
        )

    return demo

# ------------------------------------------------------------------------------
# 4) START
# ------------------------------------------------------------------------------
demo = build_app()

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