modes.py

"""Mode handlers — pure functions over QwenImageEditPlusPipeline + params dict."""

from __future__ import annotations

import random
import threading
from typing import Any

from PIL import Image

import models

# Serialize MPS inference: one GPU, one shared pipeline whose scheduler/adapter are
# swapped per request. Two concurrent runs would race that state + corrupt the
# memory peak measurement + race the calibration cache. (The app also caps queue
# concurrency to 1; this is correct-by-construction belt-and-suspenders.)
_GPU_LOCK = threading.Lock()


def _is_gpu_oom(e: BaseException) -> bool:
    """True for any GPU allocation-failure RuntimeError (MPS or CUDA).

    Covers the raw allocator OOM ("MPS backend out of memory" / "CUDA out of memory"),
    the MPS high-watermark throw (a DIFFERENT message), and CUDA OOM variants — otherwise
    those would escape the OOM-retry net and crash, defeating the never-OOM guarantee.
    """
    s = str(e).lower()
    return any(
        k in s
        for k in (
            "out of memory", "watermark", "mps allocated", "cannot allocate",
            "insufficient memory", "cuda error", "cublas",
        )
    )


def _gpu_empty_cache(torch: Any, dt: str) -> None:
    if dt == "mps":
        torch.mps.empty_cache()
    elif dt == "cuda":
        torch.cuda.empty_cache()


def _gpu_synchronize(torch: Any, dt: str) -> None:
    if dt == "mps":
        torch.mps.synchronize()
    elif dt == "cuda":
        torch.cuda.synchronize()


def _gpu_allocated_gb(torch: Any, dt: str) -> float:
    """Currently-allocated GPU memory (GB) — the resident baseline before inference."""
    if dt == "mps":
        return torch.mps.driver_allocated_memory() / (1024**3)
    if dt == "cuda":
        return torch.cuda.memory_allocated() / (1024**3)
    return 0.0


def _gpu_peak_gb(torch: Any, dt: str, baseline_gb: float) -> float:
    """Activation peak (GB) above the resident baseline, for calibration."""
    if dt == "mps":
        return max(0.0, torch.mps.driver_allocated_memory() / (1024**3) - baseline_gb)
    if dt == "cuda":
        return max(0.0, torch.cuda.max_memory_allocated() / (1024**3) - baseline_gb)
    return 0.0


def _apply_speed(pipe: Any, speed: str) -> None:
    """Configure the pipeline for Fast (Lightning LoRA) or Quality (default) mode.

    Fast:    swap in the Lightning scheduler + enable the Lightning LoRA adapter.
    Quality: restore the default scheduler + disable the LoRA adapter entirely.
    Both operations are cheap, reversible, and ZeroGPU-safe (no weight fusion/unfusion).
    """
    if speed == "Fast":
        pipe.scheduler = pipe._qie_lightning_scheduler
        # enable_lora() FIRST: a prior Quality run's disable_lora() leaves every adapter
        # layer with _disable_adapters=True, and set_adapters() only flips the *active*
        # adapter name — it does NOT re-enable the layers. Without this, Lightning is
        # silently bypassed on the long-lived MPS process (ZeroGPU hides it: each call
        # re-forks a fresh pipeline). enable_lora() is a harmless no-op when already enabled.
        pipe.enable_lora()
        pipe.set_adapters([models.LORA_ADAPTER_NAME], [1.0])
    else:
        pipe.scheduler = pipe._qie_default_scheduler
        pipe.disable_lora()


_USER_LORA_ADAPTER = "user"


def _delete_user_lora(pipe: Any) -> None:
    """Remove the user LoRA adapter (best-effort). MANDATORY on MPS: it's a single
    long-lived process with NO per-call re-fork, so a leftover adapter would silently leak
    into the next request. (On ZeroGPU the fork re-forks clean, but cleanup is harmless.)
    """
    try:
        pipe.delete_adapters(_USER_LORA_ADAPTER)
    except Exception:
        pass


def _step_callback(progress: Any, total_steps: int) -> Any:
    """Build a diffusers ``callback_on_step_end`` that drives a ``gr.Progress`` with clear
    labels ("Generating — step N/M"). Returns None when there's no progress object (CI /
    non-UI calls), so the pipeline call is unchanged off the UI path.

    Using an explicit step callback (instead of Gradio's ``track_tqdm``) avoids the
    confusing "Downloading (incomplete total…) 0/0 B" placeholder Gradio renders for the
    pipeline's internal, total-less tqdm bars during the input-encode phase.
    """
    if progress is None:
        return None

    def _cb(_pipe: Any, step: int, _timestep: Any, callback_kwargs: dict) -> dict:
        progress((step + 1) / max(1, total_steps), desc=f"Generating — step {step + 1}/{total_steps}")
        return callback_kwargs

    return _cb


