app.py

# ===== LINE 1: MUST BE SPACES FOR ZEROGPU =====
try:
    import spaces
    SPACES_AVAILABLE = True
    print("✅ ZeroGPU Mode Enabled")
except ImportError:
    SPACES_AVAILABLE = False
    print("⚠️ Running without Spaces optimization")

import sys
import os
import re
import random
import gc
import logging
import time
import io
import threading
from datetime import datetime

import torch
from PIL import Image
import gradio as gr
from huggingface_hub import hf_hub_download, HfApi, CommitOperationAdd
from huggingface_hub.utils import HfHubHTTPError

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("NextGenEngine")
logger.info("🚀 JustinD Fiction Lab boot (21+ · multi-select stacks · optional quad chefs)")

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8")

BACKUP_DATASET_ID = "JustinD/generation-images"
LAST_GEN_TIME = {}
COOLDOWN_SECONDS = 12
LOG_VISITOR_UPLOADS = os.environ.get("LOG_VISITOR_UPLOADS", "0").strip().lower() in {"1", "true", "yes"}
MAX_STACK_BATCH_PRODUCT = 6

HF_REPO_ID = "AriaVale2/my-safetensors-backup"

# Single-file chefs (Asian realism + proven NSFW pantry)
MODEL_REPOS = {}

# Newer Diffusers Asian NSFW / realism chefs (2025-2026)
DIFFUSERS_MODELS = {
    "WAI NSFW Illustrious v150": "John6666/wai-nsfw-illustrious-sdxl-v150-sdxl",
    "Janku v5 NSFW Illustrious": "John6666/janku-v5-nsfw-trained-noobai-rou-wei-illustrious-xl-v50-sdxl",
    "YAAI WAI NSFW Illustrious": "John6666/yaai-yet-another-anime-illustrious-model-full-fine-tune-based-on-wai-nsfw-v20-sdxl",
    "PornWorks Real Porn Illustrious": "John6666/porn-works-real-porn-photo-realistic-nsfw-sdxl-pony-illustrious-chekpoint-illustrious-sdxl",
    "Hardcore Asian Cosplay Illustrious": "John6666/hardcore-asian-cosplay-porn-il-v11-sdxl",
    "PornMaster Pro Noob Illustrious": "John6666/pornmaster-pro-noobv15vae-sdxl",
    "PornMaster Anime Illustrious": "John6666/pornmaster-anime-illustrious-v2-sdxl",
    "Pornzilla Hentai Illustrious": "John6666/pornzilla-hentai-v45-sdxl",
    "Reed NSFW Illustrious v3": "John6666/reed-nsfw-illustrious-sdxl-v30-il-sdxl",
    "Unearthly NSFW Illustrious": "John6666/unearthly-nsfw-illustrious-sdxl-20250610b-sdxl",
    "Perfection Asian Illustrious NSFW": "John6666/perfection-asian-ilxl-illustrious-xl-nsfw-sfw-checkpoint-v10-sdxl",
    "Ramthrusts NSFW Illustrious": "John6666/ramthrusts-nsfw-red-shift-mix-v199-sdxl",
    "Lustify NSFW GGWP v7": "John6666/lustify-sdxl-nsfw-checkpoint-ggwp-v7-sdxl",
    "Lustify NSFW Endgame": "John6666/lustify-sdxl-nsfwsfw-endgame-sdxl",
    "Lustify NSFW One Last Time": "John6666/lustify-sdxl-nsfw-checkpoint-olt-one-last-time-sdxl",
    "Spicy Realism NSFW Mix v3": "John6666/spicy-realism-nsfw-mix-v30-sdxl",
    "Fucktastic Real Porn Pony v5.1": "John6666/fucktastic-real-checkpoint-pony-pdxl-porn-realistic-nsfw-sfw-51-sdxl",
    "OhMySister Beautiful Asian Woman": "John6666/oms-oh-my-sister-beautiful-asian-woman-v10-sdxl",
    "Asian E-Girl NSFW XL": "John6666/asian-egirl-xl-realistic-uwu-egirl-tattoo-model-sdxl-nsfw-v10-sdxl",
}

SDXL_CHOICES = list(DIFFUSERS_MODELS.keys()) + list(MODEL_REPOS.keys())
SKIP_CHEF = "(skip)"
CHEF_OPTIONAL_CHOICES = [SKIP_CHEF] + SDXL_CHOICES

# LoRA pantry from known SDXL repos. Fail soft if download/load breaks.
LORA_NONE = "None"
LORA_REPOS = {
    LORA_NONE: None,
    "sexy-slider (Soft Domme heat)": ("ntc-ai/SDXL-LoRA-slider.sexy", "sexy.safetensors"),
    "oily-bodystocking": ("Muapi/oily-bodystocking-sdxl-and-pony", "oily-bodystocking-sdxl-and-pony.safetensors"),
    "super-skinny-body": (
        "Muapi/super-skinny-body-for-pony-illustrious-sdxl-and-sd1.5",
        "super-skinny-body-for-pony-illustrious-sdxl-and-sd1.5.safetensors",
    ),
    "glass-body": ("Muapi/transparent-glass-body-lora-1.5-sdxl", "transparent-glass-body-lora-1.5-sdxl.safetensors"),
}
LORA_CHOICES = list(LORA_REPOS.keys())

