Harden image enhancement runtime and test coverage

.github/workflows/release.yml

@@ -142,15 +142,50 @@ jobs:
               exit 1
             fi
           done
-          sha256sum -- *.exe *.dmg *.AppImage > SHA256SUMS.txt
-          sha256sum --check SHA256SUMS.txt
+          python <<'PY'
+          from hashlib import sha256
+          from pathlib import Path
+
+          files = []
+          for extension in ("exe", "dmg", "AppImage"):
+              files.extend(sorted(Path(".").glob(f"*.{extension}")))
+
+          with Path("SHA256SUMS.txt").open("w", encoding="utf-8") as manifest:
+              for path in files:
+                  manifest.write(f"{path.name}\t{sha256(path.read_bytes()).hexdigest()}\n")
+
+          entries = {}
+          with Path("SHA256SUMS.txt").open("r", encoding="utf-8") as manifest:
+              for line in manifest:
+                  name, digest = line.rstrip("\n").split("\t", 1)
+                  entries[name] = digest
+
+          for path in files:
+              digest = sha256(path.read_bytes()).hexdigest()
+              if entries.get(path.name) != digest:
+                  raise SystemExit(f"Checksum mismatch for {path.name}")
+          PY
 
       - name: Publish complete prerelease
         working-directory: artifacts
         shell: bash
         run: |
