seganybridge.py

"""
Script to generate Meta Segment Anything masks.

Adapted from:
https://github.com/facebookresearch/segment-anything-2
https://github.com/facebookresearch/segment-anything

Author: Shrinivas Kulkarni

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License along
with this program; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
"""

import contextlib
import glob
import math
import os
import struct
import sys
import threading
import time


def _logical_cpu_count():
    """CPUs this process may actually run on, not what the box happens to have.

    sched_getaffinity is what respects taskset/cgroup pinning; os.cpu_count()
    reports every CPU on the machine even when we're confined to two of them.
    """
    try:
        return len(os.sched_getaffinity(0))
    except (AttributeError, OSError):
        return os.cpu_count() or 1


def _physical_cpu_count(logical):
    """Best-effort count of physical cores, i.e. ignoring SMT/Hyper-Threading.

    Sizing a BLAS/OpenMP pool by *logical* CPUs is actively counterproductive
    for the dense matmuls SAM spends its time in: SMT siblings share one
    core's vector units and its L1/L2, so the second thread on a core adds
    cache pressure and synchronisation cost without adding arithmetic
    throughput. Stdlib-only on purpose — this module must import cleanly in
    whatever bare interpreter the user pointed the plug-in at, so no psutil.
    Anything we can't determine falls back to the logical count, which is
    exactly the old behaviour.
    """
    count = None
    try:
        if sys.platform.startswith("linux"):
            siblings = set()
            cpuBase = "/sys/devices/system/cpu"
            for entry in os.listdir(cpuBase):
                sibPath = os.path.join(
                    cpuBase, entry, "topology", "thread_siblings_list"
                )
                try:
                    with open(sibPath) as f:
                        siblings.add(f.read().strip())
                except OSError:
                    continue
            count = len(siblings) or None
        elif sys.platform == "darwin":
            import subprocess

            out = subprocess.run(
                ["sysctl", "-n", "hw.physicalcpu"],
                capture_output=True,
                text=True,
                timeout=2,
            )
            count = int(out.stdout.strip())
    except Exception:
        count = None
    # A bogus reading (0, or more "physical" cores than we can even run on)
    # must never quietly halve the thread pool.
    if count is None or not 1 <= count <= logical:
        return logical
    return count


# These BLAS-level thread-pool env vars only take effect if set BEFORE
# numpy/torch/cv2 first touch the underlying OpenBLAS/MKL library, so this
# has to happen here, ahead of those imports. We only set a value the user
# hasn't already pinned themselves (e.g. via a shell profile).
_LOGICAL_CPU_COUNT = _logical_cpu_count()
_CPU_COUNT = _physical_cpu_count(_LOGICAL_CPU_COUNT)
for _env_var in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"):
    os.environ.setdefault(_env_var, str(_CPU_COUNT))

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

# None of SAM1 (`segment_anything`), SAM2 (`sam2`) or SAM3 (`transformers`)
# are imported here at module level — each is imported lazily inside its
# own Strategy.load_model(), and only that one family's package is ever
# required for a given run. This used to import SAM1 and SAM2 eagerly on
# the assumption that the installer "always installs both together,
# cheaply" — but one-click setup only ever installs whichever single
# family is recommended for the machine (see installer.py), and Custom
# install's SAM2 build can fail independently (needs a C/C++ toolchain)
# while leaving SAM1 usable. Eager imports meant either package being
# absent broke BOTH families, including the one that was actually
# installed and working — this is exactly what SAM3 already avoided by
# being lazy, so SAM1/SAM2 now follow the same pattern.

# --- Progress reporting ----------------------------------------------------
#
# GIMP invokes this script as a subprocess and streams its stdout back to the
# user, line by line, as it arrives (see seganyplugin.py). Every long-running
# step below MUST print *something* every few seconds, otherwise the plug-in
# has nothing to show and the whole thing looks hung even though it is
# working perfectly fine — this was the single biggest cause of "GIMP froze,
# nothing ever happens": a multi-minute CPU-bound call with zero stdout.


def stage(name):
    print(f"[stage] {name}", flush=True)


# --- Determinate progress ---------------------------------------------------
#
# The plug-in used to show a bar that only pulsed: motion, but no
# information — after 40 seconds you still couldn't tell whether you were
# nearly done or a third of the way in. These percentages give it something
# real to draw. The phase boundaries below are a rough time budget of a run,
# not equal slices: on a large image the grid search dwarfs everything else,
# so it gets most of the bar.
#
# The bridge owns 0-90% and the plug-in spends the last 10% building layers,
# so the bar only ever moves forwards across the whole operation. Within
# that, the bands are sized by measured time, not by how important a step
# feels: saving masks is a rounding error next to the grid search, so it
# gets a rounding error's worth of bar. Get this wrong in the other
# direction and the time estimate inherits the error.
PCT_MODEL_LOADED = 8.0
PCT_IMAGE_READY = 12.0
PCT_SEGMENT_START = 13.0
# Box/Selection/Text can't be subdivided the way a grid search can: their
# cost is one image-encoder pass, which is opaque from out here. Reporting
# the boundary between "encoding" and "predicting" is the one honest split
# available, and the encoder is the overwhelming majority of it.
PCT_ENCODED = 65.0
PCT_SEGMENT_END = 86.0
PCT_SAVE_END = 90.0


def progress_pct(pct, label):
    """Emit a machine-readable progress line for the plug-in's progress bar.

    Kept deliberately separate from the [progress] heartbeat lines: those
    are for a human reading a terminal, this one is parsed.
    """
    print(f"[progress-pct] {min(pct, 100.0):.1f} {label}", flush=True)


def _expected_batch_count(mask_generator, crop_n_layers, points_per_batch):
    """How many decoder batches an Auto pass will run, or None if unknown.

    One batch per points_per_batch grid points, per crop. Crop layer i holds
    (2**i)**2 crops (generate_crop_boxes' own documented rule) and draws its
    points from point_grids[i], which the generator hands us directly — so
    this stays right even when the grid is downscaled per layer.
    """
    grids = getattr(mask_generator, "point_grids", None)
    if not grids or points_per_batch <= 0:
        return None
    total = 0
    for layer in range(crop_n_layers + 1):
        crops = 1 if layer == 0 else (2**layer) ** 2
        grid = grids[min(layer, len(grids) - 1)]
        total += crops * math.ceil(len(grid) / points_per_batch)
    return total or None


