app.py

import os
import gradio as gr
import numpy as np
import spaces
import torch
import random
import zipfile
import io
import tempfile
from PIL import Image
from diffusers import FluxKontextPipeline, FluxImg2ImgPipeline, FluxTransformer2DModel

MAX_SEED = np.iinfo(np.int32).max

# --- Load Kontext pipeline (identity tab) ---
print("Loading Kontext pipeline...")
kontext_pipe = FluxKontextPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Kontext-dev",
    torch_dtype=torch.bfloat16
).to("cuda")
print("Kontext pipeline loaded.")

# --- Load FLUX.1-dev transformer for explicit tab ---
# Shares text encoders, tokenizers, VAE, and scheduler with Kontext
print("Loading FLUX.1-dev transformer...")
dev_transformer = FluxTransformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    torch_dtype=torch.bfloat16
)
print("FLUX.1-dev transformer loaded.")

explicit_pipe = FluxImg2ImgPipeline(
    transformer=dev_transformer,
    text_encoder=kontext_pipe.text_encoder,
    text_encoder_2=kontext_pipe.text_encoder_2,
    tokenizer=kontext_pipe.tokenizer,
    tokenizer_2=kontext_pipe.tokenizer_2,
    vae=kontext_pipe.vae,
    scheduler=kontext_pipe.scheduler,
).to("cuda")
print("Explicit pipeline assembled.")

# --- Load NSFW LoRAs on explicit pipeline ---
EXPLICIT_ADAPTERS = []
try:
    explicit_pipe.load_lora_weights(
        "Heartsync/Flux-NSFW-uncensored",
        weight_name="lora.safetensors",
        adapter_name="nsfw_1"
    )
    EXPLICIT_ADAPTERS.append("nsfw_1")
    print("Loaded Heartsync NSFW LoRA on explicit pipeline")
except Exception as e:
    print(f"Heartsync LoRA failed: {e}")

try:
    explicit_pipe.load_lora_weights(
        "lustlyai/Flux_Lustly.ai_Uncensored_nsfw_v1",
        weight_name="flux_lustly-ai_v1.safetensors",
        adapter_name="nsfw_2"
    )
    EXPLICIT_ADAPTERS.append("nsfw_2")
    print("Loaded Lustly NSFW LoRA on explicit pipeline")
except Exception as e:
    print(f"Lustly LoRA failed: {e}")

if EXPLICIT_ADAPTERS:
    explicit_pipe.set_adapters(EXPLICIT_ADAPTERS, adapter_weights=[1.0] * len(EXPLICIT_ADAPTERS))
    print(f"Explicit NSFW LoRAs active: {EXPLICIT_ADAPTERS}")

# --- Also load LoRAs on Kontext for semi-nude shots ---
KONTEXT_ADAPTERS = []
try:
    kontext_pipe.load_lora_weights(
        "Heartsync/Flux-NSFW-uncensored",
        weight_name="lora.safetensors",
        adapter_name="nsfw_1"
    )
    KONTEXT_ADAPTERS.append("nsfw_1")
except Exception as e:
    print(f"Kontext Heartsync LoRA failed: {e}")

# --- Realism LoRA (fights the plastic FLUX look) on both pipelines ---
# Shakker-Labs add-details: no trigger word, natural skin, recommended scale 1.0
KONTEXT_REALISM = []
EXPLICIT_REALISM = []
for pipe, bucket, label in ((kontext_pipe, KONTEXT_REALISM, "Kontext"),
                            (explicit_pipe, EXPLICIT_REALISM, "Explicit")):
    try:
        pipe.load_lora_weights(
            "Shakker-Labs/FLUX.1-dev-LoRA-add-details",
            weight_name="FLUX-dev-lora-add_details.safetensors",
            adapter_name="realism",
        )
        bucket.append("realism")
        print(f"Loaded realism LoRA on {label} pipeline")
    except Exception as e:
        print(f"{label} realism LoRA failed: {e}")

