"""
Waifu-Inpaint-XL Gradio App
----------------------------
Free-GPU-friendly inpainting UI for ShinoharaHare/Waifu-Inpaint-XL.
Works as-is on: HF Spaces (ZeroGPU), Kaggle Notebooks, Google Colab.

Setup:
    pip install -r requirements.txt
    huggingface-cli login   # needed once, model is gated

Run:
    python app.py
"""

import spaces  # MUST be imported before torch/anything CUDA-related, ZeroGPU requirement
import os
import torch
import gradio as gr
from diffusers import StableDiffusionXLInpaintPipeline
from PIL import Image

MODEL_ID = "ShinoharaHare/Waifu-Inpaint-XL"
DTYPE = torch.float16

# Load once at startup. Moving to 'cuda' here is fine under ZeroGPU -- the actual
# GPU device is only allocated when a @spaces.GPU-decorated function is called.
pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=DTYPE,
    use_safetensors=True,
)
pipe.to("cuda")
pipe.enable_vae_slicing()
pipe.enable_attention_slicing()


@spaces.GPU(duration=60)  # seconds of GPU time requested per call; raise if you increase steps/variations
def run_inpaint(
    editor_value,          # gr.ImageEditor output: {"background":..., "layers":[...], "composite":...}
    prompt,
    negative_prompt,
    steps,
    guidance,
    num_variations,
    seed,
):
    if editor_value is None or editor_value.get("background") is None:
        raise gr.Error("Upload an image first.")

    base_image = editor_value["background"].convert("RGB")

    # Build mask from the drawn layer (painted area = white = inpaint region)
    if not editor_value.get("layers"):
        raise gr.Error("Paint over the area you want to inpaint (use the brush tool).")
    mask_layer = editor_value["layers"][0]
    mask = mask_layer.split()[-1].convert("L")  # alpha channel -> grayscale mask

    results = []
    base_seed = int(seed) if seed >= 0 else torch.seed()
    for i in range(int(num_variations)):
        gen = torch.Generator(device="cuda").manual_seed(base_seed + i)
        out = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt or None,
            image=base_image,
            mask_image=mask,
            num_inference_steps=int(steps),
            guidance_scale=float(guidance),
            height=base_image.height,
            width=base_image.width,
            generator=gen,
        ).images[0]
        results.append(out)

    return results


with gr.Blocks(title="Waifu-Inpaint-XL") as demo:
    gr.Markdown("## Waifu-Inpaint-XL — paint a mask, describe the change, generate")

    with gr.Row():
        with gr.Column():
            editor = gr.ImageEditor(
                label="Upload image, then paint the mask (brush tool)",
                type="pil",
                brush=gr.Brush(colors=["#ffffff"], default_size=25),
            )
            prompt = gr.Textbox(label="Prompt", placeholder="orange striped sweater, red sparkle eyes")
            negative_prompt = gr.Textbox(label="Negative prompt (optional)", value="blurry, low quality, extra limbs")
            with gr.Row():
                steps = gr.Slider(10, 50, value=28, step=1, label="Steps")
                guidance = gr.Slider(1, 12, value=5.0, step=0.5, label="Guidance scale")
            with gr.Row():
                num_variations = gr.Slider(1, 6, value=1, step=1, label="Variations to generate")
                seed = gr.Number(value=-1, label="Seed (-1 = random)")
            run_btn = gr.Button("Generate", variant="primary")

        with gr.Column():
            gallery = gr.Gallery(label="Results", columns=3, height=500)

    run_btn.click(
        fn=run_inpaint,
        inputs=[editor, prompt, negative_prompt, steps, guidance, num_variations, seed],
        outputs=gallery,
    )

if __name__ == "__main__":
    # Spaces already serves a public URL -- do NOT pass share=True here (errors on Spaces).
    demo.launch()