@contextlib.contextmanager
def batch_progress(mask_generator, crop_n_layers, points_per_batch):
    """Report real progress out of the automatic mask generator.

    generate() is one opaque call, but internally it loops over point
    batches through _process_batch — the same name in both SAM1's and
    SAM2's generator. Wrapping it on the *instance* (never on the library)
    turns that loop into a progress signal without touching either package,
    and if a future version renames it we simply fall back to the old
    indeterminate bar instead of breaking segmentation.
    """
    total = _expected_batch_count(mask_generator, crop_n_layers, points_per_batch)
    original = getattr(mask_generator, "_process_batch", None)
    if total is None or original is None:
        dbg("per-batch progress unavailable for this generator — "
            "the progress bar will stay indeterminate")
        yield
        return

    dbg(f"expecting {total} decoder batch(es)")
    done = 0

    def counted(*args, **kwargs):
        nonlocal done
        result = original(*args, **kwargs)
        done += 1
        span = PCT_SEGMENT_END - PCT_SEGMENT_START
        progress_pct(
            PCT_SEGMENT_START + span * min(done / total, 1.0),
            f"segmenting — batch {min(done, total)}/{total}",
        )
        return result

    mask_generator._process_batch = counted
    try:
        yield
    finally:
        # Drop the instance attribute so the class's own method shows again.
        try:
            del mask_generator._process_batch
        except AttributeError:
            pass


def dbg(message):
    """One line per meaningful step, always on.

    The plug-in streams this straight into its progress window and into
    GIMP's Error Console, so when a run dies or hangs the last line printed
    is the answer to "where did it stop?" — which beats reproducing the
    whole thing under a debugger inside GIMP.
    """
    print(f"[debug] {message}", flush=True)


@contextlib.contextmanager
def heartbeat(label, interval=3.0):
    """Print a progress line every `interval` seconds while a long call runs.

    Runs in a daemon thread so it can report progress even while the main
    thread is stuck inside a native (C++/torch) call that never returns to
    the Python interpreter until it's done — those calls still release the
    GIL for the bulk of their work, so the timer thread keeps ticking.
    """
    stop = threading.Event()
    t0 = time.time()

    def _tick():
        while not stop.wait(interval):
            print(f"[progress] {label}: {time.time() - t0:.0f}s elapsed", flush=True)

    t = threading.Thread(target=_tick, daemon=True)
    t.start()
    try:
        yield
    finally:
        stop.set()
        t.join(timeout=1)
        print(f"[progress] {label}: done in {time.time() - t0:.1f}s", flush=True)


# --- Device selection --------------------------------------------------------
#
# Deliberately generic: NVIDIA (CUDA), AMD (ROCm builds of torch report
# through the same torch.cuda.* API), Apple Silicon (MPS) and a CPU fallback
# that is tuned to actually use every core — by default torch sometimes
# leaves threads on the table in containerized/virtualized environments.


def enable_tf32():
    """Let Ampere-and-newer NVIDIA GPUs run fp32 matmuls on the TensorFloat-32
    path: same exponent range as fp32 with a 10-bit mantissa, several times
    the throughput, and no visible difference in a segmentation mask. torch
    2.9 moved this behind a new fp32_precision API and deprecated the
    allow_tf32 flags, so try the new spelling first and fall back."""
    try:
        torch.backends.cuda.matmul.fp32_precision = "tf32"
        torch.backends.cudnn.conv.fp32_precision = "tf32"
        return
    except AttributeError:
        pass
    try:
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
    except AttributeError:
        pass


def pick_device():
    if torch.cuda.is_available():
        name = torch.cuda.get_device_name(0)
        enable_tf32()
        return torch.device("cuda"), f"CUDA/ROCm GPU ({name})"
    mps = getattr(torch.backends, "mps", None)
    if mps is not None and mps.is_available():
        return torch.device("mps"), "Apple Silicon (MPS)"
    torch.set_num_threads(_CPU_COUNT)
    # OpenCV's resize/decode are memory-bound rather than vector-unit-bound,
    # so unlike the BLAS pool they do still get something out of SMT threads.
    cv2.setNumThreads(_LOGICAL_CPU_COUNT)
    smt = (
        ""
        if _CPU_COUNT == _LOGICAL_CPU_COUNT
        else f", ignoring {_LOGICAL_CPU_COUNT - _CPU_COUNT} SMT siblings"
    )
    return torch.device("cpu"), f"CPU ({_CPU_COUNT} threads{smt})"


# The whole fast inference path (autocast + inference_mode) is switched off
# globally and permanently the first time it is implicated in a failure, and
# the run is retried on the old plain-fp32/no_grad path — see
# run_with_fp32_fallback() in main(). One flag for both because the two
# failure modes are indistinguishable from out here: both surface as a
# RuntimeError from somewhere inside the model.
_FAST_PATH = True


def disable_fast_path():
    global _FAST_PATH
    _FAST_PATH = False


def autocast_dtype(device):
    """The dtype to run the forward pass in, or None for full precision.

    CUDA only: this is where autocast is a large, well-trodden win (Meta's
    own SAM2 notebooks run exactly this way) and where bf16/fp16 map onto
    Tensor Cores. MPS autocast support is still partial and CPU autocast
    without AMX is usually a slowdown, so both stay in fp32.
    """
    if device.type != "cuda":
        return None
    try:
        return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    except Exception:
        return torch.float16


@contextlib.contextmanager
def inference_ctx(device):
    """Wrap a forward pass: inference_mode plus autocast where it pays off.

    inference_mode is a strictly cheaper no_grad — it also drops the
    version-counter and view bookkeeping every tensor otherwise carries —
    and is safe here because nothing this script produces is ever fed back
    into autograd. Whatever a model does with its tensors has to stay
    *inside* this block though: an in-place write to an inference tensor
    after the block has exited is an error, which is why callers wrap
    post-processing too, not just the model call.

    The image encoder and the mask decoder must likewise share one autocast
    region: the encoder's output is cached (SamPredictor.features) and
    reused by predict(), so leaving the region in between would hand fp16
    features to fp32 weights.
    """
    if not _FAST_PATH:
        with torch.no_grad():
            yield
        return

    dtype = autocast_dtype(device)
    with torch.inference_mode():
        if dtype is None:
            yield
        else:
            with torch.autocast(device_type=device.type, dtype=dtype):
                yield


