Upload 2 files

app.py

@@ -6,19 +6,33 @@ import random
 import uuid
 import zipfile
 import threading
-from typing import Iterable
 
 import gradio as gr
-import numpy as np
 import spaces
 import torch
 from PIL import Image
 
+# ── Local modules — single source of truth for each concern ─────────────────
+from config import (
+    MODEL_VARIANT,
+    MODEL_REPO,
+    MAX_SEED,
+    MAX_LORA_SLOTS,
+    ENABLE_LOGGING,
+)
+from ui_theme import orange_red_theme
+from upscale import UPSCALE_MODELS, apply_realesrgan
+from lora_registry import (
+    LORA_STYLES,
+    LOADED_ADAPTERS,
+    get_selectable_styles,
+    get_style_by_title,
+    update_weight_sliders,
+    add_custom_lora,
+)
 from logging_utils import log_inference
-
 from image_utils import (
     fix_orientation,
-    compute_base_dimensions,
     compute_canvas_dimensions,
     fit_to_canvas,
     on_base_image_change,
@@ -35,7 +49,6 @@ from image_utils import (
     push_pil_to_base,
     push_pil_to_reference,
 )
-
 from control_tools import (
     generate_depthmap,
     detect_pose,
@@ -51,244 +64,23 @@ from control_tools import (
     OPENPOSE_KEYPOINT_NAMES,
 )
 
-# requirements.txt: spandrel, pillow-heif
-
-MODEL_VARIANT = os.environ.get("MODEL_VARIANT", "9B")
+# ── Model load ──────────────────────────────────────────────────────────────
+# Pipeline class depends on MODEL_VARIANT and is the only thing here that
+# can't live in config.py (config must stay torch/diffusers-free).
 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"
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ── Theme ────────────────────────────────────────────────────────────────────
-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.gray, secondary_hue=colors.orange_red,
-                 neutral_hue=colors.slate, text_size=sizes.text_lg,
-                 font=(fonts.GoogleFont("Outfit"), "Arial", "sans-serif"),
-                 font_mono=(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 models, FACE_SWAP_PROMPT, LORA_STYLES — UNCHANGED from previous step
-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_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": "InstaPic",
-        "adapter_name": "InstaPic V3",
-        "repo": "loras",
-        "weights": "InstaPic V3.safetensors",
-        "default_prompt": "instapic",
-        "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 = set()
-
-
-def get_all_styles(dynamic_loras):
-    return list(LORA_STYLES) + list((dynamic_loras or {}).values())
-
-def get_selectable_styles(dynamic_loras):
-    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 s in get_all_styles(dynamic_loras):
-        if s["title"] == title:
-            return s
-    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"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}")
 
 
 # ── UI helper callbacks ──────────────────────────────────────────────────────
 
-def update_weight_sliders(selected_titles, dynamic_loras):
-    selected = [get_style_by_title(t, dynamic_loras) for t in (selected_titles or [])
-                if get_style_by_title(t, dynamic_loras)]
-    updates = []
-    for i in range(MAX_LORA_SLOTS):
-        if i < len(selected):
-            s = selected[i]
-            updates.append(gr.update(visible=True,
-                                     label=f"{s['title']} — weight",
-                                     value=s.get("default_weight", 1.0)))
-        else:
-            updates.append(gr.update(visible=False, value=1.0))
-    prompts = [s.get("default_prompt") for s in selected if s.get("default_prompt")]
-    if prompts:
-        return updates + [gr.update(value="\n\n".join(prompts), visible=True)]
-    return updates + [gr.update(value="", visible=False)]
-
-
 def on_canvas_mode_change(mode):
     """Custom W/H sliders only relevant when mode == Custom."""
     is_custom = (mode == "Custom")
@@ -317,65 +109,16 @@ def on_gallery_select(evt: gr.SelectData, gallery_value):
     return item[0] if isinstance(item, (list, tuple)) else item
 
 
-# ── Upscaler (tiled) — unchanged ─────────────────────────────────────────────
-
-def _upscale_tiled(model_fn, img_t, tile=512, overlap=32):
-    _, c, h, w = img_t.shape
-    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)
-    ys = sorted(set(list(range(0, max(h - tile, 0), step)) + [max(h - tile, 0)]))
-    xs = sorted(set(list(range(0, max(w - tile, 0), step)) + [max(w - tile, 0)]))
-    for y0 in ys:
-        for x0 in xs:
-            y1 = min(y0 + tile, h); x1 = min(x0 + tile, w)
-            with torch.no_grad():
-                out = model_fn(img_t[:, :, y0:y1, x0:x1]).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, model_key):
-    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)
-    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 ──────────────────────────────────────────────────────────────────
 
 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":
