sd-webui-progressive-growing/scripts/progressive_growing_always.py

"""sd-webui-progressive-growing  ·  extension-only rewrite
===========================================================

Architecture contract
---------------------
  UI  →  process()  →  patch wrapper  →  version fn(p, plan, ...)  →  samples
                                       ↓ on any conflict/disable
                                    _ORIG_SAMPLE(...)

Files
-----
  This single file is self-contained for easy drop-in.
  Split into lib_progressive/ when the project grows beyond ~600 lines.

Versions
--------
  v1 (exact)     – original code 1:1, kept as reference / regression baseline
  v2 (safe)      – validation, dedup stages, conflict guards, predictable output
  v3 (fast)      – minimal refinement, no VAE decode on intermediate stages
  v4 (balanced)  – refinement only on stages >= BALANCED_REFINE_THRESHOLD (linear scale, not area)
  v5 (latent)    – pure latent upscale, zero refinement
"""

from __future__ import annotations

import contextlib
import math
import gradio as gr
import numpy as np
import torch

from modules import scripts, sd_samplers, devices
from modules import processing as processing_mod
from modules.processing import (
    create_random_tensors,
    decode_latent_batch,
    opt_C,
    opt_f,
)

# ─────────────────────────────────────────────────────────────────────────────
# Constants
# ─────────────────────────────────────────────────────────────────────────────

BALANCED_REFINE_THRESHOLD = 0.70   # v4: only refine stages whose linear scale >= this
ADAPTIVE_REFINE_THRESHOLD = 0.60   # v6: threshold for long-plan (5+ stages) adaptive refinement

LATENT_INTERP_MODES   = ["bicubic", "bilinear", "nearest", "area"]
LATENT_INTERP_DEFAULT = "bicubic"   # best general-purpose default; nearest is fastest

MIN_STAGES_AFTER_DEDUP    = 1      # if dedup collapses plan to 1 size, skip progressive entirely

AUTO_STAGE_JUMP_DEFAULT   = 1.35   # target linear growth factor per stage
AUTO_STAGE_MAX_DEFAULT    = 6      # hard cap on auto-computed stage count

REFINE_STEP_MODES   = ["uniform", "late-heavy", "final-heavy"]
REFINE_STEP_DEFAULT = "uniform"    # matches legacy behaviour; switch to late-heavy once comfortable

# ── Sampler compatibility profiles ───────────────────────────────────────────
#
# Progressive growing runs each stage sampler on a *smaller* latent than the
# final target.  Two classes of problem arise:
#
#   1. Custom sigma schedulers (e.g. "CosineExponential blend") compute sigmas
#      that are spatially shaped – [steps+1, H, W] – by reading p.width/p.height
#      at sampler.sample() time.  With stage dims still at final size the sigmas
#      mismatch the stage latent → RuntimeError on the first sigma operation.
#
#   2. img2img contexts: p.init_latent / p.mask / p.nmask are full-size tensors.
#      KDiffusionSampler.sample() copies them into model_wrap_cfg before
#      launching the sampler loop, causing shape mismatches inside the loop.
#
# The "smea" profile wraps every sampler.sample() / sample_img2img() call with
# _stage_sampler_context(), which temporarily sets p.width/p.height to stage
# dims and rescales init_latent/mask/nmask.  This directly mirrors the
# _Rescaler pattern already used inside sd-webui-smea itself.
#
# SMEA_SAMPLER_PATTERNS – families that need _stage_sampler_context.
# Each entry is a lowercase substring of p.sampler_name.
SMEA_SAMPLER_PATTERNS: list[str] = [
    "euler dy",       # Euler Dy, Euler Dy koishi-star
    "euler smea",     # Euler Smea, Euler Smea Dy, Euler Smea Max, all multi-*
    "euler h max",    # Euler h max a/b/c/…
    "euler max",      # Euler Max, Max1b … Max4f  (catches all Max* variants)
    "kohaku_lonyu",   # Kohaku_LoNyu_Yog
    "tcd",            # TCD / TCD Euler a
]

# INCOMPATIBLE_SAMPLER_PATTERNS – families with no viable compatibility path yet.
# Unlike SMEA, these are blocked entirely until a specific adapter is built.
# Empty for now; add entries here when a truly irreconcilable sampler is found.
INCOMPATIBLE_SAMPLER_PATTERNS: list[str] = []


def _get_sampler_profile(p) -> tuple[str, str]:
    """
    Classify the active sampler into one of three profiles:

      'standard'    – normal k-diffusion; no special handling required.
      'smea'        – sd-webui-smea family; needs _stage_sampler_context.
      'unsupported' – no compatibility path yet; progressive will be skipped.

    Returns (profile_name, reason).  reason is '' unless profile == 'unsupported'.
    """
    name = (getattr(p, 'sampler_name', '') or '').lower()

    for pattern in INCOMPATIBLE_SAMPLER_PATTERNS:
        if pattern in name:
            return 'unsupported', (
                f'disabled – sampler "{p.sampler_name}" has no compatibility '
                f'profile yet. Use a standard k-diffusion sampler or the SMEA family.'
            )

    for pattern in SMEA_SAMPLER_PATTERNS:
        if pattern in name:
            return 'smea', ''

    return 'standard', ''


