Studio: preserve foreign gallery files, force safetensors on remote ControlNets, and close dataset/seed/GPU gaps
Gallery clear/delete now scope to Studio-owned files: image_gallery and video_gallery skip PNGs / MP4s without a readable recipe (a hand-dropped or orphan file the listing already hides), so clear() and a guessed-id delete no longer destroy files the gallery never surfaced. Remote ControlNets now force use_safetensors: a bare owner/name reaches from_pretrained without the base trust gate, and the Hub scan fails open when unavailable, so requiring safetensors closes the pickle deserialization vector. POSIX uninstall now stops resident sd-server / sd-cli under an owned sd.cpp root before removing the tree (marker-gated), mirroring the Windows stop-before-delete scan; a live native server no longer survives unlinking its binary. Diffusion dataset containment: the training-start read path and the discovery picker route bare names through the protected resolver, so a symlinked dataset is rejected / not advertised like the caption/delete routes already do. Uploads gain the inference decode guard (oversized real images 400 before OOMing the trainer) and dataset upload/caption/delete/import are blocked with 409 while a diffusion run is active. JSONL readers (trainer + routes) tolerate non-object JSON and invalid UTF-8 instead of raising AttributeError / 500. LoRA family compatibility is enforced in the shared resolver, not only the picker, so a direct API client cannot apply a mismatched-family adapter. GPU arbiter gains release_if so the image/video unload idle-check and release are atomic against a concurrent same-owner load's registration. Native batch recipes persist the base batch_seed and restore replays from it, so a native batch_index>0 image no longer advances its seed twice. FLUX.2-klein selects its sd.cpp text encoder by variant (4B -> Qwen3-4B, 9B -> Qwen3-8B) instead of the single family default.
This commit is contained in:
parent
b127256eb4
commit
5eef2f4003
23 changed files with 548 additions and 33 deletions
|
|
@ -35,6 +35,37 @@ _pkill_escape() {
|
|||
printf '%s' "$1" | sed -e 's:[][\\.^$*+?{|}()/]:\\&:g'
|
||||
}
|
||||
|
||||
# Owned sd.cpp roots (default + custom siblings), each gated on the install-time
|
||||
# owner marker. Native diffusion builds beside a custom/env root at
|
||||
# <parent>/stable-diffusion.cpp (find_sd_cpp_binary resolves from
|
||||
# UNSLOTH_STUDIO_HOME.parent) and at $HOME/.unsloth/stable-diffusion.cpp by default.
|
||||
# The marker is mandatory so we never stop a user-managed sd-server from an
|
||||
# unrelated checkout that happens to sit at one of these paths.
|
||||
_owned_sd_cpp_roots() {
|
||||
_default_sd="$HOME/.unsloth/stable-diffusion.cpp"
|
||||
[ -f "$_default_sd/.unsloth-studio-owned" ] && printf '%s\n' "$_default_sd"
|
||||
_custom_studio_roots 2>/dev/null | while IFS= read -r _root; do
|
||||
[ -n "$_root" ] || continue
|
||||
_sd_root="$(dirname "$_root")/stable-diffusion.cpp"
|
||||
[ -f "$_sd_root/.unsloth-studio-owned" ] && printf '%s\n' "$_sd_root"
|
||||
done
|
||||
}
|
||||
|
||||
# pkill resident sd-server / sd-cli whose executable lives under an owned sd.cpp
|
||||
# root, BEFORE that tree is removed below: a live native server keeps running
|
||||
# after its binary is unlinked. Anchored on the owned root so an unrelated
|
||||
# checkout's sd-server is never matched.
|
||||
_stop_owned_sd_cpp_processes() {
|
||||
_signal="$1"
|
||||
command -v pkill >/dev/null 2>&1 || return 0
|
||||
_owned_sd_cpp_roots | while IFS= read -r _root; do
|
||||
[ -n "$_root" ] || continue
|
||||
[ -d "$_root" ] || continue
|
||||
_re=$(_pkill_escape "$_root")
|
||||
pkill "-$_signal" -f "^${_re}/([^ ]*/)?sd-(server|cli)( |\$)" 2>/dev/null || true
|
||||
done
|
||||
}
|
||||
|
||||
_pkill_studio() {
|
||||
# Prefer PID files written by _spawn_terminal so we only touch our own installs.
|
||||
for _data_dir in "$HOME/.local/share/unsloth" $(_custom_studio_data_dirs); do
|
||||
|
|
@ -80,6 +111,12 @@ $_roots_from_conf"
|
|||
pkill -KILL -f "$_pat" 2>/dev/null || true
|
||||
done
|
||||
done
|
||||
|
||||
# Native diffusion servers (sd-server / sd-cli) survive unlinking their binary,
|
||||
# so stop the ones under an owned sd.cpp root before those trees are removed.
|
||||
_stop_owned_sd_cpp_processes TERM
|
||||
sleep 0.5
|
||||
_stop_owned_sd_cpp_processes KILL
|
||||
}
|
||||
|
||||
_remove_path() {
|
||||
|
|
|
|||
|
|
@ -1908,7 +1908,11 @@ class DiffusionBackend:
|
|||
# resolve_controlnet accepts a bare owner/name without the base trust gate, and
|
||||
# from_pretrained deserializes it (a malicious pickle would execute), so run the same
|
||||
# Hub malware preflight the chat/export loaders use. A local dir is exempt (fail-open).
|
||||
if not getattr(resolved_cn, "is_local", False):
|
||||
# The preflight fails OPEN when the Hub scan is unavailable (offline / missing metadata),
|
||||
# so for a remote repo also force safetensors below: that closes the pickle RCE vector
|
||||
# even when the scan could not run.
|
||||
remote_cn = not getattr(resolved_cn, "is_local", False)
|
||||
if remote_cn:
|
||||
from utils.security import evaluate_file_security
|
||||
_cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None)
|
||||
if _cn_fs.blocked:
|
||||
|
|
@ -1924,10 +1928,18 @@ class DiffusionBackend:
|
|||
# state.dtype is the display string ("bfloat16"), not a torch.dtype; pass the real
|
||||
# dtype so diffusers loads at the base compute dtype, not float32 (extra VRAM).
|
||||
cn_dtype = getattr(torch, str(state.dtype).replace("torch.", ""), None)
|
||||
# Force safetensors for an untrusted remote repo: a bare owner/name reaches here without
|
||||
# the base trust gate, and if the Hub scan failed open above, an embedded pickle would
|
||||
# still deserialize on load. Requiring safetensors refuses that vector (curated
|
||||
# ControlNets are all safetensors). A local dir the user chose is exempt.
|
||||
cn_from_pretrained_kwargs: dict[str, Any] = {}
|
||||
if remote_cn:
|
||||
cn_from_pretrained_kwargs["use_safetensors"] = True
|
||||
cn_model = getattr(diffusers, model_cls_name).from_pretrained(
|
||||
resolved_cn.path,
|
||||
torch_dtype = cn_dtype,
|
||||
token = state.hf_token or None, # blank -> anonymous
|
||||
**cn_from_pretrained_kwargs,
|
||||
)
|
||||
if cancel.is_set():
|
||||
# An unload raced the blocking download; bail BEFORE placement so we don't
|
||||
|
|
@ -2029,7 +2041,12 @@ class DiffusionBackend:
|
|||
"for GGUF models."
|
||||
)
|
||||
|
||||
resolved = diffusion_lora.resolve_specs(specs, hf_token = state.hf_token, cancel_event = cancel)
|
||||
resolved = diffusion_lora.resolve_specs(
|
||||
specs,
|
||||
family = getattr(state.family, "name", None),
|
||||
hf_token = state.hf_token,
|
||||
cancel_event = cancel,
|
||||
)
|
||||
# diffusers load_lora_weights takes safetensors only; reject a .gguf adapter as a clean 400.
|
||||
bad = [r.id for r in resolved if r.fmt != "safetensors"]
|
||||
if bad:
|
||||
|
|
|
|||
|
|
@ -470,6 +470,37 @@ def family_sd_cpp_supported(fam: DiffusionFamily) -> bool:
|
|||
return bool(fam.sd_cpp_vae and fam.sd_cpp_text_encoders)
|
||||
|
||||
|
||||
# FLUX.2-klein ships two variants that use DIFFERENT text encoders: the 4B transformer pairs with
|
||||
# Qwen3-4B and the 9B transformer pairs with Qwen3-8B (Black Forest Labs model cards; a mismatched
|
||||
# encoder fails with tensor-shape / dtype errors deep in sd-cli). The family table carries only one
|
||||
# default, so the native (sd.cpp) encoder is selected per variant from the load identity below.
|
||||
_FLUX2_KLEIN_9B_SD_CPP_TEXT_ENCODERS = (
|
||||
(
|
||||
"Comfy-Org/vae-text-encorder-for-flux-klein-9b",
|
||||
"split_files/text_encoders/qwen_3_8b.safetensors",
|
||||
"llm",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def sd_cpp_text_encoders_for(
|
||||
fam: DiffusionFamily,
|
||||
repo_id: Optional[str] = None,
|
||||
gguf_filename: Optional[str] = None,
|
||||
) -> tuple[tuple[str, str, str], ...]:
|
||||
"""The sd.cpp text encoders for a specific load.
|
||||
|
||||
FLUX.2-klein selects its encoder by variant (the 9B transformer needs Qwen3-8B, the 4B needs the
|
||||
family default Qwen3-4B); every other family returns its static table. Keyed on the load identity
|
||||
(repo id + GGUF filename) so ``unsloth/FLUX.2-klein-9B-GGUF`` and a local ``*klein-9B*.gguf`` both
|
||||
resolve to the 8B encoder."""
|
||||
if fam.name == "flux.2-klein":
|
||||
identity = f"{repo_id or ''}/{gguf_filename or ''}".lower()
|
||||
if "klein-9b" in identity or "klein_9b" in identity:
|
||||
return _FLUX2_KLEIN_9B_SD_CPP_TEXT_ENCODERS
|
||||
return fam.sd_cpp_text_encoders
|
||||
|
||||
|
||||
def resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path:
|
||||
"""Resolve ``gguf_filename`` (user-supplied) to a file under ``repo_root``, rejecting absolute
|
||||
paths and ``..`` escapes."""
|
||||
|
|
|
|||
|
|
@ -203,6 +203,7 @@ def resolve_one(
|
|||
spec_id: str,
|
||||
weight: float,
|
||||
*,
|
||||
family: Optional[str] = None,
|
||||
hf_token: Optional[str] = None,
|
||||
cancel_event: Optional[threading.Event] = None,
|
||||
) -> ResolvedLora:
|
||||
|
|
@ -211,11 +212,22 @@ def resolve_one(
|
|||
Accepts a catalog/local id or a bare HF repo id (``owner/name[:weight_file.safetensors]``).
|
||||
Downloads hub weights via the xet-fallback helper. Raises FileNotFoundError/ValueError on an
|
||||
unresolvable/unsupported id, which the caller maps to a 400.
|
||||
|
||||
``family`` (the loaded model family) enforces catalog family tags HERE, not only in the picker
|
||||
(``list_loras``): a LoRA is architecture-specific, so a direct API client sending an id tagged
|
||||
for another family would otherwise load it through the wrong pipeline. Mirrors the ControlNet
|
||||
resolver's family gate. An untagged catalog entry (empty ``families``) stays unrestricted.
|
||||
"""
|
||||
# An empty/whitespace token triggers an auth error instead of anonymous access; normalise to None.
|
||||
hf_token = hf_token.strip() if hf_token and hf_token.strip() else None
|
||||
entry = _catalog_by_id().get(spec_id)
|
||||
if entry is not None:
|
||||
req_fam = (family or "").strip().lower()
|
||||
if entry.families and req_fam and req_fam not in {f.lower() for f in entry.families}:
|
||||
raise ValueError(
|
||||
f"LoRA '{spec_id}' is for {', '.join(entry.families)}, not the loaded "
|
||||
f"'{family}' model family; pick a LoRA built for this family."
|
||||
)
|
||||
if entry.source == "local":
|
||||
path = entry.local_path or ""
|
||||
if not path or not os.path.exists(path):
|
||||
|
|
@ -289,13 +301,16 @@ def _scrub_hub_url(msg: str) -> str:
|
|||
def resolve_specs(
|
||||
specs: list[tuple[str, float]],
|
||||
*,
|
||||
family: Optional[str] = None,
|
||||
hf_token: Optional[str] = None,
|
||||
cancel_event: Optional[threading.Event] = None,
|
||||
) -> list[ResolvedLora]:
|
||||
"""Resolve request (id, weight) pairs, dropping zero-weight entries.
|
||||
|
||||
Maps the named not-found/gated Hub errors to a 400 (URL scrubbed); does NOT catch the base
|
||||
HfHubHTTPError, so a Hub 5xx stays a 500. A mid-download cancel maps to a 409."""
|
||||
``family`` is the loaded model family; it enforces catalog family tags in :func:`resolve_one`
|
||||
for direct API callers, not only the UI picker. Maps the named not-found/gated Hub errors to a
|
||||
400 (URL scrubbed); does NOT catch the base HfHubHTTPError, so a Hub 5xx stays a 500. A
|
||||
mid-download cancel maps to a 409."""
|
||||
from huggingface_hub.errors import (
|
||||
EntryNotFoundError,
|
||||
GatedRepoError,
|
||||
|
|
@ -308,7 +323,11 @@ def resolve_specs(
|
|||
for spec_id, weight in specs:
|
||||
if weight == 0:
|
||||
continue
|
||||
out.append(resolve_one(spec_id, weight, hf_token = hf_token, cancel_event = cancel_event))
|
||||
out.append(
|
||||
resolve_one(
|
||||
spec_id, weight, family = family, hf_token = hf_token, cancel_event = cancel_event
|
||||
)
|
||||
)
|
||||
except (
|
||||
FileNotFoundError,
|
||||
RepositoryNotFoundError,
|
||||
|
|
|
|||
|
|
@ -94,5 +94,23 @@ def release(owner: str) -> None:
|
|||
_owner = None
|
||||
|
||||
|
||||
def release_if(owner: str, predicate: Callable[[], bool]) -> bool:
|
||||
"""Drop ``owner``'s claim only if it still holds it AND ``predicate()`` is true, atomically.
|
||||
|
||||
A slow unload's "nothing resident / no load in flight" check and the ``release`` must not
|
||||
straddle a concurrent same-owner load: that load's ``acquire_for(register=...)`` re-registers
|
||||
ownership UNDER this lock, so a plain check-then-``release`` could pass the stale check and then
|
||||
clear the newer claim (``release`` is owner-guarded but identity-less). Evaluating the predicate
|
||||
under the lock closes that window -- the load's register runs either fully before or fully after.
|
||||
``predicate`` must be quick (it holds the lock) and must not re-enter the arbiter. Returns True
|
||||
iff ownership was dropped."""
|
||||
global _owner
|
||||
with _lock:
|
||||
if _owner != owner or not predicate():
|
||||
return False
|
||||
_owner = None
|
||||
return True
|
||||
|
||||
|
||||
def current_owner() -> Optional[str]:
|
||||
return _owner
|
||||
|
|
|
|||
|
|
@ -188,6 +188,11 @@ def delete(image_id: str) -> bool:
|
|||
path = image_path(image_id)
|
||||
if path is None:
|
||||
return False
|
||||
# Only delete files we actually own (a readable recipe chunk). A foreign PNG that a caller
|
||||
# dropped in by hand is invisible to list_images, so deleting it here on a guessed id would
|
||||
# silently destroy a file the gallery never claimed.
|
||||
if _read_meta(path) is None:
|
||||
return False
|
||||
try:
|
||||
path.unlink()
|
||||
return True
|
||||
|
|
@ -197,13 +202,18 @@ def delete(image_id: str) -> bool:
|
|||
|
||||
|
||||
def clear() -> int:
|
||||
"""Delete every gallery PNG; return how many were removed."""
|
||||
"""Delete every Studio-owned gallery PNG; return how many were removed.
|
||||
|
||||
Preserves foreign PNGs (no readable recipe chunk): list_images already hides them, so clear
|
||||
must not silently destroy files the gallery never surfaced."""
|
||||
removed = 0
|
||||
try:
|
||||
paths = list(gallery_dir().glob("*.png"))
|
||||
except OSError:
|
||||
return 0
|
||||
for path in paths:
|
||||
if _read_meta(path) is None: # foreign / not ours
|
||||
continue
|
||||
try:
|
||||
path.unlink()
|
||||
removed += 1
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ from core.inference.diffusion_families import (
|
|||
family_sd_cpp_supported,
|
||||
resolve_base_repo,
|
||||
resolve_local_gguf_child,
|
||||
sd_cpp_text_encoders_for,
|
||||
supported_family_names,
|
||||
)
|
||||
from core.inference.diffusion_memory import (
|
||||
|
|
@ -599,7 +600,9 @@ class SdCppDiffusionBackend:
|
|||
specs: list[tuple[str, str, str]] = [(repo_id, gguf_filename, "diffusion_model")]
|
||||
if fam.sd_cpp_vae:
|
||||
specs.append((fam.sd_cpp_vae[0], fam.sd_cpp_vae[1], "vae"))
|
||||
for terepo, tefile, kind in fam.sd_cpp_text_encoders:
|
||||
# Select the text encoder per variant (FLUX.2-klein 4B->Qwen3-4B, 9B->Qwen3-8B) from the
|
||||
# load identity, not the family's single default, so a 9B GGUF fetches the right encoder.
|
||||
for terepo, tefile, kind in sd_cpp_text_encoders_for(fam, repo_id, gguf_filename):
|
||||
specs.append((terepo, tefile, kind))
|
||||
return specs
|
||||
|
||||
|
|
@ -694,7 +697,12 @@ class SdCppDiffusionBackend:
|
|||
repos = [state.repo_id, state.base_repo]
|
||||
if fam.sd_cpp_vae:
|
||||
repos.append(fam.sd_cpp_vae[0])
|
||||
repos.extend(terepo for terepo, _f, _k in fam.sd_cpp_text_encoders)
|
||||
# Same per-variant encoder selection as _asset_specs, keyed on the loaded repo id, so the
|
||||
# cache-deletion guard protects the encoder repo this load actually downloaded (the 9B
|
||||
# variant's Qwen3-8B, not the 4B default).
|
||||
repos.extend(
|
||||
terepo for terepo, _f, _k in sd_cpp_text_encoders_for(fam, state.repo_id)
|
||||
)
|
||||
return tuple(dict.fromkeys(r for r in repos if r))
|
||||
|
||||
# ── Generate ───────────────────────────────────────────────────────────
|
||||
|
|
@ -793,7 +801,10 @@ class SdCppDiffusionBackend:
|
|||
"sd.cpp engine."
|
||||
)
|
||||
lora_resolved = diffusion_lora.resolve_specs(
|
||||
active_loras, hf_token = state.hf_token, cancel_event = cancel
|
||||
active_loras,
|
||||
family = state.family.name,
|
||||
hf_token = state.hf_token,
|
||||
cancel_event = cancel,
|
||||
)
|
||||
if state.mode == "server" and state.server is not None:
|
||||
images, seeds = self._generate_server(
|
||||
|
|
|
|||
|
|
@ -229,10 +229,15 @@ def list_videos(
|
|||
|
||||
|
||||
def delete(video_id: str) -> bool:
|
||||
"""Remove both files of a pair; True if the MP4 existed."""
|
||||
"""Remove both files of an owned pair; True if the MP4 existed and was ours."""
|
||||
path = video_path(video_id)
|
||||
if path is None:
|
||||
return False
|
||||
# Only delete a pair we actually own (a readable sidecar). A foreign / orphan MP4 is invisible
|
||||
# to list_videos, so deleting it here on a guessed id would silently destroy a file the gallery
|
||||
# never claimed.
|
||||
if _read_meta(_sidecar_path(video_id)) is None:
|
||||
return False
|
||||
# Delete the MP4 FIRST: if the sidecar were dropped first and the mp4 unlink then failed (lock /
|
||||
# permission), the still-present mp4 would vanish from the gallery with no retry. mp4-first means
|
||||
# the worst case is an orphaned sidecar, which list_videos ignores.
|
||||
|
|
@ -250,13 +255,18 @@ def delete(video_id: str) -> bool:
|
|||
|
||||
|
||||
def clear() -> int:
|
||||
"""Delete every gallery pair; return how many videos were removed."""
|
||||
"""Delete every Studio-owned gallery pair; return how many videos were removed.
|
||||
|
||||
Preserves foreign / orphan MP4s (no readable sidecar): list_videos already hides them, so clear
|
||||
must not silently destroy files the gallery never surfaced."""
|
||||
removed = 0
|
||||
try:
|
||||
paths = list(gallery_dir().glob("*.mp4"))
|
||||
except OSError:
|
||||
return 0
|
||||
for path in paths:
|
||||
if _read_meta(_sidecar_path(path.stem)) is None: # orphan / not ours
|
||||
continue
|
||||
# mp4 first; if it can't be unlinked, leave the sidecar so the video stays listable.
|
||||
try:
|
||||
path.unlink()
|
||||
|
|
|
|||
|
|
@ -688,13 +688,22 @@ def discover_image_caption_pairs(
|
|||
meta_path = root / meta_name
|
||||
if not meta_path.is_file():
|
||||
continue
|
||||
for line in meta_path.read_text(encoding = "utf-8").splitlines():
|
||||
# Tolerate a bad upload: invalid UTF-8 in the file, or a line that is valid JSON but not an
|
||||
# object (``[]`` / ``null`` / a string / a number). Neither should crash the trainer -- the
|
||||
# record is simply skipped so the instance_prompt fallback still applies.
|
||||
try:
|
||||
meta_lines = meta_path.read_text(encoding = "utf-8").splitlines()
|
||||
except (OSError, UnicodeError):
|
||||
continue
|
||||
for line in meta_lines:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
row = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
continue
|
||||
if not isinstance(row, dict):
|
||||
continue
|
||||
key = row.get("file_name") or row.get("image") or row.get("file")
|
||||
if key and caption_column in row:
|
||||
|
|
|
|||
|
|
@ -2094,7 +2094,15 @@ class GalleryImage(BaseModel):
|
|||
height: int = Field(..., description = "Image height")
|
||||
steps: int = Field(..., description = "Denoising steps")
|
||||
guidance: float = Field(..., description = "Guidance scale")
|
||||
seed: int = Field(..., description = "Seed used")
|
||||
seed: int = Field(..., description = "Seed used for THIS image")
|
||||
batch_seed: Optional[int] = Field(
|
||||
None,
|
||||
description = (
|
||||
"Base seed the batch was launched with. The native engine derives per-image seeds as "
|
||||
"base + index, so restore must replay from this base, not from the derived per-image "
|
||||
"seed; older records without it fall back to seed."
|
||||
),
|
||||
)
|
||||
batch_index: int = Field(0, description = "Position within its batch (0-based)")
|
||||
batch_size: int = Field(
|
||||
1, description = "Batch size used; with batch_index it lets restore replay this image"
|
||||
|
|
|
|||
|
|
@ -14474,6 +14474,12 @@ async def generate_diffusion_image(
|
|||
"steps": request.steps,
|
||||
"guidance": request.guidance,
|
||||
"seed": seed,
|
||||
# The base seed the batch launched with. The native engine derives per-image
|
||||
# seeds as base + index, so ``seed`` above is already advanced for index>0;
|
||||
# restore must replay from this base (with batch_size) or it would advance a
|
||||
# second time and reproduce a different image. Diffusers shares one seed, so
|
||||
# base == seed there.
|
||||
"batch_seed": result["seed"],
|
||||
# Position within the batch (shared timestamp), so the export filename
|
||||
# stays unique.
|
||||
"batch_index": index,
|
||||
|
|
@ -14578,18 +14584,23 @@ async def clear_gallery_images(current_subject: str = Depends(get_current_subjec
|
|||
@studio_router.post("/images/unload", response_model = DiffusionStatusResponse)
|
||||
async def unload_diffusion_model(current_subject: str = Depends(get_current_subject)):
|
||||
from core.inference.diffusion_engine_router import annotate_status, get_active_diffusion_engine
|
||||
from core.inference.gpu_arbiter import release, DIFFUSION
|
||||
from core.inference.gpu_arbiter import release_if, DIFFUSION
|
||||
|
||||
status_dict = await asyncio.to_thread(get_active_diffusion_engine().unload)
|
||||
# Drop DIFFUSION ownership only if nothing is resident AND no new load is in flight: a
|
||||
# concurrent /images/load that re-acquired DIFFUSION while this (slow) unload ran must keep
|
||||
# ownership, or a later chat load would see no owner, skip eviction, and OOM the newly
|
||||
# resident pipeline. An in-flight load has is_loaded False for its whole download/finalize
|
||||
# window, so gate on loading_repo_ids() too, not just committed state. release() is
|
||||
# owner-guarded and identity-less, so an unconditional release would clear the newer claim.
|
||||
# window, so gate on loading_repo_ids() too, not just committed state. The idle check and the
|
||||
# release must be ATOMIC (release_if): the load's acquire_for register runs under the same
|
||||
# arbiter lock, so a plain check-then-release could pass the stale check and then clear the
|
||||
# newer claim.
|
||||
engine = get_active_diffusion_engine()
|
||||
if not engine.loading_repo_ids() and not engine.is_loaded:
|
||||
release(DIFFUSION)
|
||||
await asyncio.to_thread(
|
||||
release_if,
|
||||
DIFFUSION,
|
||||
lambda: not engine.loading_repo_ids() and not engine.is_loaded,
|
||||
)
|
||||
return DiffusionStatusResponse(**annotate_status(status_dict))
|
||||
|
||||
|
||||
|
|
@ -14775,7 +14786,12 @@ async def openai_image_generations(
|
|||
if per_image_seeds and index < len(per_image_seeds)
|
||||
else result["seed"]
|
||||
)
|
||||
record = image_gallery.save(image, {**recipe, "batch_index": index, "seed": seed})
|
||||
# batch_seed is the base the native engine derives per-image seeds from (base + index),
|
||||
# so restore replays from it rather than double-advancing the derived seed above.
|
||||
record = image_gallery.save(
|
||||
image,
|
||||
{**recipe, "batch_index": index, "seed": seed, "batch_seed": result["seed"]},
|
||||
)
|
||||
if want_b64:
|
||||
encoded = image_gallery.image_b64(record["id"])
|
||||
if encoded is None: # vanished between write and read — fail the call
|
||||
|
|
|
|||
|
|
@ -1137,6 +1137,23 @@ def _diffusion_training_active() -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _require_diffusion_dataset_mutable() -> None:
|
||||
"""Reject a dataset mutation while a diffusion run is active.
|
||||
|
||||
The trainer enumerates and (when the latent cache is off / over budget) re-opens dataset images
|
||||
during the loop, so uploading, importing, captioning, or deleting underneath it makes the run
|
||||
nondeterministic or raises a FileNotFoundError mid-step. Best-effort: a service-import failure
|
||||
fails open (never blocks a mutation on an unknowable state), matching the start interlock."""
|
||||
if _diffusion_training_active():
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = (
|
||||
"Training images cannot be changed while diffusion training is active. "
|
||||
"Stop the run before uploading, importing, editing captions, or deleting images."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _free_gpu_for_diffusion_training() -> None:
|
||||
"""Free GPU residents before the diffusion trainer spawns its own SDXL pipeline.
|
||||
|
||||
|
|
@ -1254,8 +1271,13 @@ def _resolve_diffusion_data_dir(raw: str) -> Path:
|
|||
# Single component and not ".." -> joining under datasets_root() cannot escape it.
|
||||
if not p.is_absolute() and len(p.parts) == 1 and p.parts[0] != "..":
|
||||
direct = datasets_root() / value
|
||||
if direct.is_dir():
|
||||
return direct
|
||||
# Route a bare image-dataset name through the SAME protected resolver the
|
||||
# caption/delete/labeling CRUD routes use, so a name -> external-directory symlink is
|
||||
# rejected here too (is_dir() follows the link, so a plain is_dir() check would train on
|
||||
# files outside the datasets root). Include a broken symlink so it is rejected, not
|
||||
# silently passed through to resolve_dataset_path.
|
||||
if direct.is_dir() or direct.is_symlink():
|
||||
return _resolve_dataset_folder(value)
|
||||
return resolve_dataset_path(raw)
|
||||
|
||||
|
||||
|
|
@ -1554,7 +1576,12 @@ async def diffusion_training_info(current_subject: str = Depends(get_current_sub
|
|||
# Skip hidden dirs: never user datasets, and an in-progress example import stages
|
||||
# into a dot-prefixed sibling that must not surface as a dataset.
|
||||
children = sorted(
|
||||
p for p in root.iterdir() if p.is_dir() and not p.name.startswith(".")
|
||||
p
|
||||
for p in root.iterdir()
|
||||
# Skip symlinked dirs: the CRUD resolver (_resolve_dataset_folder) rejects a
|
||||
# symlinked dataset, so discovery must not advertise one as selectable (an external
|
||||
# directory the read/caption/delete routes would then refuse).
|
||||
if p.is_dir() and not p.is_symlink() and not p.name.startswith(".")
|
||||
)
|
||||
except OSError:
|
||||
children = []
|
||||
|
|
@ -1615,6 +1642,7 @@ async def upload_diffusion_dataset(
|
|||
|
||||
from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label
|
||||
|
||||
_require_diffusion_dataset_mutable()
|
||||
cleaned = _clean_diffusion_dataset_name(name)
|
||||
# Run the same symlink + root-containment check as the read/caption/delete endpoints before any
|
||||
# write, so a name -> external-directory symlink can't make the staged upload write outside root.
|
||||
|
|
@ -1732,6 +1760,13 @@ async def upload_diffusion_dataset(
|
|||
),
|
||||
)
|
||||
out.write(chunk)
|
||||
# Reject a decompression bomb before commit: the byte-limit above passes a small, highly
|
||||
# compressible PNG whose decoded pixels are huge, and the trainer later decodes every
|
||||
# image in full when building its latent cache. Mirror the inference decode guard
|
||||
# (diffusion._decode_b64_image) and bound each image's dimensions from the header, BEFORE
|
||||
# any pixel decompression, so an oversized upload 400s here rather than OOMing the run.
|
||||
if Path(filename).suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS:
|
||||
_validate_uploaded_training_image(tmp, filename)
|
||||
uploaded += 1
|
||||
# Commit every staged file as one transaction. A plain replace loop is not atomic across
|
||||
# files: a mid-loop tmp.replace(dest) failure leaves earlier destinations already overwritten
|
||||
|
|
@ -1823,6 +1858,37 @@ def _resolve_dataset_folder(name: str, *, must_exist: bool = True) -> Path:
|
|||
return folder
|
||||
|
||||
|
||||
# Bound each uploaded training image's dimensions (matches the inference decode guard's 4096px
|
||||
# per-side limit in diffusion._decode_b64_image). A small, highly compressible PNG can smuggle huge
|
||||
# pixel dimensions past the byte limit and OOM the trainer when it decodes the image for its latent
|
||||
# cache, so reject an over-limit image from the header before any pixel decompression.
|
||||
_MAX_TRAINING_IMAGE_SIDE = 4096
|
||||
|
||||
|
||||
def _validate_uploaded_training_image(path: Path, original_name: str) -> None:
|
||||
"""Reject an uploaded training image whose decoded dimensions exceed the per-side limit.
|
||||
|
||||
Reads only the image header (never img.load()), so a crafted small-payload / huge-dimension file
|
||||
is caught before it can spike memory. Scoped to the decompression-bomb vector only: bytes PIL
|
||||
cannot identify are left as-is (the upload contract accepts arbitrary bytes under an image
|
||||
extension), so this changes behaviour solely for oversized real images."""
|
||||
from PIL import Image, UnidentifiedImageError
|
||||
|
||||
try:
|
||||
with Image.open(path) as image:
|
||||
width, height = image.size
|
||||
except (OSError, UnidentifiedImageError, ValueError):
|
||||
return # not a decodable image -> not a decompression bomb; leave the existing contract
|
||||
if width > _MAX_TRAINING_IMAGE_SIDE or height > _MAX_TRAINING_IMAGE_SIDE:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = (
|
||||
f"Image '{original_name}' is too large ({width}x{height}); maximum is "
|
||||
f"{_MAX_TRAINING_IMAGE_SIDE}px per side."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _safe_dataset_image_path(folder: Path, filename: str) -> Path:
|
||||
"""Resolve ``filename`` to an image path strictly inside ``folder``. Rejects any path
|
||||
separators / traversal / null bytes and non-image extensions."""
|
||||
|
|
@ -1851,9 +1917,12 @@ def _load_metadata_captions(folder: Path) -> dict[str, str]:
|
|||
meta_path = folder / meta_name
|
||||
if not meta_path.is_file():
|
||||
continue
|
||||
# Tolerate a bad upload: invalid UTF-8 (UnicodeError, not an OSError), or a line that is
|
||||
# valid JSON but not an object (``[]`` / ``null`` / a string / a number). Neither should
|
||||
# 500 the info / labeling / caption / summary endpoints; the record is simply skipped.
|
||||
try:
|
||||
lines = meta_path.read_text(encoding = "utf-8").splitlines()
|
||||
except OSError:
|
||||
except (OSError, UnicodeError):
|
||||
continue
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
|
|
@ -1861,7 +1930,9 @@ def _load_metadata_captions(folder: Path) -> dict[str, str]:
|
|||
continue
|
||||
try:
|
||||
row = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
continue
|
||||
if not isinstance(row, dict):
|
||||
continue
|
||||
key = row.get("file_name") or row.get("image") or row.get("file")
|
||||
if key and "text" in row:
|
||||
|
|
@ -1997,6 +2068,7 @@ async def set_diffusion_dataset_caption(
|
|||
):
|
||||
"""Write (or, when blank, clear) an image's ``.txt`` caption sidecar. Returns the
|
||||
updated image record."""
|
||||
_require_diffusion_dataset_mutable()
|
||||
folder = _resolve_dataset_folder(name)
|
||||
image_path = _safe_dataset_image_path(folder, filename)
|
||||
if not image_path.is_file():
|
||||
|
|
@ -2040,6 +2112,7 @@ async def delete_diffusion_dataset_image(
|
|||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""Remove an image, its caption sidecars, and any cached thumbnails."""
|
||||
_require_diffusion_dataset_mutable()
|
||||
folder = _resolve_dataset_folder(name)
|
||||
image_path = _safe_dataset_image_path(folder, filename)
|
||||
if not image_path.is_file():
|
||||
|
|
@ -2304,6 +2377,7 @@ async def import_diffusion_dataset_example(
|
|||
"""Materialize a curated example dataset into a Studio dataset folder (images + .txt
|
||||
captions), ready to train. Idempotent: a folder that already holds images is returned
|
||||
as-is rather than re-downloaded."""
|
||||
_require_diffusion_dataset_mutable()
|
||||
entry = _example_by_id(body.id)
|
||||
folder = _resolve_dataset_folder(body.name or entry["id"], must_exist = False)
|
||||
|
||||
|
|
|
|||
|
|
@ -219,16 +219,22 @@ async def video_status(current_subject: str = Depends(get_current_subject)):
|
|||
|
||||
@router.post("/video/unload", response_model = VideoStatusResponse)
|
||||
async def unload_video_model(current_subject: str = Depends(get_current_subject)):
|
||||
from core.inference.gpu_arbiter import VIDEO, release
|
||||
from core.inference.gpu_arbiter import VIDEO, release_if
|
||||
from core.inference.video import get_video_backend
|
||||
|
||||
backend = get_video_backend()
|
||||
status_dict = await asyncio.to_thread(backend.unload)
|
||||
# Drop VIDEO ownership only if nothing is resident AND no new load is in flight: a concurrent
|
||||
# /video/load that re-acquired VIDEO must keep ownership (release() is owner-guarded but
|
||||
# identity-less, so an unconditional release would clear the newer claim). Mirrors the images route.
|
||||
if not backend.loading_repo_ids() and not backend.status()["loaded"]:
|
||||
release(VIDEO)
|
||||
# /video/load that re-acquired VIDEO must keep ownership (release is owner-guarded but
|
||||
# identity-less, so an unconditional release would clear the newer claim). The idle check and
|
||||
# the release must be ATOMIC (release_if): the load's acquire_for register runs under the same
|
||||
# arbiter lock, so a plain check-then-release could pass the stale check and then clear the
|
||||
# newer claim. Mirrors the images route.
|
||||
await asyncio.to_thread(
|
||||
release_if,
|
||||
VIDEO,
|
||||
lambda: not backend.loading_repo_ids() and not backend.status()["loaded"],
|
||||
)
|
||||
return VideoStatusResponse(**status_dict)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -235,9 +235,11 @@ class _FakeCNModel:
|
|||
path,
|
||||
torch_dtype = None,
|
||||
token = None,
|
||||
use_safetensors = None,
|
||||
):
|
||||
m = cls()
|
||||
m.path = path
|
||||
m.use_safetensors = use_safetensors
|
||||
return m
|
||||
|
||||
def to(self, device):
|
||||
|
|
@ -303,6 +305,9 @@ def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
|
|||
p1 = b._controlnet_pipe(st, resolved, threading.Event())
|
||||
assert isinstance(p1, _FakeCNPipe) and isinstance(p1.controlnet, _FakeCNModel)
|
||||
assert p1.controlnet.path == "repo/id" and p1.controlnet.device == "cpu"
|
||||
# A remote (non-local) ControlNet must force safetensors so a pickle can't deserialize even if
|
||||
# the Hub scan failed open.
|
||||
assert p1.controlnet.use_safetensors is True
|
||||
# cached: same id -> same model + same pipe, no reload.
|
||||
p2 = b._controlnet_pipe(st, resolved, threading.Event())
|
||||
assert p2 is p1
|
||||
|
|
@ -327,9 +332,12 @@ def test_controlnet_pipe_blocks_flagged_remote_repo(monkeypatch):
|
|||
path,
|
||||
torch_dtype = None,
|
||||
token = None,
|
||||
use_safetensors = None,
|
||||
):
|
||||
loaded["called"] = True
|
||||
return super().from_pretrained(path, torch_dtype = torch_dtype, token = token)
|
||||
return super().from_pretrained(
|
||||
path, torch_dtype = torch_dtype, token = token, use_safetensors = use_safetensors
|
||||
)
|
||||
|
||||
mod = _fake_diffusers()
|
||||
mod.FluxControlNetModel = _TrapModel
|
||||
|
|
@ -370,6 +378,8 @@ def test_controlnet_pipe_skips_scan_for_local_dir(monkeypatch, tmp_path):
|
|||
resolved = dc.ResolvedControlNet("my-cn", str(tmp_path), is_local = True)
|
||||
p = b._controlnet_pipe(st, resolved, threading.Event())
|
||||
assert isinstance(p, _FakeCNPipe)
|
||||
# A local dir the user chose is exempt from the forced-safetensors gate (may be .bin).
|
||||
assert p.controlnet.use_safetensors is None
|
||||
|
||||
|
||||
def test_controlnet_pipe_rejects_family_without_classes():
|
||||
|
|
|
|||
|
|
@ -174,6 +174,24 @@ def test_resolve_one_local_and_unknown(tmp_path, monkeypatch):
|
|||
dl.resolve_one("does-not-exist", 1.0)
|
||||
|
||||
|
||||
def test_resolve_one_rejects_cross_family_catalog_entry(tmp_path, monkeypatch):
|
||||
# A family-tagged adapter must be rejected in the resolver (not just the UI picker) when the
|
||||
# loaded model is a different family, so a direct API client cannot apply a mismatched LoRA.
|
||||
d = tmp_path / "loras"
|
||||
d.mkdir()
|
||||
(d / "krea-style.safetensors").write_bytes(b"x")
|
||||
(d / "krea-style.json").write_text('{"families": ["krea-2"]}', encoding = "utf-8")
|
||||
monkeypatch.setattr(dl, "loras_dir", lambda: d)
|
||||
# Same family: resolves.
|
||||
r = dl.resolve_one("krea-style", 0.7, family = "krea-2")
|
||||
assert r.path.endswith("krea-style.safetensors")
|
||||
# Wrong family: rejected before any download / apply.
|
||||
with pytest.raises(ValueError):
|
||||
dl.resolve_one("krea-style", 0.7, family = "flux.1")
|
||||
# No family context (e.g. legacy caller): unrestricted.
|
||||
assert dl.resolve_one("krea-style", 0.7).path.endswith("krea-style.safetensors")
|
||||
|
||||
|
||||
def test_resolve_specs_drops_zero_weight(tmp_path, monkeypatch):
|
||||
d = tmp_path / "loras"
|
||||
d.mkdir()
|
||||
|
|
|
|||
|
|
@ -103,6 +103,25 @@ def test_discover_captions_jsonl_and_image_key(tmp_path):
|
|||
assert discover_image_caption_pairs(tmp_path) == [(str(tmp_path / "x.png"), "hi")]
|
||||
|
||||
|
||||
def test_discover_tolerates_non_object_and_invalid_utf8_jsonl(tmp_path):
|
||||
# A metadata.jsonl line that is valid JSON but not an object ([]/null/string/number) or malformed
|
||||
# must be skipped per-line rather than crash the trainer in .get(); a valid row still resolves.
|
||||
_touch(tmp_path / "x.png")
|
||||
(tmp_path / "metadata.jsonl").write_text(
|
||||
"[]\nnull\n\"str\"\n123\n{not json\n"
|
||||
+ json.dumps({"file_name": "x.png", "text": "hi"})
|
||||
+ "\n",
|
||||
encoding = "utf-8",
|
||||
)
|
||||
assert discover_image_caption_pairs(tmp_path) == [(str(tmp_path / "x.png"), "hi")]
|
||||
# Invalid UTF-8 in the metadata file must not raise; the file is skipped (image falls back to
|
||||
# the instance prompt).
|
||||
_touch(tmp_path / "y.png")
|
||||
(tmp_path / "captions.jsonl").write_bytes(b"\xff\xfe not utf-8\n")
|
||||
pairs = dict(discover_image_caption_pairs(tmp_path, instance_prompt = "fallback"))
|
||||
assert pairs[str(tmp_path / "y.png")] == "fallback"
|
||||
|
||||
|
||||
def test_discover_custom_caption_column(tmp_path):
|
||||
_touch(tmp_path / "x.png")
|
||||
(tmp_path / "metadata.jsonl").write_text(
|
||||
|
|
|
|||
|
|
@ -1000,6 +1000,116 @@ def test_diffusion_info_empty_sidecar_shadows_metadata_caption(client, dataset_r
|
|||
assert summary["caption_count"] == 2
|
||||
|
||||
|
||||
def _png_bytes(width: int, height: int) -> bytes:
|
||||
import io
|
||||
|
||||
pytest.importorskip("PIL")
|
||||
from PIL import Image
|
||||
|
||||
buf = io.BytesIO()
|
||||
Image.new("RGB", (width, height), (1, 2, 3)).save(buf, format = "PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def test_diffusion_dataset_upload_rejects_oversized_image(client, dataset_roots):
|
||||
# A decompression bomb: a small compressible PNG with huge dimensions passes the byte limit but
|
||||
# would OOM the trainer on decode. It must 400 at upload (dimension check reads only the header).
|
||||
pytest.importorskip("PIL")
|
||||
big = _png_bytes(5000, 64) # > 4096 per side
|
||||
r = client.post(
|
||||
"/api/train/diffusion/dataset",
|
||||
data = {"name": "bomb"},
|
||||
files = [("files", ("huge.png", big, "image/png"))],
|
||||
)
|
||||
assert r.status_code == 400, r.text
|
||||
assert "too large" in r.json()["detail"]
|
||||
# An in-bounds real image still uploads fine.
|
||||
ok = _png_bytes(64, 64)
|
||||
r2 = client.post(
|
||||
"/api/train/diffusion/dataset",
|
||||
data = {"name": "bomb"},
|
||||
files = [("files", ("ok.png", ok, "image/png"))],
|
||||
)
|
||||
assert r2.status_code == 200, r2.text
|
||||
|
||||
|
||||
def test_diffusion_info_tolerates_non_object_jsonl(client, dataset_roots):
|
||||
# A metadata.jsonl line that is valid JSON but not an object ([]/null/string/number) or malformed
|
||||
# must be skipped per-line, not 500 the info endpoint; a valid row in the same file still counts.
|
||||
ds_root, _ = dataset_roots
|
||||
folder = ds_root / "weird-meta"
|
||||
folder.mkdir()
|
||||
(folder / "a.png").write_bytes(b"x")
|
||||
(folder / "metadata.jsonl").write_bytes(
|
||||
b"[]\n"
|
||||
b"null\n"
|
||||
b'"just a string"\n'
|
||||
b"123\n"
|
||||
b"{not json\n"
|
||||
+ json.dumps({"file_name": "a.png", "text": "cap a"}).encode("utf-8")
|
||||
+ b"\n"
|
||||
)
|
||||
r = client.get("/api/train/diffusion/info")
|
||||
assert r.status_code == 200, r.text
|
||||
summary = next(d for d in r.json()["datasets"] if d["name"] == "weird-meta")
|
||||
assert summary["caption_count"] == 1
|
||||
|
||||
|
||||
def test_diffusion_info_tolerates_invalid_utf8_jsonl(client, dataset_roots):
|
||||
# Invalid UTF-8 in a metadata file must not 500 the info endpoint; the file is skipped, not decoded.
|
||||
ds_root, _ = dataset_roots
|
||||
folder = ds_root / "bad-utf8-meta"
|
||||
folder.mkdir()
|
||||
(folder / "a.png").write_bytes(b"x")
|
||||
(folder / "metadata.jsonl").write_bytes(b"\xff\xfe not valid utf-8\n")
|
||||
r = client.get("/api/train/diffusion/info")
|
||||
assert r.status_code == 200, r.text
|
||||
summary = next(d for d in r.json()["datasets"] if d["name"] == "bad-utf8-meta")
|
||||
assert summary["caption_count"] == 0
|
||||
|
||||
|
||||
def test_diffusion_dataset_mutations_blocked_while_training_active(client, dataset_roots):
|
||||
ds_root, _ = dataset_roots
|
||||
folder = ds_root / "locked"
|
||||
folder.mkdir()
|
||||
folder.joinpath("a.png").write_bytes(b"x")
|
||||
# Flip the fake diffusion service to active.
|
||||
client._fake._running = True
|
||||
# Upload, caption, delete, and example-import must all 409 while a run is active.
|
||||
up = client.post(
|
||||
"/api/train/diffusion/dataset",
|
||||
data = {"name": "locked"},
|
||||
files = [("files", ("b.png", b"x", "image/png"))],
|
||||
)
|
||||
assert up.status_code == 409, up.text
|
||||
cap = client.put(
|
||||
"/api/train/diffusion/dataset/locked/caption/a.png", json = {"caption": "hi"}
|
||||
)
|
||||
assert cap.status_code == 409, cap.text
|
||||
dele = client.delete("/api/train/diffusion/dataset/locked/image/a.png")
|
||||
assert dele.status_code == 409, dele.text
|
||||
imp = client.post("/api/train/diffusion/dataset/import-example", json = {"id": "anything"})
|
||||
assert imp.status_code == 409, imp.text
|
||||
|
||||
|
||||
def test_diffusion_info_skips_symlinked_dataset_dir(client, dataset_roots):
|
||||
# A directory symlink under the datasets root must not be advertised as a dataset (the CRUD
|
||||
# resolver rejects symlinked datasets, so discovery must agree).
|
||||
import os
|
||||
|
||||
ds_root, _ = dataset_roots
|
||||
outside = ds_root.parent / "outside-images"
|
||||
outside.mkdir()
|
||||
(outside / "a.png").write_bytes(b"x")
|
||||
try:
|
||||
os.symlink(outside, ds_root / "linked")
|
||||
except OSError:
|
||||
pytest.skip("symlinks unavailable")
|
||||
r = client.get("/api/train/diffusion/info")
|
||||
assert r.status_code == 200, r.text
|
||||
assert "linked" not in [d["name"] for d in r.json()["datasets"]]
|
||||
|
||||
|
||||
def test_diffusion_dataset_upload_rejects_traversal_names(client, dataset_roots):
|
||||
for bad in ("../evil", "a/b", ".hidden", " "):
|
||||
r = client.post(
|
||||
|
|
|
|||
|
|
@ -106,6 +106,35 @@ def test_evict_chat_unloads_a_still_loading_chat_backend(monkeypatch):
|
|||
assert unloaded == [True] # still-loading chat backend was unloaded, not skipped
|
||||
|
||||
|
||||
def test_release_if_drops_only_when_predicate_true(calls):
|
||||
arb.acquire_for(arb.DIFFUSION)
|
||||
# Predicate false -> ownership kept.
|
||||
assert arb.release_if(arb.DIFFUSION, lambda: False) is False
|
||||
assert arb.current_owner() == arb.DIFFUSION
|
||||
# Predicate true -> ownership dropped.
|
||||
assert arb.release_if(arb.DIFFUSION, lambda: True) is True
|
||||
assert arb.current_owner() is None
|
||||
|
||||
|
||||
def test_release_if_by_non_owner_is_noop(calls):
|
||||
arb.acquire_for(arb.CHAT)
|
||||
# Predicate is never even consulted for a non-owner; ownership is untouched.
|
||||
consulted: list[bool] = []
|
||||
assert arb.release_if(arb.DIFFUSION, lambda: consulted.append(True) or True) is False
|
||||
assert consulted == []
|
||||
assert arb.current_owner() == arb.CHAT
|
||||
|
||||
|
||||
def test_release_if_predicate_sees_a_reregistered_same_owner_load(calls):
|
||||
# The race release_if closes: a slow unload's predicate must observe a concurrent same-owner
|
||||
# load that re-registered ownership, and NOT drop the newer claim. Simulate the re-register by
|
||||
# having the predicate report a load now in flight; ownership must stay with DIFFUSION.
|
||||
arb.acquire_for(arb.DIFFUSION)
|
||||
loading = {"in_flight": True}
|
||||
assert arb.release_if(arb.DIFFUSION, lambda: not loading["in_flight"]) is False
|
||||
assert arb.current_owner() == arb.DIFFUSION
|
||||
|
||||
|
||||
def test_register_runs_under_ownership_and_returns_result(calls):
|
||||
# A register callback runs after ownership transfers (owner already set) and its
|
||||
# return value is forwarded -- the route uses this to register the in-flight load.
|
||||
|
|
|
|||
|
|
@ -103,6 +103,24 @@ def test_delete_and_clear():
|
|||
assert gallery.list_images() == []
|
||||
|
||||
|
||||
def test_clear_preserves_foreign_png():
|
||||
# A hand-dropped PNG with no recipe chunk is invisible to list_images; clear must not destroy it.
|
||||
foreign = gallery.gallery_dir() / "family-photo.png"
|
||||
_img().save(foreign, format = "PNG")
|
||||
gallery.save(_img(), _meta(prompt = "ours"))
|
||||
assert gallery.clear() == 1
|
||||
assert foreign.exists()
|
||||
assert gallery.list_images() == []
|
||||
|
||||
|
||||
def test_delete_ignores_foreign_png():
|
||||
# A per-id delete must refuse a file we do not own (no readable recipe chunk).
|
||||
foreign = gallery.gallery_dir() / "family-photo.png"
|
||||
_img().save(foreign, format = "PNG")
|
||||
assert gallery.delete("family-photo") is False
|
||||
assert foreign.exists()
|
||||
|
||||
|
||||
def test_image_path_rejects_unsafe_ids():
|
||||
# Traversal / bad chars never resolve to a path.
|
||||
assert gallery.image_path("../../etc/passwd") is None
|
||||
|
|
|
|||
|
|
@ -183,6 +183,26 @@ def test_asset_specs_cover_required_files(fam_name, expect_kinds):
|
|||
assert tr[0] == "unsloth/x-GGUF" and tr[1] == "x-Q4_K_M.gguf"
|
||||
|
||||
|
||||
def test_asset_specs_flux2_klein_selects_encoder_by_variant():
|
||||
# FLUX.2-klein 4B pairs with Qwen3-4B, 9B with Qwen3-8B; the encoder must be chosen from the
|
||||
# load identity, not the family's single default (a mismatched encoder fails deep in sd-cli).
|
||||
b = SdCppDiffusionBackend(engine = _FakeEngine())
|
||||
fam = detect_family("flux.2-klein")
|
||||
|
||||
specs_4b = b._asset_specs("unsloth/FLUX.2-klein-4B-GGUF", "FLUX.2-klein-4B-Q4_K_M.gguf", fam)
|
||||
te_4b = [(r, f) for r, f, k in specs_4b if k == "llm"]
|
||||
assert te_4b == [("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors")]
|
||||
|
||||
specs_9b = b._asset_specs("unsloth/FLUX.2-klein-9B-GGUF", "FLUX.2-klein-9B-Q4_K_M.gguf", fam)
|
||||
te_9b = [(r, f) for r, f, k in specs_9b if k == "llm"]
|
||||
assert te_9b == [
|
||||
(
|
||||
"Comfy-Org/vae-text-encorder-for-flux-klein-9b",
|
||||
"split_files/text_encoders/qwen_3_8b.safetensors",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
# ── guidance mapping ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
|
@ -295,6 +315,7 @@ def test_generate_publishes_progress_before_lora_resolution(monkeypatch):
|
|||
def _resolve(
|
||||
active,
|
||||
*,
|
||||
family = None,
|
||||
hf_token = None,
|
||||
cancel_event = None,
|
||||
):
|
||||
|
|
|
|||
|
|
@ -158,6 +158,25 @@ def test_clear_returns_count():
|
|||
assert list(gallery.gallery_dir().glob("*.json")) == []
|
||||
|
||||
|
||||
def test_clear_preserves_orphan_mp4():
|
||||
# An orphan / foreign MP4 (no readable sidecar) is invisible to list_videos; clear must not
|
||||
# destroy it while removing the owned pair.
|
||||
foreign = gallery.gallery_dir() / "recording.mp4"
|
||||
foreign.write_bytes(_mp4())
|
||||
gallery.save(_mp4(), _meta(prompt = "ours"))
|
||||
assert gallery.clear() == 1
|
||||
assert foreign.exists()
|
||||
assert gallery.list_videos() == []
|
||||
|
||||
|
||||
def test_delete_ignores_orphan_mp4():
|
||||
# A per-id delete must refuse an MP4 we do not own (no readable sidecar).
|
||||
foreign = gallery.gallery_dir() / "recording.mp4"
|
||||
foreign.write_bytes(_mp4())
|
||||
assert gallery.delete("recording") is False
|
||||
assert foreign.exists()
|
||||
|
||||
|
||||
def test_list_skips_orphan_mp4_without_sidecar():
|
||||
# An MP4 with no readable json sidecar (a hand-dropped file) is not a record.
|
||||
orphan = gallery.gallery_dir() / "orphan.mp4"
|
||||
|
|
|
|||
|
|
@ -164,6 +164,7 @@ export interface GalleryImage {
|
|||
steps: number;
|
||||
guidance: number;
|
||||
seed: number;
|
||||
batch_seed?: number | null;
|
||||
batch_index: number;
|
||||
batch_size: number;
|
||||
model: string | null;
|
||||
|
|
|
|||
|
|
@ -1287,10 +1287,14 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
setNegativePrompt(image.guidance > 0 ? (image.negative_prompt ?? "") : "");
|
||||
setSteps(image.steps);
|
||||
setGuidance(image.guidance);
|
||||
setSeed(String(image.seed));
|
||||
// Restore from the BASE batch seed, not this image's own seed. The native engine derives
|
||||
// per-image seeds as base + index, so replaying with the derived seed AND the original
|
||||
// batch_size would advance a second time and reproduce a different image. Diffusers shares one
|
||||
// seed, so batch_seed == seed there; older records without batch_seed fall back to seed.
|
||||
setSeed(String(image.batch_seed ?? image.seed));
|
||||
setWidth(image.width);
|
||||
setHeight(image.height);
|
||||
// The batch shared one seed, so image batch_index>0 only reproduces by replaying the
|
||||
// The batch shared one base seed, so image batch_index>0 only reproduces by replaying the
|
||||
// whole batch: restore the batch size too (older recipes without it default to 1).
|
||||
setBatchSize(image.batch_size ?? 1);
|
||||
// Restore the LoRA selection. The recipe stores each adapter as an "id:weight" string;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue