mage_flow/pipeline.py

"""MageFlow text-to-image + image-edit inference pipeline.

Self-contained MageFlow t2i / edit inference: load a HuggingFace diffusers-style
repo (model_index.json + transformer/ vae/ scheduler/), then generate or edit
images. No training/eval deps.

Both ``generate_images`` and ``generate_edits`` support PACKED multi-resolution
inference: several samples (each at its own resolution) are concatenated into a
single varlen sequence and processed in one transformer forward per denoise
step. Per-sample ``cu_seqlens`` (inside the flash-attn varlen kernel) isolate
samples, exactly mirroring training-time packing. These packed functions are the
sole implementation — the single-image case is just a pack of size 1, exposed
via the ``MageFlowPipeline.generate`` / ``.edit`` convenience methods.
"""

from __future__ import annotations

import json
import os
import random

import torch
from einops import rearrange
from PIL import Image

from diffusers import FlowMatchEulerDiscreteScheduler

from .models.mage_flow import MageFlowModel, ModelConfig
from .models.utils import PROMPT_TEMPLATE, get_noise, unpack
from .models.modules.mage_text import make_refusal_image
from .models.modules.mage_latent import encode_noise, resolve_gs_key


# ---------------------------------------------------------------------------
# Scheduler — diffusers FlowMatchEulerDiscreteScheduler
# ---------------------------------------------------------------------------
def build_scheduler(num_steps: int, device=None, shift: float = 6.0):
    """Construct a diffusers ``FlowMatchEulerDiscreteScheduler`` whose sigma
    schedule reproduces our default preset exactly.

    The base sigmas ``linspace(1, 1/num_steps, num_steps)`` fed to
    ``set_timesteps`` are run through the scheduler's built-in static shift
    ``shift·s/(1+(shift-1)·s)`` and a terminal 0 is appended — the static-shift 
    schedule (the only supported schedule).
    """
    scheduler = FlowMatchEulerDiscreteScheduler(
        num_train_timesteps=1000, shift=shift, use_dynamic_shifting=False)
    base_sigmas = torch.linspace(1.0, 1.0 / num_steps, num_steps).tolist()
    scheduler.set_timesteps(sigmas=base_sigmas, device=device)
    return scheduler


def _get_scheduler(model, steps, device, static_shift):
    scheduler = getattr(model, "scheduler", None)
    if scheduler is None:
        return build_scheduler(steps, device=device,
                               shift=(static_shift if static_shift is not None else 6.0))
    if static_shift is not None:
        scheduler.set_shift(static_shift)
    scheduler.set_timesteps(sigmas=torch.linspace(1.0, 1.0 / steps, steps).tolist(), device=device)
    return scheduler


# ---------------------------------------------------------------------------
# Small helpers
# ---------------------------------------------------------------------------
def _template_info(name: str | None) -> dict:
    name = name or "mage-flow"
    if name not in PROMPT_TEMPLATE:
        raise ValueError(f"Unknown prompt template: {name}")
    return PROMPT_TEMPLATE[name]


def _as_list(val, default, n):
    """Broadcast a scalar/None to a length-n list, or validate a given list."""
    if val is None:
        return [default] * n
    if isinstance(val, (list, tuple)):
        if len(val) != n:
            raise ValueError(f"expected {n} values, got {len(val)}")
        return list(val)
    return [val] * n


def _lens_to_cu(lens, device):
    """Sequence lengths -> cumulative cu_seqlens [0, l0, l0+l1, ...] (int32)."""
    t = torch.tensor(lens, device=device, dtype=torch.int32)
    return torch.cat([torch.zeros(1, dtype=torch.int32, device=device),
                      torch.cumsum(t, dim=0, dtype=torch.int32)])


