Merge remote-tracking branch 'origin/diffusion-lora' into diffusion-controlnet

This commit is contained in:
Daniel Han 2026-07-02 02:27:40 +00:00
commit c3196cb8bd
16 changed files with 2286 additions and 1408 deletions

View file

@ -43,9 +43,10 @@ False positives:
examples and `>>>` doctests cannot trip a finding. Residual findings that
are genuine library behavior (a HTTP client reading HF_TOKEN, a vendored
test fixture) are suppressed via a reviewed baseline allowlist, matched on
(package, basename(file), check). A NEW kind of finding in an already-listed
file is a different check and still fails. This mirrors the Hugging Face Hub
approach (ClamAV/picklescan: low-FP, signature/structural, surface status).
(package, package-relative file, check, evidence hash). A new check, or
changed flagged code under the same check, reopens the finding; version
bumps and line shifts do not. This mirrors the Hugging Face Hub approach
(ClamAV/picklescan: low-FP, signature/structural, surface status).
Exit codes:
0 -- no non-baselined CRITICAL or HIGH findings (or --write-baseline)
@ -55,6 +56,8 @@ Exit codes:
import argparse
import atexit
import bisect
import hashlib
import io
import json
import os
@ -156,6 +159,9 @@ RE_EMBEDDED_KEYS = re.compile(
re.DOTALL,
)
# Full PEM block (BEGIN..END), used to pin a multiline key body in evidence.
RE_PEM_BLOCK = re.compile(r"-----BEGIN[^\n]*KEY-----.*?-----END[^\n]*KEY-----", re.DOTALL)
# Cloud metadata / IMDS endpoints
RE_CLOUD_METADATA = re.compile(
r"169\.254\.169\.254" # AWS/Azure/GCP IMDS
@ -476,22 +482,26 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]:
# Large base64 blob
if RE_LARGE_BLOB.search(content):
blob = RE_LARGE_BLOB.search(content).group()
# Digest every blob (not just the first 120 chars, and not just the
# first blob), so a later payload that keeps the prefix or appends a
# second encoded blob reopens.
blob, digest = _blob_digest(content)
findings.append(
Finding(
CRITICAL,
package,
filename,
f".pth has large base64-like blob ({len(blob)} chars)",
blob[:120] + "...",
f"{blob[:120]}... sha256:{digest}",
)
)
# Catch-all: any import line in .pth if nothing else triggered
# Catch-all: any import line in .pth if nothing else triggered. Bind every
# line through a digest so an appended/swapped import reopens the key, but cap
# the displayed text so a large .pth of benign-looking imports cannot dump up
# to the archive member cap into the logs or baseline JSON.
if not findings and import_lines:
evidence = "\n".join(import_lines[:5])
if len(import_lines) > 5:
evidence += f"\n... ({len(import_lines)} import lines total)"
evidence = _cap_line("\n".join(import_lines))
findings.append(
Finding(
HIGH,
@ -505,13 +515,15 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]:
# Unusually large executable .pth (litellm's was 34 KB; legit ones are <100 bytes)
size = len(content)
if size > 500 and import_lines:
# Pin the content so a different payload of the same size/import count reopens.
digest = hashlib.sha256(content.encode("utf-8", "replace")).hexdigest()
findings.append(
Finding(
HIGH,
package,
filename,
f"Unusually large executable .pth ({size} bytes)",
f"{len(import_lines)} import line(s) in {size}-byte .pth file",
f"{len(import_lines)} import line(s) in {size}-byte .pth file sha256:{digest}",
)
)
@ -629,6 +641,13 @@ def _hidden_payload_findings(
removed = "".join(o if o != s else " " for o, s in zip(original, code))
out = []
# The visible exec/eval line is what makes the hidden string executable, so
# bind it into every finding's evidence: otherwise a reviewed false positive
# that keeps the same hidden text but flips a harmless `eval("1+1")` to
# `exec(__doc__)` (now running the payload) keeps the same key and stays
# suppressed. Taken from `stripped` (real code), where the exec/eval lives.
trigger = _extract_evidence(stripped, RE_EXEC_EVAL)
def _hidden(pat):
# Carrier present in a blanked region but NOT in real code. A carrier in
# real code is already caught by the normal check, so restricting to
@ -643,7 +662,7 @@ def _hidden_payload_findings(
package,
filename,
"exec/eval with payload hidden in a docstring/string",
f"{label}: {_extract_evidence(removed, pat)}",
f"exec: {trigger}\n{label}: {_extract_evidence(removed, pat)}",
)
)
# Fetch-then-run dropper: a network call AND an os/subprocess exec that both
@ -657,7 +676,9 @@ def _hidden_payload_findings(
package,
filename,
"exec/eval with hidden network+exec payload",
f"network+exec: {_extract_evidence(removed, RE_SUBPROCESS)}",
f"exec: {trigger}\n"
f"network+exec: {_extract_evidence(removed, RE_NETWORK)} | "
f"{_extract_evidence(removed, RE_SUBPROCESS)}",
)
)
return out
@ -717,14 +738,19 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
# openssl encryption + network/key material (encrypted exfiltration)
if has_openssl_cli and (has_network or has_keys):
# Bind whichever side(s) co-occur so a changed endpoint or key reopens.
evidence = [f"OpenSSL: {_extract_evidence(content, RE_OPENSSL_CLI)}"]
if has_network:
evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}")
if has_keys:
evidence.append(f"Key: {_embedded_key_evidence(content)}")
findings.append(
Finding(
CRITICAL,
package,
filename,
"openssl encryption + network/key material (encrypted exfiltration)",
f"OpenSSL: {_extract_evidence(content, RE_OPENSSL_CLI)}\n"
f"Network: {_extract_evidence(content, RE_NETWORK)}",
"\n".join(evidence),
)
)
@ -896,6 +922,10 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
# Obfuscated payload: base64 + exec/eval + large blob
if has_base64 and has_exec_eval and has_blob:
# Digest every blob too: a payload may sit on a separate line from the
# decode call, and a second encoded blob may be appended later, so
# binding only the base64/exec lines or the first blob would miss it.
_, blob_digest = _blob_digest(content)
findings.append(
Finding(
HIGH,
@ -903,7 +933,8 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
filename,
"base64 decode + exec/eval + large encoded blob",
f"Base64: {_extract_evidence(content, RE_BASE64)}\n"
f"Exec: {_extract_evidence(content, RE_EXEC_EVAL)}",
f"Exec: {_extract_evidence(content, RE_EXEC_EVAL)}\n"
f"Blob: sha256:{blob_digest}",
)
)
@ -928,32 +959,48 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
package,
filename,
"Embedded cryptographic key + network calls (encrypted exfil pattern)",
f"Key: {_extract_evidence(content, RE_EMBEDDED_KEYS)}\n"
f"Key: {_embedded_key_evidence(content)}\n"
f"Network: {_extract_evidence(content, RE_NETWORK)}",
)
)
# Anti-analysis + any other suspicious pattern
if has_anti and (has_network or has_subprocess or has_exec_eval):
# Bind the suspicious side too so a changed payload reopens.
evidence = [f"Anti: {_extract_evidence(content, RE_ANTI_ANALYSIS)}"]
if has_network:
evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}")
if has_subprocess:
evidence.append(f"Subprocess: {_extract_evidence(content, RE_SUBPROCESS)}")
if has_exec_eval:
evidence.append(f"Exec: {_extract_evidence(content, RE_EXEC_EVAL)}")
findings.append(
Finding(
HIGH,
package,
filename,
"Anti-analysis/sandbox evasion + suspicious behavior",
f"Anti: {_extract_evidence(content, RE_ANTI_ANALYSIS)}",
"\n".join(evidence),
)
)
# DNS exfiltration with dynamic hostnames
if has_dns_exfil and (has_base64 or has_network or has_creds):
# Bind the co-occurring side so a changed exfil channel reopens.
evidence = [f"DNS: {_extract_evidence(content, RE_DNS_EXFIL)}"]
if has_base64:
evidence.append(f"Base64: {_extract_evidence(content, RE_BASE64)}")
if has_network:
evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}")
if has_creds:
evidence.append(f"Creds: {_extract_evidence(content, RE_CRED_ACCESS)}")
findings.append(
Finding(
HIGH,
package,
filename,
"DNS exfiltration / tunneling patterns",
_extract_evidence(content, RE_DNS_EXFIL),
"\n".join(evidence),
)
)
@ -1064,7 +1111,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
package,
filename,
"Embedded cryptographic key material",
_extract_evidence(content, RE_EMBEDDED_KEYS),
_embedded_key_evidence(content),
)
)
@ -1107,39 +1154,349 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
return findings
_MAX_MULTILINE_LINES = 12
# How far a single matched call is followed over its bracket continuations. A call
# that genuinely closes is bound all the way to its real close, up to the hard
# limit, so a ``requests.post(`` with many option/header lines before ``data=``
# binds its whole argument list in the digest and a changed payload on a late
# continuation line reopens (a 40-line soft cap would hash only the first 40 lines
# and let a later ``data=``/headers change ride the baseline key). A bracket that
# never closes within the hard limit is a miscount (a multi-line string the
# single-line blanker cannot mask) or a stray opener, so it is bound only to the
# soft cap and cannot swallow unrelated code.
_MAX_CALL_LINES = 40 # soft cap: how far a NEVER-closing opener is followed
_MAX_CALL_HARD_LINES = 200 # hard cap: how far a closing call is followed to bind it
# Cap a single rendered line. A short line is shown verbatim; a long (e.g.
# minified one-liner) line is shown as a bounded prefix plus a sha256 of the full
# line, so a packed payload cannot dump unbounded content into the evidence and
# baseline while a change past the cutoff still changes the digest and reopens the
# finding. The npm scanner bounds its snippets the same way.
_MAX_LINE_CHARS = 200
# Cap on recorded spans in one evidence string; beyond it the remaining spans are
# folded into a digest so a file with thousands of matching lines cannot build a
# multi-megabyte evidence blob, while an added/removed span past the cap still
# changes the key. Comfortably above the largest real baseline entry.
_MAX_EVIDENCE_SPANS = 96
def _cap_line(code: str) -> str:
"""Bound a single line's displayed code: return it verbatim when short, else a
``_MAX_LINE_CHARS`` prefix plus a digest of the whole line so the tail is still
pinned (fail-closed) without recording the entire line."""
if len(code) <= _MAX_LINE_CHARS:
return code
digest = hashlib.sha256(code.encode("utf-8", "replace")).hexdigest()
return f"{code[:_MAX_LINE_CHARS]} sha256:{digest}"
_PY_TRIPLE = ("'''", '"""')
def _ends_with_odd_backslash(s: str) -> bool:
"""True if ``s`` ends with an odd run of backslashes, i.e. a trailing
backslash that escapes the newline (a string/line continuation) rather than a
literal ``\\\\`` pair."""
return (len(s) - len(s.rstrip("\\"))) % 2 == 1
# Single-line quoted string literal; blanks complete one-line strings (the legacy
# view) so the single-line and multi-line blanked spans can be unioned below.
_RE_STR_LITERAL = re.compile(r"'(?:[^'\\]|\\.)*'|\"(?:[^\"\\]|\\.)*\"")
def _blank_code_strings(lines: list[str]) -> list[str]:
"""Replace string contents (single- and triple-quoted, escapes honoured) with
spaces across ``lines``, keeping the line count and every bracket OUTSIDE a
string intact. Bracket counting then never miscounts a ``)`` that lives inside
a string -- including a triple-quoted string spanning several lines, which a
per-line regex cannot blank."""
out: list[str] = []
in_triple: str | None = None # active ''' or \"\"\" delimiter, or None
in_string: str | None = None # active ' or " continued via a trailing backslash
for line in lines:
buf: list[str] = []
i, n = 0, len(line)
while i < n:
if in_triple is not None:
end = line.find(in_triple, i)
if end == -1:
buf.append(" " * (n - i))
i = n
else:
buf.append(" " * (end - i + 3))
i = end + 3
in_triple = None
continue
if in_string is not None:
# A single-/double-quoted string continued onto this line by a
# backslash-escaped newline. Resume blanking until its closing quote;
# if this line also ends on an odd trailing backslash the string
# continues again, otherwise it closes (or is unterminated) here. A
# per-line regex blanker cannot see this, so a `)` on the
# continuation line would otherwise be counted as code and close the
# call early -- dropping the URL/body lines that follow.
j, closed = i, False
while j < n:
if line[j] == "\\":
j += 2
continue
if line[j] == in_string:
j += 1
closed = True
break
j += 1
buf.append(" " * (min(j, n) - i))
if closed:
in_string = None
i = j
else:
i = n
if not _ends_with_odd_backslash(line):
in_string = None # unterminated without continuation; stop
continue
ch = line[i]
if ch in "'\"":
if line[i : i + 3] in _PY_TRIPLE:
delim = line[i : i + 3]
end = line.find(delim, i + 3)
if end == -1: # opens a triple string that runs past this line
buf.append(" " * (n - i))
in_triple = delim
i = n
else:
buf.append(" " * (end - i + 3))
i = end + 3
continue
j = i + 1 # single-line string; skip to its closing quote
closed = False
while j < n:
if line[j] == "\\":
j += 2
continue
if line[j] == ch:
j += 1
closed = True
break
j += 1
buf.append(" " * (min(j, n) - i))
if closed:
i = j
else:
# Ran off the line without closing: an odd trailing backslash
# escapes the newline and continues the string onto the next
# line, so remember the quote; otherwise it is just unterminated.
i = n
if _ends_with_odd_backslash(line):
in_string = ch
continue
buf.append(ch)
i += 1
out.append("".join(buf))
return out
_RE_BRACKETS = re.compile(r"[()\[\]{}]")
_OPENERS = frozenset("([{")
def _bracket_lr(line: str) -> tuple[int, int]:
"""Order-aware bracket reduction of one already-string-blanked line: ``(L, R)``
where ``L`` is the count of closers with no opener earlier on the line (they
need an opener to the LEFT / a prior line) and ``R`` is the count of openers
with no closer later on the line (they need a closer to the RIGHT / a later
line). A plain net count (opens minus closes) collapses order and so masks a
trailing opener that follows leading closers on the same line, e.g.
``]; requests.post(`` nets to 0 and hides the ``(`` that opens the flagged
call; tracking the running minimum keeps that opener visible so the call's
argument lines still bind. Only bracket characters are walked (pulled out with
one C-level regex pass) so a long minified line stays cheap."""
depth = 0
low = 0
for ch in _RE_BRACKETS.findall(line):
if ch in _OPENERS:
depth += 1
else:
depth -= 1
if depth < low:
low = depth
return -low, depth - low
def _scan_line_end(view: list[str], start: int) -> int:
"""1-based line where the statement at ``start`` closes its brackets in
``view`` (one blanked view of the file). A call that closes is followed to its
real close up to ``_MAX_CALL_HARD_LINES`` so its whole argument list binds; a
bracket that never closes within that hard limit (a stray/miscounted opener) is
bound only to the ``_MAX_CALL_LINES`` soft cap so it cannot swallow the file.
Brackets are applied in order via ``_bracket_lr`` (leading closers clamp at 0)
so a closer that precedes the opener on the same line does not cancel it."""
depth = 0
hard = min(len(view), start + _MAX_CALL_HARD_LINES - 1)
for j in range(start, hard + 1):
ln = view[j - 1]
left, right = _bracket_lr(ln)
depth = max(0, depth - left) + right
if ln.rstrip().endswith("\\"):
continue # explicit backslash continuation: the call (e.g. its `(` and
# URL/body) is on the next physical line, so do not close here
if depth <= 0:
return j
# Never closed within the hard limit: bind only the soft cap so a stray opener
# cannot bind a giant unrelated span.
return min(len(view), start + _MAX_CALL_LINES - 1)
def _logical_line_end(sl_blanked: list[str], ml_blanked: list[str], start: int) -> int:
"""1-based line where the statement opened at ``start`` closes, so a multi-line
call binds its argument lines (a changed URL/body on a continuation line
reopens, not just the API line). Returns the LARGER of the spans found in the
single-line-blanked view (legacy: a payload embedded inside a string still
counts, so its brackets bind the call) and the multi-line-blanked view (a
bracket inside a triple-quoted string argument no longer closes the call
early). Taking the union never shrinks the bound span below either view, so
neither blanking strategy can drop a continuation line a malicious change
relies on."""
return max(_scan_line_end(sl_blanked, start), _scan_line_end(ml_blanked, start))
def _extract_evidence(
content: str,
pattern: re.Pattern,
max_matches: int = 3,
max_matches: int = 0,
) -> str:
"""Pull matching lines as evidence snippets.
"""Pull matching lines as evidence snippets (``max_matches=0`` means all).
Falls back to a whole-content search when the pattern only matches across
line boundaries (several IOC regexes use ``re.DOTALL``). Without this an
anti-analysis / archive-staging finding could report empty evidence, making
the baseline entry impossible to review.
Records every matching line in full, not a truncated sample, so an extra
match (or extra code on a long line) appended to an already-flagged file
changes the evidence and the baseline key instead of riding the first few.
Leading whitespace is kept so a flagged line moved out of a guarded block
reads as changed. Each single-line match is extended over bracket
continuations so a multi-line call binds its argument lines too. Cross-line
matches the per-line scan cannot see (DOTALL IOC regexes, or a multi-line
construct appended under a check that already had a one-line match) are
recorded afterwards, so an added multiline payload reopens the finding. A
pathological greedy span is bounded to its head line plus a digest of the
rest.
"""
lines = content.splitlines()
matches = []
sl_blanked = [_RE_STR_LITERAL.sub("", ln) for ln in lines]
ml_blanked = _blank_code_strings(lines)
out = []
seen: set[tuple[int, int]] = set()
# Overflow is streamed, not buffered: once `out` holds _MAX_EVIDENCE_SPANS
# rendered spans, every further span is folded straight into a running digest
# instead of being materialized and sliced off at the end. On a minified or
# padded file with hundreds of thousands of matching lines that keeps memory
# and work bounded to the display cap rather than the match count, while the
# digest still covers every overflow span so an over-cap payload change
# reopens. The fold reproduces _canon_evidence(" | ".join(overflow)) exactly
# (strip each span to its non-empty L<NN>-less code lines, join with "\n"), so
# the digest is identical to buffering the whole list and canonicalizing once.
overflow_count = 0
overflow_hash = hashlib.sha256()
overflow_started = False
def _emit(rendered: str) -> None:
nonlocal overflow_count, overflow_started
if len(out) < _MAX_EVIDENCE_SPANS:
out.append(rendered)
return
overflow_count += 1
for piece in _RE_EVIDENCE_SPLIT.split(rendered):
piece = _RE_EVIDENCE_PREFIX.sub("", piece, count = 1).rstrip()
if not piece:
continue
if overflow_started:
overflow_hash.update(b"\n")
overflow_hash.update(piece.encode("utf-8", "replace"))
overflow_started = True
def _render(start: int, end: int) -> str:
span = lines[start - 1 : end] or ["<multiline match>"]
if len(span) > _MAX_MULTILINE_LINES:
# Digest the code without the L<NN>: markers so a pure line shift of
# the same span stays stable while a code change still reopens. The
# head is truncated for display only; the span digest already binds
# its full content, so no per-line digest is needed here.
code = "\n".join(ln.rstrip() for ln in span)
digest = hashlib.sha256(code.encode("utf-8", "replace")).hexdigest()
head = span[0].rstrip()
if len(head) > _MAX_LINE_CHARS:
head = head[:_MAX_LINE_CHARS] + "..."
return f"L{start}: {head} sha256:{digest}"
return "\n".join(f"L{start + i}: {_cap_line(ln.rstrip())}" for i, ln in enumerate(span))
for i, line in enumerate(lines, 1):
if pattern.search(line):
snippet = line.strip()
if len(snippet) > 160:
snippet = snippet[:160] + "..."
matches.append(f"L{i}: {snippet}")
if len(matches) >= max_matches:
break
if matches:
return " | ".join(matches)
# Multiline (DOTALL) match: report the line where the match begins.
m = pattern.search(content)
if m:
line_no = content.count("\n", 0, m.start()) + 1
snippet = lines[line_no - 1].strip() if line_no - 1 < len(lines) else ""
if len(snippet) > 160:
snippet = snippet[:160] + "..."
return f"L{line_no}: {snippet}" if snippet else f"L{line_no}: <multiline match>"
return ""
span = (i, _logical_line_end(sl_blanked, ml_blanked, i))
if span in seen:
continue
# Only track spans while still filling the display list: past the cap
# every span is folded into the overflow digest, so growing `seen` with
# all of them would keep memory proportional to the match count (the
# behavior this cap exists to bound) on a generated file with millions
# of one-line matches. The per-line spans are unique by line number, so
# dropping them from `seen` past the cap cannot cause a missed dedup
# here; at worst the fallback re-folds an over-cap span into the same
# digest, which stays deterministic and still reopens on a change.
if len(out) < _MAX_EVIDENCE_SPANS:
seen.add(span)
_emit(_render(*span))
if max_matches and len(out) >= max_matches:
return " | ".join(out)
# Precompute newline offsets once so mapping a match offset to its 1-based line
# is O(log n) (bisect) rather than O(n) (content.count) per match; the latter
# made this fallback quadratic on a minified file with thousands of matches.
nl = [p for p, ch in enumerate(content) if ch == "\n"]
for m in pattern.finditer(content):
start = bisect.bisect_left(nl, m.start()) + 1
end = bisect.bisect_left(nl, m.end()) + 1
if end <= start or (start, end) in seen:
continue # single-line matches are already covered by the pass above
# A giant greedy DOTALL span is bound by the full digest of its content
# (via _render, which renders a >12-line span as a head line plus a sha256
# of the whole span). Binding only the anchors leaves the bridged interior
# unhashed, so an attacker could insert a new cross-line payload (a `/tmp`
# line and a later `subprocess` line, sharing no single line so the
# per-line pass never binds them) between unchanged outer anchors and keep
# the same key. Digesting the interior reopens on any such change; a pure
# line shift stays stable because the digest is over the markerless code.
if len(out) < _MAX_EVIDENCE_SPANS:
seen.add((start, end))
_emit(_render(start, end))
if max_matches and len(out) >= max_matches:
break
if overflow_count:
# The overflow digest was accumulated from the canonicalized (L<NN>:-less)
# spans as they were emitted, so a pure line shift above the overflow
# region does not change it and reopen an otherwise-unchanged finding,
# matching the per-span key's line-shift stability.
out.append(f"(+{overflow_count} more) sha256:{overflow_hash.hexdigest()}")
return " | ".join(out)
def _embedded_key_evidence(content: str) -> str:
"""Key evidence that also pins the full PEM block(s) via a digest, so a key
body swapped under the same BEGIN marker reopens the finding (single-line and
DER keys are already bound by their full matched line)."""
ev = _extract_evidence(content, RE_EMBEDDED_KEYS)
blocks = RE_PEM_BLOCK.findall(content)
if blocks:
digest = hashlib.sha256("\n".join(blocks).encode("utf-8", "replace")).hexdigest()
ev = f"{ev} sha256:{digest}" if ev else f"sha256:{digest}"
return ev
def _blob_digest(content: str) -> tuple[str, str]:
"""First large blob (for display) plus a digest binding EVERY large blob, so
an appended or swapped encoded payload reopens the finding rather than riding
an unchanged first blob. Assumes at least one blob is present (single-blob
files keep the prior single-blob digest, so the baseline does not drift)."""
blobs = RE_LARGE_BLOB.findall(content)
digest = hashlib.sha256("\n".join(blobs).encode("utf-8", "replace")).hexdigest()
return blobs[0], digest
# Non-Python checkers
@ -1189,7 +1546,8 @@ def check_js_file(content: str, filename: str, package: str) -> list[Finding]:
package,
filename,
"JS embeds credential regexes AND makes network calls (stealer)",
_extract_evidence(content, RE_TOKEN_REGEX),
f"Token: {_extract_evidence(content, RE_TOKEN_REGEX)}\n"
f"Network: {_extract_evidence(content, RE_NETWORK)}",
)
)
if has_workflow_inj:
@ -1202,18 +1560,31 @@ def check_js_file(content: str, filename: str, package: str) -> list[Finding]:
_extract_evidence(content, RE_WORKFLOW_INJECT),
)
)
if is_large and not findings:
findings.append(
Finding(
HIGH,
package,
filename,
# Size stays in evidence, not the check label, so the baseline key
# does not drift when a wheel's bundle grows by a few KB.
"Python wheel ships large JS bundle (uncommon; manually review)",
f"{len(content) // 1024} KB JS bundle",
# Pin the whole file's content digest to EVERY JS finding (not just large
# bundles). _extract_evidence blanks only Python string forms before counting
# brackets, so a JS backtick template literal that contains `)` can close a
# call's span early and omit the option/body lines that follow; binding the
# full content means a change to those omitted lines still reopens instead of
# riding the matched-line evidence. A large bundle with no other heuristic is a
# standalone HIGH.
if findings or is_large:
digest = hashlib.sha256(content.encode("utf-8", "replace")).hexdigest()
if findings:
for f in findings:
f.evidence = f"{f.evidence} bundle-sha256:{digest}"
else:
findings.append(
Finding(
HIGH,
package,
filename,
# Size stays out of the check label (from main) so the baseline
# key does not drift when a benign bundle grows; the full-content
# digest below still binds the bytes so a payload swap reopens.
"Python wheel ships large JS bundle (uncommon; manually review)",
f"sha256: {digest}",
)
)
)
return findings
@ -1233,6 +1604,12 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding]
if RE_DEV_TOOL_HIJACK.search(content) and (
RE_NETWORK.search(content) or RE_SUBPROCESS.search(content)
):
# Bind the hook AND the network/exec signal so a changed exfil reopens.
evidence = [f"Hook: {_extract_evidence(content, RE_DEV_TOOL_HIJACK)}"]
if RE_NETWORK.search(content):
evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}")
if RE_SUBPROCESS.search(content):
evidence.append(f"Exec: {_extract_evidence(content, RE_SUBPROCESS)}")
findings.append(
Finding(
CRITICAL,
@ -1240,7 +1617,7 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding]
filename,
"Shell installs developer-tool persistence hook (.bashrc / "
"profile.d / vscode tasks) AND has network or exec",
_extract_evidence(content, RE_DEV_TOOL_HIJACK),
"\n".join(evidence),
)
)
if RE_TOKEN_REGEX.search(content) and RE_NETWORK.search(content):
@ -1250,7 +1627,8 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding]
package,
filename,
"Shell embeds credential regexes AND makes network calls",
_extract_evidence(content, RE_TOKEN_REGEX),
f"Token: {_extract_evidence(content, RE_TOKEN_REGEX)}\n"
f"Network: {_extract_evidence(content, RE_NETWORK)}",
)
)
if RE_WORKFLOW_INJECT.search(content):
@ -2517,9 +2895,9 @@ def _find_requirements_files(root: str) -> list[str]:
# Baseline allowlist: triaged known-good CRITICAL/HIGH findings so the gate can
# enforce without drowning in legitimate-library noise. Matched on
# ``(package, basename(filename), check)`` -- not evidence text -- so a version
# bump does not reopen a finding, but a *new* kind of finding in a listed file
# is a different check and still fails. Regenerate with ``--write-baseline``.
# (package, package-relative file, check, evidence hash); the hash strips
# ``L<NN>:`` markers so version bumps and line shifts do not reopen an entry,
# but changed flagged code does. Regenerate with ``--write-baseline``.
_DEFAULT_BASELINE_PATH = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "scan_packages_baseline.json"
@ -2546,16 +2924,54 @@ def _relpath_in_package(filename: str) -> str:
return _RE_SDIST_ROOT.sub("", filename, count = 1)
def _finding_key(f: Finding) -> tuple[str, str, str]:
"""Stable allowlist key: normalized package, package-relative path, check.
# Evidence joins matched spans with " | " and a newline between labelled groups,
# each span tagged "L<NN>: ". Split only on those real delimiters (a " | " before
# a marker, or a newline), never on a bare "|" -- matched code may contain a
# bitwise-or or union type. The prefix strips only a genuine leading marker, an
# optional "Label: " then "L<NN>: "; a marker-like "L<NN>:" inside raw code (e.g.
# a .pth import line) has no leading marker and is left intact.
_RE_EVIDENCE_SPLIT = re.compile(r" \| (?=L\d+:)|\n")
_RE_EVIDENCE_PREFIX = re.compile(r"^(?:[A-Za-z][A-Za-z0-9 _/+.-]*:\s*)?L\d+:\s?")
The package-relative path (not just basename) keeps the key stable across
version bumps while still distinguishing same-named files like ``utils.py``.
def _canon_evidence(evidence: str) -> str:
"""Matched code lines in discovery order (markers removed), duplicates kept.
Splits evidence on its real span delimiters, drops each span's leading
label / line-number marker, and keeps the code with its indentation. Line
shifts are absorbed by stripping the L<NN>: markers, not by sorting, so order
stays significant: reordering matched lines (executable context, e.g. the
arguments of a multi-line call) reopens the finding. Keeping duplicates means
an appended identical occurrence still changes the key."""
spans = []
for s in _RE_EVIDENCE_SPLIT.split(evidence or ""):
s = _RE_EVIDENCE_PREFIX.sub("", s, count = 1).rstrip()
if s:
spans.append(s)
return "\n".join(spans)
def _evidence_hash(evidence: str) -> str:
"""Stable digest of the canonical matched evidence."""
return hashlib.sha256(_canon_evidence(evidence).encode("utf-8", "replace")).hexdigest()
def _finding_key(f: Finding) -> tuple[str, str, str, str]:
"""Allowlist key: package, package-relative path, check, evidence hash.
The evidence hash is over the set of matched code, so the key survives version
bumps, line shifts and reordering but reopens when the flagged code changes --
so a future payload in a baselined file/check is not auto-suppressed.
"""
return (_norm_pkg(f.package), _relpath_in_package(f.filename), f.check)
return (
_norm_pkg(f.package),
_relpath_in_package(f.filename),
f.check,
_evidence_hash(f.evidence),
)
def _load_baseline(path: str) -> set[tuple[str, str, str]]:
def _load_baseline(path: str) -> set[tuple[str, str, str, str]]:
"""Load an allowlist JSON into a set of match keys. Missing file -> empty."""
try:
with open(path, "r", encoding = "utf-8") as fh:
@ -2565,19 +2981,47 @@ def _load_baseline(path: str) -> set[tuple[str, str, str]]:
except (OSError, json.JSONDecodeError) as exc:
print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr)
return set()
keys: set[tuple[str, str, str]] = set()
for e in data.get("entries", []):
if not isinstance(data, dict):
print(f" [WARN] baseline {path} is not a JSON object", file = sys.stderr)
return set()
entries = data.get("entries", [])
if not isinstance(entries, list):
print(f" [WARN] baseline {path} entries is not a list", file = sys.stderr)
return set()
keys: set[tuple[str, str, str, str]] = set()
legacy = 0
for e in entries:
if not isinstance(e, dict):
continue
try:
keys.add((_norm_pkg(e["package"]), _relpath_in_package(e["file"]), e["check"]))
# Use the reviewed hash; else recompute it from the stored evidence.
evidence_hash = e.get("evidence_hash") or _evidence_hash(e.get("evidence") or "")
if not e.get("evidence_hash"):
legacy += 1
keys.add(
(
_norm_pkg(e["package"]),
_relpath_in_package(e["file"]),
e["check"],
evidence_hash,
)
)
except (KeyError, TypeError):
continue
if legacy:
print(
f" [WARN] baseline {path}: {legacy} entries lack evidence_hash and may "
f"not suppress until regenerated with --write-baseline (findings reopen "
f"rather than risk hiding changed code under a coarse key)",
file = sys.stderr,
)
return keys
def _write_baseline(path: str, findings: list[Finding]) -> None:
"""Persist CRITICAL/HIGH findings as an allowlist for human triage."""
entries = []
seen: set[tuple[str, str, str]] = set()
seen: set[tuple[str, str, str, str]] = set()
for f in sorted(findings, key = lambda f: SEVERITY_ORDER.get(f.severity, 99)):
if f.severity not in (CRITICAL, HIGH):
continue
@ -2591,15 +3035,18 @@ def _write_baseline(path: str, findings: list[Finding]) -> None:
"file": _relpath_in_package(f.filename),
"check": f.check,
"severity": f.severity,
"evidence": f.evidence[:240],
"evidence": f.evidence,
"evidence_hash": _evidence_hash(f.evidence),
}
)
doc = {
"_comment": (
"scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding "
"manually judged benign. Matched on (package, package-relative file, "
"check); evidence/severity are for review only. Regenerate with "
"--write-baseline AFTER reviewing every line."
"check, evidence_hash); evidence_hash is over the matched code with "
"L<NN>: markers stripped, so version bumps and line shifts do not "
"reopen an entry but changed code does. severity and evidence are for "
"review only. Regenerate with --write-baseline AFTER reviewing every line."
),
"version": 1,
"entries": entries,
@ -2611,7 +3058,7 @@ def _write_baseline(path: str, findings: list[Finding]) -> None:
def _partition_baseline(
findings: list[Finding], baseline: set[tuple[str, str, str]]
findings: list[Finding], baseline: set[tuple[str, str, str, str]]
) -> tuple[list[Finding], list[Finding]]:
"""Split findings into (active, suppressed) by allowlist membership."""
if not baseline:

File diff suppressed because one or more lines are too long

View file

@ -489,6 +489,9 @@ class DiffusionBackend:
model_kind: Optional[str] = None,
) -> dict[str, Any]:
"""Validate, then run the (slow) load on a daemon thread. Returns at once."""
# A blank token (the Studio default when none is configured) must mean
# "anonymous", not an explicit empty credential the Hub rejects with 401.
hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None
fam = self.validate_load_request(
repo_id,
gguf_filename = gguf_filename,
@ -617,6 +620,18 @@ class DiffusionBackend:
fraction = min(downloaded / expected, 1.0) if expected > 0 else 0.0
return _progress("downloading", downloaded, expected, fraction)
def loading_repo_ids(self) -> tuple[str, ...]:
"""Repo ids an in-flight background load is downloading (empty when idle).
The delete-cached guard needs this: during a load ``status()["loaded"]`` is
still False, but deleting the target repo (or its companion base) would yank
blobs and snapshot files from under the download/assembly."""
with self._lock:
loading = self._loading
if loading is None or loading.error is not None:
return ()
return tuple(r for r in (loading.repo_id, loading.base_repo) if r)
@staticmethod
def _estimate_download_bytes(
repo_id: str,
@ -702,7 +717,10 @@ class DiffusionBackend:
_load_token: Optional[int] = None,
) -> dict[str, Any]:
# Validate first (cheap, no torch/diffusers) so a direct call with a bad
# family fails with ValueError even in a no-diffusers runtime.
# family fails with ValueError even in a no-diffusers runtime. Sanitize the
# token here too (direct callers bypass begin_load): a blank string must
# load anonymously, not 401 as an explicit empty credential.
hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None
fam = self.validate_load_request(
repo_id,
gguf_filename = gguf_filename,
@ -726,12 +744,18 @@ class DiffusionBackend:
# The cancel makes that wait ~one step (or the rest of the denoise for a
# pipeline that ignores the step callback).
with self._lock:
# Bail BEFORE signalling any cancel if this load was already superseded (an
# unload/eviction or a newer load bumped the token while we were resolving /
# downloading). Otherwise a stale worker would abort an unrelated, still-live
# generation from the CURRENT model and only then discover it has nothing to do.
if _load_token is not None and _load_token != self._load_token:
raise RuntimeError("Diffusion load was cancelled.")
if self._active_generate_cancel is not None:
self._active_generate_cancel.set()
with self._generate_lock:
with self._lock:
# Bail before the (slow, VRAM-heavy) build if an unload/eviction or a
# newer load superseded this one while we were resolving/downloading.
# Re-check under the generate lock: a newer load/unload may have superseded
# this one while we waited for the in-flight denoise to exit.
if _load_token is not None and _load_token != self._load_token:
raise RuntimeError("Diffusion load was cancelled.")
@ -800,6 +824,12 @@ class DiffusionBackend:
)
pipe = None
transformer_quant_engaged = None
# Drop the exception (and its traceback) BEFORE clearing the cache:
# exc.__traceback__ keeps _load_dense_quant_pipeline's frame -- and
# thus its partially-built dense bf16 transformer/pipe -- alive, so
# clear_gpu_cache() could not otherwise reclaim that VRAM before the
# GGUF build (the OOM-fallback path this cleanup exists for).
del exc
clear_gpu_cache()
if pipe is None:
@ -1363,6 +1393,26 @@ class DiffusionBackend:
raise ValueError(f"Failed to apply LoRA: {exc}") from exc
pipe._unsloth_loras = desired
@staticmethod
def _reset_step_cache(pipe: Any) -> None:
"""Clear the transformer's stateful step cache (FBCache) before a generation.
diffusers keys FBCache residuals by cache context ("cond"/"uncond") on the
long-lived transformer, and neither the pipeline nor the context exit resets
them (``StateManager`` only clears via ``reset_stateful_hooks``, which no
pipeline calls). This backend reuses one resident pipe across generations, so
without a reset the next generation's first step compares its first-block
residual against the PREVIOUS request's -- a tensor-shape mismatch when the
resolution/batch changed, or a stale-cache reuse otherwise. Best-effort: a
transformer without the hook (uncached load) is a silent no-op."""
transformer = getattr(pipe, "transformer", None)
reset = getattr(transformer, "reset_stateful_hooks", None)
if callable(reset):
try:
reset()
except Exception: # noqa: BLE001 — reset is best-effort, never fail a generation
pass
def generate(
self,
*,
@ -1671,6 +1721,13 @@ class DiffusionBackend:
if "callback_on_step_end" in call_params:
kwargs["callback_on_step_end"] = _on_step
# Start each generation from a clean step cache: FBCache residuals from
# a prior request on this resident pipe would otherwise be compared
# against this generation's first step (shape mismatch on a resolution/
# batch change, or stale reuse). No-op when no cache is engaged.
if state.transformer_cache:
self._reset_step_cache(state.pipe)
self._gen = gen
try:
# inference_mode is strictly faster than the no_grad diffusers

View file

@ -434,20 +434,20 @@ def apply_memory_plan(
# is in the low-VRAM situation where the decode-time spike can OOM, so turn VAE
# tiling on now (if not already engaged) to cap it.
nonlocal tiling_engaged
pipe.enable_model_cpu_offload()
pipe.enable_model_cpu_offload(device = device)
if not tiling_engaged:
tiling_engaged = _enable_vae_saver(pipe, "enable_vae_tiling", "enable_tiling", logger)
policy = plan.offload_policy
if policy == OFFLOAD_MODEL:
pipe.enable_model_cpu_offload()
pipe.enable_model_cpu_offload(device = device)
elif policy == OFFLOAD_GROUP:
if not _apply_group_offload(pipe, device, logger):
_fallback_to_model_offload()
policy = OFFLOAD_MODEL
elif policy == OFFLOAD_SEQUENTIAL:
try:
pipe.enable_sequential_cpu_offload()
pipe.enable_sequential_cpu_offload(device = device)
except Exception as exc: # noqa: BLE001 — keep the model loadable
if logger is not None:
logger.warning(

View file

@ -193,6 +193,15 @@ def load_prequantized_transformer(
transformer.load_state_dict(state_dict, strict = True, assign = True)
transformer = transformer.to(device)
# Built via from_config (not from_pretrained), so it starts in TRAIN mode; the
# dense and GGUF paths load through from_pretrained, which diffusers documents as
# returning an eval()'d module. Match that here so any train/eval-sensitive layer
# (e.g. dropout) can't make prequant inference nondeterministic or diverge from
# the other load paths.
try:
transformer.eval()
except Exception: # noqa: BLE001 — eval() is best-effort
pass
try: # diagnostic marker, mirrors the runtime-quant path
transformer._unsloth_runtime_quant = scheme
except Exception: # noqa: BLE001 — marker is best-effort

View file

@ -207,8 +207,15 @@ def build_sd_cpp_command(
"""
if not files.diffusion_model:
raise ValueError("diffusion_model path is required")
if not str(params.prompt).strip():
# ``(prompt or "")`` so a None prompt is rejected here rather than slipping past
# ``str(None)`` == "None" (truthy) and landing in argv as a literal "None".
if not (params.prompt or "").strip():
raise ValueError("prompt is required")
# sd-cli inpaint needs the source image too: a --mask with no --init-img is an
# invalid invocation (sd-cli has nothing to inpaint into), so reject it here with a
# clear error instead of emitting a command that fails deep in sd-cli.
if params.mask and not params.init_img:
raise ValueError("init_img is required when mask is set (inpaint needs a source image)")
cmd: list[str] = [binary, "--mode", DEFAULT_MODE, "--diffusion-model", files.diffusion_model]
for flag, value in (

View file

@ -457,6 +457,16 @@ class SdCppDiffusionBackend:
fraction = min(downloaded / expected, 1.0) if expected > 0 else 0.0
return _progress("downloading", downloaded, expected, fraction)
def loading_repo_ids(self) -> tuple[str, ...]:
"""Repo ids an in-flight background load is downloading (empty when idle).
Mirrors the diffusers backend so the delete-cached guard can query whichever
engine is active without caring which one it got."""
with self._lock:
loading = self._loading
if loading is None or loading.error is not None:
return ()
return tuple(r for r in (loading.repo_id, loading.base_repo) if r)
# ── Generate ───────────────────────────────────────────────────────────
def generate(

View file

@ -1878,8 +1878,11 @@ class DiffusionGenerateRequest(BaseModel):
)
steps: int = Field(9, ge = 1, le = 100, description = "Number of denoising steps")
guidance: float = Field(0.0, ge = 0.0, le = 20.0, description = "Classifier-free guidance scale")
# le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript
# rounds integers above Number.MAX_SAFE_INTEGER -- a restored recipe would then
# generate a different image. Random seeds are already masked to this range.
seed: Optional[int] = Field(
None, ge = 0, le = 2**64 - 1, description = "Seed for reproducibility (random if omitted)"
None, ge = 0, le = 2**53 - 1, description = "Seed for reproducibility (random if omitted)"
)
batch_size: int = Field(
1, ge = 1, le = 32, description = "Images generated in one forward pass (VRAM-heavy)"

View file

@ -27,6 +27,10 @@ class CachedModelRepo(BaseModel):
repo_id: str
size_bytes: int
last_modified: Optional[float] = None
# "text-to-image" for cached diffusers image repos; response_model would silently
# drop the value the handler sets, letting image-only repos pass the chat picker's
# task gate.
task: Optional[str] = None
class CachedModelsResponse(BaseModel):
@ -3399,8 +3403,13 @@ async def delete_cached_model(
# delete guard is otherwise chat-only, so its GGUF could be removed from
# under a live pipeline. Repo-level match, like the chat guards above.
try:
from core.inference.diffusion import get_diffusion_backend
diffusion_status = get_diffusion_backend().status()
# The ACTIVE engine (diffusers or native sd_cpp): on a native selection the
# diffusers singleton reports unloaded while sd-cli still generates from the
# cached GGUF, so checking it alone would let the files be deleted mid-use.
from core.inference.diffusion_engine_router import get_active_diffusion_engine
engine = get_active_diffusion_engine()
diffusion_status = engine.status()
if diffusion_status.get("loaded") and diffusion_status.get("repo_id"):
loaded_id = str(diffusion_status["repo_id"]).lower()
if loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()):
@ -3408,6 +3417,17 @@ async def delete_cached_model(
status_code = 400,
detail = "Unload the model before deleting",
)
# Also refuse while a background image load is DOWNLOADING this repo (or its
# companion base): status().loaded is still False in that window, but deleting
# would remove blobs from under the in-flight download/assembly.
loading_ids = getattr(engine, "loading_repo_ids", tuple)()
for lid in loading_ids:
lid = str(lid).lower()
if lid == repo_id.lower() or lid.startswith(repo_id.lower()):
raise HTTPException(
status_code = 400,
detail = "An Images model load is using this repo; wait for it to finish",
)
except HTTPException:
raise
except Exception:

View file

@ -372,9 +372,14 @@ async def start_training(
# release the arbiter so it doesn't think the gone pipeline owns
# the GPU. Must precede the chat block, which early-returns.
from core.inference import gpu_arbiter
from core.inference.diffusion import get_diffusion_backend
from core.inference.diffusion_engine_router import (
get_active_diffusion_engine,
)
diffusion = get_diffusion_backend()
# The ACTIVE engine, not the diffusers singleton: on a native
# (sd_cpp) selection the diffusers backend reports unloaded while
# the native engine still holds model state / a live generation.
diffusion = get_active_diffusion_engine()
if diffusion.is_loaded:
logger.info(
"Unloading diffusion (Images) model to free GPU memory for training"

View file

@ -170,11 +170,13 @@ class _FakePipe:
self.moved_to = device
return self
def enable_model_cpu_offload(self) -> None:
def enable_model_cpu_offload(self, device = None) -> None:
self.offloaded = True
self.offload_device = device
def enable_sequential_cpu_offload(self) -> None:
def enable_sequential_cpu_offload(self, device = None) -> None:
self.sequential_offloaded = True
self.offload_device = device
def enable_vae_tiling(self) -> None:
self.vae_tiled = True
@ -1149,6 +1151,29 @@ def test_unload_cancels_in_flight_load(fake_runtime):
)
def test_superseded_load_does_not_cancel_live_generation(fake_runtime):
# A superseded background load (its token was bumped by a newer load/unload) that
# finally reaches load_pipeline must bail WITHOUT signalling the current model's
# in-flight generation: the token check has to run before the cancel is set, or a
# stale worker aborts an unrelated, still-live denoise.
import threading as _threading
backend = DiffusionBackend()
fam = detect_family("unsloth/Z-Image-Turbo-GGUF")
live_cancel = _threading.Event()
backend._active_generate_cancel = live_cancel # a generation from the CURRENT model
token = 11
backend._load_token = token + 1 # this load has already been superseded
with pytest.raises(RuntimeError, match = "cancelled"):
backend.load_pipeline(
"unsloth/Z-Image-Turbo-GGUF",
gguf_filename = "z-image-turbo-Q4_K_S.gguf",
base_repo = fam.base_repo,
_load_token = token,
)
assert not live_cancel.is_set() # the live generation was left untouched
def test_pick_dtype_bf16_only_on_ampere(fake_runtime, monkeypatch):
# BF16 only on Ampere+ (cc >= 8); pre-Ampere cards must fall back to FP16.
torch = sys.modules["torch"]
@ -1798,3 +1823,41 @@ def test_transformer_quant_skipped_when_plan_offloads(fake_runtime, tmp_path, mo
assert status["transformer_quant"] is None
assert status["offload_policy"] == "model"
assert _FakeTransformer.last["path"] # GGUF path used
def test_reset_step_cache_helper_is_best_effort():
# Calls the transformer's reset hook when present.
calls = []
pipe = types.SimpleNamespace(
transformer = types.SimpleNamespace(reset_stateful_hooks = lambda: calls.append(True))
)
DiffusionBackend._reset_step_cache(pipe)
assert calls == [True]
# No transformer, or a transformer without the hook -> silent no-op (never raises).
DiffusionBackend._reset_step_cache(types.SimpleNamespace())
DiffusionBackend._reset_step_cache(types.SimpleNamespace(transformer = object()))
def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path):
# FBCache residuals live on the resident transformer across generations, so each
# generate() must reset the stateful cache first -- but only when a cache is engaged.
(tmp_path / "model.gguf").write_bytes(b"weights")
backend = DiffusionBackend()
backend.load_pipeline(
str(tmp_path),
gguf_filename = "model.gguf",
base_repo = "base/repo",
family_override = "z-image",
)
resets = []
backend._state.pipe.transformer = types.SimpleNamespace(
reset_stateful_hooks = lambda: resets.append(True)
)
# No cache engaged (transformer_cache is None) -> reset must NOT run.
backend.generate(prompt = "a sloth")
assert resets == []
# Engage a cache; every subsequent generation resets the stateful cache first.
object.__setattr__(backend._state, "transformer_cache", "fbcache")
backend.generate(prompt = "a sloth")
backend.generate(prompt = "another sloth")
assert resets == [True, True]

View file

@ -327,16 +327,19 @@ def test_snapshot_never_raises_on_probe_failure(monkeypatch):
class _RecordingPipe:
def __init__(self) -> None:
self.calls: list[str] = []
self.offload_device = None
def to(self, device):
self.calls.append(f"to:{device}")
return self
def enable_model_cpu_offload(self):
def enable_model_cpu_offload(self, device = None):
self.calls.append("model_offload")
self.offload_device = device
def enable_sequential_cpu_offload(self):
def enable_sequential_cpu_offload(self, device = None):
self.calls.append("sequential_offload")
self.offload_device = device
def enable_vae_tiling(self):
self.calls.append("vae_tiling")
@ -385,6 +388,16 @@ def test_apply_model_offload_engages_offload_and_tiling():
assert "to:cuda" not in pipe.calls # offload owns placement; never both
assert "vae_tiling" in pipe.calls and "vae_slicing" in pipe.calls
assert effective == OFFLOAD_MODEL and tiled is True
assert pipe.offload_device == "cuda" # device threaded to enable_model_cpu_offload
def test_apply_model_offload_passes_target_device():
# enable_model_cpu_offload defaults to CUDA in diffusers; on a non-CUDA accelerator
# (e.g. Intel XPU, which this backend supports) the target device must be forwarded
# or diffusers offloads to the wrong backend and the load fails.
pipe = _RecordingPipe()
apply_memory_plan(pipe, _plan(OFFLOAD_MODEL, tiling = False), device = "xpu")
assert pipe.offload_device == "xpu"
def test_apply_vae_tiling_falls_back_to_vae_submodule():
@ -404,7 +417,7 @@ def test_apply_vae_tiling_falls_back_to_vae_submodule():
def _slice(self):
self.vae.sliced = True
def enable_model_cpu_offload(self):
def enable_model_cpu_offload(self, device = None):
self.offloaded = True
pipe = _VaeOnly()
@ -439,13 +452,14 @@ def test_apply_sequential_offload():
)
assert "sequential_offload" in pipe.calls and "to:cuda" not in pipe.calls
assert effective == OFFLOAD_SEQUENTIAL
assert pipe.offload_device == "cuda" # device threaded to sequential offload too
def test_apply_sequential_falls_back_to_model_offload_when_unsupported():
# Sequential offload is unreliable for GGUF on some diffusers versions; the
# applier must fall back to whole-module offload and report what actually ran.
class _NoSeqPipe(_RecordingPipe):
def enable_sequential_cpu_offload(self):
def enable_sequential_cpu_offload(self, device = None):
raise RuntimeError("sequential offload not supported for this transformer")
pipe = _NoSeqPipe()

View file

@ -66,6 +66,7 @@ class _FakeTransformer:
def __init__(self):
self.assigned = None
self.moved = None
self.eval_called = False
@classmethod
def load_config(cls, base, **kw):
@ -101,6 +102,10 @@ class _FakeTransformer:
self.moved = device
return self
def eval(self):
self.eval_called = True
return self
def _stub_torch_accelerate(
monkeypatch,
@ -182,6 +187,14 @@ def test_load_meta_init_and_assign(monkeypatch, tmp_path):
assert t._unsloth_runtime_quant == "fp8"
def test_load_puts_transformer_in_eval_mode(monkeypatch, tmp_path):
# Built via from_config (not from_pretrained), so the loader must eval() it to match
# the dense/GGUF paths; otherwise train-mode dropout makes inference nondeterministic.
t = _load(monkeypatch, tmp_path, _good_ckpt())
assert t is not None
assert t.eval_called is True
def test_load_missing_file_is_none(monkeypatch, tmp_path):
assert _load(monkeypatch, tmp_path, _good_ckpt(), exists = False) is None

View file

@ -222,6 +222,25 @@ def test_build_inpaint_adds_mask():
assert _pair(cmd, "--init-img") == "/in/src.png"
def test_build_inpaint_mask_without_init_img_rejected():
# sd-cli inpaint needs a source image; a --mask with no --init-img is invalid argv,
# so the builder must reject it up front instead of emitting a doomed command.
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
params = SdCppGenParams(prompt = "x", mask = "/in/mask.png")
with pytest.raises(ValueError, match = "init_img is required"):
build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
def test_build_rejects_none_prompt():
# A None prompt must be rejected, not coerced to the literal string "None" and
# forwarded into argv.
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
with pytest.raises(ValueError, match = "prompt is required"):
build_sd_cpp_command(
"/bin/sd-cli", files, SdCppGenParams(prompt = None), output_path = "/o.png"
)
def test_build_edit_repeats_ref_image():
files = SdCppModelFiles(diffusion_model = "/m/flux.gguf")
params = SdCppGenParams(prompt = "add a hat", ref_images = ("/r/a.png", "/r/b.png"))

View file

@ -1203,10 +1203,7 @@ export function AppSidebar() {
icon={PaintBrush02Icon}
label={t("shell.navigation.images")}
active={pathname === "/images" || pathname.startsWith("/images/")}
disabled={chatOnly}
tooltip={trainExportDisabledHint}
onClick={() => {
if (chatOnly) return;
navigate({ to: "/images" });
closeMobileIfOpen();
}}

View file

@ -1256,13 +1256,15 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
}, [refreshStatus, dismissLoadToast, pollLoadProgress]);
const handleLoad = useCallback(
// Resolves true when the background load STARTED (callers may revert
// optimistic picker state on false); poll outcomes are handled internally.
async (
repoId: string,
opts: {
kind: "gguf" | "single_file" | "pipeline";
filename?: string;
},
) => {
): Promise<boolean> => {
// Cancel any prior poll loop so two can't run at once.
if (pollTimer.current) clearTimeout(pollTimer.current);
setBusy("loading");
@ -1296,9 +1298,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
toast.error(err instanceof Error ? err.message : "Failed to start load");
setBusy(null);
void refreshStatus();
return;
return false;
}
void pollLoadProgress();
return true;
},
[
pollLoadProgress,
@ -1346,13 +1349,19 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
void handleLoad(id, { kind: spec.kind, filename: spec.filename });
return;
}
// GGUF quant pick from the variant expander.
// GGUF quant pick from the variant expander. Optimistic for instant picker
// feedback, but revert if the load fails to START (400/409/network): the
// selector must not advertise a quant that is not the loaded one. Poll-phase
// failures re-sync via refreshStatus.
if (meta.ggufVariant && meta.ggufFilename) {
const prevQuant = quant;
setQuant(meta.ggufVariant);
const dq = defaultsFor(id);
setSteps(dq.steps);
setGuidance(dq.guidance);
void handleLoad(id, { kind: "gguf", filename: meta.ggufFilename });
void handleLoad(id, { kind: "gguf", filename: meta.ggufFilename }).then((started) => {
if (!started) setQuant(prevQuant);
});
return;
}
// A direct local .gguf file picked without a variant isn't wired for Images.
@ -1370,7 +1379,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
setGuidance(d.guidance);
void handleLoad(id, { kind: "pipeline" });
},
[busy, handleLoad],
[busy, handleLoad, quant],
);
const handleUnload = useCallback(async () => {
@ -1530,7 +1539,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
height: h,
steps,
guidance,
seed: baseSeed + i,
// Offset runs by the batch size: the native engine seeds image j of a
// run at seed+j, so a +1 run offset would regenerate the previous run's
// batch-mates. Unique per image on both engines, reproducible via recipes.
seed: baseSeed + i * batchSize,
batch_size: batchSize,
// Transform/Inpaint/Extend send the source image (+ mask for inpaint/extend) and
// a denoise strength, resolved above. The backend derives output size from the