Update app.py

app.py

@@ -30,7 +30,7 @@ logging.basicConfig(
     handlers=[logging.StreamHandler(sys.stdout)],
 )
 logger = logging.getLogger("NextGenEngine")
-logger.info("🚀 NSFW East/SE Asian Lab boot (21+ · Adult SDXL / FLUX / SD3.5)")
+logger.info("🚀 NSFW East/SE Asian Lab boot (21+ · Multi-chef stacks)")
 
 if hasattr(sys.stdout, "reconfigure"):
     sys.stdout.reconfigure(encoding="utf-8")
@@ -42,11 +42,6 @@ LAST_GEN_TIME = {}
 COOLDOWN_SECONDS = 12
 LOG_VISITOR_UPLOADS = os.environ.get("LOG_VISITOR_UPLOADS", "0").strip().lower() in {"1", "true", "yes"}
 
-ENGINE_FLUX = "FLUX.1 Schnell (latest fast)"
-ENGINE_SD35 = "SD 3.5 Medium (gated / optional)"
-ENGINE_SDXL = "Adult SDXL (NSFW chefs)"
-ENGINE_CHOICES = [ENGINE_SDXL, ENGINE_FLUX, ENGINE_SD35]
-
 HF_REPO_ID = "AriaVale2/my-safetensors-backup"
 
 # Single-file chefs (Asian realism + proven NSFW pantry)
@@ -111,10 +106,59 @@ DIFFUSERS_MODELS = {
 }
 
 SDXL_CHOICES = list(DIFFUSERS_MODELS.keys()) + list(MODEL_REPOS.keys())
-DEFAULT_SDXL = "WAI NSFW Illustrious v150"
+
+# LoRA pantry from known SDXL repos (patch_app.py pattern). 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. Every gen = Chef A txt2img then Chef B img2img remix.
+# Default stack = filthiest (AI + AI x AI) with sexy-slider LoRA.
+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"
+
+STACKS = {
+    STACK_PHOTO_ILLUSTRIOUS: {
+        "chef_a": "PornWorks Real Porn v04",
+        "chef_b": "WAI NSFW Illustrious v150",
+        "lora": LORA_NONE,
+        "blend": 0.50,
+    },
+    STACK_HARDCORE_FACE: {
+        "chef_a": "Hardcore Asian Cosplay Illustrious",
+        "chef_b": "OhMySister Beautiful Asian Woman",
+        "lora": LORA_NONE,
+        "blend": 0.60,
+    },
+    STACK_ANIME_REAL: {
+        "chef_a": "PornWorks Anime Desire NSFW",
+        "chef_b": "PornWorks Real Porn v04",
+        "lora": LORA_NONE,
+        "blend": 0.55,
+    },
+    STACK_AI_AI_AI: {
+        "chef_a": "Lustify NSFW GGWP v7",
+        "chef_b": "WAI NSFW Illustrious v150",
+        "lora": "sexy-slider (Soft Domme heat)",
+        "blend": 0.55,
+    },
+}
+STACK_CHOICES = list(STACKS.keys())
+DEFAULT_STACK = STACK_AI_AI_AI
 
 # Style packs stack via CheckboxGroup. Testing Tits is the heat base; scene packs layer on top.