def _make_divisible_by_16(size: int) -> int:
    return max(16, 16 * (size // 16))


def _compute_aspect_ratio_size(pil_img: Image.Image, max_size: int):
    """Longest side = ``max_size``, short side from aspect ratio, both /16."""
    w, h = pil_img.size
    if h >= w:
        new_h, new_w = max_size, int(round(w * max_size / h))
    else:
        new_w, new_h = max_size, int(round(h * max_size / w))
    return _make_divisible_by_16(new_h), _make_divisible_by_16(new_w)


def _edit_target_size(pil_img: Image.Image, max_size, height, width):
    """Output (H, W) for an edit sample, derived from its PRIMARY reference.

    Precedence: explicit ``height`` AND ``width`` (custom size) > ``max_size``
    (longest side, short side by aspect ratio) > the source image's own size.
    All rounded down to a multiple of 16.
    """
    if height and width:
        return _make_divisible_by_16(height), _make_divisible_by_16(width)
    if max_size:
        return _compute_aspect_ratio_size(pil_img, max_size)
    # Nothing specified: keep the source resolution (its own longest side).
    return _compute_aspect_ratio_size(pil_img, max(pil_img.size))


def _decode_one(model, tokens, height, width, dev):
    """Unpack one sample's image tokens [1, H*W, C] and VAE-decode to a PIL image."""
    with torch.autocast(device_type=dev.type, dtype=torch.bfloat16):
        out = model.vae.decode(unpack(tokens.float(), height, width))
    out = rearrange(out.clamp(-1, 1), "b c h w -> b h w c")
    out = (127.5 * (out + 1.0)).cpu().byte().numpy()
    return Image.fromarray(out[0])


def _build_pack_ctx(img_ids, img_cu, img_shapes, img_lens, txt, txt_cu, txt_mask, vec,
                    neg_txt, neg_cu, neg_mask, neg_vec, cfg, renormalization, batch_cfg, device):
    """Precompute the static per-step transformer inputs for a packed batch.

    When a negative branch is present and ``batch_cfg`` is True, the conditional
    and unconditional passes are fused into ONE varlen forward: the image tokens
    are duplicated (cond copy + uncond copy) and the positive/negative texts are
    concatenated, so cond sample i and uncond sample i become two independent
    varlen segments processed in a single kernel launch. flash_attn_varlen_func
    keeps every segment isolated via cu_seqlens, so this is numerically identical
    to two separate forwards — just one launch instead of two.
    """
    na = len(img_lens)
    ctx = {
        "na": na, "cfg": cfg, "renorm": renormalization, "batch_cfg": batch_cfg,
        "has_neg": neg_txt is not None,
        "img_ids": img_ids, "img_cu": img_cu, "img_shapes": img_shapes,
        "img_max": int(max(img_lens)),
        "txt": txt, "txt_ids": torch.zeros(1, txt.shape[1], 3, device=device),
        "txt_cu": txt_cu, "txt_mask": txt_mask, "vec": vec,
        "txt_max": int((txt_cu[1:] - txt_cu[:-1]).max().item()),
    }
    if neg_txt is None:
        return ctx
    ctx.update({
        "neg_txt": neg_txt, "neg_ids": torch.zeros(1, neg_txt.shape[1], 3, device=device),
        "neg_cu": neg_cu, "neg_mask": neg_mask, "neg_vec": neg_vec,
        "neg_max": int((neg_cu[1:] - neg_cu[:-1]).max().item()),
    })
    if batch_cfg:
        # Duplicate image segments (cond then uncond) and concat pos+neg text.
        d_txt = torch.cat([txt, neg_txt], dim=1)
        pos_lens = (txt_cu[1:] - txt_cu[:-1]).tolist()
        neg_lens = (neg_cu[1:] - neg_cu[:-1]).tolist()
        ctx.update({
            "d_img_ids": torch.cat([img_ids, img_ids], dim=1),
            "d_img_cu": _lens_to_cu(list(img_lens) + list(img_lens), device),
            "d_img_shapes": [img_shapes[0] + img_shapes[0]],
            "d_txt": d_txt,
            "d_txt_ids": torch.zeros(1, d_txt.shape[1], 3, device=device),
            "d_txt_cu": _lens_to_cu(pos_lens + neg_lens, device),
            "d_txt_mask": torch.ones(1, d_txt.shape[1], device=device),
            "d_vec": torch.cat([vec, neg_vec], dim=0),
            "d_txt_max": int(max(pos_lens + neg_lens)),
        })
    return ctx


def _velocity(transformer, img, ctx, sigma):
    """CFG-combined image-token velocity for a packed batch at noise level ``sigma``.

    Returns [1, sum_img_len, C] in the conditional sample order. When
    ``batch_cfg`` is set the cond+uncond passes share a single fused varlen
    forward; otherwise they are two forwards.
    """
    dev = img.device
    na = ctx["na"]

    def _fwd(x, n, img_ids, img_cu, img_max, img_shapes, txt, txt_ids, txt_cu, txt_mask, txt_max, vec):
        t_vec = torch.full((n,), sigma, dtype=x.dtype, device=dev)
        return transformer(img=x, txt=txt, timesteps=t_vec, img_shapes=img_shapes,
                           img_cu_seqlens=img_cu, txt_cu_seqlens=txt_cu)

    if not ctx["has_neg"]:
        return _fwd(img, na, ctx["img_ids"], ctx["img_cu"], ctx["img_max"], ctx["img_shapes"],
                    ctx["txt"], ctx["txt_ids"], ctx["txt_cu"], ctx["txt_mask"], ctx["txt_max"], ctx["vec"])

    if ctx["batch_cfg"]:
        n_img = img.shape[1]
        out = _fwd(torch.cat([img, img], dim=1), 2 * na,
                   ctx["d_img_ids"], ctx["d_img_cu"], ctx["img_max"], ctx["d_img_shapes"],
                   ctx["d_txt"], ctx["d_txt_ids"], ctx["d_txt_cu"], ctx["d_txt_mask"], ctx["d_txt_max"], ctx["d_vec"])
        cond, unc = out[:, :n_img, :], out[:, n_img:, :]
    else:
        cond = _fwd(img, na, ctx["img_ids"], ctx["img_cu"], ctx["img_max"], ctx["img_shapes"],
                    ctx["txt"], ctx["txt_ids"], ctx["txt_cu"], ctx["txt_mask"], ctx["txt_max"], ctx["vec"])
        unc = _fwd(img, na, ctx["img_ids"], ctx["img_cu"], ctx["img_max"], ctx["img_shapes"],
                   ctx["neg_txt"], ctx["neg_ids"], ctx["neg_cu"], ctx["neg_mask"], ctx["neg_max"], ctx["neg_vec"])

    cfg = ctx["cfg"]
    if ctx["renorm"]:
        # CFG renormalization: rescale the guided velocity per token back to the
        # conditional velocity's norm (reduces oversaturation at high cfg).
        comb = unc + cfg * (cond - unc)
        return comb * (torch.norm(cond, dim=-1, keepdim=True) /
                       (torch.norm(comb, dim=-1, keepdim=True) + 1e-6))
    return unc + cfg * (cond - unc)


def _encode_texts_packed(model, prompts, template, drop_idx, device):
    """Encode a LIST of templated text-only prompts in ONE packed varlen forward
    (``TextEncoder.forward`` — varlen cu_seqlens isolates each prompt,
    verified zero cross-contamination). Returns (txt_flat [ΣLi, D], vec [N, D],
    per-prompt token lengths list)."""
    tokenizer = model.txt_enc.tokenizer
    max_len = model.txt_enc.tokenizer_max_length + drop_idx
    ids_list = [
        tokenizer(template.format(p), max_length=max_len, truncation=True,
                  return_tensors="pt").input_ids.squeeze(0)
        for p in prompts
    ]
    input_ids = torch.cat(ids_list).to(device)
    cu_seqlens = _lens_to_cu([int(t.numel()) for t in ids_list], device)
    res = model.txt_enc(
        input_ids, cu_seqlens, drop_idx_override=drop_idx)
    return res["txt"], res["vec"], res["txt_seq_lens"].tolist()


def _slice_packed(txt_flat, vec, lens, start, count, device):
    """Format a contiguous ``count``-prompt slice (starting at prompt ``start``) of a
    packed text encode into the (txt [1, ΣL, D], cu_seqlens, ones-mask, vec [count, D])
    tuple that ``_build_pack_ctx`` consumes."""
    seg_lens = lens[start:start + count]
    tok_start = sum(lens[:start])
    tok_end = tok_start + sum(seg_lens)
    txt = txt_flat[tok_start:tok_end].reshape(1, -1, txt_flat.shape[-1]).to(device)
    return (txt, _lens_to_cu(seg_lens, device),
            torch.ones(1, txt.shape[1], device=device), vec[start:start + count].to(device))


# ---------------------------------------------------------------------------
# Text-to-image (packed, multi-resolution)
# ---------------------------------------------------------------------------
@torch.no_grad()
def generate_images(model, prompts, neg_prompts=None, seeds=None, steps=30, cfg=5.0,
                    heights=None, widths=None, device="cuda",
                    prompt_template="mage-flow", static_shift=None,
                    gs_key=None,
                    renormalization=False, batch_cfg=True):
    """Generate one image per prompt. Prompts may request DIFFERENT resolutions;
    all are packed into a single varlen forward per denoise step — samples are
    kept isolated by ``flash_attn_varlen_func`` via per-sample ``cu_seqlens`` (no
    cross-sample attention), mirroring training-time packing. When ``cfg > 1`` and
    ``batch_cfg`` is set, the positive and negative passes are fused into that
    same varlen forward. Returns a list of PIL images aligned with ``prompts``.
    """
    if isinstance(prompts, str):
        prompts = [prompts]
    n = len(prompts)
    neg_prompts = _as_list(neg_prompts, " ", n)
    seeds = _as_list(seeds, 42, n)
    heights = _as_list(heights, 1024, n)
    widths = _as_list(widths, 1024, n)
    info = _template_info(prompt_template)
    template = info.get("template", "{}")
    drop_idx = int(info.get("start_idx", 0))
    dev = torch.device(device)

    # Content-policy gate per sample (MANDATORY — runs on the same text-encoder
    # weights as conditioning, no opt-out). Violating prompts get a refusal
    # placeholder and are dropped from the pack.
    results = [None] * n
    active = []
    for i in range(n):
        if seeds[i] == -1:
            seeds[i] = random.randint(0, 2**32 - 1)
        verdict = model.txt_enc.screen_text(prompts[i])
        if verdict.violates:
            h_, w_ = _make_divisible_by_16(heights[i]), _make_divisible_by_16(widths[i])
            results[i] = make_refusal_image(verdict, height=h_, width=w_)
            continue
        active.append(i)
    if not active:
        return results

    gs_key_int = resolve_gs_key(gs_key)
    # Per-sample noise tokens + position ids + shapes (MageVAE: flatten, no packing).
    ch = model.vae.latent_channels
    img_list, ids_list, lens, shapes, hw = [], [], [], [], []
    for i in active:
        h_, w_ = _make_divisible_by_16(heights[i]), _make_divisible_by_16(widths[i])
        torch.manual_seed(seeds[i])
        x = get_noise(num_samples=1, channel=ch, height=h_, width=w_,
                      device=dev, dtype=torch.bfloat16, seed=seeds[i])
        # Distribution-preserving watermark in the initial noise (same shape,
        # still ~N(0,1)); detect by inverting the flow ODE back to noise.
        x = encode_noise(tuple(x.shape[1:]), key=gs_key_int,
                         seed=seeds[i], device=dev, dtype=torch.bfloat16)
        _, _, gh, gw = x.shape
        img_list.append(rearrange(x, "b c h w -> b (h w) c")[0])
        ids = torch.zeros(gh, gw, 3, device=dev)
        ids[..., 1] = ids[..., 1] + torch.arange(gh, device=dev)[:, None]
        ids[..., 2] = ids[..., 2] + torch.arange(gw, device=dev)[None, :]
        ids_list.append(rearrange(ids, "h w c -> (h w) c"))
        lens.append(gh * gw); shapes.append((1, gh, gw)); hw.append((h_, w_))
    img = torch.cat(img_list, 0).unsqueeze(0)
    img_ids = torch.cat(ids_list, 0).unsqueeze(0)
    img_cu = _lens_to_cu(lens, dev)
    img_shapes = [shapes]

    # Packed text: positive prompts AND (for CFG) negative prompts are encoded
    # TOGETHER in ONE varlen forward, then split back — cu_seqlens keeps every
    # prompt isolated (verified zero cross-contamination).
    pos_prompts = [prompts[i] for i in active]
    na = len(active)
    use_neg = cfg > 1.0 and any(neg_prompts[i] for i in active)
    if use_neg:
        neg_list = [neg_prompts[i] or " " for i in active]
        txt_flat, vec_all, lens_t = _encode_texts_packed(
            model, pos_prompts + neg_list, template, drop_idx, dev)
        txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, na, dev)
        neg_txt, neg_cu, neg_mask, neg_vec = _slice_packed(txt_flat, vec_all, lens_t, na, na, dev)
    else:
        txt_flat, vec_all, lens_t = _encode_texts_packed(model, pos_prompts, template, drop_idx, dev)
        txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, na, dev)
        neg_txt = neg_cu = neg_mask = neg_vec = None

    ctx = _build_pack_ctx(img_ids, img_cu, img_shapes, lens, txt, txt_cu, txt_mask, vec,
                          neg_txt, neg_cu, neg_mask, neg_vec, cfg, renormalization, batch_cfg, dev)
    scheduler = _get_scheduler(model, steps, device, static_shift)
    for si, t in enumerate(scheduler.timesteps):
        pred = _velocity(model.transformer, img, ctx, scheduler.sigmas[si].item())
        img = scheduler.step(pred, t, img, return_dict=False)[0]

    off = 0
    for k, i in enumerate(active):
        L = lens[k]
        h_, w_ = hw[k]
        results[i] = _decode_one(model, img[:, off:off + L, :], h_, w_, dev)
        off += L
    return results


