diff --git a/conf/base.config b/conf/base.config
index e1668aaf..45b62604 100644
--- a/conf/base.config
+++ b/conf/base.config
@@ -16,6 +16,15 @@ process {
 
     // resourceLimits = [ cpus: 192, memory: 750.GB, time: 72.h ]
 
+    // Retry signal-induced exits and "killed without exit code" cases:
+    //   130..145    = signal exits (SIGINT=130, SIGKILL=137, SIGTERM=143, etc.)
+    //   104         = ECONNRESET (transient network failures during stage-in/out)
+    //   2147483647  = Integer.MAX_VALUE, Nextflow's sentinel for tasks that died
+    //                 before writing .exitcode (Nextflow surfaces this as
+    //                 "terminated for an unknown reason -- Likely it has been
+    //                 terminated by the external system"). Common on AWS Batch
+    //                 spot capacity, kubernetes preemption, and grid-scheduler
+    //                 cancellations. See nextflow docs/aws.md for the AWS case.
     errorStrategy = { task.exitStatus in ((130..145) + 104 + 2147483647) ? 'retry' : 'finish' }
     maxRetries    = 3
     maxErrors     = '-1'
diff --git a/conf/modules.config b/conf/modules.config
index 22726919..bb5a8786 100644
--- a/conf/modules.config
+++ b/conf/modules.config
@@ -157,6 +157,7 @@ process {
         ext.filter_method  = params.patch_filter_method ?: null
         ext.iqr_multiplier = params.patch_filter_iqr_multiplier
         ext.z_threshold    = params.patch_filter_z_threshold
+        ext.args           = { "--min-transcripts-per-cell ${params.baysor_tiling_min_transcripts_per_cell}" }
         publishDir = [
             path: { "${params.outdir}/${meta.id}/xenium_patch" },
             mode: params.publish_dir_mode,
@@ -307,7 +308,7 @@ process {
             params.stardist_prob_thresh != null ? "--prob_thresh ${params.stardist_prob_thresh}" : "",
             params.stardist_nms_thresh != null ? "--nms_thresh ${params.stardist_nms_thresh}" : "",
             params.stardist_n_tiles != null ? "--n_tiles ${params.stardist_n_tiles}" : "",
-        ].findAll().join(' ')}
+        ].join(' ').trim()}
     }
 
     withName: 'STARDIST_NUCLEI' {
diff --git a/modules/local/segger/create_dataset/main.nf b/modules/local/segger/create_dataset/main.nf
index 06848c41..81320eff 100644
--- a/modules/local/segger/create_dataset/main.nf
+++ b/modules/local/segger/create_dataset/main.nf
@@ -22,7 +22,6 @@ process SEGGER_CREATE_DATASET {
     }
 
     def args = task.ext.args ?: ''
-    def script_path = "/workspace/segger_dev/src/segger/cli/create_dataset_fast.py"
     prefix = task.ext.prefix ?: "${meta.id}"
 
     // check for platform values
@@ -31,193 +30,17 @@ process SEGGER_CREATE_DATASET {
     }
 
     """
-    # Set numba cache directory to avoid caching issues in container
     export NUMBA_CACHE_DIR=\$PWD/.numba_cache
     mkdir -p \$NUMBA_CACHE_DIR
 
-    # Create local bundle directory with symlinks to all original files
-    # This is necessary because input files from S3/Fusion are read-only
-    # Use absolute paths to avoid broken relative symlinks
-    mkdir -p bundle_local
-    for item in ${base_dir}/*; do
-        # Resolve to absolute path (follow any symlinks)
-        abs_path=\$(readlink -f "\$item" 2>/dev/null || realpath "\$item" 2>/dev/null || echo "\$item")
-        basename=\$(basename "\$item")
-        ln -sf "\$abs_path" "bundle_local/\$basename"
-    done
-
-    # Segger expects nucleus_boundaries.parquet but Xenium bundles have cell_boundaries.parquet
-    # Create the symlink if nucleus_boundaries doesn't exist but cell_boundaries does
-    if [ ! -e "bundle_local/nucleus_boundaries.parquet" ] && [ -e "bundle_local/cell_boundaries.parquet" ]; then
-        echo "Creating nucleus_boundaries.parquet symlink from cell_boundaries.parquet"
-        cell_bounds_path=\$(readlink -f "bundle_local/cell_boundaries.parquet" 2>/dev/null || realpath "bundle_local/cell_boundaries.parquet" 2>/dev/null)
-        ln -sf "\$cell_bounds_path" bundle_local/nucleus_boundaries.parquet
-    fi
-
-    # List bundle contents for debugging
-    echo "Bundle contents:"
-    ls -la bundle_local/
-
-    # Fix: Add parquet column statistics for segger
-    echo "Adding statistics to parquet files..."
-    python3 - << 'PYEOF'
-import pyarrow.parquet as pq
-import os
-
-def add_stats(inp, out):
-    if not os.path.exists(inp):
-        print(f"  Skip {inp}")
-        return
-    t = pq.read_table(inp)
-    pq.write_table(t, out, write_statistics=True, compression='snappy')
-    print(f"  Done {os.path.basename(inp)} ({len(t)} rows)")
-
-os.makedirs('bundle_stats', exist_ok=True)
-for f in ['transcripts.parquet', 'nucleus_boundaries.parquet']:
-    add_stats(f'bundle_local/{f}', f'bundle_stats/{f}')
-
-for item in os.listdir('bundle_local'):
-    s, d = f'bundle_local/{item}', f'bundle_stats/{item}'
-    if not os.path.exists(d):
-        os.symlink(os.path.realpath(s), d)
-print("Done")
-
-# Debug: Check overlaps_nucleus column data
-print("")
-print("=== Debugging overlaps_nucleus data ===")
-import pyarrow.compute as pc
-
-tx = pq.read_table('bundle_stats/transcripts.parquet')
-bd = pq.read_table('bundle_stats/nucleus_boundaries.parquet')
-
-if 'overlaps_nucleus' in tx.column_names:
-    col = tx.column('overlaps_nucleus')
-    print(f"overlaps_nucleus dtype: {col.type}")
-    unique_vals = pc.unique(col)
-    print(f"overlaps_nucleus unique values: {unique_vals.to_pylist()[:10]}")
-    val_counts = pc.value_counts(col)
-    print(f"overlaps_nucleus value_counts: {val_counts.to_pylist()}")
-else:
-    print("WARNING: overlaps_nucleus column NOT FOUND in transcripts.parquet")
-
-# Check cell_id overlap between transcripts and boundaries
-if 'cell_id' in tx.column_names and 'cell_id' in bd.column_names:
-    tx_cells = set(pc.unique(tx.column('cell_id')).to_pylist())
-    bd_cells = set(pc.unique(bd.column('cell_id')).to_pylist())
-    overlap = tx_cells & bd_cells
-    print("")
-    print(f"Transcripts unique cell_ids: {len(tx_cells)}")
-    print(f"Boundaries unique cell_ids: {len(bd_cells)}")
-    print(f"Overlapping cell_ids: {len(overlap)}")
-
-print("=== End Debug ===")
-PYEOF
-    ls -la bundle_stats/
-
-    python3 ${script_path} \\
-        --base_dir bundle_stats \\
-        --data_dir ${prefix} \\
-        --sample_type ${params.format} \\
-        --tile_width ${params.tile_width} \\
-        --tile_height ${params.tile_height} \\
-        --n_workers ${task.cpus} \\
+    run_create_dataset.py \\
+        --bundle-dir ${base_dir} \\
+        --output-dir ${prefix} \\
+        --sample-type ${params.format} \\
+        --tile-width ${params.tile_width} \\
+        --tile-height ${params.tile_height} \\
+        --n-workers ${task.cpus} \\
         ${args}
-
-    # Verify tiles were created and show distribution
-    echo "Dataset split (before fix):"
-    echo "  train_tiles: \$(ls ${prefix}/train_tiles/processed/ 2>/dev/null | wc -l) files"
-    echo "  val_tiles: \$(ls ${prefix}/val_tiles/processed/ 2>/dev/null | wc -l) files"
-    echo "  test_tiles: \$(ls ${prefix}/test_tiles/processed/ 2>/dev/null | wc -l) files"
-
-    # Workaround: segger commit 0787167 has a bug where all tiles go to test_tiles
-    # regardless of test_prob/val_prob settings. Move ONLY trainable tiles (those with
-    # edge_label_index) from test_tiles to train_tiles.
-    # Tiles without tx-belongs-bd edges don't have edge_label_index and cannot be used for training.
-    train_count=\$(ls ${prefix}/train_tiles/processed/ 2>/dev/null | wc -l)
-    test_count=\$(ls ${prefix}/test_tiles/processed/ 2>/dev/null | wc -l)
-
-    if [ "\$train_count" -eq 0 ] && [ "\$test_count" -gt 0 ]; then
-        echo "Applying workaround: filtering trainable tiles from test_tiles (segger split bug)"
-        export SEGGER_PREFIX="${prefix}"
-        python3 - << 'PYEOF'
-import torch
-import os
-import shutil
-
-prefix = os.environ['SEGGER_PREFIX']
-test_dir = f"{prefix}/test_tiles/processed"
-train_dir = f"{prefix}/train_tiles/processed"
-
-moved = 0
-skipped = 0
-
-for f in os.listdir(test_dir):
-    if not f.endswith('.pt'):
-        continue
-    fpath = os.path.join(test_dir, f)
-    try:
-        tile = torch.load(fpath, weights_only=False)
-        edge_store = tile['tx', 'belongs', 'bd']
-        # Check if edge_label_index exists and has data
-        if hasattr(edge_store, 'edge_label_index') and edge_store.edge_label_index.numel() > 0:
-            shutil.move(fpath, os.path.join(train_dir, f))
-            moved += 1
-        else:
-            skipped += 1
-    except Exception as e:
-        print(f"Warning: Could not process {f}: {e}")
-        skipped += 1
-
-print(f"Moved {moved} trainable tiles to train_tiles")
-print(f"Skipped {skipped} test-only tiles (no edge_label_index)")
-PYEOF
-    fi
-
-    echo "Dataset split (after fix):"
-    echo "  train_tiles: \$(ls ${prefix}/train_tiles/processed/ 2>/dev/null | wc -l) files"
-    echo "  val_tiles: \$(ls ${prefix}/val_tiles/processed/ 2>/dev/null | wc -l) files"
-    echo "  test_tiles: \$(ls ${prefix}/test_tiles/processed/ 2>/dev/null | wc -l) files"
-
-    train_tiles_dir="${prefix}/train_tiles/processed"
-    if [ ! -d "\$train_tiles_dir" ] || [ -z "\$(ls -A \$train_tiles_dir 2>/dev/null)" ]; then
-        echo "ERROR: No trainable tiles were created in \$train_tiles_dir"
-        echo "This usually means no transcripts overlap with nucleus boundaries in the dataset."
-        echo "Check if the Xenium bundle contains valid overlaps_nucleus data in transcripts.parquet."
-        exit 1
-    fi
-    echo "Successfully created \$(ls \$train_tiles_dir | wc -l) trainable tiles"
-
-    # Workaround: Segger's get_polygon_props() produces NaN boundary features (bd.x)
-    # when polygon geometries have zero area or index misalignment during GeoDataFrame
-    # construction. Replace NaN bd.x with zeros so BCEWithLogitsLoss doesn't propagate NaN.
-    export SEGGER_PREFIX="${prefix}"
-    python3 - << 'PYEOF'
-import torch
-import os
-
-prefix = os.environ['SEGGER_PREFIX']
-fixed = 0
-total = 0
-
-for split in ['train_tiles', 'test_tiles', 'val_tiles']:
-    tile_dir = f"{prefix}/{split}/processed"
-    if not os.path.isdir(tile_dir):
-        continue
-    for f in os.listdir(tile_dir):
-        if not f.endswith('.pt'):
-            continue
-        total += 1
-        fpath = os.path.join(tile_dir, f)
-        tile = torch.load(fpath, weights_only=False)
-        bd_x = tile['bd'].x
-        if bd_x.isnan().any():
-            tile['bd'].x = torch.nan_to_num(bd_x, nan=0.0)
-            torch.save(tile, fpath)
-            fixed += 1
-
-print(f"Fixed NaN bd.x in {fixed}/{total} tiles")
-PYEOF
-
     """
 
     stub:
diff --git a/modules/local/segger/create_dataset/resources/usr/bin/run_create_dataset.py b/modules/local/segger/create_dataset/resources/usr/bin/run_create_dataset.py
new file mode 100755
index 00000000..c73ab006
--- /dev/null
+++ b/modules/local/segger/create_dataset/resources/usr/bin/run_create_dataset.py
@@ -0,0 +1,253 @@
+#!/usr/bin/env python3
+"""
+Run segger create_dataset with spatialxe-specific preprocessing and workarounds.
+
+Wraps segger's create_dataset_fast.py with:
+  - bundle_local symlink prep (handles read-only S3/Fusion mounts)
+  - parquet column statistics (segger needs these)
+  - WORKAROUND: filter trainable tiles from test_tiles when segger commit 0787167 mis-splits
+  - WORKAROUND: replace NaN bd.x with zeros after get_polygon_props produces NaN
+
+Each WORKAROUND should be removable when the upstream segger bug is fixed.
+"""
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+# imports for actual work (used in functions below)
+import pyarrow.parquet as pq
+import pyarrow.compute as pc
+import torch
+
+
+SEGGER_CLI = "/workspace/segger_dev/src/segger/cli/create_dataset_fast.py"
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--bundle-dir", required=True)
+    p.add_argument("--output-dir", required=True)
+    p.add_argument("--sample-type", required=True, choices=["xenium"])
+    p.add_argument("--tile-width", type=int, required=True)
+    p.add_argument("--tile-height", type=int, required=True)
+    p.add_argument("--n-workers", type=int, required=True)
+    # remaining args forwarded to segger CLI
+    args, extra = p.parse_known_args()
+    return args, extra
+
+
+def prepare_bundle(bundle_dir):
+    """Create local bundle dir with absolute symlinks (S3/Fusion read-only-safe)."""
+    Path("bundle_local").mkdir(exist_ok=True)
+    for item in Path(bundle_dir).iterdir():
+        try:
+            abs_path = item.resolve()
+        except Exception:
+            abs_path = item
+        target = Path("bundle_local") / item.name
+        if target.exists() or target.is_symlink():
+            target.unlink()
+        target.symlink_to(abs_path)
+
+    # Segger expects nucleus_boundaries.parquet but Xenium bundles have cell_boundaries.parquet
+    nb = Path("bundle_local/nucleus_boundaries.parquet")
+    cb = Path("bundle_local/cell_boundaries.parquet")
+    if not nb.exists() and cb.exists():
+        print(
+            "Creating nucleus_boundaries.parquet symlink from cell_boundaries.parquet"
+        )
+        nb.symlink_to(cb.resolve())
+
+    print("Bundle contents:")
+    for item in sorted(Path("bundle_local").iterdir()):
+        print(f"  {item.name}")
+
+
+def add_parquet_stats():
+    """Rewrite key parquet files with column statistics (segger requires them)."""
+    Path("bundle_stats").mkdir(exist_ok=True)
+    for fname in ["transcripts.parquet", "nucleus_boundaries.parquet"]:
+        src = Path("bundle_local") / fname
+        dst = Path("bundle_stats") / fname
+        if not src.exists():
+            print(f"  Skip {src}")
+            continue
+        t = pq.read_table(str(src))
+        pq.write_table(t, str(dst), write_statistics=True, compression="snappy")
+        print(f"  Done {fname} ({len(t)} rows)")
+
+    # Symlink everything else from bundle_local into bundle_stats
+    for item in Path("bundle_local").iterdir():
+        dst = Path("bundle_stats") / item.name
+        if not dst.exists():
+            dst.symlink_to(item.resolve())
+
+    # Debug: check overlaps_nucleus column in transcripts
+    print("\n=== Debugging overlaps_nucleus data ===")
+    tx = pq.read_table("bundle_stats/transcripts.parquet")
+    bd = pq.read_table("bundle_stats/nucleus_boundaries.parquet")
+    if "overlaps_nucleus" in tx.column_names:
+        col = tx.column("overlaps_nucleus")
+        print(f"overlaps_nucleus dtype: {col.type}")
+        unique_vals = pc.unique(col)
+        print(f"overlaps_nucleus unique values: {unique_vals.to_pylist()[:10]}")
+        val_counts = pc.value_counts(col)
+        print(f"overlaps_nucleus value_counts: {val_counts.to_pylist()}")
+    else:
+        print("WARNING: overlaps_nucleus column NOT FOUND in transcripts.parquet")
+
+    if "cell_id" in tx.column_names and "cell_id" in bd.column_names:
+        tx_cells = set(pc.unique(tx.column("cell_id")).to_pylist())
+        bd_cells = set(pc.unique(bd.column("cell_id")).to_pylist())
+        overlap = tx_cells & bd_cells
+        print(f"Transcripts unique cell_ids: {len(tx_cells)}")
+        print(f"Boundaries unique cell_ids: {len(bd_cells)}")
+        print(f"Overlapping cell_ids: {len(overlap)}")
+    print("=== End Debug ===\n")
+
+
+def run_segger_cli(args, extra):
+    cmd = [
+        "python3",
+        SEGGER_CLI,
+        "--base_dir",
+        "bundle_stats",
+        "--data_dir",
+        args.output_dir,
+        "--sample_type",
+        args.sample_type,
+        "--tile_width",
+        str(args.tile_width),
+        "--tile_height",
+        str(args.tile_height),
+        "--n_workers",
+        str(args.n_workers),
+        *extra,
+    ]
+    print(f"Running: {' '.join(cmd)}")
+    result = subprocess.run(cmd)
+    if result.returncode != 0:
+        sys.exit(result.returncode)
+
+
+def filter_trainable_tiles_if_needed(prefix):
+    """
+    WORKAROUND: segger commit 0787167 has a bug where all tiles end up in test_tiles
+    regardless of test_prob/val_prob settings. Move ONLY trainable tiles (those with
+    edge_label_index) from test_tiles to train_tiles.
+
+    Remove this function once segger >= 0.1.x is bumped with the upstream fix.
+    """
+    train_dir = Path(prefix) / "train_tiles" / "processed"
+    test_dir = Path(prefix) / "test_tiles" / "processed"
+    val_dir = Path(prefix) / "val_tiles" / "processed"
+
+    train_count = len(list(train_dir.iterdir())) if train_dir.exists() else 0
+    test_count = len(list(test_dir.iterdir())) if test_dir.exists() else 0
+    val_count = len(list(val_dir.iterdir())) if val_dir.exists() else 0
+    print(
+        f"Dataset split (before fix): train={train_count} val={val_count} test={test_count}"
+    )
+
+    if train_count == 0 and test_count > 0:
+        print(
+            "Applying workaround: filtering trainable tiles from test_tiles (segger split bug)"
+        )
+        moved = 0
+        skipped = 0
+        for tile_path in list(test_dir.iterdir()):
+            if not tile_path.name.endswith(".pt"):
+                continue
+            try:
+                tile = torch.load(str(tile_path), weights_only=False)
+                edge_store = tile["tx", "belongs", "bd"]
+                if (
+                    hasattr(edge_store, "edge_label_index")
+                    and edge_store.edge_label_index.numel() > 0
+                ):
+                    shutil.move(str(tile_path), str(train_dir / tile_path.name))
+                    moved += 1
+                else:
+                    skipped += 1
+            except Exception as e:
+                print(f"Warning: Could not process {tile_path.name}: {e}")
+                skipped += 1
+        print(f"Moved {moved} trainable tiles to train_tiles")
+        print(f"Skipped {skipped} test-only tiles (no edge_label_index)")
+
+    train_count = len(list(train_dir.iterdir())) if train_dir.exists() else 0
+    test_count = len(list(test_dir.iterdir())) if test_dir.exists() else 0
+    val_count = len(list(val_dir.iterdir())) if val_dir.exists() else 0
+    print(
+        f"Dataset split (after fix): train={train_count} val={val_count} test={test_count}"
+    )
+
+    if train_count == 0:
+        print(f"ERROR: No trainable tiles were created in {train_dir}", file=sys.stderr)
+        print(
+            "This usually means no transcripts overlap with nucleus boundaries in the dataset.",
+            file=sys.stderr,
+        )
+        print(
+            "Check if the Xenium bundle contains valid overlaps_nucleus data in transcripts.parquet.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    print(f"Successfully created {train_count} trainable tiles")
+
+
+def fix_bd_x_nan(prefix):
+    """
+    WORKAROUND: segger's get_polygon_props() produces NaN boundary features (bd.x)
+    when polygon geometries have zero area or index misalignment during GeoDataFrame
+    construction. Replace NaN bd.x with zeros so BCEWithLogitsLoss doesn't propagate NaN.
+
+    Remove this function once segger >= 0.1.x is bumped with the upstream fix.
+    """
+    fixed = 0
+    total = 0
+    for split in ["train_tiles", "test_tiles", "val_tiles"]:
+        tile_dir = Path(prefix) / split / "processed"
+        if not tile_dir.is_dir():
+            continue
+        for tile_path in tile_dir.iterdir():
+            if not tile_path.name.endswith(".pt"):
+                continue
+            total += 1
+            tile = torch.load(str(tile_path), weights_only=False)
+            bd_x = tile["bd"].x
+            if bd_x.isnan().any():
+                tile["bd"].x = torch.nan_to_num(bd_x, nan=0.0)
+                torch.save(tile, str(tile_path))
+                fixed += 1
+    print(f"Fixed NaN bd.x in {fixed}/{total} tiles")
+
+
+def main():
+    args, extra = parse_args()
+
+    # Ensure numba cache dir is writable (env var should be set by caller, but belt-and-suspenders)
+    os.environ.setdefault("NUMBA_CACHE_DIR", os.path.join(os.getcwd(), ".numba_cache"))
+    os.makedirs(os.environ["NUMBA_CACHE_DIR"], exist_ok=True)
+
+    prepare_bundle(args.bundle_dir)
+    print("Adding statistics to parquet files...")
+    add_parquet_stats()
+
+    # Sanity-check bundle_stats
+    print("bundle_stats contents:")
+    for item in sorted(Path("bundle_stats").iterdir()):
+        print(f"  {item.name}")
+
+    run_segger_cli(args, extra)
+
+    filter_trainable_tiles_if_needed(args.output_dir)
+    fix_bd_x_nan(args.output_dir)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/modules/local/segger/predict/main.nf b/modules/local/segger/predict/main.nf
index a7aec75d..0da7a594 100644
--- a/modules/local/segger/predict/main.nf
+++ b/modules/local/segger/predict/main.nf
@@ -25,52 +25,17 @@ process SEGGER_PREDICT {
     }
 
     def args = task.ext.args ?: ''
-    def script_path = "/workspace/segger_dev/src/segger/cli/predict_fast.py"
     prefix = task.ext.prefix ?: "${meta.id}"
     """
-    # Limit cupy GPU memory to 80% so PyTorch has headroom for graph attention ops
-    export CUPY_GPU_MEMORY_LIMIT="80%"
-    # Belt-and-suspenders: ensure PyTorch uses expandable segments (also set in env {} block)
-    export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True,max_split_size_mb:512"
-
-    # Set numba cache directory to avoid caching issues in container
-    export NUMBA_CACHE_DIR=\$PWD/.numba_cache
-    mkdir -p \$NUMBA_CACHE_DIR
-
-    # GPU detection logging
-    echo "=== GPU Detection (SEGGER_PREDICT) ==="
-    nvidia-smi 2>/dev/null && echo "GPU available: yes" || echo "GPU available: no (nvidia-smi failed)"
-    python3 -c "import torch; print(f'PyTorch CUDA available: {torch.cuda.is_available()}'); print(f'CUDA device count: {torch.cuda.device_count()}')" 2>/dev/null || echo "PyTorch CUDA check failed"
-    echo "======================================"
-
-    # Use all available GPUs (autocast reduces VRAM ~50%, so multi-GPU is safe)
-    GPU_IDS=\$(python3 -c "
-import torch
-n = torch.cuda.device_count()
-print(','.join(str(i) for i in range(n)) if n > 0 else '0')
-" 2>/dev/null || echo "0")
-    echo "Using GPUs: \$GPU_IDS"
-
-    # Patch predict_parquet.py at runtime (avoids Docker rebuild)
-    PRED_PY=\$(python3 -c "import segger.prediction.predict_parquet as m; print(m.__file__)")
-
-    # 1. Add torch.no_grad() to disable gradient graphs during inference (~30-50% VRAM savings)
-    sed -i 's/with cp.cuda.Device(gpu_id):/with cp.cuda.Device(gpu_id), torch.no_grad():/' "\$PRED_PY"
-
-    # 2. Seed random for deterministic GPU assignment (avoids stochastic OOM)
-    sed -i 's/gpu_id = random.choice(gpu_ids)/random.seed(0); gpu_id = random.choice(gpu_ids)/' "\$PRED_PY"
-    echo "Patched \$PRED_PY: torch.no_grad() + round-robin GPU assignment"
-
-    python3 ${script_path} \\
-        --models_dir ${models_dir} \\
-        --segger_data_dir ${segger_dataset} \\
-        --transcripts_file ${transcripts} \\
-        --benchmarks_dir benchmarks_dir \\
-        --batch_size ${params.batch_size_predict} \\
-        --use_cc ${params.cc_analysis} \\
-        --knn_method ${params.segger_knn_method} \\
-        --num_workers ${task.cpus} \\
-        --gpu_ids \$GPU_IDS \\
+    run_predict.py \\
+        --models-dir ${models_dir} \\
+        --segger-data-dir ${segger_dataset} \\
+        --transcripts-file ${transcripts} \\
+        --benchmarks-dir benchmarks_dir \\
+        --batch-size ${params.batch_size_predict} \\
+        --use-cc ${params.cc_analysis} \\
+        --knn-method ${params.segger_knn_method} \\
+        --num-workers ${task.cpus} \\
         ${args}
     """
 
diff --git a/modules/local/segger/predict/resources/usr/bin/run_predict.py b/modules/local/segger/predict/resources/usr/bin/run_predict.py
new file mode 100755
index 00000000..56a77ffc
--- /dev/null
+++ b/modules/local/segger/predict/resources/usr/bin/run_predict.py
@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+"""
+Run segger predict with spatialxe-specific preprocessing.
+
+Wraps segger's predict_fast.py with:
+  - GPU enumeration (replaces inline python3 -c torch check)
+  - WORKAROUND: patch predict_parquet.py at runtime to add torch.no_grad() for ~30-50% VRAM savings
+  - WORKAROUND: seed random.choice for deterministic GPU assignment (avoids stochastic OOM)
+
+Both WORKAROUNDs should be removable once the patches are upstreamed to segger.
+"""
+
+import argparse
+import os
+import subprocess
+import sys
+
+
+SEGGER_CLI = "/workspace/segger_dev/src/segger/cli/predict_fast.py"
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--models-dir", required=True)
+    p.add_argument("--segger-data-dir", required=True)
+    p.add_argument("--transcripts-file", required=True)
+    p.add_argument("--benchmarks-dir", required=True)
+    p.add_argument("--batch-size", type=int, required=True)
+    p.add_argument("--use-cc", required=True)
+    p.add_argument("--knn-method", required=True)
+    p.add_argument("--num-workers", type=int, required=True)
+    args, extra = p.parse_known_args()
+    return args, extra
+
+
+def detect_gpus():
+    """Return comma-separated list of available CUDA device ids (or "0" if none)."""
+    import torch
+
+    print("=== GPU Detection (SEGGER_PREDICT) ===")
+    print(f"PyTorch CUDA available: {torch.cuda.is_available()}")
+    n = torch.cuda.device_count()
+    print(f"CUDA device count: {n}")
+    print("======================================")
+    if n > 0:
+        return ",".join(str(i) for i in range(n))
+    return "0"
+
+
+def patch_predict_parquet():
+    """
+    WORKAROUND: patch segger.prediction.predict_parquet at runtime.
+
+    Avoids rebuilding the segger Docker image. Two patches:
+      1. Add torch.no_grad() to disable gradient graphs during inference (~30-50% VRAM savings).
+      2. Seed random for deterministic GPU assignment (avoids stochastic OOM).
+
+    Remove this function once the patches are upstreamed to segger.
+    """
+    import segger.prediction.predict_parquet as m
+
+    pred_py = m.__file__
+    print(f"Patching {pred_py}: torch.no_grad() + round-robin GPU assignment")
+    # Use sed via subprocess for in-place edit (matches the original behavior exactly)
+    subprocess.run(
+        [
+            "sed",
+            "-i",
+            "s/with cp.cuda.Device(gpu_id):/with cp.cuda.Device(gpu_id), torch.no_grad():/",
+            pred_py,
+        ],
+        check=True,
+    )
+    subprocess.run(
+        [
+            "sed",
+            "-i",
+            "s/gpu_id = random.choice(gpu_ids)/random.seed(0); gpu_id = random.choice(gpu_ids)/",
+            pred_py,
+        ],
+        check=True,
+    )
+
+
+def run_segger_cli(args, extra, gpu_ids):
+    cmd = [
+        "python3",
+        SEGGER_CLI,
+        "--models_dir",
+        args.models_dir,
+        "--segger_data_dir",
+        args.segger_data_dir,
+        "--transcripts_file",
+        args.transcripts_file,
+        "--benchmarks_dir",
+        args.benchmarks_dir,
+        "--batch_size",
+        str(args.batch_size),
+        "--use_cc",
+        str(args.use_cc),
+        "--knn_method",
+        args.knn_method,
+        "--num_workers",
+        str(args.num_workers),
+        "--gpu_ids",
+        gpu_ids,
+        *extra,
+    ]
+    print(f"Running: {' '.join(cmd)}")
+    result = subprocess.run(cmd)
+    if result.returncode != 0:
+        sys.exit(result.returncode)
+
+
+def main():
+    args, extra = parse_args()
+
+    # Limit cupy GPU memory to 80% so PyTorch has headroom for graph attention ops
+    os.environ.setdefault("CUPY_GPU_MEMORY_LIMIT", "80%")
+    # Belt-and-suspenders: ensure PyTorch uses expandable segments
+    os.environ.setdefault(
+        "PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True,max_split_size_mb:512"
+    )
+    # Numba cache directory
+    os.environ.setdefault("NUMBA_CACHE_DIR", os.path.join(os.getcwd(), ".numba_cache"))
+    os.makedirs(os.environ["NUMBA_CACHE_DIR"], exist_ok=True)
+
+    gpu_ids = detect_gpus()
+    print(f"Using GPUs: {gpu_ids}")
+
+    patch_predict_parquet()
+
+    run_segger_cli(args, extra, gpu_ids)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/modules/local/utility/Dockerfile b/modules/local/utility/Dockerfile
deleted file mode 100644
index 79213ebb..00000000
--- a/modules/local/utility/Dockerfile
+++ /dev/null
@@ -1,11 +0,0 @@
-FROM mambaorg/micromamba:1.5.10-noble
-COPY --chown=$MAMBA_USER:$MAMBA_USER conda.yml /tmp/conda.yml
-RUN micromamba install -y -n base -f /tmp/conda.yml \
-    && micromamba install -y -n base conda-forge::procps-ng \
-    && micromamba env export --name base --explicit > environment.lock \
-    && echo ">> CONDA_LOCK_START" \
-    && cat environment.lock \
-    && echo "<< CONDA_LOCK_END" \
-    && micromamba clean -a -y
-USER root
-ENV PATH="$MAMBA_ROOT_PREFIX/bin:$PATH"
diff --git a/modules/local/utility/convert_mask_uint32/main.nf b/modules/local/utility/convert_mask_uint32/main.nf
index fd2a3bff..3f0333a7 100644
--- a/modules/local/utility/convert_mask_uint32/main.nf
+++ b/modules/local/utility/convert_mask_uint32/main.nf
@@ -35,16 +35,9 @@ process CONVERT_MASK_UINT32 {
     script:
     prefix = task.ext.prefix ?: "${meta.id}"
     """
-    python3 - ${mask} ${prefix}_uint32_mask.tif <<'PYEOF'
-import sys, tifffile, numpy as np
-
-mask_path, output_path = sys.argv[1], sys.argv[2]
-mask = tifffile.imread(mask_path)
-print(f'Input dtype: {mask.dtype}, shape: {mask.shape}, labels: {mask.max()}')
-tifffile.imwrite(output_path, mask.astype(np.uint32))
-print(f'Output dtype: uint32')
-PYEOF
-
+    convert_mask_uint32.py \\
+        --input ${mask} \\
+        --output ${prefix}_uint32_mask.tif
     """
 
     stub:
diff --git a/modules/local/utility/convert_mask_uint32/resources/usr/bin/convert_mask_uint32.py b/modules/local/utility/convert_mask_uint32/resources/usr/bin/convert_mask_uint32.py
new file mode 100755
index 00000000..955ad4b7
--- /dev/null
+++ b/modules/local/utility/convert_mask_uint32/resources/usr/bin/convert_mask_uint32.py
@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+"""
+Convert a segmentation mask TIFF to uint32 dtype.
+
+XeniumRanger import-segmentation requires uint32 masks, but upstream
+segmenters (e.g. StarDist) often emit int32 labels. This script reads
+the input mask, casts it to uint32, and writes the result.
+"""
+
+import argparse
+
+import numpy as np
+import tifffile
+
+
+def convert_mask_to_uint32(input_path: str, output_path: str) -> None:
+    """
+    Read a mask TIFF, cast to uint32, and write to output_path.
+
+    Args:
+        input_path: Path to input mask TIFF (any integer dtype).
+        output_path: Path where the uint32 mask will be written.
+    """
+    mask = tifffile.imread(input_path)
+    print(f"Input dtype: {mask.dtype}, shape: {mask.shape}, labels: {mask.max()}")
+    tifffile.imwrite(output_path, mask.astype(np.uint32))
+    print("Output dtype: uint32")
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Convert a segmentation mask TIFF to uint32 dtype."
+    )
+    parser.add_argument(
+        "--input", required=True, help="Path to input mask TIFF"
+    )
+    parser.add_argument(
+        "--output", required=True, help="Path where uint32 mask will be written"
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    convert_mask_to_uint32(input_path=args.input, output_path=args.output)
diff --git a/modules/local/utility/downscale_morphology/main.nf b/modules/local/utility/downscale_morphology/main.nf
index 197b3618..edaf3d67 100644
--- a/modules/local/utility/downscale_morphology/main.nf
+++ b/modules/local/utility/downscale_morphology/main.nf
@@ -41,55 +41,11 @@ process DOWNSCALE_MORPHOLOGY {
     def diam_mean = 30
     prefix = task.ext.prefix ?: "${meta.id}"
     """
-    mkdir -p ${prefix}
-
-    python3 -c "
-import tifffile, numpy as np, json
-from skimage.transform import resize
-
-diameter = ${diameter}
-diam_mean = ${diam_mean}
-scale = min(diameter / diam_mean, 1.0)  # clamp to prevent upscaling
-
-img = tifffile.imread('${image}')
-print(f'Original: {img.shape}, dtype={img.dtype}, ndim={img.ndim}')
-
-# Handle multichannel OME-TIFFs: shape can be (H, W), (C, H, W), or (Z, C, H, W)
-if img.ndim == 2:
-    orig_h, orig_w = img.shape
-    # Floor of 256px: cellpose network requires minimum spatial dimensions
-    new_h = max(int(orig_h * scale), 256)
-    new_w = max(int(orig_w * scale), 256)
-    output_shape = (new_h, new_w)
-elif img.ndim == 3:
-    orig_h, orig_w = img.shape[1], img.shape[2]
-    new_h = max(int(orig_h * scale), 256)
-    new_w = max(int(orig_w * scale), 256)
-    output_shape = (img.shape[0], new_h, new_w)
-else:
-    orig_h, orig_w = img.shape[-2], img.shape[-1]
-    new_h = max(int(orig_h * scale), 256)
-    new_w = max(int(orig_w * scale), 256)
-    output_shape = img.shape[:-2] + (new_h, new_w)
-
-print(f'Downscaling by {scale:.3f}: ({orig_h}, {orig_w}) -> ({new_h}, {new_w})')
-
-img_ds = resize(img, output_shape, order=3, preserve_range=True, anti_aliasing=True)
-img_ds = img_ds.astype(img.dtype)
-
-tifffile.imwrite('${prefix}/downscaled.tif', img_ds, compression='zlib')
-json.dump({
-    'scale': scale,
-    'orig_h': orig_h,
-    'orig_w': orig_w,
-    'new_h': new_h,
-    'new_w': new_w,
-    'diameter': diameter,
-    'diam_mean': diam_mean
-}, open('${prefix}/scale_info.json', 'w'))
-print(f'Done: downscaled.tif written, shape={img_ds.shape}')
-"
-
+    downscale_morphology.py \\
+        --image ${image} \\
+        --diameter ${diameter} \\
+        --diam-mean ${diam_mean} \\
+        --prefix ${prefix}
     """
 
     stub:
diff --git a/modules/local/utility/downscale_morphology/resources/usr/bin/downscale_morphology.py b/modules/local/utility/downscale_morphology/resources/usr/bin/downscale_morphology.py
new file mode 100755
index 00000000..8544ecf3
--- /dev/null
+++ b/modules/local/utility/downscale_morphology/resources/usr/bin/downscale_morphology.py
@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+"""
+Pre-downscale a morphology image for Cellpose.
+
+Reduces image dimensions by a scale factor so that Cellpose's internal
+rescaling (diam_mean / diameter) does not exceed GPU/CPU memory. The
+scale factor defaults to diameter / diam_mean (e.g., 9 / 30 = 0.3).
+After downscaling, Cellpose should run with --diameter equal to
+diam_mean (no further internal rescaling).
+
+Outputs:
+    {prefix}/downscaled.tif - Downscaled image at the same dtype as input.
+    {prefix}/scale_info.json - Scale factor and original/new dimensions.
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import tifffile
+from skimage.transform import resize
+
+# Cellpose network requires a minimum spatial size of 256 px.
+MIN_DIM = 256
+
+
+def downscale_image(
+    image_path: str, diameter: float, diam_mean: float, prefix: str
+) -> None:
+    """
+    Downscale image so Cellpose can run with diameter == diam_mean.
+
+    Args:
+        image_path: Path to morphology TIFF (2D, 3D, or 4D).
+        diameter: Target object diameter (used to compute scale).
+        diam_mean: Cellpose model's mean diameter assumption.
+        prefix: Output directory.
+    """
+    scale = min(diameter / diam_mean, 1.0)  # clamp to prevent upscaling
+
+    img = tifffile.imread(image_path)
+    print(f"Original: {img.shape}, dtype={img.dtype}, ndim={img.ndim}")
+
+    # Handle multichannel OME-TIFFs: shape can be (H, W), (C, H, W), or (Z, C, H, W)
+    if img.ndim == 2:
+        orig_h, orig_w = img.shape
+        new_h = max(int(orig_h * scale), MIN_DIM)
+        new_w = max(int(orig_w * scale), MIN_DIM)
+        output_shape = (new_h, new_w)
+    elif img.ndim == 3:
+        orig_h, orig_w = img.shape[1], img.shape[2]
+        new_h = max(int(orig_h * scale), MIN_DIM)
+        new_w = max(int(orig_w * scale), MIN_DIM)
+        output_shape = (img.shape[0], new_h, new_w)
+    else:
+        orig_h, orig_w = img.shape[-2], img.shape[-1]
+        new_h = max(int(orig_h * scale), MIN_DIM)
+        new_w = max(int(orig_w * scale), MIN_DIM)
+        output_shape = img.shape[:-2] + (new_h, new_w)
+
+    print(f"Downscaling by {scale:.3f}: ({orig_h}, {orig_w}) -> ({new_h}, {new_w})")
+
+    img_ds = resize(img, output_shape, order=3, preserve_range=True, anti_aliasing=True)
+    img_ds = img_ds.astype(img.dtype)
+
+    out_dir = Path(prefix)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    tifffile.imwrite(str(out_dir / "downscaled.tif"), img_ds, compression="zlib")
+
+    info = {
+        "scale": scale,
+        "orig_h": orig_h,
+        "orig_w": orig_w,
+        "new_h": new_h,
+        "new_w": new_w,
+        "diameter": diameter,
+        "diam_mean": diam_mean,
+    }
+    with open(out_dir / "scale_info.json", "w") as f:
+        json.dump(info, f)
+    print(f"Done: downscaled.tif written, shape={img_ds.shape}")
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Pre-downscale a morphology image for Cellpose."
+    )
+    parser.add_argument("--image", required=True, help="Morphology TIFF input")
+    parser.add_argument("--diameter", type=float, required=True, help="Target object diameter")
+    parser.add_argument("--diam-mean", type=float, required=True, help="Cellpose model diam_mean")
+    parser.add_argument("--prefix", required=True, help="Output directory")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    downscale_image(
+        image_path=args.image,
+        diameter=args.diameter,
+        diam_mean=args.diam_mean,
+        prefix=args.prefix,
+    )
diff --git a/modules/local/utility/extract_dapi/main.nf b/modules/local/utility/extract_dapi/main.nf
index adbc6d87..79cce91f 100644
--- a/modules/local/utility/extract_dapi/main.nf
+++ b/modules/local/utility/extract_dapi/main.nf
@@ -36,22 +36,10 @@ process EXTRACT_DAPI {
     prefix = task.ext.prefix ?: "${meta.id}"
     def channel_index = task.ext.channel_index ?: 0
     """
-    python3 - ${image} ${prefix}_dapi.tif ${channel_index} <<'PYEOF'
-import sys, tifffile, numpy as np
-
-image_path, output_path, channel_idx = sys.argv[1], sys.argv[2], int(sys.argv[3])
-img = tifffile.imread(image_path)
-orig_shape = img.shape
-
-if img.ndim == 3:
-    img = img[channel_idx]
-elif img.ndim == 4:
-    img = img[0, channel_idx]
-
-tifffile.imwrite(output_path, img)
-print(f'Input shape: {orig_shape} -> extracted channel {channel_idx}: {img.shape}')
-PYEOF
-
+    extract_dapi.py \\
+        --input ${image} \\
+        --output ${prefix}_dapi.tif \\
+        --channel-index ${channel_index}
     """
 
     stub:
diff --git a/modules/local/utility/extract_dapi/resources/usr/bin/extract_dapi.py b/modules/local/utility/extract_dapi/resources/usr/bin/extract_dapi.py
new file mode 100755
index 00000000..3d60f563
--- /dev/null
+++ b/modules/local/utility/extract_dapi/resources/usr/bin/extract_dapi.py
@@ -0,0 +1,60 @@
+#!/usr/bin/env python3
+"""
+Extract a single channel (e.g., DAPI) from a multi-channel OME-TIFF.
+
+Xenium morphology_focus.ome.tif has multiple channels (DAPI, boundary,
+interior). Single-channel segmenters such as StarDist 2D_versatile_fluo
+expect one channel as input. This script reads the input image, slices
+the requested channel, and writes the result.
+"""
+
+import argparse
+
+import tifffile
+
+
+def extract_channel(input_path: str, output_path: str, channel_index: int) -> None:
+    """
+    Read an OME-TIFF, extract a single channel, and write the result.
+
+    Args:
+        input_path: Path to multi-channel OME-TIFF morphology image.
+        output_path: Path where the single-channel TIFF will be written.
+        channel_index: Index of the channel to extract.
+    """
+    img = tifffile.imread(input_path)
+    orig_shape = img.shape
+
+    if img.ndim == 3:
+        img = img[channel_index]
+    elif img.ndim == 4:
+        img = img[0, channel_index]
+
+    tifffile.imwrite(output_path, img)
+    print(f"Input shape: {orig_shape} -> extracted channel {channel_index}: {img.shape}")
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Extract a single channel from a multi-channel OME-TIFF."
+    )
+    parser.add_argument(
+        "--input", required=True, help="Path to multi-channel OME-TIFF morphology image"
+    )
+    parser.add_argument(
+        "--output", required=True, help="Path where the single-channel TIFF will be written"
+    )
+    parser.add_argument(
+        "--channel-index", type=int, default=0, help="Channel index to extract (default: 0)"
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    extract_channel(
+        input_path=args.input,
+        output_path=args.output,
+        channel_index=args.channel_index,
+    )
diff --git a/modules/local/utility/upscale_mask/main.nf b/modules/local/utility/upscale_mask/main.nf
index ab1409c9..a201abf1 100644
--- a/modules/local/utility/upscale_mask/main.nf
+++ b/modules/local/utility/upscale_mask/main.nf
@@ -35,28 +35,10 @@ process UPSCALE_MASK {
     script:
     prefix = task.ext.prefix ?: "${meta.id}"
     """
-    mkdir -p ${prefix}
-
-    python3 -c "
-import tifffile, numpy as np, json
-from PIL import Image
-
-info = json.load(open('${scale_info}'))
-orig_h, orig_w = info['orig_h'], info['orig_w']
-
-mask = tifffile.imread('${mask}')
-print(f'Mask: {mask.shape}, dtype={mask.dtype}, unique cells: {len(np.unique(mask)) - 1}')
-print(f'Upscaling to ({orig_h}, {orig_w})')
-
-pil_mask = Image.fromarray(mask)
-pil_mask = pil_mask.resize((orig_w, orig_h), Image.NEAREST)
-mask_up = np.array(pil_mask, dtype=mask.dtype)
-
-out_name = '${prefix}/upscaled_${mask.baseName}.tif'
-tifffile.imwrite(out_name, mask_up, compression='zlib')
-print(f'Done: {out_name}, unique cells: {len(np.unique(mask_up)) - 1}')
-"
-
+    upscale_mask.py \\
+        --mask ${mask} \\
+        --scale-info ${scale_info} \\
+        --prefix ${prefix}
     """
 
     stub:
diff --git a/modules/local/utility/upscale_mask/resources/usr/bin/upscale_mask.py b/modules/local/utility/upscale_mask/resources/usr/bin/upscale_mask.py
new file mode 100755
index 00000000..6cc1694e
--- /dev/null
+++ b/modules/local/utility/upscale_mask/resources/usr/bin/upscale_mask.py
@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+"""
+Restore Cellpose masks to original image resolution.
+
+Uses nearest-neighbor interpolation to upscale segmentation masks back
+to the original image dimensions recorded in scale_info.json (produced
+by downscale_morphology.py).
+
+Output: {prefix}/upscaled_{mask_basename}.tif
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import tifffile
+from PIL import Image
+
+
+def upscale_mask(mask_path: str, scale_info_path: str, prefix: str) -> None:
+    """
+    Read a downscaled mask and upscale it to original dimensions.
+
+    Args:
+        mask_path: Path to downscaled segmentation mask TIFF.
+        scale_info_path: Path to scale_info.json from downscale_morphology.
+        prefix: Output directory.
+    """
+    with open(scale_info_path) as f:
+        info = json.load(f)
+    orig_h, orig_w = info["orig_h"], info["orig_w"]
+
+    mask = tifffile.imread(mask_path)
+    print(
+        f"Mask: {mask.shape}, dtype={mask.dtype}, "
+        f"unique cells: {len(np.unique(mask)) - 1}"
+    )
+    print(f"Upscaling to ({orig_h}, {orig_w})")
+
+    pil_mask = Image.fromarray(mask)
+    pil_mask = pil_mask.resize((orig_w, orig_h), Image.NEAREST)
+    mask_up = np.array(pil_mask, dtype=mask.dtype)
+
+    out_dir = Path(prefix)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    base = Path(mask_path).stem
+    out_name = out_dir / f"upscaled_{base}.tif"
+    tifffile.imwrite(str(out_name), mask_up, compression="zlib")
+    print(
+        f"Done: {out_name}, unique cells: {len(np.unique(mask_up)) - 1}"
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Upscale a Cellpose mask back to original resolution."
+    )
+    parser.add_argument("--mask", required=True, help="Downscaled mask TIFF")
+    parser.add_argument("--scale-info", required=True, help="scale_info.json from downscale step")
+    parser.add_argument("--prefix", required=True, help="Output directory")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    upscale_mask(
+        mask_path=args.mask,
+        scale_info_path=args.scale_info,
+        prefix=args.prefix,
+    )
diff --git a/modules/local/xenium_patch/stitch/main.nf b/modules/local/xenium_patch/stitch/main.nf
index 3d523971..d805a0f5 100644
--- a/modules/local/xenium_patch/stitch/main.nf
+++ b/modules/local/xenium_patch/stitch/main.nf
@@ -33,70 +33,18 @@ process XENIUM_PATCH_STITCH {
     task.ext.when == null || task.ext.when
 
     script:
+    def args = task.ext.args ?: ''
     """
     stitch_transcripts.py \\
         --patches ${patches} \\
         --output output \\
-        --min-transcripts-per-cell ${params.baysor_tiling_min_transcripts_per_cell}
-
-    # Post-process: ensure all GeoJSON geometries are Polygon.
-    # make_valid() and solve_conflicts() can produce MultiPolygon,
-    # MultiLineString, or GeometryCollection — XeniumRanger rejects these.
-    # Dropped cells must also be removed from the transcript CSV.
-    python3 -c "
-import csv, json, shapely
-from shapely.geometry import mapping, shape
-
-geojson_path = 'output/xr-cell-polygons.geojson'
-csv_path = 'output/xr-transcript-metadata.csv'
-
-with open(geojson_path) as f:
-    data = json.load(f)
-
-clean = []
-dropped_cells = set()
-for feat in data['features']:
-    geom = shape(feat['geometry'])
-    if not geom.is_valid:
-        geom = shapely.make_valid(geom)
-    poly = None
-    if geom.geom_type == 'Polygon':
-        poly = geom
-    elif geom.geom_type == 'MultiPolygon':
-        poly = max(geom.geoms, key=lambda g: g.area)
-    elif geom.geom_type == 'GeometryCollection':
-        polys = [g for g in geom.geoms if g.geom_type == 'Polygon']
-        if polys:
-            poly = max(polys, key=lambda g: g.area)
-    if poly is not None and not poly.is_empty:
-        feat['geometry'] = mapping(poly)
-        clean.append(feat)
-    else:
-        cell_id = feat.get('id') or feat.get('properties', {}).get('cell_id', '')
-        dropped_cells.add(str(cell_id))
-
-print(f'GeoJSON: {len(clean)} kept, {len(dropped_cells)} dropped: {dropped_cells}')
-data['features'] = clean
-with open(geojson_path, 'w') as f:
-    json.dump(data, f)
-
-if dropped_cells:
-    with open(csv_path) as f:
-        reader = csv.DictReader(f)
-        rows = list(reader)
-    reassigned = 0
-    for row in rows:
-        if row['cell'] in dropped_cells:
-            row['cell'] = ''
-            row['is_noise'] = '1'
-            reassigned += 1
-    with open(csv_path, 'w', newline='') as f:
-        writer = csv.DictWriter(f, fieldnames=reader.fieldnames)
-        writer.writeheader()
-        writer.writerows(rows)
-    print(f'CSV: {reassigned} transcripts reassigned to UNASSIGNED')
-"
+        ${args}
 
+    # Post-process: ensure all GeoJSON geometries are Polygon and
+    # reconcile dropped cells in the transcript CSV.
+    stitch_postprocess.py \\
+        --geojson output/xr-cell-polygons.geojson \\
+        --csv output/xr-transcript-metadata.csv
     """
 
     stub:
diff --git a/modules/local/xenium_patch/stitch/resources/usr/bin/stitch_postprocess.py b/modules/local/xenium_patch/stitch/resources/usr/bin/stitch_postprocess.py
new file mode 100755
index 00000000..7144b1ac
--- /dev/null
+++ b/modules/local/xenium_patch/stitch/resources/usr/bin/stitch_postprocess.py
@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+"""
+Post-process stitched per-patch segmentation outputs.
+
+Ensures every GeoJSON feature is a single Polygon: make_valid() and
+sopa.solve_conflicts() can produce MultiPolygon, MultiLineString, or
+GeometryCollection geometries that XeniumRanger rejects. Cells dropped
+during cleanup are also reassigned to UNASSIGNED in the transcript CSV
+so the two outputs stay consistent.
+"""
+
+import argparse
+import csv
+import json
+
+import shapely
+from shapely.geometry import mapping, shape
+
+
+def clean_geojson(geojson_path: str) -> set:
+    """
+    Force every feature to a single valid Polygon.
+
+    Returns the set of cell ids whose features were dropped.
+    """
+    with open(geojson_path) as f:
+        data = json.load(f)
+
+    clean = []
+    dropped_cells = set()
+    for feat in data["features"]:
+        geom = shape(feat["geometry"])
+        if not geom.is_valid:
+            geom = shapely.make_valid(geom)
+        poly = None
+        if geom.geom_type == "Polygon":
+            poly = geom
+        elif geom.geom_type == "MultiPolygon":
+            poly = max(geom.geoms, key=lambda g: g.area)
+        elif geom.geom_type == "GeometryCollection":
+            polys = [g for g in geom.geoms if g.geom_type == "Polygon"]
+            if polys:
+                poly = max(polys, key=lambda g: g.area)
+        if poly is not None and not poly.is_empty:
+            feat["geometry"] = mapping(poly)
+            clean.append(feat)
+        else:
+            cell_id = feat.get("id") or feat.get("properties", {}).get("cell_id", "")
+            dropped_cells.add(str(cell_id))
+
+    print(f"GeoJSON: {len(clean)} kept, {len(dropped_cells)} dropped: {dropped_cells}")
+    data["features"] = clean
+    with open(geojson_path, "w") as f:
+        json.dump(data, f)
+
+    return dropped_cells
+
+
+def reassign_dropped(csv_path: str, dropped_cells: set) -> None:
+    """
+    Reassign transcripts of dropped cells to UNASSIGNED in the CSV.
+    """
+    if not dropped_cells:
+        return
+
+    with open(csv_path) as f:
+        reader = csv.DictReader(f)
+        fieldnames = reader.fieldnames
+        rows = list(reader)
+
+    reassigned = 0
+    for row in rows:
+        if row["cell"] in dropped_cells:
+            row["cell"] = ""
+            row["is_noise"] = "1"
+            reassigned += 1
+
+    with open(csv_path, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        writer.writerows(rows)
+    print(f"CSV: {reassigned} transcripts reassigned to UNASSIGNED")
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Clean stitched GeoJSON polygons and reconcile transcript CSV."
+    )
+    parser.add_argument("--geojson", required=True, help="Path to xr-cell-polygons.geojson")
+    parser.add_argument("--csv", required=True, help="Path to xr-transcript-metadata.csv")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    dropped = clean_geojson(args.geojson)
+    reassign_dropped(args.csv, dropped)