def points_per_batch_for(device):
    """How many grid points to run through the mask decoder in one forward
    pass. On GPU, bigger batches keep thousands of cores fed at once — the
    library default (64) already undersells a modern GPU — but the batch is
    also what dominates peak VRAM in Auto mode, so it is only raised past
    128 on cards with the memory to absorb it. On CPU there's no such thing
    as an idle "core" to feed this way: torch dispatches each matmul to
    every thread via set_num_threads() above regardless of batch size, so a
    smaller batch just means lower peak memory with no speed penalty, which
    matters more on typically memory-constrained CPU boxes.
    """
    if device.type == "cuda":
        try:
            vramGb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
        except Exception:
            vramGb = 0
        return 256 if vramGb >= 12 else 128
    return 128 if device.type == "mps" else 32


# --- Utility Functions ---


def packBoolArray(filepath, arr):
    """Serialise a 2-D mask as an 8-byte big-endian (rows, cols) header
    followed by one continuous LSB-first bit stream.

    The bit stream is deliberately NOT byte-aligned per row: bits run
    straight from the end of one row into the next, and only the very last
    byte of the whole mask is zero-padded. seganyplugin.py's decoder assumes
    exactly that, so `np.packbits(arr, axis=-1)` is the wrong call here — it
    pads every row up to a byte boundary and would silently corrupt every
    mask whose width isn't a multiple of 8. Packing the flattened array is
    what reproduces the original per-bit Python loop byte for byte, while
    doing it in vectorised C (~2.1s -> ~3ms for a 12 MP mask).
    """
    arr = np.asarray(arr)
    num_rows, num_cols = arr.shape
    if arr.dtype.kind not in "biu":
        # SAM2's image predictor returns its masks as float32 0.0/1.0, not
        # bool — and np.packbits only accepts integer or boolean input,
        # where the per-bit Python loop this replaced just tested each value
        # for truthiness. Reproduce that, rather than let a whole model
        # family fail to save anything.
        arr = arr != 0
    packed = np.packbits(arr.reshape(-1), bitorder="little").tobytes()
    with open(filepath, "wb") as f:
        f.write(struct.pack(">II", num_rows, num_cols))
        f.write(packed)


def saveMask(filepath, maskArr, formatBinary):
    if formatBinary:
        packBoolArray(filepath, maskArr)
    else:
        with open(filepath, "w") as f:
            for row in maskArr:
                f.write("".join(str(int(val)) for val in row) + "\n")


def saveMasks(masks, saveFileNoExt, formatBinary, transform=None):
    """Persist each mask, optionally passing it through `transform` first.

    `transform` is how the Auto path gets its low-res masks back up to the
    image's real size (see main()) without a detour through disk: the masks
    are already right here as numpy arrays.
    """
    t0 = time.time()
    count = len(masks)
    dbg(f"saving {count} mask(s) -> {saveFileNoExt}N.seg "
        f"(binary={formatBinary}, upscaling={'yes' if transform else 'no'})")
    span = PCT_SAVE_END - PCT_SEGMENT_END
    for i, mask in enumerate(masks):
        filepath = saveFileNoExt + str(i) + ".seg"
        if transform is not None:
            mask = transform(mask)
        saveMask(filepath, mask, formatBinary)
        progress_pct(
            PCT_SEGMENT_END + span * ((i + 1) / count),
            f"saving mask {i + 1}/{count}",
        )
    # len(), not truthiness: `masks` is often a numpy (N, H, W) array, and
    # `not masks` on one of those raises "truth value is ambiguous".
    if len(masks) == 0:
        dbg("WARNING: the model returned zero masks — the plug-in will "
            "create no layers")
    else:
        shape = np.asarray(masks[0]).shape if transform is None else "upscaled"
        dbg(f"saved {len(masks)} mask(s) in {time.time() - t0:.2f}s "
            f"(first mask {shape})")


# Magic + big-endian (rows, cols, channels) — see exportImageForBridge() in
# seganyplugin.py, which writes this.
RAW_IMAGE_MAGIC = b"GSAMRAW1"
RAW_IMAGE_HEADER_LEN = len(RAW_IMAGE_MAGIC) + 12


def readInputImage(filepath):
    """Load the image the plug-in handed over, as an RGB uint8 array.

    The plug-in prefers dumping raw pixels: GIMP already holds them in
    exactly this layout, so writing them straight out skips a full PNG
    encode on GIMP's own (UI-blocking) thread and the matching decode here.
    It falls back to writing a PNG whenever that path doesn't work, so both
    have to be readable — the magic header is what tells them apart.
    """
    t0 = time.time()
    try:
        with open(filepath, "rb") as f:
            header = f.read(RAW_IMAGE_HEADER_LEN)
            if header[: len(RAW_IMAGE_MAGIC)] == RAW_IMAGE_MAGIC:
                rows, cols, channels = struct.unpack(
                    ">III", header[len(RAW_IMAGE_MAGIC) :]
                )
                pixels = np.fromfile(f, dtype=np.uint8)
                if pixels.size != rows * cols * channels:
                    print(
                        f"Error: raw image {filepath} holds {pixels.size} bytes, "
                        f"expected {rows * cols * channels}"
                    )
                    return None
                dbg(f"input image: raw hand-off, {cols}x{rows}x{channels}, "
                    f"read in {time.time() - t0:.2f}s")
                return pixels.reshape(rows, cols, channels)
    except OSError as e:
        print(f"Error: could not open {filepath}: {e}")
        return None

    image = cv2.imread(filepath)
    if image is None:
        return None
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    dbg(f"input image: PNG fallback, {image.shape[1]}x{image.shape[0]}, "
        f"decoded in {time.time() - t0:.2f}s")
    return image


