feat: ability to unload captio/review model after use

Author: Specthram

server/runners/generate.py

@@ -108,6 +108,10 @@ def run(progress: Progress) -> dict:
         findings = 0
         if params.review_after:
             findings = _review_all(progress, params, keys)
+        if params.unload_after and loader.is_model_loaded():
+            progress(sub="freeing the model…")
+            for status, _done in loader.unload_model():
+                progress(sub=status)
         return {"done": total, "grounded": grounded, "findings": findings}
 
     return run

server/runners/review_run.py

@@ -129,6 +129,18 @@ def _judge_rule(rule, target, seed) -> dict | None:
     return caption_judge.judged_finding(target["caption"], verdict)
 
 
+def _unload_resident(progress: Progress) -> None:
+    """Free whatever model is still resident (the unload-after toggle)."""
+    # pylint: disable=import-outside-toplevel
+    from src import loader
+
+    if not loader.is_model_loaded():
+        return
+    progress(sub="freeing the model…")
+    for status, _done in loader.unload_model():
+        progress(sub=status)
+
+
 def review_run_body(params):
     """Return a job body running a rule-based review over a dataset.
 
@@ -193,6 +205,8 @@ def run(progress: Progress) -> dict:
             )
 
         storage.close_review_run(run_id, state["found"])
+        if params.unload_after:
+            _unload_resident(progress)
         return {"reviewed": len(targets), "findings": state["found"]}
 
     return run

server/schemas.py

@@ -132,6 +132,8 @@ class GenerateBody(BaseModel):
     # Off = caption only media whose caption is still empty; on (default)
     # regenerates every targeted media, already-captioned ones included.
     recaption: bool = True
+    # Free the VRAM when the job is done (default: keep the model resident).
+    unload_after: bool = False
 
 
 class DeployBody(BaseModel):
@@ -199,6 +201,8 @@ class ReviewRunBody(BaseModel):
     scope: str = "all"  # all | selection | flagged | single
     rule_ids: list[int] | None = None  # None = every enabled rule
     seed: int | None = None
+    # Free the VRAM when the run is done (default: keep the model resident).
+    unload_after: bool = False
 
 
 class ReviewDecideBody(BaseModel):

tests/test_api_caption_flow.py

@@ -448,3 +448,41 @@ def test_generate_skips_filled_captions_when_recaption_off(
     assert result["done"] == 1
     assert storage.read_caption(dataset_id, filled, "txt") == "already there"
     assert storage.read_caption(dataset_id, empty, "txt") == "fresh"
+
+
+def test_generate_unloads_model_when_asked(
+    store_db, thumb_cache_dir, tmp_path, monkeypatch
+):
+    """With ``unload_after``, the job frees the resident model at the end."""
+    # pylint: disable=unused-argument
+    from server.runners import generate as generate_runner
+    from server.schemas import GenerateBody
+    from src import loader
+
+    Image.new("RGB", (64, 64), (219, 68, 55)).save(tmp_path / "a.png")
+    library_id = store.create_library("fixtures", str(tmp_path))
+    store.scan_library(library_id)
+    dataset_id = store.create_dataset("shapes")
+    for row in store.list_library_media():
+        store.add_media_to_dataset(dataset_id, row["id"])
+
+    unloaded = []
+    monkeypatch.setattr(
+        generate_runner,
+        "_caption_one",
+        lambda path, params, seed, profile: "text",
+    )
+    monkeypatch.setattr(loader, "is_model_loaded", lambda: True)
+    monkeypatch.setattr(
+        loader,
+        "unload_model",
+        lambda: unloaded.append(1) or iter([("freed", True)]),
+    )
+    params = GenerateBody(
+        dataset_id=dataset_id,
+        caption_type="txt",
+        prompt="p",
+        unload_after=True,
+    )
+    generate_runner.generate_body(params)(lambda **kwargs: None)
+    assert unloaded == [1]

web/src/api/hooks.ts

@@ -764,6 +764,8 @@ export interface GenerateVars {
   ground_after: boolean;
   /** Off = only caption media whose caption is still empty. */
   recaption: boolean;
+  /** On = free the VRAM when the job ends (default: model stays loaded). */
+  unload_after?: boolean;
 }
 
 export function useGenerate() {
@@ -890,6 +892,7 @@ export function useRunReview() {
       scope: string;
       rule_ids?: number[] | null;
       seed?: number | null;
+      unload_after?: boolean;
     }) => api.post<{ job_id: string }>("/review/run", vars),
   });
 }

web/src/components/organisms/BatchBar.tsx

@@ -51,6 +51,7 @@ export function BatchBar() {
         : null,
       ground_after: gen.groundAfter,
       recaption: gen.recaption,
+      unload_after: gen.unloadAfter,
     });
 
   const addTagToAll = async () => {

web/src/components/organisms/CaptionDetailPanel.tsx

@@ -532,6 +532,7 @@ function RegenerateButton({
         ground_after: gen.groundAfter,
         // An explicit per-media regenerate always rewrites, filled or not.
         recaption: true,
+        unload_after: gen.unloadAfter,
       },
       { onSuccess: (data) => job.start(data.job_id) },
     );

web/src/components/organisms/CaptionLeftPanel.tsx

@@ -140,6 +140,7 @@ export function CaptionLeftPanel() {
         review_judge_profile_id: gen.reviewAfter ? data.judge_id : null,
         ground_after: gen.groundAfter,
         recaption: gen.recaption,
+        unload_after: gen.unloadAfter,
       },
       {
         onSuccess: (result) => {
@@ -334,6 +335,16 @@ export function CaptionLeftPanel() {
             Only media whose caption is still empty will be captioned.
           </div>
         )}
+        <label style={checkboxRow}>
+          <input
+            type="checkbox"
+            checked={gen.unloadAfter}
+            onChange={(event) =>
+              gen.set({ unloadAfter: event.target.checked })
+            }
+          />
+          Unload the model after the job
+        </label>
         <label style={checkboxRow}>
           <input
             type="checkbox"

web/src/components/organisms/ReviewView.tsx

@@ -49,6 +49,7 @@ export function ReviewView() {
   const runReview = useRunReview();
   const profiles = useProfiles();
   const [scope, setScope] = useState("all");
+  const [unloadAfter, setUnloadAfter] = useState(false);
   const [jobId, setJobId] = useState<string | null>(null);
   const job = useJobsStore((state) =>
     jobId ? state.jobs[jobId] : undefined,
@@ -105,6 +106,7 @@ export function ReviewView() {
         media_ids: mediaIds,
         judge_profile_id: profiles.data?.judge_id ?? null,
         scope,
+        unload_after: unloadAfter,
       },
       { onSuccess: (data) => setJobId(data.job_id) },
     );
@@ -125,6 +127,8 @@ export function ReviewView() {
         running={running}
         job={job}
         onRun={run}
+        unloadAfter={unloadAfter}
+        onUnloadAfter={setUnloadAfter}
       />
       <Queue
         datasetId={datasetId}
@@ -156,6 +160,8 @@ function Rail({
   running,
   job,
   onRun,
+  unloadAfter,
+  onUnloadAfter,
 }: {
   datasetId: number;
   rules: ReviewRule[];
@@ -165,6 +171,8 @@ function Rail({
   running: boolean;
   job: JobLike | undefined;
   onRun: () => void;
+  unloadAfter: boolean;
+  onUnloadAfter: (value: boolean) => void;
 }) {
   const profiles = useProfiles();
   const createRule = useCreateReviewRule();
@@ -299,6 +307,14 @@ function Rail({
             ▶ Run review
           </Button>
         )}
+        <label style={checkRow}>
+          <input
+            type="checkbox"
+            checked={unloadAfter}
+            onChange={(event) => onUnloadAfter(event.target.checked)}
+          />
+          Unload the model after the run
+        </label>
         <p style={hintStyle}>
           Text-only rules run without loading images; vision rules load each
           one. Locked captions are skipped. Judge: {judgeLabel}.

web/src/store/captionStore.ts

@@ -19,6 +19,8 @@ interface CaptionState {
   reviewAfter: boolean;
   /** Off = only caption media whose caption is still empty. */
   recaption: boolean;
+  /** On = free the VRAM when the job ends (off keeps the model loaded). */
+  unloadAfter: boolean;
   /** Chain a SigLIP grounding pass on every freshly written caption. */
   groundAfter: boolean;
   locked: Set<string>;
@@ -33,6 +35,7 @@ export const useCaptionStore = create<CaptionState>((set) => ({
   seed: "-1",
   reviewAfter: false,
   recaption: true,
+  unloadAfter: false,
   groundAfter: false,
   locked: new Set(),