app.py

import os
import gc
import gradio as gr
import numpy as np
import spaces
import torch
import random
import uuid
from PIL import Image, ImageOps
from typing import Iterable
import time
import threading
from logging_utils import log_inference

# ── requirements.txt should include: ─────────────────────────────────────────
#   spandrel
# ─────────────────────────────────────────────────────────────────────────────

MODEL_VARIANT = os.environ.get("MODEL_VARIANT", "9B")
if MODEL_VARIANT == "9B-KV":
    from diffusers import Flux2KleinKVPipeline as _PipeClass
    _MODEL_REPO = "black-forest-labs/FLUX.2-klein-9b-kv"
else:
    from diffusers import Flux2KleinPipeline as _PipeClass
    _MODEL_REPO = "black-forest-labs/FLUX.2-klein-9B"
    MODEL_VARIANT = "9B"  # normalise any typos back to default
from huggingface_hub import hf_hub_download

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

colors.orange_red = colors.Color(
    name="orange_red", c50="#FFF0E5", c100="#FFE0CC", c200="#FFC299", c300="#FFA366",
    c400="#FF8533", c500="#FF4500", c600="#E63E00", c700="#CC3700", c800="#B33000",
    c900="#992900", c950="#802200",
)

class OrangeRedTheme(Soft):
    def __init__(
        self, *, primary_hue: colors.Color | str = colors.gray,
        secondary_hue: colors.Color | str = colors.orange_red,
        neutral_hue: colors.Color | str = colors.slate, text_size: sizes.Size | str = sizes.text_lg,
        font: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
        ),
        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue, secondary_hue=secondary_hue, neutral_hue=neutral_hue,
            text_size=text_size, font=font, font_mono=font_mono,
        )
        super().set(
            background_fill_primary="*primary_50",
            background_fill_primary_dark="*primary_900",
            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
            button_primary_text_color="white",
            button_primary_text_color_hover="white",
            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
            slider_color="*secondary_500",
            slider_color_dark="*secondary_600",
            block_title_text_weight="600", block_border_width="3px",
            block_shadow="*shadow_drop_lg", button_primary_shadow="*shadow_drop_lg",
            button_large_padding="11px", color_accent_soft="*primary_100",
            block_label_background_fill="*primary_200",
        )

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

# ── Upscaler model registry ───────────────────────────────────────────────────
# Add new entries here to extend the dropdown — spandrel handles any ESRGAN-
# family .pth file automatically. scale= is the integer output multiplier.
UPSCALE_MODELS = {
    "None": {
        "scale": None, "file": None, "url": None,
    },
    "2× — RealESRGAN (balanced)": {
        "scale": 2,
        "file": "RealESRGAN_x2plus.pth",
        "url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
    },
    "4× — RealESRGAN (balanced)": {
        "scale": 4,
        "file": "RealESRGAN_x4plus.pth",
        "url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
    },
    "4× — UltraSharp (crisp)": {
        "scale": 4,
        "file": "4x-UltraSharpV2.pth",
        "url": "https://huggingface.co/Kim2091/UltraSharpV2/resolve/main/4x-UltraSharpV2.pth",
    },
    "4× — Remacri (natural)": {
        "scale": 4,
        "file": "4x_foolhardy_Remacri.pth",
        "url": "https://huggingface.co/FacehugmanIII/4x_foolhardy_Remacri/resolve/main/4x_foolhardy_Remacri.pth",
    },
    "4× — Nomos2 HQ DAT2 (Photography)": {
        "scale": 4,
        "file": "4xNomos2_hq_dat2.pth",
        "url": "https://github.com/Phhofm/models/releases/download/4xNomos2_hq_dat2/4xNomos2_hq_dat2.pth",
    },    
}
# ─────────────────────────────────────────────────────────────────────────────

FACE_SWAP_PROMPT = """head_swap: start with Picture 1 as the base image, keeping its lighting, environment, and background. Remove the head from Picture 1 completely and replace it with the head from Picture 2.
FROM PICTURE 1 (strictly preserve):
- Scene: lighting conditions, shadows, highlights, color temperature, environment, background
- Head positioning: exact rotation angle, tilt, direction the head is facing
- Expression: facial expression, micro-expressions, eye gaze direction, mouth position, emotion
FROM PICTURE 2 (strictly preserve identity):
- Facial structure: face shape, bone structure, jawline, chin
- All facial features: eye color, eye shape, nose structure, lip shape and fullness, eyebrows
- Hair: color, style, texture, hairline
- Skin: texture, tone, complexion
The replaced head must seamlessly match Picture 1's lighting and expression while maintaining the complete identity from Picture 2. High quality, photorealistic, sharp details, 4k."""

# MAX number of simultaneous LoRA weight sliders to pre-render in the UI
MAX_LORA_SLOTS = 6