def prepareWorkImage(cvImage, maxDim):
    """Scale the image down to `maxDim` on its longest side, if needed.

    Returns (workImage, promptScale, maskTransform). promptScale converts
    caller-supplied coordinates (boxes, click points — always in original
    image pixels) into the working image's pixels; maskTransform puts the
    resulting masks back at full size, or is None when nothing was scaled.

    Worth knowing what this does and doesn't buy: SAM's image encoder
    resizes its input to a fixed size internally regardless, so this is not
    what makes the encoder cheaper. What it does make cheaper is everything
    measured in *original image pixels* — the per-mask upsampling of the
    decoder's low-res logits, the mask post-processing, and in Auto mode the
    per-mask area filtering — which on a 36 MP photo is the dominant cost.
    """
    if maxDim <= 0 or max(cvImage.shape[:2]) <= maxDim:
        return cvImage, 1.0, None

    originalShape = cvImage.shape[:2]  # (h, w)
    scale = maxDim / max(originalShape)
    newSize = (
        max(1, int(round(originalShape[1] * scale))),
        max(1, int(round(originalShape[0] * scale))),
    )
    print(
        f"Downscaling {originalShape[1]}x{originalShape[0]} -> "
        f"{newSize[0]}x{newSize[1]} for a faster pass "
        "(masks are upscaled back before saving)",
        flush=True,
    )
    workImage = cv2.resize(cvImage, newSize, interpolation=cv2.INTER_AREA)
    # Derive the prompt scale from the size actually produced, not from the
    # requested ratio: those differ by up to half a pixel after rounding,
    # and a box that drifts is a box that selects the wrong thing.
    promptScale = (newSize[0] / originalShape[1], newSize[1] / originalShape[0])
    return (
        workImage,
        promptScale,
        lambda mask: resizeMaskToOriginal(mask, originalShape),
    )


def scalePrompt(coords, promptScale):
    """Map [x, y, ...] pairs from original pixels into working pixels."""
    if promptScale == 1.0 or coords is None:
        return coords
    sx, sy = promptScale
    return [v * (sx if i % 2 == 0 else sy) for i, v in enumerate(coords)]


def resizeMaskToOriginal(mask, targetShape):
    """Nearest-neighbour resize a mask back to (h, w) = targetShape.

    Returns uint8 (0/1) rather than bool: nothing downstream needs a real
    bool array — np.packbits treats any non-zero as 1 and the ASCII writer
    goes through int() — so the extra full-array cast is pure waste.
    """
    h, w = targetShape
    return cv2.resize(
        np.asarray(mask, dtype=np.uint8), (w, h), interpolation=cv2.INTER_NEAREST
    )


# --- Strategy Pattern Implementation ---
#
# Segmentation "resolution" for Auto mode is dominated by points_per_side:
# the automatic mask generator runs one decoder pass PER GRID POINT (32x32 =
# 1024 decoder calls at the old hardcoded default), so this — not the input
# image's pixel size — is what makes Auto mode take seconds vs. tens of
# minutes on a CPU. Both SAM1 and SAM2 share the same generator API, so the
# same knobs apply to both; previously SAM1 silently ignored the "Resolution"
# dropdown entirely and always ran the heaviest possible grid.
POINTS_PER_SIDE_BY_RES = {"Low": 8, "Medium": 16, "High": 32}
DEFAULT_POINTS_PER_SIDE = POINTS_PER_SIDE_BY_RES["Medium"]


class SegmentationStrategy:
    # Assigned by main() as soon as the model is loaded. Every strategy needs
    # it to know whether its forward passes can run under autocast, and the
    # SAM1/SAM2 predictors don't otherwise carry the device back to us.
    device = torch.device("cpu")

    def get_model_type_from_filename(self, model_filename):
        raise NotImplementedError

    def load_model(self, checkPtFilePath, modelType, device):
        raise NotImplementedError

    def segment_auto(self, sam, cvImage, saveFileNoExt, formatBinary, **kwargs):
        raise NotImplementedError

    def segment_box(self, sam, cvImage, maskType, boxCos, saveFileNoExt,
                    formatBinary, promptScale=1.0, maskTransform=None):
        raise NotImplementedError

    def segment_sel(
        self, sam, cvImage, maskType, selFile, boxCos, saveFileNoExt,
        formatBinary, promptScale=1.0, maskTransform=None
    ):
        raise NotImplementedError

    def run_test(self, sam):
        raise NotImplementedError

    def cleanup(self):
        pass


class SAM1Strategy(SegmentationStrategy):
    MODEL_TYPE_LOOKUP = {
        "sam_vit_h_4b8939": "vit_h",
        "sam_vit_l_0b3195": "vit_l",
        "sam_vit_b_01ec64": "vit_b",
    }

    def __init__(self):
        self._sam_model_registry = None
        self._SamAutomaticMaskGenerator = None
        self._SamPredictor = None

    def get_model_type_from_filename(self, model_filename):
        filename_stem = os.path.splitext(model_filename)[0]
        model_type = self.MODEL_TYPE_LOOKUP.get(filename_stem)
        if model_type:
            print(f"Auto-detected SAM1 model type: {model_type}")
            return model_type
        else:
            print(
                f"Error: Could not auto-detect model type from SAM1 filename: {model_filename}"
            )
            print(
                f"Please use one of the following file names: {list(self.MODEL_TYPE_LOOKUP.keys())}"
            )
            return None

    def load_model(self, checkPtFilePath, modelType, device):
        try:
            from segment_anything import (
                sam_model_registry,
                SamAutomaticMaskGenerator as SamAutomaticMaskGenerator_SAM1,
                SamPredictor,
            )
        except ImportError as e:
            print(
                f"Error: 'segment_anything' isn't installed in this interpreter ({e}). "
                "Run: pip install git+https://github.com/facebookresearch/segment-anything.git"
            )
            return None
        self._sam_model_registry = sam_model_registry
        self._SamAutomaticMaskGenerator = SamAutomaticMaskGenerator_SAM1
        self._SamPredictor = SamPredictor
        try:
            sam = sam_model_registry[modelType](checkpoint=checkPtFilePath)
            sam.to(device=device)
            print(f"SAM1 Model loaded successfully on {device}!")
            return sam
        except Exception as e:
            print(f"Error loading SAM1 model: {e}")
            return None

    def segment_auto(self, sam, cvImage, saveFileNoExt, formatBinary, **kwargs):
        points_per_side = POINTS_PER_SIDE_BY_RES.get(
            kwargs.get("segRes"), DEFAULT_POINTS_PER_SIDE
        )
        mask_generator = self._SamAutomaticMaskGenerator(
            sam,
            points_per_side=points_per_side,
            points_per_batch=points_per_batch_for(kwargs.get("device") or torch.device("cpu")),
            crop_n_layers=kwargs.get("cropNLayers", 0),
            min_mask_region_area=kwargs.get("minMaskArea", 0),
        )
        with heartbeat(f"segmenting (grid {points_per_side}x{points_per_side})"):
            with batch_progress(
                mask_generator,
                kwargs.get("cropNLayers", 0),
                points_per_batch_for(self.device),
            ):
                with inference_ctx(self.device):
                    masks = mask_generator.generate(cvImage)
        masks = [mask["segmentation"] for mask in masks]
        saveMasks(
            masks, saveFileNoExt, formatBinary, transform=kwargs.get("maskTransform")
        )

    def segment_box(self, sam, cvImage, maskType, boxCos, saveFileNoExt,
                    formatBinary, promptScale=1.0, maskTransform=None):
        predictor = self._SamPredictor(sam)
        input_box = np.array(scalePrompt(boxCos, promptScale))
        with inference_ctx(self.device):
            progress_pct(PCT_SEGMENT_START, "encoding the image")
            predictor.set_image(cvImage)
            progress_pct(PCT_ENCODED, "predicting the mask")
            masks, _, _ = predictor.predict(
                point_coords=None,
                point_labels=None,
                box=input_box,
                multimask_output=(maskType == "Multiple"),
            )
        saveMasks(masks, saveFileNoExt, formatBinary, transform=maskTransform)

    def segment_sel(
        self, sam, cvImage, maskType, selFile, boxCos, saveFileNoExt,
        formatBinary, promptScale=1.0, maskTransform=None
    ):
        pts = []
        with open(selFile, "r") as f:
            lines = f.readlines()
            for line in lines:
                cos = line.split(" ")
                pts.append(scalePrompt([int(cos[0]), int(cos[1])], promptScale))
        predictor = self._SamPredictor(sam)
        input_point = np.array(pts)
        input_label = np.array([1] * len(input_point))
        input_box = np.array(scalePrompt(boxCos, promptScale)) if boxCos else None
        with inference_ctx(self.device):
            progress_pct(PCT_SEGMENT_START, "encoding the image")
            predictor.set_image(cvImage)
            progress_pct(PCT_ENCODED, "predicting the mask")
            masks, _, _ = predictor.predict(
                point_coords=input_point,
                point_labels=input_label,
                box=input_box,
                multimask_output=(maskType == "Multiple"),
            )
        saveMasks(masks, saveFileNoExt, formatBinary, transform=maskTransform)

    def run_test(self, sam):
        npArr = np.zeros((50, 50), np.uint8)
        cvImage = cv2.cvtColor(npArr, cv2.COLOR_GRAY2BGR)
        predictor = self._SamPredictor(sam)
        input_box = np.array([10, 10, 20, 20])
        with inference_ctx(self.device):
            predictor.set_image(cvImage)
            predictor.predict(
                point_coords=None,
                point_labels=None,
                box=input_box,
                multimask_output=False,
            )


