import os
import gc
import random
import gradio as gr
import numpy as np
import spaces
import torch
from typing import Iterable
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

# ── Theme ──────────────────────────────────────────────────────────────────────
colors.steel_blue = colors.Color(
    name="steel_blue",
    c50="#EBF3F8", c100="#D3E5F0", c200="#A8CCE1", c300="#7DB3D2",
    c400="#529AC3", c500="#4682B4", c600="#3E72A0", c700="#36638C",
    c800="#2E5378", c900="#264364", c950="#1E3450",
)

class SteelBlueTheme(Soft):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.gray,
        secondary_hue: colors.Color | str = colors.steel_blue,
        neutral_hue: colors.Color | str = colors.slate,
        text_size: sizes.Size | str = sizes.text_lg,
        font: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
        ),
        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue, secondary_hue=secondary_hue,
            neutral_hue=neutral_hue, text_size=text_size, font=font, font_mono=font_mono,
        )
        super().set(
            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
            button_primary_text_color="white",
            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
            slider_color="*secondary_500",
            block_title_text_weight="600",
            block_border_width="3px",
            block_shadow="*shadow_drop_lg",
        )

steel_blue_theme = SteelBlueTheme()

# ── Device / dtype ─────────────────────────────────────────────────────────────
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype  = torch.bfloat16

print("CUDA available:", torch.cuda.is_available())
print("Using device:", device)

# ── Model loading (local qwenimage package + FA3) ──────────────────────────────
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2509",
    transformer=QwenImageTransformer2DModel.from_pretrained(
        "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V4",
        torch_dtype=dtype,
        device_map="cuda",
    ),
    torch_dtype=dtype,
).to(device)

# ── OOM FIX: Enable VAE tiling and slicing to bound VRAM usage ─────────────────
pipe.vae.enable_tiling(tile_sample_min_width=256, tile_sample_min_height=256)
pipe.vae.enable_slicing()
# ───────────────────────────────────────────────────────────────────────────────

try:
    pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
    print("Flash Attention 3 Processor set successfully.")
except Exception as e:
    print(f"Warning: Could not set FA3 processor: {e}")