@contextlib.contextmanager
def _stage_sampler_context(p, stage_w: int, stage_h: int):
    """
    Temporarily adapt p's size-dependent state to stage dimensions.

    Three adaptations, all restored in finally:

    1. p.width / p.height -> stage dims.

    2. p.init_latent / p.mask / p.nmask -> rescaled to stage latent dims.
       KDiffusionSampler.sample() copies these into model_wrap_cfg before
       the sampler loop.

    3. p.sampler_noise_scheduler_override -> wrapped to resize spatial sigmas.
       SMEA's process() hook sets this override with full-size dims captured
       in a closure.  KDiffusionSampler.sample() calls the override DIRECTLY
       (bypassing get_sigmas entirely) when it is set, so any spatial sigma
       tensor [T,H,W] or [T,C,H,W] it returns will have full-size H/W dims.
       We wrap the override here so its output is resized to stage_hw before
       the sampler loop sees it.
    """
    import torch.nn.functional as _F

    stage_hw = (stage_h // opt_f, stage_w // opt_f)

    # save
    orig_w, orig_h   = p.width, p.height
    orig_init        = getattr(p, 'init_latent', None)
    orig_mask        = getattr(p, 'mask', None)
    orig_nmask       = getattr(p, 'nmask', None)
    orig_override    = getattr(p, 'sampler_noise_scheduler_override', None)
    # p.x is used by KDiffusionSampler.initialize() to seed BrownianTreeNoiseSampler.
    # If left at full resolution, noise_sampler returns full-size tensors that mismatch
    # the stage-sized latent x inside the sampler loop.  Resize to stage dims so any
    # noise_sampler built from p.x during initialize() works at the correct scale.
    orig_px          = getattr(p, 'x', None)

    # adapt dims
    p.width, p.height = stage_w, stage_h

    if orig_init is not None:
        p.init_latent = _F.interpolate(orig_init, size=stage_hw, mode='nearest-exact')
    if orig_mask is not None:
        p.mask = _F.interpolate(
            orig_mask.unsqueeze(0), size=stage_hw, mode='nearest-exact'
        ).squeeze(0)
    if orig_nmask is not None:
        p.nmask = _F.interpolate(
            orig_nmask.unsqueeze(0), size=stage_hw, mode='nearest-exact'
        ).squeeze(0)
    if orig_px is not None:
        try:
            p.x = _F.interpolate(orig_px, size=stage_hw, mode='nearest-exact')
        except Exception:
            pass

    # wrap noise scheduler override to resize spatial sigmas to stage dims
    if orig_override is not None:
        def _stage_override(steps, _orig=orig_override, _hw=stage_hw):
            sigs = _orig(steps)
            if not isinstance(sigs, torch.Tensor) or sigs.ndim < 3:
                return sigs
            try:
                if sigs.ndim == 3:      # [T, H, W]
                    sigs = _F.interpolate(
                        sigs.unsqueeze(1), size=_hw, mode='nearest-exact'
                    ).squeeze(1)
                elif sigs.ndim >= 4:    # [T, C, H, W]
                    sigs = _F.interpolate(sigs, size=_hw, mode='nearest-exact')
            except Exception:
                pass
            return sigs
        p.sampler_noise_scheduler_override = _stage_override

    try:
        yield
    finally:
        p.width, p.height = orig_w, orig_h
        if orig_init  is not None: p.init_latent = orig_init
        if orig_mask  is not None: p.mask         = orig_mask
        if orig_nmask is not None: p.nmask         = orig_nmask
        if orig_px    is not None: p.x             = orig_px
        p.sampler_noise_scheduler_override = orig_override



def _resize_stage_sigmas(sigmas, stage_hw):
    """Resize spatial sigma outputs to stage latent dims where possible.

    Supports tensors and lists/tuples of tensors. Leaves plain 1D schedules
    unchanged. Used to re-wrap any scheduler override that hooks may install
    *after* our stage context is entered.
    """
    import torch.nn.functional as _F

    def _resize_tensor(t):
        if not isinstance(t, torch.Tensor):
            return t
        try:
            if t.ndim == 3:      # [T, H, W]
                return _F.interpolate(t.unsqueeze(1), size=stage_hw, mode='nearest-exact').squeeze(1)
            if t.ndim == 4:      # [T, C, H, W]
                return _F.interpolate(t, size=stage_hw, mode='nearest-exact')
            if t.ndim == 2:      # [H, W] per-step tensor
                return _F.interpolate(t.unsqueeze(0).unsqueeze(0), size=stage_hw, mode='nearest-exact').squeeze(0).squeeze(0)
        except Exception:
            return t
        return t

    if isinstance(sigmas, torch.Tensor):
        return _resize_tensor(sigmas)
    if isinstance(sigmas, list):
        return [_resize_stage_sigmas(s, stage_hw) for s in sigmas]
    if isinstance(sigmas, tuple):
        return tuple(_resize_stage_sigmas(s, stage_hw) for s in sigmas)
    return sigmas


def _rewrap_scheduler_override_for_stage(p, stage_w: int, stage_h: int) -> None:
    """Re-wrap the currently active override so its outputs respect stage dims.

    Some hooks can overwrite p.sampler_noise_scheduler_override after
    _stage_sampler_context() has already wrapped the original override. Call
    this after hooks to re-apply stage-aware resizing to the current override.
    """
    override = getattr(p, 'sampler_noise_scheduler_override', None)
    if override is None or getattr(override, '_pg_stage_wrapped', False):
        return

    stage_hw = (stage_h // opt_f, stage_w // opt_f)

    def _wrapped(steps, _orig=override, _hw=stage_hw):
        return _resize_stage_sigmas(_orig(steps), _hw)

    _wrapped._pg_stage_wrapped = True
    p.sampler_noise_scheduler_override = _wrapped


def _patch_sampler_func_sigmas_to_1d(p) -> None:
    """Wrap p.sampler.func so sigmas reaching SMEA sampler are always 1D.

    Handles tensors, lists/tuples of tensors, and numpy-like values. This is a
    fallback shim for custom schedulers that return spatial sigmas in forms not
    consumed safely by SMEA Euler Max* samplers.
    """
    sampler = getattr(p, 'sampler', None)
    if sampler is None:
        return
    orig_func = getattr(sampler, 'func', None)
    if orig_func is None or getattr(orig_func, '_pg_sigmas_wrapped', False):
        return

    def _collapse_any(sigmas, like: torch.Tensor):
        dev = like.device
        dt = like.dtype

        def _collapse_tensor(t: torch.Tensor) -> torch.Tensor:
            t = t.to(device=dev, dtype=dt)
            if t.ndim == 0:
                return t.reshape(1)
            if t.ndim == 1:
                return t
            return t.mean(dim=tuple(range(1, t.ndim))).flatten()

        if isinstance(sigmas, torch.Tensor):
            return _collapse_tensor(sigmas)
        if isinstance(sigmas, (list, tuple)):
            vals = []
            for s in sigmas:
                if isinstance(s, torch.Tensor):
                    vals.append(_collapse_tensor(s).mean().reshape(()))
                else:
                    try:
                        vals.append(torch.tensor(float(s), device=dev, dtype=dt).reshape(()))
                    except Exception:
                        return sigmas
            return torch.stack(vals)
        try:
            t = torch.as_tensor(sigmas, device=dev, dtype=dt)
            return _collapse_tensor(t)
        except Exception:
            return sigmas

    def _safe_func(model, x, sigmas, *args, **kwargs):
        return orig_func(model, x, _collapse_any(sigmas, x), *args, **kwargs)

    _safe_func._pg_sigmas_wrapped = True
    sampler.func = _safe_func


# ─────────────────────────────────────────────────────────────────────────────
# Stage planner
# ─────────────────────────────────────────────────────────────────────────────

def _snap(v: float, factor: int = opt_f) -> int:
    """Round v down to nearest multiple of factor (min = factor)."""
    v_int = max(factor, int(v))
    return max(factor, (v_int // factor) * factor)


class StagePlan:
    """Validated, deduplicated list of (width, height) latent sizes."""

    def __init__(self, sizes: list[tuple[int, int]], final_w: int, final_h: int):
        self.sizes    = sizes       # [(w, h), ...], last entry == (final_w, final_h)
        self.final_w  = final_w
        self.final_h  = final_h
        self.n_stages = len(sizes)

    def __repr__(self) -> str:
        parts = [f"{w}×{h}" for w, h in self.sizes]
        return " → ".join(parts)


def _auto_n_steps(min_s: float, max_s: float,
                  target_jump: float, max_stages: int) -> int:
    """
    Compute stage count so each step grows the linear scale by ~target_jump.

    Formula:  n = 1 + ceil( log(max_s / min_s) / log(target_jump) )
    Clamped to [2, max_stages].

    Examples with target_jump=1.35:
      0.25 → 1.0  : ≈ 6 stages
      0.50 → 1.0  : ≈ 3 stages
      0.75 → 1.0  : ≈ 2 stages
    """
    growth  = max_s / max(min_s, 1e-6)
    n_steps = 1 + math.ceil(math.log(max(growth, 1.0)) / math.log(max(target_jump, 1.001)))
    return max(2, min(max_stages, n_steps))


def build_plan(p) -> StagePlan:
    """
    Build a deduplicated stage plan from p's progressive-growing attributes.

    When p.progressive_growing_auto_stages is True the stage count is computed
    automatically from min/max scale and p.progressive_growing_auto_jump
    (target linear growth per step), capped at p.progressive_growing_auto_max.
    The manual Stages slider is ignored in auto mode.

    Guarantees:
      - min_scale <= max_scale (swapped silently)
      - last stage == (p.width, p.height) snapped to opt_f
      - duplicate sizes removed
      - if only one unique size remains: plan.n_stages == 1
        (caller should fall back to normal sampling)
    """
    min_s = float(getattr(p, 'progressive_growing_min_scale', 0.25))
    max_s = float(getattr(p, 'progressive_growing_max_scale', 1.0))

    if min_s > max_s:
        min_s, max_s = max_s, min_s

    auto = bool(getattr(p, 'progressive_growing_auto_stages', False))
    if auto:
        target_jump = float(getattr(p, 'progressive_growing_auto_jump', AUTO_STAGE_JUMP_DEFAULT))
        max_stages  = int(getattr(p, 'progressive_growing_auto_max',  AUTO_STAGE_MAX_DEFAULT))
        n_steps = _auto_n_steps(min_s, max_s, target_jump, max_stages)
    else:
        n_steps = max(2, int(getattr(p, 'progressive_growing_steps', 4)))

    scales  = np.linspace(min_s, max_s, n_steps)
    final_w = _snap(p.width)
    final_h = _snap(p.height)

    raw_sizes: list[tuple[int, int]] = []
    for s in scales:
        raw_sizes.append((_snap(p.width * s), _snap(p.height * s)))

    # deduplicate while preserving order
    seen: set[tuple[int, int]] = set()
    unique: list[tuple[int, int]] = []
    for sz in raw_sizes:
        if sz not in seen:
            seen.add(sz)
            unique.append(sz)

    # force the last entry to be the true final size
    if not unique or unique[-1] != (final_w, final_h):
        if (final_w, final_h) in seen:
            unique = [sz for sz in unique if sz != (final_w, final_h)]
        unique.append((final_w, final_h))

    plan           = StagePlan(unique, final_w, final_h)
    plan.auto_mode = auto          # carry metadata for infotext
    plan.requested = n_steps       # stages before dedup
    return plan


# ─────────────────────────────────────────────────────────────────────────────
# Conflict / capability guards
# ─────────────────────────────────────────────────────────────────────────────

def _should_use_progressive(p):
    """
    Return (True, '', plan) if progressive should run, else (False, reason, None).
    Building the plan here avoids a second build_plan() call inside the version fn.
    Reasons appear in infotext via extra_generation_params.
    """
    if not getattr(p, 'enable_progressive_growing', False):
        return False, '', None

    if getattr(p, 'enable_hr', False):
        return False, 'disabled – incompatible with Hires. fix', None

    profile, reason = _get_sampler_profile(p)
    if profile == 'unsupported':
        return False, reason, None

    plan = build_plan(p)
    if plan.n_stages <= MIN_STAGES_AFTER_DEDUP:
        return False, f'disabled – all stages collapsed to one size ({plan.final_w}×{plan.final_h})', None

    return True, '', plan


def _write_params(p, plan: StagePlan, refine_policy: str, refine_step_policy: str) -> None:
    """
    Write run parameters into infotext.

    refine_policy       – effective policy for *which* stages refine
                          e.g. 'all stages', 'final stage only', 'none'
    refine_step_policy  – effective step budget policy actually used
                          e.g. 'uniform', 'late-heavy (v1 fixed)', 'none'
                          Callers must pass the *actual* behaviour, not the
                          UI dropdown value, so v1 / v5 / checkbox-off are honest.

    If the refinement checkbox is off, both refine fields collapse to 'off (checkbox)'.
    """
    refinement_on = getattr(p, 'progressive_growing_refinement', True)
    try:
        ep = p.extra_generation_params
        ep['PG']              = getattr(p, 'progressive_growing_version', 'v2 (safe)')
        ep['PG plan']         = str(plan)
        ep['PG stages']       = plan.n_stages
        ep['PG refine']       = refine_policy       if refinement_on else 'off (checkbox)'
        ep['PG refine steps'] = refine_step_policy  if refinement_on else 'off (checkbox)'
        ep['PG interp']       = getattr(p, 'progressive_growing_interp_mode', LATENT_INTERP_DEFAULT)
        # auto-stage metadata
        if getattr(plan, 'auto_mode', False):
            ep['PG stages mode'] = 'auto'
            ep['PG target jump'] = getattr(p, 'progressive_growing_auto_jump', AUTO_STAGE_JUMP_DEFAULT)
            ep['PG requested']   = getattr(plan, 'requested', plan.n_stages)
            ep['PG actual']      = plan.n_stages
        else:
            ep['PG stages mode'] = 'manual'
    except Exception:
        pass


# ─────────────────────────────────────────────────────────────────────────────
# Shared helpers
# ─────────────────────────────────────────────────────────────────────────────

def _initial_latent(p, w: int, h: int, seeds, subseeds, subseed_strength):
    return create_random_tensors(
        (opt_C, h // opt_f, w // opt_f),
        seeds,
        subseeds=subseeds,
        subseed_strength=subseed_strength,
        seed_resize_from_h=p.seed_resize_from_h,
        seed_resize_from_w=p.seed_resize_from_w,
        p=p,
    )


def _upscale_latent(samples: torch.Tensor, w: int, h: int, interp_mode: str = LATENT_INTERP_DEFAULT) -> torch.Tensor:
    """
    Upscale latent to (h // opt_f, w // opt_f) using the requested interpolation.

    Supported modes mirror torch.nn.functional.interpolate:
      bicubic  – smooth, best for photographic content (default)
      bilinear – slightly faster, softer result
      nearest  – fastest, hard edges; useful for pixel art / very structured images
      area     – anti-aliased average pooling; suited for large downscales if
                 you manually set min_scale > max_scale outside build_plan()
                 (build_plan() itself always swaps them, so area rarely fires
                 in normal use — kept for completeness)

    align_corners=False for bicubic/bilinear matches PyTorch convention and
    the behaviour of v1 (exact), avoiding edge-pixel drift on upscale.
    """
    align = interp_mode in ("bicubic", "bilinear")
    return torch.nn.functional.interpolate(
        samples,
        size=(h // opt_f, w // opt_f),
        mode=interp_mode,
        align_corners=False if align else None,
    )


def _make_noise(p, shape, seeds, subseeds, subseed_strength):
    return create_random_tensors(
        shape,
        seeds,
        subseeds=subseeds,
        subseed_strength=subseed_strength,
        seed_resize_from_h=p.seed_resize_from_h,
        seed_resize_from_w=p.seed_resize_from_w,
        p=p,
    )


def _call_script_hooks(p, samples) -> None:
    """
    Fire process_before_every_sampling on all registered scripts, if available.

    In stock A1111 this hook lets ControlNet, ADetailer, and other extensions
    inject per-sampling adjustments (e.g. attention maps, masks).  Calling it
    before every sampler.sample / sampler.sample_img2img invocation keeps
    progressive growing compatible with those extensions.

    Guarded so the extension degrades gracefully on forks that lack the hook.
    """
    scripts_obj = getattr(p, 'scripts', None)
    if scripts_obj is None:
        return
    hook = getattr(scripts_obj, 'process_before_every_sampling', None)
    if hook is None:
        return
    try:
        hook(p, samples)
    except Exception:
        pass


def _get_refine_steps(p, refine_idx: int, n_refine: int, mode: str) -> int:
    """
    Compute the step budget for a single refinement pass.

    Parameters
    ----------
    p           : processing object, provides p.steps (total generation steps)
    refine_idx  : 0-based index among refine passes that will actually run
                  (0 = earliest/smallest stage being refined)
    n_refine    : total number of refine passes that will run this generation
    mode        : one of REFINE_STEP_MODES

    Modes
    -----
    uniform      current behaviour: p.steps // total_refine_stages, equal for all
    late-heavy   weights grow towards the final pass.  Weights for n passes:
                   n=1 → [1.0]
                   n=2 → [0.4, 0.6]
                   n=3 → [0.2, 0.3, 0.5]
                   n=4+ → geometric series r=1.5 normalised to 1.0
    final-heavy  almost all budget to the last pass:
                   non-final passes each get floor(p.steps * 0.08)
                   final pass gets whatever remains, minimum 1
    """
    total = max(1, p.steps)
    n     = max(1, n_refine)
    idx   = max(0, min(refine_idx, n - 1))

    if mode == 'final-heavy':
        if idx < n - 1:
            return max(1, int(total * 0.08))
        non_final_total = max(1, int(total * 0.08)) * (n - 1)
        return max(1, total - non_final_total)

    if mode == 'late-heavy':
        if n == 1:
            weights = [1.0]
        elif n == 2:
            weights = [0.4, 0.6]
        elif n == 3:
            weights = [0.2, 0.3, 0.5]
        else:
            # geometric series with ratio 1.5
            r = 1.5
            raw = [r ** i for i in range(n)]
            s   = sum(raw)
            weights = [v / s for v in raw]
        steps = max(1, round(total * weights[idx]))
        return steps

    # uniform (default / legacy)
    return max(1, total // n)


def _resize_model_wrap_cfg(p, stage_w: int, stage_h: int) -> None:
    """
    Resize spatial state inside p.sampler.model_wrap_cfg to match stage dims.

    KDiffusionSampler.sample() calls get_scalings() and copies model_wrap_cfg.*
    into the sampler loop *before* the first step.  Any full-size spatial tensor
    still living there will mismatch the stage-sized latent `x` on the first
    sigma operation (e.g. `x - eps * (sigma_hat**2 - sigmas[i]**2)**0.5`).

    This must be called *after* create_sampler() and *inside* _stage_sampler_context
    so the resize is coherent with the p.width/p.height override.

    Only non-None tensors are touched; all errors are silently swallowed so a
    fork with a different model_wrap_cfg structure cannot crash the whole pass.
    """
    import torch.nn.functional as _F
    stage_hw = (stage_h // opt_f, stage_w // opt_f)
    mw = getattr(getattr(p, 'sampler', None), 'model_wrap_cfg', None)
    if mw is None:
        return
    for name in ('init_latent', 'mask', 'nmask'):
        t = getattr(mw, name, None)
        if t is None:
            continue
        try:
            if name == 'init_latent':
                setattr(mw, name, _F.interpolate(t, size=stage_hw, mode='nearest-exact'))
            else:
                setattr(mw, name,
                        _F.interpolate(t.unsqueeze(0), size=stage_hw,
                                       mode='nearest-exact').squeeze(0))
        except Exception:
            pass


def _stage_sample_txt2img(p, x, conditioning, unconditional_conditioning,
                          image_cond, stage_w: int, stage_h: int,
                          sampler_profile: str) -> torch.Tensor:
    """Run a single txt2img sampler pass on a stage-sized latent.

    Recommended SMEA compatibility path:
      - enter _stage_sampler_context so p.* becomes stage-sized
      - create a fresh sampler inside the context
      - patch model_wrap_cfg to stage dims
      - run hooks
      - re-wrap any scheduler override hooks may have installed
      - evict cached sigmas and patch sampler.func so any spatial/list sigmas
        are collapsed to a plain 1D schedule before entering SMEA Euler Max*
    """
    if sampler_profile == 'smea':
        with _stage_sampler_context(p, stage_w, stage_h):
            p.sampler = sd_samplers.create_sampler(p.sampler_name, p.sd_model)
            _resize_model_wrap_cfg(p, stage_w, stage_h)
            _call_script_hooks(p, x)
            _rewrap_scheduler_override_for_stage(p, stage_w, stage_h)
            try:
                extra = getattr(p.sampler, 'sampler_extra_args', {})
                if isinstance(extra, dict):
                    extra.pop('sigmas', None)
                    # noise_sampler may have been built from full-size p.x by
                    # a previous initialize() call (e.g. adept-sampler-v5 stores
                    # the original and passes it through).  Removing it here
                    # forces A1111's initialize() to recreate it from the
                    # stage-sized p.x that _stage_sampler_context set up,
                    # preventing "size of tensor a (26) must match tensor b (104)"
                    # in ancestral / TCD steps.
                    extra.pop('noise_sampler', None)
            except Exception:
                pass
            _patch_sampler_func_sigmas_to_1d(p)
            # Kill any remaining override so A1111's sample() uses standard
            # get_sigmas(p, steps) — which respects stage-sized p.width/p.height
            # and returns plain 1D sigmas.  This eliminates the spatial-sigma
            # class of crashes regardless of which scheduler extension is active.
            p.sampler_noise_scheduler_override = None
            return p.sampler.sample(
                p, x, conditioning, unconditional_conditioning,
                image_conditioning=image_cond,
            )
    else:
        # Standard samplers also need p.width/p.height, p.x, and the override
        # wrapped to stage dims.  Without this, Noise Sync's spatial override
        # ([T, H_full, W_full]) produces full-size sigma_up that mismatches the
        # stage-sized latent x.
        with _stage_sampler_context(p, stage_w, stage_h):
            p.sampler_noise_scheduler_override = None
            _call_script_hooks(p, x)
            # For standard samplers the sampler object is reused across stages.
            # A1111's initialize() may skip recreating noise_sampler if it is
            # already present in sampler_extra_args (built from full-size p.x
            # on the first call).  Drop it here so initialize() always builds a
            # fresh one from the stage-sized p.x set by _stage_sampler_context.
            try:
                extra = getattr(p.sampler, 'sampler_extra_args', {})
                if isinstance(extra, dict):
                    extra.pop('sigmas', None)
                    extra.pop('noise_sampler', None)
            except Exception:
                pass
            return p.sampler.sample(
                p, x, conditioning, unconditional_conditioning,
                image_conditioning=image_cond,
            )


def _stage_sample_img2img(p, samples, noise, conditioning, unconditional_conditioning,
                          stage_w: int, stage_h: int,
                          steps: int, sampler_profile: str) -> torch.Tensor:
    """Run a single img2img refinement pass on a stage-sized latent.

    Mirrors _stage_sample_txt2img() so txt2img and refinement use the same
    sampler lifecycle and sigma-compatibility logic for SMEA-family samplers.
    """
    if sampler_profile == 'smea':
        with _stage_sampler_context(p, stage_w, stage_h):
            p.sampler = sd_samplers.create_sampler(p.sampler_name, p.sd_model)
            _resize_model_wrap_cfg(p, stage_w, stage_h)
            _call_script_hooks(p, samples)
            _rewrap_scheduler_override_for_stage(p, stage_w, stage_h)
            try:
                extra = getattr(p.sampler, 'sampler_extra_args', {})
                if isinstance(extra, dict):
                    extra.pop('sigmas', None)
                    extra.pop('noise_sampler', None)
            except Exception:
                pass
            _patch_sampler_func_sigmas_to_1d(p)
            # Same as txt2img: kill override so get_sigmas uses stage dims.
            p.sampler_noise_scheduler_override = None
            return p.sampler.sample_img2img(
                p, samples, noise,
                conditioning, unconditional_conditioning,
                steps=steps,
                image_conditioning=p.image_conditioning,
            )
    else:
        with _stage_sampler_context(p, stage_w, stage_h):
            p.sampler_noise_scheduler_override = None
            _call_script_hooks(p, samples)
            try:
                extra = getattr(p.sampler, 'sampler_extra_args', {})
                if isinstance(extra, dict):
                    extra.pop('sigmas', None)
                    extra.pop('noise_sampler', None)
            except Exception:
                pass
            return p.sampler.sample_img2img(
                p, samples, noise,
                conditioning, unconditional_conditioning,
                steps=steps,
                image_conditioning=p.image_conditioning,
            )


def _run_img2img_refinement(p, samples, seeds, subseeds, subseed_strength,
                             conditioning, unconditional_conditioning,
                             refine_idx: int, n_refine: int,
                             refine_step_mode: str,
                             sampler_profile: str = 'standard') -> torch.Tensor:
    """
    Full VAE-decode refinement.
    Decodes latents, re-encodes as img2img conditioning, then delegates to
    _stage_sample_img2img which handles sampler lifecycle per profile.

    refine_idx / n_refine / refine_step_mode feed into _get_refine_steps so
    the step budget reflects the chosen policy rather than always uniform.
    """
    steps_for_refinement = _get_refine_steps(p, refine_idx, n_refine, refine_step_mode)
    noise = _make_noise(p, samples.shape[1:], seeds, subseeds, subseed_strength)

    decoded = decode_latent_batch(p.sd_model, samples,
                                  target_device=devices.cpu, check_for_nans=True)
    decoded = torch.stack(decoded).float()
    decoded = torch.clamp((decoded + 1.0) / 2.0, 0.0, 1.0)
    source_img = decoded * 2.0 - 1.0

    p.image_conditioning = p.img2img_image_conditioning(source_img, samples)

    # stage dims derived from the current latent (samples is already stage-sized)
    stage_w = samples.shape[3] * opt_f
    stage_h = samples.shape[2] * opt_f

    return _stage_sample_img2img(
        p, samples, noise, conditioning, unconditional_conditioning,
        stage_w, stage_h, steps_for_refinement, sampler_profile,
    )


# ─────────────────────────────────────────────────────────────────────────────
# Version implementations
# ─────────────────────────────────────────────────────────────────────────────

def sample_v1_exact(p, conditioning, unconditional_conditioning,
                    seeds, subseeds, subseed_strength, prompts, plan: StagePlan):
    """
    v1 (exact) – original implementation, preserved 1:1 as a regression baseline.
    Does NOT use the pre-built (deduped) plan for sampling; rebuilds the raw
    stage list from min_scale/max_scale/steps exactly as the original did.

    For infotext honesty, a v1-native StagePlan is constructed from the same
    raw stage list so PG plan reflects what v1 actually ran, not the deduped plan
    that the wrapper built for guard purposes.
    """
    min_scale        = float(p.progressive_growing_min_scale)
    max_scale        = float(p.progressive_growing_max_scale)
    resolution_steps = np.linspace(min_scale, max_scale, int(p.progressive_growing_steps))

    def _snap_v1(v):
        v_int = int(v)
        v_int = max(opt_f, v_int)
        v_int = (v_int // opt_f) * opt_f
        return max(opt_f, v_int)

    # build the honest v1 stage list (duplicates preserved, no forced final snap)
    v1_sizes = [
        (_snap_v1(p.width * s), _snap_v1(p.height * s))
        for s in resolution_steps
    ]
    v1_plan = StagePlan(
        sizes=v1_sizes,
        final_w=v1_sizes[-1][0],
        final_h=v1_sizes[-1][1],
    )

    _write_params(p, v1_plan, refine_policy='all stages',
                  refine_step_policy='uniform (v1 fixed)')
    sampler_profile, _ = _get_sampler_profile(p)
    # v1 always uses bicubic regardless of the UI dropdown
    try:
        p.extra_generation_params['PG interp'] = 'bicubic (v1 fixed)'
        if sampler_profile != 'standard':
            p.extra_generation_params['PG sampler profile'] = sampler_profile
        # v1 builds its stage list from the manual Stages slider, not from
        # auto stage count – override the fields _write_params may have set
        # so infotext describes what v1 actually ran, not what the user expected
        if getattr(p, 'progressive_growing_auto_stages', False):
            p.extra_generation_params['PG stages mode'] = 'manual (v1 legacy – auto ignored)'\
            
            p.extra_generation_params.pop('PG target jump', None)
            p.extra_generation_params.pop('PG requested',   None)
            p.extra_generation_params.pop('PG actual',      None)
    except Exception:
        pass

    initial_width, initial_height = v1_sizes[0]

    x = create_random_tensors(
        (opt_C, initial_height // opt_f, initial_width // opt_f),
        seeds, subseeds=subseeds, subseed_strength=subseed_strength,
        seed_resize_from_h=p.seed_resize_from_h,
        seed_resize_from_w=p.seed_resize_from_w, p=p,
    )

    image_cond = p.txt2img_image_conditioning(x, width=initial_width, height=initial_height)

    samples = _stage_sample_txt2img(
        p, x, conditioning, unconditional_conditioning,
        image_cond, initial_width, initial_height, sampler_profile,
    )

    total_stages = len(resolution_steps)

    for i in range(1, total_stages):
        target_width, target_height = v1_sizes[i]

        samples = torch.nn.functional.interpolate(
            samples,
            size=(target_height // opt_f, target_width // opt_f),
            mode='bicubic', align_corners=False,
        )

        if p.progressive_growing_refinement:
            samples = _run_img2img_refinement(
                p, samples, seeds, subseeds, subseed_strength,
                conditioning, unconditional_conditioning,
                refine_idx=i - 1,
                n_refine=total_stages - 1,
                refine_step_mode='uniform',   # v1 legacy always uniform
                sampler_profile=sampler_profile,
            )

    return samples


# ─────────────────────────────────────────────────────────────────────────────

def _run_plan_with_policy(p, plan: StagePlan, conditioning, unconditional_conditioning,
                           seeds, subseeds, subseed_strength,
                           refine_predicate) -> torch.Tensor:
    """
    Shared loop for v2 / v4 / v5 / v6.

    refine_predicate(stage_idx, w, h, plan) -> bool
        Called for every stage after the first; return True to run refinement.

    Reads p.progressive_growing_interp_mode for latent upscale (default: bicubic).
    Reads p.progressive_growing_refine_step_mode for per-pass step budget.

    Pre-computes the number of refine passes so _get_refine_steps can allocate
    a budget that accounts for the total workload (needed by late-heavy / final-heavy).

    Sampler compatibility
    --------------------
    _get_sampler_profile() is called once to determine the active profile.
    'smea' profile wraps every sampler.sample() / sample_img2img() call with
    _stage_sampler_context(), which temporarily sets p.width/p.height to stage
    dims and rescales p.init_latent/mask/nmask.  This fixes spatial-sigma
    schedulers and model-wrapper state mismatches without touching the rest of
    the pipeline.  The profile name is written to extra_generation_params as
    'PG sampler profile' when non-standard.
    """
    first_w, first_h = plan.sizes[0]
    interp_mode      = getattr(p, 'progressive_growing_interp_mode',    LATENT_INTERP_DEFAULT)
    step_mode        = getattr(p, 'progressive_growing_refine_step_mode', REFINE_STEP_DEFAULT)
    do_refine        = getattr(p, 'progressive_growing_refinement', True)
    sampler_profile, _ = _get_sampler_profile(p)

    # record non-standard profile in infotext
    if sampler_profile != 'standard':
        try:
            p.extra_generation_params['PG sampler profile'] = sampler_profile
        except Exception:
            pass

    # pre-count how many stages will actually refine so budget can be distributed
    n_refine = sum(
        1 for i, (w, h) in enumerate(plan.sizes[1:], start=1)
        if do_refine and refine_predicate(i, w, h, plan)
    ) if do_refine else 0

    x = _initial_latent(p, first_w, first_h, seeds, subseeds, subseed_strength)
    image_cond = p.txt2img_image_conditioning(x, width=first_w, height=first_h)

    samples = _stage_sample_txt2img(
        p, x, conditioning, unconditional_conditioning,
        image_cond, first_w, first_h, sampler_profile,
    )

    refine_idx = 0
    for i, (w, h) in enumerate(plan.sizes[1:], start=1):
        samples = _upscale_latent(samples, w, h, interp_mode=interp_mode)

        if do_refine and refine_predicate(i, w, h, plan):
            samples = _run_img2img_refinement(
                p, samples, seeds, subseeds, subseed_strength,
                conditioning, unconditional_conditioning,
                refine_idx=refine_idx,
                n_refine=max(1, n_refine),
                refine_step_mode=step_mode,
                sampler_profile=sampler_profile,
            )
            refine_idx += 1

    return samples


def sample_v2_safe(p, conditioning, unconditional_conditioning,
                   seeds, subseeds, subseed_strength, prompts, plan: StagePlan):
    """
    v2 (safe) – validated plan, dedup, correct per-stage conditioning size, script hooks.
    Behaves identically to v1 where plans agree; differs only in edge cases.
    """
    _write_params(p, plan, refine_policy='all stages',
                  refine_step_policy=getattr(p, 'progressive_growing_refine_step_mode', REFINE_STEP_DEFAULT))

    def always(i, w, h, plan):
        return True

    return _run_plan_with_policy(
        p, plan, conditioning, unconditional_conditioning,
        seeds, subseeds, subseed_strength,
        refine_predicate=always,
    )


def sample_v3_fast(p, conditioning, unconditional_conditioning,
                   seeds, subseeds, subseed_strength, prompts, plan: StagePlan):
    """
    v3 (fast) – refinement only on the final stage.
    Avoids expensive VAE decode on every intermediate stage.
    Best for quick iteration or large stage counts.
    """
    _write_params(p, plan, refine_policy='final stage only',
                  refine_step_policy=getattr(p, 'progressive_growing_refine_step_mode', REFINE_STEP_DEFAULT))

    def only_last(i, w, h, plan):
        return i == plan.n_stages - 1

    return _run_plan_with_policy(
        p, plan, conditioning, unconditional_conditioning,
        seeds, subseeds, subseed_strength,
        refine_predicate=only_last,
    )


def sample_v4_balanced(p, conditioning, unconditional_conditioning,
                        seeds, subseeds, subseed_strength, prompts, plan: StagePlan):
    """
    v4 (balanced) – refinement on stages whose linear scale >= BALANCED_REFINE_THRESHOLD.

    The threshold is a *linear* dimension ratio (sqrt of area ratio), so
    BALANCED_REFINE_THRESHOLD=0.70 means: refine when the stage side is >= 70 %
    of the final side, which corresponds to ≈ 49 % of the final pixel area.
    This avoids refinement on cheap small upscales while still covering the
    detail-sensitive near-final stages.
    """
    _write_params(p, plan, refine_policy=f'linear scale >= {BALANCED_REFINE_THRESHOLD}',
                  refine_step_policy=getattr(p, 'progressive_growing_refine_step_mode', REFINE_STEP_DEFAULT))

    final_area = plan.final_w * plan.final_h

    def large_stages_only(i, w, h, plan):
        # linear scale = sqrt(stage_area / final_area)
        linear_scale = math.sqrt((w * h) / final_area) if final_area > 0 else 0.0
        return linear_scale >= BALANCED_REFINE_THRESHOLD

    return _run_plan_with_policy(
        p, plan, conditioning, unconditional_conditioning,
        seeds, subseeds, subseed_strength,
        refine_predicate=large_stages_only,
    )


def sample_v5_latent(p, conditioning, unconditional_conditioning,
                     seeds, subseeds, subseed_strength, prompts, plan: StagePlan):
    """
    v5 (latent only) – pure latent upscale, no refinement at any stage.
    Fastest mode; useful to study the raw effect of latent-space growing.
    """
    _write_params(p, plan, refine_policy='none', refine_step_policy='none')

    def never(i, w, h, plan):
        return False

    return _run_plan_with_policy(
        p, plan, conditioning, unconditional_conditioning,
        seeds, subseeds, subseed_strength,
        refine_predicate=never,
    )


def sample_v6_adaptive(p, conditioning, unconditional_conditioning,
                       seeds, subseeds, subseed_strength, prompts, plan: StagePlan):
    """
    v6 (adaptive) – refinement policy scales with stage count.

      2 stages   → refine every post-initial stage (progressive would be pointless otherwise)
      3–4 stages → refine the last 2 stages only
      5+ stages  → refine stages whose linear scale >= ADAPTIVE_REFINE_THRESHOLD,
                   plus always force-refine the final stage

    Sits between v3 (final only) and v2 (all stages): cheaper than v2 on long
    plans, more thorough than v3 on short ones.
    """
    _write_params(p, plan, refine_policy=f'adaptive (threshold >= {ADAPTIVE_REFINE_THRESHOLD:.2f})',
                  refine_step_policy=getattr(p, 'progressive_growing_refine_step_mode', REFINE_STEP_DEFAULT))

    final_area = plan.final_w * plan.final_h

    def adaptive(i, w, h, plan):
        n = plan.n_stages
        if n <= 2:
            # only 1 post-initial stage – always refine it
            return True
        if n <= 4:
            # short plan: refine the last 2 stages
            return i >= n - 2
        # long plan: threshold + guaranteed final
        linear_scale = math.sqrt((w * h) / final_area) if final_area > 0 else 0.0
        return i == n - 1 or linear_scale >= ADAPTIVE_REFINE_THRESHOLD

    return _run_plan_with_policy(
        p, plan, conditioning, unconditional_conditioning,
        seeds, subseeds, subseed_strength,
        refine_predicate=adaptive,
    )


# ─────────────────────────────────────────────────────────────────────────────
# Version registry
# ─────────────────────────────────────────────────────────────────────────────

_VERSIONS: dict[str, callable] = {
    "v2 (safe)":     sample_v2_safe,
    "v3 (fast)":     sample_v3_fast,
    "v4 (balanced)": sample_v4_balanced,
    "v5 (latent)":   sample_v5_latent,
    "v6 (adaptive)": sample_v6_adaptive,
    "v1 (exact)":    sample_v1_exact,     # legacy / regression reference
}

_DEFAULT_VERSION = "v2 (safe)"


# ─────────────────────────────────────────────────────────────────────────────
# Monkey-patch – thin wrapper, applied exactly once
# ─────────────────────────────────────────────────────────────────────────────

_PATCHED    = False
_ORIG_SAMPLE = None


def _apply_patch_once() -> None:
    """Patch StableDiffusionProcessingTxt2Img.sample once at first use."""

    global _PATCHED, _ORIG_SAMPLE
    if _PATCHED:
        return

    cls = getattr(processing_mod, 'StableDiffusionProcessingTxt2Img', None)
    if cls is None:
        return

    if getattr(cls, '_pg_ext_patched', False):
        _PATCHED = True
        return

    _ORIG_SAMPLE = cls.sample

    def _sample_wrapper(self,
                         conditioning, unconditional_conditioning,
                         seeds, subseeds, subseed_strength, prompts):

        ok, reason, plan = _should_use_progressive(self)

        if ok:
            # For 'standard' profile: create sampler once here (mirrors what the
            # original sample() does internally).
            # For 'smea' profile: sampler is created fresh per stage-pass inside
            # _stage_sample_txt2img / _stage_sample_img2img so that all
            # size-dependent state (sigma schedules, spatial caches) is
            # initialised against stage dims, not the final target size.
            profile, _ = _get_sampler_profile(self)
            if profile == 'standard':
                self.sampler = sd_samplers.create_sampler(self.sampler_name, self.sd_model)

            ver = getattr(self, 'progressive_growing_version', _DEFAULT_VERSION)
            fn  = _VERSIONS.get(ver, sample_v2_safe)
            # plan was already built by _should_use_progressive – pass it through
            # so version fns don't call build_plan() a second time
            return fn(self, conditioning, unconditional_conditioning,
                      seeds, subseeds, subseed_strength, prompts, plan)

        # record skip reason if there was one
        if reason:
            try:
                self.extra_generation_params['PG skip'] = reason
            except Exception:
                pass

        return _ORIG_SAMPLE(self, conditioning, unconditional_conditioning,
                             seeds, subseeds, subseed_strength, prompts)

    cls.sample       = _sample_wrapper
    cls._pg_ext_patched = True
    _PATCHED = True


# ─────────────────────────────────────────────────────────────────────────────
# Always-visible UI script
# ─────────────────────────────────────────────────────────────────────────────

class ProgressiveGrowingAlwaysVisible(scripts.Script):

    def title(self):
        return "Progressive Growing"

    def show(self, is_img2img):
        return scripts.AlwaysVisible if not is_img2img else False

    def ui(self, is_img2img):
        with gr.Accordion("Progressive Growing", open=False):
            enabled = gr.Checkbox(value=False, label="Enable")
            version = gr.Dropdown(
                choices=list(_VERSIONS.keys()),
                value=_DEFAULT_VERSION,
                label="Mode",
            )

            with gr.Row():
                min_scale = gr.Slider(
                    minimum=0.1, maximum=1.0, step=0.05,
                    value=0.25, label="Min scale",
                )
                max_scale = gr.Slider(
                    minimum=0.1, maximum=1.0, step=0.05,
                    value=1.0,  label="Max scale",
                )

            stages = gr.Slider(
                minimum=2, maximum=16, step=1,
                value=4, label="Stages (ignored when Auto is on)",
            )
            with gr.Row():
                auto_stages = gr.Checkbox(value=False, label="Auto stage count")
                auto_jump   = gr.Slider(
                    minimum=1.1, maximum=2.0, step=0.05,
                    value=AUTO_STAGE_JUMP_DEFAULT,
                    label="Target jump per stage",
                )
                auto_max    = gr.Slider(
                    minimum=2, maximum=12, step=1,
                    value=AUTO_STAGE_MAX_DEFAULT,
                    label="Max auto stages",
                )
            with gr.Row():
                refinement = gr.Checkbox(value=True, label="Refinement between stages")
                refine_step_mode = gr.Dropdown(
                    choices=REFINE_STEP_MODES,
                    value=REFINE_STEP_DEFAULT,
                    label="Refinement step budget",
                )
                interp_mode = gr.Dropdown(
                    choices=LATENT_INTERP_MODES,
                    value=LATENT_INTERP_DEFAULT,
                    label="Latent upscale interpolation",
                )

            gr.Markdown(
                "**Modes**\n"
                "- **v2 (safe)** – validated, deduped stages, refinement at every stage *(default)*\n"
                "- **v3 (fast)** – refinement only on the final stage; fastest, no VAE mid-pass\n"
                "- **v4 (balanced)** – refinement when stage side ≥ 70 % of final side (≈ 49 % of area)\n"
                "- **v5 (latent)** – pure latent upscale, zero refinement\n"
                "- **v6 (adaptive)** – 2 stages: all; 3–4 stages: last 2; 5+ stages: threshold + final\n"
                "- **v1 (exact)** – original implementation, kept for regression comparison; "
                "always uses manual Stages, ignores Auto stage count\n\n"
                "**Auto stage count** – ignores Stages slider; computes count so each upscale grows "
                "the latent side by ~Target jump (1.35 ≈ 35 %). 0.25→1.0 ≈ 6 stages, 0.5→1.0 ≈ 3 stages.\n\n"
                "**Interpolation** – bicubic: smooth/photo; bilinear: softer; nearest: fastest/pixel-art; area: anti-aliased\n\n"
                "**Refinement budget** – uniform: equal steps per pass (legacy); "
                "late-heavy: growing budget towards final pass; "
                "final-heavy: minimal steps on all but the last refinement pass\n\n"
                "⚠ Incompatible with **Hires. fix** — enabling both disables Progressive Growing."
            )

        return [enabled, version, min_scale, max_scale,
                stages, auto_stages, auto_jump, auto_max,
                refinement, refine_step_mode, interp_mode]

    def process(self, p,
                enabled, version, min_scale, max_scale,
                stages, auto_stages, auto_jump, auto_max,
                refinement, refine_step_mode, interp_mode):

        _apply_patch_once()

        p.enable_progressive_growing            = bool(enabled)
        p.progressive_growing_version           = str(version)
        p.progressive_growing_min_scale         = float(min_scale)
        p.progressive_growing_max_scale         = float(max_scale)
        p.progressive_growing_steps             = int(stages)
        p.progressive_growing_auto_stages       = bool(auto_stages)
        p.progressive_growing_auto_jump         = float(auto_jump)
        p.progressive_growing_auto_max          = int(auto_max)
        p.progressive_growing_refinement        = bool(refinement)
        p.progressive_growing_refine_step_mode  = str(refine_step_mode) if refine_step_mode in REFINE_STEP_MODES else REFINE_STEP_DEFAULT
        p.progressive_growing_interp_mode       = str(interp_mode) if interp_mode in LATENT_INTERP_MODES else LATENT_INTERP_DEFAULT