Compare commits
94 commits
main
...
studio-dif
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
07b0cf7d2c | ||
|
|
ca68fd5d13 | ||
|
|
029ca741b4 | ||
|
|
784a9ed71c | ||
|
|
f5186e2b35 | ||
|
|
d0f4bb5165 | ||
|
|
e30c5ed386 | ||
|
|
aeba18dc6d | ||
|
|
09ca2b27d3 | ||
|
|
081377fd30 | ||
|
|
4e1c622d20 | ||
|
|
e3ce1c818e | ||
|
|
a1bec65961 | ||
|
|
90b51cc5c5 | ||
|
|
089749465c | ||
|
|
5350d4cc65 | ||
|
|
cae37123c9 | ||
|
|
91e3a281d8 | ||
|
|
3b60d40f92 | ||
|
|
b8152a5cec | ||
|
|
bec81b882d | ||
|
|
1f5f13c986 | ||
|
|
760bd38dda | ||
|
|
c4c9e2aeec | ||
|
|
79da5d910d | ||
|
|
6c528fb013 | ||
|
|
e17aea6c81 | ||
|
|
65ea3a2c81 | ||
|
|
fd7d334d10 | ||
|
|
4785f76fa2 | ||
|
|
7b5fe1cf10 | ||
|
|
3df9386ac7 | ||
|
|
48740c2664 | ||
|
|
0a7fe59a37 | ||
|
|
c6c4378f38 | ||
|
|
09c51147a9 | ||
|
|
63f3faf022 | ||
|
|
04bd9b2da5 | ||
|
|
ff3bad37fe | ||
|
|
c520a473ae | ||
|
|
c20ed25ec6 | ||
|
|
369573b784 | ||
|
|
da27143520 | ||
|
|
72ec67034c | ||
|
|
e2f41e4069 | ||
|
|
6ac67571dd | ||
|
|
e948a9601c | ||
|
|
2ef9b0e09f | ||
|
|
7c8f1eb40d | ||
|
|
05184ad15a | ||
|
|
2f9bb6929e | ||
|
|
59aa75b8ff | ||
|
|
a9b3d1a672 | ||
|
|
e03ed3dd29 | ||
|
|
f501ab8fc8 | ||
|
|
54adfdff53 | ||
|
|
ff98c6d160 | ||
|
|
ae41bfdbfd | ||
|
|
d8b785a4e2 | ||
|
|
8b8980a607 | ||
|
|
921c60232e | ||
|
|
4b1b149c0b | ||
|
|
1698b66eb1 | ||
|
|
641cdcc13a | ||
|
|
b34fc6258f | ||
|
|
1193c8144a | ||
|
|
0ae20554dc | ||
|
|
c1f9aac510 | ||
|
|
92eccc3627 | ||
|
|
fa8efafcd8 | ||
|
|
0fd9e90c45 | ||
|
|
04de106e49 | ||
|
|
8858104b29 | ||
|
|
f06895b73e | ||
|
|
f3f3f06dc1 | ||
|
|
fb0a31b9a5 | ||
|
|
ec507c5da7 | ||
|
|
0f3ed08351 | ||
|
|
18b50c1f8a | ||
|
|
f44b55c796 | ||
|
|
fbf06bec7a | ||
|
|
65f7a2680d | ||
|
|
0f9b19bb56 | ||
|
|
1601b7828f | ||
|
|
6089720c0c | ||
|
|
8c10cf5f16 | ||
|
|
8074a2b67b | ||
|
|
faa6822039 | ||
|
|
d6f2a238aa | ||
|
|
669964f52c | ||
|
|
bf5c4ac90b | ||
|
|
f8504e3f3c | ||
|
|
a08686cc46 | ||
|
|
b2b660f76f |
29 changed files with 7836 additions and 172 deletions
|
|
@ -81,9 +81,30 @@ huggingfacenotorch = [
|
||||||
"datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0",
|
"datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0",
|
||||||
"accelerate>=0.34.1",
|
"accelerate>=0.34.1",
|
||||||
"peft>=0.18.0,!=0.11.0",
|
"peft>=0.18.0,!=0.11.0",
|
||||||
|
# Round 33 P1: reverted the round-26 hub>=1.3.0 floor. studio.txt
|
||||||
|
# forces hub==0.36.2 to match the transformers 4.57.6 pin in
|
||||||
|
# extras-no-deps.txt; the 1.3.0 floor here was internally
|
||||||
|
# inconsistent and reviewers reproduced the resolver conflict.
|
||||||
|
# Align with the colab-new extra's 0.34.0 floor (line 610). The
|
||||||
|
# transformers-5.x is_offline_mode concern that motivated the
|
||||||
|
# original bump never triggers because transformers is pinned at
|
||||||
|
# 4.57.6 on the supported install path.
|
||||||
"huggingface_hub>=0.34.0",
|
"huggingface_hub>=0.34.0",
|
||||||
"hf_transfer",
|
"hf_transfer",
|
||||||
"diffusers",
|
# Studio Images page depends on Flux2KleinPipeline /
|
||||||
|
# Flux2Pipeline, both shipped in diffusers>=0.37.0. Floor was
|
||||||
|
# missing here so a `pip install unsloth[huggingfacenotorch]`
|
||||||
|
# could resolve to 0.36.0 and fail at runtime when the default
|
||||||
|
# curated FLUX.2 klein model loads.
|
||||||
|
"diffusers>=0.37.0",
|
||||||
|
# diffusers.GGUFQuantizationConfig + from_single_file rely on
|
||||||
|
# the standalone gguf package at runtime. Floor at 0.10.0 to
|
||||||
|
# match the diffusers requirement; older gguf releases raise
|
||||||
|
# at load time. Studio Images default curated picker is
|
||||||
|
# GGUF-only so this must install with the public
|
||||||
|
# huggingfacenotorch extra; missing / under-pinned it makes
|
||||||
|
# /api/inference/images/load 500.
|
||||||
|
"gguf>=0.10.0",
|
||||||
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",
|
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",
|
||||||
"trl>=0.18.2,!=0.19.0,<=0.24.0",
|
"trl>=0.18.2,!=0.19.0,<=0.24.0",
|
||||||
"sentence-transformers",
|
"sentence-transformers",
|
||||||
|
|
|
||||||
2009
studio/backend/core/inference/diffusion.py
Normal file
2009
studio/backend/core/inference/diffusion.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -612,6 +612,18 @@ class LlamaCppBackend:
|
||||||
self._process: Optional[subprocess.Popen] = None
|
self._process: Optional[subprocess.Popen] = None
|
||||||
self._port: Optional[int] = None
|
self._port: Optional[int] = None
|
||||||
self._model_identifier: Optional[str] = None
|
self._model_identifier: Optional[str] = None
|
||||||
|
# Pending-load identifier: set BEFORE _download_gguf starts and
|
||||||
|
# cleared after the load finishes (success or failure). Delete
|
||||||
|
# guards and cross-workload handoff helpers read it via
|
||||||
|
# ``loading_model_identifier`` so a multi-GB HF download cannot
|
||||||
|
# have its cache rmtree'd or be ignored by /images/load,
|
||||||
|
# /training/start, /export/load while it is still resolving.
|
||||||
|
# ``_loading_hf_variant`` mirrors the same lifetime so the
|
||||||
|
# per-variant delete guard at routes/models.py:/delete-finetuned
|
||||||
|
# compares against the NEW variant rather than the previous
|
||||||
|
# loaded ``hf_variant`` (round 15 P1 #2).
|
||||||
|
self._loading_model_identifier: Optional[str] = None
|
||||||
|
self._loading_hf_variant: Optional[str] = None
|
||||||
self._gguf_path: Optional[str] = None
|
self._gguf_path: Optional[str] = None
|
||||||
self._hf_repo: Optional[str] = None
|
self._hf_repo: Optional[str] = None
|
||||||
self._hf_variant: Optional[str] = None
|
self._hf_variant: Optional[str] = None
|
||||||
|
|
@ -713,6 +725,33 @@ class LlamaCppBackend:
|
||||||
def model_identifier(self) -> Optional[str]:
|
def model_identifier(self) -> Optional[str]:
|
||||||
return self._model_identifier
|
return self._model_identifier
|
||||||
|
|
||||||
|
@property
|
||||||
|
def loading_model_identifier(self) -> Optional[str]:
|
||||||
|
"""Identifier of a load currently in progress, or None.
|
||||||
|
|
||||||
|
Populated while ``_download_gguf`` is fetching the GGUF for a
|
||||||
|
new ``load_model`` call. Cleared in the surrounding
|
||||||
|
``finally`` block, so a failed load leaves it None. Delete
|
||||||
|
guards in ``routes/models.py`` and handoff helpers in
|
||||||
|
``routes/inference.py`` consult this so a long HF download
|
||||||
|
cannot have its destination rmtree'd or be ignored by a
|
||||||
|
concurrent /images/load that thinks llama-server is idle."""
|
||||||
|
return self._loading_model_identifier
|
||||||
|
|
||||||
|
@property
|
||||||
|
def loading_hf_variant(self) -> Optional[str]:
|
||||||
|
"""``hf_variant`` of the load currently in progress, or None.
|
||||||
|
|
||||||
|
Mirrors ``loading_model_identifier``'s lifetime so the
|
||||||
|
per-variant delete guards (routes/models.py /delete-cached and
|
||||||
|
/delete-finetuned) can compare against the NEW variant rather
|
||||||
|
than the previously-loaded one (round 15 P1 #2). Without this,
|
||||||
|
a directory with Q4 loaded and Q8 loading would still see the
|
||||||
|
stale Q4 ``hf_variant``, and a Q8 delete would be wrongly
|
||||||
|
allowed even though Q8 is being downloaded into the same
|
||||||
|
directory."""
|
||||||
|
return self._loading_hf_variant
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def is_vision(self) -> bool:
|
def is_vision(self) -> bool:
|
||||||
return self._is_vision
|
return self._is_vision
|
||||||
|
|
@ -2599,7 +2638,68 @@ class LlamaCppBackend:
|
||||||
# Serialise the whole load so concurrent /load calls never
|
# Serialise the whole load so concurrent /load calls never
|
||||||
# leave two llama-server processes alive (#5401 / #5161). Does
|
# leave two llama-server processes alive (#5401 / #5161). Does
|
||||||
# not block /unload, /status, /load-progress.
|
# not block /unload, /status, /load-progress.
|
||||||
|
#
|
||||||
|
# Publish ``_loading_model_identifier`` + ``_loading_hf_variant``
|
||||||
|
# AFTER acquiring ``_serial_load_lock``. Round 15 P1 #1: the
|
||||||
|
# previous round 14 version set them outside the lock so a
|
||||||
|
# second queued ``load_model`` would overwrite or clear the
|
||||||
|
# identifier of the load currently holding the lock, breaking
|
||||||
|
# the delete-safety and GPU handoff guards. Cleared in
|
||||||
|
# ``finally`` so failure / cancellation leaves the pending
|
||||||
|
# state empty. Round 15 P1 #2 added ``_loading_hf_variant``
|
||||||
|
# so per-variant delete guards can compare against the
|
||||||
|
# NEW variant rather than the previous loaded one.
|
||||||
with self._serial_load_lock:
|
with self._serial_load_lock:
|
||||||
|
self._loading_model_identifier = model_identifier
|
||||||
|
self._loading_hf_variant = hf_variant
|
||||||
|
try:
|
||||||
|
return self._load_model_impl_locked(
|
||||||
|
gguf_path = gguf_path,
|
||||||
|
mmproj_path = mmproj_path,
|
||||||
|
hf_repo = hf_repo,
|
||||||
|
hf_variant = hf_variant,
|
||||||
|
hf_token = hf_token,
|
||||||
|
model_identifier = model_identifier,
|
||||||
|
is_vision = is_vision,
|
||||||
|
n_ctx = n_ctx,
|
||||||
|
chat_template_override = chat_template_override,
|
||||||
|
cache_type_kv = cache_type_kv,
|
||||||
|
speculative_type = speculative_type,
|
||||||
|
spec_draft_n_max = spec_draft_n_max,
|
||||||
|
n_threads = n_threads,
|
||||||
|
n_gpu_layers = n_gpu_layers,
|
||||||
|
n_parallel = n_parallel,
|
||||||
|
extra_args = extra_args,
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
self._loading_model_identifier = None
|
||||||
|
self._loading_hf_variant = None
|
||||||
|
|
||||||
|
def _load_model_impl_locked(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
gguf_path: Optional[str] = None,
|
||||||
|
mmproj_path: Optional[str] = None,
|
||||||
|
hf_repo: Optional[str] = None,
|
||||||
|
hf_variant: Optional[str] = None,
|
||||||
|
hf_token: Optional[str] = None,
|
||||||
|
model_identifier: str,
|
||||||
|
is_vision: bool = False,
|
||||||
|
n_ctx: int = 4096,
|
||||||
|
chat_template_override: Optional[str] = None,
|
||||||
|
cache_type_kv: Optional[str] = None,
|
||||||
|
speculative_type: Optional[str] = None,
|
||||||
|
spec_draft_n_max: Optional[int] = None,
|
||||||
|
n_threads: Optional[int] = None,
|
||||||
|
n_gpu_layers: Optional[int] = None,
|
||||||
|
n_parallel: int = 1,
|
||||||
|
extra_args: Optional[List[str]] = None,
|
||||||
|
) -> bool:
|
||||||
|
"""Internal body of ``load_model``. The caller is responsible
|
||||||
|
for holding ``_serial_load_lock`` and for publishing /
|
||||||
|
clearing ``_loading_model_identifier`` + ``_loading_hf_variant``
|
||||||
|
in the surrounding try/finally."""
|
||||||
|
if True:
|
||||||
# Duplicate /load that raced past the route-level check
|
# Duplicate /load that raced past the route-level check
|
||||||
# (the first one hadn't published _healthy=True yet). If the
|
# (the first one hadn't published _healthy=True yet). If the
|
||||||
# live server already satisfies this request, do nothing.
|
# live server already satisfies this request, do nothing.
|
||||||
|
|
|
||||||
|
|
@ -293,6 +293,70 @@ app = FastAPI(
|
||||||
lifespan = lifespan,
|
lifespan = lifespan,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Validation error scrubber ────────────────────────────────────
|
||||||
|
# Round 16 P2 #10: FastAPI's default RequestValidationError handler
|
||||||
|
# echoes the rejected ``input`` value back in the 422 body. A
|
||||||
|
# request like
|
||||||
|
# {"repo_id": "https://hf_token@huggingface.co/owner/repo"}
|
||||||
|
# is rejected by ``DiffusionLoadRequest._no_embedded_hf_tokens``,
|
||||||
|
# but the rejected URL would still appear in the response payload,
|
||||||
|
# leaking the token to the browser console / network log. Wrap the
|
||||||
|
# handler so any ``hf_xxxxx`` substring is replaced with
|
||||||
|
# ``<redacted>`` before serialisation. Scoped to the response body
|
||||||
|
# only; the underlying validator behaviour is unchanged.
|
||||||
|
from fastapi.exceptions import RequestValidationError as _RequestValidationError # noqa: E402
|
||||||
|
from fastapi.encoders import jsonable_encoder as _jsonable_encoder # noqa: E402
|
||||||
|
from fastapi.responses import JSONResponse as _JSONResponse # noqa: E402
|
||||||
|
import re as _re_validation # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
_HF_TOKEN_VALIDATION_RE = _re_validation.compile(r"hf_[A-Za-z0-9]{20,}")
|
||||||
|
|
||||||
|
|
||||||
|
def _scrub_validation_obj(value):
|
||||||
|
"""Recursively scrub ``hf_xxxxx`` tokens out of a value tree.
|
||||||
|
|
||||||
|
Pydantic v2 nests raw ``ValueError`` (and other ``BaseException``)
|
||||||
|
instances under ``ctx.error``. Convert them to scrubbed strings
|
||||||
|
here; otherwise the default ``JSONResponse`` serializer raises
|
||||||
|
``TypeError: Object of type ValueError is not JSON serializable``
|
||||||
|
and the 422 turns into a 500 (round 17 P1 #1). Tuples become
|
||||||
|
lists so the downstream JSON encoder accepts them.
|
||||||
|
"""
|
||||||
|
if isinstance(value, str):
|
||||||
|
return _HF_TOKEN_VALIDATION_RE.sub("<redacted>", value)
|
||||||
|
if isinstance(value, BaseException):
|
||||||
|
return _scrub_validation_obj(str(value))
|
||||||
|
if isinstance(value, tuple):
|
||||||
|
return [_scrub_validation_obj(v) for v in value]
|
||||||
|
if isinstance(value, list):
|
||||||
|
return [_scrub_validation_obj(v) for v in value]
|
||||||
|
if isinstance(value, dict):
|
||||||
|
# Round 21 P2 #7: pydantic surfaces ``input`` for ``string_type``
|
||||||
|
# validation errors verbatim, including dict KEYS like
|
||||||
|
# ``{"hf_xxxxx": "owner/repo"}``. Scrub string keys too so the
|
||||||
|
# token does not leak through the 422 response body.
|
||||||
|
return {
|
||||||
|
(
|
||||||
|
_scrub_validation_obj(k) if isinstance(k, str) else k
|
||||||
|
): _scrub_validation_obj(v)
|
||||||
|
for k, v in value.items()
|
||||||
|
}
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
|
@app.exception_handler(_RequestValidationError)
|
||||||
|
async def _validation_error_scrubbing_handler(request, exc):
|
||||||
|
# ``jsonable_encoder`` walks the scrubbed payload one more time
|
||||||
|
# to convert anything else Pydantic v2 surfaces (URL objects,
|
||||||
|
# Path objects, Url instances, etc.) into JSON-safe primitives.
|
||||||
|
return _JSONResponse(
|
||||||
|
status_code = 422,
|
||||||
|
content = _jsonable_encoder({"detail": _scrub_validation_obj(exc.errors())}),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
# Initialize structured logging
|
# Initialize structured logging
|
||||||
from loggers.config import LogConfig
|
from loggers.config import LogConfig
|
||||||
from loggers.handlers import LoggingMiddleware
|
from loggers.handlers import LoggingMiddleware
|
||||||
|
|
|
||||||
|
|
@ -9,7 +9,13 @@ from __future__ import annotations
|
||||||
|
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from pydantic import BaseModel, Field, model_validator
|
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||||
|
|
||||||
|
# Round 23 P1 #5: identifier hardening reused from the chat models
|
||||||
|
# so /api/data_recipe/publish rejects control characters and
|
||||||
|
# URL-form ``hf_xxxxx`` tokens in ``repo_id`` before they reach
|
||||||
|
# log lines or the HF API.
|
||||||
|
from models.inference import _no_control_chars, _reject_embedded_hf_token
|
||||||
|
|
||||||
|
|
||||||
class RecipePayload(BaseModel):
|
class RecipePayload(BaseModel):
|
||||||
|
|
@ -60,6 +66,16 @@ class PublishDatasetRequest(BaseModel):
|
||||||
description = "Execution artifact path captured by the UI for completed runs",
|
description = "Execution artifact path captured by the UI for completed runs",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@field_validator("repo_id")
|
||||||
|
@classmethod
|
||||||
|
def _no_repo_id_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
@field_validator("repo_id")
|
||||||
|
@classmethod
|
||||||
|
def _no_repo_id_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
|
|
||||||
class PublishDatasetResponse(BaseModel):
|
class PublishDatasetResponse(BaseModel):
|
||||||
success: bool = True
|
success: bool = True
|
||||||
|
|
@ -74,6 +90,20 @@ class SeedInspectRequest(BaseModel):
|
||||||
split: str | None = "train"
|
split: str | None = "train"
|
||||||
preview_size: int = Field(default = 10, ge = 1, le = 50)
|
preview_size: int = Field(default = 10, ge = 1, le = 50)
|
||||||
|
|
||||||
|
# Round 26 P1 #11: dataset_name reaches HF + log/echo paths, so
|
||||||
|
# mirror the hardening other dataset request models already do.
|
||||||
|
# Round 27 P1 #7: split and subset also flow into HF dataset
|
||||||
|
# APIs / errors and must be guarded the same way.
|
||||||
|
@field_validator("dataset_name", "subset", "split")
|
||||||
|
@classmethod
|
||||||
|
def _no_dataset_name_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
@field_validator("dataset_name", "subset", "split")
|
||||||
|
@classmethod
|
||||||
|
def _no_dataset_name_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
|
|
||||||
class SeedInspectUploadRequest(BaseModel):
|
class SeedInspectUploadRequest(BaseModel):
|
||||||
# Legacy single-file flow (mutually exclusive with file_ids)
|
# Legacy single-file flow (mutually exclusive with file_ids)
|
||||||
|
|
@ -89,6 +119,37 @@ class SeedInspectUploadRequest(BaseModel):
|
||||||
unstructured_chunk_size: int | None = Field(default = None, ge = 1, le = 20000)
|
unstructured_chunk_size: int | None = Field(default = None, ge = 1, le = 20000)
|
||||||
unstructured_chunk_overlap: int | None = Field(default = None, ge = 0, le = 20000)
|
unstructured_chunk_overlap: int | None = Field(default = None, ge = 0, le = 20000)
|
||||||
|
|
||||||
|
# Round 30 P1 #6: filename / file_names are reflected as dataset
|
||||||
|
# names + error/log messages; harden them the same way the sibling
|
||||||
|
# SeedInspectRequest hardens dataset_name.
|
||||||
|
@field_validator("filename")
|
||||||
|
@classmethod
|
||||||
|
def _no_filename_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
@field_validator("filename")
|
||||||
|
@classmethod
|
||||||
|
def _no_filename_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
|
@field_validator("file_names")
|
||||||
|
@classmethod
|
||||||
|
def _no_file_names_control_chars(cls, v):
|
||||||
|
if v is None:
|
||||||
|
return v
|
||||||
|
for i, entry in enumerate(v):
|
||||||
|
_no_control_chars(entry, f"file_names[{i}]")
|
||||||
|
return v
|
||||||
|
|
||||||
|
@field_validator("file_names")
|
||||||
|
@classmethod
|
||||||
|
def _no_file_names_embedded_hf_tokens(cls, v):
|
||||||
|
if v is None:
|
||||||
|
return v
|
||||||
|
for i, entry in enumerate(v):
|
||||||
|
_reject_embedded_hf_token(entry, f"file_names[{i}]")
|
||||||
|
return v
|
||||||
|
|
||||||
@model_validator(mode = "after")
|
@model_validator(mode = "after")
|
||||||
def _check_mutual_exclusivity(self) -> "SeedInspectUploadRequest":
|
def _check_mutual_exclusivity(self) -> "SeedInspectUploadRequest":
|
||||||
has_legacy = self.content_base64 is not None
|
has_legacy = self.content_base64 is not None
|
||||||
|
|
|
||||||
|
|
@ -7,7 +7,12 @@ Dataset-related Pydantic models for API requests and responses.
|
||||||
|
|
||||||
from typing import Any, Dict, List, Optional
|
from typing import Any, Dict, List, Optional
|
||||||
|
|
||||||
from pydantic import BaseModel, Field, model_validator
|
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||||
|
|
||||||
|
# Round 24 P1 #11: reuse the chat / diffusion / export identifier
|
||||||
|
# hardening so dataset routes also reject control characters and
|
||||||
|
# URL-embedded HF tokens in user-controlled identifiers.
|
||||||
|
from models.inference import _no_control_chars, _reject_embedded_hf_token
|
||||||
|
|
||||||
|
|
||||||
class CheckFormatRequest(BaseModel):
|
class CheckFormatRequest(BaseModel):
|
||||||
|
|
@ -27,6 +32,18 @@ class CheckFormatRequest(BaseModel):
|
||||||
values.setdefault("train_split", values.pop("split"))
|
values.setdefault("train_split", values.pop("split"))
|
||||||
return values
|
return values
|
||||||
|
|
||||||
|
# Round 27 P1 #6: subset / train_split also flow into HF dataset
|
||||||
|
# APIs and errors/responses, so they need the same hardening.
|
||||||
|
@field_validator("dataset_name", "subset", "train_split")
|
||||||
|
@classmethod
|
||||||
|
def _no_dataset_name_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
@field_validator("dataset_name", "subset", "train_split")
|
||||||
|
@classmethod
|
||||||
|
def _no_dataset_name_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
|
|
||||||
class CheckFormatResponse(BaseModel):
|
class CheckFormatResponse(BaseModel):
|
||||||
"""Response for dataset format check"""
|
"""Response for dataset format check"""
|
||||||
|
|
@ -57,6 +74,16 @@ class AiAssistMappingRequest(BaseModel):
|
||||||
model_name: Optional[str] = None
|
model_name: Optional[str] = None
|
||||||
model_type: Optional[str] = None
|
model_type: Optional[str] = None
|
||||||
|
|
||||||
|
@field_validator("dataset_name", "model_name")
|
||||||
|
@classmethod
|
||||||
|
def _no_identifier_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
@field_validator("dataset_name", "model_name")
|
||||||
|
@classmethod
|
||||||
|
def _no_identifier_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
|
|
||||||
class AiAssistMappingResponse(BaseModel):
|
class AiAssistMappingResponse(BaseModel):
|
||||||
"""Response from LLM-assisted column classification and conversion advice."""
|
"""Response from LLM-assisted column classification and conversion advice."""
|
||||||
|
|
|
||||||
|
|
@ -10,6 +10,13 @@ from pathlib import Path
|
||||||
from pydantic import BaseModel, Field, field_validator
|
from pydantic import BaseModel, Field, field_validator
|
||||||
from typing import List, Optional, Literal, Dict, Any
|
from typing import List, Optional, Literal, Dict, Any
|
||||||
|
|
||||||
|
# Round 23 P1 #1 / #2 / #6: reuse the chat identifier validators
|
||||||
|
# so export requests reject newline / tab / control characters and
|
||||||
|
# URL-form ``hf_xxxxx`` tokens in any user-supplied identifier
|
||||||
|
# (Hub ``repo_id``, ``base_model_id``, the local
|
||||||
|
# ``checkpoint_path``) that flows into log lines or HF API calls.
|
||||||
|
from models.inference import _no_control_chars, _reject_embedded_hf_token
|
||||||
|
|
||||||
|
|
||||||
def _validate_save_directory(value: str) -> str:
|
def _validate_save_directory(value: str) -> str:
|
||||||
"""Reject save_directory values that escape the export root."""
|
"""Reject save_directory values that escape the export root."""
|
||||||
|
|
@ -18,9 +25,17 @@ def _validate_save_directory(value: str) -> str:
|
||||||
raw = str(value).strip()
|
raw = str(value).strip()
|
||||||
if not raw:
|
if not raw:
|
||||||
raise ValueError("save_directory must not be empty")
|
raise ValueError("save_directory must not be empty")
|
||||||
|
# save_directory is logged verbatim by merged / base / GGUF export
|
||||||
|
# flows after resolution, so reject embedded HF tokens at the same
|
||||||
|
# boundary as the sibling identifier fields on export requests.
|
||||||
|
_reject_embedded_hf_token(raw, "save_directory")
|
||||||
if "\x00" in raw:
|
if "\x00" in raw:
|
||||||
raise ValueError("save_directory may not contain null bytes")
|
raise ValueError("save_directory may not contain null bytes")
|
||||||
if any(ch in raw for ch in ("\r", "\n")):
|
# Round 32 P1: reject ALL ASCII control characters (including
|
||||||
|
# TAB / VT / FF) so a caller cannot smuggle log-line breaks or
|
||||||
|
# subprocess argv splitters past the export worker. The earlier
|
||||||
|
# CR / LF check missed every other C0 byte.
|
||||||
|
if any(ord(ch) < 0x20 or ord(ch) == 0x7F for ch in raw):
|
||||||
raise ValueError("save_directory may not contain control characters")
|
raise ValueError("save_directory may not contain control characters")
|
||||||
if len(raw) > 255:
|
if len(raw) > 255:
|
||||||
raise ValueError("save_directory must be <= 255 characters")
|
raise ValueError("save_directory must be <= 255 characters")
|
||||||
|
|
@ -54,6 +69,19 @@ class LoadCheckpointRequest(BaseModel):
|
||||||
description = "Allow loading models with custom code. Only enable for checkpoints/base models you trust.",
|
description = "Allow loading models with custom code. Only enable for checkpoints/base models you trust.",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Round 23 P1 #6: ``checkpoint_path`` is logged verbatim by the
|
||||||
|
# export route. Apply the same control-char + embedded-token
|
||||||
|
# rejection the chat / diffusion / training request models use.
|
||||||
|
@field_validator("checkpoint_path")
|
||||||
|
@classmethod
|
||||||
|
def _no_checkpoint_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
@field_validator("checkpoint_path")
|
||||||
|
@classmethod
|
||||||
|
def _no_checkpoint_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
|
|
||||||
class ExportStatusResponse(BaseModel):
|
class ExportStatusResponse(BaseModel):
|
||||||
"""Current export backend status."""
|
"""Current export backend status."""
|
||||||
|
|
@ -117,6 +145,20 @@ class ExportCommonOptions(BaseModel):
|
||||||
description = "HuggingFace model ID of the base model (for model card metadata)",
|
description = "HuggingFace model ID of the base model (for model card metadata)",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Round 23 P1 #1: ``repo_id`` (Hub destination) and
|
||||||
|
# ``base_model_id`` (model card metadata) both feed log lines
|
||||||
|
# and the HF API. Reject control characters and URL-form
|
||||||
|
# ``hf_xxxxx`` tokens before they reach those sinks.
|
||||||
|
@field_validator("repo_id", "base_model_id")
|
||||||
|
@classmethod
|
||||||
|
def _no_identifier_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
@field_validator("repo_id", "base_model_id")
|
||||||
|
@classmethod
|
||||||
|
def _no_identifier_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
|
|
||||||
class ExportMergedModelRequest(ExportCommonOptions):
|
class ExportMergedModelRequest(ExportCommonOptions):
|
||||||
"""Request for exporting a merged PEFT model."""
|
"""Request for exporting a merged PEFT model."""
|
||||||
|
|
@ -163,6 +205,35 @@ class ExportGGUFRequest(BaseModel):
|
||||||
description = "Hugging Face token for GGUF upload",
|
description = "Hugging Face token for GGUF upload",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Round 23 P1 #2: GGUF export endpoint defines its own
|
||||||
|
# ``repo_id`` (does not inherit from ExportCommonOptions), so
|
||||||
|
# the chat-style hardening needs to be applied here separately.
|
||||||
|
# ``quantization_method`` is forwarded to the export worker
|
||||||
|
# command line, so it gets the control-char check too even
|
||||||
|
# though it does not normally carry tokens.
|
||||||
|
@field_validator("repo_id")
|
||||||
|
@classmethod
|
||||||
|
def _no_repo_id_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
@field_validator("repo_id")
|
||||||
|
@classmethod
|
||||||
|
def _no_repo_id_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
|
@field_validator("quantization_method")
|
||||||
|
@classmethod
|
||||||
|
def _no_quantization_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
# Round 30 P1 #5: quantization_method is forwarded into worker
|
||||||
|
# command lines and reflected in error / success text, so also
|
||||||
|
# reject embedded HF tokens to mirror the repo_id hardening.
|
||||||
|
@field_validator("quantization_method")
|
||||||
|
@classmethod
|
||||||
|
def _no_quantization_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
|
|
||||||
class ExportLoRAAdapterRequest(ExportCommonOptions):
|
class ExportLoRAAdapterRequest(ExportCommonOptions):
|
||||||
"""Request for exporting only the LoRA adapter (not merged)."""
|
"""Request for exporting only the LoRA adapter (not merged)."""
|
||||||
|
|
|
||||||
|
|
@ -60,6 +60,29 @@ class LoadRequest(BaseModel):
|
||||||
return None
|
return None
|
||||||
return value
|
return value
|
||||||
|
|
||||||
|
# Round 20 P1 #5: extend the diffusion-side identifier hardening
|
||||||
|
# (round 5 P2 / round 15 P1 #5) to the chat LoadRequest. Newline
|
||||||
|
# / tab / control characters in ``model_path`` or ``gguf_variant``
|
||||||
|
# would otherwise be echoed verbatim into structured-log lines
|
||||||
|
# ("Loading model %s") and let a caller smuggle in fake log
|
||||||
|
# entries, and an embedded ``hf_...`` token in a URL-form path
|
||||||
|
# would leak the credential into the same log sinks the
|
||||||
|
# diffusion route already redacts.
|
||||||
|
@field_validator("model_path", "gguf_variant")
|
||||||
|
@classmethod
|
||||||
|
def _no_identifier_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
# Round 21 P1 #1: also reject embedded HF tokens in
|
||||||
|
# ``gguf_variant``. A caller can pass a variant string like
|
||||||
|
# ``Q4_K_M-hf_xxxxxxxx`` that flows into log sinks via the
|
||||||
|
# GGUF resolver path; without this only ``model_path`` was
|
||||||
|
# protected.
|
||||||
|
@field_validator("model_path", "gguf_variant")
|
||||||
|
@classmethod
|
||||||
|
def _no_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
cache_type_kv: Optional[str] = Field(
|
cache_type_kv: Optional[str] = Field(
|
||||||
None,
|
None,
|
||||||
description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')",
|
description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')",
|
||||||
|
|
@ -104,12 +127,47 @@ class LoadRequest(BaseModel):
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Round 28 P1 #13: each entry is forwarded verbatim to a logged
|
||||||
|
# subprocess command line and reflected in errors. Reject control
|
||||||
|
# chars and embedded HF tokens for every list entry; allow None.
|
||||||
|
@field_validator("llama_extra_args")
|
||||||
|
@classmethod
|
||||||
|
def _no_extra_args_control_chars(cls, v):
|
||||||
|
if v is None:
|
||||||
|
return v
|
||||||
|
for i, entry in enumerate(v):
|
||||||
|
_no_control_chars(entry, f"llama_extra_args[{i}]")
|
||||||
|
return v
|
||||||
|
|
||||||
|
@field_validator("llama_extra_args")
|
||||||
|
@classmethod
|
||||||
|
def _no_extra_args_embedded_hf_tokens(cls, v):
|
||||||
|
if v is None:
|
||||||
|
return v
|
||||||
|
for i, entry in enumerate(v):
|
||||||
|
_reject_embedded_hf_token(entry, f"llama_extra_args[{i}]")
|
||||||
|
return v
|
||||||
|
|
||||||
|
|
||||||
class UnloadRequest(BaseModel):
|
class UnloadRequest(BaseModel):
|
||||||
"""Request to unload a model"""
|
"""Request to unload a model"""
|
||||||
|
|
||||||
model_path: str = Field(..., description = "Model identifier to unload")
|
model_path: str = Field(..., description = "Model identifier to unload")
|
||||||
|
|
||||||
|
# Round 20 P1 #5: mirror the LoadRequest identifier hardening so
|
||||||
|
# /api/inference/unload also rejects control characters and
|
||||||
|
# URL-embedded HF tokens before the path reaches structured log
|
||||||
|
# sinks.
|
||||||
|
@field_validator("model_path")
|
||||||
|
@classmethod
|
||||||
|
def _no_identifier_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
@field_validator("model_path")
|
||||||
|
@classmethod
|
||||||
|
def _no_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
|
|
||||||
class ValidateModelRequest(BaseModel):
|
class ValidateModelRequest(BaseModel):
|
||||||
"""
|
"""
|
||||||
|
|
@ -130,6 +188,22 @@ class ValidateModelRequest(BaseModel):
|
||||||
None, description = "GGUF quantization variant (e.g. 'Q4_K_M')"
|
None, description = "GGUF quantization variant (e.g. 'Q4_K_M')"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Round 20 P1 #5: same identifier hardening as LoadRequest /
|
||||||
|
# UnloadRequest. /api/inference/validate flows directly into
|
||||||
|
# ``ModelConfig.from_identifier`` and the resulting log lines, so
|
||||||
|
# control characters and embedded HF tokens must not survive.
|
||||||
|
@field_validator("model_path", "gguf_variant")
|
||||||
|
@classmethod
|
||||||
|
def _no_identifier_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
# Round 21 P1 #2: extend embedded-token rejection to
|
||||||
|
# ``gguf_variant`` here too (mirrors LoadRequest).
|
||||||
|
@field_validator("model_path", "gguf_variant")
|
||||||
|
@classmethod
|
||||||
|
def _no_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
|
|
||||||
class ValidateModelResponse(BaseModel):
|
class ValidateModelResponse(BaseModel):
|
||||||
"""
|
"""
|
||||||
|
|
@ -1421,3 +1495,193 @@ class AnthropicMessagesResponse(BaseModel):
|
||||||
stop_reason: Optional[str] = None
|
stop_reason: Optional[str] = None
|
||||||
stop_sequence: Optional[str] = None
|
stop_sequence: Optional[str] = None
|
||||||
usage: AnthropicUsage = Field(default_factory = AnthropicUsage)
|
usage: AnthropicUsage = Field(default_factory = AnthropicUsage)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Diffusion image generation ────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
def _no_control_chars(value: Optional[str], field_name: str) -> Optional[str]:
|
||||||
|
"""Reject newlines, tabs, and other ASCII control chars in
|
||||||
|
identifiers that get logged before HF validates them.
|
||||||
|
|
||||||
|
Authenticated callers could otherwise inject ``\\n`` / ``\\r`` /
|
||||||
|
``\\t`` / 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). Tab is included in the reject set because some logging
|
||||||
|
sinks split fields on tab; allowing it would still let an
|
||||||
|
attacker forge fake columns.
|
||||||
|
"""
|
||||||
|
if value is None:
|
||||||
|
return value
|
||||||
|
for ch in value:
|
||||||
|
if ch == "\x7f" or ord(ch) < 0x20:
|
||||||
|
raise ValueError(
|
||||||
|
f"{field_name} contains control characters; use a plain "
|
||||||
|
"Hugging Face repo / file name."
|
||||||
|
)
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
|
import re as _re
|
||||||
|
|
||||||
|
_EMBEDDED_HF_TOKEN_RE = _re.compile(r"hf_[A-Za-z0-9]{20,}")
|
||||||
|
|
||||||
|
|
||||||
|
def _reject_embedded_hf_token(value: Optional[str], field_name: str) -> Optional[str]:
|
||||||
|
"""Refuse identifiers that contain an embedded ``hf_xxx`` token.
|
||||||
|
|
||||||
|
Round 15 P1 #5: ``repo_id`` and ``base_repo`` accept URL-style
|
||||||
|
strings (``https://hf_token@huggingface.co/owner/repo``). The
|
||||||
|
token would otherwise be stored in ``self._repo_id`` and echoed
|
||||||
|
back through ``status()`` to every authenticated browser session.
|
||||||
|
Log redaction (``_redact_hf_tokens``) covers the logger sink, but
|
||||||
|
the public status payload also needed to refuse the input. Use
|
||||||
|
the dedicated ``hf_token`` field for authentication.
|
||||||
|
"""
|
||||||
|
if value is not None and _EMBEDDED_HF_TOKEN_RE.search(value):
|
||||||
|
raise ValueError(
|
||||||
|
f"{field_name} must not embed a Hugging Face token; "
|
||||||
|
"pass it via the dedicated hf_token field instead."
|
||||||
|
)
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
|
class DiffusionLoadRequest(BaseModel):
|
||||||
|
"""Load a diffusion image-generation model.
|
||||||
|
|
||||||
|
repo_id is the HF repo (either GGUF-only or full diffusers layout).
|
||||||
|
gguf_filename selects the quant when repo_id is a GGUF repo.
|
||||||
|
base_repo overrides the auto-picked diffusers base used for the
|
||||||
|
VAE / text encoders when loading a GGUF-only repo.
|
||||||
|
"""
|
||||||
|
|
||||||
|
# repo_id and base_repo are HF Hub identifiers in this release.
|
||||||
|
# Local-path support is gated behind a frontend / Tauri
|
||||||
|
# ``load-diffusion-model`` directory lease producer that has not
|
||||||
|
# shipped yet (round 32 P1 #3 in the PR reviewer trail). The
|
||||||
|
# 1024-char cap matches POSIX PATH_MAX so future local-path
|
||||||
|
# support can flip on without re-validating the field width.
|
||||||
|
repo_id: str = Field(
|
||||||
|
...,
|
||||||
|
min_length = 1,
|
||||||
|
max_length = 1024,
|
||||||
|
description = (
|
||||||
|
"HF repo id (owner/name). Local filesystem paths are reserved "
|
||||||
|
"for a future native-lease flow and currently rejected by the "
|
||||||
|
"route's _looks_like_local_diffusion_path guard."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
# Round 30 P1 #4: chat /api/inference/load gates native local paths
|
||||||
|
# through a signed native_path_lease grant before the backend
|
||||||
|
# touches the filesystem. Mirror that here so /api/inference/images/
|
||||||
|
# load cannot be used as an authenticated probe for arbitrary
|
||||||
|
# local directories. Optional; Hub ids (no leading slash / tilde)
|
||||||
|
# skip the lease check entirely. The Images UI does not yet
|
||||||
|
# surface a local-path picker, so callers that omit this field
|
||||||
|
# always get the Hub-id code path.
|
||||||
|
native_path_lease: Optional[str] = Field(
|
||||||
|
None,
|
||||||
|
description = "Frontend-visible signed native path grant for a local repo_id",
|
||||||
|
)
|
||||||
|
gguf_filename: Optional[str] = Field(
|
||||||
|
None,
|
||||||
|
max_length = 512,
|
||||||
|
description = "GGUF filename inside repo_id (Q4_K_S, Q8_0, ...)",
|
||||||
|
)
|
||||||
|
base_repo: Optional[str] = Field(
|
||||||
|
None,
|
||||||
|
max_length = 1024,
|
||||||
|
description = (
|
||||||
|
"Diffusers base repo (HF id) for VAE + text encoders. Local "
|
||||||
|
"paths are gated on the same future native-lease flow as "
|
||||||
|
"repo_id."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
base_repo_native_path_lease: Optional[str] = Field(
|
||||||
|
None,
|
||||||
|
description = "Frontend-visible signed native path grant for a local base_repo",
|
||||||
|
)
|
||||||
|
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(
|
||||||
|
None, description = "HuggingFace token for gated models"
|
||||||
|
)
|
||||||
|
enable_model_cpu_offload: bool = Field(
|
||||||
|
True,
|
||||||
|
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)
|
||||||
|
|
||||||
|
@field_validator("repo_id", "gguf_filename", "base_repo")
|
||||||
|
@classmethod
|
||||||
|
def _no_embedded_hf_tokens(cls, v, info):
|
||||||
|
# Round 17 P2 #12: ``gguf_filename`` is forwarded to the
|
||||||
|
# backend and stored on ``DiffusionBackend._gguf_filename``,
|
||||||
|
# which is later surfaced via ``status()`` / log lines. If a
|
||||||
|
# user pastes a URL-form quant path like
|
||||||
|
# ``https://hf_xxxxx@huggingface.co/.../flux.gguf`` we drop
|
||||||
|
# the embedded credential before it can leak.
|
||||||
|
return _reject_embedded_hf_token(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."""
|
||||||
|
|
||||||
|
prompt: str = Field(..., min_length = 1, max_length = 4000)
|
||||||
|
negative_prompt: Optional[str] = Field(None, max_length = 4000)
|
||||||
|
num_inference_steps: int = Field(24, ge = 1, le = 200)
|
||||||
|
guidance_scale: float = Field(3.5, ge = 0.0, le = 20.0)
|
||||||
|
width: int = Field(1024, ge = 64, le = 2048)
|
||||||
|
height: int = Field(1024, ge = 64, le = 2048)
|
||||||
|
seed: Optional[int] = Field(
|
||||||
|
None,
|
||||||
|
ge = _SEED_MIN,
|
||||||
|
le = _SEED_MAX,
|
||||||
|
description = "Deterministic seed for reproducible outputs",
|
||||||
|
)
|
||||||
|
|
||||||
|
@field_validator("width", "height")
|
||||||
|
@classmethod
|
||||||
|
def _multiple_of_eight(cls, v: int) -> int:
|
||||||
|
if v % 8:
|
||||||
|
raise ValueError("width and height must be multiples of 8")
|
||||||
|
return v
|
||||||
|
|
||||||
|
|
||||||
|
class DiffusionGenerateResponse(BaseModel):
|
||||||
|
image_b64: str = Field(..., description = "Base64-encoded PNG")
|
||||||
|
image_mime: str = "image/png"
|
||||||
|
width: int
|
||||||
|
height: int
|
||||||
|
num_inference_steps: int
|
||||||
|
guidance_scale: float
|
||||||
|
# ``seed`` ships as a JSON number for backwards compatibility with
|
||||||
|
# the gallery and existing API consumers, but JavaScript rounds
|
||||||
|
# integers above Number.MAX_SAFE_INTEGER on JSON.parse so seeds
|
||||||
|
# bigger than 2**53 would render different from the value the
|
||||||
|
# backend actually used. ``seed_str`` is the exact decimal
|
||||||
|
# representation; the frontend reads it for reproducibility and
|
||||||
|
# falls back to ``seed`` when not supplied.
|
||||||
|
seed: Optional[int] = None
|
||||||
|
seed_str: Optional[str] = None
|
||||||
|
duration_ms: int
|
||||||
|
model: Optional[str] = None
|
||||||
|
family: Optional[str] = None
|
||||||
|
|
|
||||||
|
|
@ -5,9 +5,11 @@
|
||||||
Pydantic schemas for Model Management API
|
Pydantic schemas for Model Management API
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from pydantic import BaseModel, Field
|
from pydantic import BaseModel, Field, field_validator
|
||||||
from typing import Optional, List, Dict, Any, Literal
|
from typing import Optional, List, Dict, Any, Literal
|
||||||
|
|
||||||
|
from models.inference import _no_control_chars, _reject_embedded_hf_token
|
||||||
|
|
||||||
ModelType = Literal["text", "vision", "audio", "embeddings"]
|
ModelType = Literal["text", "vision", "audio", "embeddings"]
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -206,6 +208,19 @@ class AddScanFolderRequest(BaseModel):
|
||||||
..., description = "Absolute or relative directory path to scan for models"
|
..., description = "Absolute or relative directory path to scan for models"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# path is reflected back in /scan-folders error details and logged
|
||||||
|
# via add_scan_folder_endpoint when the directory is missing, so
|
||||||
|
# apply the same identifier hardening used on other logged paths.
|
||||||
|
@field_validator("path")
|
||||||
|
@classmethod
|
||||||
|
def _no_path_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
@field_validator("path")
|
||||||
|
@classmethod
|
||||||
|
def _no_path_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
|
|
||||||
class ScanFolderInfo(BaseModel):
|
class ScanFolderInfo(BaseModel):
|
||||||
"""A registered custom model scan folder."""
|
"""A registered custom model scan folder."""
|
||||||
|
|
|
||||||
|
|
@ -8,6 +8,13 @@ Pydantic schemas for Training API
|
||||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||||
from typing import Any, Optional, List, Dict, Literal
|
from typing import Any, Optional, List, Dict, Literal
|
||||||
|
|
||||||
|
# Round 22 P1 #1: reuse the chat / diffusion identifier validators
|
||||||
|
# so /api/training/start rejects newline / tab / control characters
|
||||||
|
# and URL-form ``hf_xxxxx`` tokens in ``model_name``. Without these
|
||||||
|
# a caller could log-line-smuggle through "Loading model %s" lines
|
||||||
|
# and leak the bearer token into structured-log sinks.
|
||||||
|
from models.inference import _no_control_chars, _reject_embedded_hf_token
|
||||||
|
|
||||||
|
|
||||||
_MAX_BATCH_SIZE = 4096
|
_MAX_BATCH_SIZE = 4096
|
||||||
_MAX_GRAD_ACCUM = 4096
|
_MAX_GRAD_ACCUM = 4096
|
||||||
|
|
@ -49,6 +56,52 @@ class TrainingStartRequest(BaseModel):
|
||||||
model_name: str = Field(
|
model_name: str = Field(
|
||||||
..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')"
|
..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Identifier hardening: extended progressively across analogous
|
||||||
|
# request models. format_type is copied into training_kwargs and
|
||||||
|
# written into trainer log lines, so it shares the same boundary.
|
||||||
|
@field_validator(
|
||||||
|
"model_name",
|
||||||
|
"hf_dataset",
|
||||||
|
"subset",
|
||||||
|
"train_split",
|
||||||
|
"eval_split",
|
||||||
|
"format_type",
|
||||||
|
)
|
||||||
|
@classmethod
|
||||||
|
def _no_model_name_control_chars(cls, v, info):
|
||||||
|
return _no_control_chars(v, info.field_name)
|
||||||
|
|
||||||
|
@field_validator(
|
||||||
|
"model_name",
|
||||||
|
"hf_dataset",
|
||||||
|
"subset",
|
||||||
|
"train_split",
|
||||||
|
"eval_split",
|
||||||
|
"format_type",
|
||||||
|
)
|
||||||
|
@classmethod
|
||||||
|
def _no_model_name_embedded_hf_tokens(cls, v, info):
|
||||||
|
return _reject_embedded_hf_token(v, info.field_name)
|
||||||
|
|
||||||
|
# local_datasets / local_eval_datasets are user-controlled lists
|
||||||
|
# reflected back in /api/training/start error details when
|
||||||
|
# _validate_local_dataset_paths fails, so the same control-char +
|
||||||
|
# embedded-token guards apply per entry.
|
||||||
|
@field_validator("local_datasets", "local_eval_datasets")
|
||||||
|
@classmethod
|
||||||
|
def _no_local_dataset_control_chars(cls, v, info):
|
||||||
|
for i, entry in enumerate(v or []):
|
||||||
|
_no_control_chars(entry, f"{info.field_name}[{i}]")
|
||||||
|
return v
|
||||||
|
|
||||||
|
@field_validator("local_datasets", "local_eval_datasets")
|
||||||
|
@classmethod
|
||||||
|
def _no_local_dataset_embedded_hf_tokens(cls, v, info):
|
||||||
|
for i, entry in enumerate(v or []):
|
||||||
|
_reject_embedded_hf_token(entry, f"{info.field_name}[{i}]")
|
||||||
|
return v
|
||||||
|
|
||||||
training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = (
|
training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = (
|
||||||
Field(
|
Field(
|
||||||
...,
|
...,
|
||||||
|
|
|
||||||
|
|
@ -43,15 +43,44 @@ safetensors>=0.4.3
|
||||||
datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0
|
datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0
|
||||||
accelerate>=0.34.1
|
accelerate>=0.34.1
|
||||||
peft>=0.18.0,!=0.11.0
|
peft>=0.18.0,!=0.11.0
|
||||||
|
# Round 33 P1: reverted the round-26 hub>=1.3.0 floor to the
|
||||||
|
# pre-PR >=0.34.0 floor. Studio's install_python_stack later
|
||||||
|
# forces hub==0.36.2 via studio.txt (constrained by
|
||||||
|
# transformers==4.57.6 in extras-no-deps.txt), so the 1.3.0
|
||||||
|
# floor was internally inconsistent. extras-no-deps holds
|
||||||
|
# transformers at 4.x, so the transformers-5.x is_offline_mode
|
||||||
|
# concern that motivated the original bump never actually
|
||||||
|
# triggers on the supported install path.
|
||||||
|
# Round 34 P1: the line itself must stay because install.sh
|
||||||
|
# --no-torch installs THIS file with --no-deps and does not run
|
||||||
|
# studio.txt afterward; without the package line a no-torch
|
||||||
|
# install ends with no huggingface_hub at all and the new
|
||||||
|
# diffusion / chat GGUF paths fail with ModuleNotFoundError.
|
||||||
|
# Verified live on B200: hub 0.36.2 + transformers 4.57.6 +
|
||||||
|
# diffusers 0.37.1 imports Flux2KleinPipeline cleanly and runs
|
||||||
|
# end-to-end image generation.
|
||||||
huggingface_hub>=0.34.0
|
huggingface_hub>=0.34.0
|
||||||
hf_transfer
|
hf_transfer
|
||||||
diffusers
|
# Floor 0.37.0 introduces Flux2KleinPipeline + Flux2Pipeline which the
|
||||||
|
# Studio Images page imports for the default curated picker.
|
||||||
|
diffusers>=0.37.0
|
||||||
|
# Required by diffusers.GGUFQuantizationConfig (used by the Images page
|
||||||
|
# to load FLUX.2 / FLUX.1 / Qwen-Image GGUFs from the Hub). Floor at
|
||||||
|
# 0.10.0 to match the diffusers requirement; older gguf releases raise
|
||||||
|
# at single-file load time.
|
||||||
|
gguf>=0.10.0
|
||||||
|
|
||||||
# Transitive deps required because this file is installed with --no-deps.
|
# Transitive deps required because this file is installed with --no-deps.
|
||||||
# Without these, `from transformers import AutoConfig` fails at import time.
|
# Without these, `from transformers import AutoConfig` fails at import time.
|
||||||
regex
|
regex
|
||||||
typing_extensions
|
typing_extensions
|
||||||
filelock
|
filelock
|
||||||
|
# `requests` and its urllib3/charset chain are required by huggingface_hub's
|
||||||
|
# blob downloader; diffusers + GGUFQuantizationConfig 500 on first
|
||||||
|
# /api/inference/images/load otherwise.
|
||||||
|
requests
|
||||||
|
urllib3
|
||||||
|
charset_normalizer
|
||||||
httpx
|
httpx
|
||||||
httpcore
|
httpcore
|
||||||
certifi
|
certifi
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,11 @@
|
||||||
# Studio UI backend dependencies
|
# Studio UI backend dependencies
|
||||||
typer
|
typer
|
||||||
fastapi
|
fastapi
|
||||||
|
# Required by FastAPI's multipart upload route validation
|
||||||
|
# (routes/datasets.py uploads files via UploadFile/File). Without
|
||||||
|
# this, importing the routes package raises RuntimeError on startup
|
||||||
|
# and CPU-only test environments fail before any test runs.
|
||||||
|
python-multipart
|
||||||
uvicorn
|
uvicorn
|
||||||
pydantic
|
pydantic
|
||||||
packaging
|
packaging
|
||||||
|
|
@ -18,3 +23,10 @@ diceware
|
||||||
ddgs
|
ddgs
|
||||||
cryptography>=42.0.0
|
cryptography>=42.0.0
|
||||||
httpx>=0.27.0
|
httpx>=0.27.0
|
||||||
|
# Studio Images page runtime. Flux2KleinPipeline / Flux2Pipeline /
|
||||||
|
# QwenImagePipeline / StableDiffusion3Pipeline are available in
|
||||||
|
# diffusers>=0.37.0, and GGUFQuantizationConfig requires the gguf
|
||||||
|
# package (round 20 P1 #4: fresh standard Studio installs failed on
|
||||||
|
# /images/load because these were only listed in the extras files).
|
||||||
|
diffusers>=0.37.0
|
||||||
|
gguf>=0.10.0
|
||||||
|
|
|
||||||
|
|
@ -433,6 +433,20 @@ async def upload_unstructured_file(
|
||||||
tracked_ids = [fid.strip() for fid in existing_file_ids.split(",") if fid.strip()]
|
tracked_ids = [fid.strip() for fid in existing_file_ids.split(",") if fid.strip()]
|
||||||
|
|
||||||
original_filename = file.filename or "upload"
|
original_filename = file.filename or "upload"
|
||||||
|
# Round 33 P1 #7: file.filename is reflected back to the client,
|
||||||
|
# persisted in the meta JSON, and echoed by error paths. Mirror
|
||||||
|
# the SeedInspectUploadRequest.filename hardening so a multipart
|
||||||
|
# upload cannot smuggle control characters or URL-form HF tokens
|
||||||
|
# through the path the JSON variant already rejects. Import
|
||||||
|
# locally to avoid a routes -> models cycle.
|
||||||
|
from models.inference import _no_control_chars, _reject_embedded_hf_token
|
||||||
|
|
||||||
|
try:
|
||||||
|
_no_control_chars(original_filename, "filename")
|
||||||
|
_reject_embedded_hf_token(original_filename, "filename")
|
||||||
|
except ValueError as exc:
|
||||||
|
raise HTTPException(status_code = 400, detail = str(exc)) from exc
|
||||||
|
|
||||||
ext = Path(original_filename).suffix.lower()
|
ext = Path(original_filename).suffix.lower()
|
||||||
if ext not in UNSTRUCTURED_ALLOWED_EXTS:
|
if ext not in UNSTRUCTURED_ALLOWED_EXTS:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
|
|
|
||||||
|
|
@ -68,11 +68,27 @@ if str(backend_path) not in sys.path:
|
||||||
# Import dataset utilities
|
# Import dataset utilities
|
||||||
from utils.datasets import check_dataset_format
|
from utils.datasets import check_dataset_format
|
||||||
from auth.authentication import get_current_subject
|
from auth.authentication import get_current_subject
|
||||||
|
from models.inference import _no_control_chars, _reject_embedded_hf_token
|
||||||
|
|
||||||
router = APIRouter()
|
router = APIRouter()
|
||||||
logger = get_logger(__name__)
|
logger = get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _validate_logged_identifier(value: str, field_name: str) -> str:
|
||||||
|
"""Round 25 P1 #1: mirror the helper in routes/models.py so the
|
||||||
|
dataset ``/download-progress`` route never reaches logger/cache
|
||||||
|
paths with control characters or embedded HF tokens. Token-shaped
|
||||||
|
strings like ``owner/hf_abcdefghij0123456789`` would otherwise pass
|
||||||
|
the cheap ``_is_valid_repo_id`` regex and end up in warning logs.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
value = _no_control_chars(value, field_name)
|
||||||
|
value = _reject_embedded_hf_token(value, field_name)
|
||||||
|
except ValueError as exc:
|
||||||
|
raise HTTPException(status_code = 422, detail = str(exc)) from exc
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
from models.datasets import (
|
from models.datasets import (
|
||||||
AiAssistMappingRequest,
|
AiAssistMappingRequest,
|
||||||
AiAssistMappingResponse,
|
AiAssistMappingResponse,
|
||||||
|
|
@ -320,7 +336,20 @@ async def upload_dataset(
|
||||||
file: UploadFile,
|
file: UploadFile,
|
||||||
current_subject: str = Depends(get_current_subject),
|
current_subject: str = Depends(get_current_subject),
|
||||||
) -> UploadDatasetResponse:
|
) -> UploadDatasetResponse:
|
||||||
filename = _sanitize_filename(file.filename or "dataset_upload")
|
# Validate the raw multipart filename BEFORE sanitization so smuggled
|
||||||
|
# control characters and embedded HF tokens are rejected at the same
|
||||||
|
# boundary as the JSON path; sanitizing first would silently strip
|
||||||
|
# control chars and let raw inputs pass the validator.
|
||||||
|
raw_filename = file.filename or "dataset_upload"
|
||||||
|
from models.inference import _no_control_chars, _reject_embedded_hf_token
|
||||||
|
|
||||||
|
try:
|
||||||
|
_no_control_chars(raw_filename, "filename")
|
||||||
|
_reject_embedded_hf_token(raw_filename, "filename")
|
||||||
|
except ValueError as exc:
|
||||||
|
raise HTTPException(status_code = 400, detail = str(exc)) from exc
|
||||||
|
|
||||||
|
filename = _sanitize_filename(raw_filename)
|
||||||
ext = Path(filename).suffix.lower()
|
ext = Path(filename).suffix.lower()
|
||||||
if ext not in LOCAL_UPLOAD_EXTS:
|
if ext not in LOCAL_UPLOAD_EXTS:
|
||||||
allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS))
|
allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS))
|
||||||
|
|
@ -370,6 +399,11 @@ async def get_dataset_download_progress(
|
||||||
bytes are observable here. Returns ``cache_path`` so the UI can
|
bytes are observable here. Returns ``cache_path`` so the UI can
|
||||||
show users where the dataset blobs landed on disk.
|
show users where the dataset blobs landed on disk.
|
||||||
"""
|
"""
|
||||||
|
# Round 25 P1 #1: harden ``repo_id`` before it reaches the
|
||||||
|
# ``logger.warning`` line at the bottom (or any future log/cache
|
||||||
|
# path). Matches ``GET /api/models/download-progress`` which
|
||||||
|
# already validates the same parameter in round 24.
|
||||||
|
repo_id = _validate_logged_identifier(repo_id, "repo_id")
|
||||||
_empty = {
|
_empty = {
|
||||||
"downloaded_bytes": 0,
|
"downloaded_bytes": 0,
|
||||||
"expected_bytes": 0,
|
"expected_bytes": 0,
|
||||||
|
|
|
||||||
|
|
@ -50,6 +50,109 @@ router = APIRouter()
|
||||||
logger = get_logger(__name__)
|
logger = get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
import contextlib
|
||||||
|
|
||||||
|
|
||||||
|
def _raise_if_training_active_for_export() -> None:
|
||||||
|
"""409 if a training run is in flight; 503 if status check itself
|
||||||
|
raises. Mirrors the load_checkpoint guard so /export/* and /cleanup
|
||||||
|
never tear down or alter export state while training is using the
|
||||||
|
GPU. Missing core.training is treated as 'no tracker'."""
|
||||||
|
try:
|
||||||
|
from core.training import get_training_backend # type: ignore
|
||||||
|
except Exception as e:
|
||||||
|
logger.debug("core.training not importable, skipping training guard: %s", e)
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
trn = get_training_backend()
|
||||||
|
active = trn.is_training_active()
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning("Could not verify training status before export op: %s", e)
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 503,
|
||||||
|
detail = (
|
||||||
|
"Could not verify training status before the export "
|
||||||
|
"operation. Try again."
|
||||||
|
),
|
||||||
|
) from e
|
||||||
|
if active:
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 409,
|
||||||
|
detail = (
|
||||||
|
"Training is currently active. Stop the training run "
|
||||||
|
"before starting an export operation."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _raise_if_export_active_for_export() -> None:
|
||||||
|
"""409 if another export job is already running; 503 if the status
|
||||||
|
check itself raises. Backends without is_export_active() are
|
||||||
|
treated as 'no tracker available' to stay compatible with mocked
|
||||||
|
backends in tests."""
|
||||||
|
backend = get_export_backend()
|
||||||
|
is_export_active_fn = getattr(backend, "is_export_active", None)
|
||||||
|
if is_export_active_fn is None:
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
export_is_active = bool(is_export_active_fn())
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning("Could not verify export status before export op: %s", e)
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 503,
|
||||||
|
detail = (
|
||||||
|
"Could not verify export status before starting the "
|
||||||
|
"export operation. Try again."
|
||||||
|
),
|
||||||
|
) from e
|
||||||
|
if export_is_active:
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 409,
|
||||||
|
detail = (
|
||||||
|
"An export job is currently active. Wait for it to "
|
||||||
|
"finish before starting another export operation."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@contextlib.asynccontextmanager
|
||||||
|
async def _export_public_window():
|
||||||
|
"""Publish the public-load window across an /export/* operation.
|
||||||
|
|
||||||
|
backend.export_*() runs in a worker thread and does not flip
|
||||||
|
``_export_active = True`` until the worker actually starts; during
|
||||||
|
that gap window another workload that calls ``_release_export_for``
|
||||||
|
would see ``is_export_active() == False`` and tear down the export
|
||||||
|
subprocess. Mirror the load_checkpoint guard so the pending counter
|
||||||
|
is set for the whole export call, and the helper-busy preflight
|
||||||
|
refuses if AI Assist is mid-handoff.
|
||||||
|
|
||||||
|
Also refuses 409 if training or another export is already active so
|
||||||
|
a queued /export/{merged,base,gguf,lora} or /cleanup cannot
|
||||||
|
double-own the GPU with a running training / export job (round 41
|
||||||
|
consensus: load_checkpoint already runs these checks but /export/*
|
||||||
|
and /cleanup were skipping them).
|
||||||
|
"""
|
||||||
|
from routes.inference import (
|
||||||
|
_clear_public_load_window,
|
||||||
|
_raise_if_helper_advisor_busy,
|
||||||
|
)
|
||||||
|
|
||||||
|
export_window_published = False
|
||||||
|
try:
|
||||||
|
_raise_if_training_active_for_export()
|
||||||
|
_raise_if_export_active_for_export()
|
||||||
|
_raise_if_helper_advisor_busy("export")
|
||||||
|
export_window_published = True
|
||||||
|
yield
|
||||||
|
finally:
|
||||||
|
if export_window_published:
|
||||||
|
try:
|
||||||
|
_clear_public_load_window("export")
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
@router.post("/load-checkpoint", response_model = ExportOperationResponse)
|
@router.post("/load-checkpoint", response_model = ExportOperationResponse)
|
||||||
async def load_checkpoint(
|
async def load_checkpoint(
|
||||||
request: LoadCheckpointRequest,
|
request: LoadCheckpointRequest,
|
||||||
|
|
@ -60,50 +163,123 @@ async def load_checkpoint(
|
||||||
|
|
||||||
Wraps ExportBackend.load_checkpoint.
|
Wraps ExportBackend.load_checkpoint.
|
||||||
"""
|
"""
|
||||||
|
# Round 30 P1 #8: track whether we published a public-load pending
|
||||||
|
# entry so the outer finally clears it on either success or
|
||||||
|
# failure path.
|
||||||
|
export_load_window_published = False
|
||||||
try:
|
try:
|
||||||
# Version switching is handled automatically by the subprocess-based
|
# Version switching is handled automatically by the subprocess-based
|
||||||
# export backend — no need for ensure_transformers_version() here.
|
# export backend — no need for ensure_transformers_version() here.
|
||||||
|
|
||||||
# Free GPU memory: shut down any running inference/training subprocesses
|
# Symmetric lifecycle guard: refuse to load an export
|
||||||
# before loading the export checkpoint (they'd compete for VRAM).
|
# checkpoint while training is active so we do not silently
|
||||||
|
# terminate someone's long-running training job and possibly
|
||||||
|
# fail the export load on top of that. Mirrors the
|
||||||
|
# _raise_if_training_active checks in routes/inference.py for
|
||||||
|
# chat and /images/load.
|
||||||
|
# Run BEFORE the chat / inference / diffusion unload helpers
|
||||||
|
# below: otherwise a 409 from this guard would still leave
|
||||||
|
# the user's chat / inference / diffusion GPU owners freed
|
||||||
|
# for nothing, which is the asymmetry round 7 review #5
|
||||||
|
# flagged. Fail-CLOSED (503) when the training backend is
|
||||||
|
# importable but its status check raises.
|
||||||
try:
|
try:
|
||||||
from core.inference import get_inference_backend
|
from core.training import get_training_backend # type: ignore
|
||||||
|
except Exception as e:
|
||||||
inf = get_inference_backend()
|
logger.debug(
|
||||||
if inf.active_model_name:
|
"core.training not importable, skipping export training guard: %s",
|
||||||
logger.info(
|
e,
|
||||||
"Unloading inference model '%s' to free GPU memory for export",
|
)
|
||||||
inf.active_model_name,
|
else:
|
||||||
|
try:
|
||||||
|
trn = get_training_backend()
|
||||||
|
active = trn.is_training_active()
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(
|
||||||
|
"Could not verify training status before export load: %s", e
|
||||||
|
)
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 503,
|
||||||
|
detail = (
|
||||||
|
"Could not verify training status before loading "
|
||||||
|
"an export checkpoint. Try again."
|
||||||
|
),
|
||||||
|
) from e
|
||||||
|
if active:
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 409,
|
||||||
|
detail = (
|
||||||
|
"Training is currently active. Stop the training "
|
||||||
|
"run before loading an export checkpoint."
|
||||||
|
),
|
||||||
)
|
)
|
||||||
inf._shutdown_subprocess()
|
|
||||||
inf.active_model_name = None
|
|
||||||
inf.models.clear()
|
|
||||||
except Exception as e:
|
|
||||||
logger.warning("Could not unload inference model: %s", e)
|
|
||||||
|
|
||||||
try:
|
|
||||||
from core.training import get_training_backend
|
|
||||||
|
|
||||||
trn = get_training_backend()
|
|
||||||
if trn.is_training_active():
|
|
||||||
logger.info("Stopping active training to free GPU memory for export")
|
|
||||||
trn.stop_training()
|
|
||||||
# Wait for training subprocess to actually exit before proceeding,
|
|
||||||
# otherwise it may still hold GPU memory when export tries to load.
|
|
||||||
for _ in range(60): # up to 30s
|
|
||||||
if not trn.is_training_active():
|
|
||||||
break
|
|
||||||
import time
|
|
||||||
|
|
||||||
time.sleep(0.5)
|
|
||||||
else:
|
|
||||||
logger.warning(
|
|
||||||
"Training subprocess did not exit within 30s, proceeding anyway"
|
|
||||||
)
|
|
||||||
except Exception as e:
|
|
||||||
logger.warning("Could not stop training: %s", e)
|
|
||||||
|
|
||||||
backend = get_export_backend()
|
backend = get_export_backend()
|
||||||
|
# Refuse to reload the export checkpoint while an export job
|
||||||
|
# is still running. ``ExportBackend.load_checkpoint`` would
|
||||||
|
# terminate the running subprocess in order to spawn a new
|
||||||
|
# one, silently corrupting the partial output the user is
|
||||||
|
# waiting on (round 13 P1 #1). Runs BEFORE the chat /
|
||||||
|
# diffusion unloads below: a 409 from this guard must not
|
||||||
|
# leave the user's chat or diffusion GPU owners freed for
|
||||||
|
# nothing (round 14 P1 #1). ``is_export_active`` may be
|
||||||
|
# absent on older / mocked backends; treat missing as "no
|
||||||
|
# async-job tracker available" and skip rather than
|
||||||
|
# fail-closed.
|
||||||
|
is_export_active_fn = getattr(backend, "is_export_active", None)
|
||||||
|
if is_export_active_fn is not None:
|
||||||
|
try:
|
||||||
|
export_is_active = bool(is_export_active_fn())
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(
|
||||||
|
"Could not verify export status before export load: %s", e
|
||||||
|
)
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 503,
|
||||||
|
detail = (
|
||||||
|
"Could not verify export status before loading "
|
||||||
|
"an export checkpoint. Try again."
|
||||||
|
),
|
||||||
|
) from e
|
||||||
|
if export_is_active:
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 409,
|
||||||
|
detail = (
|
||||||
|
"An export job is currently active. Stop the "
|
||||||
|
"export job before loading another checkpoint."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Free GPU memory: shut down any chat backend before loading
|
||||||
|
# the export checkpoint. Routes the unload through the shared
|
||||||
|
# helper so we cover llama-server is_active=True and
|
||||||
|
# safetensors loading_models -- the asymmetries round 9
|
||||||
|
# reviews #1, #8, #9 flagged.
|
||||||
|
from routes.inference import (
|
||||||
|
_clear_public_load_window,
|
||||||
|
_raise_if_helper_advisor_busy,
|
||||||
|
_release_chat_for,
|
||||||
|
_release_diffusion_for,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Round 28 P1 #6: refuse before any release fires so AI Assist
|
||||||
|
# busy does not first tear down idle diffusion.
|
||||||
|
# Round 30 P1 #8: also publishes a public-load pending entry so
|
||||||
|
# a concurrent helper / advisor start cannot win the start
|
||||||
|
# lock between our snapshot and load_checkpoint flipping
|
||||||
|
# current_checkpoint / is_export_active.
|
||||||
|
_raise_if_helper_advisor_busy("export")
|
||||||
|
export_load_window_published = True
|
||||||
|
# Round 24 P1 #3: release diffusion BEFORE chat so a failing
|
||||||
|
# diffusion unload does not leave the user with no chat
|
||||||
|
# model loaded. Same reasoning as the training-start flow
|
||||||
|
# (round 18 P1 #8 / round 24 P1 #2). Earlier rounds kept the
|
||||||
|
# chat release first because the helper was best-effort;
|
||||||
|
# now that ``_release_diffusion_for`` is strict it must run
|
||||||
|
# while chat is still resident so a failure preserves it.
|
||||||
|
await _release_diffusion_for("export load")
|
||||||
|
await _release_chat_for("export")
|
||||||
|
|
||||||
# load_checkpoint spawns and waits on a subprocess and can take
|
# load_checkpoint spawns and waits on a subprocess and can take
|
||||||
# minutes. Run it in a worker thread so the event loop stays
|
# minutes. Run it in a worker thread so the event loop stays
|
||||||
# free to serve the live log SSE stream concurrently.
|
# free to serve the live log SSE stream concurrently.
|
||||||
|
|
@ -127,6 +303,18 @@ async def load_checkpoint(
|
||||||
status_code = 500,
|
status_code = 500,
|
||||||
detail = f"Failed to load checkpoint: {str(e)}",
|
detail = f"Failed to load checkpoint: {str(e)}",
|
||||||
)
|
)
|
||||||
|
finally:
|
||||||
|
# Round 30 P1 #8: clear the public-load pending entry once the
|
||||||
|
# load attempt completes (success or failure). Skipped when
|
||||||
|
# the helper-busy check itself raised so the counter stays in
|
||||||
|
# sync with publishes.
|
||||||
|
if export_load_window_published:
|
||||||
|
try:
|
||||||
|
from routes.inference import _clear_public_load_window
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
_clear_public_load_window("export")
|
||||||
|
|
||||||
|
|
||||||
@router.post("/cleanup", response_model = ExportOperationResponse)
|
@router.post("/cleanup", response_model = ExportOperationResponse)
|
||||||
|
|
@ -140,7 +328,12 @@ async def cleanup_export_memory(
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
backend = get_export_backend()
|
backend = get_export_backend()
|
||||||
success = await asyncio.to_thread(backend.cleanup_memory)
|
# Run the cleanup under the same public-load window /export/*
|
||||||
|
# uses so a queued export's handoff gap cannot race a cleanup
|
||||||
|
# call that tears down current_checkpoint. The window also
|
||||||
|
# refuses 409 if training or another export is in flight.
|
||||||
|
async with _export_public_window():
|
||||||
|
success = await asyncio.to_thread(backend.cleanup_memory)
|
||||||
|
|
||||||
if not success:
|
if not success:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
|
|
@ -211,15 +404,16 @@ async def export_merged_model(
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
backend = get_export_backend()
|
backend = get_export_backend()
|
||||||
success, message, output_path = await asyncio.to_thread(
|
async with _export_public_window():
|
||||||
backend.export_merged_model,
|
success, message, output_path = await asyncio.to_thread(
|
||||||
save_directory = request.save_directory,
|
backend.export_merged_model,
|
||||||
format_type = request.format_type,
|
save_directory = request.save_directory,
|
||||||
push_to_hub = request.push_to_hub,
|
format_type = request.format_type,
|
||||||
repo_id = request.repo_id,
|
push_to_hub = request.push_to_hub,
|
||||||
hf_token = request.hf_token,
|
repo_id = request.repo_id,
|
||||||
private = request.private,
|
hf_token = request.hf_token,
|
||||||
)
|
private = request.private,
|
||||||
|
)
|
||||||
|
|
||||||
if not success:
|
if not success:
|
||||||
raise HTTPException(status_code = 400, detail = message)
|
raise HTTPException(status_code = 400, detail = message)
|
||||||
|
|
@ -251,15 +445,16 @@ async def export_base_model(
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
backend = get_export_backend()
|
backend = get_export_backend()
|
||||||
success, message, output_path = await asyncio.to_thread(
|
async with _export_public_window():
|
||||||
backend.export_base_model,
|
success, message, output_path = await asyncio.to_thread(
|
||||||
save_directory = request.save_directory,
|
backend.export_base_model,
|
||||||
push_to_hub = request.push_to_hub,
|
save_directory = request.save_directory,
|
||||||
repo_id = request.repo_id,
|
push_to_hub = request.push_to_hub,
|
||||||
hf_token = request.hf_token,
|
repo_id = request.repo_id,
|
||||||
private = request.private,
|
hf_token = request.hf_token,
|
||||||
base_model_id = request.base_model_id,
|
private = request.private,
|
||||||
)
|
base_model_id = request.base_model_id,
|
||||||
|
)
|
||||||
|
|
||||||
if not success:
|
if not success:
|
||||||
raise HTTPException(status_code = 400, detail = message)
|
raise HTTPException(status_code = 400, detail = message)
|
||||||
|
|
@ -291,14 +486,15 @@ async def export_gguf(
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
backend = get_export_backend()
|
backend = get_export_backend()
|
||||||
success, message, output_path = await asyncio.to_thread(
|
async with _export_public_window():
|
||||||
backend.export_gguf,
|
success, message, output_path = await asyncio.to_thread(
|
||||||
save_directory = request.save_directory,
|
backend.export_gguf,
|
||||||
quantization_method = request.quantization_method,
|
save_directory = request.save_directory,
|
||||||
push_to_hub = request.push_to_hub,
|
quantization_method = request.quantization_method,
|
||||||
repo_id = request.repo_id,
|
push_to_hub = request.push_to_hub,
|
||||||
hf_token = request.hf_token,
|
repo_id = request.repo_id,
|
||||||
)
|
hf_token = request.hf_token,
|
||||||
|
)
|
||||||
|
|
||||||
if not success:
|
if not success:
|
||||||
raise HTTPException(status_code = 400, detail = message)
|
raise HTTPException(status_code = 400, detail = message)
|
||||||
|
|
@ -330,14 +526,15 @@ async def export_lora_adapter(
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
backend = get_export_backend()
|
backend = get_export_backend()
|
||||||
success, message, output_path = await asyncio.to_thread(
|
async with _export_public_window():
|
||||||
backend.export_lora_adapter,
|
success, message, output_path = await asyncio.to_thread(
|
||||||
save_directory = request.save_directory,
|
backend.export_lora_adapter,
|
||||||
push_to_hub = request.push_to_hub,
|
save_directory = request.save_directory,
|
||||||
repo_id = request.repo_id,
|
push_to_hub = request.push_to_hub,
|
||||||
hf_token = request.hf_token,
|
repo_id = request.repo_id,
|
||||||
private = request.private,
|
hf_token = request.hf_token,
|
||||||
)
|
private = request.private,
|
||||||
|
)
|
||||||
|
|
||||||
if not success:
|
if not success:
|
||||||
raise HTTPException(status_code = 400, detail = message)
|
raise HTTPException(status_code = 400, detail = message)
|
||||||
|
|
|
||||||
File diff suppressed because it is too large
Load diff
|
|
@ -134,11 +134,30 @@ from models.responses import (
|
||||||
VisionCheckResponse,
|
VisionCheckResponse,
|
||||||
EmbeddingCheckResponse,
|
EmbeddingCheckResponse,
|
||||||
)
|
)
|
||||||
|
from models.inference import _no_control_chars, _reject_embedded_hf_token
|
||||||
|
|
||||||
router = APIRouter()
|
router = APIRouter()
|
||||||
logger = get_logger(__name__)
|
logger = get_logger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def _validate_logged_identifier(value: str, field_name: str) -> str:
|
||||||
|
"""Round 23 P1 #7 / #8 / #9 / #10: path / query parameters that
|
||||||
|
flow into ``logger.info("... %s", value)`` lines were the last
|
||||||
|
unguarded entry points. Newline / tab / control characters let
|
||||||
|
a caller smuggle forged log entries; URL-form ``hf_xxxxx``
|
||||||
|
tokens would leak into structured-log sinks. Mirror the
|
||||||
|
request-body validators by running both checks here and
|
||||||
|
mapping the validator's ``ValueError`` to HTTP 422 so the
|
||||||
|
client sees the same shape as a Pydantic validation failure.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
value = _no_control_chars(value, field_name)
|
||||||
|
value = _reject_embedded_hf_token(value, field_name)
|
||||||
|
except ValueError as exc:
|
||||||
|
raise HTTPException(status_code = 422, detail = str(exc)) from exc
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
def derive_model_type(
|
def derive_model_type(
|
||||||
is_vision: bool, audio_type: Optional[str], is_embedding: bool = False
|
is_vision: bool, audio_type: Optional[str], is_embedding: bool = False
|
||||||
) -> ModelType:
|
) -> ModelType:
|
||||||
|
|
@ -1571,6 +1590,7 @@ async def get_model_config(
|
||||||
|
|
||||||
This endpoint wraps the backend load_model_defaults function.
|
This endpoint wraps the backend load_model_defaults function.
|
||||||
"""
|
"""
|
||||||
|
model_name = _validate_logged_identifier(model_name, "model_name")
|
||||||
try:
|
try:
|
||||||
if not is_local_path(model_name):
|
if not is_local_path(model_name):
|
||||||
resolved = resolve_cached_repo_id_case(model_name)
|
resolved = resolve_cached_repo_id_case(model_name)
|
||||||
|
|
@ -1580,7 +1600,11 @@ async def get_model_config(
|
||||||
resolved,
|
resolved,
|
||||||
model_name,
|
model_name,
|
||||||
)
|
)
|
||||||
model_name = resolved
|
# Round 23 P1 #7: re-validate the cache-resolved value
|
||||||
|
# (case-only resolver should be a no-op for these
|
||||||
|
# checks, but defend in depth in case the resolver
|
||||||
|
# ever broadens its match heuristic).
|
||||||
|
model_name = _validate_logged_identifier(resolved, "model_name")
|
||||||
|
|
||||||
logger.info(f"Getting model config for: {model_name}")
|
logger.info(f"Getting model config for: {model_name}")
|
||||||
from utils.models.model_config import detect_audio_type
|
from utils.models.model_config import detect_audio_type
|
||||||
|
|
@ -1709,6 +1733,53 @@ def _is_path_under(path: Path, root: Path) -> bool:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def _diffusion_owned_targets(diff_status: dict) -> list[tuple[str, str | None]]:
|
||||||
|
"""Return ``(owned_repo_or_path, owned_gguf_filename)`` pairs for
|
||||||
|
every diffusion target the backend currently holds.
|
||||||
|
|
||||||
|
Pairs the active / pending repo with the active / pending GGUF
|
||||||
|
filename (not the UI-facing collapsed ``gguf_filename``) so the
|
||||||
|
per-variant delete guards know which quant is actually owned by
|
||||||
|
each repo. Without this pairing, a swap in progress (active
|
||||||
|
``Q4_K_S``, pending ``Q8_0``) collapsed both to the pending
|
||||||
|
variant and the active ``Q4_K_S`` GGUF could be deleted while
|
||||||
|
still mmap'd by the resident pipeline (round 13 P1 #3-5).
|
||||||
|
|
||||||
|
Base repos are paired with ``None`` for the GGUF: the base /
|
||||||
|
component repo is loaded whole via ``from_pretrained`` and has no
|
||||||
|
per-variant delete to take advantage of.
|
||||||
|
"""
|
||||||
|
return [
|
||||||
|
(
|
||||||
|
diff_status.get("active_repo_id") or "",
|
||||||
|
diff_status.get("active_gguf_filename"),
|
||||||
|
),
|
||||||
|
(diff_status.get("active_base_repo") or "", None),
|
||||||
|
(
|
||||||
|
diff_status.get("pending_repo_id") or "",
|
||||||
|
diff_status.get("pending_gguf_filename"),
|
||||||
|
),
|
||||||
|
(diff_status.get("pending_base_repo") or "", None),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def _variant_delete_is_safe_for_owned_gguf(
|
||||||
|
requested_variant: str | None,
|
||||||
|
owned_gguf_filename: str | None,
|
||||||
|
) -> bool:
|
||||||
|
"""True iff a per-variant delete for ``requested_variant`` against
|
||||||
|
a repo that owns ``owned_gguf_filename`` cannot remove the owned
|
||||||
|
file.
|
||||||
|
|
||||||
|
Returns False (i.e. unsafe -> block the delete) when either
|
||||||
|
argument is missing so a NULL owned filename or a full-repo delete
|
||||||
|
(no variant) does not accidentally pass the guard."""
|
||||||
|
if not requested_variant or not owned_gguf_filename:
|
||||||
|
return False
|
||||||
|
loaded_label = (_extract_quant_label(owned_gguf_filename.lower()) or "").lower()
|
||||||
|
return bool(loaded_label and loaded_label != requested_variant.lower())
|
||||||
|
|
||||||
|
|
||||||
def _is_path_under_lexically(path: Path, root: Path) -> bool:
|
def _is_path_under_lexically(path: Path, root: Path) -> bool:
|
||||||
"""Check containment without resolving the final path's symlink target."""
|
"""Check containment without resolving the final path's symlink target."""
|
||||||
try:
|
try:
|
||||||
|
|
@ -1724,7 +1795,15 @@ def _loaded_model_matches_deleted_path(active_model: str, deleted_path: Path) ->
|
||||||
try:
|
try:
|
||||||
active = Path(active_model).expanduser().resolve()
|
active = Path(active_model).expanduser().resolve()
|
||||||
target = deleted_path.resolve()
|
target = deleted_path.resolve()
|
||||||
return active == target or (target.is_dir() and active.is_relative_to(target))
|
# Round 27 P1 #8: match bidirectionally so deleting a child
|
||||||
|
# directory of a loaded local model (e.g. .../my-flux/text_encoder
|
||||||
|
# while .../my-flux is loaded) also trips the guard. Mirrors
|
||||||
|
# the diffusion delete-guard pattern.
|
||||||
|
return (
|
||||||
|
active == target
|
||||||
|
or (target.is_dir() and active.is_relative_to(target))
|
||||||
|
or (active.is_dir() and target.is_relative_to(active))
|
||||||
|
)
|
||||||
except (OSError, RuntimeError, ValueError) as e:
|
except (OSError, RuntimeError, ValueError) as e:
|
||||||
logger.debug(
|
logger.debug(
|
||||||
"Could not resolve loaded/deleted model paths; falling back to string comparison: %s",
|
"Could not resolve loaded/deleted model paths; falling back to string comparison: %s",
|
||||||
|
|
@ -1732,8 +1811,10 @@ def _loaded_model_matches_deleted_path(active_model: str, deleted_path: Path) ->
|
||||||
)
|
)
|
||||||
active_lower = active_model.lower()
|
active_lower = active_model.lower()
|
||||||
target_lower = str(deleted_path).lower()
|
target_lower = str(deleted_path).lower()
|
||||||
return active_lower == target_lower or active_lower.startswith(
|
return (
|
||||||
f"{target_lower}{os.sep}"
|
active_lower == target_lower
|
||||||
|
or active_lower.startswith(f"{target_lower}{os.sep}")
|
||||||
|
or target_lower.startswith(f"{active_lower}{os.sep}")
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -1805,6 +1886,14 @@ async def delete_finetuned_model(
|
||||||
Only paths under Studio's outputs/exports roots are accepted. Exported
|
Only paths under Studio's outputs/exports roots are accepted. Exported
|
||||||
GGUF entries can delete one quantization variant at a time.
|
GGUF entries can delete one quantization variant at a time.
|
||||||
"""
|
"""
|
||||||
|
# Round 24 P1 #7 + P2 #13: harden both ``model_path`` and
|
||||||
|
# ``gguf_variant`` for control characters and embedded HF
|
||||||
|
# tokens, mirroring the chat / diffusion / training request
|
||||||
|
# validators. Both fields end up in logger.info(...) lines.
|
||||||
|
model_path = _validate_logged_identifier(model_path, "model_path")
|
||||||
|
if gguf_variant is not None:
|
||||||
|
gguf_variant = _validate_logged_identifier(gguf_variant, "gguf_variant")
|
||||||
|
|
||||||
if source not in {"training", "exported"}:
|
if source not in {"training", "exported"}:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code = 400,
|
status_code = 400,
|
||||||
|
|
@ -1893,6 +1982,32 @@ async def delete_finetuned_model(
|
||||||
from routes.inference import get_llama_cpp_backend
|
from routes.inference import get_llama_cpp_backend
|
||||||
|
|
||||||
llama_backend = get_llama_cpp_backend()
|
llama_backend = get_llama_cpp_backend()
|
||||||
|
# Pending HF GGUF download targeting this path: round 14 P1 #3.
|
||||||
|
# ``loading_model_identifier`` is set before the download starts
|
||||||
|
# and cleared after the subprocess settles, so the user cannot
|
||||||
|
# rmtree the directory llama.cpp is writing into mid-flight.
|
||||||
|
# Round 15 P1 #2: compare against ``loading_hf_variant`` (the
|
||||||
|
# variant being downloaded) rather than ``hf_variant`` (the
|
||||||
|
# PREVIOUS loaded variant, which is stale until the new load
|
||||||
|
# completes its late-metadata update).
|
||||||
|
loading_identifier = getattr(llama_backend, "loading_model_identifier", None)
|
||||||
|
loading_variant = getattr(llama_backend, "loading_hf_variant", None)
|
||||||
|
if (
|
||||||
|
loading_identifier
|
||||||
|
and _loaded_model_matches_deleted_path(
|
||||||
|
loading_identifier,
|
||||||
|
target_path,
|
||||||
|
)
|
||||||
|
and (
|
||||||
|
not gguf_variant
|
||||||
|
or not loading_variant
|
||||||
|
or loading_variant.lower() == gguf_variant.lower()
|
||||||
|
)
|
||||||
|
):
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 409,
|
||||||
|
detail = "Cannot delete a model while it is loading",
|
||||||
|
)
|
||||||
if (
|
if (
|
||||||
llama_backend.is_active
|
llama_backend.is_active
|
||||||
and not llama_backend.is_loaded
|
and not llama_backend.is_loaded
|
||||||
|
|
@ -1968,6 +2083,77 @@ async def delete_finetuned_model(
|
||||||
detail = "Could not verify model load status before deleting",
|
detail = "Could not verify model load status before deleting",
|
||||||
) from e
|
) 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.
|
||||||
|
# is_loading is also blocked: status() exposes pending_repo_id /
|
||||||
|
# pending_base_repo during the load window so deletes during a
|
||||||
|
# mid-flight from_pretrained are refused. During a swap we still
|
||||||
|
# see the previous load's active_repo_id, so every owned path is
|
||||||
|
# checked rather than just the UI-facing one.
|
||||||
|
# Block both DIRECTIONS:
|
||||||
|
# * loaded path is the same as target (or a parent), and
|
||||||
|
# * loaded path is a child of target (so the user cannot rmtree
|
||||||
|
# a parent directory that contains the pipeline's mmap'd file).
|
||||||
|
# Fail-CLOSED on exception (503) like the llama.cpp / safetensors
|
||||||
|
# guards above: an unverifiable diffusion state means we cannot
|
||||||
|
# confirm the target is safe to rmtree.
|
||||||
|
try:
|
||||||
|
from core.inference.diffusion import get_diffusion_backend
|
||||||
|
|
||||||
|
diff_backend = get_diffusion_backend()
|
||||||
|
# include_internal=True so we can iterate active_*/pending_*
|
||||||
|
# raw paths against ``target_path`` (round 16 P1 #5).
|
||||||
|
diff_status = diff_backend.status(include_internal = True)
|
||||||
|
if diff_status.get("is_loaded") or diff_status.get("is_loading"):
|
||||||
|
target_str = str(target_path)
|
||||||
|
# Pair each owned repo / path with the GGUF variant it
|
||||||
|
# actually owns (round 13 P1 #5). For a swap in flight
|
||||||
|
# (active Q4_K_S, pending Q8_0) the active variant must
|
||||||
|
# NOT be deleted just because the pending variant uses
|
||||||
|
# a different quant.
|
||||||
|
for candidate, owned_gguf in _diffusion_owned_targets(diff_status):
|
||||||
|
if not candidate:
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
candidate_resolved = Path(candidate).expanduser().resolve()
|
||||||
|
except Exception:
|
||||||
|
continue
|
||||||
|
# Relative paths (the user can do
|
||||||
|
# `/images/load repo_id=exports/my-flux`) are still
|
||||||
|
# legitimate path candidates; resolve against the
|
||||||
|
# backend cwd so they can be compared with the
|
||||||
|
# absolute ``target_path``. Round 8 review #11.
|
||||||
|
overlaps = (
|
||||||
|
candidate_resolved == target_path
|
||||||
|
or str(candidate_resolved) == target_str
|
||||||
|
or _is_path_under(candidate_resolved, target_path)
|
||||||
|
or _is_path_under(target_path, candidate_resolved)
|
||||||
|
)
|
||||||
|
if not overlaps:
|
||||||
|
continue
|
||||||
|
if export_type == "gguf" and _variant_delete_is_safe_for_owned_gguf(
|
||||||
|
gguf_variant,
|
||||||
|
owned_gguf,
|
||||||
|
):
|
||||||
|
continue
|
||||||
|
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
|
||||||
|
)
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 503,
|
||||||
|
detail = "Could not verify diffusion load status before deleting",
|
||||||
|
) from e
|
||||||
|
|
||||||
try:
|
try:
|
||||||
if export_type == "gguf" and gguf_variant:
|
if export_type == "gguf" and gguf_variant:
|
||||||
if not target_path.is_dir():
|
if not target_path.is_dir():
|
||||||
|
|
@ -2043,6 +2229,9 @@ async def get_lora_base_model(
|
||||||
|
|
||||||
This endpoint wraps the backend get_base_model_from_lora function.
|
This endpoint wraps the backend get_base_model_from_lora function.
|
||||||
"""
|
"""
|
||||||
|
# Round 26 P1 #12: lora_path is echoed back in 404 detail and logs;
|
||||||
|
# harden it the same way other reflected identifiers are.
|
||||||
|
lora_path = _validate_logged_identifier(lora_path, "lora_path")
|
||||||
try:
|
try:
|
||||||
base_model = get_base_model_from_lora(lora_path)
|
base_model = get_base_model_from_lora(lora_path)
|
||||||
|
|
||||||
|
|
@ -2076,6 +2265,7 @@ async def check_vision_model(
|
||||||
|
|
||||||
This endpoint wraps the backend is_vision_model function.
|
This endpoint wraps the backend is_vision_model function.
|
||||||
"""
|
"""
|
||||||
|
model_name = _validate_logged_identifier(model_name, "model_name")
|
||||||
try:
|
try:
|
||||||
logger.info(f"Checking if vision model: {model_name}")
|
logger.info(f"Checking if vision model: {model_name}")
|
||||||
is_vision = is_vision_model(model_name)
|
is_vision = is_vision_model(model_name)
|
||||||
|
|
@ -2104,6 +2294,7 @@ async def check_embedding_model(
|
||||||
|
|
||||||
This endpoint wraps the backend is_embedding_model function.
|
This endpoint wraps the backend is_embedding_model function.
|
||||||
"""
|
"""
|
||||||
|
model_name = _validate_logged_identifier(model_name, "model_name")
|
||||||
try:
|
try:
|
||||||
logger.info(f"Checking if embedding model: {model_name}")
|
logger.info(f"Checking if embedding model: {model_name}")
|
||||||
is_embedding = is_embedding_model(model_name, hf_token = hf_token)
|
is_embedding = is_embedding_model(model_name, hf_token = hf_token)
|
||||||
|
|
@ -2141,6 +2332,7 @@ async def get_gguf_variants(
|
||||||
with file sizes, whether the model supports vision, and the recommended
|
with file sizes, whether the model supports vision, and the recommended
|
||||||
default variant.
|
default variant.
|
||||||
"""
|
"""
|
||||||
|
repo_id = _validate_logged_identifier(repo_id, "repo_id")
|
||||||
try:
|
try:
|
||||||
from utils.models.model_config import is_local_path, list_local_gguf_variants
|
from utils.models.model_config import is_local_path, list_local_gguf_variants
|
||||||
|
|
||||||
|
|
@ -2248,6 +2440,13 @@ async def get_gguf_download_progress(
|
||||||
Tracks completed shard downloads in snapshots and in-progress downloads
|
Tracks completed shard downloads in snapshots and in-progress downloads
|
||||||
in the blobs directory (incomplete files).
|
in the blobs directory (incomplete files).
|
||||||
"""
|
"""
|
||||||
|
# Round 28 P1 #14: mirror the hardening on the generic
|
||||||
|
# /download-progress route. Both repo_id and variant are echoed
|
||||||
|
# into the cache-scan path and can reach logs on the failure
|
||||||
|
# branch via the surrounding try/except.
|
||||||
|
repo_id = _validate_logged_identifier(repo_id, "repo_id")
|
||||||
|
if variant:
|
||||||
|
variant = _validate_logged_identifier(variant, "variant")
|
||||||
try:
|
try:
|
||||||
if not _is_valid_repo_id(repo_id):
|
if not _is_valid_repo_id(repo_id):
|
||||||
return {
|
return {
|
||||||
|
|
@ -2335,6 +2534,10 @@ async def get_download_progress(
|
||||||
"progress": 0,
|
"progress": 0,
|
||||||
"cache_path": None,
|
"cache_path": None,
|
||||||
}
|
}
|
||||||
|
# Round 24 P1 #9: ``repo_id`` flows into log lines deep in
|
||||||
|
# ``_get_repo_size_cached`` on lookup failure, so the same
|
||||||
|
# hardening the request-body models use applies here too.
|
||||||
|
repo_id = _validate_logged_identifier(repo_id, "repo_id")
|
||||||
try:
|
try:
|
||||||
if not _is_valid_repo_id(repo_id):
|
if not _is_valid_repo_id(repo_id):
|
||||||
return _empty
|
return _empty
|
||||||
|
|
@ -2598,39 +2801,283 @@ async def delete_cached_model(
|
||||||
are removed (e.g. ``UD-Q4_K_XL``). Otherwise the entire repo is deleted.
|
are removed (e.g. ``UD-Q4_K_XL``). Otherwise the entire repo is deleted.
|
||||||
Refuses if the model is currently loaded for inference.
|
Refuses if the model is currently loaded for inference.
|
||||||
"""
|
"""
|
||||||
|
# Round 24 P1 #8 + #10: harden both ``repo_id`` and ``variant``
|
||||||
|
# against control characters / embedded HF tokens before they
|
||||||
|
# reach logger.info(...) lines or the HF cache scan.
|
||||||
|
repo_id = _validate_logged_identifier(repo_id, "repo_id")
|
||||||
|
if variant is not None:
|
||||||
|
variant = _validate_logged_identifier(variant, "variant")
|
||||||
if not _is_valid_repo_id(repo_id):
|
if not _is_valid_repo_id(repo_id):
|
||||||
raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
|
raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
|
||||||
|
|
||||||
# Check if model is currently loaded
|
# Round 25 P1 #2 / #3: round 15 added a path-ownership check to
|
||||||
|
# the diffusion guard below, but the llama.cpp and safetensors
|
||||||
|
# guards still only compared logical ``owner/repo`` strings to
|
||||||
|
# the loaded/loading identifier. If a chat or safetensors model
|
||||||
|
# was loaded via a LOCAL HF snapshot path (e.g. through the
|
||||||
|
# ``/load-local-path`` flow), the loaded identifier is the
|
||||||
|
# absolute snapshot path -- ``owner/repo`` never appears there,
|
||||||
|
# the guards passed, and ``DELETE /api/models/delete-cached``
|
||||||
|
# could rmtree an actively mmap'd snapshot.
|
||||||
|
#
|
||||||
|
# Build the HF cache roots for ``repo_id`` ONCE up front and reuse
|
||||||
|
# them in all three guards (llama, safetensors, diffusion). Failure
|
||||||
|
# to scan the cache fails CLOSED on the assumption that we cannot
|
||||||
|
# verify ownership safely; mirrors the diffusion path-scan guard.
|
||||||
|
needle = repo_id.lower()
|
||||||
|
cache_repo_roots: list[Path] = []
|
||||||
|
try:
|
||||||
|
for hf_cache in _all_hf_cache_scans():
|
||||||
|
for repo_info in hf_cache.repos:
|
||||||
|
if (
|
||||||
|
repo_info.repo_type == "model"
|
||||||
|
and repo_info.repo_id.lower() == needle
|
||||||
|
):
|
||||||
|
try:
|
||||||
|
cache_repo_roots.append(
|
||||||
|
Path(repo_info.repo_path).expanduser().resolve()
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
except Exception as cache_scan_exc:
|
||||||
|
logger.warning(
|
||||||
|
"Could not scan HF cache during delete guard preflight: %s",
|
||||||
|
cache_scan_exc,
|
||||||
|
)
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 503,
|
||||||
|
detail = ("Could not verify cache ownership before deleting. Try again."),
|
||||||
|
) from cache_scan_exc
|
||||||
|
|
||||||
|
def _owned_cache_path_matches(value: Optional[str], roots: list[Path]) -> bool:
|
||||||
|
"""Return True if ``value`` resolves to (or contains, or is a
|
||||||
|
child of) any of the HF cache repo roots for the target repo.
|
||||||
|
Used by the llama / safetensors guards to catch local snapshot
|
||||||
|
paths the same way the diffusion guard already does.
|
||||||
|
"""
|
||||||
|
if not value or not roots:
|
||||||
|
return False
|
||||||
|
try:
|
||||||
|
owned = Path(value).expanduser().resolve()
|
||||||
|
except Exception:
|
||||||
|
return False
|
||||||
|
for root in roots:
|
||||||
|
try:
|
||||||
|
if (
|
||||||
|
owned == root
|
||||||
|
or _is_path_under(owned, root)
|
||||||
|
or _is_path_under(root, owned)
|
||||||
|
):
|
||||||
|
return True
|
||||||
|
except Exception:
|
||||||
|
continue
|
||||||
|
return False
|
||||||
|
|
||||||
|
# Round 26 P1 #13 / #14: helper/advisor GGUF loads run on a
|
||||||
|
# PRIVATE LlamaCppBackend, so the global backend below cannot see
|
||||||
|
# them. utils/datasets/llm_assist.py publishes the active repo
|
||||||
|
# via helper_advisor_owns_repo() for exactly this guard. Fail
|
||||||
|
# closed on the variant question (block any variant of the repo)
|
||||||
|
# because helper/advisor flows do not pass a variant through.
|
||||||
|
try:
|
||||||
|
from utils.datasets.llm_assist import helper_advisor_owns_repo
|
||||||
|
|
||||||
|
if helper_advisor_owns_repo(repo_id):
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 409,
|
||||||
|
detail = "Cannot delete a model while AI Assist is using it",
|
||||||
|
)
|
||||||
|
except HTTPException:
|
||||||
|
raise
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(
|
||||||
|
"Could not check helper/advisor backend status before cache delete: %s", e
|
||||||
|
)
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 503,
|
||||||
|
detail = "Could not verify AI Assist load status before deleting cache",
|
||||||
|
) from e
|
||||||
|
|
||||||
|
# Check if model is currently loaded OR loading. is_active and
|
||||||
|
# not is_loaded means an llama-server download / startup is in
|
||||||
|
# flight; the cache delete would race the hf_hub_download / mmap.
|
||||||
|
# Fail CLOSED on exception (503) like the diffusion guard below:
|
||||||
|
# unverifiable load state means we cannot confirm the delete is
|
||||||
|
# safe.
|
||||||
try:
|
try:
|
||||||
from routes.inference import get_llama_cpp_backend
|
from routes.inference import get_llama_cpp_backend
|
||||||
|
|
||||||
llama_backend = get_llama_cpp_backend()
|
llama_backend = get_llama_cpp_backend()
|
||||||
if llama_backend.is_loaded and llama_backend.model_identifier:
|
loaded_id_raw = llama_backend.model_identifier or ""
|
||||||
loaded_id = llama_backend.model_identifier.lower()
|
loaded_id = loaded_id_raw.lower()
|
||||||
if loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()):
|
loading_id_raw = getattr(llama_backend, "loading_model_identifier", None) or ""
|
||||||
|
loading_id = loading_id_raw.lower()
|
||||||
|
loading_variant = (
|
||||||
|
getattr(llama_backend, "loading_hf_variant", None) or ""
|
||||||
|
).lower()
|
||||||
|
# Also consult the pending-load identifier: a multi-GB HF
|
||||||
|
# download stays in ``loading_model_identifier`` until the
|
||||||
|
# download completes, before ``model_identifier`` is set
|
||||||
|
# (round 13 P1 #6). Without this check the cache directory
|
||||||
|
# the download was writing into could be rmtree'd mid-flight.
|
||||||
|
# Round 16 P1 #1: pair against ``loading_hf_variant`` so a
|
||||||
|
# delete of a DIFFERENT cached quant from the same repo
|
||||||
|
# (loading Q4_K_M, deleting cached Q8_0) is allowed; only
|
||||||
|
# block when the requested variant matches what is being
|
||||||
|
# downloaded. Mirrors the /delete-finetuned pairing.
|
||||||
|
requested_variant = (variant or "").lower()
|
||||||
|
# Round 25 P1 #2: also match by HF cache snapshot path so
|
||||||
|
# local-path GGUF chat loads block the cache delete that
|
||||||
|
# owns their snapshot.
|
||||||
|
loading_matches_repo = loading_id == needle or _owned_cache_path_matches(
|
||||||
|
loading_id_raw, cache_repo_roots
|
||||||
|
)
|
||||||
|
if loading_matches_repo:
|
||||||
|
same_loading_variant = (
|
||||||
|
not requested_variant
|
||||||
|
or not loading_variant
|
||||||
|
or requested_variant == loading_variant
|
||||||
|
)
|
||||||
|
if same_loading_variant:
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 409,
|
||||||
|
detail = "Cannot delete a model while it is loading",
|
||||||
|
)
|
||||||
|
# Exact match only (case-insensitive). Prefix match would
|
||||||
|
# block deleting unrelated ``org/model`` while
|
||||||
|
# ``org/model-v2`` is loaded -- same surface the diffusion
|
||||||
|
# guard fixed in round 5. Per-variant deletes that target a
|
||||||
|
# DIFFERENT quant than the loaded one are allowed so the
|
||||||
|
# llama and diffusion paths stay symmetric (round 14 P1 #7).
|
||||||
|
loaded_matches_repo = loaded_id == needle or _owned_cache_path_matches(
|
||||||
|
loaded_id_raw, cache_repo_roots
|
||||||
|
)
|
||||||
|
if loaded_matches_repo and (
|
||||||
|
llama_backend.is_loaded or getattr(llama_backend, "is_active", False)
|
||||||
|
):
|
||||||
|
loaded_variant = (getattr(llama_backend, "hf_variant", None) or "").lower()
|
||||||
|
same_variant = (
|
||||||
|
not requested_variant
|
||||||
|
or not loaded_variant
|
||||||
|
or requested_variant == loaded_variant
|
||||||
|
)
|
||||||
|
if same_variant:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code = 400,
|
status_code = 400,
|
||||||
detail = "Unload the model before deleting",
|
detail = "Unload the model before deleting",
|
||||||
)
|
)
|
||||||
except HTTPException:
|
except HTTPException:
|
||||||
raise
|
raise
|
||||||
except Exception:
|
except Exception as e:
|
||||||
pass
|
logger.warning(
|
||||||
|
"Could not check llama.cpp backend status before cache delete: %s", e
|
||||||
|
)
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 503,
|
||||||
|
detail = "Could not verify llama.cpp load status before deleting cache",
|
||||||
|
) from e
|
||||||
|
|
||||||
try:
|
try:
|
||||||
inference_backend = get_inference_backend()
|
inference_backend = get_inference_backend()
|
||||||
if inference_backend.active_model_name:
|
loading_models = getattr(inference_backend, "loading_models", set()) or set()
|
||||||
active = inference_backend.active_model_name.lower()
|
# Loading set holds model identifiers currently being
|
||||||
if active == repo_id.lower() or active.startswith(repo_id.lower()):
|
# downloaded / instantiated; treat them like active loads
|
||||||
|
# so a delete cannot race a partial mmap.
|
||||||
|
# Exact match only on the logical ``owner/repo`` side, but
|
||||||
|
# also match local snapshot paths (round 25 P1 #3) so a
|
||||||
|
# safetensors model loaded from a local HF snapshot path
|
||||||
|
# cannot have its cache rmtree'd out from under it.
|
||||||
|
for loading_model in loading_models:
|
||||||
|
ml_raw = loading_model or ""
|
||||||
|
ml = ml_raw.lower()
|
||||||
|
if ml == needle or _owned_cache_path_matches(ml_raw, cache_repo_roots):
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 409,
|
||||||
|
detail = "Cannot delete a model while it is loading",
|
||||||
|
)
|
||||||
|
active_model_raw = inference_backend.active_model_name
|
||||||
|
if active_model_raw:
|
||||||
|
active = active_model_raw.lower()
|
||||||
|
if active == needle or _owned_cache_path_matches(
|
||||||
|
active_model_raw, cache_repo_roots
|
||||||
|
):
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code = 400,
|
status_code = 400,
|
||||||
detail = "Unload the model before deleting",
|
detail = "Unload the model before deleting",
|
||||||
)
|
)
|
||||||
except HTTPException:
|
except HTTPException:
|
||||||
raise
|
raise
|
||||||
except Exception:
|
except Exception as e:
|
||||||
pass
|
logger.warning(
|
||||||
|
"Could not check safetensors backend status before cache delete: %s", e
|
||||||
|
)
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 503,
|
||||||
|
detail = "Could not verify safetensors load status before deleting cache",
|
||||||
|
) from e
|
||||||
|
|
||||||
|
# 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.
|
||||||
|
# During a swap (model A loaded, model B loading), status()
|
||||||
|
# exposes both via ``active_*`` and ``pending_*`` so we check
|
||||||
|
# every repo the backend currently owns.
|
||||||
|
# Fail-CLOSED on exception (return 503) like the neighboring
|
||||||
|
# llama.cpp / safetensors guards: we cannot verify whether the
|
||||||
|
# delete is safe, so refuse rather than risk corrupting the
|
||||||
|
# pipeline's mmap.
|
||||||
|
try:
|
||||||
|
from core.inference.diffusion import get_diffusion_backend
|
||||||
|
|
||||||
|
diff_backend = get_diffusion_backend()
|
||||||
|
# include_internal=True so we can pair owned raw paths against
|
||||||
|
# the HF cache snapshot root (round 16 P1 #5).
|
||||||
|
diff_status = diff_backend.status(include_internal = True)
|
||||||
|
if diff_status.get("is_loaded") or diff_status.get("is_loading"):
|
||||||
|
# ``needle`` and ``cache_repo_roots`` come from the
|
||||||
|
# preflight scan above; round 25 deduplicated the
|
||||||
|
# diffusion-specific rescan and now all three guards
|
||||||
|
# share the same fail-closed cache view.
|
||||||
|
#
|
||||||
|
# Pair each owned repo with the GGUF variant it actually
|
||||||
|
# owns (active or pending) so a swap in progress does not
|
||||||
|
# collapse both quants into the pending one (round 13
|
||||||
|
# P1 #4). Per-variant delete is still allowed if the
|
||||||
|
# requested variant differs from the variant that owns
|
||||||
|
# the matched repo.
|
||||||
|
for owned_id, owned_gguf in _diffusion_owned_targets(diff_status):
|
||||||
|
if not owned_id:
|
||||||
|
continue
|
||||||
|
owned_matches_repo = owned_id.lower() == needle
|
||||||
|
if not owned_matches_repo and _owned_cache_path_matches(
|
||||||
|
owned_id, cache_repo_roots
|
||||||
|
):
|
||||||
|
owned_matches_repo = True
|
||||||
|
if not owned_matches_repo:
|
||||||
|
continue
|
||||||
|
if _variant_delete_is_safe_for_owned_gguf(variant, owned_gguf):
|
||||||
|
continue
|
||||||
|
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 status before cache delete: %s",
|
||||||
|
e,
|
||||||
|
)
|
||||||
|
raise HTTPException(
|
||||||
|
status_code = 503,
|
||||||
|
detail = "Could not verify diffusion load status before deleting cache",
|
||||||
|
) from e
|
||||||
|
|
||||||
try:
|
try:
|
||||||
cache_scans = _all_hf_cache_scans()
|
cache_scans = _all_hf_cache_scans()
|
||||||
|
|
|
||||||
|
|
@ -127,6 +127,11 @@ async def start_training(
|
||||||
This endpoint initiates training in the background and returns immediately.
|
This endpoint initiates training in the background and returns immediately.
|
||||||
Use the /status endpoint to check training progress.
|
Use the /status endpoint to check training progress.
|
||||||
"""
|
"""
|
||||||
|
# Round 30 P1 #7: track whether we published a public-load pending
|
||||||
|
# entry so the outer finally clears it on either success or
|
||||||
|
# failure (including any early HTTPException raised by the helper
|
||||||
|
# check itself).
|
||||||
|
training_load_window_published = False
|
||||||
try:
|
try:
|
||||||
logger.info(f"Starting training job with model: {request.model_name}")
|
logger.info(f"Starting training job with model: {request.model_name}")
|
||||||
|
|
||||||
|
|
@ -265,37 +270,48 @@ async def start_training(
|
||||||
)
|
)
|
||||||
training_kwargs["trust_remote_code"] = True
|
training_kwargs["trust_remote_code"] = True
|
||||||
|
|
||||||
# Free GPU memory: shut down any running inference/export subprocesses
|
# Symmetric lifecycle guard: refuse to start training while
|
||||||
# before training starts (they'd compete for VRAM otherwise)
|
# an export job is in flight. Round 10 review #1 -- the
|
||||||
try:
|
# previous code went straight to ``_release_export_for``,
|
||||||
from core.inference import get_inference_backend
|
# which would terminate the in-flight export and corrupt
|
||||||
|
# the user's output artifact. Now we 409 first; the user
|
||||||
|
# stops the export and re-submits.
|
||||||
|
from routes.inference import (
|
||||||
|
_clear_public_load_window,
|
||||||
|
_raise_if_export_active,
|
||||||
|
_raise_if_helper_advisor_busy,
|
||||||
|
_release_chat_for,
|
||||||
|
_release_diffusion_for,
|
||||||
|
_release_export_for,
|
||||||
|
)
|
||||||
|
|
||||||
inf_backend = get_inference_backend()
|
_raise_if_export_active("training")
|
||||||
if inf_backend.active_model_name:
|
# Round 28 P1 #5: refuse before any release fires so AI Assist
|
||||||
logger.info(
|
# busy does not first tear down idle diffusion/export.
|
||||||
"Unloading inference model '%s' to free GPU memory for training",
|
# Round 30 P1 #7: also publishes a public-load pending entry so
|
||||||
inf_backend.active_model_name,
|
# a concurrent helper / advisor start cannot win the start
|
||||||
)
|
# lock between our snapshot and start_training flipping
|
||||||
inf_backend._shutdown_subprocess()
|
# is_training_active. Paired clear lives in the outer
|
||||||
inf_backend.active_model_name = None
|
# ``finally`` below.
|
||||||
inf_backend.models.clear()
|
_raise_if_helper_advisor_busy("training")
|
||||||
except Exception as e:
|
training_load_window_published = True
|
||||||
logger.warning("Could not unload inference model: %s", e)
|
# Round 18 P1 #8: release settled export FIRST so an export
|
||||||
|
# cleanup failure preserves the user's currently loaded chat
|
||||||
|
# model. The previous order (chat -> export) would drop chat
|
||||||
|
# and then refuse training when a wedged idle export raised,
|
||||||
|
# leaving the user with nothing loaded.
|
||||||
|
# Round 24 P1 #2: same reasoning extended to diffusion ->
|
||||||
|
# chat. A wedged diffusion unload used to fire AFTER the chat
|
||||||
|
# backend was already gone, so the user lost both chat and
|
||||||
|
# diffusion on a single failure mode. Order is now
|
||||||
|
# export -> diffusion -> chat, with chat as the last drop so
|
||||||
|
# earlier failures preserve it.
|
||||||
|
await _release_export_for("training")
|
||||||
|
await _release_diffusion_for("training")
|
||||||
|
await _release_chat_for("training")
|
||||||
|
|
||||||
try:
|
# (Diffusion release moved above chat in round 24 P1 #2;
|
||||||
from core.export import get_export_backend
|
# the old trailing call was removed to avoid double-unload.)
|
||||||
|
|
||||||
exp_backend = get_export_backend()
|
|
||||||
if exp_backend.current_checkpoint:
|
|
||||||
logger.info(
|
|
||||||
"Shutting down export subprocess to free GPU memory for training"
|
|
||||||
)
|
|
||||||
exp_backend._shutdown_subprocess()
|
|
||||||
exp_backend.current_checkpoint = None
|
|
||||||
exp_backend.is_vision = False
|
|
||||||
exp_backend.is_peft = False
|
|
||||||
except Exception as e:
|
|
||||||
logger.warning("Could not shut down export subprocess: %s", e)
|
|
||||||
|
|
||||||
# start_training now spawns a subprocess (non-blocking)
|
# start_training now spawns a subprocess (non-blocking)
|
||||||
success = backend.start_training(job_id = job_id, **training_kwargs)
|
success = backend.start_training(job_id = job_id, **training_kwargs)
|
||||||
|
|
@ -319,12 +335,31 @@ async def start_training(
|
||||||
except ValueError as e:
|
except ValueError as e:
|
||||||
logger.warning("Rejected training GPU selection: %s", e)
|
logger.warning("Rejected training GPU selection: %s", e)
|
||||||
raise HTTPException(status_code = 400, detail = str(e))
|
raise HTTPException(status_code = 400, detail = str(e))
|
||||||
|
except HTTPException:
|
||||||
|
# Preserve the intended status code from
|
||||||
|
# _raise_if_training_active / _raise_if_export_active
|
||||||
|
# (409) and the gpu-id 400 raises above. Without this
|
||||||
|
# explicit re-raise the broad ``except Exception`` below
|
||||||
|
# converts a deliberate 409 into a 500.
|
||||||
|
raise
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error starting training: {e}", exc_info = True)
|
logger.error(f"Error starting training: {e}", exc_info = True)
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code = 500,
|
status_code = 500,
|
||||||
detail = f"Failed to start training: {str(e)}",
|
detail = f"Failed to start training: {str(e)}",
|
||||||
)
|
)
|
||||||
|
finally:
|
||||||
|
# Round 30 P1 #7: clear the public-load pending entry once the
|
||||||
|
# start attempt has finished. Skipped when the helper-busy
|
||||||
|
# check itself raised (no publish to clear) so the counter
|
||||||
|
# stays in sync with publishes.
|
||||||
|
if training_load_window_published:
|
||||||
|
try:
|
||||||
|
from routes.inference import _clear_public_load_window
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
_clear_public_load_window("training")
|
||||||
|
|
||||||
|
|
||||||
@router.post("/stop", response_model = TrainingStopResponse)
|
@router.post("/stop", response_model = TrainingStopResponse)
|
||||||
|
|
|
||||||
1694
studio/backend/tests/test_diffusion_backend.py
Normal file
1694
studio/backend/tests/test_diffusion_backend.py
Normal file
File diff suppressed because it is too large
Load diff
336
studio/backend/tests/test_diffusion_routes.py
Normal file
336
studio/backend/tests/test_diffusion_routes.py
Normal file
|
|
@ -0,0 +1,336 @@
|
||||||
|
# SPDX-License-Identifier: AGPL-3.0-only
|
||||||
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||||
|
|
||||||
|
"""Route-level tests for ``/api/inference/images/*``.
|
||||||
|
|
||||||
|
Mounts the actual ``inference_router`` on a fresh FastAPI app with the
|
||||||
|
auth dependency replaced by a stub so we exercise the same FastAPI
|
||||||
|
handlers Studio ships in production. The diffusion backend is replaced
|
||||||
|
with an in-memory stub so we don't need diffusers / GPUs to run these.
|
||||||
|
|
||||||
|
To stay runnable in a minimal CPU-only env, ``routes/inference.py``
|
||||||
|
is loaded directly via ``importlib`` so we do NOT trigger
|
||||||
|
``routes/__init__.py`` -- that file eagerly imports training /
|
||||||
|
datasets / data_recipe / export and would drag in heavy deps
|
||||||
|
(matplotlib, etc.) that the diffusion tests do not need.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import importlib.util
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from fastapi import FastAPI
|
||||||
|
from fastapi.testclient import TestClient
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
|
||||||
|
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
if str(_BACKEND_ROOT) not in sys.path:
|
||||||
|
sys.path.insert(0, str(_BACKEND_ROOT))
|
||||||
|
|
||||||
|
|
||||||
|
def _import_inference_module():
|
||||||
|
"""Load ``routes/inference.py`` without executing ``routes/__init__``.
|
||||||
|
|
||||||
|
The package init imports training / datasets / data_recipe / export
|
||||||
|
routers, which pull in matplotlib / pandas / training stack. The
|
||||||
|
diffusion tests only need the inference module so we side-step the
|
||||||
|
package import via importlib.spec_from_file_location.
|
||||||
|
"""
|
||||||
|
# If a previous test already imported routes the normal way, reuse
|
||||||
|
# the cached module instead of re-loading.
|
||||||
|
cached = sys.modules.get("routes.inference")
|
||||||
|
if cached is not None:
|
||||||
|
return cached
|
||||||
|
target = _BACKEND_ROOT / "routes" / "inference.py"
|
||||||
|
spec = importlib.util.spec_from_file_location(
|
||||||
|
"routes.inference",
|
||||||
|
target,
|
||||||
|
# We do NOT set submodule_search_locations for routes itself
|
||||||
|
# because that would re-trigger routes/__init__.py. The module
|
||||||
|
# uses relative imports sparingly; absolute imports resolve via
|
||||||
|
# sys.path[0] = backend root.
|
||||||
|
)
|
||||||
|
assert spec and spec.loader, "could not build spec for routes/inference.py"
|
||||||
|
module = importlib.util.module_from_spec(spec)
|
||||||
|
sys.modules["routes.inference"] = module
|
||||||
|
# Round 15 P3 #9: drop the half-initialised module from
|
||||||
|
# sys.modules if exec_module() raises, otherwise later tests pick
|
||||||
|
# up the poisoned entry and report a misleading AttributeError
|
||||||
|
# instead of the original ImportError.
|
||||||
|
try:
|
||||||
|
spec.loader.exec_module(module)
|
||||||
|
except Exception:
|
||||||
|
sys.modules.pop("routes.inference", None)
|
||||||
|
raise
|
||||||
|
return module
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeBackend:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self._loaded = False
|
||||||
|
self._repo: str | None = None
|
||||||
|
self.calls: list[dict] = []
|
||||||
|
|
||||||
|
@property
|
||||||
|
def is_loaded(self) -> bool:
|
||||||
|
return self._loaded
|
||||||
|
|
||||||
|
def status(self) -> dict:
|
||||||
|
return {
|
||||||
|
"is_loaded": self._loaded,
|
||||||
|
"is_loading": False,
|
||||||
|
"repo_id": self._repo,
|
||||||
|
"family": "flux.2-klein" if self._loaded else None,
|
||||||
|
"pipeline_class": "Flux2KleinPipeline" if self._loaded else None,
|
||||||
|
"base_repo": "black-forest-labs/FLUX.2-klein" if self._loaded else None,
|
||||||
|
"gguf_filename": None,
|
||||||
|
"active_repo_id": self._repo,
|
||||||
|
"active_base_repo": (
|
||||||
|
"black-forest-labs/FLUX.2-klein" if self._loaded else None
|
||||||
|
),
|
||||||
|
# Round 14: guard-facing GGUF filename is now the full
|
||||||
|
# caller-supplied value, but this fake never sets one so
|
||||||
|
# both active and pending stay None.
|
||||||
|
"active_gguf_filename": None,
|
||||||
|
"pending_repo_id": None,
|
||||||
|
"pending_base_repo": None,
|
||||||
|
"pending_gguf_filename": None,
|
||||||
|
"device": "cpu",
|
||||||
|
"dtype": "torch.bfloat16",
|
||||||
|
"loaded_at": 0,
|
||||||
|
"last_error": None,
|
||||||
|
"supported_families": [],
|
||||||
|
}
|
||||||
|
|
||||||
|
def load_model(self, repo_id, **kw):
|
||||||
|
self.calls.append({"op": "load", "repo_id": repo_id, **kw})
|
||||||
|
self._loaded = True
|
||||||
|
self._repo = repo_id
|
||||||
|
return self.status()
|
||||||
|
|
||||||
|
def unload_model(self) -> dict:
|
||||||
|
self._loaded = False
|
||||||
|
self._repo = None
|
||||||
|
return {"is_loaded": False}
|
||||||
|
|
||||||
|
def generate_image(self, **kw):
|
||||||
|
self.calls.append({"op": "generate", **kw})
|
||||||
|
return Image.new("RGB", (kw["width"], kw["height"]), color = (123, 45, 67))
|
||||||
|
|
||||||
|
def generate_image_with_metadata(self, **kw):
|
||||||
|
image = self.generate_image(**kw)
|
||||||
|
meta = {
|
||||||
|
"model": self._repo,
|
||||||
|
"family": "flux.2-klein" if self._loaded else None,
|
||||||
|
}
|
||||||
|
return image, meta
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def app_with_stub(monkeypatch):
|
||||||
|
"""Build a FastAPI app that mounts the real inference router with
|
||||||
|
auth disabled and the diffusion backend swapped for a stub."""
|
||||||
|
inf = _import_inference_module()
|
||||||
|
import core.inference.diffusion as d
|
||||||
|
|
||||||
|
stub = _FakeBackend()
|
||||||
|
# Override the singleton accessor the route uses.
|
||||||
|
monkeypatch.setattr(d, "get_diffusion_backend", lambda: stub)
|
||||||
|
monkeypatch.setattr(inf, "_get_diffusion_backend", lambda: stub)
|
||||||
|
|
||||||
|
app = FastAPI()
|
||||||
|
# Diffusion image routes live on studio_router so they are NOT
|
||||||
|
# exposed under /v1 (which would let OpenAI-compat clients
|
||||||
|
# trigger Studio-only side effects).
|
||||||
|
app.include_router(inf.router, prefix = "/api/inference")
|
||||||
|
app.include_router(inf.studio_router, prefix = "/api/inference")
|
||||||
|
# Bypass auth by overriding the dependency.
|
||||||
|
from auth.authentication import get_current_subject
|
||||||
|
|
||||||
|
app.dependency_overrides[get_current_subject] = lambda: "test-user"
|
||||||
|
|
||||||
|
return app, stub
|
||||||
|
|
||||||
|
|
||||||
|
def test_status_when_unloaded(app_with_stub):
|
||||||
|
app, _ = app_with_stub
|
||||||
|
c = TestClient(app)
|
||||||
|
r = c.get("/api/inference/images/status")
|
||||||
|
assert r.status_code == 200
|
||||||
|
body = r.json()
|
||||||
|
assert body["is_loaded"] is False
|
||||||
|
assert body["repo_id"] is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_without_load_returns_400(app_with_stub):
|
||||||
|
app, _ = app_with_stub
|
||||||
|
c = TestClient(app)
|
||||||
|
r = c.post(
|
||||||
|
"/api/inference/images/generate",
|
||||||
|
json = {"prompt": "a red sphere"},
|
||||||
|
)
|
||||||
|
assert r.status_code == 400
|
||||||
|
assert "No diffusion model" in r.json()["detail"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_then_generate_round_trip(app_with_stub):
|
||||||
|
app, stub = app_with_stub
|
||||||
|
c = TestClient(app)
|
||||||
|
|
||||||
|
r = c.post(
|
||||||
|
"/api/inference/images/load",
|
||||||
|
json = {
|
||||||
|
"repo_id": "unsloth/FLUX.2-klein-4B-GGUF",
|
||||||
|
"gguf_filename": "flux-2-klein-4b-Q4_K_S.gguf",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
assert r.status_code == 200, r.text
|
||||||
|
assert r.json()["is_loaded"] is True
|
||||||
|
|
||||||
|
r = c.post(
|
||||||
|
"/api/inference/images/generate",
|
||||||
|
json = {
|
||||||
|
"prompt": "a tiny synth-pop album cover",
|
||||||
|
"width": 256,
|
||||||
|
"height": 256,
|
||||||
|
"num_inference_steps": 4,
|
||||||
|
"seed": 7,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
assert r.status_code == 200, r.text
|
||||||
|
body = r.json()
|
||||||
|
assert body["image_b64"]
|
||||||
|
assert body["image_mime"] == "image/png"
|
||||||
|
assert body["width"] == 256
|
||||||
|
assert body["height"] == 256
|
||||||
|
assert body["seed"] == 7
|
||||||
|
assert body["duration_ms"] >= 0
|
||||||
|
|
||||||
|
# Round-trip the base64 -> PIL to confirm it is a real PNG of the
|
||||||
|
# right size and not, say, an empty string.
|
||||||
|
import base64
|
||||||
|
import io
|
||||||
|
|
||||||
|
raw = base64.b64decode(body["image_b64"])
|
||||||
|
decoded = Image.open(io.BytesIO(raw))
|
||||||
|
assert decoded.format == "PNG"
|
||||||
|
assert decoded.size == (256, 256)
|
||||||
|
|
||||||
|
# Backend stub should have recorded both calls.
|
||||||
|
ops = [c["op"] for c in stub.calls]
|
||||||
|
assert ops == ["load", "generate"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_rejects_off_grid_size(app_with_stub):
|
||||||
|
app, stub = 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", "width": 513, "height": 512},
|
||||||
|
)
|
||||||
|
# Pydantic v2 wraps validator errors in 422 by default.
|
||||||
|
assert r.status_code in (400, 422), r.text
|
||||||
|
|
||||||
|
|
||||||
|
def test_unload_clears_state(app_with_stub):
|
||||||
|
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/unload")
|
||||||
|
assert r.status_code == 200
|
||||||
|
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_embedded_hf_token(app_with_stub):
|
||||||
|
"""Round 15 P1 #5: URL-embedded ``hf_xxxxx`` tokens in repo_id /
|
||||||
|
base_repo must be rejected with 422 so they never reach
|
||||||
|
``self._repo_id`` and get echoed back by ``status()``."""
|
||||||
|
app, _ = app_with_stub
|
||||||
|
c = TestClient(app)
|
||||||
|
r = c.post(
|
||||||
|
"/api/inference/images/load",
|
||||||
|
json = {
|
||||||
|
"repo_id": "https://hf_abcdefghij0123456789@huggingface.co/owner/repo",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
assert r.status_code == 422, r.text
|
||||||
|
body = r.json()
|
||||||
|
text = repr(body).lower()
|
||||||
|
assert "hf_token" in text or "embed" in text
|
||||||
|
# base_repo is also rejected.
|
||||||
|
r = c.post(
|
||||||
|
"/api/inference/images/load",
|
||||||
|
json = {
|
||||||
|
"repo_id": "owner/repo",
|
||||||
|
"gguf_filename": "x.gguf",
|
||||||
|
"base_repo": "https://hf_abcdefghij0123456789@huggingface.co/base/repo",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
assert r.status_code == 422, r.text
|
||||||
|
|
||||||
|
|
||||||
|
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
|
||||||
|
|
@ -66,15 +66,24 @@ from core.inference import llama_cpp as llama_cpp_module
|
||||||
|
|
||||||
|
|
||||||
def _load_model_source() -> str:
|
def _load_model_source() -> str:
|
||||||
"""Return the source of ``LlamaCppBackend.load_model``.
|
"""Return the source of ``LlamaCppBackend.load_model`` PLUS the
|
||||||
|
internal ``_load_model_impl_locked`` body it delegates to.
|
||||||
|
|
||||||
Using ``inspect.getsource`` instead of reading the file directly
|
Studio's diffusion PR split ``load_model`` into a thin wrapper
|
||||||
scopes the assertions to the function that actually launches
|
that publishes ``_loading_model_identifier`` under
|
||||||
llama-server, so neither the presence check nor the location check
|
``_serial_load_lock`` and an inner ``_load_model_impl_locked``
|
||||||
can be fooled by a stray occurrence of ``"--no-context-shift"``
|
body that actually spawns llama-server. The launch flags and the
|
||||||
elsewhere in the module.
|
``_wait_for_vram_settle`` call now live in the inner method, so
|
||||||
|
inspecting only ``load_model`` would miss them. Concatenating the
|
||||||
|
two sources keeps these source-inspection regression tests
|
||||||
|
working without weakening the scope (we still only look at the
|
||||||
|
two load entry points, not the entire module).
|
||||||
"""
|
"""
|
||||||
return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model)
|
parts = [inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model)]
|
||||||
|
impl = getattr(llama_cpp_module.LlamaCppBackend, "_load_model_impl_locked", None)
|
||||||
|
if impl is not None:
|
||||||
|
parts.append(inspect.getsource(impl))
|
||||||
|
return "\n".join(parts)
|
||||||
|
|
||||||
|
|
||||||
def test_no_context_shift_is_in_load_model():
|
def test_no_context_shift_is_in_load_model():
|
||||||
|
|
|
||||||
|
|
@ -271,10 +271,19 @@ def test_load_model_calls_helper_outside_lock_and_uses_last_kill_timestamp():
|
||||||
"""Pin the call site: outside Phase 3 lock, gated on the timestamp,
|
"""Pin the call site: outside Phase 3 lock, gated on the timestamp,
|
||||||
no ``had_live_process`` in-band flag regression. Mirrors the
|
no ``had_live_process`` in-band flag regression. Mirrors the
|
||||||
``inspect.getsource`` pattern from ``test_llama_cpp_no_context_shift``.
|
``inspect.getsource`` pattern from ``test_llama_cpp_no_context_shift``.
|
||||||
|
|
||||||
|
Studio's diffusion PR split ``load_model`` into a thin wrapper +
|
||||||
|
``_load_model_impl_locked`` that actually launches llama-server, so
|
||||||
|
look at both sources to keep the assertions scoped to the load entry
|
||||||
|
points and not the entire module.
|
||||||
"""
|
"""
|
||||||
import inspect
|
import inspect
|
||||||
|
|
||||||
src = inspect.getsource(LlamaCppBackend.load_model)
|
parts = [inspect.getsource(LlamaCppBackend.load_model)]
|
||||||
|
impl = getattr(LlamaCppBackend, "_load_model_impl_locked", None)
|
||||||
|
if impl is not None:
|
||||||
|
parts.append(inspect.getsource(impl))
|
||||||
|
src = "\n".join(parts)
|
||||||
assert "_wait_for_vram_settle" in src
|
assert "_wait_for_vram_settle" in src
|
||||||
assert "since_kill" in src
|
assert "since_kill" in src
|
||||||
assert "self._last_kill_monotonic" in src
|
assert "self._last_kill_monotonic" in src
|
||||||
|
|
|
||||||
|
|
@ -18,7 +18,9 @@ import logging
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
import textwrap
|
import textwrap
|
||||||
|
import threading
|
||||||
import time
|
import time
|
||||||
|
from collections import Counter
|
||||||
from itertools import islice
|
from itertools import islice
|
||||||
from typing import Any, Optional
|
from typing import Any, Optional
|
||||||
|
|
||||||
|
|
@ -31,6 +33,137 @@ DEFAULT_HELPER_MODEL_VARIANT = "UD-Q4_K_XL"
|
||||||
|
|
||||||
README_MAX_CHARS = 1500
|
README_MAX_CHARS = 1500
|
||||||
|
|
||||||
|
# Round 26 P1 #13 / #14: helper/advisor run on PRIVATE LlamaCppBackend
|
||||||
|
# instances. Expose loading repo ids through thread-safe Counters so
|
||||||
|
# DELETE /api/models/delete-cached can block while a helper or
|
||||||
|
# advisor still owns the cache.
|
||||||
|
#
|
||||||
|
# Round 28 P1 #2: split into CACHE vs GPU refcounts. precache_helper_gguf
|
||||||
|
# downloads files (cache ownership) without occupying VRAM (GPU
|
||||||
|
# ownership), so collapsing them caused the public GPU handoffs to
|
||||||
|
# 503 during a background precache that did not need the GPU.
|
||||||
|
# * CACHE: blocks delete-cache for any active downloader / loader
|
||||||
|
# * GPU : blocks public chat / training / export / diffusion loads
|
||||||
|
_HELPER_ADVISOR_CACHE_REFCOUNT: Counter[str] = Counter()
|
||||||
|
_HELPER_ADVISOR_GPU_REFCOUNT: Counter[str] = Counter()
|
||||||
|
# Round 30 P1 #7-#10: counter of public GPU workloads (chat /
|
||||||
|
# diffusion / training / export) that have passed the helper-busy
|
||||||
|
# snapshot but have not yet flipped their public ownership flags
|
||||||
|
# (``llama.is_loaded`` / ``loading_model_identifier`` /
|
||||||
|
# ``current_checkpoint`` / ``is_training_active``). Helper / advisor
|
||||||
|
# starts consult this so they cannot win the start lock and race a
|
||||||
|
# public load that already destroyed the previous owner.
|
||||||
|
_PUBLIC_LOAD_PENDING_COUNT: Counter[str] = Counter()
|
||||||
|
_HELPER_ADVISOR_LOCK = threading.Lock()
|
||||||
|
# Round 28 P1 #7 / #8 / #10: serialize helper / advisor STARTS so two
|
||||||
|
# concurrent invocations cannot both pass the busy precheck before
|
||||||
|
# either registers. Held only across the precheck + register window,
|
||||||
|
# not across the full helper run.
|
||||||
|
# Round 30 P1 #7-#10: public GPU loads also enter under this lock to
|
||||||
|
# publish their pending counter so a concurrent helper / advisor
|
||||||
|
# start sees the pending public owner and refuses VRAM.
|
||||||
|
_HELPER_ADVISOR_START_LOCK = threading.Lock()
|
||||||
|
|
||||||
|
|
||||||
|
def helper_advisor_owns_repo(repo_id: str) -> bool:
|
||||||
|
"""Return True if any helper/advisor activity (precache OR live
|
||||||
|
helper / advisor load) currently owns this HF repo id."""
|
||||||
|
if not repo_id:
|
||||||
|
return False
|
||||||
|
needle = repo_id.lower()
|
||||||
|
with _HELPER_ADVISOR_LOCK:
|
||||||
|
return _HELPER_ADVISOR_CACHE_REFCOUNT.get(needle, 0) > 0
|
||||||
|
|
||||||
|
|
||||||
|
def helper_advisor_busy() -> bool:
|
||||||
|
"""True if any helper/advisor load is currently OCCUPYING THE GPU.
|
||||||
|
Round 28 P1 #2: must not return True for a precache-only download
|
||||||
|
(it owns disk cache, not VRAM)."""
|
||||||
|
with _HELPER_ADVISOR_LOCK:
|
||||||
|
return sum(_HELPER_ADVISOR_GPU_REFCOUNT.values()) > 0
|
||||||
|
|
||||||
|
|
||||||
|
def _register_helper_advisor_repo(repo_id: str, *, gpu_owner: bool = True) -> None:
|
||||||
|
"""Register a helper/advisor activity. Set ``gpu_owner=False`` for
|
||||||
|
precache-only downloads that need cache-delete protection but do
|
||||||
|
not load weights into VRAM."""
|
||||||
|
if not repo_id:
|
||||||
|
return
|
||||||
|
needle = repo_id.lower()
|
||||||
|
with _HELPER_ADVISOR_LOCK:
|
||||||
|
_HELPER_ADVISOR_CACHE_REFCOUNT[needle] += 1
|
||||||
|
if gpu_owner:
|
||||||
|
_HELPER_ADVISOR_GPU_REFCOUNT[needle] += 1
|
||||||
|
|
||||||
|
|
||||||
|
def _unregister_helper_advisor_repo(repo_id: str, *, gpu_owner: bool = True) -> None:
|
||||||
|
if not repo_id:
|
||||||
|
return
|
||||||
|
needle = repo_id.lower()
|
||||||
|
with _HELPER_ADVISOR_LOCK:
|
||||||
|
_HELPER_ADVISOR_CACHE_REFCOUNT[needle] -= 1
|
||||||
|
if _HELPER_ADVISOR_CACHE_REFCOUNT[needle] <= 0:
|
||||||
|
_HELPER_ADVISOR_CACHE_REFCOUNT.pop(needle, None)
|
||||||
|
if gpu_owner:
|
||||||
|
_HELPER_ADVISOR_GPU_REFCOUNT[needle] -= 1
|
||||||
|
if _HELPER_ADVISOR_GPU_REFCOUNT[needle] <= 0:
|
||||||
|
_HELPER_ADVISOR_GPU_REFCOUNT.pop(needle, None)
|
||||||
|
|
||||||
|
|
||||||
|
def _publish_public_load_pending(workload: str) -> None:
|
||||||
|
"""Mark a public GPU workload as mid-handoff. Must be called under
|
||||||
|
``_HELPER_ADVISOR_START_LOCK`` immediately after the helper-busy
|
||||||
|
snapshot succeeded (round 30 P1 #7-#10)."""
|
||||||
|
if not workload:
|
||||||
|
return
|
||||||
|
needle = workload.lower()
|
||||||
|
with _HELPER_ADVISOR_LOCK:
|
||||||
|
_PUBLIC_LOAD_PENDING_COUNT[needle] += 1
|
||||||
|
|
||||||
|
|
||||||
|
def _release_public_load_pending(workload: str) -> None:
|
||||||
|
"""Decrement the pending public-load counter once per matched
|
||||||
|
publish. Safe to call in finally even if the load failed."""
|
||||||
|
if not workload:
|
||||||
|
return
|
||||||
|
needle = workload.lower()
|
||||||
|
with _HELPER_ADVISOR_LOCK:
|
||||||
|
_PUBLIC_LOAD_PENDING_COUNT[needle] -= 1
|
||||||
|
if _PUBLIC_LOAD_PENDING_COUNT[needle] <= 0:
|
||||||
|
_PUBLIC_LOAD_PENDING_COUNT.pop(needle, None)
|
||||||
|
|
||||||
|
|
||||||
|
def public_load_pending(*, excluding: str | None = None) -> bool:
|
||||||
|
"""True if any public GPU workload has passed its helper-busy
|
||||||
|
snapshot but not yet flipped its public ownership flags. Helper /
|
||||||
|
advisor starts treat this as busy so they cannot race a public
|
||||||
|
load mid-handoff.
|
||||||
|
|
||||||
|
Round 38 P1: ``excluding`` lets a route-wrapped backend call
|
||||||
|
skip the marker its own route layer already published (e.g. the
|
||||||
|
diffusion route publishes ``diffusion`` before calling into
|
||||||
|
``backend.load_model``, which publishes ``diffusion-backend`` --
|
||||||
|
the backend should ignore its own ``diffusion`` marker so the
|
||||||
|
parity check does not self-block) while still seeing every
|
||||||
|
OTHER in-flight public workload."""
|
||||||
|
ignored = excluding.lower() if excluding else None
|
||||||
|
with _HELPER_ADVISOR_LOCK:
|
||||||
|
return any(
|
||||||
|
count > 0 and workload != ignored
|
||||||
|
for workload, count in _PUBLIC_LOAD_PENDING_COUNT.items()
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def public_load_pending_for(workload: str) -> bool:
|
||||||
|
"""True if a specific public GPU workload is mid-handoff. Used by
|
||||||
|
release helpers to refuse a destructive teardown while the matching
|
||||||
|
/export/* or /chat /load_* route is still in its publish window."""
|
||||||
|
if not workload:
|
||||||
|
return False
|
||||||
|
needle = workload.lower()
|
||||||
|
with _HELPER_ADVISOR_LOCK:
|
||||||
|
return _PUBLIC_LOAD_PENDING_COUNT.get(needle, 0) > 0
|
||||||
|
|
||||||
|
|
||||||
def _strip_think_tags(text: str) -> str:
|
def _strip_think_tags(text: str) -> str:
|
||||||
"""Strip <think>...</think> reasoning blocks emitted by some models.
|
"""Strip <think>...</think> reasoning blocks emitted by some models.
|
||||||
|
|
@ -72,6 +205,12 @@ def precache_helper_gguf():
|
||||||
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
|
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Round 27 P1 #4: register the repo so DELETE /api/models/delete-cached
|
||||||
|
# cannot rmtree the cache directory while we are mid-download.
|
||||||
|
# Round 28 P1 #2: precache only downloads files; it does NOT occupy
|
||||||
|
# VRAM. Use gpu_owner=False so helper_advisor_busy() does not block
|
||||||
|
# public GPU workloads during a background pre-cache.
|
||||||
|
_register_helper_advisor_repo(repo, gpu_owner = False)
|
||||||
try:
|
try:
|
||||||
from huggingface_hub import HfApi, hf_hub_download
|
from huggingface_hub import HfApi, hf_hub_download
|
||||||
from huggingface_hub.utils import disable_progress_bars, enable_progress_bars
|
from huggingface_hub.utils import disable_progress_bars, enable_progress_bars
|
||||||
|
|
@ -103,12 +242,143 @@ def precache_helper_gguf():
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"Failed to pre-cache helper GGUF: {e}")
|
logger.warning(f"Failed to pre-cache helper GGUF: {e}")
|
||||||
finally:
|
finally:
|
||||||
|
_unregister_helper_advisor_repo(repo, gpu_owner = False)
|
||||||
try:
|
try:
|
||||||
enable_progress_bars()
|
enable_progress_bars()
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def _diffusion_image_model_busy() -> bool:
|
||||||
|
"""Round 22 P1 #2 / #3: helper / advisor GGUFs share VRAM with
|
||||||
|
the Images page diffusion pipeline. Public chat / training /
|
||||||
|
export routes call the strict ``_release_diffusion_for`` helper
|
||||||
|
before allocating, but these dataset-side helpers used to load
|
||||||
|
llama-server directly with no diffusion guard at all. Skip the
|
||||||
|
helper GGUF when ``DiffusionBackend.status()`` reports loaded /
|
||||||
|
loading so we do not double-own VRAM. Fail closed (treat as
|
||||||
|
busy) on any status() error to preserve the resident image
|
||||||
|
model rather than racing it for memory.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
from core.inference.diffusion import get_diffusion_backend
|
||||||
|
except Exception:
|
||||||
|
return False
|
||||||
|
try:
|
||||||
|
status = get_diffusion_backend().status()
|
||||||
|
except Exception:
|
||||||
|
return True
|
||||||
|
return bool(status.get("is_loaded") or status.get("is_loading"))
|
||||||
|
|
||||||
|
|
||||||
|
def _gpu_workload_busy_for_helper() -> bool:
|
||||||
|
"""Round 23 P1 #3 / #4: the diffusion-only guard from round 22
|
||||||
|
let the helper / advisor GGUF run on top of a live training run
|
||||||
|
or a resident export checkpoint. Extend the busy check to those
|
||||||
|
workloads too so any GPU owner (Images, Training, Export)
|
||||||
|
blocks the helper instead of double-owning VRAM. Each step
|
||||||
|
fails closed: an unverifiable status counts as busy so the
|
||||||
|
user's primary workload is preserved over the optional helper.
|
||||||
|
|
||||||
|
Round 24 P1 #1: extended to also catch a Chat-backend GPU owner.
|
||||||
|
The helper GGUF used to run on top of a loaded GGUF chat model
|
||||||
|
(llama-server) or safetensors chat model and OOM their shared
|
||||||
|
GPU; mirror the diffusion check by inspecting llama
|
||||||
|
``is_loaded`` / ``is_active`` / ``loading_model_identifier`` and
|
||||||
|
safetensors ``active_model_name`` / ``loading_models``.
|
||||||
|
|
||||||
|
Round 28 P1 #9: also catch another helper / advisor that already
|
||||||
|
owns a private LlamaCppBackend. Without this two concurrent
|
||||||
|
helpers could both pass the precheck and OOM each other.
|
||||||
|
"""
|
||||||
|
if helper_advisor_busy():
|
||||||
|
logger.info(
|
||||||
|
"Skipping helper GGUF while another helper/advisor is using the GPU"
|
||||||
|
)
|
||||||
|
return True
|
||||||
|
# Round 30 P1 #7-#10: a public GPU load (chat / diffusion / training /
|
||||||
|
# export) that has passed its busy snapshot but not yet flipped its
|
||||||
|
# public ownership flags is still mid-handoff. Refuse so the helper
|
||||||
|
# does not race it for VRAM after the previous owner was torn down.
|
||||||
|
if public_load_pending():
|
||||||
|
logger.info("Skipping helper GGUF while a public GPU load is mid-handoff")
|
||||||
|
return True
|
||||||
|
if _diffusion_image_model_busy():
|
||||||
|
return True
|
||||||
|
|
||||||
|
try:
|
||||||
|
from routes.inference import get_llama_cpp_backend
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
llama = get_llama_cpp_backend()
|
||||||
|
if (
|
||||||
|
getattr(llama, "is_loaded", False)
|
||||||
|
or getattr(llama, "is_active", False)
|
||||||
|
or getattr(llama, "loading_model_identifier", None)
|
||||||
|
):
|
||||||
|
logger.info(
|
||||||
|
"Skipping helper GGUF while a GGUF chat model is loaded/loading"
|
||||||
|
)
|
||||||
|
return True
|
||||||
|
except Exception:
|
||||||
|
logger.info(
|
||||||
|
"Skipping helper GGUF because llama-server status is unavailable"
|
||||||
|
)
|
||||||
|
return True
|
||||||
|
|
||||||
|
try:
|
||||||
|
from core.inference import get_inference_backend
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
inf = get_inference_backend()
|
||||||
|
active = getattr(inf, "active_model_name", None)
|
||||||
|
loading = set(getattr(inf, "loading_models", set()) or set())
|
||||||
|
if active or loading:
|
||||||
|
logger.info(
|
||||||
|
"Skipping helper GGUF while a safetensors chat model is loaded/loading"
|
||||||
|
)
|
||||||
|
return True
|
||||||
|
except Exception:
|
||||||
|
logger.info(
|
||||||
|
"Skipping helper GGUF because safetensors chat status is unavailable"
|
||||||
|
)
|
||||||
|
return True
|
||||||
|
|
||||||
|
try:
|
||||||
|
from core.training import get_training_backend
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
if get_training_backend().is_training_active():
|
||||||
|
logger.info("Skipping helper GGUF while training is active")
|
||||||
|
return True
|
||||||
|
except Exception:
|
||||||
|
logger.info("Skipping helper GGUF because training status is unavailable")
|
||||||
|
return True
|
||||||
|
|
||||||
|
try:
|
||||||
|
from core.export import get_export_backend
|
||||||
|
except Exception:
|
||||||
|
return False
|
||||||
|
|
||||||
|
try:
|
||||||
|
exp = get_export_backend()
|
||||||
|
is_active = getattr(exp, "is_export_active", None)
|
||||||
|
if (is_active and is_active()) or getattr(exp, "current_checkpoint", None):
|
||||||
|
logger.info("Skipping helper GGUF while export owns the GPU")
|
||||||
|
return True
|
||||||
|
except Exception:
|
||||||
|
logger.info("Skipping helper GGUF because export status is unavailable")
|
||||||
|
return True
|
||||||
|
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
|
def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
|
||||||
"""
|
"""
|
||||||
Load helper model, run one chat completion, unload.
|
Load helper model, run one chat completion, unload.
|
||||||
|
|
@ -118,13 +388,28 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
|
||||||
if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"):
|
if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"):
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
# Round 23 P1 #3: round 22 only guarded against a busy
|
||||||
|
# diffusion pipeline. Training / export own the same GPU too,
|
||||||
|
# so use the broader helper that gates on all three workloads.
|
||||||
|
# Round 28 P1 #7 / #10: serialize the busy check + register pair
|
||||||
|
# so two concurrent helper invocations cannot both pass the
|
||||||
|
# precheck before either registers and then OOM each other.
|
||||||
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
|
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
|
||||||
variant = os.environ.get(
|
variant = os.environ.get(
|
||||||
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
|
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
|
||||||
)
|
)
|
||||||
|
with _HELPER_ADVISOR_START_LOCK:
|
||||||
|
if _gpu_workload_busy_for_helper():
|
||||||
|
return None
|
||||||
|
_register_helper_advisor_repo(repo)
|
||||||
backend = None
|
backend = None
|
||||||
try:
|
try:
|
||||||
|
# Round 26 P1 #1 / #3 / #13 / #14: use a PRIVATE backend so the
|
||||||
|
# helper can never preempt or be preempted by the user's
|
||||||
|
# chat backend and cannot accidentally unload it in finally.
|
||||||
|
# The active repo is published via _register_helper_advisor_repo
|
||||||
|
# above so DELETE /api/models/delete-cached can still block the
|
||||||
|
# cache rmtree while the helper is downloading or mmap'ing.
|
||||||
from core.inference.llama_cpp import LlamaCppBackend
|
from core.inference.llama_cpp import LlamaCppBackend
|
||||||
|
|
||||||
backend = LlamaCppBackend()
|
backend = LlamaCppBackend()
|
||||||
|
|
@ -176,6 +461,7 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
|
||||||
logger.info("Helper model unloaded")
|
logger.info("Helper model unloaded")
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
_unregister_helper_advisor_repo(repo)
|
||||||
|
|
||||||
|
|
||||||
# ─── Public API ───────────────────────────────────────────────────────
|
# ─── Public API ───────────────────────────────────────────────────────
|
||||||
|
|
@ -508,13 +794,26 @@ def _run_multi_pass_advisor(
|
||||||
if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"):
|
if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"):
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
# Round 23 P1 #4: extend the round 22 diffusion-only check to
|
||||||
|
# training + export so the advisor cannot race the user's
|
||||||
|
# active workload for GPU memory.
|
||||||
|
# Round 28 P1 #8 / #10: serialize the precheck + register pair so
|
||||||
|
# two concurrent advisor invocations cannot both pass before
|
||||||
|
# either registers and then OOM each other.
|
||||||
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
|
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
|
||||||
variant = os.environ.get(
|
variant = os.environ.get(
|
||||||
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
|
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
|
||||||
)
|
)
|
||||||
|
with _HELPER_ADVISOR_START_LOCK:
|
||||||
|
if _gpu_workload_busy_for_helper():
|
||||||
|
return None
|
||||||
|
_register_helper_advisor_repo(repo)
|
||||||
backend = None
|
backend = None
|
||||||
try:
|
try:
|
||||||
|
# Round 26 P1 #2 / #4 / #13 / #14: mirror ``_run_with_helper``
|
||||||
|
# and use a PRIVATE backend. Round 25's global-backend swap
|
||||||
|
# introduced chat-evict races and finally-eviction bugs.
|
||||||
|
# The registry above keeps delete-cache safe.
|
||||||
from core.inference.llama_cpp import LlamaCppBackend
|
from core.inference.llama_cpp import LlamaCppBackend
|
||||||
|
|
||||||
backend = LlamaCppBackend()
|
backend = LlamaCppBackend()
|
||||||
|
|
@ -849,6 +1148,7 @@ def _run_multi_pass_advisor(
|
||||||
logger.info("Advisor model unloaded")
|
logger.info("Advisor model unloaded")
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
_unregister_helper_advisor_repo(repo)
|
||||||
|
|
||||||
|
|
||||||
def llm_conversion_advisor(
|
def llm_conversion_advisor(
|
||||||
|
|
|
||||||
|
|
@ -9,6 +9,7 @@ import { Route as dataRecipeRoute } from "./routes/data-recipes.$recipeId";
|
||||||
import { Route as chatRoute } from "./routes/chat";
|
import { Route as chatRoute } from "./routes/chat";
|
||||||
import { Route as exportRoute } from "./routes/export";
|
import { Route as exportRoute } from "./routes/export";
|
||||||
import { Route as gridTestRoute } from "./routes/grid-test";
|
import { Route as gridTestRoute } from "./routes/grid-test";
|
||||||
|
import { Route as imagesRoute } from "./routes/images";
|
||||||
import { Route as indexRoute } from "./routes/index";
|
import { Route as indexRoute } from "./routes/index";
|
||||||
import { Route as loginRoute } from "./routes/login";
|
import { Route as loginRoute } from "./routes/login";
|
||||||
import { Route as onboardingRoute } from "./routes/onboarding";
|
import { Route as onboardingRoute } from "./routes/onboarding";
|
||||||
|
|
@ -26,6 +27,7 @@ const routeTree = rootRoute.addChildren([
|
||||||
studioRoute,
|
studioRoute,
|
||||||
chatRoute,
|
chatRoute,
|
||||||
exportRoute,
|
exportRoute,
|
||||||
|
imagesRoute,
|
||||||
dataRecipesRoute,
|
dataRecipesRoute,
|
||||||
dataRecipeRoute,
|
dataRecipeRoute,
|
||||||
]);
|
]);
|
||||||
|
|
|
||||||
21
studio/frontend/src/app/routes/images.tsx
Normal file
21
studio/frontend/src/app/routes/images.tsx
Normal file
|
|
@ -0,0 +1,21 @@
|
||||||
|
// SPDX-License-Identifier: AGPL-3.0-only
|
||||||
|
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||||
|
|
||||||
|
import { createRoute } from "@tanstack/react-router";
|
||||||
|
import { lazy } from "react";
|
||||||
|
import { requireAuth } from "../auth-guards";
|
||||||
|
import { Route as rootRoute } from "./__root";
|
||||||
|
|
||||||
|
const ImagesPage = lazy(() =>
|
||||||
|
import("@/features/images").then((m) => ({
|
||||||
|
default: m.ImagesPage,
|
||||||
|
})),
|
||||||
|
);
|
||||||
|
|
||||||
|
export const Route = createRoute({
|
||||||
|
getParentRoute: () => rootRoute,
|
||||||
|
path: "/images",
|
||||||
|
staticData: { title: "Images" },
|
||||||
|
beforeLoad: () => requireAuth(),
|
||||||
|
component: ImagesPage,
|
||||||
|
});
|
||||||
|
|
@ -50,6 +50,7 @@ import {
|
||||||
Globe02Icon,
|
Globe02Icon,
|
||||||
HelpCircleIcon,
|
HelpCircleIcon,
|
||||||
Logout01Icon,
|
Logout01Icon,
|
||||||
|
PaintBrush02Icon,
|
||||||
Search01Icon,
|
Search01Icon,
|
||||||
PowerIcon,
|
PowerIcon,
|
||||||
PencilEdit02Icon,
|
PencilEdit02Icon,
|
||||||
|
|
@ -497,6 +498,18 @@ export function AppSidebar() {
|
||||||
}}
|
}}
|
||||||
/>
|
/>
|
||||||
|
|
||||||
|
<NavItem
|
||||||
|
icon={PaintBrush02Icon}
|
||||||
|
label="Images"
|
||||||
|
active={pathname === "/images" || pathname.startsWith("/images/")}
|
||||||
|
disabled={chatOnly}
|
||||||
|
onClick={() => {
|
||||||
|
if (chatOnly) return;
|
||||||
|
navigate({ to: "/images" });
|
||||||
|
closeMobileIfOpen();
|
||||||
|
}}
|
||||||
|
/>
|
||||||
|
|
||||||
<NavItem
|
<NavItem
|
||||||
icon={DownloadSquare01Icon}
|
icon={DownloadSquare01Icon}
|
||||||
label="Export"
|
label="Export"
|
||||||
|
|
|
||||||
138
studio/frontend/src/features/images/api.ts
Normal file
138
studio/frontend/src/features/images/api.ts
Normal file
|
|
@ -0,0 +1,138 @@
|
||||||
|
// SPDX-License-Identifier: AGPL-3.0-only
|
||||||
|
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||||
|
|
||||||
|
// Thin client for the diffusion image-generation routes exposed by
|
||||||
|
// studio/backend/routes/inference.py (images/load, images/generate,
|
||||||
|
// images/status, images/unload). Mirrors the shape returned by
|
||||||
|
// DiffusionBackend.status() and DiffusionGenerateResponse so the
|
||||||
|
// page can render results without re-deriving fields client-side.
|
||||||
|
|
||||||
|
import { authFetch } from "@/features/auth";
|
||||||
|
import { readFastApiError } from "@/lib/format-fastapi-error";
|
||||||
|
|
||||||
|
export interface DiffusionFamily {
|
||||||
|
name: string;
|
||||||
|
pipeline_class: string;
|
||||||
|
base_repo: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface DiffusionStatus {
|
||||||
|
is_loaded: boolean;
|
||||||
|
is_loading: boolean;
|
||||||
|
repo_id: string | null;
|
||||||
|
family: string | null;
|
||||||
|
pipeline_class: string | null;
|
||||||
|
base_repo: string | null;
|
||||||
|
gguf_filename: string | null;
|
||||||
|
device: string | null;
|
||||||
|
dtype: string | null;
|
||||||
|
loaded_at: number | null;
|
||||||
|
last_error: string | null;
|
||||||
|
supported_families: DiffusionFamily[];
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface DiffusionLoadRequest {
|
||||||
|
repo_id: string;
|
||||||
|
gguf_filename?: string;
|
||||||
|
base_repo?: string;
|
||||||
|
family?: string;
|
||||||
|
hf_token?: string;
|
||||||
|
enable_model_cpu_offload?: boolean;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface DiffusionGenerateRequest {
|
||||||
|
prompt: string;
|
||||||
|
negative_prompt?: string;
|
||||||
|
num_inference_steps?: number;
|
||||||
|
guidance_scale?: number;
|
||||||
|
width?: number;
|
||||||
|
height?: number;
|
||||||
|
// bigint when the seed exceeds Number.MAX_SAFE_INTEGER, otherwise
|
||||||
|
// number. The wire format is always a JSON integer; see
|
||||||
|
// ``stringifyWithBigInt`` below.
|
||||||
|
seed?: number | bigint;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface DiffusionGenerateResponse {
|
||||||
|
image_b64: string;
|
||||||
|
image_mime: string;
|
||||||
|
width: number;
|
||||||
|
height: number;
|
||||||
|
num_inference_steps: number;
|
||||||
|
guidance_scale: number;
|
||||||
|
/**
|
||||||
|
* Numeric seed. Safe ONLY for values <= Number.MAX_SAFE_INTEGER.
|
||||||
|
* For larger seeds, prefer ``seed_str`` (full-precision decimal).
|
||||||
|
*/
|
||||||
|
seed: number | null;
|
||||||
|
/** Decimal string with full uint64 precision. Use this for display
|
||||||
|
* and reproduction when the user pastes the seed back in. */
|
||||||
|
seed_str: string | null;
|
||||||
|
duration_ms: number;
|
||||||
|
model: string | null;
|
||||||
|
family: string | null;
|
||||||
|
}
|
||||||
|
|
||||||
|
async function parseJson<T>(res: Response): Promise<T> {
|
||||||
|
if (!res.ok) throw new Error(await readFastApiError(res));
|
||||||
|
return (await res.json()) as T;
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function fetchDiffusionStatus(): Promise<DiffusionStatus> {
|
||||||
|
return parseJson<DiffusionStatus>(
|
||||||
|
await authFetch("/api/inference/images/status"),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function loadDiffusionModel(
|
||||||
|
payload: DiffusionLoadRequest,
|
||||||
|
): Promise<DiffusionStatus> {
|
||||||
|
return parseJson<DiffusionStatus>(
|
||||||
|
await authFetch("/api/inference/images/load", {
|
||||||
|
method: "POST",
|
||||||
|
headers: { "Content-Type": "application/json" },
|
||||||
|
body: JSON.stringify(payload),
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function unloadDiffusionModel(): Promise<{ is_loaded: boolean }> {
|
||||||
|
return parseJson<{ is_loaded: boolean }>(
|
||||||
|
await authFetch("/api/inference/images/unload", { method: "POST" }),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/** JSON.stringify cannot serialise BigInt directly. Pull the seed
|
||||||
|
* BigInt out, stringify the rest of the payload normally, then
|
||||||
|
* splice the seed's decimal digits back into the JSON literal at the
|
||||||
|
* exact ``"seed":<int>`` slot.
|
||||||
|
*
|
||||||
|
* Avoids the previous regex-over-JSON approach, which could be
|
||||||
|
* tripped by a user-supplied prompt that exactly matched the
|
||||||
|
* sentinel string. With this approach the only thing we touch is
|
||||||
|
* the literal ``"seed":<number>`` substring we wrote ourselves.
|
||||||
|
*/
|
||||||
|
function stringifyWithBigInt(value: DiffusionGenerateRequest): string {
|
||||||
|
const { seed, ...rest } = value;
|
||||||
|
if (typeof seed !== "bigint") {
|
||||||
|
return JSON.stringify(value);
|
||||||
|
}
|
||||||
|
// Serialise the rest without seed, then inject the seed at the end
|
||||||
|
// of the object literal as a JSON integer. Strip the trailing "}"
|
||||||
|
// and re-append once the field is added.
|
||||||
|
const base = JSON.stringify(rest);
|
||||||
|
const inner = base.length === 2 /* '{}' */ ? "" : base.slice(1, -1) + ",";
|
||||||
|
return `{${inner}"seed":${seed.toString()}}`;
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function generateDiffusionImage(
|
||||||
|
payload: DiffusionGenerateRequest,
|
||||||
|
): Promise<DiffusionGenerateResponse> {
|
||||||
|
return parseJson<DiffusionGenerateResponse>(
|
||||||
|
await authFetch("/api/inference/images/generate", {
|
||||||
|
method: "POST",
|
||||||
|
headers: { "Content-Type": "application/json" },
|
||||||
|
body: stringifyWithBigInt(payload),
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
}
|
||||||
620
studio/frontend/src/features/images/images-page.tsx
Normal file
620
studio/frontend/src/features/images/images-page.tsx
Normal file
|
|
@ -0,0 +1,620 @@
|
||||||
|
// SPDX-License-Identifier: AGPL-3.0-only
|
||||||
|
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||||
|
|
||||||
|
import { Button } from "@/components/ui/button";
|
||||||
|
import { Input } from "@/components/ui/input";
|
||||||
|
import { Label } from "@/components/ui/label";
|
||||||
|
import {
|
||||||
|
Select,
|
||||||
|
SelectContent,
|
||||||
|
SelectItem,
|
||||||
|
SelectTrigger,
|
||||||
|
SelectValue,
|
||||||
|
} from "@/components/ui/select";
|
||||||
|
import { SectionCard } from "@/components/section-card";
|
||||||
|
import { Slider } from "@/components/ui/slider";
|
||||||
|
import { Spinner } from "@/components/ui/spinner";
|
||||||
|
import { Textarea } from "@/components/ui/textarea";
|
||||||
|
import { toast } from "@/lib/toast";
|
||||||
|
import { PaintBrush02Icon, SparklesIcon, GpuIcon } from "@hugeicons/core-free-icons";
|
||||||
|
import { HugeiconsIcon } from "@hugeicons/react";
|
||||||
|
import {
|
||||||
|
fetchDiffusionStatus,
|
||||||
|
generateDiffusionImage,
|
||||||
|
loadDiffusionModel,
|
||||||
|
unloadDiffusionModel,
|
||||||
|
type DiffusionGenerateResponse,
|
||||||
|
type DiffusionStatus,
|
||||||
|
} from "./api";
|
||||||
|
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
|
||||||
|
|
||||||
|
// Curated short list of working diffusion GGUFs. Picked to span
|
||||||
|
// size + license so any GPU class has at least one viable option:
|
||||||
|
// FLUX.2 klein 4B -> ~13 GB VRAM with Q4_K_S, Apache 2.0
|
||||||
|
// FLUX.2 klein 9B -> ~17 GB VRAM, FLUX [klein] non-commercial (gated)
|
||||||
|
// FLUX.2 dev -> ~24+ GB VRAM, FLUX [dev] non-commercial (gated)
|
||||||
|
// FLUX.1 dev -> ~12 GB VRAM, older but widely tested (gated)
|
||||||
|
//
|
||||||
|
// Filenames mirror the Hub canonical case (lowercase 'flux-2-klein-4b')
|
||||||
|
// and base_repo is set explicitly so the backend never falls back to the
|
||||||
|
// family default. The CLI on the backend can load anything supported by
|
||||||
|
// detect_family(); this list just keeps the picker compact for the v1 UI.
|
||||||
|
const CURATED_MODELS: Array<{
|
||||||
|
label: string;
|
||||||
|
repo_id: string;
|
||||||
|
default_gguf: string;
|
||||||
|
base_repo: string;
|
||||||
|
family: string;
|
||||||
|
notes: string;
|
||||||
|
}> = [
|
||||||
|
{
|
||||||
|
label: "FLUX.2 klein base 4B (Q4_K_S, Apache 2.0)",
|
||||||
|
repo_id: "unsloth/FLUX.2-klein-base-4B-GGUF",
|
||||||
|
default_gguf: "flux-2-klein-base-4b-Q4_K_S.gguf",
|
||||||
|
base_repo: "black-forest-labs/FLUX.2-klein-base-4B",
|
||||||
|
family: "flux.2-klein",
|
||||||
|
notes: "13 GB VRAM, fastest. Apache 2.0, ungated.",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
label: "FLUX.2 klein 4B (Q4_K_S, distilled)",
|
||||||
|
repo_id: "unsloth/FLUX.2-klein-4B-GGUF",
|
||||||
|
default_gguf: "flux-2-klein-4b-Q4_K_S.gguf",
|
||||||
|
// Distilled GGUF must pair with the distilled base, not the Base
|
||||||
|
// checkpoint. The Hub model card for the GGUF lists
|
||||||
|
// base_model: black-forest-labs/FLUX.2-klein-4B.
|
||||||
|
base_repo: "black-forest-labs/FLUX.2-klein-4B",
|
||||||
|
family: "flux.2-klein",
|
||||||
|
notes: "13 GB VRAM. Distilled klein 4B. Requires HF access to FLUX.2 klein 4B.",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
label: "FLUX.2 klein 9B (Q4_K_S, gated)",
|
||||||
|
repo_id: "unsloth/FLUX.2-klein-9B-GGUF",
|
||||||
|
default_gguf: "flux-2-klein-9b-Q4_K_S.gguf",
|
||||||
|
base_repo: "black-forest-labs/FLUX.2-klein-9B",
|
||||||
|
family: "flux.2-klein",
|
||||||
|
notes: "17 GB VRAM. Higher quality distilled. Requires HF access to FLUX.2 klein 9B.",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
label: "FLUX.2 dev (Q4_K_S, gated)",
|
||||||
|
repo_id: "unsloth/FLUX.2-dev-GGUF",
|
||||||
|
default_gguf: "flux2-dev-Q4_K_S.gguf",
|
||||||
|
base_repo: "black-forest-labs/FLUX.2-dev",
|
||||||
|
family: "flux.2",
|
||||||
|
notes: "24+ GB VRAM. Requires HF access to FLUX.2 dev.",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
label: "FLUX.1 dev (Q4_K_S, city96, gated)",
|
||||||
|
repo_id: "city96/FLUX.1-dev-gguf",
|
||||||
|
default_gguf: "flux1-dev-Q4_K_S.gguf",
|
||||||
|
base_repo: "black-forest-labs/FLUX.1-dev",
|
||||||
|
family: "flux.1",
|
||||||
|
notes: "12 GB VRAM. Older but widely tested. Requires HF access to FLUX.1 dev.",
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
const DEFAULT_PRESET = CURATED_MODELS[0];
|
||||||
|
|
||||||
|
const RESOLUTION_PRESETS: Array<{ label: string; w: number; h: number }> = [
|
||||||
|
{ label: "Square 1024", w: 1024, h: 1024 },
|
||||||
|
{ label: "Square 768", w: 768, h: 768 },
|
||||||
|
{ label: "Square 512", w: 512, h: 512 },
|
||||||
|
{ label: "Portrait 832x1216", w: 832, h: 1216 },
|
||||||
|
{ label: "Landscape 1216x832", w: 1216, h: 832 },
|
||||||
|
];
|
||||||
|
|
||||||
|
export function ImagesPage() {
|
||||||
|
const [status, setStatus] = useState<DiffusionStatus | null>(null);
|
||||||
|
const [refreshingStatus, setRefreshingStatus] = useState(false);
|
||||||
|
const [busy, setBusy] = useState<"idle" | "loading" | "unloading" | "generating">("idle");
|
||||||
|
|
||||||
|
const [presetIndex, setPresetIndex] = useState(0);
|
||||||
|
const [customRepoId, setCustomRepoId] = useState("");
|
||||||
|
const [customGguf, setCustomGguf] = useState("");
|
||||||
|
const [customBaseRepo, setCustomBaseRepo] = useState("");
|
||||||
|
const [customFamily, setCustomFamily] = useState<string>("auto");
|
||||||
|
const [useCustom, setUseCustom] = useState(false);
|
||||||
|
const [hfToken, setHfToken] = useState("");
|
||||||
|
|
||||||
|
const [prompt, setPrompt] = useState("a tiny ginger sloth coding in a sunlit treehouse, photorealistic");
|
||||||
|
const [negativePrompt, setNegativePrompt] = useState("");
|
||||||
|
const [steps, setSteps] = useState(24);
|
||||||
|
const [guidance, setGuidance] = useState(3.5);
|
||||||
|
const [resolutionIdx, setResolutionIdx] = useState(0);
|
||||||
|
const [seed, setSeed] = useState<string>("");
|
||||||
|
|
||||||
|
const [results, setResults] = useState<DiffusionGenerateResponse[]>([]);
|
||||||
|
const lastErrorRef = useRef<string | null>(null);
|
||||||
|
|
||||||
|
const preset = CURATED_MODELS[presetIndex] ?? DEFAULT_PRESET;
|
||||||
|
const resolution = RESOLUTION_PRESETS[resolutionIdx];
|
||||||
|
|
||||||
|
// Round 30 P2 #12: split the fetch from the spinner toggle so the
|
||||||
|
// mount + auto-poll effects can call the fetch without the
|
||||||
|
// synchronous setRefreshingStatus(true) that tripped
|
||||||
|
// react-hooks/set-state-in-effect.
|
||||||
|
const fetchAndUpdateStatus = useCallback(async () => {
|
||||||
|
try {
|
||||||
|
const next = await fetchDiffusionStatus();
|
||||||
|
setStatus(next);
|
||||||
|
} catch (err) {
|
||||||
|
const msg = err instanceof Error ? err.message : String(err);
|
||||||
|
if (lastErrorRef.current !== msg) {
|
||||||
|
lastErrorRef.current = msg;
|
||||||
|
toast.error("Could not fetch image-model status", { description: msg });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}, []);
|
||||||
|
|
||||||
|
const refreshStatus = useCallback(async () => {
|
||||||
|
setRefreshingStatus(true);
|
||||||
|
try {
|
||||||
|
await fetchAndUpdateStatus();
|
||||||
|
} finally {
|
||||||
|
setRefreshingStatus(false);
|
||||||
|
}
|
||||||
|
}, [fetchAndUpdateStatus]);
|
||||||
|
|
||||||
|
useEffect(() => {
|
||||||
|
// Defer the mount fetch out of the synchronous effect body so the
|
||||||
|
// setStatus call inside fetchAndUpdateStatus does not trip the
|
||||||
|
// react-hooks/set-state-in-effect rule.
|
||||||
|
const id = window.setTimeout(() => {
|
||||||
|
void fetchAndUpdateStatus();
|
||||||
|
}, 0);
|
||||||
|
return () => window.clearTimeout(id);
|
||||||
|
}, [fetchAndUpdateStatus]);
|
||||||
|
|
||||||
|
// Round 27 P2: when the backend is mid-load (is_loading=true) the
|
||||||
|
// status label froze at "Loading..." until the user clicked
|
||||||
|
// Refresh. Auto-poll every 2 s while a load is in flight so the
|
||||||
|
// UI tracks real backend progress.
|
||||||
|
useEffect(() => {
|
||||||
|
if (!status?.is_loading) return;
|
||||||
|
const id = window.setInterval(() => {
|
||||||
|
void fetchAndUpdateStatus();
|
||||||
|
}, 2000);
|
||||||
|
return () => window.clearInterval(id);
|
||||||
|
}, [status?.is_loading, fetchAndUpdateStatus]);
|
||||||
|
|
||||||
|
const handleLoad = useCallback(async () => {
|
||||||
|
setBusy("loading");
|
||||||
|
try {
|
||||||
|
const repo = useCustom ? customRepoId.trim() : preset.repo_id;
|
||||||
|
const gguf = useCustom ? customGguf.trim() || undefined : preset.default_gguf;
|
||||||
|
// 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
|
||||||
|
// now also lets the user pin one because private / mirrored
|
||||||
|
// GGUFs (e.g. a 9B klein transformer) would otherwise fall
|
||||||
|
// back to the family-default 4B base and 500 on load. Empty
|
||||||
|
// string still falls back to the backend's smart-base /
|
||||||
|
// repo-id defaults.
|
||||||
|
const baseRepo = useCustom
|
||||||
|
? customBaseRepo.trim() || undefined
|
||||||
|
: preset.base_repo;
|
||||||
|
if (!repo) {
|
||||||
|
toast.error("Pick a model first");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
const next = await loadDiffusionModel({
|
||||||
|
repo_id: repo,
|
||||||
|
gguf_filename: gguf,
|
||||||
|
base_repo: baseRepo,
|
||||||
|
family,
|
||||||
|
hf_token: hfToken.trim() || undefined,
|
||||||
|
});
|
||||||
|
setStatus(next);
|
||||||
|
toast.success("Loaded image model", { description: next.repo_id ?? undefined });
|
||||||
|
} catch (err) {
|
||||||
|
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, customBaseRepo, customFamily, preset, hfToken, refreshStatus]);
|
||||||
|
|
||||||
|
const handleUnload = useCallback(async () => {
|
||||||
|
setBusy("unloading");
|
||||||
|
try {
|
||||||
|
await unloadDiffusionModel();
|
||||||
|
await refreshStatus();
|
||||||
|
} catch (err) {
|
||||||
|
toast.error("Failed to unload image model", {
|
||||||
|
description: err instanceof Error ? err.message : String(err),
|
||||||
|
});
|
||||||
|
// Round 27 P2: a partial unload (subprocess refused to terminate,
|
||||||
|
// 503 from the backend) used to leave the UI showing the old
|
||||||
|
// "Loaded:" label even though the backend state was half torn
|
||||||
|
// down. Refresh so the button states match reality (mirrors
|
||||||
|
// handleLoad above which always re-fetches on catch).
|
||||||
|
await refreshStatus();
|
||||||
|
} finally {
|
||||||
|
setBusy("idle");
|
||||||
|
}
|
||||||
|
}, [refreshStatus]);
|
||||||
|
|
||||||
|
const handleGenerate = useCallback(async () => {
|
||||||
|
if (!prompt.trim()) {
|
||||||
|
toast.error("Prompt is empty");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
setBusy("generating");
|
||||||
|
try {
|
||||||
|
// Reject non-integer seeds and clamp to the [-2^63, 2^64 - 1]
|
||||||
|
// range the backend's torch.Generator can actually pack. JSON
|
||||||
|
// serialises BigInts as plain integers, so we keep the wire
|
||||||
|
// format compatible and avoid the Number(seed) precision loss
|
||||||
|
// (>= 2^53 silently rounds, producing a different image than
|
||||||
|
// the seed the user typed). When the seed fits a safe integer
|
||||||
|
// it goes through unchanged; larger seeds ride along as their
|
||||||
|
// BigInt-derived string via the wire-format BigInt JSON helper
|
||||||
|
// in the api layer.
|
||||||
|
const seedStr = seed.trim();
|
||||||
|
let parsedSeed: number | bigint | undefined;
|
||||||
|
if (seedStr) {
|
||||||
|
if (!/^-?\d+$/.test(seedStr)) {
|
||||||
|
toast.error("Seed must be an integer");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let big: bigint;
|
||||||
|
try {
|
||||||
|
big = BigInt(seedStr);
|
||||||
|
} catch {
|
||||||
|
toast.error("Seed must be an integer");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
const SEED_MIN = -(BigInt(2) ** BigInt(63));
|
||||||
|
const SEED_MAX = BigInt(2) ** BigInt(64) - BigInt(1);
|
||||||
|
if (big < SEED_MIN || big > SEED_MAX) {
|
||||||
|
toast.error(
|
||||||
|
"Seed must be in [-2^63, 2^64 - 1] (the torch.Generator range)",
|
||||||
|
);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
// Use a plain Number when it fits a safe integer so the
|
||||||
|
// existing api.ts JSON serialiser does not break on BigInt;
|
||||||
|
// otherwise pass the BigInt and let api.ts emit it as a JSON
|
||||||
|
// number via a custom replacer.
|
||||||
|
const SAFE_MAX = BigInt(Number.MAX_SAFE_INTEGER);
|
||||||
|
const SAFE_MIN = -SAFE_MAX;
|
||||||
|
parsedSeed = big >= SAFE_MIN && big <= SAFE_MAX ? Number(big) : big;
|
||||||
|
}
|
||||||
|
const out = await generateDiffusionImage({
|
||||||
|
prompt,
|
||||||
|
negative_prompt: negativePrompt.trim() || undefined,
|
||||||
|
num_inference_steps: steps,
|
||||||
|
guidance_scale: guidance,
|
||||||
|
width: resolution.w,
|
||||||
|
height: resolution.h,
|
||||||
|
seed: parsedSeed,
|
||||||
|
});
|
||||||
|
setResults((prev) => [out, ...prev].slice(0, 12));
|
||||||
|
} catch (err) {
|
||||||
|
toast.error("Image generation failed", {
|
||||||
|
description: err instanceof Error ? err.message : String(err),
|
||||||
|
});
|
||||||
|
} finally {
|
||||||
|
setBusy("idle");
|
||||||
|
}
|
||||||
|
}, [prompt, negativePrompt, steps, guidance, resolution, seed]);
|
||||||
|
|
||||||
|
const statusLabel = useMemo(() => {
|
||||||
|
if (!status) return refreshingStatus ? "Checking..." : "Not loaded";
|
||||||
|
if (status.is_loading) return "Loading...";
|
||||||
|
if (status.is_loaded) {
|
||||||
|
const dev = status.device ? ` on ${status.device}` : "";
|
||||||
|
return `Loaded: ${status.repo_id ?? "(unknown)"} (${status.family ?? "unknown"})${dev}`;
|
||||||
|
}
|
||||||
|
return "Not loaded";
|
||||||
|
}, [status, refreshingStatus]);
|
||||||
|
|
||||||
|
// FLUX.2 / FLUX.2 klein pipelines do NOT accept negative_prompt and
|
||||||
|
// would 500 if we sent one through. The backend strips the field
|
||||||
|
// defensively but hiding it client-side keeps the UI honest.
|
||||||
|
// Round 29 P2 #12: also honour the user-picked customFamily when no
|
||||||
|
// model is loaded yet, so a Custom HF repo with family flux.2 /
|
||||||
|
// flux.2-klein hides the negative-prompt field correctly.
|
||||||
|
const supportsNegativePrompt = useMemo(() => {
|
||||||
|
const family = status?.family;
|
||||||
|
if (!family) {
|
||||||
|
let candidate: string | undefined;
|
||||||
|
if (useCustom) {
|
||||||
|
candidate = customFamily === "auto" ? undefined : customFamily;
|
||||||
|
} else {
|
||||||
|
candidate = preset.family;
|
||||||
|
}
|
||||||
|
if (!candidate) return true;
|
||||||
|
return !candidate.startsWith("flux.2");
|
||||||
|
}
|
||||||
|
return !family.startsWith("flux.2");
|
||||||
|
}, [status, useCustom, customFamily, preset.family]);
|
||||||
|
|
||||||
|
return (
|
||||||
|
<div className="flex flex-1 flex-col gap-4 overflow-y-auto p-4 sm:p-6">
|
||||||
|
<SectionCard
|
||||||
|
icon={<HugeiconsIcon icon={GpuIcon} className="size-5" strokeWidth={1.5} />}
|
||||||
|
title="Local image generation"
|
||||||
|
description={
|
||||||
|
"Run diffusion GGUFs from Hugging Face on your own GPU. " +
|
||||||
|
"Pick a curated FLUX.2 model or paste any unsloth/* GGUF repo."
|
||||||
|
}
|
||||||
|
>
|
||||||
|
<div className="flex flex-col gap-3">
|
||||||
|
<div className="flex flex-col gap-2">
|
||||||
|
<Label>Model</Label>
|
||||||
|
<Select
|
||||||
|
value={useCustom ? "custom" : String(presetIndex)}
|
||||||
|
onValueChange={(v) => {
|
||||||
|
if (v === "custom") {
|
||||||
|
setUseCustom(true);
|
||||||
|
} else {
|
||||||
|
setUseCustom(false);
|
||||||
|
setPresetIndex(Number(v));
|
||||||
|
}
|
||||||
|
}}
|
||||||
|
>
|
||||||
|
<SelectTrigger>
|
||||||
|
<SelectValue placeholder="Pick a model" />
|
||||||
|
</SelectTrigger>
|
||||||
|
<SelectContent>
|
||||||
|
{CURATED_MODELS.map((m, idx) => (
|
||||||
|
<SelectItem key={m.repo_id} value={String(idx)}>
|
||||||
|
{m.label}
|
||||||
|
</SelectItem>
|
||||||
|
))}
|
||||||
|
<SelectItem value="custom">Custom HF repo...</SelectItem>
|
||||||
|
</SelectContent>
|
||||||
|
</Select>
|
||||||
|
{!useCustom && (
|
||||||
|
<p className="text-xs text-muted-foreground">{preset.notes}</p>
|
||||||
|
)}
|
||||||
|
</div>
|
||||||
|
|
||||||
|
{useCustom && (
|
||||||
|
<div className="flex flex-col gap-2">
|
||||||
|
<Label>HF repo id</Label>
|
||||||
|
<Input
|
||||||
|
value={customRepoId}
|
||||||
|
onChange={(e) => setCustomRepoId(e.target.value)}
|
||||||
|
placeholder="unsloth/FLUX.2-klein-4B-GGUF"
|
||||||
|
/>
|
||||||
|
<Label>GGUF filename (optional)</Label>
|
||||||
|
<Input
|
||||||
|
value={customGguf}
|
||||||
|
onChange={(e) => setCustomGguf(e.target.value)}
|
||||||
|
placeholder="FLUX.2-klein-4B-Q4_K_S.gguf"
|
||||||
|
/>
|
||||||
|
<Label>Base diffusers repo (optional)</Label>
|
||||||
|
<Input
|
||||||
|
value={customBaseRepo}
|
||||||
|
onChange={(e) => setCustomBaseRepo(e.target.value)}
|
||||||
|
placeholder="black-forest-labs/FLUX.2-klein-9B"
|
||||||
|
/>
|
||||||
|
<p className="text-xs text-muted-foreground">
|
||||||
|
{"Optional. Defaults to the family base. Set this when "}
|
||||||
|
{"your GGUF expects a non-default base (for example a 9B "}
|
||||||
|
{"transformer that would otherwise fall back to a 4B base)."}
|
||||||
|
</p>
|
||||||
|
<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>
|
||||||
|
)}
|
||||||
|
|
||||||
|
<div className="flex flex-col gap-2">
|
||||||
|
<Label>Hugging Face token (only for gated repos)</Label>
|
||||||
|
<Input
|
||||||
|
type="password"
|
||||||
|
value={hfToken}
|
||||||
|
onChange={(e) => setHfToken(e.target.value)}
|
||||||
|
placeholder="hf_..."
|
||||||
|
autoComplete="off"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div className="flex flex-wrap items-center gap-2">
|
||||||
|
<Button
|
||||||
|
onClick={handleLoad}
|
||||||
|
disabled={busy !== "idle"}
|
||||||
|
data-testid="diffusion-load"
|
||||||
|
>
|
||||||
|
{busy === "loading" ? <Spinner className="mr-2 size-4" /> : null}
|
||||||
|
Load model
|
||||||
|
</Button>
|
||||||
|
<Button
|
||||||
|
variant="outline"
|
||||||
|
onClick={handleUnload}
|
||||||
|
disabled={busy !== "idle" || !status?.is_loaded}
|
||||||
|
data-testid="diffusion-unload"
|
||||||
|
>
|
||||||
|
Unload
|
||||||
|
</Button>
|
||||||
|
<Button
|
||||||
|
variant="ghost"
|
||||||
|
onClick={() => void refreshStatus()}
|
||||||
|
disabled={refreshingStatus}
|
||||||
|
>
|
||||||
|
Refresh status
|
||||||
|
</Button>
|
||||||
|
<span
|
||||||
|
className="ml-auto text-xs text-muted-foreground"
|
||||||
|
data-testid="diffusion-status"
|
||||||
|
>
|
||||||
|
{statusLabel}
|
||||||
|
</span>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
</SectionCard>
|
||||||
|
|
||||||
|
<SectionCard
|
||||||
|
icon={<HugeiconsIcon icon={PaintBrush02Icon} className="size-5" strokeWidth={1.5} />}
|
||||||
|
title="Prompt"
|
||||||
|
description="The pipeline runs on the GPU you launched Unsloth Studio on."
|
||||||
|
>
|
||||||
|
<div className="flex flex-col gap-3">
|
||||||
|
<div className="flex flex-col gap-1">
|
||||||
|
<Label htmlFor="diffusion-prompt">Prompt</Label>
|
||||||
|
<Textarea
|
||||||
|
id="diffusion-prompt"
|
||||||
|
value={prompt}
|
||||||
|
onChange={(e) => setPrompt(e.target.value)}
|
||||||
|
rows={3}
|
||||||
|
data-testid="diffusion-prompt"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
{supportsNegativePrompt ? (
|
||||||
|
<div className="flex flex-col gap-1">
|
||||||
|
<Label htmlFor="diffusion-negative">Negative prompt (optional)</Label>
|
||||||
|
<Textarea
|
||||||
|
id="diffusion-negative"
|
||||||
|
value={negativePrompt}
|
||||||
|
onChange={(e) => setNegativePrompt(e.target.value)}
|
||||||
|
rows={2}
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
) : (
|
||||||
|
<p className="text-xs text-muted-foreground">
|
||||||
|
{"FLUX.2 and FLUX.2 klein do not accept a negative prompt. "}
|
||||||
|
{"Steer the output via the main prompt instead."}
|
||||||
|
</p>
|
||||||
|
)}
|
||||||
|
|
||||||
|
<div className="grid grid-cols-1 gap-3 sm:grid-cols-3">
|
||||||
|
<div className="flex flex-col gap-1">
|
||||||
|
<Label>Resolution</Label>
|
||||||
|
<Select
|
||||||
|
value={String(resolutionIdx)}
|
||||||
|
onValueChange={(v) => setResolutionIdx(Number(v))}
|
||||||
|
>
|
||||||
|
<SelectTrigger>
|
||||||
|
<SelectValue />
|
||||||
|
</SelectTrigger>
|
||||||
|
<SelectContent>
|
||||||
|
{RESOLUTION_PRESETS.map((r, idx) => (
|
||||||
|
<SelectItem key={r.label} value={String(idx)}>
|
||||||
|
{r.label}
|
||||||
|
</SelectItem>
|
||||||
|
))}
|
||||||
|
</SelectContent>
|
||||||
|
</Select>
|
||||||
|
</div>
|
||||||
|
<div className="flex flex-col gap-1">
|
||||||
|
<Label>Steps: {steps}</Label>
|
||||||
|
<Slider
|
||||||
|
aria-label="Inference steps"
|
||||||
|
min={1}
|
||||||
|
max={60}
|
||||||
|
step={1}
|
||||||
|
value={[steps]}
|
||||||
|
onValueChange={(v) => setSteps(v[0] ?? steps)}
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<div className="flex flex-col gap-1">
|
||||||
|
<Label>Guidance: {guidance.toFixed(1)}</Label>
|
||||||
|
<Slider
|
||||||
|
aria-label="Guidance scale"
|
||||||
|
min={0}
|
||||||
|
max={15}
|
||||||
|
step={0.1}
|
||||||
|
value={[guidance]}
|
||||||
|
onValueChange={(v) => setGuidance(v[0] ?? guidance)}
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div className="flex flex-col gap-1">
|
||||||
|
<Label htmlFor="diffusion-seed">Seed (optional)</Label>
|
||||||
|
<Input
|
||||||
|
id="diffusion-seed"
|
||||||
|
value={seed}
|
||||||
|
onChange={(e) => setSeed(e.target.value)}
|
||||||
|
placeholder="leave empty for random"
|
||||||
|
inputMode="numeric"
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div>
|
||||||
|
<Button
|
||||||
|
size="lg"
|
||||||
|
onClick={handleGenerate}
|
||||||
|
disabled={busy !== "idle" || !status?.is_loaded}
|
||||||
|
data-testid="diffusion-generate"
|
||||||
|
>
|
||||||
|
{busy === "generating" ? <Spinner className="mr-2 size-4" /> : null}
|
||||||
|
Generate image
|
||||||
|
</Button>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
</SectionCard>
|
||||||
|
|
||||||
|
{results.length > 0 && (
|
||||||
|
<SectionCard
|
||||||
|
icon={<HugeiconsIcon icon={SparklesIcon} className="size-5" strokeWidth={1.5} />}
|
||||||
|
title="Results"
|
||||||
|
description="Most recent first."
|
||||||
|
>
|
||||||
|
<div className="grid grid-cols-1 gap-4 sm:grid-cols-2 lg:grid-cols-3">
|
||||||
|
{results.map((r, idx) => (
|
||||||
|
<figure key={idx} className="flex flex-col gap-2">
|
||||||
|
<img
|
||||||
|
src={`data:${r.image_mime};base64,${r.image_b64}`}
|
||||||
|
alt={`Generated image ${idx + 1}`}
|
||||||
|
// h-auto + object-contain so portrait / landscape
|
||||||
|
// outputs render at their true aspect ratio instead
|
||||||
|
// of being cropped into a square thumbnail.
|
||||||
|
className="h-auto w-full rounded-md border border-border object-contain"
|
||||||
|
data-testid="diffusion-result-image"
|
||||||
|
/>
|
||||||
|
<figcaption className="text-xs text-muted-foreground">
|
||||||
|
{r.width}x{r.height} - {r.num_inference_steps} steps - g={(r.guidance_scale ?? 0).toFixed(1)}
|
||||||
|
{/* Prefer seed_str (full uint64 precision) since the
|
||||||
|
numeric seed gets rounded by JSON.parse above
|
||||||
|
Number.MAX_SAFE_INTEGER and would otherwise
|
||||||
|
display a value that does not reproduce. */}
|
||||||
|
{r.seed_str
|
||||||
|
? ` - seed ${r.seed_str}`
|
||||||
|
: r.seed !== null && r.seed !== undefined
|
||||||
|
? ` - seed ${r.seed}`
|
||||||
|
: ""} -
|
||||||
|
{` ${(r.duration_ms / 1000).toFixed(1)}s`}
|
||||||
|
</figcaption>
|
||||||
|
</figure>
|
||||||
|
))}
|
||||||
|
</div>
|
||||||
|
</SectionCard>
|
||||||
|
)}
|
||||||
|
</div>
|
||||||
|
);
|
||||||
|
}
|
||||||
5
studio/frontend/src/features/images/index.ts
Normal file
5
studio/frontend/src/features/images/index.ts
Normal file
|
|
@ -0,0 +1,5 @@
|
||||||
|
// SPDX-License-Identifier: AGPL-3.0-only
|
||||||
|
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||||
|
|
||||||
|
export { ImagesPage } from "./images-page";
|
||||||
|
export * from "./api";
|
||||||
Loading…
Add table
Add a link
Reference in a new issue