LORA_STYLES = [
    {
        "image": "https://huggingface.co/spaces/prithivMLmods/FLUX.2-Klein-LoRA-Studio/resolve/main/examples/image.webp",
        "title": "None",
        "adapter_name": None,
        "repo": None,
        "weights": None,
        "default_prompt": None,
        "default_weight": 1.0,
    },
    {
        "title": "Klein-Delight-Style",
        "adapter_name": "klein-delight",
        "repo": "linoyts/Flux2-Klein-Delight-LoRA",
        "weights": "pytorch_lora_weights.safetensors",
        "default_prompt": "Relight the image to remove all existing lighting conditions and replace them with neutral, uniform illumination. Apply soft, evenly distributed lighting with no directional shadows, no harsh highlights, and no dramatic contrast. Maintain the original identity of all subjects exactly—preserve facial structure, skin tone, proportions, expressions, hair, clothing, and textures. Do not alter pose, camera angle, background geometry, or image composition. Lighting should appear balanced, and studio-neutral, similar to diffuse overcast or a soft lightbox setup. Ensure consistent exposure across the entire image with realistic depth and subtle shading only where necessary for form.",
        "default_weight": 1.0,
    },
    {
        "title": "Klein-Consistency",
        "adapter_name": "klein-consistency",
        "repo": "dx8152/Flux2-Klein-9B-Consistency",
        "weights": "Klein-consistency.safetensors",
        "default_prompt": None,
        "default_weight": 0.3,
    },
    {
        "title": "Best-Face-Swap",
        "adapter_name": "face-swap",
        "repo": "Alissonerdx/BFS-Best-Face-Swap",
        "weights": "bfs_head_v1_flux-klein_9b_step3750_rank64.safetensors",
        "default_prompt": FACE_SWAP_PROMPT,
        "default_weight": 1.0,
    },
    {
        "title": "NSFW v2",
        "adapter_name": "nsfw-v2",
        "repo": "diroverflo/FLux_Klein_9B_NSFW",
        "weights": "Flux Klein - NSFW v2.safetensors",
        "default_prompt": None,
        "default_weight": 1.0,
    },
    {
        "title": "Ultimate Upscaler Klein-9b",
        "adapter_name": "Ultimate Upscaler",
        "repo": "loras",
        "weights": "Flux2-Klein-Image-RestoreV1.safetensors",
        "default_prompt": "restore the image quality, remove any compression artefacts, remove any haze and soft edges, enrich the original with new intricate detail in all textures and surfaces creating a professional photorealistic photograph with natural lighting and skin texture.",
        "default_weight": 1.0,
    },    
    {
        "title": "High Resolution",
        "adapter_name": "High Resolution",
        "repo": "loras",
        "weights": "HighResolution9B.safetensors",
        "default_prompt": "High Resolution",
        "default_weight": 1.0,
    },       
    {
        "title": "Female Asshole",
        "adapter_name": "Female Asshole",
        "repo": "loras",
        "weights": "femaleasshole-f2-klein-9b.safetensors",
        "default_prompt": None,
        "default_weight": 1.0,
    },        
    {
        "title": "Realistic Nudes",
        "adapter_name": "Realistic Nudes",
        "repo": "loras",
        "weights": "realistic_nudes_klein_v3.safetensors",
        "default_prompt": None,
        "default_weight": 1.0,
    },            
    {
        "title": "Perky Pointy Puffy Breasts",
        "adapter_name": "Perky Pointy Puffy Breasts",
        "repo": "loras",
        "weights": "PerkyPointyPuffy_v1.1_small_pointy_breasts_large_puffy_nipples.safetensors",
        "default_prompt": "Small pointy breasts with large puffy nipples",
        "default_weight": 1.0,
    },       
    {
        "title": "Flat Chested",
        "adapter_name": "Flat Chested",
        "repo": "loras",
        "weights": "Flux2-Klein-9b-FlatChested-v1.safetensors",
        "default_prompt": "flat chested",
        "default_weight": 1.5,
    },           
    {
        "title": "Controllight",
        "adapter_name": "Controllight",
        "repo": "ControlLight/ControlLight",
        "weights": "controllight.safetensors",
        "default_prompt": None,
        "default_weight": 1.0,
    },      
    {
        "title": "RefControl - Depth",
        "adapter_name": "RefConDep",
        "repo": "thedeoxen/refcontrol-FLUX.2-klein-9B-reference-depth-lora",
        "weights": "flux2_klein_9b_refcontrol_depth.safetensors",
        "default_prompt": "refcontrol",
        "default_weight": 1.0,
    },         
    {
        "title": "RefControl - Pose",
        "adapter_name": "RefConPos",
        "repo": "thedeoxen/refcontrol-FLUX.2-klein-9B-reference-pose-lora",
        "weights": "refcontrol_v2_poses.safetensors",
        "default_prompt": "apply pose from image 1 with reference from image 2",
        "default_weight": 1.0,
    },    
]

# LOADED_ADAPTERS is the only piece of LoRA state that's legitimately global:
# it just tracks which adapter names have been loaded onto the shared `pipe`
# at least once, so we don't re-download/re-attach weights on every call.
# It is NOT used to decide what a user can select — that comes from each
# session's own `dynamic_loras_state` (a gr.State dict), so one user's custom
# HuggingFace LoRA never shows up in another user's checklist.
LOADED_ADAPTERS = set()


def get_all_styles(dynamic_loras):
    all_styles = list(LORA_STYLES)
    for dynamic_lora in (dynamic_loras or {}).values():
        all_styles.append(dynamic_lora)
    return all_styles

def get_selectable_styles(dynamic_loras):
    """Return all styles except 'None' for the checkbox selector."""
    return [s for s in get_all_styles(dynamic_loras) if s["adapter_name"] is not None]

def get_style_by_title(title, dynamic_loras):
    for style in get_all_styles(dynamic_loras):
        if style["title"] == title:
            return style
    return None