def _run(pipe: Any, params: dict[str, Any], progress: Any = None) -> tuple[Image.Image, dict[str, Any]]:
    """Run inference and return (output_image, metadata).

    Validates that at least one image is present, applies the speed mode, resolves
    dimensions and seed, then calls the pipeline. On MPS an OOM preflight runs first:
    it sizes the request against a live memory budget and auto-degrades down a ladder
    (resolution -> refs -> drop CFG double-pass -> Quality->Fast) so inference never
    OOMs; the math + any degradation are surfaced in meta. The CUDA/CPU path is
    unchanged. ``progress`` (optional gr.Progress) drives a clean step bar.
    """
    images: list[Any] = params["images"]
    if not images:
        raise ValueError("at least one image is required; params['images'] is empty")

    import torch  # deferred — CI has no torch installed

    # Normalize device: pipe.device may be a str or a torch.device object.
    device = str(getattr(pipe, "device", "cpu"))
    device_type = device.split(":")[0]
    # ZeroGPU quirk: after the spaces pack/restore, pipe.device can read "cpu" at the
    # START of the call even though execution happens on a real CUDA GPU inside the
    # @spaces.GPU fork. Trust the Space environment over the (stale) device property so
    # the CUDA activation preflight + never-OOM degrade path actually run on ZeroGPU
    # (otherwise the simple/no-budget path is taken and a heavy request can OOM-SIGKILL).
    if device_type not in ("mps", "cuda") and models.on_spaces():
        device_type = "cuda"
    is_gpu = device_type in ("mps", "cuda")

    mode = params["mode"]
    speed = params["speed"]
    steps = int(params["steps"])
    true_cfg = float(params["true_cfg"])
    prompt = params["prompt"]
    negative_prompt = params["negative_prompt"]

    seed_in = params["seed"]
    seed = random.randint(0, 2**32 - 1) if seed_in < 0 else int(seed_in)

    # --- CPU / mocked path (no GPU memory budgeting needed) -----------------------------
    if not is_gpu:
        _apply_speed(pipe, speed)
        w, h = models.fit_dimensions(images[0])
        gen = torch.Generator(device).manual_seed(seed)
        if progress is not None:
            progress(0.0, desc="Encoding inputs…")
        out = pipe(
            image=images,
            prompt=prompt,
            negative_prompt=negative_prompt,
            true_cfg_scale=true_cfg,
            num_inference_steps=steps,
            height=h,
            width=w,
            generator=gen,
            callback_on_step_end=_step_callback(progress, steps),
        )
        meta = {
            "mode": mode, "speed": speed, "steps": steps, "true_cfg": true_cfg,
            "seed": seed, "width": w, "height": h, "num_inputs": len(images),
        }
        return out.images[0], meta

    # --- GPU path (MPS or CUDA/ZeroGPU): serialized, OOM-preflight + reactive degrade ----
    # Same activation-centric never-OOM logic on both: the ~58 GB model is already resident,
    # so only the activation must fit the free budget (MPS: free unified RAM capped by the
    # working set; CUDA/ZeroGPU: free VRAM via torch.cuda.mem_get_info). On ZeroGPU xlarge
    # (96 GB) Edit fits but heavy Compose can exceed -> auto-degrade (the user's mandate).
    import gc

    import memory

    with _GPU_LOCK:
        gc.collect()
        _gpu_empty_cache(torch, device_type)
        base_w, base_h = images[0].size
        n_ref0 = max(0, len(images) - 1)
        # User LoRA (Quality-only): a local .safetensors path resolved by the handler (off the
        # GPU clock). Ignored in Fast — the Lightning 4-step distillation fights a content LoRA.
        lora_path = params.get("lora_path")
        _lw = params.get("lora_weight")
        lora_weight = 0.9 if _lw is None else float(_lw)  # explicit None check keeps a valid 0.0
        lora_requested = bool(lora_path) and speed == "Quality"

        plan = memory.plan_request(
            device_type, mode, base_w, base_h, n_ref0, speed, steps, true_cfg, lora_requested
        )
        if plan["refused"]:
            raise RuntimeError(plan["note"])

        w, h = plan["width"], plan["height"]
        steps, true_cfg, speed, n_ref = plan["steps"], plan["true_cfg"], plan["speed"], plan["n_ref"]
        degrades = list(plan["degrades"])

        # The user LoRA applies only if the FINAL (post-degrade) speed is Quality — if the
        # preflight degraded Quality->Fast, drop it (else record_peak/penalize would mis-key
        # the Fast peak into the has_user_lora cache slot). Recomputed AFTER the plan.
        lora_active = bool(lora_path) and speed == "Quality"

        out = None
        last_err: Exception | None = None
        # Tracks whether the "user" adapter actually got loaded — drives the finally cleanup
        # INDEPENDENTLY of lora_active (which can flip False mid-loop on an OOM degrade to Fast).
        # If cleanup keyed off lora_active, a Quality->Fast degrade would leak the loaded adapter.
        user_lora_loaded = False
        try:
            # Load the user LoRA once (adapter "user"), INSIDE the try so the finally below
            # always cleans it up — even if load_lora_weights raises (no adapter leak on the
            # long-lived MPS process). Re-activated per attempt AFTER _apply_speed (Quality's
            # disable_lora() turns all adapters off). delete-before-load clears a stale one.
            if lora_active:
                _delete_user_lora(pipe)
                pipe.load_lora_weights(lora_path, adapter_name=_USER_LORA_ADAPTER)
                user_lora_loaded = True
            for _attempt in range(8):  # bounded reactive degrade — hard never-OOM guarantee
                imgs = images[: n_ref + 1]
                _apply_speed(pipe, speed)
                if lora_active and speed == "Quality":
                    # _apply_speed's disable_lora() (Quality branch) left every adapter
                    # layer disabled; re-enable before activating the user adapter, else
                    # set_adapters() flips the active name but the layers stay bypassed and
                    # the LoRA has ZERO effect (output identical to no-LoRA). See _apply_speed.
                    pipe.enable_lora()
                    pipe.set_adapters([_USER_LORA_ADAPTER], [lora_weight])
                gen = torch.Generator("cpu").manual_seed(seed)  # CPU generator is safe on MPS + CUDA
                baseline_gb = _gpu_allocated_gb(torch, device_type)
                if device_type == "cuda":
                    torch.cuda.reset_peak_memory_stats()
                if progress is not None:
                    progress(0.0, desc="Encoding inputs…")
                try:
                    out = pipe(
                        image=imgs,
                        prompt=prompt,
                        negative_prompt=negative_prompt,
                        true_cfg_scale=true_cfg,
                        num_inference_steps=steps,
                        height=h,
                        width=w,
                        generator=gen,
                        callback_on_step_end=_step_callback(progress, steps),
                    )
                    break
                except RuntimeError as e:
                    if not _is_gpu_oom(e):
                        raise
                    last_err = e
                    del gen
                    gc.collect()
                    _gpu_synchronize(torch, device_type)
                    _gpu_empty_cache(torch, device_type)
                    # Self-correct AFTER cleanup: at the OOM instant the heap is at its sticky
                    # peak, so the budget would read ~0 -> no-op penalty. A clean heap returns the
                    # real budget the config overflowed, so the next request degrades preemptively.
                    memory.penalize(
                        device_type, h, w, n_ref, true_cfg > 1.0,
                        memory.activation_budget_gb(device_type), lora_active,
                    )
                    nxt = memory.step_down(w, h, n_ref, speed, true_cfg, steps, base_w, base_h, mode)
                    if nxt is None:
                        raise RuntimeError(f"GPU OOM at the smallest config and cannot degrade further: {e}") from e
                    w, h, n_ref = nxt["width"], nxt["height"], nxt["n_ref"]
                    speed, true_cfg, steps = nxt["speed"], nxt["true_cfg"], nxt["steps"]
                    # A degrade to Fast drops the (Quality-only) user LoRA: the next attempt's
                    # _apply_speed switches to Lightning and the guard below skips set_adapters.
                    # Recompute so record_peak/meta key the Fast peak correctly (NOT the LoRA slot).
                    lora_active = bool(lora_path) and speed == "Quality"
                    degrades.append(f"OOM-retry->{w}x{h} {speed} {n_ref + 1}img")
            if out is None:  # pragma: no cover - defensive
                raise RuntimeError(f"GPU inference failed after retries: {last_err}")

            peak_gb = _gpu_peak_gb(torch, device_type, baseline_gb)
            memory.record_peak(device_type, mode, h, w, n_ref, true_cfg > 1.0, peak_gb, lora_active)

            meta = {
                "mode": mode, "speed": speed, "steps": steps, "true_cfg": true_cfg,
                "seed": seed, "width": w, "height": h, "num_inputs": n_ref + 1,
                "preflight": plan["note"], "budget_gb": plan["budget_gb"],
                "need_gb": plan["need_gb"], "measured_peak_gb": round(peak_gb, 1),
            }
            if degrades:
                meta["degrades"] = degrades
            if lora_active and speed == "Quality":
                meta["lora"] = {"weight": lora_weight, "file": lora_path.rsplit("/", 1)[-1]}
            return out.images[0], meta
        finally:
            # Cleanup runs even on error — MANDATORY on MPS (long-lived process, no re-fork):
            # a lingering "user" adapter would leak into the next request. Keyed off
            # user_lora_loaded (not lora_active) so an OOM degrade-to-Fast still cleans up.
            if user_lora_loaded:
                _delete_user_lora(pipe)
            gc.collect()
            _gpu_empty_cache(torch, device_type)


def call_edit(pipe: Any, params: dict[str, Any], progress: Any = None) -> tuple[Image.Image, dict[str, Any]]:
    """Edit mode: single input image + instruction -> edited image.

    Expects params["images"] == [target_image].
    """
    p = dict(params)
    p["mode"] = "edit"
    return _run(pipe, p, progress)


def call_compose(pipe: Any, params: dict[str, Any], progress: Any = None) -> tuple[Image.Image, dict[str, Any]]:
    """Compose mode: target image + up to 2 optional reference images -> composed edit.

    None slots in params["images"] are dropped before the pipeline call so the
    pipeline always receives a contiguous list of 1..3 real images.
    """
    p = dict(params)
    p["mode"] = "compose"
    p["images"] = [img for img in params["images"] if img is not None]
    return _run(pipe, p, progress)


DISPATCH: dict[str, Any] = {
    "edit": call_edit,
    "compose": call_compose,
}