# `spaces` MUST be imported before torch / any CUDA-initializing package.
try:
    import spaces  # HF ZeroGPU
    GPU = spaces.GPU
except ImportError:  # local fallback
    def GPU(*a, **k):
        def deco(f):
            return f
        return deco if not (a and callable(a[0])) else a[0]

import gradio as gr
import numpy as np
import cv2
import torch
from PIL import Image
from diffusers import QwenImageEditPlusPipeline

MAGENTA_PROMPT = (
    "Change only the wall color to pure magenta. Render the walls as a completely flat, "
    "uniform, solid color with no lighting, no shading,no shadows, no highlights, and no "
    "texture on the walls. Every wall in the room must be the exact same magenta — identical "
    "hue, identical saturation, and identical brightness across all walls, with no variation "
    "between different walls or surfaces. Do not make one wall darker, lighter, or more "
    "saturated than another. Keep everything else exactly the same — furniture, floor, ceiling, "
    "windows, lighting, shadows, textures, and objects unchanged. Do not move or alter any objects."
)

# HSV range for magenta extraction (note: OpenCV hue is 0-179)
HSV_LOW = (125, 91, 90)
HSV_HIGH = (180, 255, 255)


def _build_pipe():
    # On ZeroGPU, build and move to CUDA at GLOBAL scope. ZeroGPU patches
    # torch so module-level .to("cuda") is captured and replayed in the
    # worker. Doing .to("cuda") INSIDE the @GPU function trips an NVML
    # allocator assert, so it must happen here at import time.
    pipe = QwenImageEditPlusPipeline.from_pretrained(
        "Qwen/Qwen-Image-Edit-2511", torch_dtype=torch.bfloat16
    )
    pipe.set_progress_bar_config(disable=None)
    pipe.to("cuda")
    return pipe


PIPE = _build_pipe()


@GPU(duration=120)
def edit_to_magenta(image, steps, true_cfg, seed):
    gen = torch.Generator(device="cuda").manual_seed(int(seed))
    out = PIPE(
        image=[image.convert("RGB")],
        prompt=MAGENTA_PROMPT,
        negative_prompt=" ",
        num_inference_steps=int(steps),
        true_cfg_scale=float(true_cfg),
        generator=gen,
    ).images[0]
    return out


def extract_mask(edited):
    bgr = cv2.cvtColor(np.array(edited.convert("RGB")), cv2.COLOR_RGB2BGR)
    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv, HSV_LOW, HSV_HIGH)
    return mask


def overlay_mask(original, mask):
    base = np.array(original.convert("RGB"))
    h, w = base.shape[:2]
    # The edit (and thus the mask) may differ in resolution from the
    # uploaded image; resize the mask to match before indexing.
    if mask.shape[:2] != (h, w):
        mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
    overlay = base.copy()
    overlay[mask > 0] = (0, 255, 0)
    blended = cv2.addWeighted(base, 0.6, overlay, 0.4, 0)
    return Image.fromarray(blended)


def run(image, steps, true_cfg, seed):
    if image is None:
        raise gr.Error("Please upload an image.")
    edited = edit_to_magenta(image, steps, true_cfg, seed)
    mask = extract_mask(edited)
    overlay = overlay_mask(image, mask)
    return edited, Image.fromarray(mask), overlay


with gr.Blocks(title="Wall Mask Extractor — Qwen-Image-Edit-2511") as demo:
    gr.Markdown(
        "# Wall Mask Extractor\n"
        "Recolors walls to flat magenta with **Qwen-Image-Edit-2511**, then extracts a "
        "binary wall mask via HSV thresholding."
    )
    with gr.Row():
        with gr.Column():
            inp = gr.Image(type="pil", label="Input room image")
            steps = gr.Slider(8, 50, value=40, step=1, label="Inference steps")
            true_cfg = gr.Slider(1.0, 8.0, value=4.0, step=0.5, label="True CFG scale")
            seed = gr.Number(value=0, precision=0, label="Seed")
            btn = gr.Button("Generate mask", variant="primary")
        with gr.Column():
            out_edit = gr.Image(label="Magenta wall edit")
            out_mask = gr.Image(label="Binary wall mask")
            out_overlay = gr.Image(label="Mask overlay")

    btn.click(run, [inp, steps, true_cfg, seed], [out_edit, out_mask, out_overlay])

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