def get_style_by_adapter_name(adapter_name, dynamic_loras):
    for style in get_all_styles(dynamic_loras):
        if style["adapter_name"] == adapter_name:
            return style
    return None

print(f"Loading FLUX.2 Klein {MODEL_VARIANT} from {_MODEL_REPO}...")
pipe = _PipeClass.from_pretrained(_MODEL_REPO, torch_dtype=torch.bfloat16).to(device)
print(f"Model loaded successfully: FLUX.2 Klein {MODEL_VARIANT}")

def compute_base_dimensions(image):
    """Scale image so longest side fits within 1024, aligned to 16px grid."""
    if image is None:
        return 1024, 1024
    original_width, original_height = image.size
    scale = min(1024 / original_width, 1024 / original_height)
    new_width = (int(original_width * scale) // 16) * 16
    new_height = (int(original_height * scale) // 16) * 16
    return new_width, new_height

# Keep old name as alias so nothing inside infer breaks if called directly
update_dimensions_on_upload = compute_base_dimensions

def fix_orientation(img):
    """Correct EXIF orientation in-place (returns a new image).
    Browsers/OS viewers respect the EXIF orientation tag, but PIL does not
    apply it by default — so an image can *look* correctly oriented
    everywhere you view it while the actual pixel data fed to the model is
    rotated 90/180/270°. This normalises the pixels to match what you see."""
    if img is None:
        return None
    return ImageOps.exif_transpose(img)

def on_base_image_change(img):
    """Update the size/dimension preview when the base image changes."""
    if img is None:
        return "*No base image uploaded yet*"
    try:
        pil_img = img if isinstance(img, Image.Image) else Image.open(img)
        orig_w, orig_h = pil_img.size
        base_w, base_h = compute_base_dimensions(pil_img)
        return f"Input: **{orig_w} × {orig_h}** px → Output (pre-upscale): **{base_w} × {base_h}** px"
    except Exception as e:
        return f"*Could not read dimensions: {e}*"

def on_reference_change(images):
    """Update the reference image count info when the gallery changes."""
    if not images:
        return "📷 No reference images"
    count = len(images)
    return f"📷 {count} reference image{'s' if count != 1 else ''} uploaded"

def process_images(base_image, reference_images):
    """Build the ordered list of PIL images sent to the pipeline:
    [base_image] + reference_images. EXIF orientation is corrected on every
    image, regardless of whether it came in via direct upload or the editor."""
    pil_images = []
    if base_image is not None:
        try:
            img = base_image if isinstance(base_image, Image.Image) else Image.open(base_image)
            pil_images.append(fix_orientation(img).convert("RGB"))
        except Exception as e:
            print(f"Skipping invalid base image: {e}")
    for item in (reference_images or []):
        try:
            path_or_img = item[0] if isinstance(item, (tuple, list)) else item
            if isinstance(path_or_img, Image.Image):
                img = path_or_img
            elif isinstance(path_or_img, str):
                img = Image.open(path_or_img)
            else:
                img = Image.open(path_or_img.name)
            pil_images.append(fix_orientation(img).convert("RGB"))
        except Exception as e:
            print(f"Skipping invalid reference image: {e}")
    return pil_images

def fix_editor_orientation(editor_value):
    """Runs on upload into the Image Editor — corrects EXIF orientation on
    the background layer before the user starts cropping, so what they see
    in the crop box matches the actual pixel data."""
    if not editor_value or editor_value.get("background") is None:
        return editor_value
    bg = editor_value["background"]
    if isinstance(bg, Image.Image):
        editor_value["background"] = fix_orientation(bg)
    return editor_value

def _editor_composite(editor_value):
    if not editor_value or editor_value.get("composite") is None:
        raise gr.Error("Upload and crop an image in the editor first.")
    composite = editor_value["composite"]
    if isinstance(composite, np.ndarray):
        composite = Image.fromarray(composite)
    return composite.convert("RGB")

def send_editor_to_base(editor_value):
    composite = _editor_composite(editor_value)
    gr.Info("Sent to Base Image")
    return composite

def send_editor_to_reference(editor_value, current_gallery):
    composite = _editor_composite(editor_value)
    current = list(current_gallery or [])
    current.append(composite)
    gr.Info("Added to Reference Images")
    return current

# ── LoRA selection → update weight sliders + LoRA prompt display ──────────────

def update_weight_sliders(selected_titles, dynamic_loras):
    """
    Given a list of selected LoRA titles, return updated visibility/label/value
    for each of the MAX_LORA_SLOTS pre-rendered sliders, plus an update for the
    LoRA prompt display box (visible when any selected LoRA has a default prompt).

    The user's own prompt is never touched here — LoRA prompts are shown
    separately and appended at inference time.
    """
    selected_styles = [get_style_by_title(t, dynamic_loras) for t in (selected_titles or []) if get_style_by_title(t, dynamic_loras)]
    slider_updates = []
    for i in range(MAX_LORA_SLOTS):
        if i < len(selected_styles):
            style = selected_styles[i]
            slider_updates.append(gr.update(
                visible=True,
                label=f"{style['title']} — weight",
                value=style.get("default_weight", 1.0),
            ))
        else:
            slider_updates.append(gr.update(visible=False, value=1.0))

    # Collect default prompts from ALL selected LoRAs (not just one)
    lora_prompts = [s.get("default_prompt") for s in selected_styles if s.get("default_prompt")]
    if lora_prompts:
        combined = "\n\n".join(lora_prompts)
        lora_prompt_update = gr.update(value=combined, visible=True)
    else:
        lora_prompt_update = gr.update(value="", visible=False)

    return slider_updates + [lora_prompt_update]

# ── Inference ─────────────────────────────────────────────────────────────────

def _upscale_tiled(model_fn, img_t: torch.Tensor, tile: int = 512, overlap: int = 32) -> torch.Tensor:
    """
    Run the upscaler model patch-by-patch to cap peak VRAM usage.
    Overlapping tiles are averaged so there are no seam artifacts.
    Output is assembled on CPU so only one tile lives on GPU at a time.
    """
    _, c, h, w = img_t.shape

    # Probe output scale with a tiny crop rather than hard-coding it
    with torch.no_grad():
        probe = model_fn(img_t[:, :, :min(4, h), :min(4, w)])
    scale = probe.shape[-1] // min(4, w)
    del probe
    torch.cuda.empty_cache()

    out_h, out_w = h * scale, w * scale
    canvas  = torch.zeros(1, c, out_h, out_w, dtype=torch.float32)
    weights = torch.zeros(1, 1, out_h, out_w, dtype=torch.float32)

    step = max(tile - overlap, 1)
    # Ensure the last tile always reaches the edge
    ys = list(range(0, max(h - tile, 0), step)) + [max(h - tile, 0)]
    xs = list(range(0, max(w - tile, 0), step)) + [max(w - tile, 0)]
    ys = sorted(set(ys))
    xs = sorted(set(xs))

    for y0 in ys:
        for x0 in xs:
            y1 = min(y0 + tile, h)
            x1 = min(x0 + tile, w)
            patch = img_t[:, :, y0:y1, x0:x1]
            with torch.no_grad():
                out = model_fn(patch).cpu().float()
            oy0, ox0 = y0 * scale, x0 * scale
            oy1, ox1 = y1 * scale, x1 * scale
            canvas[:, :, oy0:oy1, ox0:ox1]  += out
            weights[:, :, oy0:oy1, ox0:ox1] += 1.0

    return (canvas / weights.clamp(min=1)).clamp(0, 1)


def apply_realesrgan(image: Image.Image, model_key: str) -> Image.Image:
    """
    Upscale a PIL image using the model selected in UPSCALE_MODELS.
    Weights are downloaded once to /tmp/realesrgan_weights/ and reused.
    """
    cfg = UPSCALE_MODELS[model_key]

    try:
        from spandrel import ImageModelDescriptor, ModelLoader
    except ImportError:
        raise gr.Error("spandrel is not installed. Add 'spandrel' to requirements.txt.")

    import urllib.request

    cache_dir = "/tmp/realesrgan_weights"
    os.makedirs(cache_dir, exist_ok=True)
    cache_path = os.path.join(cache_dir, cfg["file"])

    if not os.path.exists(cache_path):
        print(f"Downloading {cfg['file']}…")
        urllib.request.urlretrieve(cfg["url"], cache_path)

    model = ModelLoader().load_from_file(cache_path)
    if not isinstance(model, ImageModelDescriptor):
        raise gr.Error(f"Loaded model is not a single-image descriptor: {type(model)}")

    sr_model = model.model.to(device).eval()

    img_np = np.array(image).astype(np.float32) / 255.0
    img_t  = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).to(device)

    out_t  = _upscale_tiled(sr_model, img_t, tile=512, overlap=32)

    # Free upscaler weights immediately — FLUX pipeline stays resident
    del sr_model, img_t
    gc.collect()
    torch.cuda.empty_cache()

    out_np = out_t.squeeze(0).permute(1, 2, 0).numpy()
    return Image.fromarray((out_np * 255).astype(np.uint8))


# ── Logging helper — fires a daemon thread so logging never blocks the user ──

def _spawn_log(pil_images, result_image, prompt, seed, steps, guidance_scale,
               width, height, duration, success, error="",
               lora_titles=None, lora_weights=None, upscale_factor="None",
               lora_prompt_text=""):
    
    if os.environ.get("ENABLE_LOGGING", "No").strip().lower() != "yes":
        return
    threading.Thread(
        target=log_inference,
        args=(pil_images, result_image, prompt, seed, steps, guidance_scale,
              width, height, duration, success, error),
        kwargs={"lora_titles": lora_titles, "lora_weights": lora_weights,
                "upscale_factor": upscale_factor, "lora_prompt_text": lora_prompt_text},
        daemon=True,
    ).start()


def infer(
    base_image,
    reference_images,
    prompt,
    lora_prompt_text,
    custom_prompt_text,
    selected_titles,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    upscale_factor,
    dynamic_loras,
    *slider_values,
    progress=gr.Progress(track_tqdm=True)
):
    """CPU-side wrapper — handles pre/post processing and logging.
    The actual GPU work is delegated to _infer_gpu so that the
    _spawn_log daemon thread survives after the @spaces.GPU subprocess exits."""
    gc.collect()
    torch.cuda.empty_cache()

    # Guard against stale MCP schema sending the old output_scale float value
    if not isinstance(upscale_factor, str) or upscale_factor not in UPSCALE_MODELS:
        upscale_factor = "None"

    if base_image is None:
        raise gr.Error("Please upload a base image.")

    pil_images = process_images(base_image, reference_images)
    if not pil_images:
        raise gr.Error("Could not process uploaded images.")

    selected_titles = selected_titles or []

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

    # Collect LoRA metadata for logging before entering GPU
    active_styles = [get_style_by_title(t, dynamic_loras) for t in selected_titles
                     if get_style_by_title(t, dynamic_loras) and get_style_by_title(t, dynamic_loras)["adapter_name"] is not None]
    log_lora_titles = [s["title"] for s in active_styles]
    log_lora_weights = [float(w) for w in slider_values[:len(active_styles)]]
    log_lora_prompt = (lora_prompt_text or "").strip()

    t0 = time.perf_counter()
    try:
        image, seed, width, height = _infer_gpu(
            pil_images, prompt, lora_prompt_text, custom_prompt_text, selected_titles,
            seed, guidance_scale, steps, upscale_factor, dynamic_loras,
            *slider_values, progress=progress
        )
        duration = time.perf_counter() - t0
        # _spawn_log runs on CPU main process so the daemon thread survives
        _spawn_log(pil_images, image, prompt, seed, steps, guidance_scale,
                   width, height, duration, True,
                   lora_titles=log_lora_titles, lora_weights=log_lora_weights,
                   upscale_factor=upscale_factor, lora_prompt_text=log_lora_prompt)
        return image, str(seed)
    except Exception as e:
        duration = time.perf_counter() - t0
        _spawn_log(pil_images, None, prompt, seed, steps, guidance_scale,
                   0, 0, duration, False, str(e),
                   lora_titles=log_lora_titles, lora_weights=log_lora_weights,
                   upscale_factor=upscale_factor, lora_prompt_text=log_lora_prompt)
        raise


@spaces.GPU
def _infer_gpu(
    pil_images,
    prompt,
    lora_prompt_text,
    custom_prompt_text,
    selected_titles,
    seed,
    guidance_scale,
    steps,
    upscale_factor,
    dynamic_loras,
    *slider_values,
    progress=gr.Progress(track_tqdm=True)
):
    """GPU-side inference — all LoRA loading, pipe() calls, and upscaling happen here."""

    # Face-swap validation
    if "Best-Face-Swap" in selected_titles:
        if len(pil_images) < 2:
            raise gr.Error("Face Swap requires 2 images: a Base image and one Reference image (used as the face source).")
        if len(pil_images) > 2:
            gr.Warning("Face Swap uses only the Base image and the first Reference image.")
            pil_images = pil_images[:2]

    # Build active adapter list (skip None)
    active_styles = [get_style_by_title(t, dynamic_loras) for t in selected_titles if get_style_by_title(t, dynamic_loras) and get_style_by_title(t, dynamic_loras)["adapter_name"] is not None]
    weights = list(slider_values[:len(active_styles)])

    if not active_styles:
        print("No LoRA selected — running base model.")
        pipe.disable_lora()
    else:
        for i, style in enumerate(active_styles):
            adapter_name = style["adapter_name"]
            if adapter_name not in LOADED_ADAPTERS:
                print(f"Loading adapter: {style['title']}")
                try:
                    pipe.load_lora_weights(
                        style["repo"],
                        weight_name=style["weights"],
                        adapter_name=adapter_name,
                    )
                    LOADED_ADAPTERS.add(adapter_name)
                except Exception as e:
                    raise gr.Error(f"Failed to load {style['title']}: {e}")

        adapter_names = [s["adapter_name"] for s in active_styles]
        adapter_weights = [float(w) for w in weights]
        print(f"Activating adapters: {list(zip(adapter_names, adapter_weights))}")
        pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)


    # Combine user prompt with LoRA default prompts and any selected custom prompts
    user_prompt = (prompt or "").strip()
    lora_extra = (lora_prompt_text or "").strip()
    custom_extra = (custom_prompt_text or "").strip()
    parts = [p for p in [user_prompt, lora_extra, custom_extra] if p]
    full_prompt = "\n".join(parts)

    # Always generate at base 1024-capped resolution
    width, height = compute_base_dimensions(pil_images[0])
    print(f"Generating at: {width} × {height} px")

    processed_images = [img.resize((width, height), Image.LANCZOS).convert("RGB") for img in pil_images]
    image_input = processed_images if len(processed_images) > 1 else processed_images[0]

    try:
        # Flux2KleinKVPipeline is a distilled model and does not accept
        # guidance_scale — only pass it for the non-KV 9B variant.
        call_kwargs = dict(
            image=image_input,
            prompt=full_prompt,
            width=width,
            height=height,
            num_inference_steps=steps,
            generator=torch.Generator(device=device).manual_seed(seed),
        )
        if MODEL_VARIANT != "9B-KV":
            call_kwargs["guidance_scale"] = guidance_scale
        image = pipe(**call_kwargs).images[0]
    except Exception as e:
        raise gr.Error(f"Inference failed: {e}")

    # ── Post-generation upscaling ─────────────────────────────────────────────
    if upscale_factor and upscale_factor != "None":
        scale_int = UPSCALE_MODELS[upscale_factor]["scale"]
        print(f"Upscaling {scale_int}× with {upscale_factor}…")
        # Flush FLUX pipeline residuals before loading upscaler weights
        gc.collect()
        torch.cuda.synchronize()
        torch.cuda.empty_cache()
        try:
            image = apply_realesrgan(image, upscale_factor)
            print(f"Upscaled to: {image.width} × {image.height} px")
        except Exception as e:
            gr.Warning(f"Upscaling failed, returning 1024px result: {e}")
    # ─────────────────────────────────────────────────────────────────────────

    gc.collect()
    torch.cuda.empty_cache()
    return image, seed, width, height