# Named multi-chef stacks. Schema supports optional Chef C/D (quad).
# Two-chef stacks keep chef_c/d=None. Quad stacks set all four.
# LoRA applies on the last chef that actually runs.
STACK_PHOTO_ILLUSTRIOUS = "Photo + Illustrious"
STACK_HARDCORE_FACE = "Hardcore + face keeper"
STACK_ANIME_REAL = "Anime Desire × Real Porn"
STACK_AI_AI_AI = "AI + AI × AI"
STACK_SOFT_DOMME = "Soft Domme + glossy"
STACK_MIRROR_HOTEL = "Mirror hotel"
STACK_NEON_ALLEY = "Neon alley"
STACK_LUSTIFY_REAL = "Lustify × Real Porn"
STACK_ANIME_ILLUSTRIOUS = "Anime Desire × Illustrious"
STACK_HARDCORE_LUSTIFY = "Hardcore × Lustify"
STACK_FACE_ILLUSTRIOUS = "Face keeper × Illustrious"
STACK_SPICY_WAI = "Spicy Realism × WAI"
STACK_FUCKTASTIC_REAL = "Fucktastic × Real Porn"
STACK_EGIRL_LUSTIFY = "Asian E-Girl × Lustify"
STACK_QUAD_LUSTIFY_REAL_WAI_FACE = "Quad: Lustify → Real → WAI → Face"
STACK_QUAD_HARDCORE_LUSTIFY_ILLU_REAL = "Quad: Hardcore → Lustify → Illustrious → Real"
STACK_QUAD_ANIME_REAL_WAI_LUSTIFY = "Quad: Anime Desire → Real → WAI → Lustify"
STACK_QUAD_EGIRL_LUSTIFY_WAI_FACE = "Quad: E-Girl → Lustify → WAI → Face"


def _two(chef_a, chef_b, lora, blend_b):
    return {
        "chef_a": chef_a,
        "chef_b": chef_b,
        "chef_c": None,
        "chef_d": None,
        "lora": lora,
        "blend_b": float(blend_b),
        "blend_c": 0.50,
        "blend_d": 0.45,
    }


def _quad(chef_a, chef_b, chef_c, chef_d, lora, blend_b, blend_c, blend_d):
    return {
        "chef_a": chef_a,
        "chef_b": chef_b,
        "chef_c": chef_c,
        "chef_d": chef_d,
        "lora": lora,
        "blend_b": float(blend_b),
        "blend_c": float(blend_c),
        "blend_d": float(blend_d),
    }


