JasonHoku-ComfyUI-Ultimate-Auto-Sampler-Config-Grid-Testing-Suite/upscale_runner.py

"""
Dashboard Upscale Runner
Background thread for upscaling images triggered from the dashboard.
Loads models via ComfyUI APIs, runs upscale pipeline chains,
saves results and updates manifests.

Cancellation (Roadmap 93): a job carries a threading.Event. It is checked at every
granularity the runner controls (per image, per pipeline, per expanded step, per combo,
before the SeedVR2 / Florence2 calls) and INSIDE sampling through a cloned model whose
unet-function wrapper raises UpscaleCancelled once the event is set. The process-wide
comfy.model_management interrupt flag is deliberately NOT used: it would also kill a
Generator run in the main queue and a leftover True would abort the next queued prompt.
SeedVR2 runs as an external node call, so it can only be cancelled between steps.
"""

import math
import os
import re
import time
import json
import threading
import uuid

import torch
import numpy as np
from PIL import Image as PILImage

from .model_loader import load_checkpoint, load_vae_by_name
from .batch_encoding import encode_prompt_with_combinators, encode_pair_cached
from .regional_prompt import append_to_base
from .image_generation import upscale_image, decode_latent_with_vae, create_image_metadata
from .manifest_utils import save_manifest, base_dims_of
from . import image_metrics
from . import comfy_nodes_workflow

import comfy.model_management as mm

# Active upscale jobs (module-level state)
_upscale_jobs = {}
_jobs_lock = threading.RLock()

ALREADY_RUNNING_MSG = "An upscale job is already running for this session"
_ACTIVE_STATES = ("running", "cancelling")
_VALID_MODES = ("hires_only", "model_only", "model_then_hires", "seedvr2", "florence2_hires")


class UpscaleCancelled(Exception):
    """Raised inside the upscale thread when the user cancels the job."""


def _check_cancel(job):
    ev = job.get("cancel")
    if ev is not None and ev.is_set():
        raise UpscaleCancelled()


# =============================================================================
# CONFIG VALIDATION (pure-Python mirror of resources/logic_upscale_edit.js)
# =============================================================================
_NUM_RE = re.compile(r"^[+-]?(\d+\.?\d*|\.\d+)([eE][+-]?\d+)?$")


def _num(v):
    """Strict number: int/float or a plain decimal string; anything else (including 'x', '', None) -> None."""
    if isinstance(v, bool) or v is None:
        return None
    if isinstance(v, (int, float)):
        return float(v) if math.isfinite(v) else None
    if isinstance(v, str):
        s = v.strip()
        return float(s) if _NUM_RE.match(s) else None
    return None


def _num_list(v):
    """('1.5, 2' -> ([1.5, 2.0], bad)) ; empty entries are skipped like the runner does."""
    parts = [p.strip() for p in str(v).split(",")]
    parts = [p for p in parts if p]
    vals, bad = [], False
    for p in parts:
        n = _num(p)
        if n is None:
            bad = True
        else:
            vals.append(n)
    return parts, vals, bad


def _is_int(n):
    return n is not None and n == int(n)


def _falsy(v):
    return v is None or v is False or v == 0 or v == ""


