From b464c0a57ff240d08ecb1d2fcc59c07087398081 Mon Sep 17 00:00:00 2001
From: Jonathan Schwartz <jtschwar@gmail.com>
Date: Mon, 6 Jul 2026 09:39:18 -0700
Subject: [PATCH 1/6] improve stability for model exploration with slurm

---
 octopi/entry_points/run_optuna.py |  20 +++-
 octopi/pytorch/pg_server.py       | 187 ++++++++++++++++++++++++++++++
 octopi/pytorch/search_helper.py   | 159 ++++++++++++++++++++-----
 octopi/pytorch/submit_search.py   |  94 ++++++++++++++-
 pyproject.toml                    |   1 +
 5 files changed, 423 insertions(+), 38 deletions(-)
 create mode 100644 octopi/pytorch/pg_server.py

diff --git a/octopi/entry_points/run_optuna.py b/octopi/entry_points/run_optuna.py
index 821c0dd..35f053e 100755
--- a/octopi/entry_points/run_optuna.py
+++ b/octopi/entry_points/run_optuna.py
@@ -39,12 +39,16 @@
               help='GPU constraint to use for SLURM jobs (e.g., "a6000" or "l40,a6000")')
 @click.option('--timeout', type=int, default=4,
               help="SLURM job timeout per trial when using submitit (hours)")
+@click.option('--db-backend', type=click.Choice(['sqlite', 'postgres'], case_sensitive=False),
+              default=None,
+              help="Optuna storage backend. Default: 'postgres' with --submitit "
+                   "(safe for multi-node concurrent workers), else 'sqlite'.")
 @common.config_parameters(single_config=False)
 def cli(
     config, tomo_uris, target_uri, study_name,
     trainrunids, validaterunids, data_split, model_type, num_epochs, background_ratio,
     val_interval, ncache_tomos, best_metric, num_trials, random_seed, output,
-    submitit, njobs, cpu_constraint, gpu_constraint, timeout):
+    submitit, njobs, cpu_constraint, gpu_constraint, timeout, db_backend):
     """
     Perform model architecture search with Optuna.
     """
@@ -57,12 +61,13 @@ def cli(
         trainrunids, validaterunids, data_split, model_type, background_ratio,
         num_epochs, val_interval, ncache_tomos, best_metric, num_trials, random_seed, output,
         submitit=submitit, njobs=njobs, cpu_constraint=cpu_constraint, gpu_constraint=gpu_constraint, timeout=timeout,
+        db_backend=db_backend,
     )
 
-def run_model_explore(config, tomo_uris, target_info, study_name, 
+def run_model_explore(config, tomo_uris, target_info, study_name,
         trainrunids, validaterunids, data_split, model_type, background_ratio,
-        num_epochs, val_interval, ncache_tomos, best_metric, num_trials, random_seed, 
-        output, submitit, njobs, cpu_constraint, gpu_constraint, timeout):
+        num_epochs, val_interval, ncache_tomos, best_metric, num_trials, random_seed,
+        output, submitit, njobs, cpu_constraint, gpu_constraint, timeout, db_backend=None):
     """
     Run the model exploration (local GPUs or SLURM via submitit).
     """