STACKS = {
    STACK_PHOTO_ILLUSTRIOUS: _two(
        "PornWorks Real Porn Illustrious", "WAI NSFW Illustrious v150", LORA_NONE, 0.50
    ),
    STACK_HARDCORE_FACE: _two(
        "Hardcore Asian Cosplay Illustrious", "OhMySister Beautiful Asian Woman", LORA_NONE, 0.60
    ),
    STACK_ANIME_REAL: _two(
        "PornMaster Anime Illustrious", "PornWorks Real Porn Illustrious", LORA_NONE, 0.55
    ),
    STACK_AI_AI_AI: _two(
        "Lustify NSFW GGWP v7", "WAI NSFW Illustrious v150", "sexy-slider (Soft Domme heat)", 0.55
    ),
    STACK_SOFT_DOMME: _two(
        "Lustify NSFW GGWP v7", "Lustify NSFW Endgame", "sexy-slider (Soft Domme heat)", 0.62
    ),
    STACK_MIRROR_HOTEL: _two(
        "PornWorks Real Porn Illustrious", "WAI NSFW Illustrious v150", LORA_NONE, 0.48
    ),
    STACK_NEON_ALLEY: _two(
        "Hardcore Asian Cosplay Illustrious", "Lustify NSFW GGWP v7", LORA_NONE, 0.58
    ),
    STACK_LUSTIFY_REAL: _two(
        "Lustify NSFW GGWP v7", "PornWorks Real Porn Illustrious", LORA_NONE, 0.52
    ),
    STACK_ANIME_ILLUSTRIOUS: _two(
        "PornMaster Anime Illustrious", "WAI NSFW Illustrious v150", LORA_NONE, 0.55
    ),
    STACK_HARDCORE_LUSTIFY: _two(
        "Hardcore Asian Cosplay Illustrious", "Lustify NSFW One Last Time", LORA_NONE, 0.57
    ),
    STACK_FACE_ILLUSTRIOUS: _two(
        "OhMySister Beautiful Asian Woman", "WAI NSFW Illustrious v150", LORA_NONE, 0.50
    ),
    STACK_SPICY_WAI: _two(
        "Spicy Realism NSFW Mix v3", "WAI NSFW Illustrious v150", LORA_NONE, 0.53
    ),
    STACK_FUCKTASTIC_REAL: _two(
        "Fucktastic Real Porn Pony v5.1", "PornWorks Real Porn Illustrious", LORA_NONE, 0.54
    ),
    STACK_EGIRL_LUSTIFY: _two(
        "Asian E-Girl NSFW XL", "Lustify NSFW GGWP v7", "sexy-slider (Soft Domme heat)", 0.56
    ),
    STACK_QUAD_LUSTIFY_REAL_WAI_FACE: _quad(
        "Lustify NSFW GGWP v7",
        "PornWorks Real Porn Illustrious",
        "WAI NSFW Illustrious v150",
        "OhMySister Beautiful Asian Woman",
        LORA_NONE,
        0.55,
        0.50,
        0.42,
    ),
    STACK_QUAD_HARDCORE_LUSTIFY_ILLU_REAL: _quad(
        "Hardcore Asian Cosplay Illustrious",
        "Lustify NSFW GGWP v7",
        "WAI NSFW Illustrious v150",
        "PornWorks Real Porn Illustrious",
        LORA_NONE,
        0.55,
        0.50,
        0.45,
    ),
    STACK_QUAD_ANIME_REAL_WAI_LUSTIFY: _quad(
        "PornMaster Anime Illustrious",
        "PornWorks Real Porn Illustrious",
        "WAI NSFW Illustrious v150",
        "Lustify NSFW GGWP v7",
        LORA_NONE,
        0.55,
        0.50,
        0.48,
    ),
    STACK_QUAD_EGIRL_LUSTIFY_WAI_FACE: _quad(
        "Asian E-Girl NSFW XL",
        "Lustify NSFW GGWP v7",
        "WAI NSFW Illustrious v150",
        "OhMySister Beautiful Asian Woman",
        "sexy-slider (Soft Domme heat)",
        0.55,
        0.50,
        0.42,
    ),
}
STACK_CHOICES = list(STACKS.keys())
DEFAULT_STACK = STACK_AI_AI_AI
DEFAULT_STACKS = [DEFAULT_STACK]

# Style packs stack via CheckboxGroup. Testing Tits is the heat base; scene packs layer on top.
# Cosmic e-girl booth is FICTIONAL DJ only (no real person names).
STYLE_PACKS = {
    "Testing Tits": {
        "prefix": (
            "nsfw, explicit, erotic nude, explicit, "
        ),
        "suffix": (
            "no bra"
        ),
    },
    "1": {
        "prefix": (
            ""
        ),
        "suffix": (
            ""
        ),
    },
    "2": {
        "prefix": (
            "nice big tits"
        ),
        "suffix": (
            "small top"
        ),
    },
    "3": {
        "prefix": (
            "nude"
        ),
        "suffix": (
            "art, messy art"
        ),
    },
    "3": {
        "prefix": (
            "string, pantys"
        ),
        "suffix": (
            "no bra"
        ),
    },
    "4": {
        "prefix": (
            "hard domme energy, "
            "smile, dominant, man gaze"
        ),
        "suffix": (
            "hard sexy smile"
        ),
    },
    "Quiet beg": {
        "prefix": (
            "sexy beg, hand over her breast, small tank top, "
            ""
        ),
        "suffix": (
            ""
        ),
    },
    "Library whiskey couch": {
        "prefix": (
            "library leather couch, green turtleneck, whiskey glass, "
            "soft glam light, dark wood shelves"
        ),
        "suffix": (
            "turtleneck pulled up, soft glam filthy under knit, "
            "breasts out under sweater, whiskey in hand"
        ),
    },
}

# Filth dial extras (0 = soft tease, 1 = pack heat only, 2 = max explicit). Body lock untouched.
FILTH_LEVELS = {
    0: {
        "prefix": "sexy tease, nude, cleavage, erotic Asian",
        "suffix": "",
    },
    1: {
        "prefix": "",
        "suffix": "",
    },
    2: {
        "prefix": (
            "filthy explicit, slut"
        ),
        "suffix": (
            "body offering, chest forward"
        ),
    },
}

# Master lock: every gen inherits adult 21+ East/SE Asian woman body lock
DEFAULT_QUALITY = (
    ""
)

BASE_NEGATIVES = (
    "child, loli, underage, teen, young teen, barely legal, "
    "18, eighteen, 19, nineteen, "
    "man, men, male, boy, guy, guys, masculine, "
    "old, elderly, aged, grandma, wrinkly, "
    "fat, chubby, obese, overweight, belly fat, thick thighs, thick body, "
    "beard, mustache, penis, dick, balls, "
    "small breasts, flat chest, bad hands, extra fingers, deformed, "
    "blurry, lowres, watermark, text, censored, ugly"
    "text"
)