def _validate_upscale(cfg):
    """-> (errors, warnings); each entry {"path", "field", "msg"}. Same rules as validateUpscaleConfig()."""
    errors, warnings = [], []

    def err(path, field, msg):
        errors.append({"path": path, "field": field, "msg": msg})

    def warn(path, field, msg):
        warnings.append({"path": path, "field": field, "msg": msg})

    if not isinstance(cfg, dict) or not isinstance(cfg.get("pipelines"), list):
        err("", "pipelines", "Config has no pipelines")
        return errors, warnings

    any_active = False
    for pi, pipe in enumerate(cfg["pipelines"]):
        if not isinstance(pipe, dict) or pipe.get("active", True) is False:
            continue
        steps = pipe.get("steps")
        steps = steps if isinstance(steps, list) else []
        for si, step in enumerate(steps):
            if not isinstance(step, dict) or step.get("active", True) is False:
                continue
            any_active = True
            path = f"p{pi}.s{si}"
            mode = step.get("mode") or "hires_only"
            if mode not in _VALID_MODES:
                err(path, "mode", f"Unknown mode '{mode}'")
                continue
            hires = mode in ("hires_only", "model_then_hires")
            f2 = mode == "florence2_hires"
            model = mode in ("model_only", "model_then_hires")

            rep = step.get("repeat")
            if _falsy(rep):
                rep = 1
            else:
                rn = _num(rep)
                if rn is not None and rn < 1:
                    rep = 1
            rn = _num(rep)
            if not _is_int(rn) or rn < 1 or rn > 20:
                err(path, "repeat", "Repeat must be a whole number from 1 to 20")

            if hires:
                ratios = step.get("upscale_ratios")
                ratios = "1.5" if _falsy(ratios) else ratios
                parts, vals, bad = _num_list(ratios)
                if not parts:
                    err(path, "upscale_ratios", "Enter at least one ratio, e.g. 1.5, 2")
                elif bad:
                    err(path, "upscale_ratios", "Ratios must be numbers separated by commas")
                elif any(v <= 0 or v > 8 for v in vals):
                    err(path, "upscale_ratios", "Each ratio must be above 0 and at most 8")
            if hires or f2:
                dn = step.get("hires_denoise")
                dn = "0.3" if _falsy(dn) else dn
                parts, vals, bad = _num_list(dn)
                if not parts:
                    err(path, "hires_denoise", "Enter at least one denoise value, e.g. 0.3")
                elif bad:
                    err(path, "hires_denoise", "Denoise must be numbers separated by commas")
                elif any(v < 0 or v > 1 for v in vals):
                    err(path, "hires_denoise", "Each denoise value must be between 0 and 1")
                elif f2 and len(vals) > 1:
                    warn(path, "hires_denoise", "Florence2 uses only the first denoise value")
                hs = step.get("hires_steps", 0) if "hires_steps" in step else 0
                hn = _num(hs)
                if not _is_int(hn) or hn < 0 or hn > 150:
                    err(path, "hires_steps", "HiRes steps must be a whole number from 0 to 150")
            if hires and step.get("hires_tiled_sampling"):
                for key, label in (("hires_tile_width", "Tile width"), ("hires_tile_height", "Tile height")):
                    v = step.get(key)
                    v = 512 if _falsy(v) else v
                    n = _num(v)
                    if n is None or n < 64 or n % 8 != 0:
                        err(path, key, f"{label} must be a multiple of 8 and at least 64")
            if hires and step.get("tiled_vae"):
                v = step.get("tile_size")
                v = 512 if _falsy(v) else v
                n = _num(v)
                if n is None or n < 128:
                    err(path, "tile_size", "VAE tile size must be at least 128")
            if model:
                models = step.get("upscale_models")
                if not isinstance(models, list) or len([m for m in models if m]) == 0:
                    err(path, "upscale_models", "Select at least one upscale model")
            if mode == "model_only":
                v = step.get("upscale_size")
                v = "2.0" if _falsy(v) else v
                n = _num(v)
                if n is None or n <= 0:
                    err(path, "upscale_size", "Size multiplier must be a number above 0")
            if mode == "seedvr2":
                sv = step.get("seedvr2")
                if sv is not None and not isinstance(sv, dict):
                    sv = {}
                if sv is not None:
                    if not str(sv.get("dit_model") or "").strip():
                        err(path, "dit_model", "Choose a SeedVR2 DiT model")
                    if sv.get("resolution") is not None:
                        n = _num(sv.get("resolution"))
                        if n is None or n < 16:
                            err(path, "resolution", "SeedVR2 resolution must be at least 16")
            if f2:
                fl = step.get("florence2")
                if fl is not None and not isinstance(fl, dict):
                    fl = {}
                if fl is not None and not str(fl.get("text_input") or "").strip():
                    err(path, "text_input", "Say what Florence2 should detect (e.g. face)")

    if not any_active:
        err("", "pipelines", "Turn on at least one pipeline that has an active step")
    if cfg.get("hires_prompt_adjust") and not str(cfg.get("hires_prompt_text") or "").strip():
        warn("", "hires_prompt_text", "Prompt adjustment is on but the HiRes prompt is empty")
    return errors, warnings


def validate_upscale_config(cfg):
    """Server-side guard for POST /config_tester/upscale_images: a list of error strings (empty = ok)."""
    errors, _warnings = _validate_upscale(cfg)
    return [(f"{e['path']}: " if e["path"] else "") + e["msg"] for e in errors]


# =============================================================================
# JOB BOOKKEEPING
# =============================================================================
def _prune_jobs(max_age=3600, now=None):
    """Forget finished jobs older than max_age seconds (running jobs are never pruned)."""
    now = time.time() if now is None else now
    with _jobs_lock:
        for jid, job in list(_upscale_jobs.items()):
            if job.get("status") in _ACTIVE_STATES or job.get("status") == "running":
                continue
            fin = job.get("finished_at")
            if fin is not None and now - fin > max_age:
                del _upscale_jobs[jid]


def _is_active(job):
    return job.get("status") == "running"