# ---------------------------------------------------------------------------
# Image edit (packed, multi-resolution)
# ---------------------------------------------------------------------------
def _preprocess_ref_image(pil_img: Image.Image, height: int, width: int, device) -> torch.Tensor:
    """Resize an RGB reference image to (height, width) and normalize to [-1, 1]."""
    from torchvision.transforms import functional as TF
    img = pil_img.convert("RGB")
    img = TF.resize(img, [height, width], interpolation=TF.InterpolationMode.BICUBIC)
    t = TF.to_tensor(img)  # [3, H, W] in [0, 1]
    t = TF.normalize(t, [0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # -> [-1, 1]
    return t.to(device)


def _resize_long_edge(image: Image.Image, max_long_edge: int | None) -> Image.Image:
    """Cap the VL conditioning image's long edge, preserving aspect ratio.

    Matches training's data.processor._resize_long_edge (BICUBIC). Without this,
    inference feeds a full-resolution image to the Qwen-VL processor whose
    default max_pixels is far larger than 384**2 — a train/test mismatch.
    """
    if max_long_edge is None or max_long_edge <= 0:
        return image
    w, h = image.size
    long_edge = max(w, h)
    if long_edge <= max_long_edge:
        return image
    scale = max_long_edge / long_edge
    new_w = max(1, int(round(w * scale)))
    new_h = max(1, int(round(h * scale)))
    return image.resize((new_w, new_h), Image.BICUBIC)


# Fixed image placeholder used at edit training time (one per reference image).
_EDIT_IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"


def _edit_prompt_body(instruction, num_refs):
    """Training-time multi-reference prompt body: ``Image 1: <ph>Image 2: <ph>…{instruction}``."""
    prefix = "".join(f"Image {j}: {_EDIT_IMAGE_PLACEHOLDER}" for j in range(1, num_refs + 1))
    return prefix + instruction


def _encode_edits_packed(model, ref_pils_per_sample, instructions, template, drop_idx, device):
    """Encode ALL image-conditioned edit instructions in ONE packed multimodal
    varlen forward (pixel_values/image_grid_thw concatenated across samples,
    cu_seqlens isolates each). Returns (txt_flat [ΣLi, D], vec [N, D], per-sample lens)."""
    processor = model.txt_enc.processor
    ids_list, pv_list, thw_list = [], [], []
    for ref_pils, instr in zip(ref_pils_per_sample, instructions, strict=False):
        formatted = template.format(_edit_prompt_body(instr, len(ref_pils)))
        vl = processor(text=[formatted], images=list(ref_pils), padding=True, return_tensors="pt")
        vl = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in vl.items()}
        ids_list.append(vl["input_ids"].squeeze(0))
        if vl.get("pixel_values") is not None:
            pv_list.append(vl["pixel_values"]); thw_list.append(vl["image_grid_thw"])
    input_ids = torch.cat(ids_list).to(device)
    cu = _lens_to_cu([int(t.numel()) for t in ids_list], device)
    inputs = {"input_ids": input_ids, "cu_seqlens": cu}
    if pv_list:
        inputs["pixel_values"] = torch.cat(pv_list, dim=0)
        inputs["image_grid_thw"] = torch.cat(thw_list, dim=0)
    res = model.txt_enc(
        input_ids, cu, inputs=inputs, drop_idx_override=drop_idx)
    return res["txt"], res["vec"], res["txt_seq_lens"].tolist()


@torch.no_grad()
def generate_edits(model, prompts, ref_images, neg_prompts=None, seeds=None, steps=30, cfg=5.0,
                   max_size=None, heights=None, widths=None, device="cuda",
                   prompt_template="mage-flow-edit", static_shift=None,
                   gs_key=None,
                   vl_cond_long_edge=384,
                   renormalization=False, batch_cfg=True):
    """Edit reference image(s) per prompt. Each ``ref_images[i]`` may be a single
    image/path OR a list of source images (multi-image edit, like training —
    trained with up to 3, but more are accepted) — all produce ONE edited output. Each sample's
    ``[target, ref_1, …, ref_N]`` latent tokens are sequence-concatenated, and
    all samples are packed into one varlen forward per denoise step.

    Output resolution (derived from the first/primary reference of each sample):
    if both ``heights[i]`` and ``widths[i]`` are given, use them; else if
    ``max_size`` is given, the longest side is ``max_size`` and the short side
    follows the reference's aspect ratio; otherwise the output keeps the source
    image's own resolution. All references are VAE-encoded at that target size.
    Returns a list of PIL images.
    """
    if isinstance(prompts, str):
        prompts = [prompts]
        ref_images = [ref_images]
    n = len(prompts)
    neg_prompts = _as_list(neg_prompts, " ", n)
    seeds = _as_list(seeds, 42, n)
    heights = _as_list(heights, None, n)
    widths = _as_list(widths, None, n)
    info = _template_info(prompt_template)
    template = info.get("template", "{}")
    drop_idx = int(info.get("start_idx", 0))
    dev = torch.device(device)

    # Normalize each sample's references to a list of 1..3 PIL images.
    def _load_pil(r):
        if isinstance(r, str):
            r = Image.open(r)
        return r.convert("RGB")

    pils_per_sample = []
    for r in ref_images:
        refs = list(r) if isinstance(r, (list, tuple)) else [r]
        if not refs:
            raise ValueError("each edit sample needs at least one reference image")
        pils_per_sample.append([_load_pil(x) for x in refs])

    # Per-sample output resolution (from the first/primary reference) + content gate.
    results = [None] * n
    res_hw = [None] * n
    active = []
    for i in range(n):
        res_hw[i] = _edit_target_size(pils_per_sample[i][0], max_size, heights[i], widths[i])
        if seeds[i] == -1:
            seeds[i] = random.randint(0, 2**32 - 1)
        # Multimodal gate (MANDATORY): inspect the source image(s) AND the
        # instruction, so NSFW / copyrighted-character / real-public-figure
        # source photos are blocked even under an innocuous instruction.
        verdict = model.txt_enc.screen_edit(prompts[i], pils_per_sample[i])
        if verdict.violates:
            h_, w_ = res_hw[i]
            results[i] = make_refusal_image(verdict, height=h_, width=w_)
            continue
        active.append(i)
    if not active:
        return results

    gs_key_int = resolve_gs_key(gs_key)

    # Per sample: reference latent tokens (clean) + target noise tokens, plus the
    # combined [target, ref_1, …, ref_N] position ids and shapes. ``target_idx``
    # records where each sample's target tokens land in the packed sequence so we
    # can slice the velocity and step only the target portion.
    ch = model.vae.latent_channels
    targets, refs, ids_list, shape_seq, samp_lens, tgt_lens, hw = [], [], [], [], [], [], []
    target_idx_parts = []
    off = 0
    for i in active:
        h_, w_ = res_hw[i]
        torch.manual_seed(seeds[i])  # MageVAE.encode samples the posterior (global RNG)
        # All references resized to the target resolution and VAE-encoded together.
        ref_tensors = [_preprocess_ref_image(p, h_, w_, dev) for p in pils_per_sample[i]]
        ref_tok, ref_shapes, ref_ids = model.compute_vae_encodings(ref_tensors, with_ids=True)
        ref_tok = ref_tok.to(torch.bfloat16)               # [1, N*Lr, C]
        x = get_noise(num_samples=1, channel=ch, height=h_, width=w_,
                      device=dev, dtype=torch.bfloat16, seed=seeds[i])
        x = encode_noise(tuple(x.shape[1:]), key=gs_key_int,
                         seed=seeds[i], device=dev, dtype=torch.bfloat16)
        _, _, gh, gw = x.shape
        tgt = rearrange(x, "b c h w -> b (h w) c")          # [1, Lt, C]
        tgt_ids = torch.zeros(gh, gw, 3, device=dev)
        tgt_ids[..., 1] = tgt_ids[..., 1] + torch.arange(gh, device=dev)[:, None]
        tgt_ids[..., 2] = tgt_ids[..., 2] + torch.arange(gw, device=dev)[None, :]
        tgt_ids = rearrange(tgt_ids, "h w c -> (h w) c").unsqueeze(0)
        lt, lr = tgt.shape[1], ref_tok.shape[1]
        targets.append(tgt); refs.append(ref_tok)
        ids_list.append(torch.cat([tgt_ids, ref_ids.to(dev)], dim=1)[0])  # [Lt + N*Lr, 3]
        shape_seq.append((1, gh, gw))                       # target frame idx 0
        shape_seq.extend(s[0] for s in ref_shapes)          # ref_j frame idx j
        samp_lens.append(lt + lr); tgt_lens.append(lt); hw.append((h_, w_))
        target_idx_parts.append(torch.arange(off, off + lt, device=dev))
        off += lt + lr
    img_ids = torch.cat(ids_list, 0).unsqueeze(0)
    img_cu = _lens_to_cu(samp_lens, dev)
    img_shapes = [shape_seq]
    target_idx = torch.cat(target_idx_parts)

    # Packed edit text — positive AND (for CFG) negative are encoded TOGETHER in
    # ONE packed multimodal forward, then split. Both branches share the same
    # reference images; cu_seqlens isolates every sequence (zero cross-contamination).
    # The VL conditioning image's long edge is capped (default 384) to match
    # training preprocessing — the VAE path above keeps the full target resolution.
    na = len(active)
    edit_refs = [[_resize_long_edge(p, vl_cond_long_edge) for p in pils_per_sample[i]]
                 for i in active]
    if cfg > 1.0:
        pos_instr = [prompts[i] for i in active]
        neg_instr = [neg_prompts[i] or " " for i in active]
        txt_flat, vec_all, lens_t = _encode_edits_packed(
            model, edit_refs + edit_refs, pos_instr + neg_instr, template, drop_idx, dev)
        txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, na, dev)
        neg_txt, neg_cu, neg_mask, neg_vec = _slice_packed(txt_flat, vec_all, lens_t, na, na, dev)
    else:
        txt_flat, vec_all, lens_t = _encode_edits_packed(
            model, edit_refs, [prompts[i] for i in active], template, drop_idx, dev)
        txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, na, dev)
        neg_txt = neg_cu = neg_mask = neg_vec = None

    ctx = _build_pack_ctx(img_ids, img_cu, img_shapes, samp_lens, txt, txt_cu, txt_mask, vec,
                          neg_txt, neg_cu, neg_mask, neg_vec, cfg, renormalization, batch_cfg, dev)
    scheduler = _get_scheduler(model, steps, device, static_shift)
    for si, t in enumerate(scheduler.timesteps):
        parts = []
        for k in range(na):
            parts.append(targets[k]); parts.append(refs[k])
        img = torch.cat(parts, dim=1)                       # [1, sum(Lt+Lr), C], ref clean
        vel = _velocity(model.transformer, img, ctx, scheduler.sigmas[si].item())
        pred_t = vel[:, target_idx, :]                      # [1, sum Lt, C] — target tokens only
        tgt_packed = torch.cat(targets, dim=1)              # [1, sum Lt, C]
        stepped = scheduler.step(pred_t, t, tgt_packed, return_dict=False)[0]
        o = 0
        new_targets = []
        for k in range(na):
            lt = tgt_lens[k]
            new_targets.append(stepped[:, o:o + lt, :]); o += lt
        targets = new_targets

    for k, i in enumerate(active):
        h_, w_ = hw[k]
        results[i] = _decode_one(model, targets[k], h_, w_, dev)
    return results