+    """Fire-and-forget logger. ENABLE_LOGGING is read once at import time in
+    config.py — there's no per-call env lookup. Used by single, batch and
+    bulk inference paths."""
+    if not ENABLE_LOGGING:
         return
     threading.Thread(
         target=log_inference,
@@ -449,7 +192,7 @@ def _infer_gpu(
     if upscale_factor and upscale_factor != "None":
         gc.collect(); torch.cuda.synchronize(); torch.cuda.empty_cache()
         try:
-            image = apply_realesrgan(image, upscale_factor)
+            image = apply_realesrgan(image, upscale_factor, device)
         except Exception as e:
             gr.Warning(f"Upscaling failed, returning {width}×{height} result: {e}")
 
@@ -502,8 +245,6 @@ def infer(
     selected_titles = selected_titles or []
     batch_count = max(1, int(batch_count))
 
-    # Pre-plan per-iteration seed and weight overrides — keeps the loop simple
-    # and lets us put the LoRA-sweep values into PNG metadata cleanly.
     base_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
     seeds, weight_overrides = [], []
     for i in range(batch_count):
@@ -582,7 +323,12 @@ def bulk_infer(
     into the bulk output gallery as soon as they complete. Outputs and a CSV
     manifest are written to /tmp/bulk_<sid>/ with stable filenames, so a
     ZeroGPU quota wall mid-run still leaves earlier outputs grabbable from
-    the gallery and from disk."""
+    the gallery and from disk.
+
+    Logging: each image is logged individually via _spawn_log on success or
+    failure — same code path as single/batch, so bulk shows up in the dataset
+    with the same schema, just with extra bulk_index / bulk_total fields
+    embedded in the PNG metadata."""
     if not input_files:
         raise gr.Error("Upload at least one image first.")
 
@@ -619,8 +365,6 @@ def bulk_infer(
                 *slider_values, progress=progress,
             )
 
-            # Stable, predictable filename inside the session work_dir so the
-            # user can also find outputs on disk if the UI drops.
             stem = os.path.splitext(fname)[0]
             out_path = os.path.join(work_dir, f"{i:03d}_{stem}.png")
             meta = _meta_for(prompt, used_seed, steps, guidance_scale, w, h, upscale_factor,
@@ -642,16 +386,19 @@ def bulk_infer(
                        upscale_factor=upscale_factor, lora_prompt_text=lora_prompt_text or "")
 
             status = f"✅ {succeeded}/{total} done ({failed} failed)"
-            yield results, status, gr.update()   # zip not ready yet
+            yield results, status, gr.update()
         except Exception as e:
             failed += 1
             duration = time.perf_counter() - t0
             with open(manifest_path, "a", newline="") as f:
                 csv.writer(f).writerow([i, fname, "", "", "", "", False, str(e)[:300], f"{duration:.2f}"])
+            _spawn_log([], None, prompt, 0, steps, guidance_scale,
+                       0, 0, duration, False, str(e),
+                       lora_titles=log_titles, lora_weights=log_weights,
+                       upscale_factor=upscale_factor, lora_prompt_text=lora_prompt_text or "")
             yield results, f"⚠️ Image {i+1} failed: {e} | {succeeded} ok, {failed} failed", gr.update()
             continue
 
-    # Final pass: bundle a zip. ZIP_STORED because PNGs are already compressed.
     zip_path = os.path.join(work_dir, "outputs.zip")
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_STORED) as zf:
         for p in results:
@@ -661,44 +408,7 @@ def bulk_infer(
     yield results, f"🎉 Done — {succeeded}/{total} succeeded, {failed} failed", zip_path
 
 
-# ── Dynamic LoRA loader (unchanged from previous step) ───────────────────────
-
-def add_custom_lora(repo_id, weight_name, adapter_name, dynamic_loras_state):
-    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 = f"{base_name}_{uuid.uuid4().hex[:6]}"
-        while final in static_names or final in LOADED_ADAPTERS:
-            final = f"{base_name}_{uuid.uuid4().hex[:6]}"
-        dynamic_loras[final] = {
-            "image": "https://huggingface.co/spaces/prithivMLmods/FLUX.2-Klein-LoRA-Studio/resolve/main/examples/image.webp",
-            "title": f"Custom: {base_name}", "adapter_name": final,
-            "repo": repo_id, "weights": actual_weight,
-            "default_prompt": None, "default_weight": 1.0,
-        }
-        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 (unchanged) ────────────────────────────────────────
+# ── Custom prompt manager (session-local) ────────────────────────────────────
 
 def add_custom_prompt(name, text, prompts_state, counter_state):
     prompts = dict(prompts_state); counter = int(counter_state)
@@ -716,6 +426,7 @@ def add_custom_prompt(name, text, prompts_state, counter_state):
             gr.update(choices=choices), gr.update(choices=choices),
             gr.update(value="", interactive=True))
 
+
 def delete_custom_prompt(name, currently_selected, prompts_state):
     prompts = dict(prompts_state)
     msg = f"🗑️ Deleted: '{name}'" if name and name in prompts else "Nothing to delete."
@@ -727,6 +438,7 @@ def delete_custom_prompt(name, currently_selected, prompts_state):
             gr.update(choices=choices, value=new_sel),
             gr.update(choices=choices, value=None))
 
+
 def update_custom_prompt_display(selected_names, prompts_state):
     if not selected_names:
         return gr.update(value="", visible=False)
@@ -750,16 +462,15 @@ with gr.Blocks() as demo:
     custom_prompts_state = gr.State({})
     custom_prompt_counter_state = gr.State(0)
     dynamic_loras_state = gr.State({})
-    selected_output_state = gr.State(None)  # currently-selected gallery item path
+    selected_output_state = gr.State(None)
 
     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"
+        _logging_badge = "🟢 On" if ENABLE_LOGGING 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 using [FLUX.2-Klein-{MODEL_VARIANT}]({_MODEL_REPO}). "
+            f"adapters using [FLUX.2-Klein-{MODEL_VARIANT}]({MODEL_REPO}). "
             f"**Model:** `{MODEL_VARIANT}` · **Logging:** {_logging_badge}"
         )
 
@@ -835,11 +546,7 @@ with gr.Blocks() as demo:
                             sweep_max = gr.Slider(label="Sweep max weight", minimum=0.0, maximum=2.0,
                                                   step=0.05, value=1.4, visible=False)
 
-                    # ── Right column ─────────────────────────────────────────
                     with gr.Column(scale=1):
-                        # Gallery so batch runs stream in as they finish; type=
-                        # filepath so PNG metadata round-trips to Send→* buttons
-                        # and to the user's downloads.
                         output_gallery = gr.Gallery(
                             label="Output", type="filepath", columns=2, rows=2,
                             height=420, allow_preview=True, preview=True,
@@ -857,7 +564,6 @@ with gr.Blocks() as demo:
                             "Civitai / ComfyUI readable).*"
                         )
 
-                # LoRA selector + weight sliders
                 gr.Markdown("### 🎨 Select LoRA(s)")
                 lora_selector = gr.CheckboxGroup(
                     choices=[s["title"] for s in get_selectable_styles({})],
@@ -967,10 +673,8 @@ with gr.Blocks() as demo:
                     "Reference image."
                 )
 
-                # Per-tab state — kept here, not in module globals, so each
-                # session edits its own pose without crosstalk.
-                pose_source_state = gr.State(None)        # PIL of last source image
-                pose_keypoints_state = gr.State([])       # list[list[dict]]
+                pose_source_state = gr.State(None)
+                pose_keypoints_state = gr.State([])
 
                 with gr.Row():
                     with gr.Column(scale=1):
@@ -1000,8 +704,6 @@ with gr.Blocks() as demo:
 
                 with gr.Row():
                     with gr.Column(scale=1):
-                        # interactive=False so users can't accidentally upload a new image
-                        # into this slot. .select still fires for click coordinates.
                         pose_overlay = gr.Image(
                             label="Editor — click to place active joint",
                             type="pil", interactive=False, height=420, format="png",
@@ -1031,21 +733,21 @@ with gr.Blocks() as demo:
                     send_pose_ref_btn  = gr.Button("→ Send pose to Reference",
                                                    variant="primary")
                     send_pose_base_btn = gr.Button("→ Send pose to Base")
-                  
+
     # ── Event wiring ─────────────────────────────────────────────────────────
 
-    # HEIC preview fix on the main tab
     base_image.upload(fn=reencode_upload, inputs=[base_image], outputs=[base_image])
-
     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])
 
+    # update_weight_sliders is the one imported from lora_registry now.
     lora_selector.change(
         fn=update_weight_sliders,
         inputs=[lora_selector, dynamic_loras_state],
         outputs=weight_sliders + [lora_prompt_display],
     )
 
+    # add_custom_lora is also imported from lora_registry.
     add_lora_btn.click(
         fn=add_custom_lora,
         inputs=[lora_repo_id, lora_weight_name, lora_adapter_name, dynamic_loras_state],
@@ -1069,7 +771,6 @@ with gr.Blocks() as demo:
         outputs=[custom_prompt_display],
     )
 
-    # Canvas / fit / batch UI toggles
     canvas_mode.change(fn=on_canvas_mode_change, inputs=[canvas_mode],
                        outputs=[custom_width, custom_height])
     canvas_fit_mode.change(fn=on_fit_mode_change, inputs=[canvas_fit_mode],
@@ -1077,12 +778,9 @@ with gr.Blocks() as demo:
     batch_vary.change(fn=on_batch_vary_change, inputs=[batch_vary],
                       outputs=[sweep_min, sweep_max])
 
-    # Track which gallery item the user clicked, so Send→Base/Ref can use it
     output_gallery.select(fn=on_gallery_select, inputs=[output_gallery],
                           outputs=[selected_output_state])
 
-    # The Generate-tab generator. .click() returns an event we keep so the
-    # bulk Stop button can cancel mid-stream too if desired.
     run_event = run_button.click(
         fn=infer,
         inputs=[base_image, reference_images, prompt, lora_prompt_display, custom_prompt_display,
@@ -1094,8 +792,6 @@ with gr.Blocks() as demo:
     )
 
     # ── Editor tab wiring ────────────────────────────────────────────────────
-    # NOTE: do NOT add editor.upload(outputs=[editor]) — remounts the editor.
-
     heic_uploader.upload(fn=load_heic_to_editor, inputs=[heic_uploader], outputs=[editor])
 
     send_to_base_btn.click(fn=send_editor_to_base, inputs=[editor], outputs=[base_image]) \
@@ -1107,7 +803,6 @@ with gr.Blocks() as demo:
         .then(fn=on_reference_change, inputs=[reference_images], outputs=[reference_info]) \
         .then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs])
 
-    # Send-output buttons use the selected gallery item (or latest fallback)
     send_out_to_base_btn.click(
         fn=send_output_to_base,
         inputs=[selected_output_state, output_gallery],
@@ -1121,8 +816,6 @@ with gr.Blocks() as demo:
     ).then(fn=on_reference_change, inputs=[reference_images], outputs=[reference_info])
 
     # ── Bulk tab wiring ──────────────────────────────────────────────────────
-    # Reuses Generate-tab settings as inputs verbatim — single source of truth,
-    # no two-way state sync to keep in step.
     bulk_event = bulk_run_btn.click(
         fn=bulk_infer,
         inputs=[bulk_files,
@@ -1133,26 +826,19 @@ with gr.Blocks() as demo:
         outputs=[bulk_gallery, bulk_status, bulk_zip],
     )
 
-    # Stop cancels both running generators; the in-flight GPU call still
-    # finishes (ZeroGPU can't be killed mid-step), but no further iterations
-    # start. Any already-saved outputs remain in the gallery and on disk.
     bulk_stop_btn.click(fn=lambda: gr.Info("Stop requested — finishing current image."),
                         cancels=[bulk_event, run_event])
 
     # ── Depth / Pose tab wiring ──────────────────────────────────────────────
 
-    # Cache the last successfully-loaded source so re-renders after edits
-    # don't need the user to keep the upload widget populated.
     ctrl_source.change(
         fn=lambda img: img, inputs=[ctrl_source], outputs=[pose_source_state],
     )
 
-    # Depth generation
     detect_depth_btn.click(
         fn=generate_depthmap, inputs=[ctrl_source], outputs=[depth_output],
     )
 
-    # Pose detection → fills state, dropdowns, both preview images.
     def _on_detect_pose(source):
         if source is None:
             raise gr.Error("Upload a source image first.")
@@ -1175,13 +861,12 @@ with gr.Blocks() as demo:
         outputs=[pose_keypoints_state, active_person_dd, active_joint_dd,
                  pose_overlay, pose_clean],
     )
-    reset_pose_btn.click(  # same handler — re-runs detection
+    reset_pose_btn.click(
         fn=_on_detect_pose, inputs=[ctrl_source],
         outputs=[pose_keypoints_state, active_person_dd, active_joint_dd,
                  pose_overlay, pose_clean],
     )
 
-    # Insert a default standing-figure template centred in the source canvas.
     def _on_insert_blank(source):
         if source is None:
             raise gr.Error("Upload a source image first.")
@@ -1201,9 +886,6 @@ with gr.Blocks() as demo:
                  pose_overlay, pose_clean],
     )
 