# ── LoRA catalog ──────────────────────────────────────────────────────────
LORA_REPO = "wiikoo/Qwen-lora-nsfw"
LORA_CONFIGS = {
    "CockQwen_v3":              "loras/CockQwen-v3.safetensors",
    "Eva_Qwen_V3":              "loras/Eva_Qwen_V3.safetensors",
    "Facial_Cumshots_V1":       "loras/Facial_Cumshots_For_Qwen_Image_V1.safetensors",
    "HearmemanAI_V3_Breasts":   "loras/HearmemanAI_V3_Rank64_BreastsLoRA_Epoch60.safetensors",
    "HearmemanAI_V4_Breasts":   "loras/HearmemanAI_V4_Rank128_BreastsLoRA_Epoch80.safetensors",
    "InniePussy":               "loras/InniePussy.safetensors",
    "JTT2_5":                   "loras/[QWEN] JTT2_5.safetensors",
    "LumiNude01a":              "loras/LumiNude01a_CE_QWEN_AIT3k.safetensors",
    "MEXX_QWEN_TG300":          "loras/MEXX_QWEN_TG300_23.safetensors",
    "Meta4":                    "loras/Meta4.safetensors",
    "MysticXXX":                "loras/Qwen-MysticXXX-v1.safetensors",
    "Nsfw_Body_V10":            "loras/Qwen_Nsfw_Body_V10-4K.safetensors",
    "Nsfw_Body_V14":            "loras/Qwen_Nsfw_Body_V14-10K.safetensors",
    "OilySkin_V2":              "loras/Oily Skin QWEN V2-GMR.safetensors",
    "PillowHump_2509":          "loras/PillowHump_2509.safetensors",
    "PutItHere_V2":             "loras/Put it here_Qwen edit_V2.0.safetensors",
    "PutItHere_V01":            "loras/put it here_QwenEdit_V0.1.safetensors",
    "Qwen4Play_v2":             "loras/Qwen4Play_v2.safetensors",
    "QwenHentai_v3":            "loras/QwenImageHentaiPIV_v3.1.safetensors",
    "Qwen_Helm":                "loras/Qwen-Image-Helm_v0.1.safetensors",
    "Qwen_NSFW_Beta1":          "loras/Qwen-NSFW.safetensors",
    "Qwen_NSFW_Beta2":          "loras/Qwen-NSFW-Beta2.safetensors",
    "Qwen_NSFW_Beta4":          "loras/Qwen-NSFW-Beta4.safetensors",
    "Qwen_NSFW_Beta5":          "loras/Qwen-NSFW-Beta5.safetensors",
    "Qwen_Real_Nud3s":          "loras/Qwen_Real_Nud3s.safetensors",
    "Qwen_Real_PS":             "loras/Qwen-Real PS_v1_83K.safetensors",
    "QwenSnofs_v1":             "loras/qwen_snofs.safetensors",
    "QwenSnofs_v1_1":           "loras/QwenSnofs1_1.safetensors",
    "Real_Breast_Nipples":      "loras/Real Breast Nipples-QWEN-[rbn]-GMR.safetensors",
    "SendDudes":                "loras/[QWEN] SendDudes.safetensors",
    "SendNudesLite":            "loras/SendNudesLite (Qwen).safetensors",
    "SendNudesPro_Beta":        "loras/[QWEN] Send Nudes Pro - Beta v1.safetensors",
    "Ultimate_Breast_Nipples":  "loras/Ultimate Realistic Breast NIPPLES-QWEN-[rab]-GMR.safetensors",
    "ass_up_QWEN":              "loras/ass_up_QWEN.safetensors",
    "barbell_nipples_QWEN":     "loras/QWEN_jtn_barbell.safetensors",
    "bfs_v2_face":              "loras-sfw/face_swap_5500_qwen_image_edit_2509_v1.safetensors",
    "bfs_v2_focus_face":        "loras-sfw/bfs_v2_000005000.safetensors",
    "bfs_v2_head":              "loras-sfw/bfs_v2_head_000007000.safetensors",
    "big_nipples_QWEN":         "loras/big_nipples_QWEN.safetensors",
    "bumpynipples":             "loras/bumpynipples1.safetensors",
    "cmslt_cum_on_her":         "loras/cmslt_2509_2.safetensors",
    "consistence_edit_v1":      "loras-2/consistence_edit_v1.safetensors",
    "consistence_edit_v2":      "loras2/consistence_edit_v2.safetensors",
    "d33p7hroa7":               "loras/d33p7hroa7_qwen.safetensors",
    "d1ck_p3n1s_V1_1":          "loras/qwen-image_d!ck_P3N1S_LoRA_V1.1.safetensors",
    "goblin_anal_v1":           "loras/goblin_anal_v1_qwen.safetensors",
    "horseshoe_nipple_rings":   "loras/horseshoe_nipple_rings_QWEN.safetensors",
    "jib_nudity_fixer":         "loras/jib_qwen_fix_000002750.safetensors",
    "jillin":                   "loras/jillin1.safetensors",
    "male_nude":                "loras/lora_nudenan_v1.safetensors",
    "milk_juggs":               "loras/milk_juggs_QWEN.safetensors",
    "n00d_b":                   "loras/n00d-b-qwen.safetensors",
    "nsfw_adv_v1":              "loras/qwen-image_nsfw_adv_v1.0.safetensors",
    "p0ssy_lora_v1":            "loras/p0ssy_lora_v1.safetensors",
    "p3nis":                    "loras/p3nis.safetensors",
    "qwen_MCNL":                "loras/qwen_MCNL_v1.0.safetensors",
    "qwen_PENISLORA":           "loras/qwen-PENISLORA.safetensors",
    "qwen_hand_grab":           "loras/qwen_hand_grab_6000s.safetensors",
    "qwen_uncensor":            "loras/qwen_uncensor_000014928.safetensors",
    "reclining_nude":           "loras/reclining_nude_v1_000003500.safetensors",
    "remove_clothing":          "loras/qwen_image_edit_remove-clothing_v1.0.safetensors",
    "royal_treatment_V3":       "loras/royal+treatment+V3.safetensors",
    "sabi_character":           "loras-2/sabi_character_v1.safetensors",
    "snapchat_selfie":          "loras/qwen_image_snapchat.safetensors",
    "uka_qwen":                 "loras/uka_1_qwen.safetensors",
    "ultimate_realistic_breast":"loras/ultimate realistic breast.safetensors",
}