@@ -78,8 +83,15 @@ def run_model_explore(config, tomo_uris, target_info, study_name,
     # Create the model exploration directory
     os.makedirs(output, exist_ok=True)
 
+    # Default storage: postgres for multi-node submitit runs (SQLite locks on
+    # networked filesystems), sqlite for single-node local runs.
+    if db_backend is None:
+        db_backend = 'postgres' if submitit else 'sqlite'
+    print(f"🗄️  Optuna storage backend: {db_backend}\n")
+
     # Base keyword arguments for both local and submitit explorers
     base_kwargs = dict(
+        db_backend=db_backend,
         copick_config=copick_configs,
         target_name=target_info[0],
         target_user_id=target_info[1],
diff --git a/octopi/pytorch/pg_server.py b/octopi/pytorch/pg_server.py
new file mode 100644
index 0000000..b2f41a2
--- /dev/null
+++ b/octopi/pytorch/pg_server.py
@@ -0,0 +1,187 @@
+"""Ephemeral PostgreSQL server for multi-node Optuna model-explore.
+
+SQLite cannot be shared safely by concurrent SLURM jobs running on different
+nodes over a network filesystem (Lustre): WAL mode requires same-host shared
+memory, and rollback-journal locking is both slow and corruption-prone across
+nodes. Optuna's own guidance is to use a real server DB for multi-worker runs.
+
+This module lets the model-explore *supervisor* stand up a private PostgreSQL
+instance on its own node. Worker jobs on other nodes then connect over TCP via
+the ``postgresql://<host>:<port>/<db>`` URL the supervisor hands them. The data
+directory lives on shared storage so the DB survives a supervisor restart.
+
+Only one process (this server) ever touches the data-dir files, from one node,
+so there is no cross-node filesystem-locking problem — Postgres handles all
+concurrency for the many worker connections.
+"""
+
+from __future__ import annotations
+
+import atexit
+import getpass
+import os
+import signal
+import socket
+import subprocess
+import time
+
+import psycopg2
+
+
+def _find_free_port() -> int:
+    """Pick a currently-free TCP port. Small TOCTOU race, acceptable here."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("", 0))
+        return s.getsockname()[1]
+
+
+class PostgresServer:
+    """Manage the lifecycle of a private PostgreSQL cluster for one search.
+
+    Typical use (context manager)::
+
+        with PostgresServer(data_dir=f"{output}/pgdata") as pg:
+            storage_url = pg.url  # postgresql://user@host:port/optuna
+            ... run workers ...
+    """
+
+    def __init__(
+        self,
+        data_dir: str,
+        dbname: str = "optuna",
+        port: int | None = None,
+        host: str | None = None,
+        user: str | None = None,
+    ):
+        self.data_dir = os.path.abspath(data_dir)
+        self.dbname = dbname
+        self.user = user or getpass.getuser()
+        # Hostname must be resolvable from the worker nodes; on SLURM clusters
+        # the short node name (e.g. "cpu-b-6") is cluster-wide resolvable.
+        self.host = host or socket.gethostname()
+        self.port = port or _find_free_port()
+        self.log_file = os.path.join(self.data_dir, "postgres.log")
+        # Unix socket lives in a short, node-local, writable dir (NOT the Lustre
+        # data dir, whose long path would exceed the ~107-char socket limit, and
+        # NOT the default /var/run/postgresql, which users can't write to).
+        # Admin connections use loopback TCP anyway; workers use TCP over the net.
+        self.socket_dir = f"/tmp/pg_sock_{self.user}_{self.port}"
+        self._started = False
+        self._stopped = False
+
+    # -- URLs -----------------------------------------------------------------
+    @property
+    def url(self) -> str:
+        """SQLAlchemy/Optuna URL that worker nodes use to connect over TCP."""
+        return f"postgresql://{self.user}@{self.host}:{self.port}/{self.dbname}"
+
+    def _admin_dsn(self, dbname: str = "postgres") -> dict:
+        """Local (loopback) connection params for admin tasks like CREATE DATABASE."""
+        return dict(host="127.0.0.1", port=self.port, user=self.user, dbname=dbname)
+
+    # -- lifecycle ------------------------------------------------------------
+    def _initdb_if_needed(self) -> None:
+        version_marker = os.path.join(self.data_dir, "PG_VERSION")
+        if os.path.exists(version_marker):
+            return  # cluster already initialized (persisted on shared storage)
+        os.makedirs(self.data_dir, exist_ok=True)
+        print(f"[postgres] initializing cluster at {self.data_dir}", flush=True)
+        subprocess.run(
+            ["initdb", "-D", self.data_dir, "-U", self.user,
+             "--auth-local=trust", "--auth-host=trust", "--encoding=UTF8"],
+            check=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+        )
+        # Allow worker nodes to connect over TCP. Data is non-sensitive Optuna
+        # trial state on a private cluster network, so trust auth is acceptable.
+        hba = os.path.join(self.data_dir, "pg_hba.conf")
+        with open(hba, "a") as f:
+            f.write("\n# Added by octopi PostgresServer: allow cluster nodes\n")
+            f.write("host    all    all    0.0.0.0/0    trust\n")
+
+    def _clear_stale_pidfile(self) -> None:
+        """A hard-killed supervisor can leave postmaster.pid behind, blocking start."""
+        if self.is_running():
+            return
+        pid_file = os.path.join(self.data_dir, "postmaster.pid")
+        if os.path.exists(pid_file):
+            print("[postgres] removing stale postmaster.pid", flush=True)
+            os.remove(pid_file)
+
+    def is_running(self) -> bool:
+        r = subprocess.run(["pg_ctl", "-D", self.data_dir, "status"],
+                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+        return r.returncode == 0
+
+    def _create_db_if_needed(self) -> None:
+        # Wait for the postmaster to accept connections, then ensure the DB exists.
+        last_err = None
+        for _ in range(30):
+            try:
+                conn = psycopg2.connect(**self._admin_dsn(), connect_timeout=5)
+                conn.autocommit = True
+                with conn.cursor() as cur:
+                    cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (self.dbname,))
+                    if cur.fetchone() is None:
+                        cur.execute(f'CREATE DATABASE "{self.dbname}"')
+                        print(f"[postgres] created database {self.dbname!r}", flush=True)
+                conn.close()
+                return
+            except psycopg2.OperationalError as e:
+                last_err = e
+                time.sleep(1)
+        raise RuntimeError(f"Postgres did not become ready: {last_err}")
+
+    def start(self) -> "PostgresServer":
+        if self._started:
+            return self
+        self._initdb_if_needed()
+        self._clear_stale_pidfile()
+        if not self.is_running():
+            os.makedirs(self.socket_dir, exist_ok=True)
+            print(f"[postgres] starting on {self.host}:{self.port} "
+                  f"(data={self.data_dir})", flush=True)
+            subprocess.run(
+                ["pg_ctl", "-D", self.data_dir, "-l", self.log_file, "-w",
+                 "-o", f"-p {self.port} -c listen_addresses='*' "
+                       f"-c unix_socket_directories={self.socket_dir}", "start"],
+                check=True,
+            )
+        self._create_db_if_needed()
+        self._started = True
+        self._stopped = False
+        # Defensive teardown so a crashing/terminated supervisor does not orphan
+        # the postmaster. SIGKILL can't be caught; a stale pidfile is cleaned on
+        # the next start() instead.
+        atexit.register(self.stop)
+        for sig in (signal.SIGTERM, signal.SIGINT):
+            try:
+                prev = signal.getsignal(sig)
+                signal.signal(sig, self._make_signal_handler(sig, prev))
+            except (ValueError, OSError):
+                pass  # not in main thread; atexit still covers normal exit
+        print(f"[postgres] ready — storage url: {self.url}", flush=True)
+        return self
+
+    def _make_signal_handler(self, sig, prev):
+        def _handler(signum, frame):
+            self.stop()
+            if callable(prev):
+                prev(signum, frame)
+            else:
+                raise SystemExit(128 + signum)
+        return _handler
+
+    def stop(self) -> None:
+        if self._stopped or not self._started:
+            return
+        self._stopped = True
+        print("[postgres] stopping server", flush=True)
+        subprocess.run(["pg_ctl", "-D", self.data_dir, "-m", "fast", "-w", "stop"],
+                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+
+    # -- context manager ------------------------------------------------------
+    def __enter__(self) -> "PostgresServer":
+        return self.start()
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.stop()
diff --git a/octopi/pytorch/search_helper.py b/octopi/pytorch/search_helper.py
index 23747f0..84cc9b1 100644
--- a/octopi/pytorch/search_helper.py
+++ b/octopi/pytorch/search_helper.py
@@ -7,6 +7,7 @@
 import torch, mlflow, optuna, logging
 from sqlite3 import OperationalError
 from optuna.trial import TrialState
+from optuna.storages._heartbeat import get_heartbeat_thread, is_heartbeat_enabled
 from typing import Set, Tuple, List
 import torch.multiprocessing as mp
 from octopi.datasets import config
@@ -20,6 +21,7 @@
 MAX_RESTARTS_PER_GPU = 25
 CHECK_DB_EVERY_S = 120 # 2 minutes, to detect external study changes
 PRINT_EVERY_S = 120 # print submitit once per minute
+SNAPSHOT_EVERY_S = 60 # refresh the local sqlite dashboard snapshot (postgres backend only)
 
 @dataclass
 class WorkerSpec:
@@ -69,20 +71,104 @@ def _retry_optuna_db(callable_fn, *args, **kwargs):
 def count_terminal_trials(study) -> int:
     return len(study.get_trials(states=TERMINAL_STATES))
 
+def _finalize_stale_trial(study, trial_id: int) -> None:
+    """Close out one RUNNING trial that's actually dead (heartbeat stale, or
+    the process exited without ever calling tell()).
+
+    A walltime/SIGKILL is not a failure of the hyperparameters — the trial
+    usually already has a real, if incomplete, best intermediate score. We
+    record that as COMPLETE (tagged status=timed_out), matching the manual
+    convention used to reclassify the first batch of stuck trials, so the
+    sampler treats it as genuine data instead of poisoning history with FAIL.
+    Only a trial with zero intermediate values (died before any progress) has
+    no usable signal — that one is marked FAIL, and we invoke the storage's
+    failed_trial_callback ourselves (set_trial_state_values doesn't trigger
+    it) so RetryFailedTrialCallback can requeue it.
+    """
+    storage = study._storage
+    frozen = storage.get_trial(trial_id)
+    if frozen.state != TrialState.RUNNING:
+        return  # already finalized by another worker/process
+
+    if frozen.intermediate_values:
+        best_value = max(frozen.intermediate_values.values())
+        if storage.set_trial_state_values(trial_id, state=TrialState.COMPLETE, values=[best_value]):
+            storage.set_trial_user_attr(trial_id, "status", "timed_out")
+            storage.set_trial_user_attr(trial_id, "finalized_by", "auto-reclaim")
+            print(f"[reclaim] trial {frozen.number} timed out -> COMPLETE (value={best_value:.4f})", flush=True)
+    else:
+        if storage.set_trial_state_values(trial_id, state=TrialState.FAIL):
+            print(f"[reclaim] trial {frozen.number} died with no progress -> FAIL", flush=True)
+            callback = storage.get_failed_trial_callback()
+            if callback is not None:
+                callback(study, storage.get_trial(trial_id))
+
+def reclaim_stale_trials(study) -> None:
+    """Finalize any RUNNING trial whose heartbeat has gone stale (dead job).
+
+    optuna.storages.fail_stale_trials only runs automatically inside
+    study.optimize()'s internal loop (right before each ask()), and it always
+    marks stale trials FAIL. We use ask/tell directly, so nothing ever called
+    it — heartbeats were being recorded but never checked, and a killed job's
+    trial stayed RUNNING forever. We call the equivalent ourselves here, but
+    via _finalize_stale_trial so a genuine partial result isn't recorded as a
+    failure.
+    """
+    storage = study._storage
+    if not is_heartbeat_enabled(storage):
+        return
+    try:
+        stale_trial_ids = storage._get_stale_trial_ids(study._study_id)
+    except Exception as e:
+        print(f"[reclaim] could not list stale trials (will retry next cycle): {e}", flush=True)
+        return
+    for trial_id in stale_trial_ids:
+        try:
+            _finalize_stale_trial(study, trial_id)
+        except optuna.exceptions.UpdateFinishedTrialError:
+            continue  # race: another process already finalized it
+        except Exception as e:
+            print(f"[reclaim] error finalizing trial {trial_id}: {e}", flush=True)
+
+def finalize_trial_if_running(study, trial) -> None:
+    """Close out a trial still RUNNING when its process is exiting (e.g. a
+    graceful walltime signal). We use Optuna's ask/tell API, which — unlike
+    ``study.optimize`` — does not record heartbeats on its own, so without
+    this the trial would stay RUNNING forever. Cannot help a hard SIGKILL —
+    the heartbeat-based reclaim_stale_trials covers that case instead.
+    """
+    try:
+        _finalize_stale_trial(study, trial._trial_id)
+    except Exception as e:
+        # Best-effort cleanup; never mask the original exit path.
+        print(f"[finalize] could not finalize trial {getattr(trial, 'number', '?')}: {e}", flush=True)
+
 def make_storage(storage_url: str):
-    # For multi-worker: prefer Postgres/MySQL. SQLite can lock.
+    # Engine kwargs are dialect-specific: SQLite's connect_args (timeout /
+    # check_same_thread) are invalid for the psycopg2 driver, and vice versa.
+    # For multi-worker runs across nodes, prefer Postgres (SQLite locks on
+    # networked filesystems). See octopi.pytorch.pg_server.
+    if storage_url.startswith("postgresql") or storage_url.startswith("postgres"):
+        engine_kwargs = {
+            "pool_pre_ping": True,      # drop dead connections from the pool
+            "pool_size": 5,
+            "max_overflow": 10,
+            "connect_args": {"connect_timeout": 30},
+        }
+    else:  # sqlite (single-node / local runs)
+        engine_kwargs = {
+            "connect_args": {
+                "timeout": 300,  # 5 minutes to wait on a lock
+                "check_same_thread": False,  # allow multi-threaded access
+            },
+            "pool_pre_ping": True,
+        }
     return optuna.storages.RDBStorage(
         url=storage_url,
         heartbeat_interval=60,
         grace_period=600,
         failed_trial_callback=optuna.storages.RetryFailedTrialCallback(max_retry=1),
-        engine_kwargs={
-            "connect_args": {
-                "timeout": 300,  # 5 minutes timeout for lock acquisition
-                "check_same_thread": False  # Allow multi-threaded access
-            },
-            "pool_pre_ping": True  # Verify connections before using
-        }
+        engine_kwargs=engine_kwargs,
     )
 
 def get_sampler():
@@ -143,7 +229,8 @@ def gpu_worker_loop(
     # Main loop: keep asking for new trials until enough trials have been completed or something breaks
     while not stop_event.is_set():
         try:
-            # Ask for a new trial (retry on DB lock)
+            # Reclaim any trial whose job died without a heartbeat update, then ask (retry on DB lock)
+            reclaim_stale_trials(study)
             trial = _retry_optuna_db(study.ask)
 
             # Verbose to show data splits (# runs / tomograms) for only the first trial
@@ -156,15 +243,16 @@ def gpu_worker_loop(
             model_search = ModelExplorer(data_generator, submit_kwargs["model_type"], submit_kwargs["output"])            
 
             try:
-                value = model_search.objective(
-                    trial=trial,
-                    epochs=int(submit_kwargs.get("num_epochs", 100)),
-                    device=device,
-                    val_interval=int(submit_kwargs.get("val_interval", 10)),
-                    best_metric=str(submit_kwargs.get("best_metric", "avg_f1")),
-                )
-                _retry_optuna_db(study.tell, trial, value)
-                print(f"[worker {gpu_id}] COMPLETE trial={trial.number} value={value}", flush=True)
+                with get_heartbeat_thread(trial._trial_id, storage):
+                    value = model_search.objective(
+                        trial=trial,
+                        epochs=int(submit_kwargs.get("num_epochs", 100)),
+                        device=device,
+                        val_interval=int(submit_kwargs.get("val_interval", 10)),
+                        best_metric=str(submit_kwargs.get("best_metric", "avg_f1")),
+                    )
+                    _retry_optuna_db(study.tell, trial, value)
+                    print(f"[worker {gpu_id}] COMPLETE trial={trial.number} value={value}", flush=True)
 
             except optuna.TrialPruned:
                 _retry_optuna_db(study.tell, trial, state=TrialState.PRUNED)
@@ -189,6 +277,10 @@ def gpu_worker_loop(
                     pass
                 time.sleep(2)
 
+            finally:
+                # Ensure a killed/interrupted trial does not linger in RUNNING.
+                finalize_trial_if_running(study, trial)
+
         except Exception as e:
             # This catches Optuna/DB-level issues (e.g. sqlite lock)
             print(f"[worker {gpu_id}] LOOP ERROR: {e}", flush=True)
@@ -235,6 +327,8 @@ def _on_walltime(signum, frame):
         val_interval=val_interval,
         n_warmup_steps=n_warmup_steps
     )
+    # Reclaim any trial whose job died without a heartbeat update, then ask for our own.
+    reclaim_stale_trials(study)
     trial = _retry_optuna_db(study.ask)
     trial_num = trial.number
     print(f"[Trial {trial_num}] START trial", flush=True)
@@ -253,17 +347,21 @@ def _on_walltime(signum, frame):
     print(f"[Trial {trial.number}] epochs={nepochs} val_interval={val_interval}", flush=True)
     print(f"[Trial {trial.number}] pruner={type(pr).__name__} - warmup={getattr(pr,'_n_warmup_steps',None)} - interval={getattr(pr,'_interval_steps',None)}", flush=True)
 
+    # Heartbeat thread keeps this RUNNING trial's heartbeat fresh so that a
+    # hard-killed job (SIGKILL / node loss) is reclaimed by fail_stale_trials
+    # after grace_period. ask/tell does not record heartbeats on its own.
     try:
-        value = model_search.objective(
-            trial=trial,
-            epochs=int(submit_kwargs.get("num_epochs", 100)),
-            device=device,
-            val_interval=int(submit_kwargs.get("val_interval", 10)),
-            best_metric=str(submit_kwargs.get("best_metric", "avg_f1")),
-        )
-        _retry_optuna_db(study.tell, trial, value)
-        print(f"[Trial {trial_num}] COMPLETE value={value}", flush=True)
-        return value
+        with get_heartbeat_thread(trial._trial_id, storage):
+            value = model_search.objective(
+                trial=trial,
+                epochs=int(submit_kwargs.get("num_epochs", 100)),
+                device=device,
+                val_interval=int(submit_kwargs.get("val_interval", 10)),
+                best_metric=str(submit_kwargs.get("best_metric", "avg_f1")),
+            )
+            _retry_optuna_db(study.tell, trial, value)
+            print(f"[Trial {trial_num}] COMPLETE value={value}", flush=True)
+            return value
     except optuna.TrialPruned:
         _retry_optuna_db(study.tell, trial, state=TrialState.PRUNED)
         print(f"[Trial {trial_num}] PRUNED", flush=True)
@@ -285,6 +383,11 @@ def _on_walltime(signum, frame):
         except Exception:
             pass
         raise
+    finally:
+        # Graceful-shutdown backstop: if a walltime signal ended training
+        # without any tell() path running, close the trial so it does not
+        # linger in RUNNING forever.
+        finalize_trial_if_running(study, trial)
 
 #--------------------------------
 # SLURM GPU Query Verification 
diff --git a/octopi/pytorch/submit_search.py b/octopi/pytorch/submit_search.py
index c7daf01..59d66e7 100644
--- a/octopi/pytorch/submit_search.py
+++ b/octopi/pytorch/submit_search.py
@@ -25,7 +25,10 @@ def __init__(
         ntomo_cache: int = 15,
         trainRunIDs: List[str] = None, validateRunIDs: List[str] = None,
         study_name: str = 'explore',
-        background_ratio: float = 0.0
+        background_ratio: float = 0.0,
+        db_backend: str = 'sqlite',
+        pg_dbname: str = 'optuna',
+        pg_port: int = None,
     ):
         """
         Initialize the ModelSearch class for architecture search with Optuna.
@@ -67,6 +70,34 @@ def __init__(
         self.validateRunIDs = validateRunIDs
         self.data_split = data_split
         self.background_ratio = background_ratio
+        self.db_backend = db_backend
+        self.pg_dbname = pg_dbname
+        self.pg_port = pg_port
+        self._pg_server = None  # set when db_backend == 'postgres'
+
+    def _resolve_storage_url(self, output: str) -> str:
+        """Return the Optuna storage URL, starting a Postgres server if requested.
+
+        - 'sqlite'   -> file DB at {output}/trials.db (fine for single-node runs).
+        - 'postgres' -> start a private PostgreSQL server on this (supervisor)
+                        node and return a TCP URL the worker nodes can reach.
+                        Required for multi-node submitit runs; SQLite locks on
+                        networked filesystems like Lustre.
+        """
+        if self.db_backend == "postgres":
+            from octopi.pytorch.pg_server import PostgresServer
+            self._pg_server = PostgresServer(
+                data_dir=os.path.join(output, "pgdata"),
+                dbname=self.pg_dbname,
+                port=self.pg_port,
+            ).start()
+            return self._pg_server.url
+        return f"sqlite:///{output}/trials.db"
+
+    def _stop_storage_server(self) -> None:
+        """Tear down the Postgres server (no-op for sqlite)."""
+        if self._pg_server is not None:
+            self._pg_server.stop()
 
     def _setup_study_and_storage(self, study_name: str, output: str):
         """Create output dir, Optuna study, and save parameters. Returns (storage_url, storage, study, submit_kwargs)."""
@@ -75,12 +106,35 @@ def _setup_study_and_storage(self, study_name: str, output: str):
         submit_kwargs["output"] = output
         self.save_parameters(submit_kwargs, output)
         os.makedirs(output, exist_ok=True)
-        storage_url = f"sqlite:///{output}/trials.db"
+        storage_url = self._resolve_storage_url(output)
         storage = helper.make_storage(storage_url)
         study = helper.get_study(study_name, storage, self.val_interval, self.n_warmup_steps)
         mlflow.set_experiment(study_name)
         return storage_url, storage, study, submit_kwargs
 
+    def _snapshot_dashboard_db(self, storage_url: str, study_name: str, output: str) -> None:
+        """Refresh a local SQLite mirror of the study for the VS Code Optuna
+        Dashboard extension, which only opens local sqlite files (it runs
+        client-side via WASM with no TCP support, so it can't reach 'postgres'
+        directly). No-op for the sqlite backend, since trials.db is already a
+        file the extension can open. Writes to a temp file and atomically
+        renames it so a viewer never opens a half-written snapshot.
+        """
+        if self.db_backend != "postgres":
+            return
+        dest = os.path.join(output, "trails.db")
+        tmp_dest = dest + ".tmp"
+        if os.path.exists(tmp_dest):
+            os.remove(tmp_dest)
+        try:
+            optuna.copy_study(
+                from_study_name=study_name, from_storage=storage_url,
+                to_storage=f"sqlite:///{tmp_dest}",
+            )
+            os.replace(tmp_dest, dest)
+        except Exception as e:
+            print(f"[dashboard-snapshot] skipped (will retry): {e}", flush=True)
+
     def _run_worker_pool(
         self,
         storage_url: str,
@@ -111,6 +165,7 @@ def start_worker(gid: int):
             start_worker(gid)
 
         last_db_check = 0.0
+        last_snapshot = 0.0
         storage = helper.make_storage(storage_url)
         try:
             while True:
@@ -133,6 +188,7 @@ def start_worker(gid: int):
                     last_db_check = now
                     try:
                         study = optuna.load_study(study_name=study_name, storage=storage)
+                        helper.reclaim_stale_trials(study)
                         done = helper.count_terminal_trials(study)
                         print(f"[supervisor] done={done}/{self.num_trials}", flush=True)
                         if done >= self.num_trials:
@@ -147,6 +203,10 @@ def start_worker(gid: int):
                         else:
                             raise
 
+                if now - last_snapshot >= helper.SNAPSHOT_EVERY_S:
+                    last_snapshot = now
+                    self._snapshot_dashboard_db(storage_url, study_name, submit_kwargs["output"])
+
                 time.sleep(helper.POLL_S)
         except KeyboardInterrupt:
             print("[supervisor] Ctrl-C received. stopping.", flush=True)
@@ -173,9 +233,16 @@ def run_model_search(self, study_name: str = 'octopi_nas', output: str = 'explor
         storage_url, storage, _study, submit_kwargs = self._setup_study_and_storage(
             study_name, output
         )
-        self._run_worker_pool(storage_url, study_name, submit_kwargs)
-        study = optuna.load_study(study_name=study_name, storage=storage)
-        self.save_contour_plot_as_png(study, output)
+        try:
+            self._run_worker_pool(storage_url, study_name, submit_kwargs)
+            study = optuna.load_study(study_name=study_name, storage=storage)
+            self.save_contour_plot_as_png(study, output)
+            # Final snapshot so the dashboard file reflects the finished study,
+            # not just whatever the last periodic tick captured.
+            self._snapshot_dashboard_db(storage_url, study_name, output)
+        finally:
+            # Always shut down the Postgres server (no-op for sqlite).
+            self._stop_storage_server()
 
     def save_contour_plot_as_png(self, study, output):
         """
@@ -330,6 +397,8 @@ def _run_worker_pool(
 
         # Printing status every helper.PRINT_EVERY_S seconds
         last_print = 0
+        last_snapshot = 0
+        last_reclaim = 0
 
         # Start the submitit job pool
         storage = helper.make_storage(storage_url)
@@ -357,10 +426,19 @@ def submit_one():
 
             # Check study progress
             study = optuna.load_study(study_name=study_name, storage=storage)
+
+            # Reclaim trials whose job died without a heartbeat update (e.g.
+            # SIGKILL/timeout) so RetryFailedTrialCallback can requeue them.
+            # Runs on the supervisor's own cadence, independent of whether any
+            # worker happens to call ask() (which also reclaims on its own).
+            now = time.time()
+            if now - last_reclaim >= helper.CHECK_DB_EVERY_S:
+                last_reclaim = now
+                helper.reclaim_stale_trials(study)
+
             done_count = helper.count_terminal_trials(study)
 
             # Print status periodically
-            now = time.time()
             if now - last_print >= helper.PRINT_EVERY_S:
                 print(
                     f"[submitit] done={done_count}/{num_trials} running={len(running)}",
@@ -368,6 +446,10 @@ def submit_one():
                 )
                 last_print = now
 
+            if now - last_snapshot >= helper.SNAPSHOT_EVERY_S:
+                last_snapshot = now
+                self._snapshot_dashboard_db(storage_url, study_name, submit_kwargs["output"])
+
             if done_count >= num_trials and not running:
                 break
 
diff --git a/pyproject.toml b/pyproject.toml
index c2b39c7..e9a027a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -33,6 +33,7 @@ dependencies = [
     "nibabel",
     "mlflow",
     "optuna",
+    "psycopg2-binary",
     "monai",
     "pandas",
     "plotly",

From 757883581dd006ee560d7ebfb30bf4002744fccc Mon Sep 17 00:00:00 2001
From: jtschwar <jtschwar@gmail.com>
Date: Thu, 16 Jul 2026 10:53:02 -0700
Subject: [PATCH 2/6] add HF hub checkpoints, restructure extract into command
 group

---
 .gitignore                                    |   9 ++
 docs/api/inference.md                         |  39 ++++---
 docs/api/quick-start.md                       |  10 +-
 docs/getting-started/installation.md          |  12 +-
 docs/getting-started/quickstart.md            |   5 +-
 docs/user-guide/claude-code-mcp.md            |   4 +-
 docs/user-guide/data-import.md                |   5 +-
 docs/user-guide/inference.md                  |  85 +++++++++++---
 docs/user-guide/labels.md                     |  34 +++---
 octopi/entry_points/common.py                 |  11 +-
 octopi/entry_points/groups.py                 |  22 ++--
 octopi/entry_points/run_create_targets.py     |  43 ++++----
 octopi/entry_points/run_extract_mb_picks.py   |  10 +-
 octopi/entry_points/run_segment.py            |  56 ++++------
 octopi/extract/cli.py                         |  18 +++
 octopi/extract/segmentations.py               |  96 ++++++++++++++++
 octopi/main.py                                |   4 +-
 octopi/mcp/server.py                          |  43 +++++---
 octopi/processing/importers.py                |   2 +-
 .../pytorch/{segmentation.py => inference.py} |  10 ++
 octopi/utils/hub.py                           | 104 ++++++++++++++++++
 octopi/workflows.py                           |  42 ++++---
 pyproject.toml                                |   1 +
 23 files changed, 491 insertions(+), 174 deletions(-)
 create mode 100644 octopi/extract/cli.py
 create mode 100644 octopi/extract/segmentations.py
 rename octopi/pytorch/{segmentation.py => inference.py} (97%)
 create mode 100644 octopi/utils/hub.py

diff --git a/.gitignore b/.gitignore
index 42f70ac..c895181 100644
--- a/.gitignore
+++ b/.gitignore
@@ -166,3 +166,12 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# Local scratch notes / codebase-state snapshots (not for version control)
+notes/
+
+# Downloaded Hugging Face checkpoint cache
+octopi/cache/
+
+# Claude markdown files
+CLAUDE.md
\ No newline at end of file
diff --git a/docs/api/inference.md b/docs/api/inference.md
index b21768e..a22ab6f 100644
--- a/docs/api/inference.md
+++ b/docs/api/inference.md
@@ -7,9 +7,10 @@ This page covers running inference with trained octopi models, including segment
 Run trained models on tomograms to generate segmentation masks. The segmentation process takes your trained model weights and configuration to produce probability maps for each object class defined in your training targets.
 
 !!! info "Segmentation Parameters"
-    - **model_weights**: Path to your trained model weights (`.pth` file from training)
-    - **model_config**: Path to model configuration (`.yaml` file from training)
-    - **seg_info**: Tuple defining where to save segmentation results (`name`, `user_id`, `session_id`)
+    - **model_weights**: Path to your trained model weights (`.pth` file from training), or a pretrained checkpoint alias (e.g. `'tomogram-boundary'`) to auto-download from the [Hugging Face Hub](https://huggingface.co/biohub/octopi)
+    - **model_config**: Path to model configuration (`.yaml` file from training). Omit when `model_weights` is a checkpoint alias — its config is bundled and downloaded automatically
+    - **tomo_uri**: Tomogram URI in the form `"algorithm@voxel_size"` (e.g. `"denoised@10.0"`)
+    - **seg_uri**: Segmentation output URI in the form `"name:user_id/session_id"` (e.g. `"predict:octopi/1"`)
     - **use_tta**: Whether to use test-time augmentation for improved robustness
     - **run_ids**: Optional list of specific tomograms to process (None for all available)
 
@@ -27,7 +28,7 @@ The algorithm automatically applies size filtering based on object radii defined
 
     ```python
     import numpy as np
-    from octopi.pytorch.segmentation import Predictor
+    from octopi.pytorch.inference import Predictor
 
     model_weights = 'model_output/best_model.pth'
     model_config = 'model_output/model_config.yaml'
@@ -59,32 +60,42 @@ The algorithm automatically applies size filtering based on object radii defined
     config = 'eval_config.json'
     model_weights = 'model_output/best_model.pth'
     model_config = 'model_output/model_config.yaml'
-    seg_info = ['predict', 'octopi', '1']  # (name, user_id, session_id)
 
     segment(
         config=config,
-        tomo_algorithm='denoised',
-        voxel_size=10.0,
         model_weights=model_weights,
         model_config=model_config,
-        seg_info=seg_info,
+        tomo_uri='denoised@10.0',
+        seg_uri='predict:octopi/1',
         ntta=4
     )
     ```
 
+    You can also skip training entirely and segment with a pretrained checkpoint from the
+    [Hugging Face Hub](https://huggingface.co/biohub/octopi) — the weights and matching
+    config are downloaded and cached automatically, so `model_config` can be omitted:
+
+    ```python
+    segment(
+        config=config,
+        model_weights='tomogram-boundary',
+        tomo_uri='denoised@10.0',
+        seg_uri='predict:octopi/1',
+    )
+    ```
+
     <details markdown="1">
     <summary><strong>💡 segment() reference</strong></summary>
 
-    `segment(config, tomo_algorithm, voxel_size, model_weights, model_config, seg_info=['predict', 'octopi', '1'], run_ids=None, batch_size=1, ntta=4)`
+    `segment(config, model_weights, model_config=None, tomo_uri='wbp@10.0', seg_uri='predict:octopi/1', run_ids=None, batch_size=1, swbs=4, overlap=0.5, ntta=4)`
 
     **Parameters:**
 
     - `config` (str): Path to CoPick configuration file
-    - `tomo_algorithm` (str): Tomogram algorithm identifier
-    - `voxel_size` (float): Voxel spacing in Angstroms
-    - `model_weights` (str or list): Path(s) to trained model weights (.pth file(s))
-    - `model_config` (str or list): Path(s) to model configuration (.yaml file(s))
-    - `seg_info` (list): Output segmentation identifier `(name, user_id, session_id)`
+    - `model_weights` (str or list): Path(s) to trained model weights (.pth file(s)), or a pretrained checkpoint alias (e.g. `'tomogram-boundary'`)
+    - `model_config` (str or list): Path(s) to model configuration (.yaml file(s)). May be omitted when `model_weights` is a checkpoint alias
+    - `tomo_uri` (str): Tomogram URI in the form `"algorithm@voxel_size"` (default: `'wbp@10.0'`)
+    - `seg_uri` (str): Segmentation output URI in the form `"name:user_id/session_id"` (default: `'predict:octopi/1'`)
     - `run_ids` (list): Specific run IDs to process (default: None — all runs)
     - `batch_size` (int): Tomograms processed concurrently per GPU (default: 1)
     - `ntta` (int): Number of test-time augmentation rotations (default: 4; set to 1 to disable)
diff --git a/docs/api/quick-start.md b/docs/api/quick-start.md
index f3307b1..8d48dd7 100644
--- a/docs/api/quick-start.md
+++ b/docs/api/quick-start.md
@@ -91,11 +91,10 @@ from octopi.workflows import segment
 
 segment(
     config=config,
-    tomo_algorithm='denoised',
-    voxel_size=10.012,
     model_weights=f'{results_folder}/best_model.pth',
     model_config=f'{results_folder}/model_config.yaml',
-    seg_info=['predict', 'octopi', '1'],
+    tomo_uri='denoised@10.012',
+    seg_uri='predict:octopi/1',
     ntta=4  # number of test-time augmentation rotations
 )
 ```
@@ -200,11 +199,10 @@ evaluate(
     print("Step 3: Running segmentation...")
     segment(
         config=config,
-        tomo_algorithm=tomo_algorithm,
-        voxel_size=voxel_size,
         model_weights=f'{results_folder}/best_model.pth',
         model_config=f'{results_folder}/model_config.yaml',
-        seg_info=['predict', 'octopi', '1'],
+        tomo_uri=f'{tomo_algorithm}@{voxel_size}',
+        seg_uri='predict:octopi/1',
         ntta=4
     )
 
diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md
index 35cffd6..e92400d 100644
--- a/docs/getting-started/installation.md
+++ b/docs/getting-started/installation.md
@@ -50,12 +50,12 @@ The editable (-e) install ensures that local code changes are immediately reflec
     │ model-explore   Perform model architecture search with Optuna.                  │
     ╰─────────────────────────────────────────────────────────────────────────────────╯
     ╭─ Inference ─────────────────────────────────────────────────────────────────────╮
-    │ segment           Segment volumes using trained neural network models.          │
-    │ localize          Convert Segmentation Masks to 3D Particle Coordinates.        │
-    │ membrane-extract  Extract membrane-bound picks based on proximity to organelle  │
-    │                   or membrane segmentation.                                     │
-    │ evaluate          Evaluate particle localization performance against ground     │
-    │                   truth annotations.                                            │
+    │ segment    Segment volumes using trained neural network models.                 │
+    │ localize   Convert Segmentation Masks to 3D Particle Coordinates.               │
+    │ extract    Extract objects or picks from existing pipeline outputs (isolate a   │
+    │            single object's mask, or split picks by membrane proximity).         │
+    │ evaluate   Evaluate particle localization performance against ground truth      │
+    │            annotations.                                                         │
     ╰─────────────────────────────────────────────────────────────────────────────────╯
     ```
 
diff --git a/docs/getting-started/quickstart.md b/docs/getting-started/quickstart.md
index 233a459..746eac6 100644
--- a/docs/getting-started/quickstart.md
+++ b/docs/getting-started/quickstart.md
@@ -17,10 +17,9 @@ Create semantic masks for your proteins of interest using annotation metadata:
 ```bash
 octopi create-targets \
     --config config.json \
-    --tomo-alg wbp --voxel-size 10 \
+    --tomo-uri wbp@10.0 \
     --picks-user-id data-portal --picks-session-id 0 \
-    --target-session-id 1 --target-segmentation-name targets \
-    --target-user-id octopi
+    --target-uri targets:octopi/1
 ```
 
 🎯 This creates training targets for a single copick query. To produce targets from multiple coordinate queries,  refer to the [Prepare Labels](../user-guide/labels.md) section.
diff --git a/docs/user-guide/claude-code-mcp.md b/docs/user-guide/claude-code-mcp.md
index a210602..d8918f9 100644
--- a/docs/user-guide/claude-code-mcp.md
+++ b/docs/user-guide/claude-code-mcp.md
@@ -32,7 +32,7 @@ Once connected, you can describe what you want in plain language and Claude hand
 
     ---
 
-    Executes `create-targets`, `localize`, `evaluate`, and `membrane-extract` directly and reports back.
+    Executes `create-targets`, `localize`, `evaluate`, and `extract` directly and reports back.
 
 -   :material-chip:{ .lg .middle } **Hand off GPU jobs**
 
@@ -126,7 +126,7 @@ octopi create-targets \
     --target ribosome,manual,1 \
     --target virus-like-particle,tm,2 \
     --target membranes,membrane-seg,1 \
-    --voxel-size 10
+    --tomo-uri wbp@10.0
 ```
 
 ??? tip "URI shorthand for tomograms, segmentations, and picks"
diff --git a/docs/user-guide/data-import.md b/docs/user-guide/data-import.md
index 93df02b..65d4ddf 100644
--- a/docs/user-guide/data-import.md
+++ b/docs/user-guide/data-import.md
@@ -78,9 +78,7 @@ Copick provides a CLI for generating configuration files that mount either local
         --proj-description "Synaptic Vesicles collected on 24sep24"
     ```
 
-    We can define either objects that are continuous segmentations (e.g., organelles or memebranes) or coordinates. For pickable objects, we can store meta-data including the particle radius and a corresponding PDB-ID:
-        
-        - `--objects name,is_particle,radius,pdb_id`.
+    We can define either objects that are continuous segmentations (e.g., organelles or memebranes) or coordinates. For pickable objects, we can store meta-data including the particle radius and a corresponding PDB-ID: `--objects name,is_particle,radius,pdb_id`.
 
     ??? info "`copick config filesystem` parameters"
 
@@ -124,6 +122,7 @@ If you have tomograms stored locally in `*.mrc` format (e.g., from Warp, IMOD, o
 
     ```bash 
     copick add tomogram \
+        ".mrc" \
         -c /path/to/config.json \
         --tomo-type sart \
     ```
diff --git a/docs/user-guide/inference.md b/docs/user-guide/inference.md
index 6eba361..fd6bb9b 100644
--- a/docs/user-guide/inference.md
+++ b/docs/user-guide/inference.md
@@ -8,7 +8,8 @@ Octopi inference follows a systematic two-step approach:
 
 1. **Segmentation** - Apply trained model to generate 3D probability masks with test-time augmentation (TTA).
 2. **Localization** - Convert probability masks into 3D coordinates using size-based filtering.
-3. **Evaluation (Optional)** - Compare predicted coordinates against ground truth annotations.
+3. **Extraction (Optional)** - Isolate a single object from a prediction, or split picks by membrane proximity.
+4. **Evaluation (Optional)** - Compare predicted coordinates against ground truth annotations.
 
 ??? tip "Parallelism and Resource Utilization"
 
@@ -53,8 +54,8 @@ octopi segment \
 
         | Parameter | Description | Notes |
         |----------|-------------|------|
-        | `--model-config` | Model configuration file(s). | Required; comma-separated for ensembles |
-        | `--model-weights` | Model weight file(s). | Must match `--model-config` order |
+        | `--model-config` | Model configuration file(s). | Required for local weights; comma-separated for ensembles. Omit when `--model-weights` is a Hugging Face checkpoint alias |
+        | `--model-weights` | Model weight file(s), or a pretrained checkpoint alias. | Must match `--model-config` order for local files |
 
     === "Inference"
 
@@ -76,6 +77,24 @@ octopi segment \
     --seg-uri ensemble:octopi/1
 ```
 
+### Pretrained Checkpoints
+
+In development! Only one pre-trained model is available. More to come.
+
+!!! tip "Skip training with a pretrained checkpoint"
+
+    Pass a checkpoint name from the  Hugging Face Hub repo as `--model-weights`. It's downloaded and cached automatically, so `--model-config` can be omitted.
+
+    ```bash
+    octopi segment \
+        --config config.json \
+        --tomo-uri wbp@10.0 \
+        --model-weights tomogram-boundary \
+        --seg-uri predict:octopi/1
+    ```
+
+    💡 - See the [model card](https://huggingface.co/biohub/octopi) for the full list of available checkpoints.
+
 ---
 
 ## Localization 
@@ -120,13 +139,58 @@ The localization algorithm uses **particle size information** from your copick c
         | `--pick-session-id` | Session ID for particle picks. | `1` | Used for result grouping |
         | `--pick-user-id` | User ID for particle picks. | `octopi` | Used for result grouping |
 
-## Context-Aware Particle Extraction
+## Extraction (Optional)
+
+`octopi extract` is a command group for post-processing outputs you've already generated — isolating a single object from a raw multi-class prediction, or splitting picks by proximity to a membrane/organelle segmentation.
 
-This optional post-processing step splits an existing set of particle picks into two groups:
+### Isolate a Single Object
+
+Pull one object's mask out of a raw multi-class `octopi segment` prediction and save it as its own standalone segmentation.
+
+```bash
+octopi extract seg \
+    --config config.json \
+    --seg-uri predict:octopi/1 \
+    --name membranes \
+    --session-id 1
+```
+
+Octopi reads the inference log written by `octopi segment` to recover the voxel size and the object's integer label within the raw prediction, then writes a binary mask for just that object as a new segmentation.
+
+??? info "`octopi extract seg -h`"
+
+    === "Input"
+
+        | Parameter | Description | Default | Notes |
+        |----------|-------------|---------|------|
+        | `--config` | Path to the CoPick configuration file. | – | Required |
+        | `--name` | Object name to extract from the raw multi-class prediction. | – | Required. Example: `membranes` |
+        | `--seg-uri` | Source segmentation to extract from (`name:user_id/session_id`). | `predict:octopi/1` | Must be a prediction written by `octopi segment` |
+        | `--run-ids` | Specific run IDs to process. | All runs | Example: `run1,run2` |
+
+    === "Output"
+
+        | Parameter | Description | Default | Notes |
+        |----------|-------------|---------|------|
+        | `--user-id` | User ID for the extracted segmentation. | Source segmentation's user ID | |
+        | `--session-id` | Session ID for the extracted segmentation. | `1` | |
+
+### Membrane-Proximity Picks
+
+This step splits an existing set of particle picks into two groups:
 
 - **Membrane-close picks**: particles within a configurable distance threshold of a membrane/organelle segmentation.
 - **Membrane-far picks**: particles outside that threshold.
 
+```bash
+octopi extract mb-picks \
+    --config config.json \
+    --picks-uri ribosome:octopi/1 \
+    --seg-uri membrane:membrain-seg/1 \
+    --save-user-id octopi \
+    --save-session-id 10
+```
+
 For membrane-close particles, we can also **align orientations** so that each particle’s rotation is consistent with the local membrane normal (estimated from the vector between the particle and the closest organelle center).
 
 ??? tip "What this is useful for"
@@ -240,17 +304,6 @@ octopi evaluate
 
 ---
 
-## Visualization
-
-To visualize your results and validate the quality of segmentations and coordinates, refer to our interactive notebook:
-
-**📓 [Inference Notebook](https://github.com/chanzuckerberg/octopi/blob/main/notebooks/inference.ipynb)**
-
-With this notebook, we can overlay the segmentation masks or coordiantes the tomograms. 
-
-![Coordinates Visualization](../assets/coordinates.png)
-*Example of predicted particle coordinates displayed on a holdout tomogram from cryo-ET training dataset. The visualization shows Octopi's localization results overlaid on tomographic data from [DatasetID: 10440](https://cryoetdataportal.czscience.com/datasets/10440).*
-
 ## Next Steps
 
 You now have a complete workflow for applying Octopi models to new tomographic data. The inference pipeline transforms your trained models into actionable scientific results through robust segmentation, intelligent localization, and comprehensive evaluation.
diff --git a/docs/user-guide/labels.md b/docs/user-guide/labels.md
index b600e1a..ed76b50 100644
--- a/docs/user-guide/labels.md
+++ b/docs/user-guide/labels.md
@@ -6,11 +6,11 @@ We will use Copick to manage the filesystem, extract runIDs, and create spherica
 
 * **Generating Targets**: For each tomogram, we extract particle coordinates and generate spherical targets based on these coordinates, and save the target data in OME Zarr format.
 
-* **Target dimensions** are determined with an associated tomogram, (specified by the `--tomo-alg` and `--voxel-size` parameters).
+* **Target dimensions** are determined with an associated tomogram, (specified by the `--tomo-uri` parameter, e.g. `wbp@10.0`).
 
 * **Additional segmentations** like organelles and membranes can be included as continuous targets with the `--seg-target` flag. 
 
-The segmentations are saved with the associated `target-session-id`, `target-user-id` and `target-name` flags.  
+The segmentations are saved under the query specified by the `--target-uri` flag (`name:user_id/session_id`).  
 
 ## Method 1: Automated Query
 
@@ -27,9 +27,8 @@ octopi create-targets \
     --config config.json \
     --picks-user-id data-portal --picks-session-id 0 \
     --seg-target membrane \
-    --tomo-alg wbp --voxel-size 10 \
-    --target-session-id 1 --target-segmentation-name targets \
-    --target-user-id octopi
+    --tomo-uri wbp@10.0 \
+    --target-uri targets:octopi/1
 ```
 
 This command automatically finds all pickable objects associated with `data-portal` user and session `0`, plus includes membrane segmentations.
@@ -48,20 +47,22 @@ For more control, manually specify which protein types and annotation sources to
 ```bash
 octopi create-targets \
     --config config.json \
-    --target apoferritin --target beta-galactosidase,slabpick,1 \
-    --target ribosome,pytom,0 --target virus-like-particle,pytom,0 \
+    --target apoferritin --target beta-galactosidase:slabpick/1 \
+    --target ribosome:pytom/0 --target virus-like-particle:pytom/0 \
     --seg-target membrane \
-    --tomo-alg wbp --voxel-size 10 \
-    --target-session-id 1 --target-segmentation-name targets \
-    --target-user-id octopi
+    --tomo-uri wbp@10.0 \
+    --target-uri targets:octopi/1
 ```
 
 ### Target Specification Formats
 
+`--target` and `--seg-target` both accept the same query grammar:
+
 | Format | Description | Example |
 |--------|-------------|---------|
 | `protein_name` | Use default user/session from config | `apoferritin` |
-| `protein_name,user_id,session_id` | Specify source explicitly | `ribosome,pytom,0` |
+| `protein_name:user_id/session_id` | Specify source explicitly (preferred) | `ribosome:pytom/0` |
+| `protein_name,user_id,session_id` | Legacy comma form (still supported) | `ribosome,pytom,0` |
 
 ## Check Target Quality
 
@@ -79,10 +80,10 @@ This notebook shows how to load segmentation targets and overlay targets on tomo
 | Parameter | Description | Example | Required |
 |-----------|-------------|---------|----------|
 | `--config` | Path to CoPick configuration file | `config.json` | ✅ |
-| `--target` | Target specifications: "name" or "name,user_id,session_id" | `ribosome,pytom,0` | * |
+| `--target` | Pickable object target(s): "name" or "name:user_id/session_id" | `ribosome:pytom/0` | * |
 | `--picks-session-id` | Session ID for automated pick retrieval | `0` | * |
 | `--picks-user-id` | User ID for automated pick retrieval | `data-portal` | * |
-| `--seg-target` | Segmentation targets: "name" or "name,user_id,session_id" | `membrane` | ❌ |
+| `--seg-target` | Continuous segmentation target(s): "name" or "name:user_id/session_id" | `membrane` | ❌ |
 | `--run-ids` | Specific run IDs to process | `run_001,run_002` | ❌ |
 
 *Either `--target` OR both `--picks-session-id` and `--picks-user-id` must be specified.
@@ -91,17 +92,14 @@ This notebook shows how to load segmentation targets and overlay targets on tomo
 
 | Parameter | Description | Default | Example |
 |-----------|-------------|---------|---------|
-| `--tomo-alg` | Tomogram reconstruction algorithm | `wbp` | `wbp`, `denoised` |
+| `--tomo-uri` | Tomogram URI in the form `alg@voxel_size` | `wbp@10.0` | `denoised@10.0` |
 | `--radius-scale` | Scale factor for object radius | `0.8` | `0.8` (80% of original radius) |
-| `--voxel-size` | Voxel size for tomogram reconstruction | `10` | `10` (10 Ångström) |
 
 ### Output Arguments
 
 | Parameter | Description | Default | Example |
 |-----------|-------------|---------|---------|
-| `--target-segmentation-name` | Name for the target segmentation | `targets` | `my_targets` |
-| `--target-user-id` | User ID for target segmentation | `octopi` | `my_experiment` |
-| `--target-session-id` | Session ID for target segmentation | `1` | `1` |
+| `--target-uri` | Output query as `name`, `name:user_id`, or `name:user_id/session_id` | `targets:octopi/1` | `my_targets:my_experiment/1` |
 
 
  
\ No newline at end of file
diff --git a/octopi/entry_points/common.py b/octopi/entry_points/common.py
index d309732..afbf965 100755
--- a/octopi/entry_points/common.py
+++ b/octopi/entry_points/common.py
@@ -22,9 +22,11 @@ def inference_model_parameters():
     """Decorator for adding inference model parameters"""
     def decorator(f):
         f = click.option("-mw", "--model-weights", type=str, required=True,
-                        help="Path to the model weights file")(f)
-        f = click.option("-mc", "--model-config", type=str, required=True,
-                        help="Path to the model configuration file")(f)
+                        help="Path to the model weights file, or a pretrained checkpoint alias "
+                             "(e.g. 'tomogram-boundary') to auto-download from the Hugging Face Hub")(f)
+        f = click.option("-mc", "--model-config", type=str, required=False, default=None,
+                        help="Path to the model configuration file. Omit when --model-weights is a "
+                             "Hugging Face checkpoint alias; its config is bundled and downloaded automatically")(f)
         return f
     return decorator
 
@@ -86,8 +88,7 @@ def decorator(f):
                         callback=lambda ctx, param, value: parsers.parse_list(value) if value else None,
                         help="List of run IDs for prediction, e.g., run1,run2 or [run1,run2]. If not provided, all available runs will be processed.")(f)
         f = click.option('-suri', "--seg-uri", type=str, default='predict:octopi/1',
-                        callback=lambda ctx, param, value: parsers.parse_target(value) if value else value,
-                        help='Information Query to save Segmentation predictions under (e.g., "name" or "name:user_id/session_id"')(f)
+                        help='Segmentation output URI to save predictions under (e.g., "name" or "name:user_id/session_id")')(f)
         f = click.option('-swbs', "--sliding-window-batch-size", default=4, type=IntRange(min=1),
                         help="Batch size for sliding window inference")(f)
         f = click.option('--overlap', '-o', default=0.5, type=FloatRange(0.0, 1.0),
diff --git a/octopi/entry_points/groups.py b/octopi/entry_points/groups.py
index 02417fb..a61ae3f 100644
--- a/octopi/entry_points/groups.py
+++ b/octopi/entry_points/groups.py
@@ -17,11 +17,7 @@
         },
         {
             "name": "Inference",
-            "commands": ["segment", "localize", "membrane-extract", "evaluate"]
-        },
-        {
-            "name": "nnUNet",
-            "commands": ["nnunet"]
+            "commands": ["segment", "localize", "evaluate", "extract"]
         }
     ]
 }
@@ -57,11 +53,11 @@
         },
         {
             "name": "Parameters",
-            "options": ["--tomo-alg", "--radius-scale", "--voxel-size"]
+            "options": ["--tomo-uri", "--radius-scale"]
         },
         {
             "name": "Output Arguments",
-            "options": ["--target-segmentation-name", "--target-user-id", "--target-session-id"]
+            "options": ["--target-uri"]
         }
     ],
     "routines segment": [
@@ -125,7 +121,7 @@
             "options": ["--save-path"]
         }
     ],
-    "routines membrane-extract": [
+    "routines extract mb-picks": [
         {
             "name": "Input Arguments",
             "options": ["--config", "--voxel-size", "--picks-uri",
@@ -140,6 +136,16 @@
             "options": ["--save-user-id", "--save-session-id"]
         }
     ],
+    "routines extract seg": [
+        {
+            "name": "Input Arguments",
+            "options": ["--config", "--name", "--seg-uri", "--run-ids"]
+        },
+        {
+            "name": "Output Arguments",
+            "options": ["--user-id", "--session-id"]
+        }
+    ],
     "routines download": [
         {
             "name": "Input Arguments",
diff --git a/octopi/entry_points/run_create_targets.py b/octopi/entry_points/run_create_targets.py
index 1c13c63..fa8dd91 100755
--- a/octopi/entry_points/run_create_targets.py
+++ b/octopi/entry_points/run_create_targets.py
@@ -143,54 +143,57 @@ def add_segmentation_targets(
 
 @click.command('create-targets', no_args_is_help=True)
 # Output Arguments
-@click.option('-sid', '--target-session-id', type=str, default="1",
-              help="Session ID for the target segmentation")
-@click.option('-uid','--target-user-id', type=str, default="octopi",
-              help="User ID associated with the target segmentation")
-@click.option('-name', '--target-segmentation-name', type=str, default='targets',
-              help="Name for the target segmentation")
+@click.option('-turi', '--target-uri', type=str, default="targets:octopi/1",
+              callback=lambda ctx, param, value: parsers.parse_target(value),
+              help="Target query as 'name', 'name:user_id', or 'name:user_id/session_id'. Default 'targets:octopi/1'.")
 # Parameters
-@click.option('-vs', '--voxel-size', type=float, default=10,
-              help="Voxel size for tomogram reconstruction")
+@click.option('-uri', '--tomo-uri', type=str, default="wbp@10.0",
+              help="Tomogram URI for target dimensions (tomo-alg@voxel-size)")
 @click.option('-rs', '--radius-scale', type=float, default=0.7,
               help="Scale factor for object radius")
-@click.option('-alg', '--tomo-alg', type=str, default="wbp",
-              help="Tomogram reconstruction algorithm")
 # Input Arguments
-@click.option('--run-ids', type=str, default=None,
+@click.option('--run-ids', '-runs', type=str, default=None,
               callback=lambda ctx, param, value: parsers.parse_list(value) if value else None,
               help="List of run IDs")
 @click.option('--seg-target', type=str, multiple=True,
               callback=lambda ctx, param, value: [parsers.parse_target(v) for v in value] if value else [],
-              help='Segmentation targets: "name" or "name,user_id,session_id"')
-@click.option('--picks-user-id', type=str, default=None,
+              help='Continuous segmentation target(s): "name", "name:user_id/session_id", or the legacy "name,user_id,session_id"')
+@click.option('--picks-user-id', '-puid', type=str, default=None,
               help="User ID associated with the picks")
-@click.option('--picks-session-id', type=str, default=None,
+@click.option('--picks-session-id', '-psid', type=str, default=None,
               help="Session ID for the picks")
 @click.option('-t', '--target', type=str, multiple=True,
               callback=lambda ctx, param, value: [parsers.parse_target(v) for v in value] if value else None,
-              help='Target specifications: "name" or "name,user_id,session_id"')
+              help='Pickable object target(s): "name", "name:user_id/session_id", or the legacy "name,user_id,session_id"')
 @click.option('-c', '--config', type=click.Path(exists=True), required=True,
               help="Path to the CoPick configuration file")
 def cli(config, target, picks_session_id, picks_user_id, seg_target, run_ids,
-        tomo_alg, radius_scale, voxel_size,
-        target_segmentation_name, target_user_id, target_session_id):
+        tomo_uri, radius_scale, target_uri):
     """
     Generate segmentation targets from CoPick configurations.
 
     This tool allows users to specify target labels for training in two ways:
 
-    1. Manual Specification: Define a subset of pickable objects using --target name or --target name,user_id,session_id
+    1. Manual Specification: Define a subset of pickable objects using --target name or --target name:user_id/session_id
 
     2. Automated Query: Provide --picks-session-id and/or --picks-user-id to automatically retrieve all pickable objects
 
     Example Usage:
 
-        Manual: octopi create-targets --config config.json --target ribosome --target apoferritin,123,456
+        Manual: octopi create-targets --config config.json --target ribosome --target apoferritin:manual/1 --tomo-uri wbp@10.0
 
-        Automated: octopi create-targets --config config.json --picks-session-id 123 --picks-user-id 456
+        Automated: octopi create-targets --config config.json --picks-session-id 123 --picks-user-id 456 --tomo-uri wbp@10.0
     """
 
+    # Parse the Tomogram URI
+    if '@' not in tomo_uri:
+        raise ValueError("Tomogram URI must contain '@' for voxel size, e.g. 'wbp@10.0'.")
+    tomo_alg, voxel_size = tomo_uri.split('@')
+    voxel_size = float(voxel_size)
+
+    # Parse the Target URI
+    target_segmentation_name, target_user_id, target_session_id = target_uri
+
     # Print Summary To User
     print('\n⚙️ Generating Target Segmentation Masks from the Following Copick-Query:')
     if target is not None and len(target) > 0:
diff --git a/octopi/entry_points/run_extract_mb_picks.py b/octopi/entry_points/run_extract_mb_picks.py
index 0d593b6..5ea992e 100755
--- a/octopi/entry_points/run_extract_mb_picks.py
+++ b/octopi/entry_points/run_extract_mb_picks.py
@@ -62,7 +62,7 @@ def save_parameters( params_dict: dict, output_path: str ):
     io.save_parameters_yaml(params_dict, output_path)
 
 
-@click.command('membrane-extract', no_args_is_help=True)
+@click.command('mb-picks', no_args_is_help=True)
 # Output Arguments
 @click.option('-ssid','--save-session-id', type=str, required=True,
               help="Session ID to save the new picks")
@@ -74,7 +74,7 @@ def save_parameters( params_dict: dict, output_path: str ):
 @click.option('-t', '--threshold', type=str, default="1,10",
               help="Distance threshold for membrane proximity in Voxels (provide the min and max as 'min,max' if only one value is provided, it is used as max with min=1)",)
 # Input Arguments
-@click.option('-rids','--runIDs', type=str, default=None,
+@click.option('-runs','--runIDs', type=str, default=None,
               callback=lambda ctx, param, value: parsers.parse_list(value) if value else None,
               help="List of run IDs to process")
 @click.option('-suri','--seg-uri', type=str, default=None,
@@ -116,14 +116,14 @@ def cli(config, voxel_size, picks_uri, seg_uri, runids,
     Examples:
 
     # Extract membrane-bound picks with default distance threshold (1–10 voxels)
-    octopi membrane-extract -c config.json \\
+    octopi extract mb-picks -c config.json \\
         --picks-uri predictions:octopi/1 \\
         --seg-uri membrane:membrain-seg/1 \\
         --save-user-id octopi \\
         --save-session-id 1
 
     # Use a custom distance range (2–6 voxels)
-    octopi membrane-extract -c config.json \\
+    octopi extract mb-picks -c config.json \\
         --picks-uri predictions:octopi/1 \\
         --seg-uri membrane:membrain-seg/1 \\
         --threshold 2,6 \\
@@ -162,7 +162,7 @@ def run_mb_extract(
     overlay_root = io.remove_prefix(root.config.overlay_root)
     basepath = os.path.join(overlay_root, 'logs')
     os.makedirs(basepath, exist_ok=True)
-    output_yaml = f'membrane-extract_{save_user_id}_{save_session_id}.yaml'
+    output_yaml = f'extract-mb-picks_{save_user_id}_{save_session_id}.yaml'
     output_path = os.path.join(basepath, output_yaml)        
 
     # Save parameters
diff --git a/octopi/entry_points/run_segment.py b/octopi/entry_points/run_segment.py
index d6ab06f..f9ba4d8 100755
--- a/octopi/entry_points/run_segment.py
+++ b/octopi/entry_points/run_segment.py
@@ -1,14 +1,13 @@
 from octopi.entry_points import common
-from typing import List, Tuple
+from typing import List
 import rich_click as click
 
 def inference(
     config: str,
-    model_weights: str, 
+    model_weights: str,
     model_config: str,
-    seg_info: Tuple[str,str,str],
-    voxel_size: float,
-    tomo_algorithm: str,
+    tomo_uri: str,
+    seg_uri: str,
     run_ids: List[str],
     swbs: int,
     overlap: float,
@@ -20,21 +19,17 @@ def inference(
     Args:
         config (str): Path to CoPick configuration file.
         run_ids (List[str]): List of tomogram run IDs for inference.
-        model_weights (str): Path to the trained model weights file.
-        channels (List[int]): List of channel sizes for each layer.
-        strides (List[int]): List of strides for the layers.
-        res_units (int): Number of residual units for the model.
-        voxel_size (float): Voxel size for tomogram reconstruction.
-        tomo_algorithm (str): Tomogram reconstruction algorithm to use.
-        segmentation_name (str): Name for the segmentation output.
-        segmentation_user_id (str): User ID associated with the segmentation.
-        segmentation_session_id (str): Session ID for this segmentation run.
+        model_weights (str): Path to the trained model weights file, or a pretrained checkpoint
+            alias (e.g. "tomogram-boundary") to auto-download from the Hugging Face Hub.
+        model_config (str): Path to the model configuration file.
+        tomo_uri (str): Tomogram URI in the form "algorithm@voxel_size".
+        seg_uri (str): Segmentation output URI in the form "name:user_id/session_id".
     """
     from octopi.workflows import segment
 
     if ',' in model_weights:
         model_weights = model_weights.split(',')
-    if ',' in model_config:
+    if model_config and ',' in model_config:
         model_config = model_config.split(',')
     if isinstance(model_weights, list) and isinstance(model_config, list):
         if len(model_weights) != len(model_config):
@@ -47,8 +42,8 @@ def inference(
 
 
     segment(
-        config, tomo_algorithm, voxel_size,
-        model_weights, model_config, seg_info,
+        config, model_weights, model_config,
+        tomo_uri=tomo_uri, seg_uri=seg_uri,
         run_ids=run_ids, swbs = swbs, overlap=overlap, ntta=ntta
     )
 
@@ -89,6 +84,13 @@ def cli(config, tomo_uri,
         --model-config model.yaml --model-weights model.pth \\
         --seg-uri predictions:octopi/1
 
+    \b
+      # Segment with a pretrained checkpoint from the Hugging Face Hub (auto-downloaded)
+      octopi segment -c config.json \\
+        --tomo-uri wbp@10.0 \\
+        --model-weights tomogram-boundary \\
+        --seg-uri predictions:octopi/1
+
     \b
       # Segment with model ensemble (comma-separated)
       octopi segment -c config.json \\
@@ -105,27 +107,15 @@ def cli(config, tomo_uri,
         --run-ids TS_001,TS_002,TS_003
     """
     
-    # Set default values if not provided
-    seg_info = list(seg_uri)  # Convert parsed (name, user, session) tuple to list
-    if seg_info[1] is None:
-        seg_info[1] = "octopi"
-    if seg_info[2] is None:
-        seg_info[2] = "1"
-
-    # Parse the tomogram URI
-    if '@' not in tomo_uri:
-        raise ValueError("Tomogram URI must contain '@' for voxel size.")
-    tomo_alg, voxel_size = tomo_uri.split('@')
-
-    # Call the inference function with parsed arguments
+    # Call the inference function with parsed arguments; tomo_uri/seg_uri parsing
+    # happens once, inside octopi.workflows.segment().
     print('\n🚀 Starting inference with Octopi...\n')
     inference(
         config=config,
         model_weights=model_weights,
         model_config=model_config,
-        seg_info=seg_info,
-        voxel_size=float(voxel_size),
-        tomo_algorithm=tomo_alg,
+        tomo_uri=tomo_uri,
+        seg_uri=seg_uri,
         run_ids=run_ids,
         swbs=sliding_window_batch_size,
         overlap=overlap,
diff --git a/octopi/extract/cli.py b/octopi/extract/cli.py
new file mode 100644
index 0000000..d48547c
--- /dev/null
+++ b/octopi/extract/cli.py
@@ -0,0 +1,18 @@
+from octopi.extract.segmentations import cli as extract_seg_cli
+from octopi.entry_points.run_extract_mb_picks import cli as mb_picks_cli
+import rich_click as click
+
+@click.group('extract', no_args_is_help=True)
+def cli():
+    """Extract objects or picks from existing pipeline outputs.
+
+    seg: isolate a single object's mask from a multi-class `segment` prediction.
+    mb-picks: split existing picks by proximity to a membrane/organelle segmentation.
+    """
+    pass
+
+cli.add_command(extract_seg_cli)
+cli.add_command(mb_picks_cli)
+
+if __name__ == "__main__":
+    cli()
\ No newline at end of file
diff --git a/octopi/extract/segmentations.py b/octopi/extract/segmentations.py
new file mode 100644
index 0000000..b317bb6
--- /dev/null
+++ b/octopi/extract/segmentations.py
@@ -0,0 +1,96 @@
+from typing import List, Optional
+from octopi.utils import parsers
+import rich_click as click
+
+def run_extract_seg(
+    config: str,
+    object_name: str,
+    seg_uri: str,
+    user_id: Optional[str],
+    session_id: str,
+    run_ids: Optional[List[str]],
+):
+    from octopi.utils import io
+    from copick_utils.io import readers, writers
+    import numpy as np
+    import copick
+
+    # Parse the source segmentation query (the raw multi-class prediction)
+    seg_name, seg_user_id, seg_session_id = parsers.parse_target(seg_uri)
+    seg_user_id = seg_user_id or 'octopi'
+    seg_session_id = seg_session_id or '1'
+
+    # Load the inference log to recover the voxel size and object -> label mapping
+    # used when the raw prediction was generated (written by `octopi segment`).
+    seg_config = io.get_config(config, seg_name, 'segment', seg_user_id, seg_session_id)
+    voxel_size = seg_config['inputs']['voxel_size']
+    labels = seg_config['labels']
+
+    if object_name not in labels:
+        raise ValueError(
+            f"Object '{object_name}' not found in the labels for segmentation "
+            f"'{seg_name}:{seg_user_id}/{seg_session_id}'. Available objects: {sorted(labels)}"
+        )
+    label = labels[object_name]
+
+    # Default the output user ID to the source segmentation's user ID
+    user_id = user_id or seg_user_id
+
+    # Load Copick Project and get run IDs if not provided
+    root = copick.from_file(config)
+    if run_ids is None:
+        run_ids = [run.name for run in root.runs]
+
+    print(
+        f"\n🔍 Extracting '{object_name}' (label={label}) from "
+        f"'{seg_name}:{seg_user_id}/{seg_session_id}' @ {voxel_size} "
+        f"-> '{object_name}:{user_id}/{session_id}'\n"
+    )
+
+    for run_id in run_ids:
+        run = root.get_run(run_id)
+        seg = readers.segmentation(run, voxel_size, seg_name, seg_user_id, seg_session_id, verbose=False)
+        if seg is None:
+            print(f"  ⚠️  {run_id}: no '{seg_name}' segmentation found at voxel size {voxel_size} — skipping.")
+            continue
+
+        out_seg = (seg == label).astype(np.uint8)
+        writers.segmentation(run, out_seg, user_id, object_name, session_id, voxel_size, multilabel=False)
+        print(f"  ✅ {run_id}: extracted {int(out_seg.sum()):,} voxels")
+
+    print("\n✅ Extraction complete.")
+
+
+@click.command('seg', no_args_is_help=True)
+@click.option('-c', '--config', type=click.Path(exists=True), required=True,
+              help="Path to the CoPick configuration file")
+@click.option('-n', '--name', 'object_name', type=str, required=True,
+              help="Object name to extract from the raw multi-class prediction (e.g. 'membranes')")
+@click.option('-uri', '--seg-uri', type=str, default='predict:octopi/1',
+              help='Source segmentation to extract from: "name", "name:user_id", or '
+                   '"name:user_id/session_id". Default "predict:octopi/1".')
+@click.option('-uid', '--user-id', type=str, default=None,
+              help="User ID for the extracted segmentation (defaults to the source segmentation's user ID)")
+@click.option('-sid', '--session-id', type=str, default='1',
+              help="Session ID for the extracted segmentation")
+@click.option('-runs', '--run-ids', type=str, default=None,
+              callback=lambda ctx, param, value: parsers.parse_list(value) if value else None,
+              help="List of run IDs to process. Defaults to all runs.")
+def cli(config, object_name, seg_uri, user_id, session_id, run_ids):
+    """
+    Extract a single object from a raw multi-class Octopi prediction and save it as its
+    own standalone CoPick segmentation.
+
+    Reads the inference log written by `octopi segment` to recover the voxel size and the
+    object's integer label within the raw prediction, then writes a binary mask for just
+    that object as a new segmentation.
+
+    \b
+    Example:
+      octopi extract seg --config config.json --name membranes
+    """
+    run_extract_seg(config, object_name, seg_uri, user_id, session_id, run_ids)
+
+
+if __name__ == "__main__":
+    cli()
diff --git a/octopi/main.py b/octopi/main.py
index 95e0eae..f4b64d0 100644
--- a/octopi/main.py
+++ b/octopi/main.py
@@ -9,7 +9,7 @@
 from octopi.entry_points.run_segment import cli as inference
 from octopi.entry_points.run_localize import cli as localize
 from octopi.entry_points.run_evaluate import cli as evaluate
-from octopi.entry_points.run_extract_mb_picks import cli as mb_extract
+from octopi.extract.cli import cli as extract_cli
 from octopi.mcp.cli import mcp_cli
 
 @click.group(context_settings=cli_context)
@@ -25,8 +25,8 @@ def routines():
 routines.add_command(localize)
 routines.add_command(model_explore)
 routines.add_command(evaluate)
-routines.add_command(mb_extract)
 routines.add_command(mcp_cli)
+routines.add_command(extract_cli)
 
 @click.group(context_settings=cli_context)
 def slurm_routines():
diff --git a/octopi/mcp/server.py b/octopi/mcp/server.py
index d6da08d..f3e1fc3 100644
--- a/octopi/mcp/server.py
+++ b/octopi/mcp/server.py
@@ -34,18 +34,19 @@
     → --seg-uri predict:octopi/1
 
   Pick/Target URI  "name:user_id/session_id"   e.g. "ribosome:manual/1"
-    → --target ribosome,manual,1  (for create-targets source picks)
-    → --picks-uri ribosome:manual/1  (for membrane-extract)
+    → --target ribosome:manual/1  (for create-targets source picks, repeatable)
+    → --seg-target membrane:membrane-seg/1  (for create-targets continuous/segmentation targets, repeatable)
+    → --picks-uri ribosome:manual/1  (for extract mb-picks)
     → --pick-user-id manual --pick-session-id 1  (for localize output)
-    → --target-uri targets:octopi/1  (for train / model-explore target segmentation)
+    → --target-uri targets:octopi/1  (for create-targets / train / model-explore target segmentation)
 
   Multiple URIs of the same type are passed as repeated flags, e.g.:
-    --target ribosome,manual,1 --target virus-like-particle,tm,2
+    --target ribosome:manual/1 --target virus-like-particle:tm/2
 
 STEP 1 — create-targets
   Convert pick coordinates from a CoPick project into 3D segmentation masks (Zarr).
   This is a fast command — you can run it directly.
-  Key params: --config, --target (pick URI → name,user_id,session_id), --voxel-size, --radius-scale
+  Key params: --config, --target (pick URI → name,user_id,session_id), --tomo-uri, --target-uri (output seg URI), --radius-scale
 
 STEP 2 — train OR model-explore
   train: Train a 3D U-Net model on tomogram/segmentation pairs. GPU-intensive, takes hours.
@@ -59,6 +60,9 @@
 STEP 3 — segment
   Run sliding-window inference on tomograms to produce probability maps.
   Supports model ensembling (comma-separated --model-config and --model-weights paths).
+  --model-weights also accepts a pretrained checkpoint alias (e.g. 'tomogram-boundary') to
+  auto-download from the biohub/octopi Hugging Face Hub repo — in that case omit
+  --model-config, since its config is bundled and downloaded automatically.
   GPU-intensive — always suggest rather than run unless the user explicitly asks.
   Key params: --config, --model-config, --model-weights, --tomo-uri (tomo URI), --seg-uri (seg URI)
 
@@ -72,19 +76,24 @@
   Fast command — can run directly.
   Key params: --config, --ground-truth-user-id, --predict-user-id
 
-STEP 6 (optional) — membrane-extract
-  Split picks by proximity to a membrane or organelle segmentation.
-  Fast command — can run directly.
-  Key params: --config, --picks-uri (pick URI), --seg-uri (seg URI), --threshold, --save-session-id
+STEP 6 (optional) — extract
+  Command group for post-processing existing pipeline outputs. Fast — can run directly.
+  extract mb-picks: split picks by proximity to a membrane or organelle segmentation.
+    Key params: --config, --picks-uri (pick URI), --seg-uri (seg URI), --threshold, --save-session-id
+  extract seg: isolate a single object's mask from a multi-class `segment` prediction.
+    Key params: --config, --name (object name), --seg-uri (source seg URI), --session-id
 
 HOW TO RESPOND
 - Use get_command_help to look up flags before suggesting a command.
 - ALWAYS suggest long-running commands (train, model-explore, segment) as copy-pasteable code blocks.
   NEVER run them unless the user explicitly says "run it", "go ahead and run it", or "execute it".
 - Describing what they want ("I'd like to train a model") is NOT permission to run — suggest instead.
-- For fast commands (create-targets, localize, evaluate, membrane-extract), you may run them directly
+- For fast commands (create-targets, localize, evaluate, extract), you may run them directly
   if the user has provided all required parameters.
 - If the user provides all required parameters, go straight to the suggestion without asking follow-up questions.
+- Default to CLI commands (this server only knows the CLI). Only suggest the Python API
+  (octopi.workflows) if the user is clearly programming — writing a script or notebook, not
+  running a one-off command — e.g. "how do I call segment from my script/notebook."
 """,
 )
 
@@ -95,7 +104,8 @@
     ("segment", "Run sliding-window inference to produce probability maps — GPU-intensive"),
     ("localize", "Convert segmentation maps to 3D particle coordinates — fast"),
     ("evaluate", "Measure Precision/Recall/F1 against ground truth — fast"),
-    ("membrane-extract", "Split picks by membrane proximity (alias: mb-extract) — fast"),
+    ("extract mb-picks", "Split picks by membrane proximity — fast"),
+    ("extract seg", "Isolate a single object's mask from a multi-class prediction — fast"),
 ]
 
 LONG_RUNNING = {"train", "model-explore", "segment"}
@@ -112,7 +122,7 @@ def list_octopi_commands() -> dict[str, Any]:
     return {
         "success": True,
         "commands": [{"command": f"octopi {cmd}", "description": desc} for cmd, desc in OCTOPI_COMMANDS],
-        "workflow_order": ["create-targets", "train or model-explore", "segment", "localize", "evaluate (optional)", "membrane-extract (optional)"],
+        "workflow_order": ["create-targets", "train or model-explore", "segment", "localize", "evaluate (optional)", "extract (optional)"],
         "tip": "Call get_command_help with a command name (e.g. 'train') to see all options.",
     }
 
@@ -122,9 +132,10 @@ def get_command_help(command: str) -> dict[str, Any]:
     """Get the full --help output for an octopi command.
 
     Args:
-        command: Command name, e.g. 'train', 'segment', 'create-targets', 'model-explore'.
+        command: Command name, e.g. 'train', 'segment', 'create-targets', or a group
+            subcommand like 'extract seg'.
     """
-    cmd = ["octopi", command, "--help"]
+    cmd = ["octopi", *command.split(), "--help"]
     try:
         result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
         help_text = result.stdout or result.stderr
@@ -145,12 +156,12 @@ def get_command_help(command: str) -> dict[str, Any]:
 def run_octopi_command(args: list[str], working_dir: str | None = None) -> dict[str, Any]:
     """Run an octopi command and return its output.
 
-    Use this for fast commands: create-targets, localize, evaluate, membrane-extract.
+    Use this for fast commands: create-targets, localize, evaluate, extract.
     For long-running GPU jobs (train, model-explore, segment), suggest the command as a
     copy-pasteable block instead — only run them if the user explicitly asks you to.
 
     Args:
-        args: Arguments after 'octopi', e.g. ['create-targets', '--config', 'config.json', '--voxel-size', '10'].
+        args: Arguments after 'octopi', e.g. ['create-targets', '--config', 'config.json', '--tomo-uri', 'wbp@10.0'].
         working_dir: Directory to run the command in. Defaults to cwd.
     """
     cwd = working_dir or os.getcwd()
diff --git a/octopi/processing/importers.py b/octopi/processing/importers.py
index f17b1aa..e96a585 100644
--- a/octopi/processing/importers.py
+++ b/octopi/processing/importers.py
@@ -83,7 +83,7 @@ def import_tomos(
     print(f'✅ Import Complete! Imported {len(tomograms)} tomograms')
 
 
-@click.command('import')
+@click.command('import', no_args_is_help=True)
 # Input Arguments
 @click.option('-p', '--path', type=click.Path(exists=True), default=None, required=True,
               help="Path to the folder containing the tomograms")
diff --git a/octopi/pytorch/segmentation.py b/octopi/pytorch/inference.py
similarity index 97%
rename from octopi/pytorch/segmentation.py
rename to octopi/pytorch/inference.py
index 002f7e7..48de6d9 100755
--- a/octopi/pytorch/segmentation.py
+++ b/octopi/pytorch/inference.py
@@ -51,6 +51,11 @@ def __init__(self,
             rank (int, optional): Rank of the current process for distributed inference. Defaults to 0.
         """
 
+        # Resolve checkpoint aliases (e.g. "tomogram-boundary") against the HF Hub;
+        # local paths pass through unchanged.
+        from octopi.utils.hub import resolve_model_source
+        model_weights, model_config = resolve_model_source(model_weights, model_config)
+
         # Open the Copick Project
         self.config = config
         self.root = copick.from_file(config)
@@ -585,6 +590,11 @@ def __init__(
         device: Optional[str] = None,  # ignored; we choose per-process cuda device
         compile_model: bool = False,
     ):
+        # Resolve checkpoint aliases once here so every spawned worker process
+        # receives already-local paths instead of racing to download from the Hub.
+        from octopi.utils.hub import resolve_model_source
+        model_weights, model_config = resolve_model_source(model_weights, model_config)
+
         self.config = config
         self.root = copick.from_file(config)
         self.model_config = model_config
diff --git a/octopi/utils/hub.py b/octopi/utils/hub.py
new file mode 100644
index 0000000..1e1b72e
--- /dev/null
+++ b/octopi/utils/hub.py
@@ -0,0 +1,104 @@
+"""Resolve pretrained checkpoint aliases against the biohub/octopi Hugging Face Hub repo.
+
+Repo layout on the Hub is one subfolder per checkpoint, each containing a weights file and a
+matching model config, e.g. `tomogram-boundary/weights.pth` + `tomogram-boundary/config.yaml`.
+"""
+from typing import List, Optional, Tuple, Union
+import os
+
+HF_REPO_ID = "biohub/octopi"
+WEIGHTS_FILENAME = "weights.pth"
+CONFIG_FILENAME = "config.yaml"
+
+# Checkpoint alias -> subfolder name in HF_REPO_ID.
+KNOWN_MODELS = {
+    "tomogram-boundary": "tomogram-boundary",
+}
+
+
+def default_cache_dir() -> str:
+    import octopi
+
+    return os.path.join(os.path.dirname(os.path.abspath(octopi.__file__)), "cache")
+
+
+def _resolve_one(
+    weights: str, config: Optional[str], cache_dir: str
+) -> Tuple[str, str]:
+    if os.path.exists(weights):
+        if config is None:
+            raise ValueError(
+                f"--model-config is required when --model-weights is a local path ({weights})."
+            )
+        return weights, config
+
+    if weights not in KNOWN_MODELS:
+        raise ValueError(
+            f"'{weights}' is not a local file and not a known octopi checkpoint alias. "
+            f"Known aliases: {sorted(KNOWN_MODELS)}"
+        )
+
+    from huggingface_hub import snapshot_download
+
+    subfolder = KNOWN_MODELS[weights]
+    try:
+        local_dir = snapshot_download(
+            repo_id=HF_REPO_ID,
+            allow_patterns=f"{subfolder}/*",
+            cache_dir=cache_dir,
+        )
+    except Exception as e:
+        raise ValueError(
+            f"Could not download checkpoint '{weights}' from the Hugging Face Hub repo "
+            f"'{HF_REPO_ID}' (subfolder '{subfolder}'). It may not be published yet, or "
+            f"there was a network/authentication issue. Original error: {e}"
+        ) from e
+
+    resolved_weights = os.path.join(local_dir, subfolder, WEIGHTS_FILENAME)
+    resolved_config = os.path.join(local_dir, subfolder, CONFIG_FILENAME)
+    if not os.path.exists(resolved_weights) or not os.path.exists(resolved_config):
+        raise ValueError(
+            f"Checkpoint '{weights}' is registered but no files were found under "
+            f"'{subfolder}/' in the Hugging Face Hub repo '{HF_REPO_ID}'. It may not be "
+            f"uploaded yet."
+        )
+    return resolved_weights, resolved_config
+
+
+def resolve_model_source(
+    model_weights: Union[str, List[str]],
+    model_config: Optional[Union[str, List[str]]] = None,
+    cache_dir: Optional[str] = None,
+) -> Tuple[Union[str, List[str]], Union[str, List[str]]]:
+    """
+    Resolve model_weights/model_config to local file paths, downloading from the Hub
+    when a value is a known checkpoint alias rather than an existing local path.
+
+    Accepts and returns either a single str or a list (for model-soup ensembles), matching
+    the shape of the input `model_weights`. Cached downloads are skipped on repeat calls.
+    """
+    cache_dir = cache_dir or default_cache_dir()
+    os.makedirs(cache_dir, exist_ok=True)
+
+    is_list = isinstance(model_weights, list)
+    weights_list = model_weights if is_list else [model_weights]
+
+    if model_config is None:
+        config_list = [None] * len(weights_list)
+    elif isinstance(model_config, list):
+        config_list = model_config
+    else:
+        config_list = [model_config] * len(weights_list)
+
+    if len(config_list) != len(weights_list):
+        raise ValueError("Number of model configs must match number of model weights.")
+
+    resolved = [
+        _resolve_one(w, c, cache_dir) for w, c in zip(weights_list, config_list)
+    ]
+    resolved_weights = [r[0] for r in resolved]
+    resolved_config = [r[1] for r in resolved]
+
+    if not is_list:
+        return resolved_weights[0], resolved_config[0]
+    return resolved_weights, resolved_config
diff --git a/octopi/workflows.py b/octopi/workflows.py
index e765a4d..93ccd6b 100644
--- a/octopi/workflows.py
+++ b/octopi/workflows.py
@@ -3,9 +3,9 @@
 from monai.metrics import ConfusionMatrixMetric
 from octopi.models import common as builder
 from octopi.models.common import wrap_loss_for_ds
-from octopi.pytorch import segmentation
-from octopi.pytorch import trainer 
-from octopi.utils import io
+from octopi.pytorch import inference
+from octopi.pytorch import trainer
+from octopi.utils import io, parsers
 import multiprocess as mp
 from pprint import pprint
 import copick, torch, os
@@ -103,36 +103,46 @@ def train(data_generator, loss_function, batch_size = 16,
     io.save_parameters_to_yaml(model_builder, model_trainer, data_generator, cfg_name)
     io.save_results_to_csv(results, os.path.join(model_save_path, "results.csv"))
 
-def segment(config, tomo_algorithm, voxel_size, model_weights, model_config, 
-            seg_info = ['predict', 'octopi', '1'], run_ids = None, batch_size = 1,
-            swbs = 4, overlap = 0.5, ntta = 4):
+def segment(config, model_weights, model_config = None,
+            tomo_uri = 'wbp@10.0', seg_uri = 'predict:octopi/1',
+            run_ids = None, batch_size = 1, swbs = 4, overlap = 0.5, ntta = 4):
     """
     Segment a Dataset using a Trained Model or Ensemble of Models
 
     Args:
         config (str): Path to the Copick Config File
-        tomo_algorithm (str): The tomographic algorithm to use for segmentation
-        voxel_size (float): The voxel size of the data
-        model_weights (str, list): The path to the model weights or a list of paths to the model weights
-        model_config (str, list): The model configuration or a list of model configurations
-        seg_info (list): The segmentation information
+        model_weights (str, list): The path to the model weights, a pretrained checkpoint alias, or a list of these for ensemble prediction
+        model_config (str, list): The model configuration or a list of model configurations. May be omitted when model_weights is a checkpoint alias
+        tomo_uri (str): Tomogram URI in the form "algorithm@voxel_size", e.g. "wbp@10.0"
+        seg_uri (str): Segmentation output URI in the form "name:user_id/session_id", e.g. "predict:octopi/1"
         swbs (int): The sliding window batch size for inference
         overlap (float): The overlap between sliding windows for inference
         ntta (int): The number of test-time augmentations for inference
         run_ids (list): The list of run IDs to use for segmentation
     """
 
+    # Parse the Tomogram URI
+    if '@' not in tomo_uri:
+        raise ValueError(f"Tomogram URI '{tomo_uri}' must contain '@' for voxel size, e.g. 'wbp@10.0'.")
+    tomo_algorithm, voxel_size = tomo_uri.split('@')
+    voxel_size = float(voxel_size)
+
+    # Parse the Segmentation URI
+    seg_name, seg_userid, seg_sessionid = parsers.parse_target(seg_uri)
+    seg_userid = seg_userid or 'octopi'
+    seg_sessionid = seg_sessionid or '1'
+
     # Initialize the Predictor
     gpu_count = torch.cuda.device_count()
     if gpu_count > 1:
         print(f"# of GPUs Available: {gpu_count} -- Using Multi-GPU Predictor.")
-        predict = segmentation.MultiGpuPredictor(
+        predict = inference.MultiGpuPredictor(
             config, model_config, model_weights,
             sw_bs=swbs, overlap=overlap, ntta=ntta
         )
     else:
         print(f"# of GPUs Available: {gpu_count} -- Using Single-GPU Predictor.")
-        predict = segmentation.Predictor(
+        predict = inference.Predictor(
             config, model_config, model_weights,
             sw_bs=swbs, overlap=overlap, ntta=ntta
         )
@@ -143,9 +153,9 @@ def segment(config, tomo_algorithm, voxel_size, model_weights, model_config,
         num_tomos_per_batch=batch_size,
         tomo_algorithm=tomo_algorithm,
         voxel_spacing=voxel_size,
-        name=seg_info[0],
-        userid=seg_info[1],
-        sessionid=seg_info[2]
+        name=seg_name,
+        userid=seg_userid,
+        sessionid=seg_sessionid
     )
 
 def localize(config, voxel_size, seg_info, pick_user_id, pick_session_id, n_procs = 16,
diff --git a/pyproject.toml b/pyproject.toml
index e9a027a..b554f45 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -40,6 +40,7 @@ dependencies = [
     "mrcfile",
     "kaleido",
     "requests",
+    "huggingface_hub",
     "submitit",
     "torch-ema",
     "matplotlib",

From 73f2e335f4f7123b797341f7929bc0c0e4b84b4d Mon Sep 17 00:00:00 2001
From: Jonathan Schwartz <32110921+jtschwar@users.noreply.github.com>
Date: Thu, 16 Jul 2026 11:03:56 -0700
Subject: [PATCH 3/6] Potential fix for pull request finding 'CodeQL / Binding
 a socket to all network interfaces'

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
---
 octopi/pytorch/pg_server.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/octopi/pytorch/pg_server.py b/octopi/pytorch/pg_server.py
index b2f41a2..5fdf3b9 100644
--- a/octopi/pytorch/pg_server.py
+++ b/octopi/pytorch/pg_server.py
@@ -31,7 +31,7 @@
 def _find_free_port() -> int:
     """Pick a currently-free TCP port. Small TOCTOU race, acceptable here."""
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("", 0))
+        s.bind(("127.0.0.1", 0))
         return s.getsockname()[1]
 
 

From 14dd93093aebec46d0ea6605f026e02855c16464 Mon Sep 17 00:00:00 2001
From: Jonathan Schwartz <jtschwar@gmail.com>
Date: Fri, 17 Jul 2026 12:56:29 -0700
Subject: [PATCH 4/6] fix: prioritize SLURM CPU allocation in auto worker
 count, align default SLURM CPU/mem requests

---
 octopi/datasets/helpers.py        | 71 +++++++++++++++++++++++++------
 octopi/entry_points/run_optuna.py |  6 ++-
 octopi/pytorch/submit_search.py   |  4 +-
 octopi/utils/submit_slurm.py      |  8 ++--
 4 files changed, 67 insertions(+), 22 deletions(-)

diff --git a/octopi/datasets/helpers.py b/octopi/datasets/helpers.py
index 7e95a75..a97b98e 100644
--- a/octopi/datasets/helpers.py
+++ b/octopi/datasets/helpers.py
@@ -20,24 +20,67 @@ def _spatial_shape(obj):
     except Exception:
         return None
 
-def auto_num_workers(cap: int = 16, min_workers: int = 1) -> int:
+def auto_num_workers(cap: int = 16, min_workers: int = 1, reserve: int = 1) -> int:
     """
-    Pick a DataLoader worker count that respects the CPUs actually available
+    Pick a DataLoader worker count that respects the CPUs actually *allocated*
     to this process.
 
-    On Linux this uses ``os.sched_getaffinity`` so SLURM ``--cpus-per-task``
-    bindings are honoured. Elsewhere it falls back to ``os.cpu_count``. The
-    result is clamped to ``cap`` since past ~16 workers gains are usually
-    noise for a 3D UNet with cached datasets, and each worker adds ~1-2 GB
-    of RAM overhead.
-
-    The cap only clamps from above: a 4-core laptop still gets 4 workers.
+    Priority order:
+
+    1. SLURM allocation env vars — ``SLURM_CPUS_PER_TASK`` (set when
+       ``--cpus-per-task`` is given), falling back to ``SLURM_CPUS_ON_NODE``
+       (set for essentially any job step, covering jobs that size CPUs
+       differently). These are checked first because on GPU partitions the
+       cgroup frequently does *not* restrict CPU affinity, so
+       ``os.sched_getaffinity`` reports every core on the node. Trusting that
+       would spawn dozens of 100%-CPU workers on a small allocation, starving
+       co-tenants and (via copy-on-write cache drift) thrashing node memory
+       into swap.
+    2. ``os.sched_getaffinity`` — honours cgroup CPU binding when SLURM isn't
+       the scheduler (or the vars are unset).
+    3. ``os.cpu_count`` — non-Linux fallback.
+
+    ``reserve`` cores are held back for the main process (and the validation
+    loader, which sizes itself from this value). The result is clamped to
+    ``cap`` as a sanity ceiling: training is loader-bound so throughput keeps
+    rising to ~16 workers, but past that gains flatten while RAM/context-switch
+    overhead grows, and the cap also guards the auto-detect path from reading a
+    whole 64-256 core node as the worker count. The cap does NOT encode a
+    per-node "fair share" (e.g. gpu-f = 14 cores/GPU) — that is a scheduling
+    policy invisible to this process and must be set via ``--cpus-per-task``.
+
+    The cap only clamps from above: a 4-core laptop still gets ~3 workers.
+
+    ``OCTOPI_MAX_WORKERS`` overrides the count entirely (bypasses cap/reserve)
+    for non-SLURM/desktop users on large workstations who want more workers
+    than the default cap. SLURM jobs should size via ``--cpus-per-task``
+    instead, so this is intended for interactive/local use.
     """
-    try:
-        n = len(os.sched_getaffinity(0))
-    except AttributeError:
-        n = 4 # default to small number on non-slurm platforms
-    return max(min_workers, min(cap, n))
+    override = os.environ.get("OCTOPI_MAX_WORKERS")
+    if override:
+        try:
+            return max(min_workers, int(override))
+        except ValueError:
+            pass  # malformed override -> fall through to auto-detection
+
+    # Prefer SLURM's view of the allocation. Both vars are plain integers;
+    # SLURM_JOB_CPUS_PER_NODE is intentionally skipped since it can be a packed
+    # form like "12(x2)".
+    n = None
+    for var in ("SLURM_CPUS_PER_TASK", "SLURM_CPUS_ON_NODE"):
+        val = os.environ.get(var)
+        if val:
+            try:
+                n = int(val)
+                break
+            except ValueError:
+                continue
+    if n is None:
+        try:
+            n = len(os.sched_getaffinity(0))
+        except AttributeError:
+            n = os.cpu_count() or 4  # non-Linux fallback
+    return max(min_workers, min(cap, n - reserve))
 
 def parse_resolution_uris(tomo_uris) -> list[tuple[str, float]]:
     """
diff --git a/octopi/entry_points/run_optuna.py b/octopi/entry_points/run_optuna.py
index 35f053e..f9ea386 100755
--- a/octopi/entry_points/run_optuna.py
+++ b/octopi/entry_points/run_optuna.py
@@ -33,8 +33,10 @@
               help="Submit trials via SLURM (submitit) instead of local GPUs")
 @click.option('--njobs', '-nj', type=int, default=5,
               help="Number of concurrent training jobs when using submitit")
-@click.option('--cpu-constraint', '-cc', type=str, default='16,8',
-              help='Number of CPUs and mem-per-cpu to requested. (e.g., "4,16" for 4 CPUs and 16GB per CPU)')
+@click.option('--cpu-constraint', '-cc', type=str, default='12,8',
+              help='Number of CPUs and mem-per-cpu to request. (e.g., "12,8" for 12 CPUs and 8GB per CPU). '
+                   'DataLoader workers auto-scale to the CPU count (cpus-1, capped at 16), so keep them '
+                   'matched. 12 fits the fair CPU/GPU share on every node incl. H100 (14 cores/GPU).')
 @click.option('--gpu-constraint', '-gc', type=str, default=None,
               help='GPU constraint to use for SLURM jobs (e.g., "a6000" or "l40,a6000")')
 @click.option('--timeout', type=int, default=4,
diff --git a/octopi/pytorch/submit_search.py b/octopi/pytorch/submit_search.py
index 59d66e7..b4300e2 100644
--- a/octopi/pytorch/submit_search.py
+++ b/octopi/pytorch/submit_search.py
@@ -332,8 +332,8 @@ class SubmititExplorer(ExploreSubmitter):
     def __init__(
         self,
         n_concurrent_jobs: int = 5,
-        cpus_per_task: int = 4,
-        mem_per_cpu: int = 16,
+        cpus_per_task: int = 12,
+        mem_per_cpu: int = 8,
         slurm_timeout_min: int = 1080,
         gpu_constraint: str = None,
         submitit_folder: str = "submitit_logs",
diff --git a/octopi/utils/submit_slurm.py b/octopi/utils/submit_slurm.py
index e9be42b..6b4af5a 100755
--- a/octopi/utils/submit_slurm.py
+++ b/octopi/utils/submit_slurm.py
@@ -16,8 +16,8 @@ def create_shellsubmit(
 
 {slurm_gpus}
 #SBATCH --time=18:00:00
-#SBATCH --cpus-per-task=4
-#SBATCH --mem-per-cpu=16G
+#SBATCH --cpus-per-task=12
+#SBATCH --mem-per-cpu=8G
 #SBATCH --job-name={job_name}
 #SBATCH --output={output_file}
 
@@ -43,8 +43,8 @@ def create_shellsubmit_array(
     shell_script_content = f"""#!/bin/bash
 
 #SBATCH --time=18:00:00
-#SBATCH --cpus-per-task=4
-#SBATCH --mem-per-cpu=16G
+#SBATCH --cpus-per-task=12
+#SBATCH --mem-per-cpu=8G
 #SBATCH --job-name={job_name}
 #SBATCH --output={output_file}
 #SBATCH --array={job_array[0]}-{job_array[1]}

From cda40567626b199adb1f52ea7b947c779c8c5c27 Mon Sep 17 00:00:00 2001
From: jtschwar <jtschwar@gmail.com>
Date: Fri, 17 Jul 2026 16:14:45 -0700
Subject: [PATCH 5/6] correct recomendation for cpu usage

---
 docs/user-guide/training.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/user-guide/training.md b/docs/user-guide/training.md
index b4f0cdb..a2bbfa5 100644
--- a/docs/user-guide/training.md
+++ b/docs/user-guide/training.md
@@ -184,7 +184,7 @@ Octopi supports two complementary workflows:
             |----------|-------------|---------|------|
             | `--submitit` | Submit trials as independent SLURM jobs using submitit instead of running locally. | `False` | Enables HPC / multi-node execution |
             | `--njobs` | Maximum number of concurrent SLURM jobs (trials) to run at once. | `5` | Each job runs exactly one Optuna trial |
-            | `--compute-constraint` | CPU and memory request per SLURM job in the form `cpus,mem_gb`. | `4,16` | Example: `8,32` requests 8 CPUs and 32 GB RAM |
+            | `--compute-constraint` | CPU and memory request per SLURM job in the form `cpus,mem_gb`. | `12,8` | Example: `12,8` requests 12 CPUs with 8 GB RAM per core. |
             | `--timeout` | Walltime limit (hours) per SLURM job. | `4` | Jobs exceeding this limit are terminated by the scheduler |
 
             !!! example "What changes when `--submitit` is enabled?"

From 396b56bdc30eb2bf232b3441918b57d2ac21e3d5 Mon Sep 17 00:00:00 2001
From: jtschwar <jtschwar@gmail.com>
Date: Mon, 20 Jul 2026 12:04:01 -0700
Subject: [PATCH 6/6] update links to biohub org

---
 README.md                            | 4 ++--
 docs/getting-started/installation.md | 2 +-
 docs/index.md                        | 2 +-
 docs/user-guide/labels.md            | 2 +-
 pyproject.toml                       | 6 +++---
 zensical.toml                        | 6 +++---
 6 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index d267c8c..99b67fb 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,6 @@
 # OCTOPI 🐙🐙🐙
 
-[![License](https://img.shields.io/pypi/l/octopi.svg?color=green)](https://github.com/chanzuckerberg/octopi/raw/main/LICENSE)
+[![License](https://img.shields.io/pypi/l/octopi.svg?color=green)](https://github.com/Biohub/octopi/raw/main/LICENSE)
 [![PyPI](https://img.shields.io/pypi/v/octopi.svg?color=green)](https://pypi.org/project/octopi)
 [![Python Version](https://img.shields.io/pypi/pyversions/octopi.svg?color=green)](https://www.python.org/)
 
@@ -47,7 +47,7 @@ octopi
 
 ## 📚 Documentation
 
-For detailed documentation, tutorials, CLI and API reference, visit our [documentation](https://chanzuckerberg.github.io/octopi/).
+For detailed documentation, tutorials, CLI and API reference, visit our [documentation](https://biohub.github.io/octopi/).
 
 ## 🤝 Contributing
 
diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md
index e92400d..0299689 100644
--- a/docs/getting-started/installation.md
+++ b/docs/getting-started/installation.md
@@ -14,7 +14,7 @@ This will install the latest stable release along with all required dependencies
 If you want to contribute to octopi or need the latest development version, you can install from source:
 
 ```bash
-git clone https://github.com/chanzuckerberg/octopi.git
+git clone https://github.com/Biohub/octopi.git
 cd octopi
 pip install -e .
 ```
diff --git a/docs/index.md b/docs/index.md
index 11b514b..a58e5cc 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -96,4 +96,4 @@
 
 ## Getting Help
 
-Open an issue on our [GitHub repository](https://github.com/chanzuckerberg/octopi).
+Open an issue on our [GitHub repository](https://github.com/Biohub/octopi).
diff --git a/docs/user-guide/labels.md b/docs/user-guide/labels.md
index ed76b50..8936ad6 100644
--- a/docs/user-guide/labels.md
+++ b/docs/user-guide/labels.md
@@ -66,7 +66,7 @@ octopi create-targets \
 
 ## Check Target Quality
 
-To validate your training targets, refer to our interactive notebook: [Inspect Segmentation Targets](https://github.com/chanzuckerberg/octopi/blob/main/notebooks/inspect_segmentation_targets.ipynb)
+To validate your training targets, refer to our interactive notebook: [Inspect Segmentation Targets](https://github.com/Biohub/octopi/blob/main/notebooks/inspect_segmentation_targets.ipynb)
 
 This notebook shows how to load segmentation targets and overlay targets on tomograms.
 
diff --git a/pyproject.toml b/pyproject.toml
index b554f45..73dfa7a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -80,6 +80,6 @@ dependency-overrides = [
 ]
 
 [project.urls]
-Homepage = "https://github.com/chanzuckerberg/octopi"
-Documentation = "https://chanzuckerberg.github.io/octopi/"
-Issues = "https://github.com/chanzuckerberg/octopi/issues"
\ No newline at end of file
+Homepage = "https://github.com/Biohub/octopi"
+Documentation = "https://biohub.github.io/octopi/"
+Issues = "https://github.com/Biohub/octopi/issues"
\ No newline at end of file
diff --git a/zensical.toml b/zensical.toml
index b9e7aa3..772dfc8 100644
--- a/zensical.toml
+++ b/zensical.toml
@@ -1,7 +1,7 @@
 [project]
 site_name = "octopi 🐙🐙🐙"
-repo_url = "https://github.com/chanzuckerberg/octopi"
-repo_name = "chanzuckerberg/octopi"
+repo_url = "https://github.com/Biohub/octopi"
+repo_name = "Biohub/octopi"
 copyright = "2025, Jonathan Schwartz"
 extra_javascript = [
     "javascripts/mathjax.js",
@@ -79,7 +79,7 @@ toggle.name = "Switch to light mode"
 
 [[project.extra.social]]
 icon = "fontawesome/brands/github"
-link = "https://github.com/chanzuckerberg/octopi"
+link = "https://github.com/Biohub/octopi"
 
 # --- Markdown Extensions (must list all defaults when overriding any) ---