File size: 25,014 Bytes
ced6a25
7fb172f
 
ced6a25
 
 
 
 
 
 
 
7fb172f
db296ae
 
7fb172f
 
 
eb521b9
 
db296ae
ced6a25
 
 
 
db296ae
ced6a25
 
 
 
 
 
 
 
 
 
7fb172f
 
ced6a25
 
 
 
 
 
 
 
 
 
 
 
 
db296ae
eb521b9
7fb172f
eb521b9
38b7feb
7fb172f
eb521b9
ced6a25
7fb172f
ced6a25
 
 
 
 
 
 
7fb172f
ced6a25
7fb172f
ced6a25
7fb172f
ced6a25
 
7fb172f
 
 
 
38b7feb
7fb172f
 
 
 
 
 
 
 
38b7feb
7fb172f
38b7feb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7fb172f
 
 
38b7feb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7fb172f
 
3404bfb
 
 
 
7fb172f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ced6a25
7fb172f
 
3404bfb
 
7fb172f
 
 
 
 
 
 
 
 
 
 
 
38b7feb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7fb172f
 
3404bfb
 
 
 
 
 
 
 
7fb172f
 
 
 
 
 
af6ede9
7fb172f
 
 
2bc96b5
 
 
 
 
 
 
 
 
 
7fb172f
 
2bc96b5
7fb172f
 
 
 
 
 
ced6a25
 
 
 
 
 
 
2bc96b5
7fb172f
 
ced6a25
7fb172f
 
 
 
 
3404bfb
 
 
7fb172f
ced6a25
7fb172f
 
 
 
38b7feb
7fb172f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ced6a25
 
 
 
7fb172f
 
 
 
 
 
 
 
 
 
2bc96b5
7fb172f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3404bfb
 
7fb172f
 
 
 
38b7feb
7fb172f
 
2bc96b5
 
 
 
 
7fb172f
 
2bc96b5
7fb172f
 
2bc96b5
 
 
ced6a25
2bc96b5
 
ced6a25
 
7fb172f
 
db296ae
ced6a25
2bc96b5
7fb172f
2bc96b5
 
7fb172f
38b7feb
2bc96b5
 
 
 
 
 
 
ced6a25
7fb172f
ced6a25
7fb172f
 
 
 
 
 
46c11b4
ad73f53
7fb172f
46c11b4
7fb172f
ced6a25
 
99581ee
 
 
 
 
 
 
 
 
 
7fb172f
 
2bc96b5
 
 
38b7feb
7fb172f
 
 
 
 
 
 
 
2bc96b5
7fb172f
 
 
 
 
 
 
2bc96b5
 
 
 
 
 
7fb172f
2bc96b5
 
 
 
 
 
7fb172f
 
ced6a25
 
7fb172f
 
 
 
 
 
ced6a25
 
 
2bc96b5
7fb172f
 
 
 
 
 
 
 
 
 
 
 
99581ee
 
 
 
 
 
 
 
 
 
 
 
7fb172f
 
 
38b7feb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
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",
    "QwenSnofs_1.3": ("Cassanovason69/Qwensnofs", "Qwen_Snofs_1_3.safetensors"),  # ← from other repo
    "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",
    "QwenSnofs_1.3": "nsfw, nude, sex, blowjob, cum, selfie",  # same style as previous versions
    "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 successfully loaded into the pipeline.
