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:
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -263,7 +263,7 @@ The toolbar has the stack toggle, saved templates, reconnect and settings. The t
|
||||
|
||||

|
||||
|
||||
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
@@ -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
|
||||
|
||||
@@ -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
@@ -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();
|
||||
}
|
||||
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user