class SAM2Strategy(SegmentationStrategy):
    MODEL_TYPE_LOOKUP = {
        "sam2_hiera_large": "sam2_hiera_large",
        "sam2_hiera_base_plus": "sam2_hiera_base_plus",
        "sam2_hiera_small": "sam2_hiera_small",
        "sam2_hiera_tiny": "sam2_hiera_tiny",
        "sam2.1_hiera_large": "sam2_hiera_large",
        "sam2.1_hiera_base_plus": "sam2_hiera_base_plus",
        "sam2.1_hiera_small": "sam2_hiera_small",
        "sam2.1_hiera_tiny": "sam2_hiera_tiny",
    }

    def __init__(self):
        self._temp_pth_path = None
        self._build_sam2 = None
        self._SAM2ImagePredictor = None
        self._SAM2AutomaticMaskGenerator = None

    def get_model_type_from_filename(self, model_filename):
        filename_stem = os.path.splitext(model_filename)[0]
        model_type = self.MODEL_TYPE_LOOKUP.get(filename_stem)
        if model_type:
            print(f"Auto-detected SAM2 model type: {model_type}")
            return model_type
        else:
            print(
                f"Error: Could not auto-detect model type from SAM2 filename: {model_filename}"
            )
            print(
                f"Please use one of the following file names (or their .safetensors/.pt equivalents): {list(self.MODEL_TYPE_LOOKUP.keys())}"
            )
            return None

    def _convert_safetensors_to_pth(self, safetensors_path, pth_path):
        try:
            from safetensors.torch import load_file

            state_dict = load_file(safetensors_path)
            checkpoint = {"model": state_dict}
            torch.save(checkpoint, pth_path)
            return True
        except Exception as e:
            print(f"Error converting safetensors to pth: {e}")
            return False

    def load_model(self, checkPtFilePath, modelType, device):
        try:
            from sam2.build_sam import build_sam2
            from sam2.sam2_image_predictor import SAM2ImagePredictor
            from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
        except ImportError as e:
            print(
                f"Error: 'sam2' (segment-anything-2) isn't installed in this interpreter ({e}). "
                "Run: pip install git+https://github.com/facebookresearch/segment-anything-2.git"
            )
            return None
        self._build_sam2 = build_sam2
        self._SAM2ImagePredictor = SAM2ImagePredictor
        self._SAM2AutomaticMaskGenerator = SAM2AutomaticMaskGenerator

        model_configs = {
            "sam2_hiera_tiny": "sam2_hiera_t.yaml",
            "sam2_hiera_small": "sam2_hiera_s.yaml",
            "sam2_hiera_base_plus": "sam2_hiera_b+.yaml",
            "sam2_hiera_large": "sam2_hiera_l.yaml",
        }
        config_file = model_configs.get(modelType, "sam2_hiera_l.yaml")
        actual_checkpoint_path = checkPtFilePath
        if checkPtFilePath.endswith(".safetensors"):
            print("Converting safetensors to pth format...")
            self._temp_pth_path = checkPtFilePath.replace(".safetensors", "_temp.pth")
            if self._convert_safetensors_to_pth(checkPtFilePath, self._temp_pth_path):
                actual_checkpoint_path = self._temp_pth_path
                print(f"Converted to: {self._temp_pth_path}")
            else:
                print("Failed to convert safetensors file")
                return None
        try:
            sam = build_sam2(
                config_file, actual_checkpoint_path, device=str(device)
            )
            print(f"SAM2 Model loaded successfully on {device}!")
            return sam
        except Exception as e:
            print(f"Error loading SAM2 model: {e}")
            self.cleanup()
            return None

    def segment_auto(self, sam, cvImage, saveFileNoExt, formatBinary, **kwargs):
        points_per_side = POINTS_PER_SIDE_BY_RES.get(
            kwargs.get("segRes"), DEFAULT_POINTS_PER_SIDE
        )
        mask_generator = self._SAM2AutomaticMaskGenerator(
            model=sam,
            points_per_side=points_per_side,
            points_per_batch=points_per_batch_for(kwargs.get("device") or torch.device("cpu")),
            crop_n_layers=kwargs.get("cropNLayers", 0),
            min_mask_region_area=kwargs.get("minMaskArea", 0),
        )
        with heartbeat(f"segmenting (grid {points_per_side}x{points_per_side})"):
            with batch_progress(
                mask_generator,
                kwargs.get("cropNLayers", 0),
                points_per_batch_for(self.device),
            ):
                with inference_ctx(self.device):
                    masks = mask_generator.generate(cvImage)
        masks = [mask["segmentation"] for mask in masks]
        saveMasks(
            masks, saveFileNoExt, formatBinary, transform=kwargs.get("maskTransform")
        )

    def segment_box(self, sam, cvImage, maskType, boxCos, saveFileNoExt,
                    formatBinary, promptScale=1.0, maskTransform=None):
        predictor = self._SAM2ImagePredictor(sam)
        input_box = np.array(scalePrompt(boxCos, promptScale))
        with inference_ctx(self.device):
            progress_pct(PCT_SEGMENT_START, "encoding the image")
            predictor.set_image(cvImage)
            progress_pct(PCT_ENCODED, "predicting the mask")
            masks, _, _ = predictor.predict(
                point_coords=None,
                point_labels=None,
                box=input_box,
                multimask_output=(maskType == "Multiple"),
            )
        saveMasks(masks, saveFileNoExt, formatBinary, transform=maskTransform)

    def segment_sel(
        self, sam, cvImage, maskType, selFile, boxCos, saveFileNoExt,
        formatBinary, promptScale=1.0, maskTransform=None
    ):
        pts = []
        with open(selFile, "r") as f:
            lines = f.readlines()
            for line in lines:
                cos = line.split(" ")
                pts.append(scalePrompt([int(cos[0]), int(cos[1])], promptScale))
        predictor = self._SAM2ImagePredictor(sam)
        input_point = np.array(pts)
        input_label = np.array([1] * len(input_point))
        input_box = np.array(scalePrompt(boxCos, promptScale)) if boxCos else None
        with inference_ctx(self.device):
            progress_pct(PCT_SEGMENT_START, "encoding the image")
            predictor.set_image(cvImage)
            progress_pct(PCT_ENCODED, "predicting the mask")
            masks, _, _ = predictor.predict(
                point_coords=input_point,
                point_labels=input_label,
                box=input_box,
                multimask_output=(maskType == "Multiple"),
            )
        saveMasks(masks, saveFileNoExt, formatBinary, transform=maskTransform)

    def run_test(self, sam):
        npArr = np.zeros((50, 50), np.uint8)
        cvImage = cv2.cvtColor(npArr, cv2.COLOR_GRAY2BGR)
        predictor = self._SAM2ImagePredictor(sam)
        input_box = np.array([10, 10, 20, 20])
        with inference_ctx(self.device):
            predictor.set_image(cvImage)
            predictor.predict(
                point_coords=None,
                point_labels=None,
                box=input_box,
                multimask_output=False,
            )

    def cleanup(self):
        if self._temp_pth_path and os.path.exists(self._temp_pth_path):
            os.remove(self._temp_pth_path)
            print(f"Removed temporary file: {self._temp_pth_path}")