LORA_TRIGGER_WORDS = {
    "Qwen4Play_v2":             "d0gg13, c0wg1rl, r3v3rs3_c0wg1rl, m15510n4ry, bl0wj0b, penis",
    "qwen_MCNL":                "nsfw, cum_on_face, blowjob, cowgirlout, creamp1e, penis, l1ck, missionary, nipples, reversecowgirlpov, vagina",
    "remove_clothing":          "remove her clothing",
    "Qwen_Real_Nud3s":          "nud3",
    "HearmemanAI_V4_Breasts":   "large breasts, hard nipples, erect nipples",
    "HearmemanAI_V3_Breasts":   "large breasts, hard nipples, erect nipples",
    "Ultimate_Breast_Nipples":  "rab",
    "ass_up_QWEN":              "ass up showing pussy and anus",
    "PillowHump_2509":          "Pillow, Humping",
    "InniePussy":               "Innie pussy, Clean shaven, Vertical slit",
    "p0ssy_lora_v1":            "Nude",
    "CockQwen_v3":              "Erect Penis",
    "p3nis":                    "holding a p3nis",
    "qwen_PENISLORA":           "PENISLORA",
    "Facial_Cumshots_V1":       "cum",
    "bfs_v2_head":              "head swap, transfer head from image 1 to image 2",
    "bfs_v2_face":              "keep the face consistent, preserve facial identity",
    "bfs_v2_focus_face":        "head swap from Image 1 to Image 2",
    "goblin_anal_v1":           "anal penetration, spread ass",
    "d33p7hroa7":               "deepthroat, penis deep in mouth",
    "QwenHentai_v3":            "nsfw, anime style, explicit",
    "Eva_Qwen_V3":              "Eva_gothic, in a kneeling position",
    "JTT2_5":                   "massive breasts, large breasts, medium breasts, small breasts",
    "MEXX_QWEN_TG300":          "nsfw, female body",
    "OilySkin_V2":              "oilski",
    "barbell_nipples_QWEN":     "barbell nipple piercings",
    "Qwen_Helm":                "nsfw, anime style",
    "MysticXXX":                "nsfw",
    "Qwen_NSFW_Beta1":          "nsfw",
    "Qwen_NSFW_Beta2":          "nsfw",
    "Qwen_NSFW_Beta4":          "nsfw",
    "Qwen_NSFW_Beta5":          "nsfw",
    "QwenSnofs_v1":             "sex, missionary, cum, cowgirl, reverse cowgirl, selfie, snapchat selfie, prone position, spooning position, undressing",
    "QwenSnofs_v1_1":           "nsfw, nude, sex, blowjob, cum, selfie",
    "Nsfw_Body_V10":            "Hourglass figure, Hairless pussy, Hairly pussy",
    "Nsfw_Body_V14":            "SSS Waistline, Hairless pussy, Hairly pussy",
    "SendNudesLite":            "nude",
    "SendNudesPro_Beta":        "flat chest, small breasts, medium breasts, large breasts, massive breasts, big nipples",
    "SendDudes":                "Penis",
    "cmslt_cum_on_her":         "Put cum on her",
    "horseshoe_nipple_rings":   "horseshoe-ring nipple piercings, circular-barbell nipple piercings",
    "jib_nudity_fixer":         "nude, nipples, vagina",
    "jillin":                   "masturbating",
    "male_nude":                "nudeman",
    "n00d_b":                   "nude, art photography",
    "nsfw_adv_v1":              "nsfw",
    "d1ck_p3n1s_V1_1":          "P3N1S, penis",
    "qwen_uncensor":            "nsfw, cum_on_face, blowjob, cowgirlout, creamp1e, penis, l1ck, missionary, nipples, reversecowgirlpov, vagina",
    "royal_treatment_V3":       "lick ass, blowjob",
    "snapchat_selfie":          "selfie, snapchat",
    "ultimate_realistic_breast":"urb, realistic breast",
}