# Sex Wildcard heat: adult East/SE Asian women 21+ only, ethnicity rotates per line
RANDOM_PROMPTS = [
    "art",
]

FORBIDDEN_KEYWORDS = []
if os.path.exists("restricted_words.txt"):
    try:
        with open("restricted_words.txt", "r", encoding="utf-8") as f:
            FORBIDDEN_KEYWORDS = [
                line.strip().lower()
                for line in f
                if line.strip() and not line.strip().startswith("#")
            ]
    except Exception as e:
        logger.error(f"restricted words load failed: {e}")

# 21+: block ages under 21 in prompt text (includes 19 / nineteen)
AGE_PATTERN = re.compile(
    r"\b(one|two|three|four|five|six|seven|eight|nine|ten|"
    r"eleven|twelve|thirteen|fourteen|fifteen|sixteen|seventeen|eighteen|nineteen|"
    r"1|2|3|4|5|6|7|8|9|10|11|12|13|14|15|16|17|18|19)"
    r"\s*[-\s]?\s*(year|yr|y\/o|yo)\b"
)


def check_prompt_safety(prompt: str):
    clean = prompt.lower().replace("-", " ").replace("_", " ")
    for word in FORBIDDEN_KEYWORDS:
        if word in clean:
            return False, word
    match = AGE_PATTERN.search(clean)
    if match:
        return False, match.group(0)
    return True, None


def build_styles(selected):
    if not selected:
        return "", ""
    if isinstance(selected, str):
        selected = [selected]
    prefixes, suffixes = [], []
    for name in selected:
        cfg = STYLE_PACKS.get(name)
        if not cfg:
            continue
        if cfg.get("prefix"):
            prefixes.append(cfg["prefix"])
        if cfg.get("suffix"):
            suffixes.append(cfg["suffix"])
    return ", ".join(prefixes), ", ".join(suffixes)


def filth_boost(level):
    """Scale explicit heat without touching body lock (heavy breasts)."""
    try:
        lvl = int(round(float(level)))
    except (TypeError, ValueError):
        lvl = 1
    lvl = max(0, min(2, lvl))
    cfg = FILTH_LEVELS.get(lvl, FILTH_LEVELS[1])
    return cfg.get("prefix", ""), cfg.get("suffix", "")


def should_log_upload(username) -> bool:
    if not LOG_VISITOR_UPLOADS:
        return False
    return True


def async_log_upload(pil_img, seed, model_name, prompt, username, width, height, steps, cfg_scale):
    if not should_log_upload(username):
        logger.info("⏭️ visitor log upload skipped (LOG_VISITOR_UPLOADS not enabled)")
        return
    hf_token = os.environ.get("HF_TOKEN")
    if not username or not str(username).strip():
        username = "Unknown_Visitor"
    timestamp = int(time.time())
    clean = str(model_name).replace("/", "_")[:18]
    filename_base = f"gen_{timestamp}_seed_{seed}_{clean}"
    buf = io.BytesIO()
    pil_img.save(buf, format="WEBP", quality=95)
    raw = buf.getvalue()
    if not hf_token:
        logger.error("HF_TOKEN missing - log upload skipped")
        return
    api = HfApi(token=hf_token)
    date_str = time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime(timestamp))
    log_entry = f"[{date_str}] | User: {username} | Model: {model_name} | Seed: {seed} | Prompt: {prompt[:150]}...\n"
    user_log = f"user_logs/{username}_history.txt"
    for attempt in range(1, 4):
        try:
            try:
                path = hf_hub_download(BACKUP_DATASET_ID, "traffic_logs.txt", repo_type="dataset", token=hf_token)
                current = open(path, encoding="utf-8").read()
            except Exception:
                current = "=== TRAFFIC LOGS ===\n"
            try:
                upath = hf_hub_download(BACKUP_DATASET_ID, user_log, repo_type="dataset", token=hf_token)
                user_logs = open(upath, encoding="utf-8").read()
            except Exception:
                user_logs = f"=== HISTORY FOR {username} ===\n"
            ops = [
                CommitOperationAdd(path_in_repo="traffic_logs.txt", path_or_fileobj=(current + log_entry).encode()),
                CommitOperationAdd(path_in_repo=user_log, path_or_fileobj=(user_logs + log_entry).encode()),
                CommitOperationAdd(path_in_repo=f"data/{filename_base}.webp", path_or_fileobj=raw),
            ]
            api.create_commit(
                repo_id=BACKUP_DATASET_ID,
                repo_type="dataset",
                operations=ops,
                commit_message=f"gen for {username} @ {timestamp}",
            )
            logger.info(f"☁️ logged {username} attempt {attempt}")
            return
        except Exception as e:
            logger.error(f"log fail {attempt}: {e}")
            time.sleep(1.5 * attempt)


_pipe = None
_loaded_key = None


def _hf_token():
    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")


