import subprocess
import sys
import os

# ── 1. Clone Moebius repo ────────────────────────────────────────────────────
MOEBIUS_DIR = "/app/Moebius"
if not os.path.exists(MOEBIUS_DIR):
    subprocess.run(
        ["git", "clone", "https://github.com/hustvl/Moebius.git", MOEBIUS_DIR],
        check=True,
    )

# Moebius config uses relative paths — must run from its root
os.chdir(MOEBIUS_DIR)
sys.path.insert(0, MOEBIUS_DIR)

# ── 2. Download weights ──────────────────────────────────────────────────────
from huggingface_hub import hf_hub_download, snapshot_download

VARIANTS = ["pretrained", "ft_places2", "ft_celebahq", "ft_ffhq"]

for variant in VARIANTS:
    weight_path = f"weight/Moebius/{variant}/diffusion_pytorch_model.bin"
    os.makedirs(f"weight/Moebius/{variant}", exist_ok=True)
    if not os.path.exists(weight_path):
        hf_hub_download(
            repo_id="hustvl/Moebius",
            filename=f"{variant}/diffusion_pytorch_model.bin",
            local_dir="weight/Moebius",
        )

VAE_DIR = "weight/vae"
os.makedirs(VAE_DIR, exist_ok=True)
if not os.path.exists(os.path.join(VAE_DIR, "config.json")):
    snapshot_download(
        repo_id="stabilityai/sd-vae-ft-mse",
        local_dir=VAE_DIR,
        ignore_patterns=["*.msgpack", "*.h5", "flax_model*"],
    )

# ── 3. Build pipelines (one per variant) ────────────────────────────────────
import torch
from types import SimpleNamespace
from infer.utils import build_pipeline

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

pipelines = {}
for variant in VARIANTS:
    args = SimpleNamespace(
        model_config="config/model_cfg/moebius.yaml",
        model_weight=f"weight/Moebius/{variant}/diffusion_pytorch_model.bin",
        device=DEVICE,
    )
    pipelines[variant] = build_pipeline(args)

# ── 4. Gradio UI ─────────────────────────────────────────────────────────────
import gradio as gr
from PIL import Image
import numpy as np


def remove_object(image_dict, variant, image_size, num_steps, guidance_scale):
    # Force square — the lambda attention assumes h == w
    size = (image_size, image_size)
    image = image_dict["background"].convert("RGB").resize(size, Image.LANCZOS)
    mask_layer = image_dict["layers"][0].convert("RGBA").resize(size, Image.NEAREST)
    mask = Image.fromarray(np.array(mask_layer)[:, :, 3]).convert("L")

    results = pipelines[variant](
        input_image_list=[image],
        input_mask_list=[mask],
        image_size=image_size,
        num_steps=num_steps,
        guidance_scale=guidance_scale,
        paste=True,       # paste result back onto original — preserves detail outside mask
        compensate=True,  # blend edges
        mute=True,
    )
    return results[0]


VARIANT_LABELS = {
    "pretrained":  "Pretrained (general)",
    "ft_places2":  "Places2 (scenes & backgrounds)",
    "ft_celebahq": "CelebA-HQ (faces)",
    "ft_ffhq":     "FFHQ (portraits)",
}

with gr.Blocks(title="Moebius Object Removal") as demo:
    gr.Markdown(
        "## Moebius — Lightweight Object Removal\n"
        "Paint over the object you want removed, then click **Remove**.\n"
        "([Paper](https://arxiv.org/abs/2606.19195) · "
        "[GitHub](https://github.com/hustvl/Moebius))"
    )
    with gr.Row():
        with gr.Column():
            canvas = gr.ImageEditor(
                label="Upload image & paint mask",
                type="pil",
                brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed"),
            )
            variant = gr.Radio(
                choices=list(VARIANT_LABELS.values()),
                value=VARIANT_LABELS["ft_places2"],
                label="Model variant",
            )
            with gr.Accordion("Advanced", open=False):
                image_size     = gr.Slider(256, 1024, value=512, step=64,  label="Image size")
                num_steps      = gr.Slider(5,   50,   value=20,  step=1,   label="Diffusion steps")
                guidance_scale = gr.Slider(1.0, 10.0, value=4.5, step=0.5, label="Guidance scale")
            btn = gr.Button("Remove", variant="primary")
        with gr.Column():
            output = gr.Image(label="Result")

    # Map display label back to key
    label_to_key = {v: k for k, v in VARIANT_LABELS.items()}

    btn.click(
        fn=lambda img, var, sz, steps, cfg: remove_object(
            img, label_to_key[var], sz, steps, cfg
        ),
        inputs=[canvas, variant, image_size, num_steps, guidance_scale],
        outputs=output,
    )

demo.launch()