Add batch processing app.py - upload multiple images, same prompt, view/download all results

app.py

@@ -0,0 +1,400 @@
+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)