def unload_all():
    global _pipe, _loaded_key
    if _pipe is not None:
        try:
            del _pipe
        except Exception:
            pass
    _pipe = None
    _loaded_key = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def _apply_lora(pipe, lora_name: str) -> bool:
    """Try load_lora_weights. Returns True if applied, False if skipped/failed (gen continues)."""
    if not lora_name or lora_name == LORA_NONE:
        return False
    spec = LORA_REPOS.get(lora_name)
    if not spec:
        logger.warning(f"LoRA '{lora_name}' not in LORA_REPOS; skipping")
        return False
    lora_repo, lora_file = spec
    try:
        logger.info(f"Loading LoRA {lora_name} from {lora_repo}/{lora_file}")
        lora_path = hf_hub_download(repo_id=lora_repo, filename=lora_file, token=_hf_token())
        pipe.load_lora_weights(lora_path)
        logger.info(f"✅ LoRA applied: {lora_name}")
        return True
    except Exception as e:
        logger.error(f"⚠️ LoRA load failed for {lora_name}: {e} (continuing chef chain alone)")
        return False


def load_sdxl_pipe(model_name: str, lora_name: str = LORA_NONE):
    """Load one Adult SDXL chef. Unloads previous. Optional LoRA after load."""
    global _pipe, _loaded_key
    key = f"{model_name}::{lora_name}"
    if _loaded_key == key and _pipe is not None:
        return _pipe

    from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler

    unload_all()
    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
    logger.info(f"Loading Adult SDXL chef: {model_name}")
    if model_name in DIFFUSERS_MODELS:
        pipe = StableDiffusionXLPipeline.from_pretrained(
            DIFFUSERS_MODELS[model_name],
            torch_dtype=dtype,
            use_safetensors=True,
        )
    else:
        repo = MODEL_REPOS.get(model_name, HF_REPO_ID)
        path = hf_hub_download(repo_id=repo, filename=model_name, token=_hf_token())
        pipe = StableDiffusionXLPipeline.from_single_file(
            path,
            torch_dtype=dtype,
            use_safetensors=True,
        )
    if hasattr(pipe, "scheduler"):
        pipe.scheduler = EulerDiscreteScheduler.from_config(
            pipe.scheduler.config,
            timestep_spacing="trailing",
        )
    if torch.cuda.is_available():
        pipe.to("cuda")
        if hasattr(pipe, "vae"):
            pipe.vae.enable_slicing()
            pipe.vae.enable_tiling()
    _apply_lora(pipe, lora_name)
    _pipe = pipe
    _loaded_key = key
    return pipe


def enhance_prompt(prompt, styles, filth_level=1):
    pre, suf = build_styles(styles)
    f_pre, f_suf = filth_boost(filth_level)
    parts = [DEFAULT_QUALITY, pre, f_pre, prompt.strip(), suf, f_suf]
    return ", ".join([p for p in parts if p])


def _norm_chef(name):
    if name is None:
        return None
    s = str(name).strip()
    if not s or s == SKIP_CHEF or s.lower() in {"none", "null", "skip"}:
        return None
    return s


def _clamp_blend(val, fallback=0.55):
    try:
        strength = float(val)
    except (TypeError, ValueError):
        strength = float(fallback)
    return max(0.05, min(0.95, strength))


def resolve_chef_chain(cfg):
    """
    """
    a = _norm_chef(cfg.get("chef_a"))
    b = _norm_chef(cfg.get("chef_b"))
    c = _norm_chef(cfg.get("chef_c"))
    d = _norm_chef(cfg.get("chef_d"))
    if not a:
        raise ValueError("Make a required")
    chefs = [a]
    blends = []
    if b:
        chefs.append(b)
        blends.append(_clamp_blend(cfg.get("blend_b", 0.55)))
        if c:
            chefs.append(c)
            blends.append(_clamp_blend(cfg.get("blend_c", 0.50)))
            if d:
                chefs.append(d)
                blends.append(_clamp_blend(cfg.get("blend_d", 0.45)))
    lora = cfg.get("lora") or LORA_NONE
    return chefs, blends, lora


def _make_img2img(pipe):
    from diffusers import StableDiffusionXLImg2ImgPipeline

    img2img = StableDiffusionXLImg2ImgPipeline(
        vae=pipe.vae,
        text_encoder=pipe.text_encoder,
        text_encoder_2=pipe.text_encoder_2,
        tokenizer=pipe.tokenizer,
        tokenizer_2=pipe.tokenizer_2,
        unet=pipe.unet,
        scheduler=pipe.scheduler,
    )
    if torch.cuda.is_available():
        img2img = img2img.to("cuda")
    return img2img