def start_upscale_job(session_name, image_ids, upscale_config, all_favorited=False):
    """
    Start an async upscale job. Returns (job_id, error) immediately.
    The actual work runs in a background thread. Only one job may run per session
    (two jobs would both rewrite the same manifest).
    """
    import folder_paths
    _prune_jobs()
    with _jobs_lock:
        for other in _upscale_jobs.values():
            if other.get("session_name") == session_name and _is_active(other):
                return None, ALREADY_RUNNING_MSG

    benchmarks_dir = os.path.realpath(os.path.join(folder_paths.get_output_directory(), "benchmarks"))
    session_dir = os.path.join(benchmarks_dir, session_name)
    manifest_path = os.path.join(session_dir, "manifest.json")
    images_dir = os.path.join(session_dir, "images")

    if not os.path.exists(manifest_path):
        return None, "Manifest not found"

    with open(manifest_path, "r", encoding="utf-8") as f:
        manifest_data = json.load(f)

    items = manifest_data.get("items", [])
    meta = manifest_data.get("meta", {})

    # Resolve image list
    if all_favorited:
        target_items = [item for item in items if item.get("favorited")]
    else:
        id_set = set(image_ids) if image_ids else set()
        target_items = [item for item in items if item.get("id") in id_set]

    if not target_items:
        return None, "No matching images found"

    job_id = f"upscale_{uuid.uuid4().hex[:12]}"
    job = {
        "id": job_id,
        "session_name": session_name,
        "status": "running",
        "completed": 0,
        "total": len(target_items),
        "outputs": 0,
        "current_image": "",
        "current_step": "",
        "error": None,
        "cancel": threading.Event(),
        "started_at": time.time(),
        "finished_at": None,
    }
    with _jobs_lock:
        for other in _upscale_jobs.values():     # re-check: the manifest read above is not under the lock
            if other.get("session_name") == session_name and _is_active(other):
                return None, ALREADY_RUNNING_MSG
        _upscale_jobs[job_id] = job

    thread = threading.Thread(
        target=_run_upscale_thread,
        args=(job, target_items, upscale_config, meta, manifest_data, manifest_path, images_dir, session_name),
        daemon=True
    )
    thread.start()

    return job_id, None


def _public_status(job):
    status = job["status"]
    if status == "running" and job["cancel"].is_set():
        status = "cancelling"
    return status


def get_upscale_status(job_id):
    """Get status of an upscale job."""
    job = _upscale_jobs.get(job_id)
    if not job:
        return {"status": "not_found"}
    return {
        "status": _public_status(job),
        "completed": job["completed"],
        "total": job["total"],
        "outputs": job.get("outputs", 0),
        "current_image": job.get("current_image", ""),
        "current_step": job.get("current_step", ""),
        "error": job.get("error"),
    }


def list_upscale_jobs(session_name=None):
    """Summaries of the known jobs (optionally for one session), newest first."""
    with _jobs_lock:
        jobs = [j for j in _upscale_jobs.values() if session_name in (None, "", j.get("session_name"))]
    jobs.sort(key=lambda j: j.get("started_at") or 0, reverse=True)
    return [{"job_id": j["id"], "status": _public_status(j), "completed": j["completed"], "total": j["total"]}
            for j in jobs]


def cancel_upscale_job(job_id):
    """Signal an upscale job to stop. -> {"cancelled": bool, "status": str}. True only for a running job
    (calling it again while it winds down is harmless)."""
    job = _upscale_jobs.get(job_id)
    if not job:
        return {"cancelled": False, "status": "not_found"}
    if job["status"] == "running":
        job["cancel"].set()
        return {"cancelled": True, "status": "cancelling"}
    return {"cancelled": False, "status": job["status"]}


def _cancellable_model(model, cancel_event):
    """A clone of `model` whose unet function wrapper raises UpscaleCancelled once cancel_event is set.
    The raise happens between sampler steps (every model call), so a cancel lands within one step.
    An existing wrapper is chained, never replaced."""
    m = model.clone()
    prev = (getattr(m, "model_options", None) or {}).get("model_function_wrapper")

    def _wrapper(apply_model, args):
        if cancel_event.is_set():
            raise UpscaleCancelled()
        if prev is not None:
            return prev(apply_model, args)
        return apply_model(args["input"], args["timestep"], **args["c"])

    m.set_model_unet_function_wrapper(_wrapper)
    return m


def _clean_vram():
    try:
        mm.soft_empty_cache()
        import gc
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass


def _drive_upscale_items(job, target_items, process_item, on_item_done=None, on_progress=None, on_finish=None):
    """Run process_item(item, produced) over target_items with the cancel / progress semantics.

    produced is a list the callee appends every manifest entry it creates, so whatever an item
    finished before a cancel (or an error) is kept. Cancel is honoured before each item and when
    process_item raises UpscaleCancelled. Returns "complete" or "cancelled" (errors propagate after
    the partial entries were handed to on_item_done). The driver owns job status / counters."""
    total = len(target_items)
    cancelled = False
    for idx, item in enumerate(target_items):
        if job["cancel"].is_set():
            cancelled = True
            break
        produced = []
        try:
            process_item(item, produced)
        except UpscaleCancelled:
            cancelled = True
        except Exception:
            if produced:
                job["outputs"] += len(produced)
                if on_item_done:
                    on_item_done(item, produced)
            raise
        if produced:
            job["outputs"] += len(produced)
            if on_item_done:
                on_item_done(item, produced)
        if cancelled:
            break
        job["completed"] = idx + 1
        if on_progress:
            on_progress(job["completed"], total)
    state = "cancelled" if cancelled else "complete"
    if on_finish:
        try:
            on_finish(state)
        except Exception as e:
            print(f"[DashboardUpscale] ⚠️ final progress event failed: {e}")
    job["finished_at"] = time.time()
    job["status"] = state
    return state