# ---------------------------------------------------------------------------
# Flow-ODE inversion (Gaussian-Shading watermark detection)
# ---------------------------------------------------------------------------
@torch.no_grad()
def invert_to_noise(model, z0, height, width, steps=30, device="cuda",
                    prompt_template="mage-flow", static_shift=None, prompt=""):
    """Reverse the flow ODE from a clean latent ``z0`` back to the initial noise.

    This is the detection primitive for the Gaussian-Shading watermark: VAE-encode
    the image to ``z0`` (posterior MEAN — deterministic), run this to recover the
    initial noise, then read the signs via ``mage_latent.decode_bits``.

    Inversion uses an empty prompt at cfg=1 (the standard Tree-Ring /
    Gaussian-Shading setup). Reverse Euler recovers ``x_i`` from ``x_{i+1}`` with
    the velocity evaluated at the point in hand; the sign-only watermark tolerates
    the resulting approximation error (see the module's redundancy).

    Args:
        z0: clean latent ``[1, C, gh, gw]`` (e.g. the mean of ``model.vae.encode``).
    Returns:
        recovered initial-noise latent ``[1, C, gh, gw]`` (float32).
    """
    dev = torch.device(device)
    info = _template_info(prompt_template)
    template = info.get("template", "{}")
    drop_idx = int(info.get("start_idx", 0))

    z0 = z0.to(dev)
    _, ch, gh, gw = z0.shape
    img = rearrange(z0, "b c h w -> b (h w) c").to(torch.bfloat16)   # [1, gh*gw, C]

    ids = torch.zeros(gh, gw, 3, device=dev)
    ids[..., 1] = ids[..., 1] + torch.arange(gh, device=dev)[:, None]
    ids[..., 2] = ids[..., 2] + torch.arange(gw, device=dev)[None, :]
    img_ids = rearrange(ids, "h w c -> (h w) c").unsqueeze(0)
    lens = [gh * gw]
    img_cu = _lens_to_cu(lens, dev)
    img_shapes = [[(1, gh, gw)]]

    # Empty-prompt conditioning, no negative branch, cfg=1 (single forward).
    txt_flat, vec_all, lens_t = _encode_texts_packed(model, [prompt], template, drop_idx, dev)
    txt, txt_cu, txt_mask, vec = _slice_packed(txt_flat, vec_all, lens_t, 0, 1, dev)
    ctx = _build_pack_ctx(img_ids, img_cu, img_shapes, lens, txt, txt_cu, txt_mask, vec,
                          None, None, None, None, 1.0, False, False, dev)

    scheduler = _get_scheduler(model, steps, device, static_shift)
    sigmas = scheduler.sigmas
    n = len(scheduler.timesteps)
    # Forward step si: x_{si+1} = x_si + (s_{si+1}-s_si)·v(x_si, s_si).
    # Reverse it from clean (x_n, sigma 0) up to noise (x_0), using x_{si+1} as the
    # proxy for x_si at the forward eval sigma s_si.
    for si in range(n - 1, -1, -1):
        s_cur = sigmas[si].item()
        s_next = sigmas[si + 1].item()
        vel = _velocity(model.transformer, img, ctx, s_cur)
        img = img - (s_next - s_cur) * vel
    return unpack(img.float(), height, width)   # [1, C, gh, gw]