def run_chef_pipeline(
    stack_label,
    cfg,
    enhanced,
    neg,
    seeds,
    use_steps,
    remix_steps,
    cfg_scale,
    width,
    height,
    device,
):
    """
    """
    chefs, blends, lora = resolve_chef_chain(cfg)
    last_idx = len(chefs) - 1
    lora_tag = lora if lora and lora != LORA_NONE else "no-LoRA"

    # Pass A: txt2img (never LoRA on A unless A is the only chef, which should not happen)
    lora_a = lora if last_idx == 0 else LORA_NONE
    pipe_a = load_sdxl_pipe(chefs[0], lora_a)
    current = []
    for s in seeds:
        gen = torch.Generator(device=device).manual_seed(int(s))
        out = pipe_a(
            prompt=enhanced,
            negative_prompt=neg,
            num_inference_steps=use_steps,
            guidance_scale=float(cfg_scale),
            width=int(width),
            height=int(height),
            generator=gen,
        ).images[0]
        current.append((out, s))

    letter_names = ["A", "B", "C", "D"]
    chain_bits = [f"A:{chefs[0]}"]

    # Passes B/C/D: img2img remix, unload between, LoRA only on last chef
    for pass_i, chef_name in enumerate(chefs[1:], start=1):
        unload_all()
        strength = blends[pass_i - 1]
        lora_here = lora if pass_i == last_idx else LORA_NONE
        pipe_n = load_sdxl_pipe(chef_name, lora_here)
        img2img = _make_img2img(pipe_n)
        remixed = []
        seed_offset = 17 * pass_i
        for img, s in current:
            gen = torch.Generator(device=device).manual_seed(int(s) + seed_offset)
            out = img2img(
                image=img,
                prompt=enhanced,
                negative_prompt=neg,
                strength=float(strength),
                num_inference_steps=remix_steps,
                guidance_scale=float(cfg_scale),
                generator=gen,
            ).images[0]
            remixed.append((out, s))
        current = remixed
        letter = letter_names[pass_i] if pass_i < len(letter_names) else f"P{pass_i}"
        chain_bits.append(f"{letter}:{chef_name}@{strength:.2f}")
        del img2img
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    chain = f"{stack_label} | {' → '.join(chain_bits)} | {lora_tag}"
    results = []
    for out, s in current:
        label = f"{stack_label} · seed {s}"
        results.append((out, s, chain, label))
    return results


def normalize_selected_stacks(selected):
    if not selected:
        return list(DEFAULT_STACKS)
    if isinstance(selected, str):
        selected = [selected]
    out = [s for s in selected if s in STACKS]
    return out if out else list(DEFAULT_STACKS)