# Tracks which adapter names have been successfully loaded into the pipeline.
# We keep this across @spaces.GPU calls because the LoRA weights stay in RAM
# (only active GPU state is reset by ZeroGPU). This avoids re-downloading
# and re-registering the same adapters on every generation.
LOADED_ADAPTERS: set[str] = set()

# ── Helpers ────────────────────────────────────────────────────────────────────
def append_triggers(current_prompt: str, lora_name: str) -> str:
    """Append a LoRA's trigger words to the prompt (no duplicates)."""
    if lora_name == "None":
        return current_prompt
    triggers = LORA_TRIGGER_WORDS.get(lora_name, "")
    if not triggers:
        return current_prompt
    existing = {w.strip().lower() for w in current_prompt.replace(",", " ").split()}
    new_words = [w.strip() for w in triggers.split(",")
                 if w.strip().lower() not in existing and w.strip()]
    if not new_words:
        return current_prompt
    sep = ", " if current_prompt.strip() else ""
    return current_prompt.rstrip(", ") + sep + ", ".join(new_words)


def load_and_apply_stack(extra_adapters: list[str], extra_weights: list[float]):
    """Lazy-load any unseen adapters (only once per app lifetime), then
    activate exactly the requested stack for this inference.
    """
    if not extra_adapters:
        pipe.disable_lora()
        return [], []

    loaded, weights_out = [], []
    for name, weight in zip(extra_adapters, extra_weights):
        if name not in LORA_CONFIGS:
            continue
        if name not in LOADED_ADAPTERS:
            try:
                print(f"--- Loading adapter: {name} ---")
                pipe.load_lora_weights(
                    LORA_REPO,
                    weight_name=LORA_CONFIGS[name],
                    adapter_name=name,
                )
                LOADED_ADAPTERS.add(name)
            except Exception as e:
                # If it already exists from a previous call in this Space session,
                # we treat it as success (no need to reload weights).
                if "already exists" in str(e).lower() or "Adapter" in str(e):
                    print(f"Adapter '{name}' already registered — reusing.")
                    LOADED_ADAPTERS.add(name)  # ensure it's tracked
                else:
                    print(f"WARNING: Failed to load LoRA '{name}': {e}")
                    continue
        loaded.append(name)
        weights_out.append(weight)

    if loaded:
        pipe.enable_lora()
        pipe.set_adapters(loaded, adapter_weights=weights_out)
    else:
        pipe.disable_lora()

    return loaded, weights_out


def clear_lora_stack():
    """Reset all 6 LoRA dropdowns to 'None' and sliders to 0.75."""
    updates = []
    for _ in range(6):
        updates.append(gr.update(value="None"))
        updates.append(gr.update(value=0.75))
    return updates


# ── Inference ──────────────────────────────────────────────────────────────────
MAX_SEED = np.iinfo(np.int32).max
DEFAULT_NEGATIVE_PROMPT = (
    "worst quality, low quality, bad anatomy, bad hands, text, error, "
    "missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, "
    "signature, watermark, username, blurry"
)

