Upload 3 files

app.py

@@ -0,0 +1,999 @@
+import gc
+import html
+import importlib.metadata
+import os
+import random
+import sys
+import threading
+import time
+
+NOISE_MARKERS = (b"CUDACachingAllocator.cpp", b"negative_prompt")
+
+
+def _install_stderr_filter():
+    """Drops known-noisy lines from stderr (must run before torch is imported)."""
+    try:
+        real_fd = os.dup(2)
+        read_fd, write_fd = os.pipe()
+        os.dup2(write_fd, 2)
+        os.close(write_fd)
+
+        def pump():
+            while True:
+                try:
+                    chunk = os.read(read_fd, 65536)
+                except OSError:
+                    break
+                if not chunk:
+                    break
+                chunk = b"".join(
+                    ln for ln in chunk.splitlines(keepends=True)
+                    if not any(m in ln for m in NOISE_MARKERS)
+                )
+                if chunk:
+                    try:
+                        os.write(real_fd, chunk)
+                    except OSError:
+                        break
+
+        threading.Thread(target=pump, daemon=True).start()
+    except Exception as exc:
+        sys.__stdout__.write(f"[log filter] not installed: {exc}\n")
+
+
+_install_stderr_filter()
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import gradio as gr
+import torch
+import spaces
+
+import torchao.quantization as _torchao_quantization
+print(f"[versions] torchao={importlib.metadata.version('torchao')} torch={torch.__version__}")
+if not hasattr(_torchao_quantization, "FqnToConfig"):
+    class FqnToConfig:
+        pass
+    _torchao_quantization.FqnToConfig = FqnToConfig
+
+from huggingface_hub import hf_hub_download
+from PIL import Image
+from safetensors import safe_open
+from starlette.middleware import Middleware
+from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig
+
+from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
+from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
+from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
+
+from security import (
+    FloodGuard,
+    check_admin_password,
+    check_admin_tab_visibility,
+    check_user_password,
+    clear_all_admin_log,
+    delete_admin_log_entry,
+    end_session,
+    is_admin,
+    is_quota_error,
+    is_user,
+    load_admin_gallery,
+    load_rate_stats,
+    log_generation,
+    mark_browser_session,
+    rate_acquire,
+    rate_release,
+    save_input_images,
+    show_rate_message,
+)
+
+_GR6 = int(str(gr.__version__).split(".")[0]) >= 6
+_PRIVATE = {"api_visibility": "private"} if _GR6 else {"api_name": False}
+
+import inspect as _inspect
+_SHOW_PROGRESS_ON = "show_progress_on" in _inspect.signature(gr.on).parameters
+
+_original_image_preprocess = gr.Image.preprocess
+
+
+def _safe_image_preprocess(self, *args, **kwargs):
+    try:
+        return _original_image_preprocess(self, *args, **kwargs)
+    except FileNotFoundError:
+        print("[image] stale client-side file reference (likely a Space restart) — treating this input as empty.")
+        return None
+
+
+gr.Image.preprocess = _safe_image_preprocess
+
+torch.cuda.is_current_stream_capturing = lambda: False
+
+
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+if torch.cuda.is_available():
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+MODEL_QWEN = "qwen"
+
+PRESET_RESOLUTIONS = [
+    ("720×1280 (9:16)", 720, 1280),
+    ("864×1536 (9:16)", 864, 1536),
+    ("1080×1920 (9:16)", 1080, 1920),
+    ("1440×2560 (9:16)", 1440, 2560),
+    ("1600×2560 (10:16)", 1600, 2560),
+    ("1536×2728 (9:16)", 1536, 2728),
+    ("1280×720 (16:9)", 1280, 720),
+    ("1536×864 (16:9)", 1536, 864),
+    ("1920×1080 (16:9)", 1920, 1080),
+    ("2560×1440 (16:9)", 2560, 1440),
+    ("2560×1600 (16:10)", 2560, 1600),
+    ("2728×1536 (16:9)", 2728, 1536),
+    ("768×1024 (3:4)", 768, 1024),
+    ("1080×1440 (3:4)", 1080, 1440),
+    ("1536×2048 (3:4)", 1536, 2048),
+    ("1024×1536 (2:3)", 1024, 1536),
+    ("1440×1800 (4:5)", 1440, 1800),
+    ("1080×2400 (9:20)", 1080, 2400),
+    ("1024×768 (4:3)", 1024, 768),
+    ("1440×1080 (4:3)", 1440, 1080),
+    ("2048×1536 (4:3)", 2048, 1536),
+    ("1536×1024 (3:2)", 1536, 1024),
+    ("1800×1440 (5:4)", 1800, 1440),
+    ("2560×1080 (21:9)", 2560, 1080),
+]
+PRESET_RESOLUTION_MAP = {label: (w, h) for label, w, h in PRESET_RESOLUTIONS}
+PRESET_TALL = [label for label, w, h in PRESET_RESOLUTIONS if w < h]
+PRESET_WIDE = [label for label, w, h in PRESET_RESOLUTIONS if w > h]
+PRESET_ALL = PRESET_TALL + PRESET_WIDE
+PRESET_DEFAULT = "720×1280 (9:16)"
+AUTO_SCALES = ["1x", "2x", "3x", "4x"]
+AUTO_BASE_PIXELS = 1024 * 1024
+AUTO_MAX_PIXELS = 4194304
+
+from diffusers.utils import logging as diffusers_logging
+
+diffusers_logging.set_verbosity_error()
+transformer = QwenImageTransformer2DModel.from_pretrained(
+    "AST-1320/Qwen-2511-Image-To-Image-Edit-V.17-Diffusers", subfolder="transformer", dtype=dtype
+)
+diffusers_logging.set_verbosity_warning()
+gc.collect()
+
+QWEN_TRANSFORMER_MODE = "fp8_exact"
+
+
+class _FP8ExactLinear(torch.nn.Module):
+    def __init__(self, linear: torch.nn.Linear):
+        super().__init__()
+        self.in_features = linear.in_features
+        self.out_features = linear.out_features
+        self.register_buffer("weight", linear.weight.detach().to(torch.float8_e4m3fn), persistent=False)
+        self.bias = linear.bias
+
+    def forward(self, x):
+        b = self.bias.to(x.dtype) if self.bias is not None else None
+        return torch.nn.functional.linear(x, self.weight.to(x.dtype), b)
+
+
+def _convert_linears_to_fp8_exact(model) -> int:
+    """Replaces every nn.Linear with _FP8ExactLinear and returns how many were replaced
+    (0 = nothing changed). Must run on CPU tensors, BEFORE pipe.to(device)."""
+    targets = [(n, m) for n, m in model.named_modules() if isinstance(m, torch.nn.Linear)]
+    if not targets:
+        return 0
+    for i, (n, m) in enumerate(targets):
+        if i % 8 == 0:
+            w = m.weight.detach()
+            if not torch.equal(w.to(torch.float8_e4m3fn).to(w.dtype), w):
+                print(f"[qwen] exact-FP8 check failed on {n!r}: weights are not plain E4M3 values - not converting")
+                return 0
+    for n, m in targets:
+        parent_name, _, child = n.rpartition(".")
+        parent = model.get_submodule(parent_name) if parent_name else model
+        setattr(parent, child, _FP8ExactLinear(m))
+    return len(targets)
+
+
+def _enable_vae_tiling(p, tile=None):
+    vae = getattr(p, "vae", None)
+    if vae is None:
+        return
+    try:
+        vae.enable_tiling()
+        vae.enable_slicing()
+    except Exception:
+        pass
+    if tile:
+        stride = (tile * 7 // 8) // 8 * 8
+        for name, value in (
+            ("tile_sample_min_height", tile),
+            ("tile_sample_min_width", tile),
+            ("tile_sample_stride_height", stride),
+            ("tile_sample_stride_width", stride),
+        ):
+            if hasattr(vae, name):
+                setattr(vae, name, value)
+        print("[vae] tiling applied:", {k: getattr(vae, k, None) for k in ("tile_sample_min_height", "tile_sample_min_width", "tile_sample_stride_height", "tile_sample_stride_width")})
+
+
+QWEN_VAE_TILE = int(os.environ.get("QWEN_VAE_TILE", "1600"))
+
+
+pipe = QwenImageEditPlusPipeline.from_pretrained(
+    "Qwen/Qwen-Image-Edit-2511",
+    transformer=transformer,
+    torch_dtype=dtype,
+    ignore_patterns=["transformer/*"],
+)
+
+_qwen_fp8_exact = False
+if QWEN_TRANSFORMER_MODE == "fp8_exact":
+    try:
+        _n_lin = _convert_linears_to_fp8_exact(pipe.transformer)
+        _qwen_fp8_exact = _n_lin > 0
+        if _qwen_fp8_exact:
+            print(f"[qwen] exact FP8 mode: {_n_lin} Linear layers keep the checkpoint's FP8 weights (bf16 compute)")
+        else:
+            print("[qwen] exact FP8 mode unavailable - falling back to torchao dynamic FP8")
+    except Exception as _e_fp8x:
+        print(f"[qwen] exact FP8 conversion failed ({_e_fp8x}) - falling back to torchao dynamic FP8")
+
+pipe = pipe.to(device)
+
+if _qwen_fp8_exact:
+    print("[qwen] transformer: exact FP8 weights, no torchao re-quantization")
+else:
+    quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
+    print("[qwen] FP8 quantization: SUCCESS (transformer)")
+
+_enable_vae_tiling(pipe, QWEN_VAE_TILE)
+
+pipe.transformer.__class__ = QwenImageTransformer2DModel
+pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+
+gc.collect()
+
+
+MAX_SEED = 2**31 - 1
+_VALID_FORMATS = ("PNG", "JPG", "WebP")
+
+GPU_DURATION_MIN = 10
+GPU_DURATION_MAX = 120
+GPU_DURATION_STEP = 5
+GPU_DURATION_DEFAULT = 40
+_GPU_DURATION_STEPS = list(range(GPU_DURATION_MIN, GPU_DURATION_MAX + 1, GPU_DURATION_STEP))
+GPU_SIZE_OPTIONS = ("large", "xlarge")
+GPU_SIZE_DEFAULT = "large"
+
+
+def _to_int(value, lo, hi, default):
+    try:
+        v = int(value)
+    except (TypeError, ValueError, OverflowError):
+        return default
+    return max(lo, min(hi, v))
+
+
+LIMIT_GUIDANCE_MIN = 1.0
+LIMIT_GUIDANCE_MAX = 2.0
+
+
+def _to_float(value, lo, hi, default):
+    try:
+        v = float(value)
+    except (TypeError, ValueError, OverflowError):
+        return default
+    if v != v:  # NaN
+        return default
+    return max(lo, min(hi, v))
+
+
+def _round_size(width, height, multiple=8):
+    return max(multiple, (int(width) // multiple) * multiple), max(multiple, (int(height) // multiple) * multiple)
+
+
+def _auto_size(scale_choice, image_1, image_2):
+    src = image_1 if image_1 is not None else image_2
+    if src is None:
+        return None, None
+    try:
+        with Image.open(src if isinstance(src, str) else src.name) as im:
+            width, height = im.size
+    except Exception:
+        return None, None
+    factor = _to_int(str(scale_choice).rstrip("x"), 1, 4, 1)
+    pixels = min(AUTO_BASE_PIXELS * factor, AUTO_MAX_PIXELS)
+    ratio = width / height
+    out_w = (pixels * ratio) ** 0.5
+    return _round_size(int(out_w), int(out_w / ratio), multiple=16)
+
+
+def _resolve_size(mode, preset, custom_width, custom_height, auto_scale, image_1, image_2):
+    try:
+        if mode == "Preset" and preset in PRESET_RESOLUTION_MAP:
+            return _round_size(*PRESET_RESOLUTION_MAP[preset])
+        if mode == "Custom":
+            return _round_size(int(custom_width), int(custom_height))
+        if mode == "Auto":
+            return _auto_size(auto_scale, image_1, image_2)
+    except (TypeError, ValueError, OverflowError):
+        pass
+    return None, None
+
+
+def _load_image(img):
+    if not isinstance(img, Image.Image):
+        img = Image.open(img if isinstance(img, str) else img.name)
+    return img.convert("RGB")
+
+
+def _sanitize_common(seed, gpu_duration, output_format):
+    seed = _to_int(seed, 0, MAX_SEED, 0)
+    gpu_duration = _to_int(gpu_duration, GPU_DURATION_MIN, GPU_DURATION_MAX, GPU_DURATION_DEFAULT)
+    output_format = output_format if output_format in _VALID_FORMATS else "PNG"
+    return seed, gpu_duration, output_format
+
+
+def _first_path(gallery_items):
+    first = gallery_items[0]
+    return first[0] if isinstance(first, (list, tuple)) else first
+
+
+def _free_vram_cache():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def log_vram(label):
+    if not torch.cuda.is_available():
+        print(f"[VRAM] {label}: CUDA not available")
+        return
+    print(
+        f"[VRAM] {label}: allocated={torch.cuda.memory_allocated() / 1e9:.2f}GB "
+        f"reserved={torch.cuda.memory_reserved() / 1e9:.2f}GB "
+        f"peak={torch.cuda.max_memory_allocated() / 1e9:.2f}GB / "
+        f"device_total={torch.cuda.get_device_properties(0).total_memory / 1e9:.2f}GB"
+    )
+
+
+def _generate_core(
+    image_1,
+    image_2,
+    prompt,
+    seed,
+    randomize_seed,
+    true_guidance_scale,
+    num_inference_steps,
+    height,
+    width,
+    duplicate_input=True,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    pil_images = []
+    for img in (image_1, image_2):
+        if img is None:
+            continue
+        try:
+            pil_images.append(_load_image(img))
+        except Exception:
+            continue
+
+    duplicated = False
+    if duplicate_input and len(pil_images) == 1:
+        pil_images.append(pil_images[0].copy())
+        duplicated = True
+
+    res_txt = f"{width}x{height}" if (width and height) else "auto"
+    print(f"[gen] model={MODEL_QWEN} guidance={true_guidance_scale} steps={num_inference_steps} resolution={res_txt}")
+
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats(device)
+    log_vram(f"{MODEL_QWEN} before generation")
+
+    try:
+        images_pil = pipe(
+            image=pil_images or None,
+            prompt=prompt,
+            height=height,
+            width=width,
+            negative_prompt=" ",
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            true_cfg_scale=true_guidance_scale,
+            num_images_per_prompt=1,
+        ).images
+
+    except torch.cuda.OutOfMemoryError:
+        _free_vram_cache()
+        raise
+    log_vram(f"{MODEL_QWEN} after generation")
+
+    return {
+        "images": images_pil, "seed": seed, "width": width, "height": height,
+        "duplicated": duplicated, "steps": num_inference_steps,
+        "guidance": true_guidance_scale, "model": MODEL_QWEN,
+    }
+
+
+def _finalize(res, image_1, image_2, prompt, output_format):
+    seed = res["seed"]
+    input_paths = save_input_images((image_1, image_2), seed)
+
+    os.makedirs("outputs", exist_ok=True)
+    ext = {"PNG": "png", "JPG": "jpg", "WebP": "webp"}.get(output_format, "png")
+    output_paths = []
+    for idx, img in enumerate(res["images"]):
+        output_path = f"outputs/output_{seed}_{idx}_{int(time.time() * 1000)}.{ext}"
+        if ext == "jpg":
+            img.convert("RGB").save(output_path, quality=100)
+        elif ext == "webp":
+            img.save(output_path, quality=100)
+        else:
+            img.save(output_path)
+        output_paths.append(output_path)
+        log_generation(
+            path=output_path,
+            prompt=prompt,
+            seed=seed,
+            steps=res["steps"],
+            guidance=res["guidance"],
+            model=res["model"],
+            width=res["width"],
+            height=res["height"],
+            out_width=img.width,
+            out_height=img.height,
+            duplicated=res["duplicated"],
+            inputs=input_paths,
+        )
+    has_output = bool(output_paths)
+    return output_paths, seed, gr.update(visible=has_output)
+
+
+def _make_gpu_variant(duration_seconds, size):
+    def _infer_gpu_call(*args, **kwargs):
+        return _generate_core(*args, **kwargs)
+
+    _infer_gpu_call.__name__ = f"_infer_impl_gpu_{size}_{duration_seconds}s"
+    _infer_gpu_call.__qualname__ = _infer_gpu_call.__name__
+    return spaces.GPU(duration=duration_seconds, size=size)(_infer_gpu_call)
+
+
+_gpu_variants = {(s, d): _make_gpu_variant(d, s) for s in GPU_SIZE_OPTIONS for d in _GPU_DURATION_STEPS}
+
+
+def _pick_gpu_variant(gpu_duration, gpu_size=GPU_SIZE_DEFAULT):
+    requested = _to_int(gpu_duration, GPU_DURATION_MIN, GPU_DURATION_MAX, GPU_DURATION_DEFAULT)
+    chosen = next((d for d in _GPU_DURATION_STEPS if d >= requested), _GPU_DURATION_STEPS[-1])
+    size = gpu_size if gpu_size in GPU_SIZE_OPTIONS else GPU_SIZE_DEFAULT
+    gpu_fn = _gpu_variants[(size, chosen)]
+
+    def _call(image_1, image_2, prompt, *args, output_format="PNG", **kwargs):
+        res = gpu_fn(image_1, image_2, prompt, *args, **kwargs)
+        return _finalize(res, image_1, image_2, prompt, output_format)
+
+    return _call
+
+
+def _run_limited(request, seed, run):
+    if not is_user(request):
+        gr.Warning("Please log in first.")
+        return gr.update(), seed, gr.update()
+    ok, key, msg = rate_acquire(request)
+    if not ok:
+        show_rate_message(msg)
+        return gr.update(), seed, gr.update()
+    refund = False
+    try:
+        return run()
+    except Exception as exc:
+        refund = is_quota_error(exc)
+        raise
+    finally:
+        rate_release(key, refund=refund)
+
+
+def infer_qwen_only(
+    image_1,
+    image_2,
+    prompt,
+    seed,
+    randomize_seed,
+    true_guidance_scale,
+    num_inference_steps,
+    gpu_duration,
+    gpu_size,
+    resolution_mode,
+    resolution_preset,
+    custom_width,
+    custom_height,
+    auto_scale,
+    output_format,
+    duplicate_input,
+    progress=gr.Progress(track_tqdm=True),
+    request: gr.Request = None,
+):
+    if not (isinstance(prompt, str) and prompt.strip()):
+        return gr.update(), seed, gr.update()
+
+    def run():
+        clean_seed, duration, fmt = _sanitize_common(seed, gpu_duration, output_format)
+        width, height = _resolve_size(resolution_mode, resolution_preset, custom_width, custom_height, auto_scale, image_1, image_2)
+        return _pick_gpu_variant(duration, gpu_size)(
+            image_1,
+            image_2,
+            prompt if isinstance(prompt, str) else "",
+            seed=clean_seed,
+            randomize_seed=bool(randomize_seed),
+            true_guidance_scale=_to_float(true_guidance_scale, LIMIT_GUIDANCE_MIN, LIMIT_GUIDANCE_MAX, 1.0),
+            num_inference_steps=int(num_inference_steps),
+            height=height,
+            width=width,
+            output_format=fmt,
+            duplicate_input=bool(duplicate_input),
+            progress=progress,
+        )
+
+    return _run_limited(request, seed, run)
+
+
+def use_output_as_input(output_images):
+    if not output_images:
+        return gr.update()
+    return gr.update(value=_first_path(output_images))
+
+
+def _sync_result_followups(output_images):
+    return gr.update(visible=bool(output_images))
+
+
+def _clear_all_admin_images(request: gr.Request):
+    if not is_admin(request):
+        return []
+    clear_all_admin_log()
+    return load_admin_gallery(request)
+
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 1024px;
+}
+#logo-title {
+    text-align: center;
+}
+
+@keyframes ast-title-flash {
+    0% { background-position: 0% 50%; }
+    100% { background-position: 200% 50%; }
+}
+#ast-main-title {
+    font-style: italic;
+    display: inline-block;
+    background: linear-gradient(90deg, #FF8C00, #FF3CAC, #7B2FF7, #00E5FF, #FF8C00);
+    background-size: 400% 400%;
+    -webkit-background-clip: text;
+    background-clip: text;
+    -webkit-text-fill-color: transparent;
+    color: transparent;
+    animation: ast-title-flash 10s linear infinite;
+}
+/* Force a 2-column grid for the preset radio (Gradio's default flex-wrap stacks into one column on mobile). */
+.preset-grid > .wrap,
+.preset-grid .wrap {
+    display: grid !important;
+    grid-template-columns: 1fr 1fr !important;
+    gap: 8px !important;
+}
+#adv-settings:has(#res-mode input[value="Preset"]:checked) #res-custom-height,
+#adv-settings:has(#res-mode input[value="Preset"]:checked) #res-custom-width,
+#adv-settings:has(#res-mode input[value="Preset"]:checked) #res-auto-scale,
+#adv-settings:has(#res-mode input[value="Custom"]:checked) #res-preset,
+#adv-settings:has(#res-mode input[value="Custom"]:checked) #res-auto-scale,
+#adv-settings:has(#res-mode input[value="Auto"]:checked) #res-preset,
+#adv-settings:has(#res-mode input[value="Auto"]:checked) #res-custom-height,
+#adv-settings:has(#res-mode input[value="Auto"]:checked) #res-custom-width {
+    display: none !important;
+}
+/* Seed is an output component: hide Gradio's progress/status overlay so it never covers the number box. */
+#seed-slider .wrap.generating,
+#seed-slider .wrap.pending,
+#seed-slider .progress-text,
+#seed-slider .progress-level {
+    display: none !important;
+}
+.preset-grid label {
+    width: 100% !important;
+    margin: 0 !important;
+    box-sizing: border-box !important;
+}
+#res-mode .wrap {
+    display: grid !important;
+    grid-template-columns: repeat(3, 1fr) !important;
+    gap: 8px !important;
+}
+#res-mode label {
+    width: 100% !important;
+    margin: 0 !important;
+    box-sizing: border-box !important;
+    justify-content: center !important;
+}
+#output-format .wrap {
+    display: grid !important;
+    grid-template-columns: repeat(3, 1fr) !important;
+    gap: 8px !important;
+}
+#output-format label {
+    width: 100% !important;
+    margin: 0 !important;
+    box-sizing: border-box !important;
+    justify-content: center !important;
+}
+#gpu-size .wrap {
+    display: grid !important;
+    grid-template-columns: repeat(2, 1fr) !important;
+    gap: 8px !important;
+}
+#gpu-size label {
+    width: 100% !important;
+    margin: 0 !important;
+    box-sizing: border-box !important;
+    justify-content: center !important;
+}
+#gpu-size-note,
+#gpu-size-note * {
+    background: transparent !important;
+    border: none !important;
+    box-shadow: none !important;
+}
+#gpu-size-note {
+    background: var(--block-background-fill) !important;
+    padding: 14px 12px 16px 12px !important;
+    container-type: inline-size;
+}
+#gpu-size-note p {
+    margin: 0 !important;
+    font-family: inherit !important;
+    font-size: var(--checkbox-label-text-size, var(--text-md, 1rem)) !important;
+    font-size: min(var(--checkbox-label-text-size, var(--text-md, 1rem)), 5.2cqw) !important;
+    white-space: nowrap !important;
+    font-style: normal !important;
+    line-height: 1.5 !important;
+    opacity: 0.7;
+}
+#res-auto-scale .wrap {
+    display: grid !important;
+    grid-template-columns: repeat(4, minmax(0, 1fr)) !important;
+    gap: 6px !important;
+}
+#res-auto-scale label {
+    width: 100% !important;
+    min-width: 0 !important;
+    margin: 0 !important;
+    padding-left: 8px !important;
+    padding-right: 8px !important;
+    box-sizing: border-box !important;
+    justify-content: center !important;
+}
+"""
+
+_div_a = f'.preset-grid label:has(input[value="{PRESET_WIDE[0]}"])'
+_div_b = f'.preset-grid label:has(input[value="{PRESET_WIDE[1]}"])'
+css += (
+    f"{_div_a}, {_div_b} {{ position: relative !important; margin-top: 12px !important; overflow: visible !important; }}\n"
+    f"{_div_a}::before, {_div_b}::before {{ content: \"\"; position: absolute; top: -10px; height: 1px; "
+    f"background: rgba(128, 128, 128, 0.55); pointer-events: none; }}\n"
+    f"{_div_a}::before {{ left: 0; right: -4px; }}\n"
+    f"{_div_b}::before {{ left: -4px; right: 0; }}\n"
+)
+
+_theme_dir = os.path.dirname(os.path.abspath(__file__))
+_theme_pick = os.environ.get("UI_THEME", "")
+if _theme_pick == "off":
+    _theme_name = None
+elif _theme_pick:
+    _theme_name = f"theme_{_theme_pick}.css"
+else:
+    _found = sorted(f for f in os.listdir(_theme_dir) if f.startswith("theme_") and f.endswith(".css"))
+    _theme_name = _found[0] if _found else None
+if _theme_name and os.path.isfile(os.path.join(_theme_dir, _theme_name)):
+    with open(os.path.join(_theme_dir, _theme_name), encoding="utf-8") as _f:
+        css = _f.read() + css
+
+with gr.Blocks(**({} if _GR6 else {"css": css})) as demo:
+    with gr.Tabs():
+        with gr.Tab("Editor"):
+            with gr.Group(visible=True) as login_group:
+                login_password = gr.Textbox(label="Password", type="password", placeholder="Enter password")
+                login_btn = gr.Button("Login", variant="primary")
+                login_status = gr.Markdown("")
+            with gr.Column(elem_id="col-container", visible=False) as app_col:
+                gr.HTML('<div id="logo-title"><h1 id="ast-main-title">Qwen 2511 Image Edit</h1></div>')
+                with gr.Row():
+                    with gr.Column():
+                        with gr.Row():
+                            image_1 = gr.Image(label="Image 1", type="filepath", interactive=True)
+                            with gr.Column():
+                                image_2 = gr.Image(label="Image 2 (optional)", type="filepath", interactive=True)
+                                duplicate_input = gr.Checkbox(
+                                    label="Duplicate image 1",
+                                    value=False,
+                                )
+
+                        prompt = gr.Text(
+                            label="Prompt 🪄",
+                            show_label=True,
+                            placeholder="Enter your prompt here...",
+                        )
+                        run_button = gr.Button("Edit!", variant="primary", interactive=False)
+
+                        with gr.Group():
+                            gpu_size = gr.Radio(
+                                label="GPU Size",
+                                choices=[("Large", "large"), ("Xlarge", "xlarge")],
+                                value=GPU_SIZE_DEFAULT,
+                                elem_id="gpu-size",
+                            )
+                            gr.Markdown(
+                                "⚠️ Note: Xlarge is faster but uses 2× quota",
+                                elem_id="gpu-size-note",
+                            )
+
+                        with gr.Group():
+                            output_format = gr.Radio(
+                                label="Output Format",
+                                choices=["PNG", "JPG", "WebP"],
+                                value="PNG",
+                                elem_id="output-format",
+                            )
+
+                        with gr.Accordion("Advanced Settings", open=False, elem_id="adv-settings"):
+                            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, elem_id="seed-slider")
+                            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+                            with gr.Row():
+                                true_guidance_scale = gr.Slider(
+                                    label="True guidance scale",
+                                    minimum=1.0,
+                                    maximum=2.0,
+                                    step=0.1,
+                                    value=1.0,
+                                )
+                                num_inference_steps = gr.Slider(
+                                    label="Number of inference steps",
+                                    minimum=1,
+                                    maximum=20,
+                                    step=1,
+                                    value=4,
+                                )
+                                gpu_duration = gr.Slider(
+                                    label="GPU duration (seconds)",
+                                    minimum=GPU_DURATION_MIN,
+                                    maximum=GPU_DURATION_MAX,
+                                    step=GPU_DURATION_STEP,
+                                    value=GPU_DURATION_DEFAULT,
+                                )
+
+                            with gr.Group():
+                                resolution_mode = gr.Radio(
+                                    label="Resolution mode",
+                                    choices=["Auto", "Preset", "Custom"],
+                                    value="Auto",
+                                    elem_id="res-mode",
+                                )
+                                resolution_preset = gr.Radio(
+                                    label="Resolution",
+                                    choices=PRESET_ALL,
+                                    value=PRESET_DEFAULT,
+                                    elem_classes=["preset-grid"],
+                                    elem_id="res-preset",
+                                )
+                                custom_height = gr.Slider(
+                                    label="Height", minimum=256, maximum=2048, step=8, value=256,
+                                    buttons=[], elem_id="res-custom-height",
+                                )
+                                custom_width = gr.Slider(
+                                    label="Width", minimum=256, maximum=2048, step=8, value=256,
+                                    buttons=[], elem_id="res-custom-width",
+                                )
+                                auto_scale = gr.Radio(
+                                    label="Auto resolution",
+                                    choices=AUTO_SCALES,
+                                    value="1x",
+                                    elem_id="res-auto-scale",
+                                )
+
+                    with gr.Column():
+                        result = gr.Gallery(label="Result", show_label=False, type="filepath")
+                        with gr.Row():
+                            use_output_btn = gr.Button("↗️ Use as input", variant="secondary", size="sm", visible=False)
+
+            gr.HTML('<div style="height:1024px"></div>', elem_id="login-spacer")
+
+        with gr.Tab("🔒", visible=False) as admin_tab:
+            with gr.Group(visible=True) as admin_lock_group:
+                admin_password = gr.Textbox(label="Password", type="password", placeholder="Enter admin password")
+                admin_unlock_btn = gr.Button("Unlock", size="sm", variant="primary")
+                admin_status = gr.Markdown("")
+            with gr.Group(visible=False) as admin_unlocked_group:
+                with gr.Row():
+                    admin_refresh_btn = gr.Button("🔄 Refresh", size="sm")
+                    admin_clear_all_btn = gr.Button("🗑️ Clear ALL", size="sm", variant="stop")
+                admin_stats_md = gr.Markdown("")
+                admin_items_state = gr.State([])
+
+                @gr.render(inputs=[admin_items_state])
+                def _render_admin_grid(items):
+                    if not items:
+                        gr.Markdown("*No images in history.*")
+                        return
+                    for row_start in range(0, len(items), 3):
+                        with gr.Row():
+                            for item in items[row_start:row_start + 3]:
+                                is_pair = isinstance(item, (tuple, list))
+                                path = item[0] if is_pair else item
+                                caption = item[1] if is_pair and len(item) > 1 else ""
+                                inputs = list(item[2]) if is_pair and len(item) > 2 else []
+                                shown = (
+                                    [(p, f"Input {n}") for n, p in enumerate(inputs, 1)] + [(path, "Output")]
+                                    if inputs else [path]
+                                )
+                                with gr.Column(min_width=160):
+                                    gr.Gallery(
+                                        value=shown,
+                                        show_label=False,
+                                        columns=len(shown),
+                                        rows=1,
+                                        height=320,
+                                        object_fit="contain",
+                                        allow_preview=True,
+                                    )
+                                    if caption:
+                                        gr.HTML(f"<small>{html.escape(str(caption))}</small>")
+                                    del_btn = gr.Button("🗑️ Remove", size="sm", variant="stop")
+
+                                    def _delete_this(request: gr.Request, _path=path):
+                                        if not is_admin(request):
+                                            return []
+                                        delete_admin_log_entry(_path)
+                                        return load_admin_gallery(request)
+
+                                    del_btn.click(
+                                        fn=_delete_this,
+                                        inputs=None,
+                                        outputs=[admin_items_state],
+                                        **_PRIVATE,
+                                    )
+
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer_qwen_only,
+        inputs=[
+            image_1,
+            image_2,
+            prompt,
+            seed,
+            randomize_seed,
+            true_guidance_scale,
+            num_inference_steps,
+            gpu_duration,
+            gpu_size,
+            resolution_mode,
+            resolution_preset,
+            custom_width,
+            custom_height,
+            auto_scale,
+            output_format,
+            duplicate_input,
+        ],
+        outputs=[result, seed, use_output_btn],
+        **({"show_progress_on": [result]} if _SHOW_PROGRESS_ON else {}),
+    ).then(
+        fn=_sync_result_followups,
+        inputs=result,
+        outputs=use_output_btn,
+        show_progress="hidden",
+    )
+
+    prompt.change(
+        fn=None,
+        inputs=prompt,
+        outputs=run_button,
+        js="(text) => ({__type__: 'update', interactive: !!(text && text.trim())})",
+        show_progress="hidden",
+        **_PRIVATE,
+    )
+
+    use_output_btn.click(fn=use_output_as_input, inputs=[result], outputs=[image_1])
+
+    result.change(
+        fn=_sync_result_followups,
+        inputs=result,
+        outputs=use_output_btn,
+        show_progress="hidden",
+    )
+
+    gr.on(
+        triggers=[login_btn.click, login_password.submit],
+        fn=check_user_password,
+        inputs=[login_password],
+        outputs=[login_group, app_col, login_status],
+        **_PRIVATE,
+    ).then(
+        fn=None, inputs=None, outputs=None,
+        js="""() => {
+            const nudge = () => {
+                window.dispatchEvent(new Event("resize"));
+                try { if (window.parentIFrame && window.parentIFrame.size) window.parentIFrame.size(); } catch (e) {}
+            };
+            setTimeout(() => {
+                const sp = document.getElementById("login-spacer");
+                const col = document.getElementById("col-container");
+                if (sp && col && col.offsetParent !== null) sp.style.display = "none";
+                nudge();
+            }, 600);
+            [900, 1500, 3000].forEach((ms) => setTimeout(nudge, ms));
+        }""",
+        show_progress="hidden",
+        **_PRIVATE,
+    )
+
+    resolution_mode.change(
+        fn=None, inputs=None, outputs=None,
+        js="""() => {
+            const nudge = () => {
+                window.dispatchEvent(new Event("resize"));
+                try { if (window.parentIFrame && window.parentIFrame.size) window.parentIFrame.size(); } catch (e) {}
+            };
+            [0, 100, 300, 800].forEach((ms) => setTimeout(nudge, ms));
+        }""",
+        show_progress="hidden",
+        **_PRIVATE,
+    )
+
+    try:
+        demo.unload(end_session)
+    except Exception as e:
+        print(f"[auth] demo.unload not available, sessions persist until restart: {e}")
+
+    demo.load(fn=mark_browser_session, inputs=None, outputs=None, **_PRIVATE)
+    demo.load(fn=check_admin_tab_visibility, inputs=None, outputs=[admin_tab], **_PRIVATE)
+
+    admin_unlock_btn.click(
+        fn=check_admin_password,
+        inputs=[admin_password],
+        outputs=[admin_lock_group, admin_unlocked_group, admin_status],
+        **_PRIVATE,
+    ).then(
+        fn=load_admin_gallery, inputs=None, outputs=[admin_items_state], **_PRIVATE,
+    ).then(
+        fn=load_rate_stats, inputs=None, outputs=[admin_stats_md], **_PRIVATE,
+    )
+    admin_refresh_btn.click(
+        fn=load_admin_gallery, inputs=None, outputs=[admin_items_state], **_PRIVATE,
+    ).then(
+        fn=load_rate_stats, inputs=None, outputs=[admin_stats_md], **_PRIVATE,
+    )
+    admin_clear_all_btn.click(
+        fn=_clear_all_admin_images, inputs=None, outputs=[admin_items_state], **_PRIVATE,
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=int(os.environ.get("QUEUE_MAX_SIZE", "30")))
+    launch_kwargs = dict(
+        ssr_mode=False,
+        max_file_size="20mb",
+        app_kwargs={"middleware": [Middleware(FloodGuard)]},
+    )
+    if _GR6:
+        launch_kwargs["css"] = css
+        launch_kwargs["footer_links"] = ["gradio", "settings"]
+    else:
+        launch_kwargs["show_api"] = False
+    demo.launch(**launch_kwargs)

requirements.txt

@@ -7,6 +7,4 @@ torchvision==0.28.0
 torchao==0.18.0
 huggingface-hub==1.33.0
 tokenizers==0.23.3
-pillow==12.3.0
-opencv-python-headless==4.13.0.92
-mediapipe==0.10.35
+pillow==12.3.0

security.py

@@ -0,0 +1,655 @@
+import hmac
+import ipaddress
+import json
+import os
+import shutil
+import threading
+import time
+
+import gradio as gr
+
+ADMIN_PASSWORD = os.environ.get("ADMIN_PASSWORD", "")
+ADMIN_URL_KEY = os.environ.get("ADMIN_URL_KEY", "")
+ADMIN_LOG_PATH = "outputs/generation_log.jsonl"
+ADMIN_INPUT_DIR = "outputs/inputs"
+ADMIN_MAX_ENTRIES = 300
+
+
+def _session_id(request):
+    return getattr(request, "session_hash", None)
+
+
+def _read_log_lines():
+    if not os.path.exists(ADMIN_LOG_PATH):
+        return None
+    try:
+        with open(ADMIN_LOG_PATH, "r", encoding="utf-8") as f:
+            return [line.strip() for line in f if line.strip()]
+    except Exception as e:
+        print(f"[admin log] read failed: {e}")
+        return None
+
+
+def _write_log_lines(lines):
+    try:
+        with open(ADMIN_LOG_PATH, "w", encoding="utf-8") as f:
+            f.write("\n".join(lines) + ("\n" if lines else ""))
+    except Exception as e:
+        print(f"[admin log] write failed: {e}")
+
+
+def _remove_file(path):
+    try:
+        if path and os.path.exists(path):
+            os.remove(path)
+    except Exception as e:
+        print(f"[admin log] failed to remove {path}: {e}")
+
+
+def _entry_path(line):
+    try:
+        return json.loads(line).get("path")
+    except Exception:
+        return None
+
+
+def _entry_inputs(line):
+    try:
+        return list(json.loads(line).get("inputs") or [])
+    except Exception:
+        return []
+
+
+def _entry_files(line):
+    return [_entry_path(line), *_entry_inputs(line)]
+
+
+def save_input_images(sources, seed):
+    paths = []
+    try:
+        os.makedirs(ADMIN_INPUT_DIR, exist_ok=True)
+        stamp = int(time.time() * 1000)
+        for idx, src in enumerate([s for s in sources if s is not None]):
+            try:
+                src_path = src if isinstance(src, str) else getattr(src, "name", None)
+                if src_path and os.path.isfile(src_path):
+                    ext = os.path.splitext(src_path)[1].lower()[:8] or ".png"
+                    input_path = f"{ADMIN_INPUT_DIR}/input_{seed}_{idx}_{stamp}{ext}"
+                    shutil.copy2(src_path, input_path)
+                else:
+                    input_path = f"{ADMIN_INPUT_DIR}/input_{seed}_{idx}_{stamp}.png"
+                    src.save(input_path)
+                paths.append(input_path)
+            except Exception as e:
+                print(f"[admin log] failed to save input image {idx}: {e}")
+    except Exception as e:
+        print(f"[admin log] failed to save input images: {e}")
+    return paths
+
+
+def log_generation(path, prompt, seed, steps, guidance, model, width, height,
+                   out_width=None, out_height=None, duplicated=False, inputs=None):
+    try:
+        os.makedirs("outputs", exist_ok=True)
+        entry = {
+            "path": path,
+            "prompt": prompt,
+            "seed": seed,
+            "steps": steps,
+            "guidance": guidance,
+            "upscaler": model,
+            "width": width,
+            "height": height,
+            "out_width": out_width,
+            "out_height": out_height,
+            "duplicated_input": duplicated,
+            "inputs": list(inputs or []),
+            "timestamp": time.time(),
+        }
+        with open(ADMIN_LOG_PATH, "a", encoding="utf-8") as f:
+            f.write(json.dumps(entry, ensure_ascii=False) + "\n")
+        _prune_log()
+    except Exception as e:
+        print(f"[admin log] failed to write entry: {e}")
+
+
+def _prune_log():
+    lines = _read_log_lines()
+    if lines is None or len(lines) <= ADMIN_MAX_ENTRIES:
+        return
+    for line in lines[:-ADMIN_MAX_ENTRIES]:
+        for file_path in _entry_files(line):
+            _remove_file(file_path)
+    _write_log_lines(lines[-ADMIN_MAX_ENTRIES:])
+
+
+def delete_admin_log_entry(path):
+    if not path:
+        return
+    lines = _read_log_lines()
+    if lines is None:
+        return
+    kept = [line for line in lines if _entry_path(line) != path]
+    removed = [line for line in lines if _entry_path(line) == path]
+    _write_log_lines(kept)
+    for line in removed:
+        for file_path in _entry_files(line):
+            _remove_file(file_path)
+    _remove_file(path)
+
+
+def clear_all_admin_log():
+    lines = _read_log_lines()
+    if lines is None:
+        return
+    for line in lines:
+        for file_path in _entry_files(line):
+            _remove_file(file_path)
+    _write_log_lines([])
+
+
+def load_admin_gallery(request: gr.Request, limit=200):
+    if not is_admin(request):
+        return []
+    items = []
+    for line in reversed(_read_log_lines() or []):
+        try:
+            e = json.loads(line)
+        except Exception:
+            continue
+        path = e.get("path")
+        if not path or not os.path.exists(path):
+            continue
+        ts = e.get("timestamp")
+        when = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(ts)) if ts else "?"
+        w, h = e.get("width"), e.get("height")
+        res_txt = f"{w}x{h}" if (w and h) else "auto"
+        ow, oh = e.get("out_width"), e.get("out_height")
+        if ow and oh and (ow, oh) != (w, h):
+            res_txt += f" -> {ow}x{oh}"
+        caption = (
+            f"[{when}] seed={e.get('seed')} | guidance={e.get('guidance')} | "
+            f"steps={e.get('steps')} | resolution={res_txt} | {e.get('prompt', '')}"
+        )
+        inputs = [p for p in (e.get("inputs") or []) if os.path.exists(p)]
+        items.append((path, caption, inputs))
+        if len(items) >= limit:
+            break
+    return items
+
+
+_ADMIN_URL_OK_SESSIONS = set()
+_ADMIN_SESSIONS = set()
+
+
+_ADMIN_FAIL_TIMES = []
+_ADMIN_FAIL_LIMIT = 8
+_ADMIN_FAIL_WINDOW = 600
+
+
+def is_admin(request):
+    sid = _session_id(request)
+    return bool(sid) and sid in _ADMIN_SESSIONS
+
+
+def _too_many_admin_failures():
+    now = time.time()
+    _ADMIN_FAIL_TIMES[:] = [t for t in _ADMIN_FAIL_TIMES if now - t < _ADMIN_FAIL_WINDOW]
+    return len(_ADMIN_FAIL_TIMES) >= _ADMIN_FAIL_LIMIT
+
+
+def check_admin_tab_visibility(request: gr.Request):
+    sid = _session_id(request)
+    if not ADMIN_URL_KEY or not sid:
+        return gr.update(visible=False)
+    key = request.query_params.get("admin", "") if request else ""
+    ok = hmac.compare_digest(str(key).encode("utf-8"), ADMIN_URL_KEY.encode("utf-8"))
+    if ok:
+        _ADMIN_URL_OK_SESSIONS.add(sid)
+    else:
+        _ADMIN_URL_OK_SESSIONS.discard(sid)
+        _ADMIN_SESSIONS.discard(sid)
+    return gr.update(visible=ok)
+
+
+def check_admin_password(password, request: gr.Request):
+    locked = (gr.update(visible=True), gr.update(visible=False))
+    if not ADMIN_PASSWORD:
+        return (*locked, "⚠️ No ADMIN_PASSWORD secret set on this Space — admin gallery stays locked.")
+    sid = _session_id(request)
+    if not sid or sid not in _ADMIN_URL_OK_SESSIONS:
+        return (*locked, "❌ Wrong password")
+    if _too_many_admin_failures():
+        return (*locked, "⏳ Too many failed attempts — try again later.")
+    if hmac.compare_digest(str(password or "").encode("utf-8"), ADMIN_PASSWORD.encode("utf-8")):
+        _ADMIN_SESSIONS.add(sid)
+        return gr.update(visible=False), gr.update(visible=True), ""
+    _ADMIN_FAIL_TIMES.append(time.time())
+    time.sleep(1.0)
+    return (*locked, "❌ Wrong password")
+
+
+APP_PASSWORD = os.environ.get("APP_PASSWORD", "")
+_USER_SESSIONS = set()
+_USER_LOCK = threading.Lock()
+_USER_FAILS = {}
+_USER_FAILS_ALL = []
+_USER_MAX_FAILS = 5
+_USER_MAX_FAILS_ALL = 60
+_USER_WINDOW_S = 600
+_USER_FAILS_MAX_KEYS = 5000
+
+
+def is_user(request):
+    sid = _session_id(request)
+    if not sid:
+        return False
+    with _USER_LOCK:
+        return sid in _USER_SESSIONS
+
+
+def check_user_password(password, request: gr.Request):
+    locked = (gr.update(visible=True), gr.update(visible=False))
+    if not APP_PASSWORD:
+        return (*locked, "⚠️ No APP_PASSWORD secret set on this Space — app stays locked.")
+    sid = _session_id(request)
+    ckey = _client_key(request)
+    now = time.time()
+    with _USER_LOCK:
+        _USER_FAILS_ALL[:] = [t for t in _USER_FAILS_ALL if now - t < _USER_WINDOW_S]
+        count, start = _USER_FAILS.get(ckey, [0, now])
+        if now - start > _USER_WINDOW_S:
+            count, start = 0, now
+        if count >= _USER_MAX_FAILS or len(_USER_FAILS_ALL) >= _USER_MAX_FAILS_ALL:
+            return (*locked, "⏳ Too many attempts — try again later.")
+    if sid and hmac.compare_digest(str(password or "").encode("utf-8"), APP_PASSWORD.encode("utf-8")):
+        with _USER_LOCK:
+            _USER_SESSIONS.add(sid)
+            _USER_FAILS.pop(ckey, None)
+        return gr.update(visible=False), gr.update(visible=True), ""
+    with _USER_LOCK:
+        _USER_FAILS[ckey] = [count + 1, start]
+        _USER_FAILS_ALL.append(now)
+        if len(_USER_FAILS) > _USER_FAILS_MAX_KEYS:
+            for k in [k for k, v in _USER_FAILS.items() if now - v[1] > _USER_WINDOW_S]:
+                _USER_FAILS.pop(k, None)
+    time.sleep(1.0)
+    return (*locked, "❌ Wrong password")
+
+
+def end_session(request: gr.Request):
+    sid = _session_id(request)
+    if not sid:
+        return
+    with _USER_LOCK:
+        _USER_SESSIONS.discard(sid)
+    _ADMIN_SESSIONS.discard(sid)
+    _ADMIN_URL_OK_SESSIONS.discard(sid)
+
+
+_BROWSER_SESSIONS = {}
+_BROWSER_SESSIONS_MAX = 20000
+
+
+def mark_browser_session(request: gr.Request):
+    sid = _session_id(request)
+    if not sid:
+        return
+    with _RATE_LOCK:
+        _BROWSER_SESSIONS[sid] = time.time()
+        while len(_BROWSER_SESSIONS) > _BROWSER_SESSIONS_MAX:
+            _BROWSER_SESSIONS.pop(next(iter(_BROWSER_SESSIONS)), None)
+
+
+RATE_MAX_RUNS = int(os.environ.get("RATE_MAX_RUNS", "20"))
+RATE_WINDOW_S = int(os.environ.get("RATE_WINDOW_S", "600"))
+PENALTY_AFTER_BLOCKED = int(os.environ.get("PENALTY_AFTER_BLOCKED", "3"))
+LOCKOUT_STEPS_S = (300, 1800, 7200, 43200)
+OFFENCE_FORGET_S = 86400
+SPAM_WHILE_LOCKED_N = int(os.environ.get("SPAM_WHILE_LOCKED_N", "3"))
+SPAM_EXTRA_LOCK_S = int(os.environ.get("SPAM_EXTRA_LOCK_S", "300"))
+LOCKOUT_MAX_S = int(os.environ.get("LOCKOUT_MAX_S", "86400"))
+BUSY_FREE = int(os.environ.get("BUSY_FREE", "1"))
+BUSY_WINDOW_S = 60
+_RATE_LOCK = threading.Lock()
+_RATE_HITS = {}
+_RATE_BUSY = set()
+_RATE_PENALTY = {}
+_RATE_MAX_KEYS = 5000
+_RATE_STATS = {
+    "allowed": 0,
+    "blocked_not_browser": 0,
+    "blocked_busy": 0,
+    "blocked_rate": 0,
+    "blocked_locked": 0,
+    "lockouts": 0,
+    "lockouts_extended": 0,
+    "since": time.time(),
+}
+
+
+TRUSTED_PROXY_HOPS = max(1, int(os.environ.get("TRUSTED_PROXY_HOPS", "1")))
+
+
+def _normalize_ip(ip):
+    try:
+        addr = ipaddress.ip_address(str(ip).strip().strip("[]"))
+    except ValueError:
+        return ip
+    if addr.version == 6:
+        if addr.ipv4_mapped:
+            return str(addr.ipv4_mapped)
+        return str(ipaddress.ip_network(f"{addr}/64", strict=False))
+    return str(addr)
+
+
+def _pick_forwarded(fwd):
+    parts = [p.strip() for p in fwd.split(",") if p.strip()]
+    if not parts:
+        return None
+    pick = parts[-TRUSTED_PROXY_HOPS] if len(parts) >= TRUSTED_PROXY_HOPS else parts[0]
+    return _normalize_ip(pick)
+
+
+_IP_DEBUG_LEFT = int(os.environ.get("IP_DEBUG_LOG", "0"))
+
+
+def _debug_ip(fwd, host, chosen):
+    global _IP_DEBUG_LEFT
+    if _IP_DEBUG_LEFT <= 0:
+        return
+    _IP_DEBUG_LEFT -= 1
+    print(f"[ip-debug] x-forwarded-for={fwd!r} client.host={host!r} hops={TRUSTED_PROXY_HOPS} -> key={chosen!r}")
+
+
+def _client_key(request):
+    if request is None:
+        return "?"
+    host = getattr(getattr(request, "client", None), "host", None)
+    try:
+        fwd = request.headers.get("x-forwarded-for", "")
+        ip = _pick_forwarded(fwd)
+        if ip:
+            _debug_ip(fwd, host, ip)
+            return ip
+    except Exception:
+        pass
+    if host:
+        return _normalize_ip(host)
+    return _session_id(request) or "?"
+
+
+def _fmt_wait(seconds):
+    seconds = max(1, int(seconds))
+    if seconds >= 3600:
+        return f"{(seconds + 3599) // 3600} h"
+    return f"{(seconds + 59) // 60} min"
+
+
+class BlockMsg(str):
+    pass
+
+
+def _paused_msg(seconds):
+    return BlockMsg(f"🚫 Too many requests - this connection is paused. "
+                    f"Try again in about {_fmt_wait(seconds)}.")
+
+
+def show_rate_message(msg):
+    if isinstance(msg, BlockMsg):
+        raise gr.Error(msg)
+    gr.Warning(msg)
+
+
+def rate_acquire(request):
+    if is_admin(request):
+        return True, None, ""
+    key = _client_key(request)
+    sid = _session_id(request)
+    now = time.time()
+    with _RATE_LOCK:
+        if not sid or sid not in _BROWSER_SESSIONS:
+            _RATE_STATS["blocked_not_browser"] += 1
+            return False, None, "Please refresh the page and try again."
+
+        pen = _RATE_PENALTY.get(key)
+        if pen is not None and pen["level"] and now - pen["last"] > OFFENCE_FORGET_S and now >= pen["until"]:
+            _RATE_PENALTY.pop(key, None)
+            pen = None
+
+        if pen is not None and now < pen["until"]:
+            _RATE_STATS["blocked_locked"] += 1
+
+
+            pen["spam"] = pen.get("spam", 0) + 1
+            pen["last"] = now
+            if pen["spam"] >= SPAM_WHILE_LOCKED_N:
+                pen["spam"] = 0
+                pen["until"] = min(pen["until"] + SPAM_EXTRA_LOCK_S, now + LOCKOUT_MAX_S)
+                _RATE_STATS["lockouts_extended"] += 1
+                print(f"[rate] lock extended by {SPAM_EXTRA_LOCK_S}s on {key} (spam while locked)")
+            return False, None, _paused_msg(pen["until"] - now)
+
+        hits = [t for t in _RATE_HITS.get(key, []) if now - t < RATE_WINDOW_S]
+
+        if key in _RATE_BUSY:
+            _RATE_HITS[key] = hits
+            _RATE_STATS["blocked_busy"] += 1
+            pen = _RATE_PENALTY.setdefault(key, {"blocked": 0, "level": 0, "until": 0.0, "last": now})
+            busy = [t for t in pen.get("busy", []) if now - t < BUSY_WINDOW_S]
+            busy.append(now)
+            pen["busy"] = busy
+            if len(busy) > BUSY_FREE:
+                duration = LOCKOUT_STEPS_S[min(pen["level"], len(LOCKOUT_STEPS_S) - 1)]
+                pen["level"] += 1
+                pen["blocked"] = 0
+                pen["busy"] = []
+                pen["until"] = now + duration
+                pen["spam"] = 0
+                pen["last"] = now
+                _RATE_STATS["lockouts"] += 1
+                print(f"[rate] lockout #{pen['level']} for {duration}s on {key} (busy spam)")
+                return False, None, _paused_msg(duration)
+            return False, None, ("⚠️ You already have a run in progress - please wait for it to finish. "
+                                 "Warning: pressing Run again will pause this connection.")
+
+        if len(hits) >= RATE_MAX_RUNS:
+            _RATE_HITS[key] = hits
+            _RATE_STATS["blocked_rate"] += 1
+            pen = _RATE_PENALTY.setdefault(key, {"blocked": 0, "level": 0, "until": 0.0, "last": now})
+            pen["blocked"] += 1
+            if pen["blocked"] >= PENALTY_AFTER_BLOCKED:
+                duration = LOCKOUT_STEPS_S[min(pen["level"], len(LOCKOUT_STEPS_S) - 1)]
+                pen["level"] += 1
+                pen["blocked"] = 0
+                pen["until"] = now + duration
+                pen["spam"] = 0
+                pen["last"] = now
+                _RATE_STATS["lockouts"] += 1
+                print(f"[rate] lockout #{pen['level']} for {duration}s on {key}")
+                return False, None, _paused_msg(duration)
+            wait_min = max(1, int(RATE_WINDOW_S - (now - hits[0])) // 60 + 1)
+            return False, None, (f"⚠️ Rate limit reached ({RATE_MAX_RUNS} runs per "
+                                 f"{max(1, RATE_WINDOW_S // 60)} min). Try again in about {wait_min} min. "
+                                 f"Warning {pen['blocked']}/{PENALTY_AFTER_BLOCKED - 1}: "
+                                 f"pressing Run again will pause this connection.")
+
+        hits.append(now)
+        _RATE_HITS[key] = hits
+        _RATE_BUSY.add(key)
+        _RATE_STATS["allowed"] += 1
+        if len(_RATE_HITS) > _RATE_MAX_KEYS:
+            for k in [k for k, v in _RATE_HITS.items() if not v or now - v[-1] >= RATE_WINDOW_S]:
+                _RATE_HITS.pop(k, None)
+        if len(_RATE_PENALTY) > _RATE_MAX_KEYS:
+            for k in [k for k, v in _RATE_PENALTY.items()
+                      if now >= v["until"] and now - v["last"] > OFFENCE_FORGET_S]:
+                _RATE_PENALTY.pop(k, None)
+    return True, key, ""
+
+
+def rate_release(key, refund=False):
+    if key is None:
+        return
+    with _RATE_LOCK:
+        _RATE_BUSY.discard(key)
+        if refund:
+            hits = _RATE_HITS.get(key)
+            if hits:
+                hits.pop()
+                _RATE_STATS["allowed"] = max(0, _RATE_STATS["allowed"] - 1)
+
+
+def is_quota_error(exc):
+    return "quota" in str(exc).lower()
+
+
+FLOOD_MAX_JOINS = int(os.environ.get("FLOOD_MAX_JOINS", "8"))
+FLOOD_MAX_UPLOADS = int(os.environ.get("FLOOD_MAX_UPLOADS", "4"))
+FLOOD_WINDOW_S = int(os.environ.get("FLOOD_WINDOW_S", "10"))
+FLOOD_BAN_STEPS = (60, 300, 1800, 7200, 43200)
+FLOOD_STRIKE_FORGET_S = 86400
+FLOOD_ALLOW_IPS = {_normalize_ip(x.strip()) for x in os.environ.get("FLOOD_ALLOW_IPS", "").split(",") if x.strip()}
+_FLOOD_JOIN_PATHS = ("/queue/join", "/call/", "/run/")
+_FLOOD_UPLOAD_PATHS = ("/upload",)
+_FLOOD_LOCK = threading.Lock()
+_FLOOD_HITS = {}
+_FLOOD_BANNED = {}
+_FLOOD_STRIKES = {}
+_FLOOD_SPAM = {}
+_FLOOD_STATS = {"bans": 0, "extended": 0}
+
+
+def _scope_ip(scope):
+    try:
+        fwd = next((v.decode("latin-1") for k, v in scope.get("headers") or []
+                    if k == b"x-forwarded-for"), "")
+        ip = _pick_forwarded(fwd)
+        if ip:
+            return ip
+    except Exception:
+        pass
+    client = scope.get("client")
+    return _normalize_ip(client[0]) if client else "?"
+
+
+def _flood_retry_after(ip, bucket="join", limit=None):
+    if ip in FLOOD_ALLOW_IPS:
+        return 0
+    if limit is None:
+        limit = FLOOD_MAX_JOINS
+    now = time.time()
+    with _FLOOD_LOCK:
+        until = _FLOOD_BANNED.get(ip, 0.0)
+        if now < until:
+            spam = _FLOOD_SPAM.get(ip, 0) + 1
+            if spam >= SPAM_WHILE_LOCKED_N:
+                spam = 0
+                until = min(until + SPAM_EXTRA_LOCK_S, now + LOCKOUT_MAX_S)
+                _FLOOD_BANNED[ip] = until
+                _FLOOD_STATS["extended"] += 1
+                print(f"[flood] {ip} kept sending while banned - ban extended by {SPAM_EXTRA_LOCK_S}s")
+            _FLOOD_SPAM[ip] = spam
+            return int(until - now) + 1
+        if until:
+            _FLOOD_BANNED.pop(ip, None)
+            _FLOOD_SPAM.pop(ip, None)
+        hkey = (bucket, ip)
+        hits = [t for t in _FLOOD_HITS.get(hkey, []) if now - t < FLOOD_WINDOW_S]
+        hits.append(now)
+        if len(hits) > limit:
+            strikes, last = _FLOOD_STRIKES.get(ip, (0, 0.0))
+            if now - last > FLOOD_STRIKE_FORGET_S:
+                strikes = 0
+            ban = FLOOD_BAN_STEPS[min(strikes, len(FLOOD_BAN_STEPS) - 1)]
+            _FLOOD_STRIKES[ip] = (strikes + 1, now)
+            _FLOOD_BANNED[ip] = now + ban
+            _FLOOD_SPAM.pop(ip, None)
+            _FLOOD_STATS["bans"] += 1
+            _FLOOD_HITS.pop(hkey, None)
+            print(f"[flood] {ip} sent {len(hits)} {bucket} requests in {FLOOD_WINDOW_S}s - "
+                  f"blocked for {ban}s (offence #{strikes + 1})")
+            return ban
+        _FLOOD_HITS[hkey] = hits
+        if len(_FLOOD_HITS) > _RATE_MAX_KEYS:
+            for k in [k for k, v in _FLOOD_HITS.items() if not v or now - v[-1] >= FLOOD_WINDOW_S]:
+                _FLOOD_HITS.pop(k, None)
+        if len(_FLOOD_BANNED) > _RATE_MAX_KEYS:
+            for k in [k for k, v in _FLOOD_BANNED.items() if now >= v]:
+                _FLOOD_BANNED.pop(k, None)
+            for k in [k for k in _FLOOD_SPAM if k not in _FLOOD_BANNED]:
+                _FLOOD_SPAM.pop(k, None)
+        if len(_FLOOD_STRIKES) > _RATE_MAX_KEYS:
+            for k in [k for k, v in _FLOOD_STRIKES.items() if now - v[1] > FLOOD_STRIKE_FORGET_S]:
+                _FLOOD_STRIKES.pop(k, None)
+    return 0
+
+
+class FloodGuard:
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        bucket, limit = None, 0
+        if scope.get("type") == "http" and scope.get("method") == "POST":
+            path = scope.get("path", "")
+            if any(p in path for p in _FLOOD_UPLOAD_PATHS):
+                bucket, limit = "upload", FLOOD_MAX_UPLOADS
+            elif any(p in path for p in _FLOOD_JOIN_PATHS):
+                bucket, limit = "join", FLOOD_MAX_JOINS
+        if bucket:
+            try:
+                wait = _flood_retry_after(_scope_ip(scope), bucket, limit)
+            except Exception:
+                wait = 0
+            if wait:
+                body = json.dumps({"error": f"Too many requests. Try again in {wait} s."}).encode()
+                await send({"type": "http.response.start", "status": 429, "headers": [
+                    (b"content-type", b"application/json"),
+                    (b"retry-after", str(wait).encode()),
+                    (b"content-length", str(len(body)).encode()),
+                ]})
+                await send({"type": "http.response.body", "body": body})
+                return
+        await self.app(scope, receive, send)
+
+
+def _mask_key(key):
+    parts = str(key).split(".")
+    if len(parts) == 4:
+        return f"{parts[0]}.{parts[1]}.*.*"
+    return str(key)[:9] + "..."
+
+
+def load_rate_stats(request: gr.Request):
+    if not is_admin(request):
+        return ""
+    now = time.time()
+    with _RATE_LOCK:
+        st = dict(_RATE_STATS)
+        locked = [(k, v["until"] - now, v["level"]) for k, v in _RATE_PENALTY.items() if v["until"] > now]
+        sessions = len(_BROWSER_SESSIONS)
+    with _FLOOD_LOCK:
+        flood_bans = _FLOOD_STATS["bans"]
+        flood_ext = _FLOOD_STATS["extended"]
+        flood_now = sum(1 for v in _FLOOD_BANNED.values() if v > now)
+    hours = max(0.01, (now - st["since"]) / 3600)
+    lines = [
+        f"**Protection stats** (since last restart, {hours:.1f} h)",
+        "",
+        f"- Runs allowed: {st['allowed']}",
+        f"- Blocked - not a browser session: {st['blocked_not_browser']}",
+        f"- Blocked - over run limit: {st['blocked_rate']}",
+        f"- Blocked - run already in progress: {st['blocked_busy']}",
+        f"- Blocked - locked-out IP: {st['blocked_locked']}",
+        f"- Lockouts issued: {st['lockouts']}",
+        f"- Lockouts extended (spam while locked): {st['lockouts_extended']}",
+        f"- Browser sessions tracked: {sessions}",
+        f"- IPs locked right now: {len(locked)}",
+        f"- Flood-guard bans issued: {flood_bans}",
+        f"- Flood-guard bans extended (spam while banned): {flood_ext}",
+        f"- IPs flood-banned right now: {flood_now}",
+    ]
+    for k, remaining, level in sorted(locked, key=lambda x: -x[1])[:10]:
+        lines.append(f"  - `{_mask_key(k)}` - level {level}, {_fmt_wait(remaining)} left")
+    return "\n".join(lines)