Spaces:
Running on Zero
Running on Zero
| 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" | |
| ) | |
| 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 | |
| ) | |