@spaces.GPU(duration=120)
def infer(
    input_image,
    prompt,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    negative_prompt,
    *lora_params,
    progress=gr.Progress(track_tqdm=True),
):
    # ── OOM FIX: Aggressive memory cleanup before inference ──────────────────
    gc.collect()
    torch.cuda.empty_cache()
    # ─────────────────────────────────────────────────────────────────────────

    # Always start with a clean LoRA state (previous active adapters are
    # disabled by the finally block, but this is extra safety).
    pipe.disable_lora()

    if input_image is None:
        raise gr.Error("Please upload an image.")

    # ── Validate aspect ratio ────────────────────────────────────────────────
    # Extreme ratios (>4:1) produce degenerate latent shapes that crash the
    # transformer or produce garbage.  Reject early with a clear message.
    image = input_image.convert("RGB")
    w, h = image.size
    ratio = max(w, h) / max(min(w, h), 1)
    if ratio > 4.0:
        raise gr.Error(
            f"Image aspect ratio too extreme ({w}x{h}, ratio {ratio:.1f}:1). "
            "Please use an image with aspect ratio ≤ 4:1."
        )
    # ─────────────────────────────────────────────────────────────────────────

    extra_adapters, extra_weights = [], []
    for i in range(0, len(lora_params), 2):
        name, strength = lora_params[i], lora_params[i + 1]
        if name != "None" and float(strength) > 0.05:
            extra_adapters.append(name)
            extra_weights.append(float(strength))

    loaded_adapters, _ = load_and_apply_stack(extra_adapters, extra_weights)

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)

    try:
        # Pipeline's __call__ is already decorated with @torch.no_grad().
        # Do NOT use torch.inference_mode() here — it breaks LoRA in-place
        # weight scaling (scale_lora_layers / unscale_lora_layers).
        # Let the pipeline auto-calculate output dimensions from the input
        # image's aspect ratio (targeting 1MP, snapped to multiples of 32).
        result = pipe(
            image=image,
            prompt=prompt,
            negative_prompt=negative_prompt if guidance_scale > 1.0 else None,
            num_inference_steps=steps,
            generator=generator,
            true_cfg_scale=guidance_scale,
        ).images[0]
        return result, seed
    except torch.cuda.OutOfMemoryError:
        gc.collect()
        torch.cuda.empty_cache()
        raise gr.Error(
            "GPU out of memory. Try reducing inference steps or using fewer LoRAs."
        )
    except RuntimeError as e:
        if "CUDA" in str(e) or "out of memory" in str(e).lower():
            gc.collect()
            torch.cuda.empty_cache()
            raise gr.Error(f"GPU error: {e}")
        raise gr.Error(f"Inference failed: {e}")
    finally:
        # ── OOM FIX: Unload active LoRAs and clean up after each inference ───
        pipe.disable_lora()
        gc.collect()
        torch.cuda.empty_cache()
        # ─────────────────────────────────────────────────────────────────────

# ── UI ─────────────────────────────────────────────────────────────────────────
css = """
#col-container { margin: 0 auto; max-width: 980px; }
#main-title h1 { font-size: 2.25em !important; letter-spacing: -0.02em; }
.gr-button { transition: all 0.1s ease; }
.gr-button:hover { transform: translateY(-1px); box-shadow: 0 4px 12px rgba(0,0,0,0.15); }
/* Better image preview */
.contain-preview .image-frame img {
  object-fit: contain !important;
  max-height: min(78vh, 920px) !important;
  width: auto !important;
  max-width: 100% !important;
  margin: 0 auto;
  border-radius: 12px;
  box-shadow: 0 10px 30px rgba(0,0,0,0.1);
}
/* LoRA rows alignment */
.gr-row { align-items: center; }
"""

LORA_NAMES = ["None"] + sorted(LORA_CONFIGS.keys())

