LoRA training: Download LoRA on a finished run in Results

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P

app/ui/features/training/results.js

@@ -29,7 +29,18 @@ export function resultsStep({ app, ui }) {
       rec && rec.error ? c.inlineError.default({ text: rec.error }) : null,
       rec && rec.resumable && ["error", "stopped"].includes(rec.state) ? c.banner.info({ text: `The state at step ${rec.steps || 0} or earlier is kept. Restore and continue picks the training up from it with the Train step's current settings.` }) : null,
       chart, c.hint.default({ text: settings }),
-      c.toolbar.default({ items: rec && rec.resumable && !ui.run.active && !["training", "queued", "previewing", "paused"].includes(ui.run.state) ? [c.button.sm({ label: "Restore and continue", tone: "primary", title: "Carry on from the saved state", onClick: async () => { await ui.restore(chosen.run); } })] : [], trailing: [c.button.sm({ label: "Delete this record", tone: "ghost", title: "Removes the record only, not the LoRA file", onClick: async () => { await app.api.trainForgetRun(chosen.run); ui.chosenRun = ""; await ui.loadRuns(); ui.redraw(); } })] }),
+      c.toolbar.default({ items: [...(rec && rec.lora ? [download(app, rec)] : []), ...(rec && rec.resumable && !ui.run.active && !["training", "queued", "previewing", "paused"].includes(ui.run.state) ? [c.button.sm({ label: "Restore and continue", tone: "primary", title: "Carry on from the saved state", onClick: async () => { await ui.restore(chosen.run); } })] : [])], trailing: [c.button.sm({ label: "Delete this record", tone: "ghost", title: "Removes the record only, not the LoRA file", onClick: async () => { await app.api.trainForgetRun(chosen.run); ui.chosenRun = ""; await ui.loadRuns(); ui.redraw(); } })] }),
     ].filter(Boolean) })] }),
   ] });
 }
+
+/** Saves the finished LoRA through the browser, to wherever the person keeps files: an instance's own disk is gone when it is. */
+function download(app, rec) {
+  return c.button.sm({ label: "Download LoRA", tone: "primary", title: `Save ${rec.lora} to this computer`, onClick: async () => {
+    try {
+      const res = await fetch(app.api.trainLoraUrl(rec.run));
+      if (!res.ok) throw new Error((await res.json().catch(() => ({}))).why || `The server refused (${res.status}).`);
+      app.dom.saveBlob(await res.blob(), rec.lora);
+    } catch (err) { c.toast.warn({ text: err.message }); }
+  } });
+}

app/ui/features/training/training.test.js

@@ -410,3 +410,26 @@ test("Train step: the preset picked stays picked when a setting is changed after
   assert.ok(!("preset" in plan[2]));
   off();
 });
+
+test("Results: a finished run has a Download LoRA button that saves the file under its own name", async () => {
+  const f = fake();
+  const saved = [];
+  f.app.dom.saveBlob = (blob, name) => saved.push(name);
+  f.api.trainLoraUrl = (id) => `/lora/${id}`;
+  f.api.trainRuns = async () => [{ run: "r5", name: "cats", set: "cats", state: "done", steps: 10, total: 10, lora: "cats_lora.safetensors" }];
+  f.api.trainRunRecord = async () => ({ run: "r5", state: "done", steps: 10, total: 10, set: "cats", loss: [0.2], lora: "cats_lora.safetensors", settings: {} });
+  const realFetch = globalThis.fetch;
+  globalThis.fetch = async (url) => (String(url) === "/lora/r5" ? { ok: true, blob: async () => new Blob(["w"]) } : realFetch(url));
+  try {
+    const { button, off } = await open(f);
+    button.click();
+    await wait(60);
+    tab("Results").click();
+    await wait(80);
+    assert.ok(click("Download LoRA"));
+    click("Download LoRA").click();
+    await wait(60);
+    assert.deepEqual(saved, ["cats_lora.safetensors"]);
+    off();
+  } finally { globalThis.fetch = realFetch; }
+});

app/ui/shell/api.js

