import spaces  # must be imported before torch / any CUDA-touching import

import torch
import numpy as np
import gradio as gr
from PIL import Image, ImageChops
from huggingface_hub import hf_hub_download
from diffusers import DDIMScheduler
from diffusers.models import AutoencoderKL

from removal.v1_2.removal_model import build_removal_model, load_removal_model
from removal.v1_2.pipeline import RemovalSDXLPipeline_BatchMode

MODEL_CONFIG = "config/model_cfg/moebius.yaml"
NUM_EMBEDDINGS = 20
DTYPE = torch.float32

VARIANTS = {
    "Places2 (natural scenes)": "ft_places2",
    "CelebA-HQ (faces)": "ft_celebahq",
    "FFHQ (faces)": "ft_ffhq",
    "Pretrained (general)": "pretrained",
}

# Shared VAE (PixelHacker f8d4) — load on CPU, the pipeline moves it to CUDA.
vae = AutoencoderKL.from_pretrained("hustvl/PixelHacker", subfolder="vae")

# Build one pipeline per fine-tuned variant; all share the same VAE.
PIPELINES = {}
for label, subdir in VARIANTS.items():
    weight_path = hf_hub_download("hustvl/Moebius", f"{subdir}/diffusion_pytorch_model.bin")
    model = build_removal_model(MODEL_CONFIG, NUM_EMBEDDINGS)
    load_removal_model(model, weight_path, device="cpu", dtype=DTYPE)
    scheduler = DDIMScheduler(
        beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
        num_train_timesteps=1000, clip_sample=False,
    )
    PIPELINES[label] = RemovalSDXLPipeline_BatchMode(
        removal_model=model, vae=vae, scheduler=scheduler, device="cuda", dtype=DTYPE,
    )


def _extract(editor):
    """Pull the source image + a binary mask (white = region to inpaint) from an ImageEditor value."""
    if editor is None:
        raise gr.Error("Please upload an image and paint over the region to fill.")
    image = editor["background"].convert("RGB")
    layers = editor.get("layers") or []
    mask = Image.new("L", image.size, 0)
    for layer in layers:
        if layer.mode == "RGBA":
            mask = ImageChops.lighter(mask, layer.split()[-1])
        else:
            mask = ImageChops.lighter(mask, layer.convert("L"))
    mask = mask.point(lambda p: 255 if p > 10 else 0)
    if not mask.getbbox():
        raise gr.Error("The mask is empty — paint over the area you want to inpaint.")
    return image, mask


@spaces.GPU(duration=60)
def inpaint(editor, variant, num_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
    image, mask = _extract(editor)
    pipe = PIPELINES[variant]
    out = pipe(
        [image], [mask],
        image_size=512,
        num_steps=int(num_steps),
        guidance_scale=float(guidance_scale),
        noise_offset=0.0357,
        paste=True,
        compensate=False,
        retry=int(seed),
        mute=False,
    )
    return out[0]


with gr.Blocks(title="Moebius Inpainting") as demo:
    gr.Markdown(
        """# Moebius — 0.2B Lightweight Image Inpainting
        A 0.22B-parameter inpainting model (2% of FLUX.1-Fill-Dev's size) matching 10B-level quality.
        Upload an image, **paint over the region you want to fill**, pick a model variant, and run.

        [Paper](https://arxiv.org/abs/2606.19195) · [Code](https://github.com/hustvl/Moebius) · [Weights](https://huggingface.co/hustvl/Moebius)
        """
    )
    with gr.Row():
        with gr.Column():
            editor = gr.ImageEditor(
                label="Image — paint the area to inpaint",
                type="pil",
                brush=gr.Brush(colors=["#ffffff"], default_size=40, color_mode="fixed"),
                layers=False,
                sources=["upload", "clipboard"],
                height=512,
            )
            variant = gr.Dropdown(
                choices=list(VARIANTS.keys()),
                value="Places2 (natural scenes)",
                label="Model variant",
            )
            with gr.Accordion("Advanced settings", open=False):
                num_steps = gr.Slider(1, 50, value=20, step=1, label="Sampling steps")
                guidance_scale = gr.Slider(1.0, 10.0, value=2.5, step=0.1, label="Guidance scale (CFG)")
                seed = gr.Slider(0, 100000, value=0, step=1, label="Seed")
            run = gr.Button("Inpaint", variant="primary")
        with gr.Column():
            output = gr.Image(label="Result", type="pil", height=512)

    run.click(
        inpaint,
        inputs=[editor, variant, num_steps, guidance_scale, seed],
        outputs=output,
    )

demo.launch()