class SAM3Strategy(SegmentationStrategy):
    """SAM3 ("Segment Anything with Concepts") is a fundamentally different
    tool from SAM1/2: give it a short noun phrase (e.g. "car") and it finds
    every matching instance in the image, in addition to the usual box
    prompts. There is no "Auto" grid-search mode and no per-model size
    variants — Meta ships a single checkpoint.

    Deliberately built on `transformers`' own Sam3Model/Sam3Processor
    rather than Meta's standalone facebookresearch/sam3 package. That
    package's own README lists "Python 3.12+" and "a CUDA-compatible GPU"
    under Prerequisites — i.e. Meta doesn't test or support anything else.
    transformers' SAM3 is a regular HF model: it loads and runs through the
    exact same `.to(device)` path as every other model in this file,
    including plain CPU, and only needs `transformers` itself (no separate
    git clone / editable install, no Python version floor beyond whatever
    transformers itself requires). The trade-off is purely speed, not
    capability: SAM3 has an 848M-parameter architecture with a fairly heavy
    text encoder, so CPU inference is genuinely slow (several seconds per
    image) compared to a fraction of a second on a GPU — but it does work.

    The checkpoint is a gated Hugging Face snapshot (several files, not a
    lone .pth) — `checkPtFilePath` here is the local snapshot directory the
    installer downloaded (or, for convenience, could be an "org/model" hub
    id if a user already has one authenticated locally).
    """

    def get_model_type_from_filename(self, model_filename):
        return "sam3"

    def load_model(self, checkPtFilePath, modelType, device):
        try:
            from transformers import Sam3Model, Sam3Processor
        except ImportError as e:
            print(
                "Error: 'transformers' with SAM3 support isn't installed in "
                f"this interpreter ({e}). Run: pip install -U transformers"
            )
            return None

        # Name the actual contents before from_pretrained() does: the
        # snapshot is only loadable here if it carries transformers-format
        # weights. A directory holding just Meta's own .pt (what
        # facebook/sam3.1 ships) looks perfectly complete — config.json,
        # tokenizer, several GB on disk — and still cannot be loaded.
        if os.path.isdir(checkPtFilePath):
            present = sorted(os.listdir(checkPtFilePath))
            dbg(f"SAM3 snapshot at {checkPtFilePath}: {len(present)} file(s)")
            weights = [f for f in present
                       if f.endswith((".safetensors", ".bin"))]
            if not weights:
                print(
                    "Error: this SAM3 snapshot has no transformers-format "
                    f"weights (found: {', '.join(present) or 'nothing'}). "
                    "transformers needs model.safetensors — that is what "
                    "facebook/sam3 ships. facebook/sam3.1 ships only "
                    "sam3.1_multiplex.pt, which loads solely through Meta's "
                    "standalone sam3 package, not through this bridge. "
                    "Re-download with: gimpsam sam3 download --token <hf_token>"
                )
                return None
            dbg(f"SAM3 weights found: {', '.join(weights)}")

        try:
            model = Sam3Model.from_pretrained(checkPtFilePath)
            model.to(device=device)
            model.eval()
            processor = Sam3Processor.from_pretrained(checkPtFilePath)
            print(f"SAM3 model loaded successfully on {device}!")
            return {"model": model, "processor": processor, "device": device}
        except Exception as e:
            print(f"Error loading SAM3 model: {e}")
            print(
                "This is most often either (1) no Hugging Face access token "
                "configured — run the installer's SAM3 download step after "
                "being approved for facebook/sam3 — or (2) the checkpoint "
                "hasn't been fully downloaded yet."
            )
            return None

    def _run(self, sam, pilImage, **prompt_kwargs):
        model, processor, device = sam["model"], sam["processor"], sam["device"]
        inputs = processor(images=pilImage, return_tensors="pt", **prompt_kwargs)
        inputs = inputs.to(device)
        # Post-processing stays inside the block on purpose: it works on the
        # model's own output tensors and is free to do so in place, which
        # inference_mode only permits while the block is still open.
        with inference_ctx(device):
            progress_pct(PCT_SEGMENT_START, "running the model")
            outputs = model(**inputs)
            progress_pct(PCT_ENCODED, "post-processing the masks")
            results = processor.post_process_instance_segmentation(
                outputs,
                threshold=0.5,
                mask_threshold=0.5,
                target_sizes=inputs.get("original_sizes").tolist(),
            )[0]
            masks = results["masks"]
            scores = results.get("scores")
            if scores is not None:
                # Detach from the model's tensors while still inside the
                # block, so nothing downstream holds an inference tensor.
                try:
                    scores = [float(s) for s in scores]
                except (TypeError, ValueError):
                    scores = None
            try:
                masks = masks.detach().cpu().numpy()
            except AttributeError:
                masks = np.asarray(masks)
            masks = masks.astype(bool)
        if masks.ndim == 2:
            masks = masks[None, ...]
        return masks, scores

    def _pick_single(self, masks, scores):
        if len(masks) <= 1:
            return masks
        best = 0
        if scores is not None:
            try:
                best = int(np.argmax(np.asarray([float(s) for s in scores])))
            except Exception:
                best = 0
        return masks[best : best + 1]

    def segment_text(self, sam, pilImage, maskType, textPrompt, saveFileNoExt,
                     formatBinary, promptScale=1.0, maskTransform=None):
        masks, scores = self._run(sam, pilImage, text=textPrompt)
        if maskType == "Single":
            masks = self._pick_single(masks, scores)
        print(f"SAM3 found {len(masks)} instance(s) of '{textPrompt}'")
        saveMasks(list(masks), saveFileNoExt, formatBinary, transform=maskTransform)

    def segment_box(self, sam, cvImage, maskType, boxCos, saveFileNoExt,
                    formatBinary, promptScale=1.0, maskTransform=None):
        pilImage = Image.fromarray(cvImage)
        masks, scores = self._run(
            sam, pilImage, input_boxes=[[boxCos]], input_boxes_labels=[[1]]
        )
        if maskType == "Single":
            masks = self._pick_single(masks, scores)
        print(f"SAM3 found {len(masks)} mask(s) for the given box")
        saveMasks(list(masks), saveFileNoExt, formatBinary, transform=maskTransform)

    def segment_sel(
        self, sam, cvImage, maskType, selFile, boxCos, saveFileNoExt,
        formatBinary, promptScale=1.0, maskTransform=None
    ):
        raise NotImplementedError(
            "SAM3 'Selection' (freehand/point) mode isn't wired up in this "
            "bridge — use Segmentation Type 'Box' or 'Text' instead."
        )

    def segment_auto(self, sam, cvImage, saveFileNoExt, formatBinary, **kwargs):
        raise NotImplementedError(
            "SAM3 has no grid-search 'Auto' mode — use Segmentation Type "
            "'Text' with a short phrase describing what to segment (it will "
            "find every matching instance, which covers most of what "
            "'Auto' is used for with SAM1/2)."
        )

    def run_test(self, sam):
        img = Image.new("RGB", (64, 64), (255, 255, 255))
        self._run(sam, img, text="square")