if KONTEXT_ADAPTERS or KONTEXT_REALISM:
    names = KONTEXT_ADAPTERS + KONTEXT_REALISM
    kontext_pipe.set_adapters(names, adapter_weights=[1.0] * len(names))
    print(f"Kontext LoRAs active: {names}")

CANDID_SUFFIX = (
    ", shot on iPhone, candid, handheld, slightly out of focus, natural grain, "
    "hint of motion blur, imperfect framing, real photo, not posed"
)

STREET_PROMPTS = [
    "Same person, white t-shirt and jeans, mirror selfie, apartment bathroom, morning light, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, gray hoodie, coffee shop selfie, warm cafe light, slight smile, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, tank top, city sidewalk, midday sun, squinting slightly, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, gym shorts and tank top, gym selfie, sweaty, fluorescent light overhead, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, button shirt open at collar, casual restaurant booth, evening warm light, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, board shorts, poolside selfie, bright afternoon sun, wet hair, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, leather jacket, street corner at night, city lights behind, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, towel around waist, bedroom, just out of the shower, damp skin, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, baseball cap on backward, park bench, afternoon, trees in background, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, polo shirt, driving selfie, sunglasses on, car interior visible, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, flannel shirt, hiking trail, outdoors, slightly breathless, trees around, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, shirtless, beach, bright sun, sand in background, looking out at water, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, dress shirt, office bathroom mirror selfie, midday, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, track jacket, running path, earbuds in, post-run, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, oversized tee, home couch, TV glow, relaxed evening, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, long sleeve henley, corner bar, low warm light, casual, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, crop hoodie, rooftop, golden hour light, city skyline behind, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, shirtless, morning bed selfie, sheets around waist, morning light from window, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, compression shorts only, home gym, post-workout, slight pump, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, fleece vest, mountain trail, overcast sky, outdoors, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, flannel shirt, outdoor farmers market, morning crowd behind, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, muscle tee, dim music venue, colored stage lighting behind, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, cargo shorts and t-shirt, beach boardwalk, afternoon, ocean behind, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, loose linen shirt, outdoor patio, summer soft shade, relaxed, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
    "Same person, shirtless, back to camera, looking over shoulder in bathroom mirror, shot on iPhone, candid, handheld, slightly out of focus, natural grain, hint of motion blur, imperfect framing, real photo, not posed",
]