-    # ── Click-to-edit: the heart of the pose editor ──
-    # gr.SelectData on a gr.Image gives .index = (x, y) in image pixels, even
-    # when interactive=False, which is exactly what we need.
     def _on_overlay_click(evt: gr.SelectData, poses, source, person_label, joint_name):
         if not poses or source is None or evt is None or evt.index is None:
             return gr.update(), gr.update(), gr.update()
@@ -1227,8 +909,6 @@ with gr.Blocks() as demo:
         outputs=[pose_keypoints_state, pose_overlay, pose_clean],
     )
 
-    # Changing the active joint or person just re-renders the overlay so the
-    # highlight ring follows — keypoints are not mutated.
     def _on_active_change(poses, source, person_label, joint_name):
         if not poses or source is None:
             return gr.update()
@@ -1249,7 +929,6 @@ with gr.Blocks() as demo:
         outputs=[pose_overlay],
     )
 
-    # Hide / clear
     def _on_hide_active(poses, source, person_label, joint_name):
         person_idx = parse_person_idx(person_label)
         joint_idx  = joint_name_to_index(joint_name)
@@ -1283,8 +962,6 @@ with gr.Blocks() as demo:
         outputs=[pose_keypoints_state, pose_overlay, pose_clean],
     )
 
-    # Send → main tab. We reuse the existing base_image / reference_images
-    # components so users land back on the Generate tab fully wired up.
     send_depth_ref_btn.click(
         fn=push_pil_to_reference, inputs=[depth_output, reference_images],
         outputs=[reference_images],
@@ -1307,7 +984,8 @@ with gr.Blocks() as demo:
     ).then(fn=on_base_image_change, inputs=[base_image], outputs=[size_info]
     ).then(fn=lambda: gr.Tabs(selected="tab_generate"), outputs=[main_tabs])
 
