From bfaee072585a8a3be19f32b4c616544459be86b3 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 17 Oct 2026 22:35:32 +0000 Subject: [PATCH] LoRA training: Download LoRA on a finished run in Results Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01MpnF7GwNGMbngU11wjym7P --- app/ui/features/training/results.js | 13 +++++++- app/ui/features/training/training.test.js | 23 ++++++++++++++ app/ui/shell/api.js | 1 + modules/training/lora/__init__.py | 18 +++++++++++ .../lora/tests/test_plan_and_routes.py | 30 +++++++++++++++++++ 5 files changed, 84 insertions(+), 1 deletion(-) diff --git a/app/ui/features/training/results.js b/app/ui/features/training/results.js index 6b0e1301..76726fdc 100644 --- a/app/ui/features/training/results.js +++ b/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 }); } + } }); +} diff --git a/app/ui/features/training/training.test.js b/app/ui/features/training/training.test.js index 0f16cce9..db4e945a 100644 --- a/app/ui/features/training/training.test.js +++ b/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; } +}); diff --git a/app/ui/shell/api.js b/app/ui/shell/api.js index db585a44..a063dc28 100644 --- a/app/ui/shell/api.js +++ b/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) { diff --git a/modules/training/lora/__init__.py b/modules/training/lora/__init__.py index 23e01c92..68197ef8 100644 --- a/modules/training/lora/__init__.py +++ b/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"] diff --git a/modules/training/lora/tests/test_plan_and_routes.py b/modules/training/lora/tests/test_plan_and_routes.py index a55317af..8e6f8fce 100644 --- a/modules/training/lora/tests/test_plan_and_routes.py +++ b/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]