# ── Dynamic LoRA loader ───────────────────────────────────────────────────────

def add_custom_lora(repo_id, weight_name, adapter_name, dynamic_loras_state):
    """
    Adds a custom LoRA to *this session's* state only — other users never see
    it in their selector. The adapter_name attached to the shared pipe is
    still made globally unique (via a short random suffix) since the pipe
    object itself is one shared resource across all sessions.
    """
    dynamic_loras = dict(dynamic_loras_state or {})
    if not repo_id or not repo_id.strip():
        return "Please enter a valid HuggingFace repo ID.", gr.update(), dynamic_loras
    repo_id = repo_id.strip()
    requested_name = adapter_name.strip() if adapter_name and adapter_name.strip() else None
    try:
        from huggingface_hub import model_info
        info = model_info(repo_id)
        actual_weight = weight_name.strip() if weight_name and weight_name.strip() else None
        if not actual_weight:
            for name in ["pytorch_lora_weights.safetensors", "lora.safetensors", "adapter_model.safetensors"]:
                if any(f.filename == name for f in info.siblings):
                    actual_weight = name
                    break
        if not actual_weight:
            available = [f.filename for f in info.siblings if f.filename.endswith(('.safetensors', '.bin'))]
            return f"No weight found. Available: {', '.join(available) or 'None'}", gr.update(), dynamic_loras

        base_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in requested_name) if requested_name else "custom"

        static_names = {s["adapter_name"] for s in LORA_STYLES if s["adapter_name"]}
        final_adapter_name = f"{base_name}_{uuid.uuid4().hex[:6]}"
        while final_adapter_name in static_names or final_adapter_name in LOADED_ADAPTERS:
            final_adapter_name = f"{base_name}_{uuid.uuid4().hex[:6]}"

        custom_style = {
            "image": "https://huggingface.co/spaces/prithivMLmods/FLUX.2-Klein-LoRA-Studio/resolve/main/examples/image.webp",
            "title": f"Custom: {base_name}",
            "adapter_name": final_adapter_name,
            "repo": repo_id,
            "weights": actual_weight,
            "default_prompt": None,
            "default_weight": 1.0,
        }
        dynamic_loras[final_adapter_name] = custom_style
        new_choices = [s["title"] for s in get_selectable_styles(dynamic_loras)]
        return f"✅ Added: {base_name} from {repo_id}", gr.update(choices=new_choices), dynamic_loras
    except Exception as e:
        return f"❌ Failed: {str(e)}", gr.update(), dynamic_loras