def run_with_fp32_fallback(call, saveFileNoExt):
    """Run a segmentation, retrying once on the old fp32/no_grad path.

    autocast and inference_mode are the two places where this script now
    asks more of torch than it used to, and the ways they can go wrong
    (an op with no half-precision kernel on an exotic build, a model that
    writes to its own output tensors) are build- and model-specific enough
    that they can't be enumerated here. A retry costs the user one repeated
    segmentation; not retrying costs them the feature outright.

    Any .seg files from the failed attempt have to go first: the plug-in
    reads masks by index until one is missing, so leftovers from a longer
    failed run would be picked up as if they belonged to the retry.
    """
    try:
        return call()
    except RuntimeError as exc:
        if not _FAST_PATH:
            raise
        disable_fast_path()
        print(
            f"[warn] fast inference path failed ({exc}); "
            "retrying in full precision",
            flush=True,
        )

    for stale in glob.glob(saveFileNoExt + "*.seg"):
        try:
            os.remove(stale)
        except OSError:
            pass
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return call()


def main():
    """Returns a process exit status: 0 on success, non-zero on failure.

    Every failure below used to `return` bare, i.e. exit 0. The plug-in
    decides whether to show an error dialog by looking at exactly that
    status, so a missing SAM package or an unreadable checkpoint came back
    as "success" and the user got no dialog, no error, and no layers —
    silence that looks identical to "the model found nothing".
    """
    if len(sys.argv) < 3:
        print(
            "Usage: python seganybridge.py <model_type|auto> <checkpoint_path> [options]"
        )
        return 2

    t_start = time.time()
    modelType = sys.argv[1]
    checkPtFilePath = sys.argv[2]
    model_filename = os.path.basename(checkPtFilePath.rstrip(os.sep))

    dbg(f"bridge starting: python {sys.version.split()[0]}, "
        f"{len(sys.argv) - 1} argument(s)")
    dbg("argv: " + " | ".join(sys.argv[1:]))

    if model_filename.lower() == "sam3" or modelType.lower() == "sam3":
        strategy = SAM3Strategy()
    elif model_filename.lower().startswith("sam_"):
        strategy = SAM1Strategy()
    elif model_filename.lower().startswith("sam2"):
        strategy = SAM2Strategy()
    else:
        print(
            f"Error: Could not determine model family from filename: {model_filename}"
        )
        print(
            "Filename must start with 'sam_' for SAM1, 'sam2' for SAM2, or "
            "be a folder named 'sam3' for SAM3."
        )
        return 1

    if modelType.lower() == "auto":
        modelType = strategy.get_model_type_from_filename(model_filename)
        if not modelType:
            return 1

    if not os.path.exists(checkPtFilePath):
        print(f"Error: Checkpoint file not found: {checkPtFilePath}")
        return 1

    dbg(f"model family: {type(strategy).__name__}, model type: {modelType}, "
        f"checkpoint: {checkPtFilePath}")

    device, device_desc = pick_device()
    print(f"Using device: {device_desc}", flush=True)
    amp = autocast_dtype(device)
    dbg(f"inference path: inference_mode + autocast "
        f"{str(amp).replace('torch.', '') if amp else 'off (full fp32)'}; "
        f"points_per_batch={points_per_batch_for(device)}")

    stage("loading_model")
    progress_pct(2.0, "loading the model")
    with heartbeat("loading model"):
        sam = strategy.load_model(checkPtFilePath, modelType, device)
    if sam is None:
        dbg("model failed to load — see the error above; nothing else ran")
        return 1
    strategy.device = device
    dbg("model loaded")
    progress_pct(PCT_MODEL_LOADED, "reading the image")

    if len(sys.argv) == 3:
        dbg("self-test mode (no image given): running a tiny inference")
        strategy.run_test(sam)
        print("Success!!")
        strategy.cleanup()
        return 0

    ipFile = sys.argv[3]
    segType = sys.argv[4]
    maskType = sys.argv[5]
    saveFileNoExt = sys.argv[6]
    formatBinary = sys.argv[7] == "True" if len(sys.argv) > 7 else True

    stage("reading_image")
    cvImage = readInputImage(ipFile)
    if cvImage is None:
        print(f"Error: could not read the input image: {ipFile}")
        return 1

    dbg(f"mode: segType={segType}, maskType={maskType}, "
        f"masks -> {saveFileNoExt}N.seg")
    progress_pct(PCT_IMAGE_READY, "preparing to segment")

    # The working resolution applies to every prompt mode now, not just
    # Auto. It arrives in the same argv slot for Auto (11) as before; the
    # other modes take it as a trailing argument, absent in older callers.
    def requestedMaxDim(index):
        try:
            return int(sys.argv[index]) if len(sys.argv) > index else 0
        except ValueError:
            return 0

    def dispatch():
        if segType == "Auto":
            auto_kwargs = {"device": device}
            if len(sys.argv) > 8:
                auto_kwargs["segRes"] = sys.argv[8]
            if len(sys.argv) > 9:
                auto_kwargs["cropNLayers"] = int(sys.argv[9])
            if len(sys.argv) > 10:
                auto_kwargs["minMaskArea"] = int(sys.argv[10])

            # Segment at the reduced resolution, then upscale each mask back
            # to the original size on its way to disk. This used to
            # round-trip every low-res mask through a temporary ASCII .seg
            # file purely to get the masks out of segment_auto — they are
            # already numpy arrays in there, so the resize is handed down as
            # a transform instead.
            workImage, _, transform = prepareWorkImage(
                cvImage, requestedMaxDim(11)
            )
            if transform is not None:
                auto_kwargs["maskTransform"] = transform

            stage("segmenting")
            dbg(f"Auto pass on {workImage.shape[1]}x{workImage.shape[0]}, "
                f"res={auto_kwargs.get('segRes')}, "
                f"cropNLayers={auto_kwargs.get('cropNLayers', 0)}, "
                f"minMaskArea={auto_kwargs.get('minMaskArea', 0)}")
            strategy.segment_auto(
                sam, workImage, saveFileNoExt, formatBinary, **auto_kwargs
            )
        elif segType in {"Selection", "Box-Selection"}:
            selFile = sys.argv[8]
            # Slot 9 is the optional bounding box. Selection runs that have
            # no box still have to fill the slot (with an empty string) so
            # that the working resolution keeps its fixed place at slot 10.
            boxCos = (
                [float(val.strip()) for val in sys.argv[9].split(",")]
                if len(sys.argv) > 9 and sys.argv[9].strip()
                else None
            )
            workImage, promptScale, transform = prepareWorkImage(
                cvImage, requestedMaxDim(10)
            )
            stage("segmenting")
            dbg(f"Selection pass: points file={selFile}, box={boxCos}, "
                f"prompt scale={promptScale}")
            with heartbeat("segmenting selection"):
                strategy.segment_sel(
                    sam, workImage, maskType, selFile, boxCos, saveFileNoExt,
                    formatBinary, promptScale=promptScale, maskTransform=transform
                )
        elif segType == "Box":
            boxCos = [float(val.strip()) for val in sys.argv[9].split(",")]
            workImage, promptScale, transform = prepareWorkImage(
                cvImage, requestedMaxDim(10)
            )
            stage("segmenting")
            dbg(f"Box pass: box={boxCos}, prompt scale={promptScale}")
            with heartbeat("segmenting box"):
                strategy.segment_box(
                    sam, workImage, maskType, boxCos, saveFileNoExt,
                    formatBinary, promptScale=promptScale, maskTransform=transform
                )
        elif segType == "Text":
            textPrompt = sys.argv[8] if len(sys.argv) > 8 else ""
            if not textPrompt.strip():
                print('Error: "Text" segmentation type needs a non-empty prompt.')
                return 1
            # SAM3's processor resizes to its own fixed input size, so this
            # doesn't make the model itself cheaper — what it saves is the
            # per-instance mask upsampling back to full resolution, which on
            # a large photo is the difference worth having.
            workImage, _, transform = prepareWorkImage(
                cvImage, requestedMaxDim(9)
            )
            stage("segmenting")
            dbg(f'Text pass: prompt="{textPrompt}"')
            with heartbeat(f'segmenting concept "{textPrompt}"'):
                strategy.segment_text(
                    sam,
                    Image.fromarray(workImage),
                    maskType,
                    textPrompt,
                    saveFileNoExt,
                    formatBinary,
                    maskTransform=transform,
                )
        else:
            print(f"Unknown segmentation type: {segType}")
            return 1

    try:
        # `or 0`: the mode handlers above return a status only when they
        # fail, so None means "ran to completion".
        return run_with_fp32_fallback(dispatch, saveFileNoExt) or 0
    finally:
        print(f"Done! (total {time.time() - t_start:.1f}s)", flush=True)
        strategy.cleanup()


if __name__ == "__main__":
    sys.exit(main())