# ---------------------------------------------------------------------------
# High-level pipeline wrapper
# ---------------------------------------------------------------------------
class MageFlowPipeline:
    """``MageFlowPipeline.from_pretrained(repo).generate(...) / .edit(...)``.

    ``generate`` / ``edit`` are packed multi-resolution calls: they take a list
    of prompts (a single string is accepted and treated as a pack of size 1) and
    return a list of PIL images. Per-sample ``heights``/``widths``/``seeds`` are
    lists. Every prompt is screened by the text encoder's mandatory content
    gate (no opt-out); banned prompts come back as refusal placeholders
    interleaved with the real images. Real outputs always carry a Gaussian-Shading
    watermark in the initial noise (no toggle), using the configured secret key.
    """

    def __init__(self, model, device="cuda"):
        self.model = model
        self.device = device

    @classmethod
    def from_pretrained(cls, repo_dir: str, device: str = "cuda"):
        """Load a Mage-Flow diffusers-style repo (``model_index.json`` +
        ``transformer/`` ``vae/`` ``scheduler/`` ``text_encoder/``).

        ``repo_dir`` may be a local directory OR a Hugging Face Hub repo id
        (e.g. ``"microsoft/Mage-Flow-4B"``), which is downloaded and cached
        automatically on first use.
        """
        return cls(load_from_repo(repo_dir, device), device)

    def generate(self, prompts, **kw) -> list[Image.Image]:
        """Packed multi-resolution t2i. ``prompts`` is a list (or a single
        string); pass per-sample ``heights``/``widths``/``seeds`` as lists."""
        kw.setdefault("device", self.device)
        return generate_images(self.model, prompts, **kw)

    def edit(self, prompts, ref_images, **kw) -> list[Image.Image]:
        """Packed multi-resolution edit. ``prompts`` is a list (or a single
        string); each ``ref_images[i]`` is one reference or a list of references."""
        kw.setdefault("device", self.device)
        return generate_edits(self.model, prompts, ref_images, **kw)

    def invert_to_noise(self, z0, height, width, **kw):
        """Recover the initial noise from a clean latent (Gaussian-Shading detect)."""
        kw.setdefault("device", self.device)
        return invert_to_noise(self.model, z0, height, width, **kw)