# ── Custom Prompt manager ─────────────────────────────────────────────────────

def add_custom_prompt(name, text, prompts_state, counter_state):
    """
    All state is passed in and returned — never touches module-level globals.
    Each user session gets its own isolated copy via gr.State.
    """
    prompts = dict(prompts_state)   # never mutate the state object directly
    counter = int(counter_state)
    text = text.strip() if text else ""
    name = name.strip() if name else ""
    if not text:
        return "Please enter some prompt text.", prompts, counter, gr.update(), gr.update(), gr.update()
    if not name:
        counter += 1
        name = f"Prompt {counter}"
    if name in prompts:
        return f"⚠️ '{name}' already exists — choose a different name.", prompts, counter, gr.update(), gr.update(), gr.update()
    prompts[name] = text
    choices = list(prompts.keys())
    return (
        f"✅ Saved: '{name}'",
        prompts,                                # updated prompts state
        counter,                                # updated counter state
        gr.update(choices=choices),             # custom_prompt_selector
        gr.update(choices=choices),             # delete_prompt_name dropdown
        gr.update(value="", interactive=True),  # clear the name field
    )

def delete_custom_prompt(name, currently_selected, prompts_state):
    prompts = dict(prompts_state)
    if name and name in prompts:
        del prompts[name]
        msg = f"🗑️ Deleted: '{name}'"
    else:
        msg = "Nothing to delete."
    choices = list(prompts.keys())
    # Drop deleted name from active selection if it was ticked
    new_selection = [n for n in (currently_selected or []) if n in prompts]
    return (
        msg,
        prompts,                                          # updated prompts state
        gr.update(choices=choices, value=new_selection),  # custom_prompt_selector
        gr.update(choices=choices, value=None),           # delete_prompt_name
    )