# We keep this across @spaces.GPU calls because the LoRA weights stay in RAM
# (only active GPU state is reset by ZeroGPU). This avoids re-downloading
# and re-registering the same adapters on every generation.
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 (only once per app lifetime), then
    activate exactly the requested stack for this inference.
    """
    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} ---")
                cfg = LORA_CONFIGS[name]

                if isinstance(cfg, tuple):
                    # LoRA from a different Hugging Face repo
                    repo_id, weight_name = cfg
                    pipe.load_lora_weights(
                        repo_id,
                        weight_name=weight_name,
                        adapter_name=name,
                    )
                else:
                    # Normal LoRA from the main LORA_REPO
                    pipe.load_lora_weights(
                        LORA_REPO,
                        weight_name=cfg,
                        adapter_name=name,
                    )

                LOADED_ADAPTERS.add(name)
            except Exception as e:
                # If it already exists from a previous call in this Space session,
                # we treat it as success (no need to reload weights).
                if "already exists" in str(e).lower() or "Adapter" in str(e):
                    print(f"Adapter '{name}' already registered — reusing.")
                    LOADED_ADAPTERS.add(name)  # ensure it's tracked
                else:
                    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


def clear_lora_stack():
    """Reset all 6 LoRA dropdowns to 'None' and sliders to 0.75."""
    updates = []
    for _ in range(6):
        updates.append(gr.update(value="None"))
        updates.append(gr.update(value=0.75))
    return updates


# ── Inference ──────────────────────────────────────────────────────────────────
MAX_SEED = np.iinfo(np.int32).max
DEFAULT_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,
    negative_prompt,
    *lora_params,
    progress=gr.Progress(track_tqdm=True),
):
    # ── OOM FIX: Aggressive memory cleanup before inference ──────────────────
    gc.collect()
    torch.cuda.empty_cache()
    # ─────────────────────────────────────────────────────────────────────────

    # Always start with a clean LoRA state (previous active adapters are
    # disabled by the finally block, but this is extra safety).
    pipe.disable_lora()

    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 active LoRAs and clean up after each inference ───
        pipe.disable_lora()
        gc.collect()
        torch.cuda.empty_cache()
        # ─────────────────────────────────────────────────────────────────────


# ── UI ─────────────────────────────────────────────────────────────────────────
css = """
#col-container { margin: 0 auto; max-width: 980px; }
#main-title h1 { font-size: 2.25em !important; letter-spacing: -0.02em; }
.gr-button { transition: all 0.1s ease; }
.gr-button:hover { transform: translateY(-1px); box-shadow: 0 4px 12px rgba(0,0,0,0.15); }
/* Better image preview */
.contain-preview .image-frame img {
  object-fit: contain !important;
  max-height: min(78vh, 920px) !important;
  width: auto !important;
  max-width: 100% !important;
  margin: 0 auto;
  border-radius: 12px;
  box-shadow: 0 10px 30px rgba(0,0,0,0.1);
}
/* LoRA rows alignment */
.gr-row { align-items: center; }
"""

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(
            "**Fast local LoRA-powered image editing** powered by `Qwen-Image-Edit-2509` + `Qwen-Image-Edit-Rapid-AIO-V4` + 50+ community LoRAs.<br>"
            "Upload an image, describe the edit, optionally stack LoRAs, and hit **Edit Image**."
        )

        with gr.Accordion("💡 Quick Tips", open=False):
            gr.Markdown(
                "- **LoRAs auto-add trigger words** to your prompt when selected.<br>"
                "- Keep total LoRA strength under ~2.5 to avoid artifacts.<br>"
                "- Use **Randomize Seed** for variations of the same edit.<br>"
                "- For best results, start with 4-8 steps and CFG 1.0–2.0."
            )

        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():
                with gr.Column():
                    output_image = gr.Image(
                        label="Output",
                        interactive=False,
                        format="png",
                        elem_classes=["contain-preview"],
                    )
                    with gr.Row():
                        use_as_input_btn = gr.Button("↻ Use as Input", size="sm", variant="secondary")
                        clear_output_btn = gr.Button("✕ Clear", size="sm", variant="secondary")

                with gr.Accordion("➕ Extra LoRAs (optional)", open=False):
                    with gr.Row():
                        gr.Markdown("**Stack up to 6 LoRAs.** Select a LoRA to auto-fill trigger words in the prompt.")
                        clear_btn = gr.Button("🧹 Clear All", size="sm", variant="secondary", scale=0)

                    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,
                                interactive=True,
                            )
                            sl = gr.Slider(
                                0.0, 1.5, value=0.75,
                                step=0.05, label="Strength", scale=2,
                            )
                            lora_stack.extend([dd, sl])

                    clear_btn.click(
                        fn=clear_lora_stack,
                        inputs=None,
                        outputs=lora_stack,
                    )

                with gr.Accordion("⚙️ Advanced", open=False):
                    negative_prompt = gr.Textbox(
                        label="Negative Prompt",
                        value=DEFAULT_NEGATIVE_PROMPT,
                        lines=2,
                        placeholder="Customize what to avoid...",
                    )
                    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, negative_prompt]
                + 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],
        )

    # ── Output image actions ─────────────────────────────────────────────────
    use_as_input_btn.click(
        fn=lambda img: img,
        inputs=[output_image],
        outputs=[input_image],
    )
    clear_output_btn.click(
        fn=lambda: gr.update(value=None),
        inputs=None,
        outputs=[output_image],
    )

if __name__ == "__main__":
    demo.queue(max_size=30).launch(
        mcp_server=True, ssr_mode=False, show_error=True
    )