def _save_and_stamp(pil, path, fmt, meta, pack):
    """Write the image in the session's format and embed the nodes workflow (PNG/WEBP; never raises)."""
    if fmt == "png":
        pil.save(path, format="PNG")
    elif fmt in ("jpg", "jpeg"):
        pil.save(path, format="JPEG", quality=95)
    else:
        pil.save(path, format="WEBP", quality=95)
    comfy_nodes_workflow.stamp_file(path, meta, enabled=pack)


def _send(PromptServer, event, payload):
    if PromptServer is None:
        return
    try:
        PromptServer.instance.send_sync(event, payload)
    except Exception:
        pass


def _run_upscale_thread(job, target_items, upscale_config, meta, manifest_data, manifest_path, images_dir, session_name):
    """
    Background thread that runs the upscale pipeline on each image.
    Loads the session's model/VAE/CLIP, encodes prompts, runs pipeline chains.
    """
    try:
        pipelines = upscale_config.get("pipelines", [])
        if not pipelines:
            job["status"] = "error"
            job["error"] = "No upscale pipelines configured"
            job["finished_at"] = time.time()
            return

        # Resolve image format from manifest meta (written by orchestrator)
        _raw_fmt = meta.get("image_format", "webp")
        _fmt = (_raw_fmt or "webp").lower().strip()
        if _fmt not in ("webp", "png", "jpg", "jpeg"):
            _fmt = "webp"
        _ext = "jpg" if _fmt in ("jpg", "jpeg") else _fmt

        # Load model, VAE, CLIP for conditioning
        model_name = meta.get("model", "") or ""
        vae_name = meta.get("vae", "")

        # If meta.model is empty/missing (legacy manifests, manifests written before
        # meta.model was recorded), fall back to the first non-empty item.model in
        # the target set. Items always store their own model name.
        if not model_name:
            for _it in target_items:
                _it_model = (_it.get("model", "") or "").strip()
                if _it_model:
                    model_name = _it_model
                    print(f"[DashboardUpscale] meta.model is empty; using first item's model: {model_name}")
                    break

        # Validate the checkpoint exists before calling load_checkpoint. Without
        # this guard, a missing file ends up at comfy/utils.py:load_torch_file(None,
        # ...) and crashes with the cryptic "'NoneType' object has no attribute
        # 'lower'". Common cause: manifest copied from another machine where the
        # checkpoint isn't installed; or manifest entry has model=None.
        try:
            import folder_paths as _fp_check
            ckpt_path_check = _fp_check.get_full_path("checkpoints", model_name) if model_name else None
        except Exception:
            ckpt_path_check = None
        if not model_name or ckpt_path_check is None:
            job["status"] = "error"
            job["error"] = (
                f"Checkpoint '{model_name or '(empty)'}' not found in "
                f"ComfyUI/models/checkpoints/. The session manifest references it but "
                f"the file is not installed on this machine. Copy the .safetensors "
                f"file over, or regenerate the images on this machine first."
            )
            job["finished_at"] = time.time()
            print(f"[DashboardUpscale] ❌ {job['error']}")
            return

        # VRAM hygiene before load_checkpoint: the previous workflow's model handles
        # may still be Python-referenced and won't be evicted by ComfyUI's LRU on a
        # new load_checkpoint call. soft_empty_cache + gc.collect releases those
        # before we add the upscale session's checkpoint on top. Critical for
        # Florence2 mode on small-VRAM cards (otherwise prior SDXL gen + new SDXL
        # session + Florence2 = OOM during beam search).
        _clean_vram()

        print(f"[DashboardUpscale] \U0001f504 Loading model: {model_name}")
        loaded_model, loaded_clip, loaded_vae = load_checkpoint(
            target_model_name=model_name,
            ckpt_name=model_name,  # Same as target since we're loading directly
            use_remote_vae=False,
        )

        if vae_name and vae_name != "None" and vae_name != "Default":
            print(f"[DashboardUpscale] \U0001f504 Loading VAE: {vae_name}")
            loaded_vae = load_vae_by_name(vae_name)

        # One cancellable clone for the whole job: every sampler step (hires, tiled hires) checks the event.
        cancellable = _cancellable_model(loaded_model, job["cancel"])

        # Get PromptServer for sending events
        PromptServer = None
        try:
            import server as comfy_server
            PromptServer = comfy_server.PromptServer
            # Ensure last_prompt_id exists — KSampler progress callbacks require it
            if PromptServer.instance is not None and not hasattr(PromptServer.instance, 'last_prompt_id'):
                PromptServer.instance.last_prompt_id = f"dashboard_upscale_{job['id']}"
        except Exception:
            pass

        ctx = {
            "job": job, "upscale_config": upscale_config, "meta": meta, "manifest_data": manifest_data,
            "images_dir": images_dir, "session_name": session_name, "fmt": _fmt, "ext": _ext,
            "vae": loaded_vae, "model": cancellable, "clip": loaded_clip,
            "conditioning_cache": {"positive": {}, "negative": {}},
            "pack": bool((meta.get("session_settings") or {}).get("pack_nodes_metadata")),
        }
        started = time.time()
        total = len(target_items)

        def on_item_done(item, produced):
            save_manifest(manifest_path, manifest_data)
            # Send exactly the entries this item produced (they were inserted at index 0)
            _send(PromptServer, "ultimate_grid.update_data", {
                "session_name": session_name,
                "new_items": list(produced),
            })

        def on_progress(done, total_):
            _send(PromptServer, "ultimate_grid.progress", {
                "session_name": session_name,
                "current_job": done,
                "total_jobs": total_,
                "progress_pct": int((done / total_) * 100),
                "eta_str": "",
                "avg_duration": 0,
                "last_duration": 0,
                "phase": "upscaling",
            })
            print(f"[DashboardUpscale] \U0001f504 Upscaled {done}/{total_}: {job.get('current_image', '')}")

        def on_finish(state):
            done = job["completed"]
            if state == "cancelled":
                _send(PromptServer, "ultimate_grid.progress", {
                    "session_name": session_name,
                    "phase": "upscaling",
                    "state": "cancelled",
                    "current_job": done,
                    "total_jobs": total,
                    "progress_pct": int((done / total) * 100) if total else 0,
                    "eta_str": "",
                    "avg_duration": 0,
                    "last_duration": 0,
                })
            else:
                elapsed = time.time() - started
                _send(PromptServer, "ultimate_grid.progress", {
                    "session_name": session_name,
                    "phase": "upscaling",
                    "complete": True,
                    "total_generated": job.get("outputs", 0),
                    "total_elapsed": f"{int(elapsed)}s",
                    "avg_duration": round(elapsed / total, 1) if total else 0,
                })

        state = _drive_upscale_items(
            job, target_items,
            lambda item, produced: _upscale_one_item(ctx, item, produced),
            on_item_done, on_progress, on_finish)

        if state == "cancelled":
            print(f"[DashboardUpscale] \U0001f6d1 Cancelled: {job['completed']}/{total} images done, "
                  f"{job.get('outputs', 0)} new images kept")
        else:
            print(f"[DashboardUpscale] ✅ Complete: {job['completed']}/{job['total']} images upscaled")

    except Exception as e:
        import traceback
        traceback.print_exc()
        job["status"] = "error"
        job["error"] = str(e)
        job["finished_at"] = time.time()
        print(f"[DashboardUpscale] ❌ Error: {e}")
    finally:
        # Always attempt cleanup on any exit path
        _clean_vram()