@@ -101,6 +101,7 @@ export const api = {
   trainStop: () => call("POST", `${TRAIN}/run/stop`, {}),
   trainRuns: () => call("GET", `${TRAIN}/runs`).then((d) => d.runs || []),
   trainRunRecord: (id) => call("GET", `${TRAIN}/runs/${encodeURIComponent(id)}`),
+  trainLoraUrl: (id) => `${BASE}${TRAIN}/runs/${encodeURIComponent(id)}/lora`,
   trainForgetRun: (id) => call("DELETE", `${TRAIN}/runs/${encodeURIComponent(id)}`),
   /** Any graph as a ComfyUI job. -> the prompt id. */
   async queueGraph(graph) {

modules/training/lora/__init__.py

@@ -233,6 +233,24 @@ def routes(table, base, web):
         rec = jobs.record(req.match_info["run"])
         return web.json_response({**rec, "resumable": loop.has_checkpoint(rec.get("run"))}) if rec else _bad(web, "There is no such run.", 404)
 
+    @table.get(base + "/runs/{run}/lora")
+    async def _download(req):
+        from pathlib import Path
+        from urllib.parse import quote
+        rec = jobs.record(req.match_info["run"])
+        name = (rec or {}).get("lora") or ""
+        if not name or Path(name).name != name or not name.endswith(".safetensors"):
+            return _bad(web, "This run has no saved LoRA file.", 404)
+        try:
+            import folder_paths
+            folders = [Path(f) for f in folder_paths.get_folder_paths("loras")]
+        except Exception:                                  # noqa: BLE001 -- no ComfyUI folders here
+            folders = []
+        path = next((f / name for f in folders if (f / name).is_file()), None)
+        if path is None:
+            return _bad(web, f"{name} is no longer in your loras folder.", 404)
+        return web.FileResponse(path, headers={"Content-Type": "application/octet-stream", "Content-Disposition": f"attachment; filename*=UTF-8''{quote(name)}"})
+
     @table.delete(base + "/runs/{run}")
     async def _forget(req):
         run = req.match_info["run"]

modules/training/lora/tests/test_plan_and_routes.py

@@ -126,3 +126,33 @@ def test_rated_results_list_only_the_videos_that_are_still_here(tmp_path, monkey
             return await (await c.get("/m/rated")).json()
     got = asyncio.run(go())
     assert [r["media"]["filename"] for r in got["results"]] == ["here.mp4"] and got["counts"]["liked"] == 1
+
+
+def test_a_finished_run_can_be_downloaded_and_only_its_own_file(tmp_path, monkeypatch):
+    import folder_paths
+    from aiohttp import web
+    from aiohttp.test_utils import TestClient, TestServer
+    from core import config
+    from modules.training import lora
+    from modules.training.lora import jobs
+    monkeypatch.setattr(config, "TRAINING_DIR", tmp_path)
+    loras = tmp_path / "loras"
+    loras.mkdir()
+    (loras / "cats.safetensors").write_bytes(b"weights")
+    (tmp_path / "secret.safetensors").write_bytes(b"no")
+    monkeypatch.setattr(folder_paths, "get_folder_paths", lambda kind: [str(loras)])
+    runs = jobs._dir()
+    import json
+    for run, name in (("aaaaaaaaaaaa", "cats.safetensors"), ("bbbbbbbbbbbb", "../secret.safetensors"), ("cccccccccccc", "gone.safetensors"), ("dddddddddddd", "")):
+        (runs / f"{run}.json").write_text(json.dumps({"run": run, "lora": name}))
+
+    async def go():
+        app, table = web.Application(), web.RouteTableDef()
+        lora.routes(table, "/m", web)
+        app.add_routes(table)
+        async with TestClient(TestServer(app)) as c:
+            ok = await c.get("/m/runs/aaaaaaaaaaaa/lora")
+            return ok.status, ok.headers.get("Content-Disposition"), await ok.read(), [(await c.get(f"/m/runs/{r}/lora")).status for r in ("bbbbbbbbbbbb", "cccccccccccc", "dddddddddddd", "eeeeeeeeeeee")]
+    status, disposition, body, refused = asyncio.run(go())
+    assert status == 200 and body == b"weights" and "attachment" in disposition and "cats.safetensors" in disposition
+    assert refused == [404, 404, 404, 404]