def _safe_subpath(root: str, *parts: str) -> str:
    """Join ``parts`` under ``root`` and confirm the result stays inside ``root``.

    ``root`` is normalized up front; the joined path is normalized **lexically**
    (``os.path.normpath`` — symlinks are *not* followed, so a Hugging Face cache
    whose weight files are symlinks into the shared blob store still loads) and
    rejected if it escapes ``root``. This guards the user-supplied model path
    against path traversal (CWE-22 / CodeQL ``py/path-injection``).
    """
    root = os.path.realpath(root)
    full = os.path.normpath(os.path.join(root, *parts))
    if full != root and not full.startswith(root + os.sep):
        raise ValueError(
            f"Resolved path {os.path.join(*parts)!r} escapes repo directory {root!r}"
        )
    return full


def _resolve_repo_dir(repo_dir: str) -> str:
    """Return a local directory for ``repo_dir``.

    If ``repo_dir`` is an existing local path it is returned as a normalized
    absolute path; otherwise it is treated as a Hugging Face Hub repo id (e.g.
    ``microsoft/Mage-Flow``) and downloaded/cached via
    ``huggingface_hub.snapshot_download``.
    """
    candidate = os.path.realpath(repo_dir)
    if os.path.isdir(candidate):
        return candidate
    from huggingface_hub import snapshot_download
    return snapshot_download(repo_id=repo_dir)


