Merge pull request #2206 from atomantic/next/issue-2199

feat: verify safetensors integrity after Civitai/HF LoRA download

Author: atomantic

.changelog/NEXT.md

@@ -59,6 +59,7 @@
 - **Installing a video LoRA from HuggingFace now shows live download progress.** The LoRA Manager's "Install video LoRA from HuggingFace" form (and the curated video quick-install buttons) used to freeze on a static "Downloading…" label with no indication of how far along a multi-GB LoRA was — unlike image-model downloads, which show a percentage. The install now streams byte-level progress over SSE: the button shows `Downloading 42%` and a progress bar with `received / total` bytes fills as it downloads (an indeterminate pulse when the server can't read a Content-Length). Errors — including the "couldn't detect the model family" case that offers the inline "Install as LTX-Video" retry — still surface exactly as before.
 - **The LoRA Manager's Installed list is now grouped by Video / Image.** The curated "Video LoRAs" suggestion cards sat directly above a single flat "Installed" grid, so installed *image* LoRAs looked like they belonged to the video section. The Installed section is now visually divided from the suggestions and, in the "All" view, splits into labeled **Video LoRAs** / **Image LoRAs** subsections (mirroring the suggestion panel) so it's obvious which renders each LoRA applies to. The Image/Video filter still collapses it to a single scoped list.
 - **[issue-2036] Sketch & Annotation Canvas (phase 2): re-render an image guided by your annotations.** The annotate page (`/media/annotate/:mediaKey`) gained a "Re-render" action: draw over a generated image, then feed your markup back through local img2img so the marks reshape the render. A confirmation dialog names the exact local model it will run (no surprise AI calls) and lets you add an optional prompt and tune how much to change before it starts; the new render is queued and appears in Media History. Requires a local FLUX img2img runner — the action explains when one isn't available. Phase 3 (blank-canvas storyboard) remains open on #2036.
+- **[issue-2199] LoRA downloads are now verified for corruption before install.** A truncated or corrupt `.safetensors` download from Civitai or HuggingFace used to install silently and then render garbage ("mosaic") output at generate time. PortOS now checks every freshly downloaded LoRA's structure (safetensors header + size) and, when Civitai provides a SHA-256 for the file, compares it against the bytes on disk. On any mismatch the partial file is deleted and the install fails with a clear error so you re-download instead of training or rendering against corrupt weights.
 - **[issue-2196] Video Gen now takes an exact custom Width×Height.** The `/media/video` page only offered a 9-item resolution preset dropdown, so a size the presets didn't cover was only reachable via URL params, remix/continue, or aspect-snap on upload. Two numeric Width and Height inputs now sit beside the preset select — type any dimensions and the render uses them (the server still rounds each down to a multiple of 64 and enforces the per-tier pixel budget). Uploading a source image still auto-snaps the size to the image's aspect ratio by default, until you take the dimensions into your own hands.
 - **[issue-2197] First-run model downloads now show GB progress instead of raw text.** The video and image generation handlers now parse byte sizes from `DOWNLOAD:` markers and tqdm progress bars, displaying a formatted "Downloading model · first run · X.X GB / Y.Y GB" message instead of the raw protocol output. Progress frames include `downloadedBytes`/`totalBytes` numeric fields so the UI can build download progress bars. Phase tracking (`isDownloading`, `currentPhase`) ensures the formatting only applies during actual download phases.
 - **[issue-2198] A wedged local video render that produces no output is now killed automatically instead of pinning the GPU queue.** The existing completion watchdog only guards a render that finishes and then hangs in teardown — it never armed for a render that stalls *before emitting anything* (a known MLX/Metal failure mode where the job never exits and never prints), which held the serialized GPU lane until you manually canceled. A new idle-stall deadline arms at spawn and resets on every child output line (stdout or stderr); if a render produces no output for a generous window (default 10 minutes, tunable via `VIDEOGEN_IDLE_STALL_MS`) it's SIGKILLed and surfaced as a failed job, so the next queued render can start. The window is deliberately generous so a legitimately slow first-token render (weights load + Metal kernel compile) is never killed, and a manual cancel still works and always wins. (`server/services/videoGen/local.js`)