def _resize_for_model(image, max_dim=1024):
    image = image.convert("RGB")
    w, h = image.size
    if max(w, h) > max_dim:
        scale = max_dim / max(w, h)
        w, h = int(w * scale), int(h * scale)
    w, h = (w // 8) * 8, (h // 8) * 8
    return image.resize((w, h), Image.LANCZOS), w, h


@spaces.GPU(duration=120)
def generate_identity(reference_image, prompt, seed, randomize_seed, guidance_scale, steps, nsfw_strength, realism_strength, candid, progress=gr.Progress(track_tqdm=True)):
    if not prompt or not prompt.strip():
        raise gr.Error("Write a prompt first!")

    names = KONTEXT_ADAPTERS + KONTEXT_REALISM
    if names:
        weights = [nsfw_strength] * len(KONTEXT_ADAPTERS) + [realism_strength] * len(KONTEXT_REALISM)
        kontext_pipe.set_adapters(names, adapter_weights=weights)

    if candid:
        prompt = prompt.rstrip().rstrip(",") + CANDID_SUFFIX

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator("cuda").manual_seed(seed)
    kwargs = dict(
        prompt=prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=steps,
        generator=generator,
    )

    if reference_image is not None:
        reference_image, w, h = _resize_for_model(reference_image)
        kwargs.update(image=reference_image, width=w, height=h)

    result = kontext_pipe(**kwargs).images[0]
    return result, seed


@spaces.GPU(duration=120)
def generate_explicit(input_image, char_desc, prompt, seed, randomize_seed, guidance_scale, steps, strength, nsfw_strength, realism_strength, candid, progress=gr.Progress(track_tqdm=True)):
    if not prompt or not prompt.strip():
        raise gr.Error("Write a prompt first!")
    if input_image is None:
        raise gr.Error("Upload a starting image! Use an identity shot from the other tab.")

    # Prepend character description so the model anchors to this person's face/features
    full_prompt = f"{char_desc.strip()}, {prompt}" if char_desc and char_desc.strip() else prompt
    if candid:
        full_prompt = full_prompt.rstrip().rstrip(",") + CANDID_SUFFIX

    names = EXPLICIT_ADAPTERS + EXPLICIT_REALISM
    if names:
        weights = [nsfw_strength] * len(EXPLICIT_ADAPTERS) + [realism_strength] * len(EXPLICIT_REALISM)
        explicit_pipe.set_adapters(names, adapter_weights=weights)

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator("cuda").manual_seed(seed)
    input_image, w, h = _resize_for_model(input_image)

    result = explicit_pipe(
        image=input_image,
        prompt=full_prompt,
        strength=strength,
        guidance_scale=guidance_scale,
        width=w,
        height=h,
        num_inference_steps=steps,
        generator=generator,
    ).images[0]
    return result, seed


@spaces.GPU(duration=700)
def generate_batch(reference_image, count, progress=gr.Progress()):
    if reference_image is None:
        raise gr.Error("Upload a reference photo first!")
    # NSFW LoRA off (adds a polished AI look), realism LoRA on — the candid
    # prompt suffix + realism LoRA together sell the "real photo" feel
    names = KONTEXT_ADAPTERS + KONTEXT_REALISM
    if names:
        kontext_pipe.set_adapters(
            names,
            adapter_weights=[0.0] * len(KONTEXT_ADAPTERS) + [1.0] * len(KONTEXT_REALISM),
        )
    gallery_images = []
    try:
        reference_image, w, h = _resize_for_model(reference_image)
        for i in range(count):
            prompt = STREET_PROMPTS[i % len(STREET_PROMPTS)]
            seed = random.randint(0, MAX_SEED)
            generator = torch.Generator("cuda").manual_seed(seed)
            progress(i / count, desc=f"Generating {i + 1} of {count}...")
            result = kontext_pipe(
                prompt=prompt,
                image=reference_image,
                width=w,
                height=h,
                guidance_scale=2.5,
                num_inference_steps=28,
                generator=generator,
            ).images[0]
            gallery_images.append(result)
            n = len(gallery_images)
            status = f"Generating {i + 1} of {count}..."
            yield gallery_images, gallery_images, f"**{n}** image{'s' if n != 1 else ''} in training set", status, SPINNER_HTML, f"**{i + 1} of {count}**"
        done_msg = f"Done -- {count} shots added to training set."
        yield gallery_images, gallery_images, f"**{count}** image{'s' if count != 1 else ''} in training set", done_msg, "", ""
    finally:
        names = KONTEXT_ADAPTERS + KONTEXT_REALISM
        if names:
            kontext_pipe.set_adapters(names, adapter_weights=[1.0] * len(names))


def add_to_gallery(image, gallery_images):
    if image is None:
        raise gr.Error("Generate an image first!")
    if gallery_images is None:
        gallery_images = []
    gallery_images = gallery_images + [image]
    n = len(gallery_images)
    return gallery_images, gallery_images, f"**{n}** image{'s' if n != 1 else ''} in training set"


def download_zip(gallery_images):
    if not gallery_images:
        raise gr.Error("Training set is empty!")
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".zip", prefix="dreamboat_")
    with zipfile.ZipFile(tmp.name, "w", zipfile.ZIP_DEFLATED) as zf:
        for i, img in enumerate(gallery_images):
            buf = io.BytesIO()
            if isinstance(img, str):
                img = Image.open(img)
            elif isinstance(img, dict) and "name" in img:
                img = Image.open(img["name"])
            elif isinstance(img, np.ndarray):
                img = Image.fromarray(img)
            elif isinstance(img, tuple):
                # Gallery items can round-trip as (image, caption) tuples
                first = img[0]
                img = Image.open(first) if isinstance(first, str) else (
                    Image.fromarray(first) if isinstance(first, np.ndarray) else first
                )
            if not isinstance(img, Image.Image):
                continue
            img.save(buf, format="PNG")
            zf.writestr(f"training_{i:03d}.png", buf.getvalue())
    return tmp.name


