import os
import gc
import random
import numpy as np
import spaces
import torch
from typing import Iterable
from PIL import Image
from gradio import Server
from gradio.data_classes import FileData
from fastapi.responses import HTMLResponse

# ── 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-2511",
    transformer=QwenImageTransformer2DModel.from_pretrained(
        "prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19",
        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}")

# ── NSFW 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",
}

# Tracks which adapter names have been loaded into the pipeline this session.
LOADED_ADAPTERS: set[str] = set()

# ── 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 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:
                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

app = Server()

@app.api()
@spaces.GPU(duration=120)
def edit_image(
    input_image: FileData,
    prompt: str,
    seed: int,
    randomize_seed: bool,
    guidance_scale: float,
    steps: int,
    loras: list[dict] # List of { "name": str, "strength": float }
) -> FileData:
    gc.collect()
    torch.cuda.empty_cache()
    LOADED_ADAPTERS.clear()

    if input_image is None:
        raise ValueError("Please upload an image.")

    image = Image.open(input_image["path"]).convert("RGB")
    
    # Validate aspect ratio
    w, h = image.size
    ratio = max(w, h) / max(min(w, h), 1)
    if ratio > 4.0:
        raise ValueError(f"Image aspect ratio too extreme ({w}x{h}, ratio {ratio:.1f}:1).")

    extra_adapters = [l["name"] for l in loras if l["name"] != "None" and l["strength"] > 0.05]
    extra_weights = [l["strength"] for l in loras if l["name"] != "None" and l["strength"] > 0.05]

    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:
        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]
        
        output_path = input_image["path"].rsplit(".", 1)[0] + "_edited.png"
        result.save(output_path)
        return FileData(path=output_path)
    finally:
        if loaded_adapters:
            pipe.disable_lora()
        gc.collect()
        torch.cuda.empty_cache()

@app.get("/")
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    if not os.path.exists(html_path):
        return HTMLResponse("index.html not found", status_code=404)
    with open(html_path, "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())

@app.get("/loras")
async def get_loras():
    return sorted(list(LORA_CONFIGS.keys()))

demo = app

if __name__ == "__main__":
    app.launch(show_error=True)