def _upscale_one_item(ctx, item, produced):
    """Upscale one manifest image through every active pipeline. Every manifest entry created is
    inserted at index 0 of manifest_data["items"] AND appended to `produced`. Raises UpscaleCancelled."""
    import itertools as up_itertools
    import random as up_random

    job = ctx["job"]
    upscale_config = ctx["upscale_config"]
    meta = ctx["meta"]
    manifest_data = ctx["manifest_data"]
    images_dir = ctx["images_dir"]
    session_name = ctx["session_name"]
    _fmt, _ext, pack = ctx["fmt"], ctx["ext"], ctx["pack"]
    loaded_vae = ctx["vae"]
    patched_model = ctx["model"]
    patched_clip = ctx["clip"]
    conditioning_cache = ctx["conditioning_cache"]
    pipelines = upscale_config.get("pipelines", [])

    # Never inherit the source's metrics; recompute them only when the source had
    # them. This thread must not read generation-run state, so the spec is explicit.
    item_metrics_spec = image_metrics.spec_from_item(item)
    _check_cancel(job)

    # Extract filename from manifest item — may be in "filename" or parsed from "file" URL
    filename = item.get("filename", "")
    if not filename and item.get("file"):
        fname_match = re.search(r'filename=([^&]+)', item["file"])
        if fname_match:
            filename = fname_match.group(1)
    job["current_image"] = filename
    image_path = os.path.join(images_dir, filename)

    if not os.path.exists(image_path):
        print(f"[DashboardUpscale] ⚠️ File not found: {filename}, skipping")
        return

    # Load image and VAE-encode to latent
    try:
        pil_image = PILImage.open(image_path).convert("RGB")
        img_array = np.array(pil_image).astype(np.float32) / 255.0
        img_tensor = torch.from_numpy(img_array).unsqueeze(0)
        encoded = loaded_vae.encode(img_tensor[:, :, :, :3])
        result_latent = {"samples": encoded}
    except Exception as e:
        print(f"[DashboardUpscale] ⚠️ Failed to load/encode {filename}: {e}")
        return

    # Get conditioning
    pos_prompt = item.get("positive", meta.get("positive", ""))
    neg_prompt = item.get("negative", meta.get("negative", ""))
    clip_skip = item.get("clip_skip", 0)

    # One encoder load for the pair (ComfyUI re-prepares the TE on every encode otherwise)
    final_positive, final_negative = encode_pair_cached(
        patched_clip, conditioning_cache, pos_prompt, neg_prompt, clip_skip
    )

    # HiRes prompt adjustment
    hires_positive_cond = final_positive
    hires_prompt_active = False
    hires_prompt_behavior_rt = ""
    hires_prompt_text_rt = ""
    if upscale_config.get("hires_prompt_adjust") and upscale_config.get("hires_prompt_text", "").strip():
        hires_prompt_behavior_rt = upscale_config.get("hires_prompt_behavior", "append_end")
        hires_prompt_text_rt = upscale_config["hires_prompt_text"].strip()
        if hires_prompt_behavior_rt == "prepend":
            adjusted = hires_prompt_text_rt + " " + pos_prompt
        elif hires_prompt_behavior_rt == "append_end":
            adjusted = append_to_base(pos_prompt, hires_prompt_text_rt)
        elif hires_prompt_behavior_rt == "replace":
            adjusted = hires_prompt_text_rt
        else:
            adjusted = pos_prompt
        if adjusted not in conditioning_cache["positive"]:
            conditioning_cache["positive"][adjusted] = encode_prompt_with_combinators(patched_clip, adjusted, clip_skip)
        hires_positive_cond = conditioning_cache["positive"][adjusted]
        hires_prompt_active = True

    pipe_w = item.get("width", 512)
    pipe_h = item.get("height", 512)
    total_upscale_duration = 0

    for pipeline_idx, pipeline in enumerate(pipelines):
        _check_cancel(job)
        if pipeline.get("active", True) is False:
            continue
        pipeline_name = pipeline.get("name", f"Pipeline {pipeline_idx + 1}")
        pipeline_steps = pipeline.get("steps", [])
        if not pipeline_steps:
            continue

        pipe_latent = result_latent
        pipe_w_current = pipe_w
        pipe_h_current = pipe_h

        # Flatten steps with repeat
        expanded_steps = []
        for step in pipeline_steps:
            if step.get("active", True) is False:
                continue
            repeat = max(1, int(step.get("repeat", 1)))
            for _ in range(repeat):
                expanded_steps.append(step)

        for step_idx, ucfg in enumerate(expanded_steps):
            _check_cancel(job)
            mode = ucfg.get("mode", "hires_only")
            show_hires = mode in ("hires_only", "model_then_hires")
            show_model = mode in ("model_only", "model_then_hires")
            step_label = f"{pipeline_name} · step {step_idx + 1}/{len(expanded_steps)}"
            job["current_step"] = step_label

            # --- SeedVR2 upscale mode ---
            if mode == "seedvr2":
                from .image_generation import seedvr2_upscale
                sv_config = ucfg.get("seedvr2", {})
                _check_cancel(job)   # external node call: no way to interrupt it once started
                # Convert source PIL image directly (no latent decoding needed)
                pil_input = PILImage.open(os.path.join(images_dir, filename))
                result_pil, up_w, up_h, sv_duration = seedvr2_upscale(pil_input, sv_config)
                total_upscale_duration += sv_duration

                is_last_step = step_idx == len(expanded_steps) - 1
                if is_last_step:
                    # Create manifest entry, then save the upscaled image
                    upscale_id = int(time.time() * 100000) + up_random.randint(0, 1000)
                    upscaled_filename = f"img_{upscale_id}_upscaled.{_ext}"
                    _up_path = os.path.join(images_dir, upscaled_filename)

                    upscaled_meta = {
                        k: v for k, v in item.items()
                        if k not in ("id", "gen_index", "file", "filename", "upscaled", "width", "height", "duration",
                                     "upscale_source", "upscale_pipeline", "upscale_mode",
                                     "upscale_ratio", "upscale_denoise", "upscale_model", "metrics")
                    }
                    upscaled_meta.update({
                        "id": upscale_id,
                        "gen_index": len(manifest_data["items"]),
                        "file": f"/view?filename={upscaled_filename}&type=output&subfolder=benchmarks/{session_name}/images",
                        "filename": upscaled_filename,
                        "width": up_w, "height": up_h,
                        "duration": round(sv_duration, 2),
                        "upscaled": True,
                        **base_dims_of(item),
                        "upscale_source": "dashboard",
                        "upscale_pipeline": pipeline_name,
                        "upscale_mode": "seedvr2",
                        "upscale_model": sv_config.get("dit_model", ""),
                    })
                    _save_and_stamp(result_pil, _up_path, _fmt, upscaled_meta, pack)
                    image_metrics.finalize(upscaled_meta, result_pil, spec=item_metrics_spec)
                    manifest_data["items"].insert(0, upscaled_meta)
                    produced.append(upscaled_meta)
                else:
                    # For chained steps, the next step gets the PIL result
                    pil_input = result_pil
                    pipe_w_current = up_w
                    pipe_h_current = up_h
                continue  # Skip the normal combo loop

            # --- Florence2 Hi-Res Fix mode ---
            if mode == "florence2_hires":
                from .florence2_hires import (
                    run_florence2_step, preflight_florence2,
                    build_florence2_manifest_entry, build_florence2_no_detection_entry,
                )
                import numpy as _np

                # Preflight once per job (cached)
                if not job.get("_florence2_preflighted"):
                    preflight_florence2()
                    job["_florence2_preflighted"] = True

                # Merge top-level step sampling fields into a unified config dict.
                # hires_denoise on the step may be "0.45, 0.65" — Florence2 uses just one
                # value per pass; take the first.
                _hd = str(ucfg.get("hires_denoise", "0.45")).split(",")[0].strip() or "0.45"
                f2_config = {
                    **(ucfg.get("florence2", {}) or {}),
                    "hires_denoise": float(_hd),
                    "hires_steps": int(ucfg.get("hires_steps", 15) or 15),
                    "cfg": float(ucfg.get("cfg", 1.5) or 1.5),
                    "sampler": ucfg.get("sampler", "euler"),
                    "scheduler": ucfg.get("scheduler", "simple"),
                }

                # Source image: load fresh from disk as tensor (1, H, W, 3)
                src_pil = PILImage.open(os.path.join(images_dir, filename)).convert("RGB")
                src_arr = _np.array(src_pil).astype(_np.float32) / 255.0
                src_tensor = torch.from_numpy(src_arr).unsqueeze(0)

                # Per-job checkpoint+lora cache (lazy init)
                if "_florence2_ckpt_cache" not in job:
                    job["_florence2_ckpt_cache"] = {}

                _check_cancel(job)
                step_result = run_florence2_step(
                    source_image=src_tensor,
                    item=item,
                    step_config=f2_config,
                    fallback_handles=(patched_model, patched_clip, loaded_vae),
                    ckpt_cache=job["_florence2_ckpt_cache"],
                    conditioning_cache=conditioning_cache,
                    positive_prompt=pos_prompt,
                    negative_prompt=neg_prompt,
                    clip_skip=clip_skip,
                    # Pass the session's already-loaded model name so Florence2 can
                    # short-circuit duplicate loads when item.model matches it.
                    session_model_name=meta.get("model", ""),
                    # Sampling inside the step checks the job's cancel event between steps
                    model_wrap=lambda m: _cancellable_model(m, job["cancel"]),
                )

                is_last_step = step_idx == len(expanded_steps) - 1

                if step_result.get("status") == "no_detection":
                    sentinel_id = int(time.time() * 100000) + up_random.randint(0, 1000)
                    entry = build_florence2_no_detection_entry(
                        step_result, item,
                        sentinel_id=sentinel_id,
                        current_index=len(manifest_data["items"]),
                        base_dims=base_dims_of(item),
                    )
                    manifest_data["items"].insert(0, entry)
                    produced.append(entry)
                    total_upscale_duration += float(step_result.get("duration", 0))
                    continue

                total_upscale_duration += float(step_result.get("duration", 0))

                if is_last_step:
                    upscale_id = int(time.time() * 100000) + up_random.randint(0, 1000)
                    upscaled_filename = f"img_{upscale_id}_upscaled.{_ext}"
                    _up_path = os.path.join(images_dir, upscaled_filename)
                    result_pil = step_result["image_pil"]

                    entry = build_florence2_manifest_entry(
                        step_result, item,
                        session_name=session_name,
                        pipeline_name=pipeline_name,
                        upscale_id=upscale_id,
                        upscaled_filename=upscaled_filename,
                        current_index=len(manifest_data["items"]),
                        hires_denoise=f2_config["hires_denoise"],
                        base_dims=base_dims_of(item),
                    )
                    _save_and_stamp(result_pil, _up_path, _fmt, entry, pack)
                    image_metrics.finalize(entry, result_pil, spec=item_metrics_spec)
                    manifest_data["items"].insert(0, entry)
                    produced.append(entry)
                else:
                    pipe_w_current = step_result["image_width"]
                    pipe_h_current = step_result["image_height"]
                continue  # Skip the normal combo loop

            raw_ratios = str(ucfg.get("upscale_ratios", "1.5"))
            ratios = [float(r.strip()) for r in raw_ratios.split(",") if r.strip()] or [1.5]
            raw_denoise = str(ucfg.get("hires_denoise", "0.3"))
            denoises = [float(d.strip()) for d in raw_denoise.split(",") if d.strip()] or [0.3]
            models = ucfg.get("upscale_models", []) or [""]

            if show_hires and show_model:
                combos = list(up_itertools.product(models, ratios, denoises))
            elif show_hires:
                combos = list(up_itertools.product([""], ratios, denoises))
            elif show_model:
                combos = list(up_itertools.product(models, [1.0], [0.0]))
            else:
                combos = []

            for combo_idx, combo in enumerate(combos):
                _check_cancel(job)
                job["current_step"] = f"{step_label} · combo {combo_idx + 1}/{len(combos)}"
                up_model, up_ratio, up_denoise = combo
                single_config = {
                    "mode": mode, "upscale_model": up_model,
                    "upscale_ratio": up_ratio, "hires_denoise": up_denoise,
                    "hires_steps": ucfg.get("hires_steps", 0),
                    "tiled_vae": ucfg.get("tiled_vae", False),
                    "tile_size": ucfg.get("tile_size", 512),
                    "upscale_size": ucfg.get("upscale_size", "2.0"),
                    "resize_method": ucfg.get("resize_method", "bilinear"),
                    "hires_tiled_sampling": ucfg.get("hires_tiled_sampling", False),
                    "hires_tile_width": ucfg.get("hires_tile_width", 512),
                    "hires_tile_height": ucfg.get("hires_tile_height", 512),
                    "hires_mask_blur": ucfg.get("hires_mask_blur", 8),
                    "hires_tile_padding": ucfg.get("hires_tile_padding", 32),
                    "hires_force_uniform_tiles": ucfg.get("hires_force_uniform_tiles", False)
                }

                up_positive = hires_positive_cond if (show_hires and hires_prompt_active) else final_positive
                upscale_result, upscale_duration = upscale_image(
                    pipe_latent, loaded_vae, patched_model, single_config,
                    item, up_positive, final_negative, pipe_w_current, pipe_h_current
                )

                is_last_step = step_idx == len(expanded_steps) - 1
                is_last_combo = combo == combos[-1]
                is_final_output = is_last_step and is_last_combo

                up_metrics_sink = {}
                up_metrics_img = None
                if isinstance(upscale_result, dict) and "samples" in upscale_result:
                    up_metrics_img = upscaled_pil = decode_latent_with_vae(
                        loaded_vae, upscale_result["samples"],
                        metrics_sink=up_metrics_sink if (is_final_output and item_metrics_spec) else None,
                        metrics_spec=item_metrics_spec)
                    up_w, up_h = upscaled_pil.size
                    out_pil = upscaled_pil
                elif isinstance(upscale_result, PILImage.Image):
                    up_metrics_img = upscale_result
                    up_w, up_h = upscale_result.size
                    out_pil = upscale_result
                else:
                    continue

                if not is_final_output:
                    if isinstance(upscale_result, dict) and "samples" in upscale_result:
                        pipe_latent = upscale_result
                    pipe_w_current = up_w
                    pipe_h_current = up_h
                    continue

                upscale_id = int(time.time() * 100000) + up_random.randint(0, 1000)
                upscaled_filename = f"img_{upscale_id}_upscaled.{_ext}"
                _up_path = os.path.join(images_dir, upscaled_filename)

                # Create manifest entry — copy source item fields, override with upscale info
                upscaled_meta = {
                    k: v for k, v in item.items()
                    if k not in ("id", "gen_index", "file", "filename", "upscaled", "width", "height", "duration",
                                 "upscale_source", "upscale_pipeline", "upscale_mode",
                                 "upscale_ratio", "upscale_denoise", "upscale_model", "metrics")
                }
                # gen_index: place upscaled images at the end (newest) since they were just created
                # Use current manifest length so each new upscale gets the next sequential position
                upscale_gen_index = len(manifest_data["items"])

                upscaled_meta.update({
                    "id": upscale_id,
                    "gen_index": upscale_gen_index,
                    "file": f"/view?filename={upscaled_filename}&type=output&subfolder=benchmarks/{session_name}/images",
                    "filename": upscaled_filename,
                    "width": up_w,
                    "height": up_h,
                    "duration": round(upscale_duration, 2),
                    "upscaled": True,
                    **base_dims_of(item),
                    "upscale_source": "dashboard",
                    "upscale_pipeline": pipeline_name,
                    "upscale_mode": mode,
                    "upscale_ratio": up_ratio,
                    "upscale_denoise": up_denoise,
                })
                if up_model:
                    upscaled_meta["upscale_model"] = up_model
                if hires_prompt_active:
                    upscaled_meta["hires_prompt_behavior"] = hires_prompt_behavior_rt
                    upscaled_meta["hires_prompt_text"] = hires_prompt_text_rt

                _save_and_stamp(out_pil, _up_path, _fmt, upscaled_meta, pack)
                image_metrics.finalize(upscaled_meta, up_metrics_img, up_metrics_sink, spec=item_metrics_spec)
                manifest_data["items"].insert(0, upscaled_meta)
                produced.append(upscaled_meta)