def update_custom_prompt_display(selected_names, prompts_state):
    if not selected_names:
        return gr.update(value="", visible=False)
    texts = [prompts_state[n] for n in selected_names if n in prompts_state]
    if texts:
        return gr.update(value="\n\n".join(texts), visible=True)
    return gr.update(value="", visible=False)


# ── UI ────────────────────────────────────────────────────────────────────────

css = """
#col-container { margin: 0 auto; max-width: 980px; }
#main-title h1 { font-size: 2.4em !important; }
#reference_gallery .grid-wrap { min-height: 120px }
.lora-weight-row { background: var(--block-background-fill); border-radius: 8px; padding: 4px 12px; margin-bottom: 4px; }
"""

with gr.Blocks(css=css, theme=orange_red_theme) as demo:
    # Per-session state — each user gets their own isolated copy, never shared across sessions
    custom_prompts_state = gr.State({})
    custom_prompt_counter_state = gr.State(0)
    dynamic_loras_state = gr.State({})  # this session's custom HF LoRAs — never shared with other users

    with gr.Column(elem_id="col-container"):
        _logging_on = os.environ.get("ENABLE_LOGGING", "No").strip().lower() == "yes"
        _logging_badge = "🟢 On" if _logging_on else "🔴 Off"

        gr.Markdown("# **FLUX.2-Klein-LoRA-Studio**", elem_id="main-title")
        gr.Markdown(
            f"Apply one or more [LoRA](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.2-klein-9B) "
            f"adapters to your images using "
            f"[FLUX.2-Klein-{MODEL_VARIANT}]({_MODEL_REPO}).\n\n"
            f"**Model:** `{MODEL_VARIANT}` · **Logging:** {_logging_badge}"
        )
        gr.Markdown(
            "*Results may vary when combining multiple LoRAs. "
            "Uploaded images and metadata may be temporarily stored to help optimise the space — "
            "stored data is automatically deleted. "
            "Always comply with HF and model Terms of Service.*"
        )

        with gr.Tabs() as main_tabs:
            # ── Generate tab ────────────────────────────────────────────────────
            with gr.Tab("🎨 Generate", id="tab_generate"):
                with gr.Row(equal_height=False):
                    # ── Left column ──────────────────────────────────────────────
                    with gr.Column(scale=1):
                        base_image = gr.Image(
                            label="Base Image",
                            type="pil",
                            sources=["upload", "clipboard"],
                            height=260,
                            elem_id="base_image",
                        )
                        size_info = gr.Markdown("*No image uploaded yet*")

                        reference_images = gr.Gallery(
                            label="Reference Image(s) — optional",
                            type="filepath",
                            columns=2, rows=1, height=140,
                            allow_preview=True,
                            elem_id="reference_gallery",
                        )
                        reference_info = gr.Markdown("📷 No reference images")
                        gr.Markdown(
                            "*For Face Swap, the first Reference image is used as the face source.*"
                        )

                        prompt = gr.Text(
                            label="Prompt",
                            max_lines=3,
                            placeholder="Describe the edit, or leave blank for style-only LoRAs",
                        )

                        # Read-only display for LoRA default prompts — shown when any
                        # selected LoRA has a default prompt. Combined with user prompt
                        # at inference time; never overwrites what the user typed.
                        lora_prompt_display = gr.Textbox(
                            label="LoRA Default Prompts (auto-appended)",
                            info="These prompts come from your selected LoRAs and are appended to your prompt above during generation.",
                            interactive=False,
                            visible=False,
                            lines=3,
                        )

                        # Read-only display for active custom prompts — shown when any
                        # saved custom prompt is ticked in the Custom Prompts accordion.
                        custom_prompt_display = gr.Textbox(
                            label="Custom Prompts (auto-appended)",
                            info="These are your selected saved prompts, appended alongside LoRA prompts during generation.",
                            interactive=False,
                            visible=False,
                            lines=3,
                        )

                        run_button = gr.Button("▶ Generate", variant="primary", size="lg")

                        with gr.Accordion("⚙️ Advanced Settings", open=False):
                            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                            guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=1.0, visible=MODEL_VARIANT != "9B-KV")
                            steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=4, step=1)
                            upscale_factor = gr.Dropdown(
                                label="Upscale model",
                                info="Applied after generation. UltraSharp / Remacri are sharper than RealESRGAN for clean photo inputs.",
                                choices=list(UPSCALE_MODELS.keys()),
                                value="None",
                            )

                    # ── Right column ─────────────────────────────────────────────
                    with gr.Column(scale=1):
                        output_image = gr.Image(label="Output", interactive=False, format="png", height=320)
                        used_seed = gr.Textbox(
                            label="🌱 Seed used",
                            info="Re-enter this seed in Advanced Settings to reproduce the result exactly.",
                            interactive=False,
                            visible=True,
                            max_lines=1,
                        )

                # ── LoRA selector ─────────────────────────────────────────────────
                gr.Markdown("### 🎨 Select LoRA(s)")
                lora_selector = gr.CheckboxGroup(
                    choices=[s["title"] for s in get_selectable_styles({})],
                    value=[],
                    label="Active LoRAs — tick one or more",
                )

                # ── Weight sliders (pre-rendered, shown/hidden dynamically) ───────
                gr.Markdown("#### Weights for selected LoRAs")
                weight_sliders = []
                with gr.Group():
                    for i in range(MAX_LORA_SLOTS):
                        with gr.Row(elem_classes="lora-weight-row"):
                            s = gr.Slider(
                                minimum=0.0, maximum=2.0, step=0.05, value=1.0,
                                label=f"LoRA slot {i+1}",
                                visible=False,
                                interactive=True,
                            )
                            weight_sliders.append(s)

                # ── Dynamic HuggingFace loader ─────────────────────────────────────
                with gr.Accordion("➕ Load Custom LoRA from HuggingFace", open=False):
                    gr.Markdown("Add any FLUX.2-Klein-compatible LoRA from [HuggingFace Hub](https://huggingface.co/models?other=base_model:adapter:black-forest-labs/FLUX.2-klein-9B). "
                                "*Added LoRAs are only visible in your own session — other visitors won't see them.*")
                    with gr.Row():
                        lora_repo_id = gr.Textbox(label="Repo ID", placeholder="username/repo-name")
                    with gr.Row():
                        lora_weight_name = gr.Textbox(label="Weight filename (optional)", placeholder="pytorch_lora_weights.safetensors")
                        lora_adapter_name = gr.Textbox(label="Adapter name (optional)", placeholder="my-lora")
                    with gr.Row():
                        add_lora_btn = gr.Button("Add LoRA", variant="primary")
                        lora_status = gr.Textbox(label="Status", interactive=False)

                # ── Custom Prompt library ──────────────────────────────────────────
                with gr.Accordion("📝 Custom Prompts", open=False):
                    gr.Markdown(
                        "Save reusable prompt snippets for this session. "
                        "Tick any saved prompt below and it will be appended to your generation alongside any LoRA prompts. "
                        "*Prompts are stored in memory — they reset when the Space restarts or the worker recycles.*"
                    )

                    custom_prompt_selector = gr.CheckboxGroup(
                        choices=[],
                        value=[],
                        label="Saved prompts — tick to append to generation",
                    )

                    with gr.Row():
                        prompt_name_input = gr.Textbox(
                            label="Name",
                            placeholder="e.g. Skin detail enhancer",
                            scale=1,
                        )
                    with gr.Row():
                        prompt_text_input = gr.Textbox(
                            label="Prompt text",
                            placeholder="Enter the prompt snippet you want to save and reuse…",
                            lines=4,
                        )
                    with gr.Row():
                        add_prompt_btn = gr.Button("💾 Save Prompt", variant="primary")
                        prompt_status = gr.Textbox(label="Status", interactive=False, scale=2)
                    with gr.Row():
                        delete_prompt_name = gr.Dropdown(
                            label="Delete a saved prompt",
                            choices=[],
                            value=None,
                            interactive=True,
                            scale=2,
                        )
                        delete_prompt_btn = gr.Button("🗑️ Delete", variant="secondary", scale=1)

                gr.Markdown("*Experimental space — results may vary when combining multiple LoRAs.*")

            # ── Crop / Fix Image tab ─────────────────────────────────────────────
            with gr.Tab("✂️ Crop / Fix Image", id="tab_editor"):
                gr.Markdown(
                    "Upload an image to crop it, then send the result to the Base Image or add it as a "
                    "Reference image on the Generate tab. EXIF orientation (the most common cause of "
                    "images coming out sideways) is corrected automatically as soon as you upload here."
                )
                editor = gr.ImageEditor(
                    label="Editor",
                    type="pil",
                    transforms=["crop"],
                    brush=False,
                    eraser=False,
                    layers=False,
                    sources=["upload", "clipboard"],
                    height=420,
                )
                with gr.Row():
                    send_to_base_btn = gr.Button("→ Send to Base Image", variant="primary")
                    send_to_ref_btn = gr.Button("→ Add to Reference Images")

    # ── Event wiring ──────────────────────────────────────────────────────────

    base_image.change(
        fn=on_base_image_change,
        inputs=[base_image],
        outputs=[size_info],
    )

    reference_images.change(
        fn=on_reference_change,
        inputs=[reference_images],
        outputs=[reference_info],
    )

    lora_selector.change(
        fn=update_weight_sliders,
        inputs=[lora_selector, dynamic_loras_state],
        outputs=weight_sliders + [lora_prompt_display],
    )

    add_lora_btn.click(
        fn=add_custom_lora,
        inputs=[lora_repo_id, lora_weight_name, lora_adapter_name, dynamic_loras_state],
        outputs=[lora_status, lora_selector, dynamic_loras_state],
    )

    add_prompt_btn.click(
        fn=add_custom_prompt,
        inputs=[prompt_name_input, prompt_text_input, custom_prompts_state, custom_prompt_counter_state],
        outputs=[prompt_status, custom_prompts_state, custom_prompt_counter_state, custom_prompt_selector, delete_prompt_name, prompt_name_input],
    )

    delete_prompt_btn.click(
        fn=delete_custom_prompt,
        inputs=[delete_prompt_name, custom_prompt_selector, custom_prompts_state],
        outputs=[prompt_status, custom_prompts_state, custom_prompt_selector, delete_prompt_name],
    )

    custom_prompt_selector.change(
        fn=update_custom_prompt_display,
        inputs=[custom_prompt_selector, custom_prompts_state],
        outputs=[custom_prompt_display],
    )

    run_button.click(
        fn=infer,
        inputs=[base_image, reference_images, prompt, lora_prompt_display, custom_prompt_display, lora_selector, seed, randomize_seed, guidance_scale, steps, upscale_factor, dynamic_loras_state] + weight_sliders,
        outputs=[output_image, used_seed],
    )

    # ── Image Editor tab wiring ─────────────────────────────────────────────
    editor.upload(
        fn=fix_editor_orientation,
        inputs=[editor],
        outputs=[editor],
    )

    send_to_base_btn.click(
        fn=send_editor_to_base,
        inputs=[editor],
        outputs=[base_image],
    ).then(
        fn=on_base_image_change, inputs=[base_image], outputs=[size_info],
    ).then(
        fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs],
    )

    send_to_ref_btn.click(
        fn=send_editor_to_reference,
        inputs=[editor, reference_images],
        outputs=[reference_images],
    ).then(
        fn=on_reference_change, inputs=[reference_images], outputs=[reference_info],
    ).then(
        fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs],
    )

if __name__ == "__main__":
    demo.queue().launch(css=css, theme=orange_red_theme, mcp_server=True, ssr_mode=False, show_error=True)