def clear_all():
    return [], [], "**0** images in training set"


def use_result_as_ref(result_image):
    return result_image


def make_preset_fn(text):
    def fn():
        return text
    return fn


SPINNER_HTML = """
<div style="text-align:center;padding:12px 0 4px 0;">
  <div style="
    width:60px;height:60px;border-radius:50%;
    background:conic-gradient(#E40303,#FF8C00,#FFED00,#008026,#24408E,#732982,#E40303);
    animation:rainbowSpin 1s linear infinite;
    display:inline-flex;align-items:center;justify-content:center;">
    <div style="width:44px;height:44px;border-radius:50%;background:#fff;"></div>
  </div>
</div>
"""

css = """
#col-container { margin: 0 auto; max-width: 1200px; }
#header-text {
    text-align: center;
    background: linear-gradient(135deg, #E40303, #FF8C00, #FFED00, #008026, #24408E, #732982);
    -webkit-background-clip: text;
    background-clip: text;
    -webkit-text-fill-color: transparent;
    font-size: 3em;
    font-weight: 800;
    margin: 0;
    padding: 10px 0 0 0;
}
#tagline {
    text-align: center;
    color: #999;
    margin-top: 0;
    font-style: italic;
}
@keyframes rainbowSpin {
    to { transform: rotate(360deg); }
}
"""

# ============ Quick Start Scenarios (one-click: sets prompt + every setting + tells
# you what the input should be) ============
# Identity (Norma Jean / Kontext) scenarios. Kontext re-poses from a reference, so the
# input hint is about what reference photo gives the cleanest result.
IDENTITY_SCENARIOS = {
    "Face close-up (studio)": {
        "prompt": "Same person, face close-up portrait, sharp focus, looking toward camera, even studio lighting, clean solid background",
        "candid": False, "nsfw": 0.3, "realism": 1.0, "guidance": 2.5, "steps": 28,
        "hint": "a clear, well-lit photo where the FACE is visible. Front-on or a slight angle is best.",
    },
    "3/4 angle portrait": {
        "prompt": "Same person, three-quarter angle view, upper body, soft directional light from the left, subtle natural expression",
        "candid": False, "nsfw": 0.3, "realism": 1.0, "guidance": 2.5, "steps": 28,
        "hint": "any clear face photo. Head-and-shoulders works great.",
    },
    "Full body (training staple)": {
        "prompt": "Same person, full body standing, neutral relaxed pose, even studio lighting, solid gray background",
        "candid": False, "nsfw": 0.4, "realism": 1.0, "guidance": 2.5, "steps": 28,
        "hint": "a photo that shows his build if you have one. Face-only still works — the body is inferred.",
    },
    "Shirtless torso": {
        "prompt": "Same person, shirtless, upper body and torso, warm studio light, relaxed confident expression, simple background",
        "candid": False, "nsfw": 1.0, "realism": 1.0, "guidance": 2.5, "steps": 28,
        "hint": "face + upper body if possible. Kontext keeps the face and generates the torso.",
    },
    "Candid mirror selfie": {
        "prompt": "Same person, casual mirror selfie, phone in hand, bedroom or bathroom, natural window light, authentic",
        "candid": True, "nsfw": 0.5, "realism": 1.0, "guidance": 2.5, "steps": 28,
        "hint": "any clear face photo. This one leans hardest on candid realism.",
    },
    "Golden hour outdoor": {
        "prompt": "Same person, outdoors at golden hour, warm low sunlight, half body, relaxed candid feel",
        "candid": True, "nsfw": 0.4, "realism": 1.0, "guidance": 2.5, "steps": 28,
        "hint": "a clear face photo — the lighting gets replaced with golden hour.",
    },
    "Dramatic cinematic": {
        "prompt": "Same person, dramatic side lighting, moody shadows, face and torso, high contrast, cinematic",
        "candid": False, "nsfw": 0.5, "realism": 1.0, "guidance": 3.0, "steps": 30,
        "hint": "a sharp, well-exposed face photo. Dramatic lighting needs clean features to work from.",
    },
    "Over-the-shoulder / back": {
        "prompt": "Same person seen from behind, looking back over the shoulder, upper body and face visible, natural light",
        "candid": False, "nsfw": 0.6, "realism": 1.0, "guidance": 2.5, "steps": 28,
        "hint": "a front-facing face photo is fine — the model turns the pose for you.",
    },
}