+
 if __name__ == "__main__":
-    # Gradio 6.0: theme and css go here, not on Blocks()
+    # Gradio 6.0: theme and css go on launch(), not Blocks()
     demo.queue().launch(css=css, theme=orange_red_theme,
                         mcp_server=True, ssr_mode=False, show_error=True)

lora_registry.py

@@ -33,8 +33,7 @@ FROM PICTURE 2 (strictly preserve identity):
 - 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
+# NOTE: MAX_LORA_SLOTS is imported from config.py — do NOT re-declare it here.
 
 LORA_STYLES = [
     {
@@ -85,7 +84,7 @@ LORA_STYLES = [
         "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",
@@ -93,7 +92,7 @@ LORA_STYLES = [
         "weights": "HighResolution9B.safetensors",
         "default_prompt": "High Resolution",
         "default_weight": 1.0,
-    },       
+    },
     {
         "title": "InstaPic",
         "adapter_name": "InstaPic V3",
@@ -101,7 +100,7 @@ LORA_STYLES = [
         "weights": "InstaPic V3.safetensors",
         "default_prompt": "instapic",
         "default_weight": 1.0,
-    },        
+    },
     {
         "title": "Realistic Nudes",
         "adapter_name": "Realistic Nudes",
@@ -109,7 +108,7 @@ LORA_STYLES = [
         "weights": "realistic_nudes_klein_v3.safetensors",
         "default_prompt": None,
         "default_weight": 1.0,
-    },            
+    },
     {
         "title": "Perky Pointy Puffy Breasts",
         "adapter_name": "Perky Pointy Puffy Breasts",
@@ -117,7 +116,7 @@ LORA_STYLES = [
         "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",
@@ -125,7 +124,7 @@ LORA_STYLES = [
         "weights": "Flux2-Klein-9b-FlatChested-v1.safetensors",
         "default_prompt": "flat chested",
         "default_weight": 1.5,
-    },           
+    },
     {
         "title": "Controllight",
         "adapter_name": "Controllight",
@@ -133,7 +132,7 @@ LORA_STYLES = [
         "weights": "controllight.safetensors",
         "default_prompt": None,
         "default_weight": 1.0,
-    },      
+    },
     {
         "title": "RefControl - Depth",
         "adapter_name": "RefConDep",
@@ -141,7 +140,7 @@ LORA_STYLES = [
         "weights": "flux2_klein_9b_refcontrol_depth.safetensors",
         "default_prompt": "refcontrol",
         "default_weight": 1.0,
-    },         
+    },
     {
         "title": "RefControl - Pose",
         "adapter_name": "RefConPos",
@@ -149,11 +148,10 @@ LORA_STYLES = [
         "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.
@@ -218,7 +216,6 @@ def update_weight_sliders(selected_titles, dynamic_loras):
         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)