with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# **Qwen-Image-Edit • 2509**", elem_id="main-title")
        gr.Markdown(
            "**Fast local LoRA-powered image editing** powered by `Qwen-Image-Edit-2509` + `Qwen-Image-Edit-Rapid-AIO-V4` + 50+ community LoRAs.<br>"
            "Upload an image, describe the edit, optionally stack LoRAs, and hit **Edit Image**."
        )
        with gr.Accordion("💡 Quick Tips", open=False):
            gr.Markdown(
                "- **LoRAs auto-add trigger words** to your prompt when selected.<br>"
                "- Keep total LoRA strength under ~2.5 to avoid artifacts.<br>"
                "- Use **Randomize Seed** for variations of the same edit.<br>"
                "- For best results, start with 4-8 steps and CFG 1.0–2.0."
            )

        with gr.Row(equal_height=False):
            with gr.Column():
                input_image = gr.Image(
                    label="Input Image",
                    type="pil",
                    elem_classes=["contain-preview"],
                )
                prompt = gr.Textbox(
                    label="Edit Prompt",
                    placeholder="e.g. change clothing...",
                    lines=3,
                )
                run_button = gr.Button("✨ Edit Image", variant="primary", size="lg")

            with gr.Column():
                with gr.Column():
                    output_image = gr.Image(
                        label="Output",
                        interactive=False,
                        format="png",
                        elem_classes=["contain-preview"],
                    )
                    with gr.Row():
                        use_as_input_btn = gr.Button("↻ Use as Input", size="sm", variant="secondary")
                        clear_output_btn = gr.Button("✕ Clear", size="sm", variant="secondary")

                with gr.Accordion("➕ Extra LoRAs (optional)", open=False):
                    with gr.Row():
                        gr.Markdown("**Stack up to 6 LoRAs.** Select a LoRA to auto-fill trigger words in the prompt.")
                        clear_btn = gr.Button("🧹 Clear All", size="sm", variant="secondary", scale=0)
                    lora_stack = []
                    for i in range(6):
                        with gr.Row():
                            dd = gr.Dropdown(
                                choices=LORA_NAMES,
                                value="None",
                                label=f"LoRA {i + 1}",
                                scale=3,
                                interactive=True,
                            )
                            sl = gr.Slider(
                                0.0, 1.5, value=0.75,
                                step=0.05, label="Strength", scale=2,
                            )
                            lora_stack.extend([dd, sl])

                    clear_btn.click(
                        fn=clear_lora_stack,
                        inputs=None,
                        outputs=lora_stack,
                    )

                with gr.Accordion("⚙️ Advanced", open=False):
                    negative_prompt = gr.Textbox(
                        label="Negative Prompt",
                        value=DEFAULT_NEGATIVE_PROMPT,
                        lines=2,
                        placeholder="Customize what to avoid...",
                    )
                    seed = gr.Slider(
                        label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0
                    )
                    randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                    guidance_scale = gr.Slider(
                        label="CFG Scale", minimum=1.0, maximum=5.0, step=0.1, value=1.0
                    )
                    steps = gr.Slider(
                        label="Steps", minimum=1, maximum=30, step=1, value=4
                    )

    run_button.click(
        fn=infer,
        inputs=[input_image, prompt, seed, randomize_seed, guidance_scale, steps, negative_prompt]
                + lora_stack,
        outputs=[output_image, seed],
    )

    # Auto-fill trigger words when a LoRA is selected
    for i in range(0, len(lora_stack), 2):
        lora_stack[i].change(
            fn=append_triggers,
            inputs=[prompt, lora_stack[i]],
            outputs=[prompt],
        )

    # ── Output image actions ─────────────────────────────────────────────────
    use_as_input_btn.click(
        fn=lambda img: img,
        inputs=[output_image],
        outputs=[input_image],
    )

    clear_output_btn.click(
        fn=lambda: gr.update(value=None),
        inputs=None,
        outputs=[output_image],
    )

if __name__ == "__main__":
    demo.queue(max_size=30).launch(
        mcp_server=True, ssr_mode=False, show_error=True
    )