def load_from_repo(repo_dir: str, device: str = "cuda") -> MageFlowModel:
    """Load a Mage-Flow diffusers-style repo (model_index.json + transformer/
    vae/ scheduler/). Transformer weights come from the bf16 safetensors;
    VAE + text encoder are built from the sources recorded in model_index.json.

    ``repo_dir`` may be a local directory OR a Hugging Face Hub repo id (e.g.
    ``microsoft/Mage-Flow-4B``), which is downloaded/cached automatically.
    """
    from safetensors.torch import load_file
    repo_dir = _resolve_repo_dir(repo_dir)
    mi = json.load(open(_safe_subpath(repo_dir, "model_index.json")))
    tcfg = json.load(open(_safe_subpath(repo_dir, "transformer", "config.json")))
    # Keys stripped from the checkpoint config before it becomes model_structure.
    # ``schedule_mode`` is a legacy field still present in some config.json files;
    # Keys of the checkpoint config that are NOT MageFlowParams constructor args
    # (legacy/unused fields). Everything else becomes model_structure. The DiT only
    # reads: in_channels, out_channels, context_in_dim, hidden_size, num_heads,
    # depth, axes_dim, checkpoint, patch_size.
    _meta = {"_class_name", "txt_max_length", "max_sequence_length", "param_dtype",
             "packing", "schedule_mode", "static_shift", "use_time_shift",
             "rope_type", "apply_text_rotary_emb",
             "mlp_ratio", "depth_single_blocks", "theta", "qkv_bias", "guidance_embed",
             "vec_in_dim", "vec_type", "time_type", "double_block_type"}
    structure = {k: v for k, v in tcfg.items() if k not in _meta}

    def _resolve(p):
        return p if os.path.isabs(p) else _safe_subpath(repo_dir, p)

    cfg = ModelConfig(
        vae_path=_resolve(mi.get("_vae_source")),
        txt_enc_path=_resolve(mi.get("_text_encoder_path")),
        model_structure=structure,
        txt_max_length=tcfg.get("txt_max_length", 2048),
        packing=tcfg.get("packing", True),
        static_shift=tcfg.get("static_shift", 6.0),
    )
    model = MageFlowModel(cfg)
    sd = load_file(_safe_subpath(repo_dir, "transformer", "diffusion_pytorch_model.safetensors"),
                   device="cpu")
    model.transformer.load_state_dict(sd, strict=False, assign=True)
    model.to(device)
    model.transformer.to(torch.bfloat16)
    model.txt_enc.to(torch.bfloat16)
    if model.vae is not None:
        model.vae.to(torch.bfloat16)
    model.eval()
    # Diffusers FlowMatchEulerDiscreteScheduler (scheduler/scheduler_config.json).
    model.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
        _safe_subpath(repo_dir, "scheduler"))
    return model