server/lib/README.md

@@ -103,7 +103,7 @@ The barrel `server/lib/index.js` is a machine-checkable enumeration of every pub
 | `ansiStrip.js` | Streaming ANSI / control-byte stripper. |
 | `hfToken.js` | HuggingFace token resolution (settings > env > CLI). |
 | `hfErrors.js` | Parse huggingface_hub gated-access errors: `extractGatedRepo(text)` → `owner/name` (or null) for the UI's license deep-link. Shared by the image runner and LoRA trainer. Pure. |
-| `hfCache.js` | HuggingFace Hub cache inspection (`inspectModelCache(repoId)` → `{cached,sizeBytes,snapshotPath}`, `isModelCached`, `getHfCacheRoot`). Drives the inline "Available / Download" badge on the image + video gen forms. Also `verifyModelCache(repoId,{deep})` (structural safetensors-header + optional sha256 integrity check) and `repairModelCache(repoId,{deep})` (delete corrupt weight files so the download path re-fetches them) — power the "Repair model" banner. |
+| `hfCache.js` | HuggingFace Hub cache inspection (`inspectModelCache(repoId)` → `{cached,sizeBytes,snapshotPath}`, `isModelCached`, `getHfCacheRoot`). Drives the inline "Available / Download" badge on the image + video gen forms. Also `verifyModelCache(repoId,{deep})` (structural safetensors-header + optional sha256 integrity check) and `repairModelCache(repoId,{deep})` (delete corrupt weight files so the download path re-fetches them) — power the "Repair model" banner. `verifySafetensorsStructure(path,size)` is the reusable header/size structural check (reads only the header region) — also used by the Civitai/HF LoRA install path to reject truncated downloads. |
 | `hfDownload.js` | `downloadHfRepo({repo,onEvent})` returning `{promise,kill}` — spawns `scripts/hf_download_repo.py` in the FLUX.2 venv (fallback: mflux pythonPath) and emits SSE-friendly stage/progress/complete events. Powers the inline "Download" button next to the model picker. |
 | `sseHeaders.js` | `SSE_HEADERS` — the canonical SSE response headers (incl. `X-Accel-Buffering:no`) in a dependency-free module so any producer (`sseDownload.js`, `sseUtils.js`) can share them without pulling in a heavier module's transitive imports. |
 | `sseDownload.js` | `startHfDownloadStream({req,res,repo,alreadyDownloadedMessage})`, `openSseStream(res)` (`{send,safeEnd}` SSE boilerplate; uses `SSE_HEADERS` from `sseHeaders.js`) — shared SSE driver used by both image and video gen `/models/:id/download` routes. Owns the cross-route in-flight Map so a double-click (or both pages running) can't spawn two python children against the same repo. |

server/lib/hfCache.js

@@ -167,7 +167,7 @@ const SHA256_RE = /^[0-9a-f]{64}$/;
 
 // Cheap structural check for a single .safetensors file. Reads only the header
 // region (a few KB), never the tensor payload. Returns { ok, reason, ... }.
