Merge pull request #236 from ausboss/feat/lora-civitai-picture

LoRA Loader: Fetch Civitai info also saves the LoRA's picture
This commit is contained in:
ausboss
2026-10-24 18:32:00 -05:00
committed by GitHub
11 changed files with 317 additions and 43 deletions
+5 -3
View File
@@ -121,14 +121,16 @@ example_workflows/ # example workflows (regular workflow JSON, not API JSON)
A route answers with data it made (JSON, a small picture), never a file by
its path (`web.FileResponse`), and its error text goes through
`safe_route_error` (`nodes/_route_errors.py`). Besides those the pack only reads
its registered model folders (the one write there is the `.civitai.info`
file LoRA Loader's Civitai button saves beside a LoRA) and keeps its own
its registered model folders (the one write there is LoRA Loader's Civitai
button: a `.civitai.info` file and a `.preview.jpg` picture, saved beside a
LoRA, never replacing a picture) and keeps its own
settings files in ComfyUI's user folder (`user/ausboss/`). Nothing is handed to a
subprocess. The Registry bans versions for exactly this.
- The pack's server code makes no network requests: no HTTP or socket client
in shipped Python (`release_preflight.py` fails on one). The one outbound
call is LoRA Loader's click-only Civitai lookup, made by the browser to a
fixed civitai.com address with a plain SHA256 (`js/lora_loader/index.js`,
fixed civitai.com address with a plain SHA256, plus the LoRA's picture
from `image.civitai.com` on the same click (`js/lora_loader/index.js`,
`js/shared/civitai_lookup.mjs`). `release_preflight.py` allows a raw
`fetch(` only there and in `js/cache_guard`, which only talks to the pack's
own server. Graph links are made with
+10
View File
@@ -2,6 +2,16 @@
All notable changes to ComfyUI-AusBoss are documented here.
## Unreleased
- **LoRA Loader 🆎 saves the LoRA's picture from Civitai, no add-on
needed.** **Fetch Civitai info** now also saves the LoRA's first picture
next to the file, so hovering the LoRA shows it. Pictures Civitai rates R
or above are skipped unless you pick **any** under the gear menu's
**Civitai picture** option, and a picture you already have is never
replaced. The separate LoRA Previews add-on from 2.19.0 is no longer
needed.
## 2.19.0 - 2026-10-24
- **LoRA Loader 🆎 can add several LoRAs at once.** In the **+ LoRA** list,
+1 -1
View File
@@ -263,7 +263,7 @@ The toolbar has the stack toggle, saved templates, reconnect and settings. The t
![LoRA Loader absorbs a connected LoRA chain, shows strength changes while scrubbing, and restores the previous toggle selection.](assets/readme/lora-chain-demo.gif)
Each LoRA gets an info card from the file's metadata, a `.civitai.info` file next to it, and words you save yourself. The card's **Fetch Civitai info** button looks the file up on Civitai (your browser asks civitai.com, only when you click) and saves the answer as that `.civitai.info` file. Files Civitai doesn't know, such as unpublished or hidden models, come back as not found. The trigger words you pick come out of the `triggers` output. To get the LoRA's preview picture from Civitai as well, install the separate [LoRA Previews add-on](https://github.com/ausboss/ComfyUI-AusBoss-LoRA-Previews#install): it saves the picture when you fetch the info, and hovering the LoRA shows it.
Each LoRA gets an info card from the file's metadata, a `.civitai.info` file next to it, and words you save yourself. The card's **Fetch Civitai info** button looks the file up on Civitai (your browser asks civitai.com, only when you click) and saves the answer as that `.civitai.info` file. Files Civitai doesn't know, such as unpublished or hidden models, come back as not found. The trigger words you pick come out of the `triggers` output. The same click saves the LoRA's picture from Civitai next to the file, so hovering the LoRA shows it (PG-13 or milder unless you allow more in the gear menu; a picture you already have is never replaced).
In the **+ LoRA** list, tick several files to add them at once. Rows can be gathered into groups, like layer groups in an image editor: each group folds away, and has its own switch that turns all its LoRAs off while each one keeps its own setting. A Workflow Switch can turn a whole group on and off.
+6 -2
View File
@@ -34,8 +34,12 @@ follows a few fixed rules.
`<model>.civitai.info` (inside ComfyUI's LoRA folders). The pack checks
that the answer describes the same file before saving it. Nothing is sent
until you click, and the button can be turned off in the node's settings.
The pack never downloads the LoRA's preview picture; that is left to a
separate add-on you can choose to install (LoRA Previews).
The same click also has your browser download the LoRA's first picture
from Civitai's picture server (a small copy, without cookies, rated PG-13
or milder unless you allow more). The pack checks it is a picture and
saves it beside the LoRA as `<model>.preview.jpg`, never replacing a
picture that is there. The gear menu's **Civitai picture** option turns
this off.
- **No code execution.** Math Expression parses its text into an
allowlisted syntax tree of arithmetic and a few math functions; it never
runs Python source. The pack starts no subprocesses and installs nothing
+9 -7
View File
@@ -303,12 +303,13 @@ or the bare `.png` / `.jpg` / `.jpeg` / `.webp`). No sidecar, no thumbnail —
nothing appears and nothing shifts. The gear menu's **Preview thumbnails**
switch turns the behavior off.
**Get the pictures from Civitai:** install the separate
[LoRA Previews add-on](https://github.com/ausboss/ComfyUI-AusBoss-LoRA-Previews#install). With it, **Fetch Civitai info** (or
**Refresh Civitai info**) also saves the LoRA's preview picture from Civitai
next to the file, and the hover and the info card show it right away. It never
replaces a picture you already have. The pack itself downloads no pictures; the
info card links to the add-on when a LoRA has Civitai info but no picture.
**Get the pictures from Civitai:** click **Fetch Civitai info** (or
**Refresh Civitai info**) in a LoRA's info card. It also saves the LoRA's
first picture from Civitai next to the file, and the hover and the info card
show it right away. Gear menu → **Civitai picture** picks which pictures:
**PG-13** (the default) skips pictures Civitai rates R or above, **any**
takes the first picture whatever its rating, and **off** saves none. It never
replaces a picture you already have.
## Trigger words
@@ -326,7 +327,8 @@ switched on for the old one.
- **Fetch Civitai info** looks this one LoRA up on civitai.com. Nothing is
sent until you click. Your browser then sends only the file's SHA256
fingerprint, and the answer is saved next to the LoRA as
`<model>.civitai.info`. After that the button reads **Refresh Civitai
`<model>.civitai.info`, with the LoRA's picture as `<model>.preview.jpg`
(see "Get the pictures from Civitai" above). After that the button reads **Refresh Civitai
info** and the card gets a **View on Civitai** link. A LoRA Civitai does
not know (unpublished or hidden) reads **Not found on Civitai**. Gear
menu → **Civitai lookup button** removes the button. If it always reads
+56 -29
View File
@@ -82,7 +82,10 @@ import {
uniqueGroupName,
upsertTemplate,
} from "../shared/lora_stack.mjs";
import { civitaiButtonLabel, civitaiByHashUrl, parseCivitaiBody } from "../shared/civitai_lookup.mjs";
import {
MAX_PICTURE_BYTES, civitaiButtonLabel, civitaiByHashUrl, isCivitaiAddress, parseCivitaiBody,
pickCivitaiPicture, smallPictureUrl,
} from "../shared/civitai_lookup.mjs";
const NODE_CLASS = "AUSBOSS_NODES_LoraLoader";
// A workflow saved by a 1.x release holds the panel's own empty value where
@@ -173,6 +176,13 @@ const SETTINGS_SCHEMA = [
+ "asks civitai.com about that one file, from this browser, and saves "
+ "the answer next to the LoRA. Nothing is sent until you click.",
},
{
key: "civitai_picture", label: "Civitai picture", type: "choice",
options: ["off", "PG-13", "any"], default: "PG-13",
hint: "Fetch Civitai info also saves the LoRA's first picture next to "
+ "the file, so hovering the LoRA shows it. PG-13 skips pictures Civitai "
+ "rates R or above. A picture you already have is never replaced.",
},
{
key: "name_cut", label: "Long names", type: "choice",
options: ["end", "middle"], default: "end",
@@ -608,10 +618,6 @@ function closePopup() {
// ---------- hover thumbnail ----------
// Where to read about the add-on that saves a LoRA's preview picture from
// Civitai. The pack itself downloads no pictures.
const PREVIEWS_ADDON_URL = "https://github.com/ausboss/ComfyUI-AusBoss-LoRA-Previews#install";
// Raised when a preview picture was saved since the page loaded, so the
// browser asks for the picture again instead of reusing a failed load.
let thumbVersion = 0;
@@ -621,11 +627,11 @@ function thumbUrl(name) {
return api.apiURL(`/ausboss/lora/thumb?name=${encodeURIComponent(name)}${version}`);
}
// The preview add-on saved a picture: the next hover and an open card show it.
window.addEventListener("ausboss:lora-preview-saved", () => {
// A picture was saved: the next hover and the card show it.
function pictureSaved() {
thumbVersion += 1;
if (hoverThumb) delete hoverThumb.dataset.name;
});
}
// One floating preview image for the whole pack tab, shared by row hover and
// the picker. Fixed-positioned and pointer-events: none, so it can never
@@ -1384,8 +1390,9 @@ function openPicker(state, index, anchor, { keyboard = false } = {}) {
// One click, one file: the server hashes the LoRA, this browser asks Civitai
// about that hash, and the server saves the answer beside the LoRA after
// checking it describes the same file. Returns false when Civitai does not
// know the file (unpublished or hidden models answer 404).
async function lookUpOnCivitai(serverName) {
// know the file (unpublished or hidden models answer 404). `picture` is the
// Civitai picture gear option ("off", "PG-13" or "any").
async function lookUpOnCivitai(serverName, picture = "PG-13") {
const hashed = await (await api.fetchApi(
`/ausboss/lora/hash?name=${encodeURIComponent(serverName)}`
)).json();
@@ -1402,11 +1409,43 @@ async function lookUpOnCivitai(serverName) {
body: JSON.stringify({ name: serverName, info }),
})).json();
if (!saved.ok) throw new Error(saved.error || "Could not save the Civitai info.");
// An add-on can listen for this to save the preview picture.
window.dispatchEvent(new CustomEvent("ausboss:lora-civitai-saved", { detail: { name: serverName, info } }));
if (picture !== "off") {
await saveCivitaiPicture(serverName, info, picture === "any").catch((error) =>
loraToast(`Saved the Civitai info, but not its picture: ${error?.message || error}`));
}
return true;
}
// The LoRA's first picture on Civitai, saved beside it as its preview. The
// browser downloads a 450-pixel copy from Civitai's picture server without
// cookies; the server checks it is a picture and never replaces one.
async function saveCivitaiPicture(serverName, info, anyRating) {
const picture = pickCivitaiPicture(info, anyRating);
if (!picture) {
loraToast(anyRating
? "Civitai has no picture of this LoRA to save."
: "Civitai has no picture of this LoRA rated PG-13 or milder.");
return;
}
const response = await fetch(smallPictureUrl(picture.url), {
credentials: "omit", referrerPolicy: "no-referrer", cache: "no-store",
});
if (!response.ok) throw new Error(`Civitai's picture server answered ${response.status}.`);
if (!isCivitaiAddress(response.url)) throw new Error("The picture came from somewhere other than Civitai.");
if (!(response.headers.get("content-type") || "").startsWith("image/")) {
throw new Error("Civitai did not send a picture.");
}
const blob = await response.blob();
if (blob.size > MAX_PICTURE_BYTES) throw new Error("The picture is too large.");
const saved = await (await api.fetchApi(`/ausboss/lora/picture?name=${encodeURIComponent(serverName)}`, {
method: "POST",
headers: { "Content-Type": blob.type || "application/octet-stream" },
body: blob,
})).json();
if (!saved.ok) throw new Error(saved.error || "Could not save the picture.");
pictureSaved();
}
// ---------- info card ----------
function openInfo(state, index, anchor) {
@@ -1445,19 +1484,9 @@ function openInfo(state, index, anchor) {
card.append(link);
}
// Civitai info but no picture: say how to get one.
if (info.has_civitai && !info.has_preview && state.settings?.civitai_lookup !== false) {
const hint = el("div", "ausboss-lora-meta");
if (window.ausbossLoraPreviews) {
hint.textContent = "No preview picture yet. Click Refresh Civitai info to save one.";
} else {
hint.append("No preview picture yet. The LoRA Previews add-on saves Civitai's picture for you. ");
const addon = el("a", "ausboss-lora-civitai-link", "How to install ↗");
addon.href = PREVIEWS_ADDON_URL;
addon.target = "_blank";
addon.rel = "noopener noreferrer";
hint.append(addon);
}
card.append(hint);
if (info.has_civitai && !info.has_preview && state.settings?.civitai_lookup !== false
&& state.settings?.civitai_picture !== "off") {
card.append(el("div", "ausboss-lora-meta", "No preview picture yet. Click Refresh Civitai info to save one."));
}
// One click can change other chips too: the same word can sit in two
@@ -1501,7 +1530,8 @@ function openInfo(state, index, anchor) {
fetchButton.disabled = true;
fetchButton.textContent = "Fetching...";
try {
if (!(await lookUpOnCivitai(serverName))) {
const picture = info.has_preview ? "off" : (state.settings?.civitai_picture ?? "PG-13");
if (!(await lookUpOnCivitai(serverName, picture))) {
fetchButton.textContent = "Not found on Civitai";
return;
}
@@ -1591,9 +1621,6 @@ function openInfo(state, index, anchor) {
card.append(el("div", "ausboss-lora-empty", "Could not read LoRA info."));
});
};
window.addEventListener("ausboss:lora-preview-saved", (event) => {
if (event.detail?.name === serverName) load();
}, { signal: activePopup?.abort.signal });
load();
}
+65
View File
@@ -29,3 +29,68 @@ export function parseCivitaiBody(text) {
export function civitaiButtonLabel(hasCivitai) {
return hasCivitai ? "Refresh Civitai info" : "Fetch Civitai info";
}
// ---------- the LoRA's preview picture ----------
// After the lookup, the browser also downloads the LoRA's first picture from
// Civitai's picture server, and the pack's server saves it beside the LoRA.
export const PICTURE_HOST = "image.civitai.com";
export const PICTURE_WIDTH = 450;
export const MAX_PICTURE_BYTES = 8 * 1024 * 1024;
// Civitai rates each picture. The number is a set of flags: 1 = PG,
// 2 = PG-13, 4 = R, 8 = X, 16 = XXX. Older answers give a word instead.
const MILD_LEVEL = 2;
const MILD_WORDS = new Set(["none", "soft"]);
export function isMildPicture(image) {
const level = Number(image?.nsfwLevel);
if (Number.isFinite(level) && level > 0) return level <= MILD_LEVEL;
if (typeof image?.nsfw === "string") return MILD_WORDS.has(image.nsfw.toLowerCase());
if (image?.nsfw === false) return true;
return false;
}
// The first still picture in Civitai's answer that may be saved: PG and
// PG-13 only, unless any rating is allowed.
export function pickCivitaiPicture(info, anyRating = false) {
const images = Array.isArray(info?.images) ? info.images : [];
return images.find((image) =>
image && typeof image.url === "string"
&& (image.type ?? "image") === "image"
&& (anyRating || isMildPicture(image))
&& smallPictureUrl(image.url) !== null) ?? null;
}
// The address of a 450-pixel copy of the picture, or null for any address
// that is not Civitai's picture server over https.
export function smallPictureUrl(url) {
let parsed;
try {
parsed = new URL(url);
} catch {
return null;
}
if (parsed.protocol !== "https:" || parsed.hostname !== PICTURE_HOST) return null;
const parts = parsed.pathname.split("/");
const sizing = parts.findIndex((part, index) => index > 1 && index < parts.length - 1 && part.includes("="));
if (sizing >= 0) parts[sizing] = `width=${PICTURE_WIDTH}`;
else if (parts.length >= 3) parts.splice(parts.length - 1, 0, `width=${PICTURE_WIDTH}`);
else return null;
parsed.pathname = parts.join("/");
parsed.search = "";
parsed.hash = "";
return parsed.href;
}
// Where the picture really came from, after Civitai's redirect to its file
// store: Civitai's own servers, over https.
export function isCivitaiAddress(url) {
try {
const parsed = new URL(url);
return parsed.protocol === "https:"
&& (parsed.hostname === "civitai.com" || parsed.hostname.endsWith(".civitai.com"));
} catch {
return false;
}
}
+60
View File
@@ -731,6 +731,53 @@ def save_civitai_sidecar(name: str, payload: Any) -> dict[str, Any]:
return info
# Civitai's preview picture, saved beside the LoRA so the hover preview finds
# it. The browser downloads it; this server only checks and writes it.
MAX_PICTURE_BYTES = 8 * 1024 * 1024
MAX_PICTURE_PIXELS = 40_000_000
PICTURE_EDGE = 768
def save_civitai_picture(name: str, data: bytes) -> Path:
"""Save a picture as <model>.preview.jpg next to the LoRA.
The bytes come from the browser, so they are decoded here and written
again as a small JPEG with no metadata. A LoRA that already has any
preview picture keeps it: nothing is replaced, even in a race.
"""
if not isinstance(data, (bytes, bytearray)) or not data:
raise ValueError("No picture was sent.")
if len(data) > MAX_PICTURE_BYTES:
raise ValueError("The picture is too large.")
path = resolve_lora_path(name)
if find_thumbnail(name) is not None:
raise ValueError("This LoRA already has a preview picture.")
from io import BytesIO
from PIL import Image
try:
with Image.open(BytesIO(data)) as picture:
width, height = picture.size
if width * height > MAX_PICTURE_PIXELS:
raise ValueError("The picture is too large.")
picture = picture.convert("RGB")
picture.thumbnail((PICTURE_EDGE, PICTURE_EDGE))
out = BytesIO()
picture.save(out, "JPEG", quality=88)
except ValueError:
raise
except Exception as exc:
raise ValueError("That is not a picture.") from exc
target = Path(str(path.with_suffix("")) + ".preview.jpg")
try:
with open(target, "xb") as handle:
handle.write(out.getvalue())
except FileExistsError as exc:
raise ValueError("This LoRA already has a preview picture.") from exc
return target
def lora_info(name: str) -> dict[str, Any]:
path = resolve_lora_path(name)
metadata = read_safetensors_metadata(path)
@@ -815,6 +862,19 @@ def register_lora_routes() -> None:
except Exception as exc:
return web.json_response({"ok": False, "error": safe_route_error(exc)}, status=400)
@prompt_server.routes.post("/ausboss/lora/picture")
async def ausboss_lora_picture(request):
try:
if (request.content_length or 0) > MAX_PICTURE_BYTES:
raise ValueError("The picture is too large.")
data = await request.read()
await asyncio_run_in_executor(
save_civitai_picture, request.query.get("name", ""), data
)
return web.json_response({"ok": True})
except Exception as exc:
return web.json_response({"ok": False, "error": safe_route_error(exc)}, status=400)
@prompt_server.routes.post("/ausboss/lora/triggers")
async def ausboss_lora_triggers(request):
try:
+2 -1
View File
@@ -313,7 +313,8 @@ for path in shipped_source:
# The browser may call fetch() in two places only: the cache guard (it asks
# the pack's own server) and LoRA Loader's click-only Civitai lookup (a fixed
# civitai.com address with a plain SHA256). Anything else is a new outbound
# civitai.com address with a plain SHA256, then the LoRA's picture from
# image.civitai.com). Anything else is a new outbound
# call and needs the same review SECURITY.md describes.
FETCH_ALLOWED = {"js/cache_guard/index.js", "js/lora_loader/index.js"}
for path in shipped_source:
+46
View File
@@ -6,7 +6,11 @@ import {
MAX_CIVITAI_BYTES,
civitaiButtonLabel,
civitaiByHashUrl,
isCivitaiAddress,
isMildPicture,
parseCivitaiBody,
pickCivitaiPicture,
smallPictureUrl,
} from "../js/shared/civitai_lookup.mjs";
const SHA = "a".repeat(64);
@@ -35,3 +39,45 @@ test("the button says refresh once Civitai info exists", () => {
assert.equal(civitaiButtonLabel(false), "Fetch Civitai info");
assert.equal(civitaiButtonLabel(true), "Refresh Civitai info");
});
const picture = (name) => `https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/0b1c2d3e/width=1024/${name}.jpeg`;
test("only PG and PG-13 pictures are picked unless any rating is allowed", () => {
const info = { images: [
{ url: picture("video"), type: "video", nsfwLevel: 1 },
{ url: picture("r"), type: "image", nsfwLevel: 4 },
{ url: picture("pg13"), type: "image", nsfwLevel: 2 },
] };
assert.equal(pickCivitaiPicture(info).url, picture("pg13"));
assert.equal(pickCivitaiPicture(info, true).url, picture("r"));
assert.equal(pickCivitaiPicture({ images: [{ url: picture("x"), nsfwLevel: 8 }] }), null);
assert.equal(pickCivitaiPicture({}), null);
assert.equal(pickCivitaiPicture(null), null);
});
test("older answers rate pictures with a word, and an unrated picture is not mild", () => {
assert.equal(isMildPicture({ nsfw: "None" }), true);
assert.equal(isMildPicture({ nsfw: "Soft" }), true);
assert.equal(isMildPicture({ nsfw: "Mature" }), false);
assert.equal(isMildPicture({}), false);
});
test("the small copy comes only from Civitai's picture server over https", () => {
assert.equal(smallPictureUrl(picture("a")), "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/0b1c2d3e/width=450/a.jpeg");
assert.equal(smallPictureUrl("https://image.civitai.com/key/uuid/original=true/a.png"),
"https://image.civitai.com/key/uuid/width=450/a.png");
assert.equal(smallPictureUrl("https://image.civitai.com/key/uuid/a.png"), "https://image.civitai.com/key/uuid/width=450/a.png");
assert.equal(smallPictureUrl("http://image.civitai.com/key/uuid/width=1024/a.png"), null);
assert.equal(smallPictureUrl("https://image.civitai.com.example.org/key/uuid/a.png"), null);
assert.equal(smallPictureUrl("https://example.org/a.png"), null);
assert.equal(smallPictureUrl("not a url"), null);
assert.equal(pickCivitaiPicture({ images: [{ url: "https://example.org/a.png", nsfwLevel: 1 }] }), null);
});
test("after the redirect the picture must still be on a Civitai server", () => {
assert.equal(isCivitaiAddress("https://blobs-b2.civitai.com/file.jpeg"), true);
assert.equal(isCivitaiAddress("https://civitai.com/file.jpeg"), true);
assert.equal(isCivitaiAddress("http://blobs-b2.civitai.com/file.jpeg"), false);
assert.equal(isCivitaiAddress("https://evilcivitai.com/file.jpeg"), false);
assert.equal(isCivitaiAddress(""), false);
});
+57
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import contextlib
import importlib.util
import io
import json
from pathlib import Path
@@ -224,6 +225,62 @@ class LoraCivitaiSidecarTests(unittest.TestCase):
self.assertFalse(outside.with_suffix(".civitai.info").exists())
def _picture(width=1200, height=800, fmt="PNG", mode="RGB") -> bytes:
from PIL import Image
out = io.BytesIO()
Image.new(mode, (width, height)).save(out, fmt)
return out.getvalue()
@unittest.skipIf(importlib.util.find_spec("PIL") is None, "needs Pillow (run with ComfyUI's Python)")
class LoraCivitaiPictureTests(unittest.TestCase):
name = "Krea 2/candid.safetensors"
def setUp(self):
LoraCivitaiSidecarTests.setUp(self)
self.target = self.lora.parent / "candid.preview.jpg"
def tearDown(self):
LoraCivitaiSidecarTests.tearDown(self)
def test_saves_a_small_jpeg_the_hover_preview_finds(self):
from PIL import Image
saved = _lora_helpers.save_civitai_picture(self.name, _picture(fmt="WEBP", mode="RGBA"))
self.assertEqual(saved, self.target)
with Image.open(saved) as picture:
self.assertEqual(picture.format, "JPEG")
self.assertEqual(max(picture.size), _lora_helpers.PICTURE_EDGE)
self.assertFalse(picture.info.get("exif"))
self.assertEqual(_lora_helpers.find_thumbnail(self.name), self.target)
self.assertTrue(_lora_helpers.lora_info(self.name)["has_preview"])
def test_never_replaces_a_picture_that_is_there(self):
own = self.lora.parent / "candid.png"
own.write_bytes(_picture(10, 10))
with self.assertRaises(ValueError):
_lora_helpers.save_civitai_picture(self.name, _picture())
self.assertFalse(self.target.exists())
self.assertEqual(own.read_bytes(), _picture(10, 10))
def test_refuses_what_is_not_a_picture_or_too_large(self):
for data in (b"", b"not a picture", b"\x89PNG\r\n\x1a\n broken", None):
with self.assertRaises(ValueError):
_lora_helpers.save_civitai_picture(self.name, data)
with self.assertRaises(ValueError):
_lora_helpers.save_civitai_picture(self.name, b"x" * (_lora_helpers.MAX_PICTURE_BYTES + 1))
with self.assertRaises(ValueError):
_lora_helpers.save_civitai_picture(self.name, _picture(8000, 6000, mode="L"))
self.assertFalse(self.target.exists())
def test_only_next_to_a_listed_lora(self):
outside = Path(self._tmp.name) / "secret.safetensors"
outside.write_bytes(b"x")
for name in ("../../secret.safetensors", str(outside), "Krea 2/missing.safetensors"):
with self.assertRaises(Exception):
_lora_helpers.save_civitai_picture(name, _picture())
self.assertFalse((Path(self._tmp.name) / "secret.preview.jpg").exists())
class ParseStackTests(unittest.TestCase):