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 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 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. Returns the list of adapters that were successfully loaded (some may be skipped if the remote file is missing or corrupted). """ 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: 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=120) 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() # ───────────────────────────────────────────────────────────────────────── # ── ZeroGPU FIX: GPU state resets between @spaces.GPU calls, so we must # clear the adapter tracking set to force reloading on each invocation. LOADED_ADAPTERS.clear() 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 LoRAs and clean up after each inference ────────── if loaded_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 - 2509**", elem_id="main-title") gr.Markdown( "Base: `Qwen-Image-Edit-2509` + `Qwen-Image-Edit-Rapid-AIO-V4`" "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. change clothing...", 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." ) 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=30, 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 )