app.py

import os
import gc
import gradio as gr
import numpy as np
import spaces
import torch
import random
from PIL import Image
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": 1.0,
    },
    {
        "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,
    },      
]

LOADED_ADAPTERS = set()
DYNAMIC_LORAS = {}
DYNAMIC_LORA_COUNTER = 0


def get_all_styles():
    all_styles = list(LORA_STYLES)
    for dynamic_lora in DYNAMIC_LORAS.values():
        all_styles.append(dynamic_lora)
    return all_styles

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

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

def get_style_by_adapter_name(adapter_name):
    for style in get_all_styles():
        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 on_gallery_change(images):
    """Update image count info and size display when gallery changes."""
    count_info = get_image_count_info(images)
    if not images:
        return count_info, "*No image uploaded yet*"
    try:
        first = images[0]
        path = first[0] if isinstance(first, (tuple, list)) else first
        img = Image.open(path if isinstance(path, str) else path.name)
        orig_w, orig_h = img.size
        base_w, base_h = compute_base_dimensions(img)
        size_text = f"Input: **{orig_w} × {orig_h}** px → Output (pre-upscale): **{base_w} × {base_h}** px"
        return count_info, size_text
    except Exception as e:
        return count_info, f"*Could not read dimensions: {e}*"

def process_gallery_images(images):
    if not images:
        return []
    pil_images = []
    for item in images:
        try:
            path_or_img = item[0] if isinstance(item, (tuple, list)) else item
            if isinstance(path_or_img, str):
                pil_images.append(Image.open(path_or_img).convert("RGB"))
            elif isinstance(path_or_img, Image.Image):
                pil_images.append(path_or_img.convert("RGB"))
            else:
                pil_images.append(Image.open(path_or_img.name).convert("RGB"))
        except Exception as e:
            print(f"Skipping invalid image: {e}")
    return pil_images

def get_image_count_info(images):
    if not images:
        return "📷 No images uploaded"
    count = len(images)
    if count == 1:
        return "📷 1 image uploaded (Picture 1 - Base)"
    elif count == 2:
        return "📷 2 images uploaded (Picture 1 - Base, Picture 2 - Face Source)"
    return f"📷 {count} images uploaded"

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

def update_weight_sliders(selected_titles):
    """
    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) for t in (selected_titles or []) if get_style_by_title(t)]
    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(
    input_images,
    prompt,
    lora_prompt_text,
    custom_prompt_text,
    selected_titles,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    upscale_factor,
    *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 not input_images:
        raise gr.Error("Please upload at least one image.")

    pil_images = process_gallery_images(input_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) for t in selected_titles
                     if get_style_by_title(t) and get_style_by_title(t)["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,
            *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,
    *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: Picture 1 (base) and Picture 2 (face source).")
        if len(pil_images) > 2:
            gr.Warning("Face Swap uses only the first 2 images.")
            pil_images = pil_images[:2]

    # Build active adapter list (skip None)
    active_styles = [get_style_by_title(t) for t in selected_titles if get_style_by_title(t) and get_style_by_title(t)["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):
    global DYNAMIC_LORA_COUNTER
    if not repo_id or not repo_id.strip():
        return "Please enter a valid HuggingFace repo ID.", gr.update(), gr.update()
    repo_id = repo_id.strip()
    adapter_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(), gr.update()
        if not adapter_name:
            DYNAMIC_LORA_COUNTER += 1
            adapter_name = f"custom_{DYNAMIC_LORA_COUNTER}"
        else:
            adapter_name = "".join(c if c.isalnum() or c in "-_" else "_" for c in adapter_name)
        custom_style = {
            "image": "https://huggingface.co/spaces/prithivMLmods/FLUX.2-Klein-LoRA-Studio/resolve/main/examples/image.webp",
            "title": f"Custom: {adapter_name}",
            "adapter_name": adapter_name,
            "repo": repo_id,
            "weights": actual_weight,
            "default_prompt": None,
            "default_weight": 1.0,
        }
        DYNAMIC_LORAS[adapter_name] = custom_style
        new_choices = [s["title"] for s in get_selectable_styles()]
        return f"✅ Added: {adapter_name} from {repo_id}", gr.update(choices=new_choices), gr.update(choices=new_choices)
    except Exception as e:
        return f"❌ Failed: {str(e)}", gr.update(), gr.update()

# ── 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; }
#input_gallery .grid-wrap { min-height: 200px }
.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)

    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.Row(equal_height=False):
            # ── Left column ──────────────────────────────────────────────────
            with gr.Column(scale=1):
                input_images = gr.Gallery(
                    label="Upload Image(s)",
                    type="filepath",
                    columns=2, rows=1, height=260,
                    allow_preview=True,
                    elem_id="input_gallery",
                )
                image_info = gr.Markdown("📷 No images uploaded")
                size_info = gr.Markdown("*No image uploaded yet*")

                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)")
            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.*")

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

    input_images.change(
        fn=on_gallery_change,
        inputs=[input_images],
        outputs=[image_info, size_info],
    )

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

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

    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=[input_images, prompt, lora_prompt_display, custom_prompt_display, lora_selector, seed, randomize_seed, guidance_scale, steps, upscale_factor] + weight_sliders,
        outputs=[output_image, used_seed],
    )

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