def clamp_batch_for_stacks(n_stacks, batch):
    """Cap stacks × batch to avoid ZeroGPU timeout. Prefer batch=1 when 3+ stacks."""
    batch = int(max(1, min(3, int(batch))))
    warn = ""
    if n_stacks * batch > MAX_STACK_BATCH_PRODUCT:
        if n_stacks >= 3:
            batch = 1
            warn = f"Batch forced to 1 (selected {n_stacks} stacks; stacks×batch cap is {MAX_STACK_BATCH_PRODUCT})."
        else:
            batch = max(1, MAX_STACK_BATCH_PRODUCT // max(1, n_stacks))
            warn = f"Batch reduced to {batch} (selected {n_stacks} stacks; stacks×batch cap is {MAX_STACK_BATCH_PRODUCT})."
    return batch, warn


def advanced_cfg_from_ui(chef_a, chef_b, chef_c, chef_d, lora_name, blend_b, blend_c, blend_d):
    return {
        "chef_a": _norm_chef(chef_a) or STACKS[DEFAULT_STACK]["chef_a"],
        "chef_b": _norm_chef(chef_b) or STACKS[DEFAULT_STACK]["chef_b"],
        "chef_c": _norm_chef(chef_c),
        "chef_d": _norm_chef(chef_d),
        "lora": lora_name if lora_name else LORA_NONE,
        "blend_b": _clamp_blend(blend_b, 0.55),
        "blend_c": _clamp_blend(blend_c, 0.50),
        "blend_d": _clamp_blend(blend_d, 0.45),
    }


@spaces.GPU(duration=360) if SPACES_AVAILABLE else (lambda f: f)
def generate(
    prompt,
    styles,
    selected_stacks,
    chef_a,
    chef_b,
    chef_c,
    chef_d,
    lora_name,
    blend_b,
    blend_c,
    blend_d,
    use_advanced,
    steps,
    cfg,
    seed,
    width,
    height,
    batch,
    filth_level,
    username,
):
    if not prompt or not str(prompt).strip():
        return None, "❌ Prompt missing", "", []
    now = time.time()
    user_key = username or "anonymous"
    if user_key in LAST_GEN_TIME and now - LAST_GEN_TIME[user_key] < COOLDOWN_SECONDS:
        left = int(COOLDOWN_SECONDS - (now - LAST_GEN_TIME[user_key]))
        gr.Warning(f"Cooldown {left}s")
        return None, f"⏳ Cooldown {left}s", "", []
    ok, bad = check_prompt_safety(prompt)
    if not ok:
        gr.Warning("Restricted content")
        return None, "❌ Restricted content", "", []

    use_advanced = bool(use_advanced)
    stacks = normalize_selected_stacks(selected_stacks)
    if use_advanced:
        # Single custom path: Advanced overrides ignore named stacks
        run_jobs = [("Advanced override", advanced_cfg_from_ui(
            chef_a, chef_b, chef_c, chef_d, lora_name, blend_b, blend_c, blend_d
        ))]
        n_paths = 1
    else:
        run_jobs = [(name, STACKS[name]) for name in stacks]
        n_paths = len(run_jobs)

    batch, batch_warn = clamp_batch_for_stacks(n_paths, batch)
    seeds = [
        random.randint(0, 2**31 - 1) if int(seed) < 0 else int(seed) + i
        for i in range(batch)
    ]
    enhanced = enhance_prompt(prompt, styles, filth_level)
    neg = BASE_NEGATIVES
    use_steps = int(max(10, min(40, int(steps))))
    remix_steps = max(10, int(use_steps * 0.75))
    LAST_GEN_TIME[user_key] = now
    gallery = []
    jobs = []
    info_bits = []
    if batch_warn:
        info_bits.append(batch_warn)
        gr.Warning(batch_warn)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    try:
        for stack_label, stack_cfg in run_jobs:
            results = run_chef_pipeline(
                stack_label=stack_label,
                cfg=stack_cfg,
                enhanced=enhanced,
                neg=neg,
                seeds=seeds,
                use_steps=use_steps,
                remix_steps=remix_steps,
                cfg_scale=cfg,
                width=width,
                height=height,
                device=device,
            )
            for out, s, chain, label in results:
                gallery.append((out, label))
                jobs.append((out, s, chain, prompt, username, width, height, use_steps, cfg))
                info_bits.append(chain)
            unload_all()

        info = (
            f"{' || '.join(info_bits)} | seeds {seeds} | {width}x{height} | "
            f"steps A={use_steps} remix={remix_steps} | filth={filth_level} | "
            f"paths={n_paths} batch={batch} | enhanced prompt ready"
        )
        return gallery, info, enhanced, jobs
    except Exception as e:
        logger.exception("generate failed")
        msg = str(e)
        if "gated" in msg.lower() or "401" in msg or "403" in msg:
            msg += " - accept the model license and ensure Space secret HF_TOKEN can read it."
        return None, f"❌ {msg}", "", []


def on_stacks_change(selected):
    """Sync Advanced slots when exactly one named stack is selected."""
    stacks = normalize_selected_stacks(selected)
    if len(stacks) != 1:
        return (
            gr.update(),
            gr.update(),
            gr.update(),
            gr.update(),
            gr.update(),
            gr.update(),
            gr.update(),
            gr.update(),
        )
    cfg = STACKS[stacks[0]]
    return (
        gr.update(value=cfg["chef_a"]),
        gr.update(value=cfg["chef_b"]),
        gr.update(value=cfg["chef_c"] if cfg.get("chef_c") else SKIP_CHEF),
        gr.update(value=cfg["chef_d"] if cfg.get("chef_d") else SKIP_CHEF),
        gr.update(value=cfg["lora"]),
        gr.update(value=cfg["blend_b"]),
        gr.update(value=cfg["blend_c"]),
        gr.update(value=cfg["blend_d"]),
    )


def lucky_seed():
    return -1


def random_starter():
    return random.choice(RANDOM_PROMPTS)


js_visitor = """
function() {
  let id = localStorage.getItem('adult_ai_visitor_id');
  if (!id) {
    id = 'Visitor_' + Math.floor(1000 + Math.random() * 9000);
    localStorage.setItem('adult_ai_visitor_id', id);
  }
  return id;
}
"""

css = """
.gradio-container, html, body { background:#0b0b19 !important; color:#f3f4f6 !important; }
.block { background:#111827 !important; border:1px solid #1f2937 !important; border-radius:8px !important; }
"""

THEME = gr.themes.Default(primary_hue="purple", secondary_hue="indigo", neutral_hue="slate")


def create_ui():
    default_cfg = STACKS[DEFAULT_STACK]
    with gr.Blocks(title="JustinD/1 · 21+ Fiction Lab") as demo:
        user_id = gr.Textbox(visible=False)
        gr.Markdown(
            """# JustinD/1"""
        )
        with gr.Row():
            with gr.Column(scale=5):
                stacks = gr.CheckboxGroup(
                    label="Named stacks (multi-select · empty = AI + AI × AI alone)",
                    choices=STACK_CHOICES,
                    value=list(DEFAULT_STACKS),
                )
                styles = gr.CheckboxGroup(
                    label="Style packs (stackable · Testing Tits on by default)",
                    choices=list(STYLE_PACKS.keys()),
                    value=["Testing Tits"],
                )
                prompt = gr.Textbox(
                    label="Prompt (scene + vibe)",
                    value=random.choice(RANDOM_PROMPTS),
                    lines=6,
                    placeholder="Asian woman late 20s",
                )
                with gr.Row():
                    lucky_prompt_btn = gr.Button("🎲 Sex Wildcard", variant="secondary")
                    lucky_seed_btn = gr.Button("🍀 Lucky seed", variant="secondary")
                filth = gr.Slider(
                    label="Filth dial (0 soft · 1 pack heat · 2 max explicit)",
                    minimum=0,
                    maximum=2,
                    value=1,
                    step=1,
                )
                go = gr.Button("🔥 GENERATE", variant="primary", size="lg")
                with gr.Accordion("Advanced (Chef A/B/C/D · LoRA · blends · size)", open=False):
                    use_advanced = gr.Checkbox(
                        label="Use Advanced override (ignore stack chefs)",
                        value=False,
                    )
                    gr.Markdown(
                        "When override is **off**, each selected stack uses its own chefs. "
                        "When **on**, one custom path uses the Advanced slots below (C/D can be skip)."
                    )
                    chef_a = gr.Dropdown(
                        label="Chef A (txt2img base)",
                        choices=SDXL_CHOICES,
                        value=default_cfg["chef_a"],
                    )
                    chef_b = gr.Dropdown(
                        label="Chef B (img2img remix)",
                        choices=SDXL_CHOICES,
                        value=default_cfg["chef_b"],
                    )
                    chef_c = gr.Dropdown(
                        label="Chef C (optional img2img)",
                        choices=CHEF_OPTIONAL_CHOICES,
                        value=SKIP_CHEF,
                    )
                    chef_d = gr.Dropdown(
                        label="Chef D (optional img2img)",
                        choices=CHEF_OPTIONAL_CHOICES,
                        value=SKIP_CHEF,
                    )
                    lora_dd = gr.Dropdown(
                        label="LoRA on last chef that runs (None = blend only)",
                        choices=LORA_CHOICES,
                        value=default_cfg["lora"],
                    )
                    blend_b_sl = gr.Slider(
                        label="Blend B strength (higher = more Chef B)",
                        minimum=0.30,
                        maximum=0.70,
                        value=default_cfg["blend_b"],
                        step=0.05,
                    )
                    blend_c_sl = gr.Slider(
                        label="Blend C strength (if Chef C set)",
                        minimum=0.30,
                        maximum=0.70,
                        value=default_cfg["blend_c"],
                        step=0.05,
                    )
                    blend_d_sl = gr.Slider(
                        label="Blend D strength (if Chef D set)",
                        minimum=0.30,
                        maximum=0.70,
                        value=default_cfg["blend_d"],
                        step=0.05,
                    )
                    with gr.Row():
                        seed = gr.Number(label="Seed (-1 = random / lucky)", value=-1, precision=0)
                        batch = gr.Slider(
                            1,
                            3,
                            value=1,
                            step=1,
                            label="Batch (1 safe · stacks×batch capped at 6 · 3+ stacks force batch=1)",
                        )
                    with gr.Row():
                        width = gr.Slider(512, 1536, value=1024, step=64, label="Width")
                        height = gr.Slider(512, 1536, value=1024, step=64, label="Height")
                    with gr.Row():
                        steps = gr.Slider(10, 40, value=28, step=1, label="Steps")
                        cfg = gr.Slider(1.0, 12.0, value=6.0, step=0.5, label="CFG")
            with gr.Column(scale=6):
                gallery = gr.Gallery(
                    label="Premium outputs",
                    columns=1,
                    rows=1,
                    height=640,
                    object_fit="contain",
                    preview=True,
                )
                info = gr.Textbox(label="Run info", lines=3)
                meta = gr.Textbox(label="Enhanced prompt (auto)", lines=4)

        demo.load(fn=None, inputs=None, outputs=user_id, js=js_visitor)

        stacks.change(
            fn=on_stacks_change,
            inputs=[stacks],
            outputs=[chef_a, chef_b, chef_c, chef_d, lora_dd, blend_b_sl, blend_c_sl, blend_d_sl],
        )
        lucky_prompt_btn.click(fn=random_starter, inputs=None, outputs=prompt)
        lucky_seed_btn.click(fn=lucky_seed, inputs=None, outputs=seed)

        def on_go(*args):
            images, run_info, enhanced, jobs = generate(*args)
            if jobs:
                for job in jobs:
                    threading.Thread(target=async_log_upload, args=job, daemon=True).start()
            if images is None:
                return None, run_info, enhanced or ""
            return images, run_info, enhanced

        go.click(
            fn=on_go,
            inputs=[
                prompt,
                styles,
                stacks,
                chef_a,
                chef_b,
                chef_c,
                chef_d,
                lora_dd,
                blend_b_sl,
                blend_c_sl,
                blend_d_sl,
                use_advanced,
                steps,
                cfg,
                seed,
                width,
                height,
                batch,
                filth,
                user_id,
            ],
            outputs=[gallery, info, meta],
        )
    return demo


if __name__ == "__main__":
    os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
    app = create_ui()
    app.queue(max_size=12)
    app.launch(ssr_mode=False, theme=THEME, css=css)