-# No packs named "1" or "2". Cosmic e-girl booth is FICTIONAL DJ only (no real person names).
+# Cosmic e-girl booth is FICTIONAL DJ only (no real person names).
 STYLE_PACKS = {
     "Testing Tits": {
         "prefix": (
@@ -346,80 +390,57 @@ def async_log_upload(pil_img, seed, model_name, prompt, username, width, height,
             time.sleep(1.5 * attempt)
 
 
-_pipes = {}
+_pipe = None
+_loaded_key = None
 
 
-def _dtype():
-    if torch.cuda.is_available():
-        return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
-    return torch.float32
+def _hf_token():
+    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
 
 
 def unload_all():
-    global _pipes
-    _pipes.clear()
+    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 _hf_token():
-    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
-
-
-def get_flux_pipe():
-    key = "flux-schnell"
-    if key in _pipes:
-        return _pipes[key]
-    from diffusers import FluxPipeline
-    unload_all()
-    logger.info("Loading FLUX.1 Schnell…")
-    token = _hf_token()
-    pipe = FluxPipeline.from_pretrained(
-        "black-forest-labs/FLUX.1-schnell",
-        torch_dtype=_dtype(),
-        token=token,
-    )
-    if torch.cuda.is_available():
-        pipe.to("cuda")
-    _pipes[key] = pipe
-    return pipe
-
-
-def get_sd35_pipe():
-    key = "sd35-medium"
-    if key in _pipes:
-        return _pipes[key]
-    from diffusers import StableDiffusion3Pipeline
-    from huggingface_hub import login as hf_login
-    token = _hf_token()
-    if not token:
-        raise RuntimeError(
-            "SD 3.5 is gated. Add HF_TOKEN secret and accept "
-            "https://huggingface.co/stabilityai/stable-diffusion-3.5-medium"
-        )
+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:
-        hf_login(token=token, add_to_git_credential=False)
+        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.warning(f"hf_login warning: {e}")
-    unload_all()
-    logger.info("Loading SD 3.5 Medium (authenticated)…")
-    pipe = StableDiffusion3Pipeline.from_pretrained(
-        "stabilityai/stable-diffusion-3.5-medium",
-        torch_dtype=_dtype(),
-        token=token,
-    )
-    if torch.cuda.is_available():
-        pipe.to("cuda")
-    _pipes[key] = pipe
-    return pipe
+        logger.error(f"⚠️ LoRA load failed for {lora_name}: {e} (continuing A+B blend 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
 
-def get_sdxl_pipe(model_name: str):
-    key = f"sdxl:{model_name}"
-    if key in _pipes:
-        return _pipes[key]
     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}")
@@ -431,7 +452,7 @@ def get_sdxl_pipe(model_name: str):
         )
     else:
         repo = MODEL_REPOS.get(model_name, HF_REPO_ID)
-        path = hf_hub_download(repo_id=repo, filename=model_name)
+        path = hf_hub_download(repo_id=repo, filename=model_name, token=_hf_token())
         pipe = StableDiffusionXLPipeline.from_single_file(
             path,
             torch_dtype=dtype,
@@ -447,7 +468,9 @@ def get_sdxl_pipe(model_name: str):
         if hasattr(pipe, "vae"):
             pipe.vae.enable_slicing()
             pipe.vae.enable_tiling()
-    _pipes[key] = pipe
+    _apply_lora(pipe, lora_name)
+    _pipe = pipe
+    _loaded_key = key
     return pipe
 
 
@@ -457,13 +480,30 @@ def enhance_prompt(prompt, styles):
     return ", ".join([p for p in parts if p])
 
 
-@spaces.GPU(duration=180) if SPACES_AVAILABLE else (lambda f: f)
+def resolve_stack_slots(stack_name, chef_a, chef_b, lora_name, blend):
+    """Use Advanced overrides when set; fill gaps from named stack."""
+    cfg = STACKS.get(stack_name) or STACKS[DEFAULT_STACK]
+    a = chef_a if chef_a and str(chef_a).strip() else cfg["chef_a"]
+    b = chef_b if chef_b and str(chef_b).strip() else cfg["chef_b"]
+    lora = lora_name if lora_name else cfg["lora"]
+    try:
+        strength = float(blend)
+    except (TypeError, ValueError):
+        strength = float(cfg["blend"])
+    strength = max(0.05, min(0.95, strength))
+    return a, b, lora, strength
+
+
+@spaces.GPU(duration=240) if SPACES_AVAILABLE else (lambda f: f)
 def generate(
-    engine,
     prompt,
     styles,
     negative,
-    sdxl_model,
+    stack_name,
+    chef_a,
+    chef_b,
+    lora_name,
+    blend,
     steps,
     cfg,
     seed,
@@ -484,6 +524,7 @@ def generate(
     if not ok:
         gr.Warning("Restricted content")
         return None, "❌ Restricted content", "", []
+
     batch = int(max(1, min(3, int(batch))))
     seeds = [
         random.randint(0, 2**31 - 1) if int(seed) < 0 else int(seed) + i
@@ -493,69 +534,80 @@ def generate(
     neg = BASE_NEGATIVES
     if negative and str(negative).strip():
         neg = f"{BASE_NEGATIVES}, {negative.strip()}"
+
+    model_a, model_b, lora, strength = resolve_stack_slots(
+        stack_name, chef_a, chef_b, lora_name, blend
+    )
+    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 = []
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
     try:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        if engine == ENGINE_FLUX:
-            pipe = get_flux_pipe()
-            use_steps = int(max(1, min(8, int(steps) if int(steps) <= 8 else 4)))
-            flux_prompt = (
-                enhanced
-                + ", 1girl, solo, adult East Asian OR Southeast Asian woman 21+ only, "
-                "late 20s, no men, no males, no old people, "
-                "tall slim small frame tiny waist, large full heavy breasts"
-            )
-            for s in seeds:
-                gen = torch.Generator(device=device).manual_seed(int(s))
-                out = pipe(
-                    prompt=flux_prompt,
-                    num_inference_steps=use_steps,
-                    guidance_scale=0.0,
-                    width=int(width),
-                    height=int(height),
-                    generator=gen,
-                    max_sequence_length=256,
-                ).images[0]
-                label = f"FLUX · seed {s} · {datetime.now().strftime('%I:%M:%S %p')}"
-                gallery.append((out, label))
-                jobs.append((out, s, "FLUX.1-schnell", prompt, username, width, height, use_steps, 0.0))
-        elif engine == ENGINE_SD35:
-            pipe = get_sd35_pipe()
-            use_steps = int(max(15, min(40, int(steps))))
-            for s in seeds:
-                gen = torch.Generator(device=device).manual_seed(int(s))
-                out = pipe(
-                    prompt=enhanced,
-                    negative_prompt=neg,
-                    num_inference_steps=use_steps,
-                    guidance_scale=float(cfg),
-                    width=int(width),
-                    height=int(height),
-                    generator=gen,
-                ).images[0]
-                label = f"SD3.5 · seed {s} · {datetime.now().strftime('%I:%M:%S %p')}"
-                gallery.append((out, label))
-                jobs.append((out, s, "SD3.5-medium", prompt, username, width, height, use_steps, cfg))
-        else:
-            pipe = get_sdxl_pipe(sdxl_model)
-            use_steps = int(max(10, min(40, int(steps))))
-            for s in seeds:
-                gen = torch.Generator(device=device).manual_seed(int(s))
-                out = pipe(
-                    prompt=enhanced,
-                    negative_prompt=neg,
-                    num_inference_steps=use_steps,
-                    guidance_scale=float(cfg),
-                    width=int(width),
-                    height=int(height),
-                    generator=gen,
-                ).images[0]
-                label = f"SDXL · {sdxl_model[:28]} · seed {s}"
-                gallery.append((out, label))
-                jobs.append((out, s, sdxl_model, prompt, username, width, height, use_steps, cfg))
-        info = f"{engine} | seeds {seeds} | {width}x{height} | enhanced prompt ready"
+        # Pass A: txt2img with Chef A (no LoRA on A unless A==B and LoRA set later)
+        pipe_a = load_sdxl_pipe(model_a, LORA_NONE)
+        base_images = []
+        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),
+                width=int(width),
+                height=int(height),
+                generator=gen,
+            ).images[0]
+            base_images.append((out, s))
+
+        # Unload A, load B (+ optional LoRA), img2img remix
+        unload_all()
+        from diffusers import StableDiffusionXLImg2ImgPipeline
+
+        pipe_b = load_sdxl_pipe(model_b, lora)
+        img2img = StableDiffusionXLImg2ImgPipeline(
+            vae=pipe_b.vae,
+            text_encoder=pipe_b.text_encoder,
+            text_encoder_2=pipe_b.text_encoder_2,
+            tokenizer=pipe_b.tokenizer,
+            tokenizer_2=pipe_b.tokenizer_2,
+            unet=pipe_b.unet,
+            scheduler=pipe_b.scheduler,
+        )
+        if torch.cuda.is_available():
+            img2img = img2img.to("cuda")
+
+        remixed = []
+        for img, s in base_images:
+            gen = torch.Generator(device=device).manual_seed(int(s) + 17)
+            out = img2img(
+                image=img,
+                prompt=enhanced,
+                negative_prompt=neg,
+                strength=float(strength),
+                num_inference_steps=remix_steps,
+                guidance_scale=float(cfg),
+                generator=gen,
+            ).images[0]
+            remixed.append((out, s))
+        del img2img
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        lora_tag = lora if lora and lora != LORA_NONE else "no-LoRA"
+        chain = f"{stack_name} | A:{model_a} → B:{model_b}@{strength:.2f} | {lora_tag}"
+        for out, s in remixed:
+            label = f"seed {s} · {stack_name} · A→B {strength:.2f}"
+            gallery.append((out, label))
+            jobs.append((out, s, chain, prompt, username, width, height, use_steps, cfg))
+
+        info = (
+            f"{chain} | seeds {seeds} | {width}x{height} | "
+            f"steps A={use_steps} B={remix_steps} | enhanced prompt ready"
+        )
         return gallery, info, enhanced, jobs
     except Exception as e:
         logger.exception("generate failed")
@@ -565,6 +617,16 @@ def generate(
         return None, f"❌ {msg}", "", []
 
 
+def on_stack_change(stack_name):
+    cfg = STACKS.get(stack_name) or STACKS[DEFAULT_STACK]
+    return (
+        gr.update(value=cfg["chef_a"]),
+        gr.update(value=cfg["chef_b"]),
+        gr.update(value=cfg["lora"]),
+        gr.update(value=cfg["blend"]),
+    )
+
+
 js_visitor = """
 function() {
   let id = localStorage.getItem('adult_ai_visitor_id');
@@ -585,23 +647,29 @@ THEME = gr.themes.Default(primary_hue="purple", secondary_hue="indigo", neutral_
 
 
 def create_ui():
+    default_cfg = STACKS[DEFAULT_STACK]
     with gr.Blocks() as demo:
         user_id = gr.Textbox(visible=False)
-        gr.Markdown("# ⚡ 21+ Adult East/SE Asian Lab")
+        gr.Markdown("# ⚡ 21+ Adult East/SE Asian Lab · Fiction Space · Multi-chef stacks")
         gr.Markdown(
             """### 21+ only · Fiction Space · ZeroGPU · first model load is slow · queues happen
 
 **Scene + body lab for adult East / Southeast Asian women (21+, late 20s fiction).**
 Body lock: tall slim small frame, tiny waist, large/full/heavy breasts, cleavage. No thick thighs.
-Default engine: Adult SDXL. Chef list is **porn / NSFW only** (Illustrious + Pony + photo porn). Default: **WAI NSFW Illustrious v150**.
+Every gen runs **Chef A txt2img → Chef B img2img remix** (no single-chef path).
+Pick a named **stack** chip, or override slots in Advanced.
 Style packs stack: start with **Testing Tits**, layer scene packs (Cosmic e-girl booth, night club, hotel mirror, rain alley, Soft Domme, Quiet beg, library whiskey couch).
-FLUX and SD 3.5 stay optional. Visitor gens are **not** saved unless secret `LOG_VISITOR_UPLOADS=1`."""
+Visitor gens are **not** saved unless secret `LOG_VISITOR_UPLOADS=1`."""
         )
         with gr.Row():
             with gr.Column(scale=2):
-                prompt = gr.Textbox(label="Prompt (scene + body)", value=random.choice(RANDOM_PROMPTS), lines=7)
+                prompt = gr.Textbox(
+                    label="Prompt (scene + body)",
+                    value=random.choice(RANDOM_PROMPTS),
+                    lines=7,
+                )
                 negative = gr.Textbox(
-                    label="Negatives (SD3.5 / SDXL; ignored by FLUX)",
+                    label="Negatives",
                     value=BASE_NEGATIVES,
                     lines=4,
                 )
@@ -610,29 +678,53 @@ FLUX and SD 3.5 stay optional. Visitor gens are **not** saved unless secret `LOG
                     choices=list(STYLE_PACKS.keys()),
                     value=["Testing Tits"],
                 )
+                stack = gr.Radio(
+                    label="Stack (named multi-chef chips)",
+                    choices=STACK_CHOICES,
+                    value=DEFAULT_STACK,
+                )
                 go = gr.Button("GENERATE", variant="primary", size="lg")
             with gr.Column(scale=1):
-                engine = gr.Radio(
-                    label="Engine (Adult SDXL = default)",
-                    choices=ENGINE_CHOICES,
-                    value=ENGINE_SDXL,
-                )
-                sdxl_model = gr.Dropdown(
-                    label="Adult SDXL chef (only used on Adult SDXL engine)",
-                    choices=SDXL_CHOICES,
-                    value=DEFAULT_SDXL,
-                )
+                with gr.Accordion("Advanced (override stack slots)", open=False):
+                    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"],
+                    )
+                    lora_dd = gr.Dropdown(
+                        label="LoRA (None = blend only)",
+                        choices=LORA_CHOICES,
+                        value=default_cfg["lora"],
+                    )
+                    blend_sl = gr.Slider(
+                        label="Blend A→B strength (higher = more Chef B)",
+                        minimum=0.30,
+                        maximum=0.70,
+                        value=default_cfg["blend"],
+                        step=0.05,
+                    )
                 seed = gr.Number(label="Seed (-1 random)", value=-1, precision=0)
                 width = gr.Slider(512, 1536, value=1024, step=64, label="Width")
                 height = gr.Slider(512, 1536, value=1024, step=64, label="Height")
-                steps = gr.Slider(1, 40, value=28, step=1, label="Steps (SDXL~28, FLUX~4, SD3.5~28)")
-                cfg = gr.Slider(0.0, 12.0, value=6.0, step=0.5, label="CFG (SDXL~6, FLUX uses 0)")
+                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")
                 batch = gr.Slider(1, 3, value=1, step=1, label="Batch (max 3)")
         gallery = gr.Gallery(label="Outputs", columns=3, height="auto")
         info = gr.Textbox(label="Run info", visible=True)
         meta = gr.Textbox(label="Enhanced prompt", visible=True)
         demo.load(fn=None, inputs=None, outputs=user_id, js=js_visitor)
 
+        stack.change(
+            fn=on_stack_change,
+            inputs=[stack],
+            outputs=[chef_a, chef_b, lora_dd, blend_sl],
+        )
+
         def on_go(*args):
             images, run_info, enhanced, jobs = generate(*args)
             if jobs:
@@ -644,7 +736,23 @@ FLUX and SD 3.5 stay optional. Visitor gens are **not** saved unless secret `LOG
 
         go.click(
             fn=on_go,
-            inputs=[engine, prompt, styles, negative, sdxl_model, steps, cfg, seed, width, height, batch, user_id],
+            inputs=[
+                prompt,
+                styles,
+                negative,
+                stack,
+                chef_a,
+                chef_b,
+                lora_dd,
+                blend_sl,
+                steps,
+                cfg,
+                seed,
+                width,
+                height,
+                batch,
+                user_id,
+            ],
             outputs=[gallery, info, meta],
         )
     return demo