# Explicit (Marilyn / FLUX.1-dev img2img) scenarios. This is img2img: it KEEPS the
# composition of the starting image, so the input hint really matters here.
EXPLICIT_SCENARIOS = {
    "Full nude — standing": {
        "prompt": "same person, fully nude, full body standing, anatomically correct male, penis visible, studio lighting, solid neutral background, realistic, high quality",
        "strength": 0.80, "guidance": 3.5, "steps": 28, "nsfw": 1.0, "realism": 1.0, "candid": False,
        "hint": "a FULL-BODY STANDING shot. Img2img keeps the pose — start standing to get a standing nude.",
    },
    "Nude — reclining / lying down": {
        "prompt": "same person, nude, reclining on bed, relaxed, full body visible, anatomically detailed male, warm soft bedroom light, intimate, realistic",
        "strength": 0.80, "guidance": 3.5, "steps": 28, "nsfw": 1.0, "realism": 1.0, "candid": False,
        "hint": "a shot where he is ALREADY LYING DOWN or reclining. Img2img won't cleanly turn a standing photo into a lying-down one.",
    },
    "Shower / wet": {
        "prompt": "same person, nude, wet skin, standing in shower, water droplets, steam, full body, anatomically correct, natural lighting, realistic",
        "strength": 0.80, "guidance": 3.5, "steps": 28, "nsfw": 1.0, "realism": 1.0, "candid": False,
        "hint": "a STANDING full-body shot.",
    },
    "Nude from behind": {
        "prompt": "same person, nude from behind, athletic back and buttocks visible, looking over shoulder, studio lighting, realistic",
        "strength": 0.80, "guidance": 3.5, "steps": 28, "nsfw": 1.0, "realism": 1.0, "candid": False,
        "hint": "a standing shot; a back or 3/4 view works best.",
    },
    "Keep-the-face light nude": {
        "prompt": "same person, nude, natural relaxed pose, anatomically correct male, soft lighting, realistic",
        "strength": 0.60, "guidance": 3.5, "steps": 28, "nsfw": 1.0, "realism": 1.0, "candid": False,
        "hint": "your BEST identity shot. Low strength (0.60) preserves the face — also fill in the Character Description box.",
    },
    "Explicit — max freedom": {
        "prompt": "same person, nude, explicit, anatomically detailed male, realistic, high quality",
        "strength": 0.85, "guidance": 3.5, "steps": 28, "nsfw": 1.0, "realism": 1.0, "candid": False,
        "hint": "a full-body shot. High strength (0.85) changes the body most — fill in the Character Description box to hold the face.",
    },
}


def apply_identity_scenario(name):
    s = IDENTITY_SCENARIOS.get(name)
    if not s:
        return (gr.update(),) * 7
    return (
        gr.update(value=s["prompt"]),
        gr.update(value=s["candid"]),
        gr.update(value=s["nsfw"]),
        gr.update(value=s["realism"]),
        gr.update(value=s["guidance"]),
        gr.update(value=s["steps"]),
        gr.update(value=f"**Best input for this scenario:** {s['hint']}"),
    )


def apply_explicit_scenario(name):
    s = EXPLICIT_SCENARIOS.get(name)
    if not s:
        return (gr.update(),) * 7
    return (
        gr.update(value=s["prompt"]),
        gr.update(value=s["strength"]),
        gr.update(value=s["nsfw"]),
        gr.update(value=s["realism"]),
        gr.update(value=s["guidance"]),
        gr.update(value=s["steps"]),
        gr.update(value=f"**Best input for this scenario:** {s['hint']}"),
    )