-          state=$(gh api --paginate "repos/$GH_REPO/releases?per_page=100" \
-            --jq '.[] | select(.tag_name == env.RELEASE_TAG) | .draft')
+          owner=${GH_REPO%/*}
+          repo=${GH_REPO#*/}
+          state=$(gh api graphql \
+            -F owner="$owner" \
+            -F name="$repo" \
+            -F tag="$RELEASE_TAG" \
+            -f query='
+              query($owner: String!, $name: String!, $tag: String!) {
+                repository(owner: $owner, name: $name) {
+                  release(tagName: $tag) {
+                    isDraft
+                  }
+                }
+              }
+            ' \
+            --jq '.data.repository.release | if . == null then "" else .isDraft end')
           if [[ "$state" == "false" ]]; then
             echo "Release $RELEASE_TAG is already published; leaving its assets unchanged."
             exit 0

README.md

@@ -28,6 +28,8 @@ Local Forge is a fast, local-first desktop workbench for open models. Chat with
 - Persistent sessions, image attachments, search, pinning, and deletion
 - Local Diffusers and `.safetensors` / `.ckpt` model registration
 - Cancellable image generation with live step progress and Library import
+- Local image descriptions and generation prompts through Ollama vision models
+- Optional Face Fix, tiled 2x/4x UltraSharp upscaling, and saved NSFW region masks
 - Cancellable QLoRA training with live progress and standard PEFT output
 - Screenshot and reference-image imports shared by Studio and Library
 - Live RAM and NVIDIA telemetry, with unavailable data left unknown
@@ -47,4 +49,4 @@ npm run dev
 
 Opening Vite directly uses a browser preview adapter. Use the Electron window for filesystem imports, MCP servers, encrypted credential persistence, and real hardware data.
 
-See the [runtime and project guide](docs/RUNTIME_AND_PROJECT_GUIDE.md) for image and training setup, MCP servers, validation, releases, and project status.
+See the [runtime and project guide](docs/RUNTIME_AND_PROJECT_GUIDE.md) for image and training setup, MCP servers, validation, releases, and project status. See the [contribution guide](CONTRIBUTING.md), [security policy](SECURITY.md), [asset notes](docs/THIRD_PARTY_ASSETS.md), and [MIT License](LICENSE) for repository policies and project details.

docs/RUNTIME_AND_PROJECT_GUIDE.md

@@ -11,7 +11,9 @@ python -m pip install -r runtime/requirements-training.txt
 
 Select that interpreter under **Settings > Runtimes > Python runtime**. Local Forge never installs Python packages or downloads model weights automatically.
 
-Studio executes local Diffusers directories containing `model_index.json`. Registered `.safetensors` and `.ckpt` files remain visible for inventory but must be converted to Diffusers format before generation. Completed PNGs are added to Library and stored under Local Forge's user-data `outputs/images` directory.
+Studio executes local Diffusers directories containing `model_index.json`. Registered `.safetensors` and `.ckpt` files remain visible for inventory but must be converted to Diffusers format before generation. Completed PNGs are added to Library and stored under Local Forge's user-data `outputs/images` directory. Optional NSFW segmentation writes one binary mask per detected region beside the completed image.
+
+The Studio **Describe image** and **Create prompt** actions use the Ollama model currently selected in Workbench. That model must advertise Ollama's `vision` capability. Local Forge does not pre-classify or block adult images; description quality and model-level refusals depend on the selected vision model.
 
 Tune requires a local Hugging Face Transformers model directory containing `config.json`; Ollama tags and GGUF files are inference artifacts and are not offered as trainable base models. Datasets may be JSON, JSONL, or CSV and should contain either a `text` field or chat `messages`. Completed adapters use standard PEFT format under `outputs/adapters`.
 
@@ -35,17 +37,17 @@ npm run build
 
 ## Automated releases
 
-Every push to `main` (including a merged pull request) runs tests, lint, and the application build. After those checks pass, GitHub Actions packages that exact commit for Windows x64 (`.exe`), macOS universal (`.dmg`, Intel and Apple Silicon), and Linux x64 (`.AppImage`).
+Every push to `main` (including a merged pull request) runs tests, lint, and the application build. After those checks pass, GitHub Actions packages that exact commit for Windows x64 (`.exe`), one universal macOS `.dmg` that runs on both Intel and Apple Silicon, and Linux x64 (`.AppImage`).
 
 Download installers from [GitHub Releases](https://github.com/Azayzel/local-forge/releases). Each main build is an unsigned prerelease, tagged `v<package-version>-main.<test-run-number>` (for example, `v0.1.0-main.42`). The installer version matches the tag without its `v` prefix. These alpha builds are not marked as the latest stable release.
 
-All three installers and `SHA256SUMS.txt` are uploaded before the release is published. To verify a download, compare its SHA-256 hash with the matching entry in that file:
+All three installers and `SHA256SUMS.txt` are uploaded before the release is published. That file lists one tab-separated `filename<TAB>sha256` entry per installer. To verify a download, compare its SHA-256 hash with the matching entry in that file:
 
 ```powershell
 Get-FileHash .\Local-Forge-0.1.0-main.42-win-x64.exe -Algorithm SHA256
 ```
 
-Use your downloaded installer's actual filename. On Linux, download all three installers and the checksum file into one directory and run `sha256sum --check SHA256SUMS.txt`; on macOS use `shasum -a 256 -c SHA256SUMS.txt`.
+Use your downloaded installer's actual filename. On Linux, run `sha256sum ./Local-Forge-0.1.0-main.42-linux-x64.AppImage`; on macOS run `shasum -a 256 ./Local-Forge-0.1.0-main.42-mac-universal.dmg`; in either case, compare the printed hash with the matching entry in `SHA256SUMS.txt`.
 
 Windows and macOS may show security warnings: these builds are not yet signed or notarized. See [release maintenance](../CONTRIBUTING.md#release-maintenance) for setup, versioning, and retries.
 

docs/THIRD_PARTY_ASSETS.md

@@ -17,6 +17,20 @@ The files are `1200 x 1800`, `1200 x 800`, `1200 x 800`, and `1200 x 1800` respe
 
 Images imported by a user are copied into that user's private Local Forge data directory. They are never included in the source tree or project releases.
 
+## Optional enhancement models
+
+Local Forge does not bundle enhancement weights. When a user chooses Install in Studio, the app downloads an immutable, checksum-verified revision to that user's private Local Forge data directory:
+
+| Purpose | Repository | File | SHA-256 |
+| --- | --- | --- | --- |
+| 4x image upscaling | [lokCX/4x-Ultrasharp](https://huggingface.co/lokCX/4x-Ultrasharp) | `4x-UltraSharp.pth` | `a5812231fc936b42af08a5edba784195495d303d5b3248c24489ef0c4021fe01` |
+| Face detection | [Bingsu/adetailer](https://huggingface.co/Bingsu/adetailer) | `face_yolov8n.pt` | `70b640f8f60b1cf0dcc72f30caf3da9495eb2fb6509da48c53374ad6806e6a9c` |
+| Breast segmentation | [NSFW-API/NSFW_Segmentation](https://huggingface.co/NSFW-API/NSFW_Segmentation) | `nsfw-seg-breast-x.pt` | `5ad882ddaf149873be131943b373da9f15a0603c91508e2291ece83d729c8ecc` |
+| Penis segmentation | [NSFW-API/NSFW_Segmentation](https://huggingface.co/NSFW-API/NSFW_Segmentation) | `nsfw-seg-penis-x.pt` | `49d9fc8ee67d3bdee44e46bf75aeb76058f5cd074bac027b55ff071b63a32e21` |
+| Vagina segmentation | [NSFW-API/NSFW_Segmentation](https://huggingface.co/NSFW-API/NSFW_Segmentation) | `nsfw-seg-vagina-x.pt` | `f8349ae348f5cf041809c38bf38d2c80a0f3dccdb1375b97971358cde742197f` |
+
+These files remain subject to the terms published by their respective authors and are not covered by Local Forge's MIT code license.
+
 ## Fonts and icons
 
 - Manrope and Space Grotesk are bundled through Fontsource packages and distributed under the SIL Open Font License 1.1.

electron/chat.test.ts

@@ -7,7 +7,7 @@ import type {
   ChatStreamEvent,
   McpServerConfig,
 } from "../src/types";
-import { runMcpChat } from "./chat";
+import { runMcpChat, warmChatModel } from "./chat";
 import { closeMcpConnections } from "./mcp";
 
 const fixture: McpServerConfig = {
@@ -26,6 +26,41 @@ const fixture: McpServerConfig = {
 afterAll(() => closeMcpConnections());
 
 describe("MCP chat orchestration", () => {
+  it("preloads the selected model with the chat keep-alive", async () => {
+    let requestPath = "";
+    let requestBody: Record<string, unknown> = {};
+    const server = createServer(async (incoming, response) => {
+      requestPath = incoming.url ?? "";
+      const chunks: Buffer[] = [];
+      for await (const chunk of incoming) chunks.push(Buffer.from(chunk));
+      requestBody = JSON.parse(
+        Buffer.concat(chunks).toString("utf8"),
+      ) as Record<string, unknown>;
+      response.writeHead(200, { "content-type": "application/json" });
+      response.end(JSON.stringify({ done: true }));
+    });
+    await new Promise<void>((resolve) =>
+      server.listen(0, "127.0.0.1", resolve),
+    );
+    const address = server.address() as AddressInfo;
+
+    try {
+      await warmChatModel(`http://127.0.0.1:${address.port}`, "fixture-model");
+    } finally {
+      await new Promise<void>((resolve, reject) =>
+        server.close((error) => (error ? reject(error) : resolve())),
+      );
+    }
+
+    expect(requestPath).toBe("/api/generate");
+    expect(requestBody).toMatchObject({
+      model: "fixture-model",
+      prompt: "",
+      stream: false,
+      keep_alive: "30m",
+    });
+  });
+
   it("approves and executes a tool before continuing the Ollama chat", async () => {
     const requests: Array<Record<string, unknown>> = [];
     const server = createServer(async (incoming, response) => {
@@ -105,6 +140,8 @@ describe("MCP chat orchestration", () => {
     }
 
     expect(requests).toHaveLength(2);
+    expect(requests[0].think).toBe(false);
+    expect(requests[0].keep_alive).toBe("30m");
     expect(requests[0].tools).toEqual(
       expect.arrayContaining([
         expect.objectContaining({

electron/chat.ts

@@ -7,6 +7,7 @@ import type {
 import { callMcpTool, listEnabledMcpTools, type McpToolBinding } from "./mcp";
 
 const MAX_TOOL_ROUNDS = 8;
+const MODEL_KEEP_ALIVE = "30m";
 
 interface OllamaToolCall {
   function: {
@@ -40,18 +41,40 @@ interface RoundResult {
   evalDurationNs: number;
 }
 
-function runtimeEndpoint(baseUrl: string): string {
+function runtimeEndpoint(baseUrl: string, route = "/api/chat"): string {
   const url = new URL(baseUrl);
   const localHosts = new Set(["localhost", "127.0.0.1", "::1", "[::1]"]);
   if (url.protocol !== "http:" || !localHosts.has(url.hostname)) {
     throw new Error("Local Forge only connects to runtimes on this machine.");
   }
-  url.pathname = "/api/chat";
+  url.pathname = route;
   url.search = "";
   url.hash = "";
   return url.toString();
 }
 
+export async function warmChatModel(
+  baseUrl: string,
+  model: string,
+): Promise<void> {
+  if (!model.trim()) return;
+  const response = await fetch(runtimeEndpoint(baseUrl, "/api/generate"), {
+    method: "POST",
+    headers: { "content-type": "application/json" },
+    body: JSON.stringify({
+      model,
+      prompt: "",
+      stream: false,
+      keep_alive: MODEL_KEEP_ALIVE,
+    }),
+  });
+  if (!response.ok) {
+    const detail = await response.text();
+    throw new Error(detail || `Runtime returned ${response.status}.`);
+  }
+  await response.text();
+}
+
 async function readJsonLines(
   response: Response,
   onValue: (value: Record<string, unknown>) => void,
@@ -135,6 +158,8 @@ async function runRound(
       model: options.request.model,
       messages,
       options: options.request.options,
+      think: false,
+      keep_alive: MODEL_KEEP_ALIVE,
       tools: bindings.length > 0 ? toolDefinitions(bindings) : undefined,
       stream: true,
     }),

electron/content-security-policy.test.ts

@@ -0,0 +1,26 @@
+import { readFile } from "node:fs/promises";
+import path from "node:path";
+import { describe, expect, it } from "vitest";
+
+describe("content security policy", () => {
+  it("allows both Local Forge image protocols", async () => {
+    const html = await readFile(path.join(process.cwd(), "index.html"), "utf8");
+    const page = new DOMParser().parseFromString(html, "text/html");
+    const policy = page
+      .querySelector('meta[http-equiv="Content-Security-Policy"]')
+      ?.getAttribute("content");
+    const imageSources = policy
+      ?.split(";")
+      .find((directive) => directive.trim().startsWith("img-src "))
+      ?.trim()
+      .split(/\s+/)
+      .slice(1);
+
+    expect(imageSources).toEqual(
+      expect.arrayContaining([
+        "local-forge-attachment:",
+        "local-forge-output:",
+      ]),
+    );
+  });
+});

electron/enhancement-models.test.ts

@@ -0,0 +1,121 @@
+import { promises as fs } from "node:fs";
+import { createHash } from "node:crypto";
+import os from "node:os";
+import path from "node:path";
+import { afterEach, describe, expect, it } from "vitest";
+import {
+  discoverEnhancementModels,
+  downloadTrustedAsset,
+  type TrustedAsset,
+} from "./enhancement-models";
+
+const temporaryDirectories: string[] = [];
+
+async function temporaryDirectory(): Promise<string> {
+  const directory = await fs.mkdtemp(
+    path.join(os.tmpdir(), "local-forge-enhancements-"),
+  );
+  temporaryDirectories.push(directory);
+  return directory;
+}
+
+afterEach(async () => {
+  await Promise.all(
+    temporaryDirectories
+      .splice(0)
+      .map((directory) => fs.rm(directory, { recursive: true, force: true })),
+  );
+});
+
+describe("discoverEnhancementModels", () => {
+  it("discovers the Lavely enhancement model layout", async () => {
+    const root = await temporaryDirectory();
+    const upscaler = path.join(root, "upscalers", "4x-UltraSharp.pth");
+    const faceDetector = path.join(root, "face_detector", "face_yolov8n.pt");
+    const segmenterDirectory = path.join(root, "nsfw_segmentation");
+    await fs.mkdir(path.dirname(upscaler), { recursive: true });
+    await fs.mkdir(path.dirname(faceDetector), { recursive: true });
+    await fs.mkdir(segmenterDirectory, { recursive: true });
+    await fs.writeFile(upscaler, "upscaler");
+    await fs.writeFile(faceDetector, "detector");
+    await Promise.all(
+      ["breast", "penis", "vagina"].map((region) =>
+        fs.writeFile(
+          path.join(segmenterDirectory, `nsfw-seg-${region}-x.pt`),
+          region,
+        ),
+      ),
+    );
+
+    await expect(discoverEnhancementModels([root])).resolves.toEqual({
+      upscalerModelPath: upscaler,
+      faceDetectorModelPath: faceDetector,
+      nsfwSegmenterModelPath: segmenterDirectory,
+    });
+  });
+
+  it("ignores missing and incomplete model files", async () => {
+    const root = await temporaryDirectory();
+    const upscaler = path.join(root, "upscalers", "4x-UltraSharp.pth");
+    const segmenterDirectory = path.join(root, "nsfw_segmentation");
+    await fs.mkdir(path.dirname(upscaler), { recursive: true });
+    await fs.mkdir(segmenterDirectory, { recursive: true });
+    await fs.writeFile(upscaler, "");
+    await fs.writeFile(
+      path.join(segmenterDirectory, "nsfw-seg-breast-x.pt"),
+      "breast",
+    );
+    await fs.writeFile(
+      path.join(segmenterDirectory, "nsfw-seg-penis-x.pt"),
+      "penis",
+    );
+
+    await expect(discoverEnhancementModels([root])).resolves.toEqual({
+      upscalerModelPath: "",
+      faceDetectorModelPath: "",
+      nsfwSegmenterModelPath: "",
+    });
+  });
+});
+
+describe("downloadTrustedAsset", () => {
+  const payload = Buffer.from("trusted model payload");
+  const asset: TrustedAsset = {
+    directory: "upscalers",
+    filename: "model.pth",
+    url: "https://models.example/model.pth",
+    bytes: payload.byteLength,
+    sha256: createHash("sha256").update(payload).digest("hex"),
+  };
+
+  it("promotes a verified download to its final path", async () => {
+    const root = await temporaryDirectory();
+    const target = await downloadTrustedAsset(
+      root,
+      asset,
+      async () => new Response(payload),
+    );
+
+    await expect(fs.readFile(target)).resolves.toEqual(payload);
+    await expect(fs.stat(`${target}.part`)).rejects.toMatchObject({
+      code: "ENOENT",
+    });
+  });
+
+  it("rejects a corrupt download without leaving model files", async () => {
+    const root = await temporaryDirectory();
+    const target = path.join(root, asset.directory, asset.filename);
+
+    await expect(
+      downloadTrustedAsset(
+        root,
+        asset,
+        async () => new Response("not the expected model"),
+      ),
+    ).rejects.toThrow("integrity check");
+    await expect(fs.stat(target)).rejects.toMatchObject({ code: "ENOENT" });
+    await expect(fs.stat(`${target}.part`)).rejects.toMatchObject({
+      code: "ENOENT",
+    });
+  });
+});