app.py

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}")

# ── NSFW LoRA catalog ──────────────────────────────────────────────────────────
LORA_REPO = "wiikoo/Qwen-lora-nsfw"
LORA_CONFIGS = {
    "CockQwen_v3":              "loras/CockQwen-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",
    "LumiNude01a":              "loras/LumiNude01a_CE_QWEN_AIT3k.safetensors",
    "Meta4":                    "loras/Meta4.safetensors",
    "PillowHump_2509":          "loras/PillowHump_2509.safetensors",
    "Qwen4Play_v2":             "loras/Qwen4Play_v2.safetensors",
    "QwenHentai_v3":            "loras/QwenImageHentaiPIV_v3.1.safetensors",
    "Qwen_Real_Nud3s":          "loras/Qwen_Real_Nud3s.safetensors",
    "Qwen_Real_PS":             "loras/Qwen-Real PS_v1_83K.safetensors",
    "Real_Breast_Nipples":      "loras/Real Breast Nipples-QWEN-[rbn]-GMR.safetensors",
    "Ultimate_Breast_Nipples":  "loras/Ultimate Realistic Breast NIPPLES-QWEN-[rab]-GMR.safetensors",
    "ass_up_QWEN":              "loras/ass_up_QWEN.safetensors",
    "bfs_v2_face":              "loras-sfw/face_swap_5500_qwen_image_edit_2509_v1.safetensors",
    "bfs_v2_head":              "loras-sfw/bfs_v2_head_000007000.safetensors",
    "big_nipples_QWEN":         "loras/big_nipples_QWEN.safetensors",
    "bumpynipples":             "loras/bumpynipples1.safetensors",
    "consistence_edit_v1":      "loras-2/consistence_edit_v1.safetensors",
    "consistence_edit_v2":      "loras2/consistence_edit_v2.safetensors",
    "d33p7hroa7":               "loras/d33p7hroa7_qwen.safetensors",
    "goblin_anal_v1":           "loras/goblin_anal_v1_qwen.safetensors",
    "milk_juggs":               "loras/milk_juggs_QWEN.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",
    "remove_clothing":          "loras/qwen_image_edit_remove-clothing_v1.0.safetensors",
    "reclining_nude":           "loras/reclining_nude_v1_000003500.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",
    "goblin_anal_v1":           "anal penetration, spread ass",
    "d33p7hroa7":               "deepthroat, penis deep in mouth",
    "QwenHentai_v3":            "nsfw, anime style, explicit",
}

# Tracks which adapter names have been loaded into the pipeline this session.
# ZeroGPU resets VRAM on every @spaces.GPU call, so we reload as needed.
LOADED_ADAPTERS: set[str] = set()

# ── Helpers ────────────────────────────────────────────────────────────────────
def update_dimensions(image):
    """Return (width, height) scaled so the long edge is 1024 and dims are multiples of 8."""
    if image is None:
        return 1024, 1024
    w, h = image.size
    if w >= h:
        new_w = 1024
        new_h = int(1024 * h / w)
    else:
        new_h = 1024
        new_w = int(1024 * w / h)
    return (new_w // 8) * 8, (new_h // 8) * 8 
    
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, then activate the full stack."""
    if not extra_adapters:
        pipe.disable_lora()
        return

    for name in extra_adapters:
        if name not in LOADED_ADAPTERS and name in LORA_CONFIGS:
            print(f"--- Loading adapter: {name} ---")
            pipe.load_lora_weights(
                LORA_REPO,
                weight_name=LORA_CONFIGS[name],
                adapter_name=name,
            )
            LOADED_ADAPTERS.add(name)

    pipe.enable_lora()
    pipe.set_adapters(extra_adapters, adapter_weights=extra_weights)

# ── Inference ──────────────────────────────────────────────────────────────────
MAX_SEED = np.iinfo(np.int32).max
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=90)
def infer(
    input_image,
    prompt,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    *lora_params,
    progress=gr.Progress(track_tqdm=True),
):
    # ── OOM FIX: Aggressive memory cleanup before inference ──────────────────
    gc.collect()
    torch.cuda.empty_cache()
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    # ─────────────────────────────────────────────────────────────────────────

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

    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))

    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)
    image = input_image.convert("RGB")
    width  = (image.width  // 8) * 8
    height = (image.height // 8) * 8

    try:
        # ── OOM FIX: torch.no_grad() prevents gradient accumulation in VRAM ──
        with torch.no_grad():
            result = pipe(
                image=image,
                prompt=prompt,
                negative_prompt=NEGATIVE_PROMPT,
                height=height,
                width=width,
                num_inference_steps=steps,
                generator=generator,
                true_cfg_scale=guidance_scale,
            ).images[0]
        # ─────────────────────────────────────────────────────────────────────
        return result, seed
    finally:
        # ── OOM FIX: Unload LoRAs and clean up after each inference ──────────
        if extra_adapters:
            pipe.disable_lora()
        gc.collect()
        torch.cuda.empty_cache()
        # ─────────────────────────────────────────────────────────────────────

# ── UI ─────────────────────────────────────────────────────────────────────────
css = """
#col-container { margin: 0 auto; max-width: 960px; }
#main-title h1 { font-size: 2.1em !important; }
/* Preview defaults to cropping (cover-like); contain keeps the full image in frame */
#col-container .image-frame,
#col-container .image-frame > div,
.contain-preview .image-frame,
.contain-preview .image-frame > div {
  overflow: visible !important;
}
#col-container .image-frame img,
.contain-preview .image-frame img {
  object-fit: contain !important;
  height: auto !important;
  max-height: min(75vh, 900px) !important;
  width: auto !important;
  max-width: 100% !important;
  margin-inline: auto;
}
"""

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 · NSFW**", elem_id="main-title")
        gr.Markdown(
            "Base: `Qwen-Image-Edit-2509` + `Qwen-Image-Edit-Rapid-AIO-V4`\n"
            "Add optional extra LoRAs below — selecting one auto-fills its trigger words."
        )

        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. nude woman, remove clothing, detailed skin texture...",
                    lines=3,
                )
                run_button = gr.Button("✨ Edit Image", variant="primary", size="lg")

            with gr.Column():
                output_image = gr.Image(
                    label="Output",
                    interactive=False,
                    format="png",
                    elem_classes=["contain-preview"],
                )

                with gr.Accordion("➕ Extra LoRAs (optional)", open=False):
                    gr.Markdown(
                        "Stack up to 6 LoRAs. Selecting one appends its trigger words to the prompt."
                    )
                    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,
                            )
                            sl = gr.Slider(
                                0.0, 1.5, value=0.75,
                                step=0.05, label="Strength", scale=2,
                            )
                            lora_stack.extend([dd, sl])

                with gr.Accordion("⚙️ Advanced", open=False):
                    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=16, step=1, value=4
                    )

    run_button.click(
        fn=infer,
        inputs=[input_image, prompt, seed, randomize_seed, guidance_scale, steps]
                + 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],
        )

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