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)