-async function verifySafetensorsStructure(path, size) {
+export async function verifySafetensorsStructure(path, size) {
   if (size < 8) return { ok: false, reason: 'truncated-header' };
   let fd;
   try {

server/services/loras.js

@@ -25,7 +25,8 @@ import { Readable, Transform } from 'stream';
 import { pipeline } from 'stream/promises';
 import { basename, join } from 'path';
 import { ServerError } from '../lib/errorHandler.js';
-import { atomicWrite, assertSafeFilename, ensureDir, listDirectoryByExtension, PATHS } from '../lib/fileUtils.js';
+import { atomicWrite, assertSafeFilename, ensureDir, listDirectoryByExtension, sha256File, PATHS } from '../lib/fileUtils.js';
+import { verifySafetensorsStructure } from '../lib/hfCache.js';
 import { isPlainObject } from '../lib/objects.js';
 import {
   applyDownloadToken,
@@ -367,6 +368,43 @@ const downloadToFile = async (url, destPath, { fetchImpl = fetch, headers = {} ,
   });
 };
 
+// After a LoRA finishes downloading, verify the on-disk `.safetensors` before
+// we commit to it by writing the sidecar. A truncated (short download) or
+// right-size-but-wrong-bytes file otherwise installs silently and renders
+// garbage ("mosaic" output) at generate time — the leading non-Metal cause of
+// corrupt output (issue #2199; mirrors the HF-cache integrity check in
+// hfCache.js / #1324). Always runs the cheap structural header/size check
+// (reads only the header region, never the multi-GB payload); adds a deep
+// sha256 compare only when the source metadata carried a digest (Civitai
+// `file.hashes.SHA256`). On any failure the file is deleted so the next install
+// re-downloads instead of training/rendering against corrupt weights.
+const verifyDownloadedLora = async (destPath, { expectedSha256 = null, source = 'civitai' } = {}) => {
+  const label = source === 'huggingface' ? 'HuggingFace' : 'Civitai';
+  const code = source === 'huggingface' ? 'HF_LORA_CORRUPT' : 'CIVITAI_LORA_CORRUPT';
+  const st = await stat(destPath).catch(() => null);
+  const structural = await verifySafetensorsStructure(destPath, st?.size ?? 0);
+  if (!structural.ok) {
+    await rm(destPath, { force: true }).catch(() => {});
+    throw new ServerError(
+      `${label} LoRA download is corrupt (${structural.reason}) — the partial file was deleted. Retry the install.`,
+      { status: 502, code },
+    );
+  }
+  // Civitai hashes are uppercase hex; sha256File returns lowercase — compare
+  // case-insensitively and only when the digest is a well-formed sha256.
+  const want = typeof expectedSha256 === 'string' ? expectedSha256.trim().toLowerCase() : '';
+  if (/^[0-9a-f]{64}$/.test(want)) {
+    const actual = await sha256File(destPath).catch(() => null);
+    if (actual && actual.toLowerCase() !== want) {
+      await rm(destPath, { force: true }).catch(() => {});
+      throw new ServerError(
+        `${label} LoRA failed SHA-256 verification (expected ${want.slice(0, 12)}…, got ${actual.slice(0, 12)}…) — the file was deleted. Retry the install.`,
+        { status: 502, code },
+      );
+    }
+  }
+};
+
 // Install a LoRA from a Civitai URL. Returns the new sidecar JSON so the
 // client can render it immediately without a second list round-trip.
 export const installFromCivitai = async (input, { fetchImpl = fetch } = {}) => {
@@ -435,6 +473,7 @@ export const installFromCivitai = async (input, { fetchImpl = fetch } = {}) => {
     headers: { 'User-Agent': 'PortOS/civitai-installer' },
     hasApiKey: !!apiKey,
   });
+  await verifyDownloadedLora(destPath, { expectedSha256: file?.hashes?.SHA256 || null, source: 'civitai' });
 
   const sidecar = buildSidecar({ model, version, file, filename });
   await atomicWrite(sidecarPath(filename), JSON.stringify(sidecar, null, 2) + '\n');
@@ -506,6 +545,10 @@ export const installFromHuggingface = async (input, { fetchImpl = fetch, onProgr
     onProgress,
     signal,
   });
+  // HF's model metadata doesn't expose a per-file digest through pickHfLoraFile,
+  // so the structural header/size check is the integrity guard here (no deep
+  // sha256 compare available for this path).
+  await verifyDownloadedLora(destPath, { source: 'huggingface' });
 
   const sidecar = buildHfLoraSidecar({ repo, revision, file, model, family, filename });
   await atomicWrite(sidecarPath(filename), JSON.stringify(sidecar, null, 2) + '\n');