app.py

import os
import gc
import random
import tempfile
import zipfile
import gradio as gr
import numpy as np
import spaces
import torch
from typing import Iterable
from PIL import Image
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 ──────────────────────────────────────────────────────────────
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)

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 = {
    "remove_clothing": "remove her clothing",
    "Qwen_Real_Nud3s": "nud3",
    "HearmemanAI_V4_Breasts": "large breasts, hard nipples, erect nipples",
    "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",
    "OilySkin_V2": "oilski",
    "MysticXXX": "nsfw",
    "SendNudesLite": "nude",
}

LOADED_ADAPTERS: set[str] = set()

# ── Helpers ────────────────────────────────────────────────────────────────────
def append_triggers(current_prompt: str, lora_name: str) -> str:
    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]):
    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:
                pipe.load_lora_weights(LORA_REPO, weight_name=LORA_CONFIGS[name], adapter_name=name)
                LOADED_ADAPTERS.add(name)
            except Exception as e:
                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


# ── 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=300)
def infer_batch(
    input_images,
    prompt,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    *lora_params,
    progress=gr.Progress(track_tqdm=True),
):
    """Process multiple images with the same prompt, one at a time."""
    gc.collect()
    torch.cuda.empty_cache()
    LOADED_ADAPTERS.clear()

    if not input_images:
        raise gr.Error("Please upload at least one image.")

    # Parse LoRA params
    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)

    results = []
    current_seed = seed

    for idx, img_path in enumerate(progress.tqdm(input_images, desc="Processing images")):
        # Load image
        if isinstance(img_path, str):
            image = Image.open(img_path).convert("RGB")
        elif isinstance(img_path, Image.Image):
            image = img_path.convert("RGB")
        else:
            image = Image.open(img_path).convert("RGB")

        # Validate aspect ratio
        w, h = image.size
        ratio = max(w, h) / max(min(w, h), 1)
        if ratio > 4.0:
            print(f"Skipping image {idx+1}: aspect ratio too extreme ({w}x{h})")
            continue

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

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

        try:
            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]
            results.append(result)
        except torch.cuda.OutOfMemoryError:
            gc.collect()
            torch.cuda.empty_cache()
            print(f"OOM on image {idx+1}, skipping...")
            continue
        except Exception as e:
            print(f"Error on image {idx+1}: {e}")
            continue

        # Clean between images
        gc.collect()
        torch.cuda.empty_cache()

    # Cleanup
    if loaded_adapters:
        pipe.disable_lora()
    gc.collect()
    torch.cuda.empty_cache()

    if not results:
        raise gr.Error("All images failed to process.")

    # Create zip file for download
    zip_path = create_zip(results)

    return results, zip_path, f"✅ Successfully processed {len(results)}/{len(input_images)} images"


def create_zip(images: list) -> str:
    """Create a zip file of all processed images for download."""
    tmp_dir = tempfile.mkdtemp()
    zip_path = os.path.join(tmp_dir, "edited_images.zip")
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        for i, img in enumerate(images):
            img_path = os.path.join(tmp_dir, f"edited_{i+1:03d}.png")
            img.save(img_path, "PNG")
            zf.write(img_path, f"edited_{i+1:03d}.png")
    return zip_path


# ── UI ─────────────────────────────────────────────────────────────────────────
css = """
#col-container { margin: 0 auto; max-width: 1100px; }
#main-title h1 { font-size: 2.1em !important; }
"""

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 — Batch Processing**", elem_id="main-title")
        gr.Markdown(
            "Upload **multiple images**, apply the **same prompt** to all of them, "
            "view results in a gallery, and **download all as a ZIP**.\n\n"
            "Base: `Qwen-Image-Edit-2509` + `Qwen-Image-Edit-Rapid-AIO-V4` | Optional LoRA stacking"
        )

        with gr.Row():
            # LEFT: Inputs
            with gr.Column(scale=1):
                input_images = gr.File(
                    label="📁 Upload Images (select multiple)",
                    file_count="multiple",
                    file_types=["image"],
                    type="filepath",
                )
                prompt = gr.Textbox(
                    label="✏️ Edit Prompt (applied to ALL images)",
                    placeholder="e.g. remove background, change hair color to blonde, add sunglasses...",
                    lines=3,
                )
                run_button = gr.Button("🚀 Process All Images", variant="primary", size="lg")

                with gr.Accordion("➕ Extra LoRAs (optional)", open=False):
                    gr.Markdown("Stack up to 4 LoRAs.")
                    lora_stack = []
                    for i in range(4):
                        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 Settings", 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=30, step=1, value=4)

            # RIGHT: Outputs
            with gr.Column(scale=1):
                status_text = gr.Markdown("*Upload images and click Process to start*")
                output_gallery = gr.Gallery(
                    label="📸 Processed Results",
                    columns=2,
                    rows=3,
                    height="auto",
                    object_fit="contain",
                    show_download_button=True,
                )
                download_zip = gr.File(label="⬇️ Download All (ZIP)", visible=True)

    # Wire up the batch processing
    run_button.click(
        fn=infer_batch,
        inputs=[input_images, prompt, seed, randomize_seed, guidance_scale, steps] + lora_stack,
        outputs=[output_gallery, download_zip, status_text],
    )

    # Auto-fill trigger words when LoRA 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=10).launch(ssr_mode=False, show_error=True)