with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    gallery_state = gr.State([])

    with gr.Column(elem_id="col-container"):
        gr.HTML('<p id="header-text">DreamBoat</p>')
        gr.HTML('<p id="tagline">Character training image generator &#8212; make him consistent, keep him hot</p>')

        with gr.Accordion("How to Use DreamBoat", open=False):
            gr.Markdown("""
**DreamBoat generates training images for character LoRAs** using two models:

**Identity tab (Kontext)** -- Upload a reference photo. Kontext maintains the character's face and body
across different angles, lighting, and clothing levels. Best for: face shots, body shots, shirtless, varied angles.

**Explicit tab (FLUX.1-dev)** -- Takes an image from the Identity tab and transforms it with full NSFW capability.
FLUX.1-dev has weaker safety training, so the NSFW LoRAs actually work here. Best for: nude, anatomical, explicit.

**Workflow:**
1. Generate 10-15 identity shots on the Identity tab (varied angles, lighting, clothed/shirtless)
2. Take your best full-body shot and send it to the Explicit tab
3. Generate 5-10 explicit variations there
4. Download the combined training set as a zip
5. Feed into [ai-toolkit](https://github.com/ostris/ai-toolkit) for LoRA training
            """)

        # --- SHARED GALLERY (right side, always visible) ---
        with gr.Row():
            # LEFT: Tabs
            with gr.Column(scale=1):
                with gr.Tabs():
                    # === IDENTITY TAB ===
                    with gr.TabItem("Norma Jean"):
                        id_reference = gr.Image(
                            label="Reference Photo",
                            type="pil",
                            height=300,
                        )

                        with gr.Group():
                            gr.Markdown("**🚀 Quick Start — pick a scenario, everything gets set up for you**")
                            with gr.Row():
                                id_scenario_dd = gr.Dropdown(
                                    choices=list(IDENTITY_SCENARIOS.keys()),
                                    label="Scenario", value=None, scale=3,
                                )
                                id_scenario_btn = gr.Button("Apply", variant="secondary", scale=1)
                            id_scenario_hint = gr.Markdown("")

                        id_prompt = gr.Textbox(
                            label="Variation Prompt",
                            placeholder="Same person, shirtless, different angle...",
                            lines=2,
                        )
                        with gr.Row():
                            id_gen_btn = gr.Button("Generate", variant="primary", scale=2)
                            id_gen_add_btn = gr.Button("Generate & Add", variant="secondary", scale=2)

                        id_candid = gr.Checkbox(
                            label="Candid realism mode (kills the AI gloss)",
                            value=True,
                            info="Appends the iPhone/candid suffix to your prompt. Turn off for studio-look shots.",
                        )

                        with gr.Accordion("Advanced", open=False):
                            id_seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                            id_randomize = gr.Checkbox(label="Randomize seed", value=True)
                            id_guidance = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=2.5)
                            id_steps = gr.Slider(label="Steps", minimum=4, maximum=30, value=28, step=1)
                            id_nsfw = gr.Slider(
                                label="NSFW Strength", minimum=0.0, maximum=1.5, step=0.05, value=1.0,
                                info="Helps with shirtless/semi-nude. Kontext has strong safety training. High values add a polished AI look — drop toward 0 for realistic clothed shots."
                            )
                            id_realism = gr.Slider(
                                label="Realism Strength", minimum=0.0, maximum=1.5, step=0.05, value=1.0,
                                info="Shakker add-details LoRA — natural skin texture and detail."
                            )

                    # === STREET PACK TAB ===
                    with gr.TabItem("Street Pack"):
                        batch_reference = gr.Image(
                            label="Reference Photo",
                            type="pil",
                            height=300,
                        )
                        batch_count = gr.Slider(
                            label="How many shots",
                            minimum=5,
                            maximum=30,
                            value=20,
                            step=1,
                        )
                        batch_btn = gr.Button("Generate Pack", variant="primary")
                        batch_status = gr.Textbox(
                            label="Progress",
                            interactive=False,
                            lines=2,
                            value="Upload a photo and hit Generate Pack.",
                        )

                    # === EXPLICIT TAB ===
                    with gr.TabItem("Marilyn"):
                        ex_input = gr.Image(
                            label="Starting Image (use an identity shot from the other tab)",
                            type="pil",
                            height=300,
                        )
                        ex_char_desc = gr.Textbox(
                            label="Character Description (describe face/hair/features to anchor identity)",
                            placeholder="e.g. dark brown hair, brown eyes, olive skin, strong jaw, early 30s, athletic build",
                            lines=1,
                            info="Gets prepended to every prompt. The more specific, the more the output looks like your guy.",
                        )
                        with gr.Group():
                            gr.Markdown("**🚀 Quick Start — pick a scenario, everything gets set up for you**")
                            with gr.Row():
                                ex_scenario_dd = gr.Dropdown(
                                    choices=list(EXPLICIT_SCENARIOS.keys()),
                                    label="Scenario", value=None, scale=3,
                                )
                                ex_scenario_btn = gr.Button("Apply", variant="secondary", scale=1)
                            ex_scenario_hint = gr.Markdown("")

                        ex_prompt = gr.Textbox(
                            label="Explicit Prompt",
                            placeholder="same person, fully nude, anatomically correct...",
                            lines=2,
                        )
                        with gr.Row():
                            ex_gen_btn = gr.Button("Generate", variant="primary", scale=2)
                            ex_gen_add_btn = gr.Button("Generate & Add", variant="secondary", scale=2)

                        with gr.Accordion("Advanced", open=False):
                            ex_seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                            ex_randomize = gr.Checkbox(label="Randomize seed", value=True)
                            ex_guidance = gr.Slider(
                                label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=3.5,
                                info="3.5 is standard for FLUX.1-dev."
                            )
                            ex_steps = gr.Slider(label="Steps", minimum=4, maximum=50, value=28, step=1)
                            ex_strength = gr.Slider(
                                label="Transformation Strength", minimum=0.1, maximum=1.0, step=0.05, value=0.60,
                                info="Lower = more face preserved, less change. Try 0.55-0.65 to keep identity, 0.75+ for more freedom."
                            )
                            ex_nsfw = gr.Slider(
                                label="NSFW Strength", minimum=0.0, maximum=1.5, step=0.05, value=1.0,
                                info="1.0 should be enough on FLUX.1-dev. Raise if needed."
                            )
                            ex_realism = gr.Slider(
                                label="Realism Strength", minimum=0.0, maximum=1.5, step=0.05, value=1.0,
                                info="Shakker add-details LoRA — natural skin texture and detail."
                            )
                            ex_candid = gr.Checkbox(
                                label="Candid realism mode",
                                value=False,
                                info="Appends the iPhone/candid suffix. Off by default for explicit shots.",
                            )

            # RIGHT: Output + Gallery
            with gr.Column(scale=1):
                # type="pil" matters: without it Gradio passes numpy arrays to
                # add_to_gallery, and download_zip silently drops non-PIL entries
                result_image = gr.Image(label="Result", interactive=False, height=350, type="pil")
                seed_output = gr.Number(label="Seed", interactive=False)

                with gr.Row():
                    add_btn = gr.Button("Add to Set", scale=1)
                    use_as_ref_btn = gr.Button("Use as Identity Ref", scale=1)
                    send_to_explicit_btn = gr.Button("Send to Explicit", scale=1)

                batch_spinner = gr.HTML("")
                batch_progress = gr.Markdown("", elem_id="batch-progress-text")

                gr.HTML('<hr style="margin:10px 0;">')
                image_count = gr.Markdown("**0** images in training set")
                gallery = gr.Gallery(
                    label="Training Set",
                    columns=4,
                    height=250,
                    object_fit="cover",
                    show_label=False,
                )
                with gr.Row():
                    download_btn = gr.Button("Download .zip", variant="primary", scale=2)
                    clear_btn = gr.Button("Clear", variant="stop", scale=1)
                download_file = gr.File(visible=False)

        gr.HTML(
            '<p style="text-align:center;color:#666;font-size:0.85em;margin-top:16px;">'
            'DreamBoat is a companion to <a href="https://huggingface.co/spaces/Jds20001/PrideWAN">PrideWAN</a>. '
            "Train your character LoRAs here, animate them there.</p>"
        )

    # --- Quick Start scenario events ---
    # Selecting a scenario auto-fills the prompt + all settings (no Apply click needed);
    # the Apply button stays as a redundant option.
    id_scenario_outputs = [id_prompt, id_candid, id_nsfw, id_realism, id_guidance, id_steps, id_scenario_hint]
    id_scenario_dd.change(fn=apply_identity_scenario, inputs=[id_scenario_dd], outputs=id_scenario_outputs)
    id_scenario_btn.click(fn=apply_identity_scenario, inputs=[id_scenario_dd], outputs=id_scenario_outputs)

    ex_scenario_outputs = [ex_prompt, ex_strength, ex_nsfw, ex_realism, ex_guidance, ex_steps, ex_scenario_hint]
    ex_scenario_dd.change(fn=apply_explicit_scenario, inputs=[ex_scenario_dd], outputs=ex_scenario_outputs)
    ex_scenario_btn.click(fn=apply_explicit_scenario, inputs=[ex_scenario_dd], outputs=ex_scenario_outputs)

    # --- Identity tab events ---
    id_gen_btn.click(
        fn=generate_identity,
        inputs=[id_reference, id_prompt, id_seed, id_randomize, id_guidance, id_steps, id_nsfw, id_realism, id_candid],
        outputs=[result_image, seed_output],
    )
    id_prompt.submit(
        fn=generate_identity,
        inputs=[id_reference, id_prompt, id_seed, id_randomize, id_guidance, id_steps, id_nsfw, id_realism, id_candid],
        outputs=[result_image, seed_output],
    )
    id_gen_add_btn.click(
        fn=generate_identity,
        inputs=[id_reference, id_prompt, id_seed, id_randomize, id_guidance, id_steps, id_nsfw, id_realism, id_candid],
        outputs=[result_image, seed_output],
    ).then(
        fn=add_to_gallery,
        inputs=[result_image, gallery_state],
        outputs=[gallery_state, gallery, image_count],
    )

    # --- Explicit tab events ---
    ex_gen_btn.click(
        fn=generate_explicit,
        inputs=[ex_input, ex_char_desc, ex_prompt, ex_seed, ex_randomize, ex_guidance, ex_steps, ex_strength, ex_nsfw, ex_realism, ex_candid],
        outputs=[result_image, seed_output],
    )
    ex_prompt.submit(
        fn=generate_explicit,
        inputs=[ex_input, ex_char_desc, ex_prompt, ex_seed, ex_randomize, ex_guidance, ex_steps, ex_strength, ex_nsfw, ex_realism, ex_candid],
        outputs=[result_image, seed_output],
    )
    ex_gen_add_btn.click(
        fn=generate_explicit,
        inputs=[ex_input, ex_char_desc, ex_prompt, ex_seed, ex_randomize, ex_guidance, ex_steps, ex_strength, ex_nsfw, ex_realism, ex_candid],
        outputs=[result_image, seed_output],
    ).then(
        fn=add_to_gallery,
        inputs=[result_image, gallery_state],
        outputs=[gallery_state, gallery, image_count],
    )

    # --- Shared events ---
    add_btn.click(
        fn=add_to_gallery,
        inputs=[result_image, gallery_state],
        outputs=[gallery_state, gallery, image_count],
    )
    use_as_ref_btn.click(
        fn=use_result_as_ref,
        inputs=[result_image],
        outputs=[id_reference],
    )
    send_to_explicit_btn.click(
        fn=use_result_as_ref,
        inputs=[result_image],
        outputs=[ex_input],
    )
    download_btn.click(
        fn=download_zip,
        inputs=[gallery_state],
        outputs=[download_file],
    )
    clear_btn.click(
        fn=clear_all,
        outputs=[gallery_state, gallery, image_count],
    )

    # --- Street Pack batch events ---
    batch_btn.click(
        fn=generate_batch,
        inputs=[batch_reference, batch_count],
        outputs=[gallery_state, gallery, image_count, batch_status, batch_spinner, batch_progress],
        show_progress=False,
    )

demo.launch(ssr_mode=False)