Fix/adjust diffusion: round 5 lifecycle + validation hardening for PR #5754
Round 5 reviewer findings, mostly symmetric-lifecycle and input
validation gaps the earlier rounds left open.
Backend lifecycle (P1)
* routes/training.py: training start now also unloads the GGUF
llama-server subprocess; was previously only unloading the
safetensors backend, so starting training while a GGUF chat
model was loaded kept the subprocess pinned to VRAM.
* routes/inference.py: new _raise_if_training_active helper. Both
GGUF and standard chat loads, plus /api/inference/images/load,
now refuse with HTTP 409 when training is active instead of
silently stopping training to free VRAM.
* core/inference/diffusion.py: _release_other_gpu_owners_for_
diffusion no longer stops active training. The route layer
refuses the request first, so reaching the helper with training
live would only happen from programmatic backend calls; better
to surface OOM than terminate a long training run.
* core/inference/diffusion.py: BF16 dtype is now gated on
torch.cuda.is_bf16_supported. Pascal/Turing GPUs report
is_available()=True but lack BF16 ALUs; FLUX kernels then fail
inside from_pretrained. Falls back to FP16 instead of refusing.
* core/inference/diffusion.py: GGUF transformer allocation and
pipeline allocation now run AFTER releasing chat/export GPU
owners; previously from_single_file ran first and could OOM
before the intended VRAM handoff happened.
* routes/models.py: /delete-cached now also blocks delete when
diffusion is_loading=True (not just is_loaded); concurrent
delete during hf_hub_download / from_single_file would have
raced the rmtree.
* routes/models.py: /delete-finetuned now also checks the
diffusion backend before unlinking a Studio outputs/exports
path. A user who exported a FLUX LoRA locally and loaded it via
/images/load could previously rmtree the directory the
diffusion backend was reading from.
Backend correctness / safety (P2)
* core/inference/diffusion.py: _FAMILY_EXCLUDE for qwen-image now
also covers qwen_image_edit / qwenimageedit underscore spellings
so '...qwen_image_edit-GGUF' no longer misdetects as Qwen-Image.
* core/inference/diffusion.py: detect_family now scans
_FULL_REPO_FAMILIES in addition to _FAMILIES, so SDXL repos
(stabilityai/stable-diffusion-xl-base-1.0) are auto-detected
instead of failing with 'Could not infer a diffusion family'.
* core/inference/diffusion.py: generate_image now uses a separate
_generate_lock for the pipeline forward instead of holding
_lock for the whole call. status() polls and concurrent unload
requests no longer block for the full minutes-long generation.
* routes/models.py: diffusion delete guard now uses exact repo-id
match instead of prefix match; previously loading 'org/model-v2'
would block deleting unrelated cached 'org/model'.
* models/inference.py: DiffusionLoadRequest now rejects ASCII
control characters in repo_id / gguf_filename / base_repo /
family via field_validator (closes log-injection surface from
authenticated callers). Also caps lengths at 256 chars.
* models/inference.py: DiffusionGenerateRequest seed is now
bounded to the int64/uint64 range; previously a huge seed
(e.g. 2**100) passed Pydantic then crashed inside
torch.Generator.manual_seed with 'Overflow when unpacking long
long'.
Frontend (P2)
* features/images/images-page.tsx: Custom HF repo panel now
exposes a Pipeline family override dropdown; previously the
backend supported it via DiffusionLoadRequest.family but the UI
had no way to send it, so custom repos whose names did not
contain a hard-coded substring failed to load.
* features/images/images-page.tsx: handleLoad now re-fetches
status on error. The backend clears its old pipeline before
allocating the replacement; a failed swap previously left the
UI showing 'Loaded:' with Generate enabled until manual
refresh.
Tests (10 new)
* underscore qwen-image-edit exclusion + SDXL full-repo detection
* BF16 fallback when is_bf16_supported() returns False
* status() does not block while generate_image holds _generate_lock
* route layer rejects control chars in repo_id
* route layer rejects 2**100 seeds (uint64-max boundary accepted)
* route layer happy-path with negative-prompt true_cfg_scale
forwarding (Qwen/Flux) and skip-when-no-neg (distilled CFG)
This commit is contained in:
Daniel Han-Chen
parent
f3f3f06dc1
commit
f06895b73e
8 changed files with 458 additions and 48 deletions
|
|
@ -178,8 +178,22 @@ def _smart_base_repo(fam: DiffusionFamily, repo_id: str) -> str:
|
|||
# Qwen-Image-Edit. Each entry maps a family name to substrings that
|
||||
# must NOT appear anywhere in the repo id.
|
||||
_FAMILY_EXCLUDE: dict[str, tuple[str, ...]] = {
|
||||
"stable-diffusion-3": ("3.5", "3-5", "stable-diffusion-3.5"),
|
||||
"qwen-image": ("qwen-image-edit", "qwenimage-edit"),
|
||||
"stable-diffusion-3": (
|
||||
"3.5",
|
||||
"3-5",
|
||||
"3_5",
|
||||
"stable-diffusion-3.5",
|
||||
"stable_diffusion_3_5",
|
||||
),
|
||||
# All underscore / hyphen spellings that appear in Hub repo ids for
|
||||
# the *-Edit family must exclude Qwen-Image, otherwise
|
||||
# ``unsloth/qwen_image_edit-GGUF`` matches the Qwen-Image base.
|
||||
"qwen-image": (
|
||||
"qwen-image-edit",
|
||||
"qwenimage-edit",
|
||||
"qwen_image_edit",
|
||||
"qwenimageedit",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -206,7 +220,10 @@ def detect_family(
|
|||
needle = (repo_id or "").lower()
|
||||
if not needle:
|
||||
return None
|
||||
for fam in _FAMILIES:
|
||||
# Scan _FAMILIES first (GGUF-supported), then _FULL_REPO_FAMILIES
|
||||
# so a repo like ``stabilityai/stable-diffusion-xl-base-1.0`` is
|
||||
# auto-detected as SDXL instead of returning None.
|
||||
for fam in _FAMILIES + _FULL_REPO_FAMILIES:
|
||||
excludes = _FAMILY_EXCLUDE.get(fam.name, ())
|
||||
if any(e in needle for e in excludes):
|
||||
continue
|
||||
|
|
@ -243,15 +260,27 @@ class DiffusionBackend:
|
|||
|
||||
def __init__(self) -> None:
|
||||
self._pipe: Any = None
|
||||
# `_lock` protects mutations to the small state fields and the
|
||||
# pipe call inside generate_image. `_load_lock` serialises the
|
||||
# entire load_model call so two concurrent /images/load requests
|
||||
# cannot both reach pipeline_cls.from_pretrained at the same
|
||||
# time (which would double-spend VRAM and corrupt _pipe). The
|
||||
# locks are taken in order load -> state so a generation in
|
||||
# flight cannot deadlock the next load.
|
||||
# `_lock` protects mutations to the small state fields and is
|
||||
# the only lock taken by status(). It is intentionally NOT held
|
||||
# for the long pipeline forward pass: holding it for the whole
|
||||
# generate would block status() polls (frontend at 1 Hz) and
|
||||
# any concurrent unload requests for minutes at a time.
|
||||
#
|
||||
# `_load_lock` serialises the entire load_model call so two
|
||||
# concurrent /images/load requests cannot both reach
|
||||
# pipeline_cls.from_pretrained at the same time (which would
|
||||
# double-spend VRAM and corrupt _pipe).
|
||||
#
|
||||
# `_generate_lock` serialises pipeline __call__ since diffusers
|
||||
# pipelines are not thread-safe; overlapping forwards on the
|
||||
# shared pipe corrupt internal scheduler state.
|
||||
#
|
||||
# Lock order is load -> state and generate -> state (never
|
||||
# state -> load/generate) so a forward in flight cannot
|
||||
# deadlock the next load or a status poll.
|
||||
self._lock = threading.Lock()
|
||||
self._load_lock = threading.Lock()
|
||||
self._generate_lock = threading.Lock()
|
||||
self._family: Optional[DiffusionFamily] = None
|
||||
self._repo_id: Optional[str] = None
|
||||
self._gguf_path: Optional[str] = None
|
||||
|
|
@ -306,11 +335,23 @@ class DiffusionBackend:
|
|||
validated on. On macOS we use MPS in float16 to keep the pipeline
|
||||
on the Metal GPU. CPU is allowed only as a last resort because
|
||||
running FLUX on CPU is unusably slow (> 10 minutes per image).
|
||||
|
||||
BF16 is gated on ``torch.cuda.is_bf16_supported`` because the
|
||||
Pascal / Turing class (sm_60 / sm_70 / sm_75) reports
|
||||
``is_available() == True`` but lacks BF16 ALUs; FLUX kernels
|
||||
then fail inside ``from_pretrained`` or at the first denoise
|
||||
step. Those cards still work on FP16, so fall back rather than
|
||||
refuse to load.
|
||||
"""
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
return "cuda", torch.bfloat16
|
||||
bf16_ok = False
|
||||
try:
|
||||
bf16_ok = bool(torch.cuda.is_bf16_supported())
|
||||
except Exception:
|
||||
bf16_ok = False
|
||||
return "cuda", torch.bfloat16 if bf16_ok else torch.float16
|
||||
if (
|
||||
hasattr(torch, "backends")
|
||||
and getattr(torch.backends, "mps", None)
|
||||
|
|
@ -430,6 +471,23 @@ class DiffusionBackend:
|
|||
filename = gguf_filename,
|
||||
token = hf_token,
|
||||
)
|
||||
|
||||
# All cheap failure points (bad gguf_filename, missing
|
||||
# pipeline / transformer class, gated download token,
|
||||
# transient Hub error on the GGUF download) have now
|
||||
# been validated. Anything past this line allocates
|
||||
# GPU memory, so release competing GPU owners before
|
||||
# we touch from_single_file or from_pretrained:
|
||||
# * Chat backends (llama-server + safetensors) so the
|
||||
# diffusion transformer does not race them for VRAM.
|
||||
# * Export subprocess (also holds GB on the same GPU).
|
||||
# Training is *not* unloaded here: the route layer
|
||||
# refuses /images/load with HTTP 409 when training is
|
||||
# active so the user keeps their long run.
|
||||
_release_chat_backend_for_diffusion()
|
||||
_release_other_gpu_owners_for_diffusion()
|
||||
|
||||
if gguf_filename:
|
||||
quant_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype)
|
||||
# Diffusers-format GGUFs (FLUX.2 klein / Qwen-Image /
|
||||
# SD3) need the matching base repo's component config
|
||||
|
|
@ -465,15 +523,6 @@ class DiffusionBackend:
|
|||
if hf_token:
|
||||
pipe_kwargs["token"] = hf_token
|
||||
|
||||
# Cheap failure modes (bad gguf_filename, gated token,
|
||||
# transient Hub error) have all happened by now. Only
|
||||
# release the current chat backend + previous diffusion
|
||||
# pipeline right before the expensive allocation so a
|
||||
# typo does not kill the user's loaded chat model. Peak
|
||||
# VRAM still stays at one model's worth because the
|
||||
# release happens before from_pretrained.
|
||||
_release_chat_backend_for_diffusion()
|
||||
_release_other_gpu_owners_for_diffusion()
|
||||
old = self._pipe
|
||||
if old is not None:
|
||||
with self._lock:
|
||||
|
|
@ -562,9 +611,14 @@ class DiffusionBackend:
|
|||
) -> "Any":
|
||||
"""Generate a single PIL image and return it.
|
||||
|
||||
The mutex is held for the entire call: diffusion pipelines are
|
||||
not thread-safe, and overlapping ``__call__``s on a shared
|
||||
pipeline frequently corrupt their internal scheduler state.
|
||||
Concurrent generations are serialised by ``_generate_lock`` so
|
||||
diffusion pipelines (not thread-safe; overlapping ``__call__``s
|
||||
corrupt internal scheduler state) only ever run one at a time.
|
||||
The state ``_lock`` is taken only to snapshot ``_pipe`` /
|
||||
``_device`` and immediately released: holding it for the whole
|
||||
forward pass blocked ``status()`` polls and concurrent unload
|
||||
requests for the entire (minutes-long) generation, which made
|
||||
the UI feel frozen.
|
||||
"""
|
||||
if not prompt or not prompt.strip():
|
||||
raise ValueError("prompt is empty")
|
||||
|
|
@ -586,6 +640,13 @@ class DiffusionBackend:
|
|||
pipe = self._pipe
|
||||
device = self._device or "cpu"
|
||||
|
||||
# _generate_lock outside _lock: only one forward at a time, but
|
||||
# status() / unload() callers do not block on a running forward
|
||||
# pass. unload_model takes _load_lock + _lock; the pipe object
|
||||
# itself is kept alive by the local ``pipe`` reference until
|
||||
# this function returns, so a concurrent unload during forward
|
||||
# cannot free the weights from under us.
|
||||
with self._generate_lock:
|
||||
generator = None
|
||||
if seed is not None:
|
||||
# Match the device of the pipeline so determinism holds
|
||||
|
|
@ -728,20 +789,13 @@ def _release_other_gpu_owners_for_diffusion() -> None:
|
|||
except Exception as exc:
|
||||
logger.debug("export unload skipped: %s", exc)
|
||||
|
||||
# Active training subprocess
|
||||
try:
|
||||
from core.training import get_training_backend # type: ignore
|
||||
|
||||
trn = get_training_backend()
|
||||
if trn.is_training_active():
|
||||
logger.info("Stopping training subprocess before diffusion load")
|
||||
trn.stop_training()
|
||||
for _ in range(60):
|
||||
if not trn.is_training_active():
|
||||
break
|
||||
time.sleep(0.5)
|
||||
except Exception as exc:
|
||||
logger.debug("training unload skipped: %s", exc)
|
||||
# Note: active training is *not* stopped here. The route layer
|
||||
# (`_raise_if_training_active` in routes/inference.py) refuses
|
||||
# /images/load with HTTP 409 before this helper runs, so reaching
|
||||
# this point with training still active would only happen in
|
||||
# programmatic backend calls (tests, scripts). Silently terminating
|
||||
# someone's training run when the diffusion load might still fail
|
||||
# is worse than letting the load OOM and surfacing it explicitly.
|
||||
|
||||
|
||||
def _release(obj: Any) -> None:
|
||||
|
|
|
|||
|
|
@ -1426,6 +1426,27 @@ class AnthropicMessagesResponse(BaseModel):
|
|||
# ── Diffusion image generation ────────────────────────────────────
|
||||
|
||||
|
||||
def _no_control_chars(value: Optional[str], field_name: str) -> Optional[str]:
|
||||
"""Reject newlines and other ASCII control chars in identifiers
|
||||
that get logged before HF validates them.
|
||||
|
||||
Authenticated callers could otherwise inject ``\\n`` / ``\\r`` /
|
||||
NUL into ``logger.info("Loading diffusion model %s", repo_id)``
|
||||
and forge fake log lines. HF repo ids and filenames legitimately
|
||||
contain only ``[A-Za-z0-9._/-]``, so this is also a useful
|
||||
correctness check (catches accidental ``"my repo\\n"`` paste).
|
||||
"""
|
||||
if value is None:
|
||||
return value
|
||||
for ch in value:
|
||||
if ch == "\x7f" or (ord(ch) < 0x20 and ch != "\t"):
|
||||
raise ValueError(
|
||||
f"{field_name} contains control characters; use a plain "
|
||||
"Hugging Face repo / file name."
|
||||
)
|
||||
return value
|
||||
|
||||
|
||||
class DiffusionLoadRequest(BaseModel):
|
||||
"""Load a diffusion image-generation model.
|
||||
|
||||
|
|
@ -1435,16 +1456,20 @@ class DiffusionLoadRequest(BaseModel):
|
|||
VAE / text encoders when loading a GGUF-only repo.
|
||||
"""
|
||||
|
||||
repo_id: str = Field(..., description = "HF repo id")
|
||||
repo_id: str = Field(..., min_length = 1, max_length = 256, description = "HF repo id")
|
||||
gguf_filename: Optional[str] = Field(
|
||||
None, description = "GGUF filename inside repo_id (Q4_K_S, Q8_0, ...)"
|
||||
None,
|
||||
max_length = 256,
|
||||
description = "GGUF filename inside repo_id (Q4_K_S, Q8_0, ...)",
|
||||
)
|
||||
base_repo: Optional[str] = Field(
|
||||
None,
|
||||
max_length = 256,
|
||||
description = "Diffusers base repo to source VAE + text encoders from",
|
||||
)
|
||||
family: Optional[str] = Field(
|
||||
None,
|
||||
max_length = 64,
|
||||
description = "Force pipeline family: flux.2-klein | flux.2 | flux.1 | qwen-image | stable-diffusion-3 | stable-diffusion-xl",
|
||||
)
|
||||
hf_token: Optional[str] = Field(
|
||||
|
|
@ -1455,6 +1480,20 @@ class DiffusionLoadRequest(BaseModel):
|
|||
description = "Offload submodules to CPU between forwards. Trades a small speed hit for ~6 GB less VRAM on FLUX-class models.",
|
||||
)
|
||||
|
||||
@field_validator("repo_id", "gguf_filename", "base_repo", "family")
|
||||
@classmethod
|
||||
def _no_control_chars(cls, v, info):
|
||||
return _no_control_chars(v, info.field_name)
|
||||
|
||||
|
||||
# torch.Generator.manual_seed packs into signed int64; values outside
|
||||
# [-2**63, 2**63 - 1] raise ``Overflow when unpacking long long`` deep
|
||||
# in the C++ layer. uint64 is also routinely cited online so accept
|
||||
# any value the underlying RNG could store and bounce the rest at the
|
||||
# Pydantic layer with a clean error.
|
||||
_SEED_MIN = -(2 ** 63)
|
||||
_SEED_MAX = (2 ** 64) - 1
|
||||
|
||||
|
||||
class DiffusionGenerateRequest(BaseModel):
|
||||
"""Generate a single image from the currently-loaded diffusion model."""
|
||||
|
|
@ -1466,7 +1505,10 @@ class DiffusionGenerateRequest(BaseModel):
|
|||
width: int = Field(1024, ge = 64, le = 2048)
|
||||
height: int = Field(1024, ge = 64, le = 2048)
|
||||
seed: Optional[int] = Field(
|
||||
None, description = "Deterministic seed for reproducible outputs"
|
||||
None,
|
||||
ge = _SEED_MIN,
|
||||
le = _SEED_MAX,
|
||||
description = "Deterministic seed for reproducible outputs",
|
||||
)
|
||||
|
||||
@field_validator("width", "height")
|
||||
|
|
|
|||
|
|
@ -242,6 +242,35 @@ router = APIRouter()
|
|||
studio_router = APIRouter()
|
||||
|
||||
|
||||
def _raise_if_training_active(workload: str) -> None:
|
||||
"""Refuse a chat/diffusion/export load while training is active.
|
||||
|
||||
Without this guard the load path would either (a) silently stop a
|
||||
running training run via _release_other_gpu_owners_for_diffusion
|
||||
or (b) double-spend VRAM and OOM both jobs. Both are worse for the
|
||||
user than a 409 explaining why the request was refused. Best-effort
|
||||
import so unit-test backends without core.training do not 500.
|
||||
"""
|
||||
try:
|
||||
from core.training import get_training_backend # type: ignore
|
||||
except Exception:
|
||||
return
|
||||
try:
|
||||
trn = get_training_backend()
|
||||
if trn.is_training_active():
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = (
|
||||
f"Training is currently active. Stop the training run "
|
||||
f"before loading a {workload} model."
|
||||
),
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.debug("training activity check skipped: %s", exc)
|
||||
|
||||
|
||||
def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict:
|
||||
"""Classify reasoning/tool capabilities via the GGUF classifier so
|
||||
flags match across backends. gpt-oss is overridden because Harmony
|
||||
|
|
@ -737,6 +766,12 @@ async def load_model(
|
|||
detail = "gpu_ids is not supported for GGUF models yet.",
|
||||
)
|
||||
|
||||
# Symmetric lifecycle guard: refuse a chat load while
|
||||
# training is active. Diffusion and export paths refuse;
|
||||
# without this the GGUF chat load would start llama-server
|
||||
# while training still owned VRAM and double-spend it.
|
||||
_raise_if_training_active("chat")
|
||||
|
||||
llama_backend = get_llama_cpp_backend()
|
||||
unsloth_backend = get_inference_backend()
|
||||
|
||||
|
|
@ -928,6 +963,11 @@ async def load_model(
|
|||
)
|
||||
|
||||
# ── Standard path: load via Unsloth/transformers ──────────
|
||||
# Symmetric lifecycle guard: refuse a chat load while training
|
||||
# is active so we do not OOM both the training and inference
|
||||
# jobs together.
|
||||
_raise_if_training_active("chat")
|
||||
|
||||
backend = get_inference_backend()
|
||||
|
||||
# Unload any active GGUF model first
|
||||
|
|
@ -1643,6 +1683,10 @@ async def diffusion_load(
|
|||
desired ``gguf_filename``. Returns the new status payload (same
|
||||
shape as ``/images/status``).
|
||||
"""
|
||||
# Refuse before the long download starts: silently stopping a
|
||||
# running training run to free VRAM was the previous behavior and
|
||||
# left the user with no model loaded plus a dead training job.
|
||||
_raise_if_training_active("diffusion")
|
||||
backend = _get_diffusion_backend()
|
||||
try:
|
||||
status = await asyncio.get_event_loop().run_in_executor(
|
||||
|
|
|
|||
|
|
@ -1968,6 +1968,49 @@ async def delete_finetuned_model(
|
|||
detail = "Could not verify model load status before deleting",
|
||||
) from e
|
||||
|
||||
# Diffusion pipelines can also be loaded directly from a Studio
|
||||
# outputs/exports path (e.g. user fine-tuned a FLUX LoRA, exported
|
||||
# the merged repo locally, then loaded it via /images/load with a
|
||||
# local path as repo_id). Without this guard /delete-finetuned
|
||||
# could rmtree the directory the diffusion backend is reading from.
|
||||
try:
|
||||
from core.inference.diffusion import get_diffusion_backend
|
||||
|
||||
diff_backend = get_diffusion_backend()
|
||||
diff_status = diff_backend.status()
|
||||
if diff_status.get("is_loaded") or diff_status.get("is_loading"):
|
||||
diff_repo = diff_status.get("repo_id") or ""
|
||||
diff_base = diff_status.get("base_repo") or ""
|
||||
target_str = str(target_path)
|
||||
for candidate in (diff_repo, diff_base):
|
||||
if not candidate:
|
||||
continue
|
||||
try:
|
||||
candidate_path = Path(candidate).expanduser()
|
||||
except Exception:
|
||||
continue
|
||||
if not candidate_path.is_absolute():
|
||||
continue
|
||||
try:
|
||||
candidate_resolved = candidate_path.resolve()
|
||||
except Exception:
|
||||
continue
|
||||
if (
|
||||
candidate_resolved == target_path
|
||||
or str(candidate_resolved) == target_str
|
||||
or _is_path_under(candidate_resolved, target_path)
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Unload the diffusion image model before deleting",
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Could not check diffusion backend loaded model before delete: %s", e
|
||||
)
|
||||
|
||||
try:
|
||||
if export_type == "gguf" and gguf_variant:
|
||||
if not target_path.is_dir():
|
||||
|
|
@ -2632,20 +2675,25 @@ async def delete_cached_model(
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
# Also refuse to delete the cache underlying a loaded diffusion
|
||||
# pipeline. The diffusion backend mmap's the GGUF + base repo
|
||||
# weights and continues to read from the cache long after load,
|
||||
# so deleting them out from under it would corrupt generation.
|
||||
# Also refuse to delete the cache underlying a loaded or *loading*
|
||||
# diffusion pipeline. The diffusion backend mmap's the GGUF + base
|
||||
# repo weights and continues to read from the cache long after
|
||||
# load; deleting them out from under it would corrupt generation.
|
||||
# is_loading=True is also blocked because a mid-flight
|
||||
# hf_hub_download / from_single_file would race the rmtree.
|
||||
# Match exactly on repo_id (case-insensitive) instead of prefix to
|
||||
# avoid blocking unrelated deletes like "org/model" while
|
||||
# "org/model-v2" is loaded.
|
||||
try:
|
||||
from core.inference.diffusion import get_diffusion_backend
|
||||
|
||||
diff_backend = get_diffusion_backend()
|
||||
diff_status = diff_backend.status()
|
||||
if diff_status.get("is_loaded"):
|
||||
if diff_status.get("is_loaded") or diff_status.get("is_loading"):
|
||||
diff_repo = (diff_status.get("repo_id") or "").lower()
|
||||
diff_base = (diff_status.get("base_repo") or "").lower()
|
||||
needle = repo_id.lower()
|
||||
if diff_repo.startswith(needle) or diff_base.startswith(needle):
|
||||
if diff_repo == needle or diff_base == needle:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Unload the diffusion image model before deleting",
|
||||
|
|
|
|||
|
|
@ -282,6 +282,23 @@ async def start_training(
|
|||
except Exception as e:
|
||||
logger.warning("Could not unload inference model: %s", e)
|
||||
|
||||
# GGUF chat backend (llama-server subprocess). Without this,
|
||||
# starting training while a GGUF model is loaded keeps the
|
||||
# subprocess pinned to VRAM and OOMs the training job. Mirrors
|
||||
# the symmetric handoffs in routes/inference.py and
|
||||
# routes/export.py.
|
||||
try:
|
||||
from routes.inference import get_llama_cpp_backend
|
||||
|
||||
llama_backend = get_llama_cpp_backend()
|
||||
if getattr(llama_backend, "is_loaded", False):
|
||||
logger.info(
|
||||
"Unloading GGUF chat model to free GPU memory for training"
|
||||
)
|
||||
llama_backend.unload_model()
|
||||
except Exception as e:
|
||||
logger.warning("Could not unload GGUF chat model: %s", e)
|
||||
|
||||
try:
|
||||
from core.export import get_export_backend
|
||||
|
||||
|
|
|
|||
|
|
@ -127,6 +127,25 @@ def test_detect_family_qwen_image_edit_is_not_qwen_image():
|
|||
|
||||
assert detect_family("unsloth/Qwen-Image-Edit-GGUF") is None
|
||||
assert detect_family("unsloth/Qwen-Image-Edit-2509-GGUF") is None
|
||||
# Underscore spellings on the Hub must also be excluded; otherwise
|
||||
# qwen_image_edit-GGUF silently matches the base Qwen-Image family.
|
||||
assert detect_family("unsloth/qwen_image_edit-GGUF") is None
|
||||
assert detect_family("unsloth/QwenImageEdit-GGUF") is None
|
||||
|
||||
|
||||
def test_detect_family_finds_full_repo_sdxl():
|
||||
"""SDXL lives in _FULL_REPO_FAMILIES, but the auto-detector must
|
||||
still find it for ``stabilityai/stable-diffusion-xl-base-1.0`` so
|
||||
the Custom HF repo entry point does not fail with 'Could not infer
|
||||
a diffusion family' for the canonical SDXL repo."""
|
||||
from core.inference.diffusion import detect_family
|
||||
|
||||
fam = detect_family("stabilityai/stable-diffusion-xl-base-1.0")
|
||||
assert fam is not None
|
||||
assert fam.name == "stable-diffusion-xl"
|
||||
fam2 = detect_family("nerijs/sdxl-lora-test")
|
||||
assert fam2 is not None
|
||||
assert fam2.name == "stable-diffusion-xl"
|
||||
|
||||
|
||||
def test_supported_families_payload_shape():
|
||||
|
|
@ -937,3 +956,102 @@ def test_generate_image_skips_true_cfg_scale_without_negative_prompt(monkeypatch
|
|||
assert captured["guidance_scale"] == 7.5
|
||||
# Default left untouched: real CFG only activates with neg prompt.
|
||||
assert captured["true_cfg_scale"] == 4.0
|
||||
|
||||
|
||||
def test_generate_image_does_not_block_status(monkeypatch):
|
||||
"""status() must return promptly while a generation is in flight;
|
||||
holding _lock for the whole forward froze the Images UI on the
|
||||
polling endpoint for the entire (minutes long) generation."""
|
||||
import threading
|
||||
import core.inference.diffusion as d
|
||||
from PIL import Image
|
||||
|
||||
backend = d.get_diffusion_backend()
|
||||
pipe_started = threading.Event()
|
||||
pipe_release = threading.Event()
|
||||
|
||||
class _SlowPipe:
|
||||
def __call__(self, **kw):
|
||||
pipe_started.set()
|
||||
# Wait until the test releases us; status() should return
|
||||
# before this lock is released.
|
||||
pipe_release.wait(timeout = 5)
|
||||
|
||||
class _Out:
|
||||
pass
|
||||
|
||||
o = _Out()
|
||||
o.images = [Image.new("RGB", (kw["width"], kw["height"]), (1, 2, 3))]
|
||||
return o
|
||||
|
||||
backend._pipe = _SlowPipe()
|
||||
backend._device = "cpu"
|
||||
backend._family = d._FAMILIES[0]
|
||||
backend._repo_id = "stub/stub"
|
||||
|
||||
t = threading.Thread(
|
||||
target = backend.generate_image,
|
||||
kwargs = dict(
|
||||
prompt = "a sloth",
|
||||
num_inference_steps = 1,
|
||||
guidance_scale = 1.0,
|
||||
width = 64,
|
||||
height = 64,
|
||||
),
|
||||
)
|
||||
t.start()
|
||||
try:
|
||||
assert pipe_started.wait(timeout = 5)
|
||||
# Forward is in progress; status() must not block on _lock.
|
||||
completed = [False]
|
||||
|
||||
def call_status():
|
||||
backend.status()
|
||||
completed[0] = True
|
||||
|
||||
s = threading.Thread(target = call_status)
|
||||
s.start()
|
||||
s.join(timeout = 2)
|
||||
assert completed[0], "status() blocked on generate_image"
|
||||
finally:
|
||||
pipe_release.set()
|
||||
t.join(timeout = 5)
|
||||
|
||||
|
||||
def test_bf16_falls_back_to_fp16_on_old_cuda(monkeypatch):
|
||||
"""CUDA availability does not imply BF16 support; old GPUs report
|
||||
is_available()=True and is_bf16_supported()=False. The backend
|
||||
must fall back to FP16 rather than picking BF16 and failing
|
||||
deep inside from_pretrained."""
|
||||
import core.inference.diffusion as d
|
||||
|
||||
class _FakeCuda:
|
||||
@staticmethod
|
||||
def is_available():
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def is_bf16_supported():
|
||||
return False
|
||||
|
||||
class _FakeBackends:
|
||||
class mps:
|
||||
@staticmethod
|
||||
def is_available():
|
||||
return False
|
||||
|
||||
class _FakeTorch:
|
||||
cuda = _FakeCuda
|
||||
backends = _FakeBackends
|
||||
# Sentinel objects so the dtype identity comparison works.
|
||||
bfloat16 = object()
|
||||
float16 = object()
|
||||
float32 = object()
|
||||
|
||||
fake_torch = _FakeTorch()
|
||||
monkeypatch.setitem(sys.modules, "torch", fake_torch)
|
||||
|
||||
backend = d.DiffusionBackend()
|
||||
device, dtype = backend._pick_device_and_dtype()
|
||||
assert device == "cuda"
|
||||
assert dtype is fake_torch.float16
|
||||
|
|
|
|||
|
|
@ -188,3 +188,54 @@ def test_unload_clears_state(app_with_stub):
|
|||
assert r.json()["is_loaded"] is False
|
||||
r = c.get("/api/inference/images/status")
|
||||
assert r.json()["is_loaded"] is False
|
||||
|
||||
|
||||
def test_load_rejects_control_chars_in_repo_id(app_with_stub):
|
||||
"""Newline-laden repo ids must be rejected by Pydantic BEFORE the
|
||||
log line that echoes them. Catches log-injection from authenticated
|
||||
callers (issues a 422 instead of forging a fake log line)."""
|
||||
app, _ = app_with_stub
|
||||
c = TestClient(app)
|
||||
r = c.post(
|
||||
"/api/inference/images/load",
|
||||
json = {"repo_id": "owner/model\nFAKE_LOG_LINE"},
|
||||
)
|
||||
assert r.status_code == 422, r.text
|
||||
body = r.json()
|
||||
text = repr(body).lower()
|
||||
assert "control" in text or "repo_id" in text
|
||||
|
||||
|
||||
def test_generate_rejects_oversize_seed(app_with_stub):
|
||||
"""Huge seeds raise inside torch.Generator.manual_seed; Pydantic
|
||||
must clamp first with a 422 instead of a 500 traceback."""
|
||||
app, _ = app_with_stub
|
||||
c = TestClient(app)
|
||||
c.post(
|
||||
"/api/inference/images/load",
|
||||
json = {"repo_id": "unsloth/FLUX.2-klein-4B-GGUF", "gguf_filename": "x.gguf"},
|
||||
)
|
||||
r = c.post(
|
||||
"/api/inference/images/generate",
|
||||
json = {"prompt": "x", "seed": 2 ** 100},
|
||||
)
|
||||
assert r.status_code == 422, r.text
|
||||
|
||||
|
||||
def test_generate_accepts_uint64_max_seed(app_with_stub):
|
||||
"""Boundary value: 2**64 - 1 (uint64 max) is the largest seed
|
||||
torch.Generator on CPU accepts; reject would frustrate users
|
||||
who paste large seeds from other tooling."""
|
||||
app, _ = app_with_stub
|
||||
c = TestClient(app)
|
||||
c.post(
|
||||
"/api/inference/images/load",
|
||||
json = {"repo_id": "unsloth/FLUX.2-klein-4B-GGUF", "gguf_filename": "x.gguf"},
|
||||
)
|
||||
r = c.post(
|
||||
"/api/inference/images/generate",
|
||||
json = {"prompt": "x", "seed": (2 ** 64) - 1},
|
||||
)
|
||||
# The fake backend returns 200 on success; we only care that the
|
||||
# request did NOT 422 on seed bounds.
|
||||
assert r.status_code != 422, r.text
|
||||
|
|
|
|||
|
|
@ -110,6 +110,7 @@ export function ImagesPage() {
|
|||
const [presetIndex, setPresetIndex] = useState(0);
|
||||
const [customRepoId, setCustomRepoId] = useState("");
|
||||
const [customGguf, setCustomGguf] = useState("");
|
||||
const [customFamily, setCustomFamily] = useState<string>("auto");
|
||||
const [useCustom, setUseCustom] = useState(false);
|
||||
const [hfToken, setHfToken] = useState("");
|
||||
|
||||
|
|
@ -151,7 +152,15 @@ export function ImagesPage() {
|
|||
try {
|
||||
const repo = useCustom ? customRepoId.trim() : preset.repo_id;
|
||||
const gguf = useCustom ? customGguf.trim() || undefined : preset.default_gguf;
|
||||
const family = useCustom ? undefined : preset.family;
|
||||
// Custom mode lets the user pin a family explicitly because
|
||||
// detect_family is substring-based and exotic repo names (custom
|
||||
// fine-tunes, third-party mirrors) frequently fail to match.
|
||||
// "auto" leaves the override blank and lets the backend infer.
|
||||
const family = useCustom
|
||||
? customFamily === "auto"
|
||||
? undefined
|
||||
: customFamily
|
||||
: preset.family;
|
||||
// Always pass base_repo for curated entries; custom-repo mode
|
||||
// lets the backend either infer it from the family default or
|
||||
// (when no GGUF is given) treat the repo as a full diffusers
|
||||
|
|
@ -174,10 +183,15 @@ export function ImagesPage() {
|
|||
toast.error("Failed to load image model", {
|
||||
description: err instanceof Error ? err.message : String(err),
|
||||
});
|
||||
// Backend clears its old pipeline before allocating the new one;
|
||||
// a failed swap leaves status.is_loaded=false while our local
|
||||
// copy still says loaded. Re-fetch so Generate disables and the
|
||||
// user does not see a stale "Loaded:" label.
|
||||
await refreshStatus();
|
||||
} finally {
|
||||
setBusy("idle");
|
||||
}
|
||||
}, [useCustom, customRepoId, customGguf, preset, hfToken]);
|
||||
}, [useCustom, customRepoId, customGguf, customFamily, preset, hfToken, refreshStatus]);
|
||||
|
||||
const handleUnload = useCallback(async () => {
|
||||
setBusy("unloading");
|
||||
|
|
@ -320,6 +334,28 @@ export function ImagesPage() {
|
|||
onChange={(e) => setCustomGguf(e.target.value)}
|
||||
placeholder="FLUX.2-klein-4B-Q4_K_S.gguf"
|
||||
/>
|
||||
<Label>Pipeline family (override)</Label>
|
||||
<Select
|
||||
value={customFamily}
|
||||
onValueChange={setCustomFamily}
|
||||
>
|
||||
<SelectTrigger>
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="auto">Auto-detect from repo id</SelectItem>
|
||||
<SelectItem value="flux.2-klein">FLUX.2 klein</SelectItem>
|
||||
<SelectItem value="flux.2">FLUX.2</SelectItem>
|
||||
<SelectItem value="flux.1">FLUX.1</SelectItem>
|
||||
<SelectItem value="qwen-image">Qwen-Image</SelectItem>
|
||||
<SelectItem value="stable-diffusion-3">Stable Diffusion 3</SelectItem>
|
||||
<SelectItem value="stable-diffusion-xl">Stable Diffusion XL</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
<p className="text-xs text-muted-foreground">
|
||||
{"Set this when your repo name does not contain "}
|
||||
{"a recognised family substring (e.g. private fine-tunes)."}
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue