Merge remote-tracking branch 'origin/image-generation' into pr6872
# Conflicts: # studio/backend/core/inference/diffusion.py # studio/backend/routes/training.py # studio/backend/tests/test_diffusion_dataset_api.py
This commit is contained in:
Daniel Han
commit
14dc3b86a7
10 changed files with 348 additions and 19 deletions
|
|
@ -226,6 +226,22 @@
|
|||
"evidence": "FS: L528: history: list[Response] | None = None, sha256:f56272dccd651b2644aa41ef6e688e211462427aad07fef5150240ec7347446e\nNetwork: L9: import urllib.request | L1243: class _CookieCompatRequest(urllib.request.Request):",
|
||||
"evidence_hash": "b32f79e58c938680d89efa74113eeba76c9fc5aedf5de18086f93bef274c4bda"
|
||||
},
|
||||
{
|
||||
"package": "huggingface-hub",
|
||||
"file": "huggingface_hub/_sandbox.py",
|
||||
"check": "C2 polling/beaconing loop detected",
|
||||
"severity": "CRITICAL",
|
||||
"evidence": "L1179: while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38",
|
||||
"evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8"
|
||||
},
|
||||
{
|
||||
"package": "huggingface-hub",
|
||||
"file": "huggingface_hub/_sandbox.py",
|
||||
"check": "Writes to /tmp and executes (staged dropper)",
|
||||
"severity": "CRITICAL",
|
||||
"evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"",
|
||||
"evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22"
|
||||
},
|
||||
{
|
||||
"package": "huggingface-hub",
|
||||
"file": "huggingface_hub/hf_api.py",
|
||||
|
|
@ -242,6 +258,14 @@
|
|||
"evidence": "L4600: while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272",
|
||||
"evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602"
|
||||
},
|
||||
{
|
||||
"package": "huggingface-hub",
|
||||
"file": "huggingface_hub/hf_api.py",
|
||||
"check": "C2 polling/beaconing loop detected",
|
||||
"severity": "CRITICAL",
|
||||
"evidence": "L4613: while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6",
|
||||
"evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4"
|
||||
},
|
||||
{
|
||||
"package": "huggingface-hub",
|
||||
"file": "huggingface_hub/hf_api.py",
|
||||
|
|
@ -266,6 +290,14 @@
|
|||
"evidence": "L298: while True: sha256:6b8e5e569594caf7c4eca6137646dae471a7c3aae7294096cf876f30b5f90306",
|
||||
"evidence_hash": "c066cc27bce31ee7b6ce07411ee7a7d9ecfbf3aafc8848f6641fabfe522a7703"
|
||||
},
|
||||
{
|
||||
"package": "huggingface-hub",
|
||||
"file": "huggingface_hub/utils/_http.py",
|
||||
"check": "C2 polling/beaconing loop detected",
|
||||
"severity": "CRITICAL",
|
||||
"evidence": "L462: while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958",
|
||||
"evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e"
|
||||
},
|
||||
{
|
||||
"package": "ipython",
|
||||
"file": "IPython/core/debugger.py",
|
||||
|
|
|
|||
|
|
@ -437,11 +437,17 @@ class DiffusionBackend:
|
|||
base: str,
|
||||
base_files: list[str],
|
||||
hf_token: Optional[str],
|
||||
) -> None:
|
||||
) -> Optional[str]:
|
||||
"""Pre-download the GGUF + the given ``base_files`` into the HF cache,
|
||||
WITHOUT the lock and honoring ``_cancel_event``, so load_pipeline's
|
||||
from_single_file / from_pretrained hit the cache and the heavy download can
|
||||
be preempted by an unload/eviction. Raises ``RuntimeError("Cancelled")``."""
|
||||
be preempted by an unload/eviction. Raises ``RuntimeError("Cancelled")``.
|
||||
|
||||
Returns the base repo's local snapshot dir when the prefetched set includes
|
||||
the pipeline manifest, so from_pretrained can load from disk instead of
|
||||
re-sweeping the hub (its own sweep also pulls files the scoped list skips,
|
||||
e.g. the 24 GB packaged root singles in each FLUX.1 repo); None otherwise
|
||||
(estimate failure, config-only base, local repo) -> hub id as before."""
|
||||
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
|
||||
|
||||
# GGUF transformer (hub repos only; a local path is already on disk).
|
||||
|
|
@ -450,12 +456,16 @@ class DiffusionBackend:
|
|||
repo_id, gguf_filename, hf_token, cancel_event = self._cancel_event
|
||||
)
|
||||
# Base repo (VAE / text-encoder / scheduler); list comes from the estimate.
|
||||
snapshot_root: Optional[str] = None
|
||||
for rfilename in base_files:
|
||||
if self._cancel_event.is_set():
|
||||
raise RuntimeError("Cancelled")
|
||||
hf_hub_download_with_xet_fallback(
|
||||
local = hf_hub_download_with_xet_fallback(
|
||||
base, rfilename, hf_token, cancel_event = self._cancel_event
|
||||
)
|
||||
if rfilename == "model_index.json":
|
||||
snapshot_root = str(Path(local).parent)
|
||||
return snapshot_root
|
||||
|
||||
def validate_load_request(
|
||||
self,
|
||||
|
|
@ -672,7 +682,7 @@ class DiffusionBackend:
|
|||
self._loading.expected_bytes = expected
|
||||
# Download outside the lock so unload()/an eviction can preempt the
|
||||
# multi-GB pull; load_pipeline below then assembles from the cache.
|
||||
self._prefetch_files(
|
||||
kwargs["_base_local_dir"] = self._prefetch_files(
|
||||
kwargs["repo_id"],
|
||||
kwargs.get("gguf_filename"),
|
||||
base,
|
||||
|
|
@ -899,6 +909,7 @@ class DiffusionBackend:
|
|||
transformer_cache_threshold: Optional[float] = None,
|
||||
model_kind: Optional[str] = None,
|
||||
_load_token: Optional[int] = None,
|
||||
_base_local_dir: Optional[str] = None,
|
||||
) -> dict[str, Any]:
|
||||
# A blank / whitespace-only token must degrade to anonymous access, not be passed
|
||||
# as an explicit credential (from_single_file / from_pretrained / the Hub client
|
||||
|
|
@ -1015,6 +1026,7 @@ class DiffusionBackend:
|
|||
transformer_quant,
|
||||
transformer_quant_fast_accum,
|
||||
fam = fam,
|
||||
base_local_dir = _base_local_dir,
|
||||
prequant_path = transformer_prequant_path,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 — fall back to the GGUF build
|
||||
|
|
@ -1058,7 +1070,12 @@ class DiffusionBackend:
|
|||
pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype}
|
||||
if hf_token:
|
||||
pipe_kwargs["token"] = hf_token
|
||||
pipe = pipeline_cls.from_pretrained(repo_id, **pipe_kwargs)
|
||||
# The prefetched snapshot dir keeps from_pretrained off the hub:
|
||||
# its own snapshot sweep re-downloads files the scoped prefetch
|
||||
# skipped (root packaged singles, e.g. 24 GB per FLUX.1 repo).
|
||||
pipe = pipeline_cls.from_pretrained(
|
||||
_base_local_dir or repo_id, **pipe_kwargs
|
||||
)
|
||||
elif kind == "single_file" and fam.single_file_is_pipeline:
|
||||
# A single-file SDXL-style checkpoint is the WHOLE pipeline
|
||||
# (U-Net + VAE + both text encoders), not a transformer-only file,
|
||||
|
|
@ -1094,7 +1111,7 @@ class DiffusionBackend:
|
|||
pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer}
|
||||
if hf_token:
|
||||
pipe_kwargs["token"] = hf_token
|
||||
pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
|
||||
pipe = pipeline_cls.from_pretrained(_base_local_dir or base, **pipe_kwargs)
|
||||
|
||||
# Resolve the effective speed mode: GGUF models default to the
|
||||
# near-lossless `default` profile (compile is ~2.2x and sits below
|
||||
|
|
@ -1326,6 +1343,7 @@ class DiffusionBackend:
|
|||
*,
|
||||
fam: Optional[DiffusionFamily] = None,
|
||||
prequant_path: Optional[str] = None,
|
||||
base_local_dir: Optional[str] = None,
|
||||
) -> tuple[Any, str]:
|
||||
"""Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``.
|
||||
|
||||
|
|
@ -1373,7 +1391,7 @@ class DiffusionBackend:
|
|||
)
|
||||
if transformer is not None:
|
||||
pipe = self._assemble_pipe(
|
||||
pipeline_cls, base, transformer, dtype, hf_token, device
|
||||
pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir
|
||||
)
|
||||
return pipe, scheme
|
||||
|
||||
|
|
@ -1381,7 +1399,9 @@ class DiffusionBackend:
|
|||
transformer = transformer_cls.from_pretrained(
|
||||
base, subfolder = "transformer", torch_dtype = dtype, token = hf_token
|
||||
)
|
||||
pipe = self._assemble_pipe(pipeline_cls, base, transformer, dtype, hf_token, device)
|
||||
pipe = self._assemble_pipe(
|
||||
pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir
|
||||
)
|
||||
scheme = quantize_transformer(
|
||||
pipe,
|
||||
target,
|
||||
|
|
@ -1402,13 +1422,14 @@ class DiffusionBackend:
|
|||
dtype: Any,
|
||||
hf_token: Optional[str],
|
||||
device: str,
|
||||
base_local_dir: Optional[str] = None,
|
||||
) -> Any:
|
||||
"""Assemble the diffusers pipeline around ``transformer`` and place it on ``device``
|
||||
(a no-op for an already-placed pre-quantized transformer; it moves the companions)."""
|
||||
pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer}
|
||||
if hf_token:
|
||||
pipe_kwargs["token"] = hf_token
|
||||
pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
|
||||
pipe = pipeline_cls.from_pretrained(base_local_dir or base, **pipe_kwargs)
|
||||
pipe.to(device)
|
||||
return pipe
|
||||
|
||||
|
|
@ -1555,6 +1576,16 @@ class DiffusionBackend:
|
|||
self._cn_models.clear()
|
||||
self._cn_pipes.clear()
|
||||
clear_gpu_cache()
|
||||
# resolve_controlnet accepts a bare owner/name repo without the non-GGUF base
|
||||
# trust gate, and from_pretrained below downloads and deserializes it. A
|
||||
# malicious pickle .bin would execute on load, so run the same Hub malware
|
||||
# preflight the chat/export loaders use before any remote ControlNet load. A
|
||||
# local dir the user picked has no Hub scan and is exempt (fail-open there).
|
||||
if not getattr(resolved_cn, "is_local", False):
|
||||
from utils.security import evaluate_file_security
|
||||
_cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None)
|
||||
if _cn_fs.blocked:
|
||||
raise ValueError(_cn_fs.reason)
|
||||
import torch
|
||||
|
||||
# state.dtype is the display string saved at load ("bfloat16"), NOT a
|
||||
|
|
@ -1598,7 +1629,14 @@ class DiffusionBackend:
|
|||
pipe = getattr(diffusers, pipe_cls_name).from_pipe(
|
||||
state.pipe, controlnet = cn_model, torch_dtype = None
|
||||
)
|
||||
self._cn_pipes[key] = pipe
|
||||
with self._lock:
|
||||
# Same race as the model cache above: an unload/superseding load may
|
||||
# have cleared _cn_pipes while from_pipe ran; caching now would pin a
|
||||
# pipeline built around the UNLOADED base and hand it to the next load.
|
||||
if cancel.is_set() or self._state is not state:
|
||||
del pipe
|
||||
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
|
||||
self._cn_pipes[key] = pipe
|
||||
return pipe
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -507,6 +507,16 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool:
|
|||
import torch
|
||||
from diffusers.hooks import apply_group_offloading
|
||||
|
||||
# A dual-DiT pipeline (e.g. Ideogram 4's unconditional tower) carries a second
|
||||
# denoiser as large as the first; leaving it resident would defeat this tier
|
||||
# (the pair rarely fits where one alone did not). Stream every DiT and keep
|
||||
# only the genuinely smaller companions resident.
|
||||
streamed: dict[str, Any] = {"transformer": transformer}
|
||||
for extra in ("transformer_2", "unconditional_transformer"):
|
||||
module = getattr(pipe, extra, None)
|
||||
if isinstance(module, torch.nn.Module):
|
||||
streamed[extra] = module
|
||||
|
||||
onload = torch.device(device)
|
||||
use_stream = onload.type == "cuda" # overlap H2D copies with compute on CUDA
|
||||
gkwargs: dict[str, Any] = {
|
||||
|
|
@ -536,11 +546,12 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool:
|
|||
# load-time crash. The streamed transformer manages its own placement via the
|
||||
# offloading hooks applied next.
|
||||
for name, comp in getattr(pipe, "components", {}).items():
|
||||
if name == "transformer":
|
||||
if name in streamed:
|
||||
continue
|
||||
if isinstance(comp, torch.nn.Module):
|
||||
comp.to(onload)
|
||||
apply_group_offloading(transformer, **gkwargs)
|
||||
for module in streamed.values():
|
||||
apply_group_offloading(module, **gkwargs)
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001 — fall back to whole-module offload
|
||||
if logger is not None:
|
||||
|
|
|
|||
|
|
@ -1100,6 +1100,13 @@ class SdCppDiffusionBackend:
|
|||
state.server.stop()
|
||||
if pending is not None and pending is not (state.server if state else None):
|
||||
pending.stop()
|
||||
# Wait for a signalled one-shot generation to actually exit before reporting
|
||||
# unloaded: callers (the GPU arbiter, training cleanup) treat this return as
|
||||
# "the device is free", but a one-shot sd-cli child killed by the cancel above
|
||||
# unwinds under _generate_lock. A bare acquire is the exit barrier (never taken
|
||||
# while holding _lock; same pattern as DiffusionBackend.unload).
|
||||
with self._generate_lock:
|
||||
pass
|
||||
return self.status()
|
||||
|
||||
def status(self) -> dict[str, Any]:
|
||||
|
|
|
|||
|
|
@ -1177,11 +1177,29 @@ def _preflight_gated_base(base_model: str, hf_token: Optional[str]) -> None:
|
|||
|
||||
@router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse)
|
||||
async def start_diffusion_training(
|
||||
body: DiffusionTrainingStartRequest, current_subject: str = Depends(get_current_subject)
|
||||
body: DiffusionTrainingStartRequest,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
via_api_key: bool = Depends(authenticated_via_api_key),
|
||||
):
|
||||
"""Start an SDXL LoRA training job from an image + caption dataset."""
|
||||
from core.training.diffusion_training_service import get_diffusion_training_service
|
||||
|
||||
# When Studio is driven as an inference API (API-key auth), refuse to start training
|
||||
# while a request is in flight: _free_gpu_for_diffusion_training() below unloads the
|
||||
# chat backends to reclaim VRAM, which would kill the stream. Mirrors start_training so
|
||||
# a diffusion start cannot silently drop an active API inference request.
|
||||
if via_api_key is True:
|
||||
from core.inference.llama_keepwarm import other_inference_request_count
|
||||
if other_inference_request_count(current_request_counted = False) > 0:
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = (
|
||||
"Cannot start diffusion (Images) training over the API while an inference "
|
||||
"request is in progress. Wait for it to finish, or start training from the "
|
||||
"Studio UI."
|
||||
),
|
||||
)
|
||||
|
||||
# Interlock: refuse while an LLM training run holds the GPU (symmetric with the
|
||||
# diffusion check in start_training), so the two trainers never contend for VRAM.
|
||||
try:
|
||||
|
|
@ -1234,6 +1252,21 @@ async def start_diffusion_training(
|
|||
# user's loaded chat/Images model, and never surfaces as a confusing mid-load 401.
|
||||
_preflight_gated_base(config.get("base_model", ""), config.get("hf_token"))
|
||||
|
||||
# Preflight the dataset too: a missing/empty/uncaptionable data_dir otherwise
|
||||
# fails inside the spawned trainer AFTER the user's chat/Images model was
|
||||
# evicted. Same discovery the trainer runs, so the two cannot disagree.
|
||||
from core.training import diffusion_train_common as _dtc
|
||||
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
_dtc.discover_image_caption_pairs,
|
||||
config["data_dir"],
|
||||
instance_prompt = config.get("instance_prompt") or None,
|
||||
caption_column = config.get("caption_column") or "text",
|
||||
)
|
||||
except (FileNotFoundError, ValueError) as e:
|
||||
raise HTTPException(status_code = 400, detail = str(e))
|
||||
|
||||
# Free resident GPU workloads (export / Images pipeline / chat) before the trainer
|
||||
# loads its own pipeline.
|
||||
_free_gpu_for_diffusion_training()
|
||||
|
|
@ -1601,6 +1634,10 @@ async def get_diffusion_dataset_image(
|
|||
|
||||
thumbs_dir = folder / _THUMBS_DIRNAME
|
||||
thumbs_dir.mkdir(exist_ok = True)
|
||||
# Key on the full filename (stem + extension), not the stem: two images that
|
||||
# share a stem but differ by extension (sample.png / sample.jpg) would otherwise
|
||||
# collide on one cache file, and an mtime-newer cache built for the first would
|
||||
# be served for the second, showing the wrong image in the labeling grid.
|
||||
thumb_path = thumbs_dir / f"{image_path.name}_{size}.jpg"
|
||||
src_mtime = image_path.stat().st_mtime
|
||||
if thumb_path.is_file() and thumb_path.stat().st_mtime >= src_mtime:
|
||||
|
|
@ -1646,12 +1683,24 @@ async def set_diffusion_dataset_caption(
|
|||
sidecar = image_path.with_suffix(".txt")
|
||||
if caption:
|
||||
sidecar.write_text(caption, encoding = "utf-8")
|
||||
else:
|
||||
# Blank clears the sidecar; also drop a stale .caption so the image reads as
|
||||
# uncaptioned afterwards.
|
||||
sidecar.unlink(missing_ok = True)
|
||||
image_path.with_suffix(".caption").unlink(missing_ok = True)
|
||||
return _image_record(folder, image_path, _load_metadata_captions(folder))
|
||||
return _image_record(folder, image_path, _load_metadata_captions(folder))
|
||||
# Blank must actually clear. Unlinking alone would resurface this image's
|
||||
# metadata.jsonl / captions.jsonl caption (the fallback source), so when one
|
||||
# exists write an EMPTY sidecar instead: both the record reader and the
|
||||
# trainer's discovery treat an existing sidecar as authoritative even when
|
||||
# empty, which makes it a tombstone. No metadata caption -> plain cleanup.
|
||||
meta = _load_metadata_captions(folder)
|
||||
try:
|
||||
rel = image_path.relative_to(folder).as_posix()
|
||||
except ValueError:
|
||||
rel = image_path.name
|
||||
if image_path.name in meta or rel in meta:
|
||||
sidecar.write_text("", encoding = "utf-8")
|
||||
else:
|
||||
sidecar.unlink(missing_ok = True)
|
||||
image_path.with_suffix(".caption").unlink(missing_ok = True)
|
||||
return _image_record(folder, image_path, meta)
|
||||
|
||||
return await asyncio.to_thread(write)
|
||||
|
||||
|
|
@ -1674,6 +1723,9 @@ async def delete_diffusion_dataset_image(
|
|||
image_path.with_suffix(ext).unlink(missing_ok = True)
|
||||
thumbs_dir = folder / _THUMBS_DIRNAME
|
||||
if thumbs_dir.is_dir():
|
||||
# Thumbs are keyed on the full filename (stem + extension), so match that
|
||||
# here too; a stem-only glob would leave this image's thumbs behind and
|
||||
# could delete a same-stem sibling's (sample.png vs sample.jpg).
|
||||
for t in thumbs_dir.glob(f"{image_path.name}_*.jpg"):
|
||||
t.unlink(missing_ok = True)
|
||||
return {"deleted": image_path.name}
|
||||
|
|
|
|||
|
|
@ -2121,3 +2121,34 @@ def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path):
|
|||
backend.generate(prompt = "a sloth")
|
||||
backend.generate(prompt = "another sloth")
|
||||
assert resets == [True, True]
|
||||
|
||||
|
||||
def test_prefetch_returns_snapshot_dir_for_manifest(monkeypatch):
|
||||
# The prefetched pipeline manifest's directory is the local snapshot root; a
|
||||
# config-only base list (no manifest) returns None so the hub id stays in use.
|
||||
backend = DiffusionBackend()
|
||||
monkeypatch.setattr(
|
||||
"utils.hf_xet_fallback.hf_hub_download_with_xet_fallback",
|
||||
lambda repo, fn, tok, **k: f"/cache/snap/{fn}",
|
||||
)
|
||||
root = backend._prefetch_files(
|
||||
"base/repo", None, "base/repo", ["model_index.json", "vae/x.safetensors"], None
|
||||
)
|
||||
assert root == "/cache/snap"
|
||||
assert (
|
||||
backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None) is None
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path):
|
||||
# With a prefetched snapshot, from_pretrained must receive the local dir --
|
||||
# its own hub sweep would re-download the root packaged singles the scoped
|
||||
# prefetch skips (24 GB per FLUX.1 repo).
|
||||
backend = DiffusionBackend()
|
||||
backend.load_pipeline(
|
||||
"unsloth/Qwen-Image-2512-bnb-4bit",
|
||||
model_kind = "pipeline",
|
||||
_base_local_dir = str(tmp_path),
|
||||
)
|
||||
assert _FakePipeline.last["base"] == str(tmp_path)
|
||||
backend.unload()
|
||||
|
|
|
|||
|
|
@ -232,14 +232,28 @@ def _state():
|
|||
)
|
||||
|
||||
|
||||
def _allow_cn_security(monkeypatch):
|
||||
"""Stub the Hub malware preflight to allow the load (hermetic, no network)."""
|
||||
import utils.security
|
||||
monkeypatch.setattr(
|
||||
utils.security,
|
||||
"evaluate_file_security",
|
||||
lambda name, hf_token = None, **kw: types.SimpleNamespace(blocked = False, reason = ""),
|
||||
)
|
||||
|
||||
|
||||
def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
|
||||
import threading
|
||||
|
||||
from core.inference.diffusion import DiffusionBackend
|
||||
|
||||
monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
|
||||
_allow_cn_security(monkeypatch)
|
||||
b = DiffusionBackend()
|
||||
st = _state()
|
||||
# The pipe cache only commits while ``st`` is the CURRENT load (an unload racing
|
||||
# from_pipe must not repopulate the cache), so mirror the loaded invariant.
|
||||
b._state = st
|
||||
resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False)
|
||||
p1 = b._controlnet_pipe(st, resolved, threading.Event())
|
||||
assert isinstance(p1, _FakeCNPipe) and isinstance(p1.controlnet, _FakeCNModel)
|
||||
|
|
@ -250,6 +264,69 @@ def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
|
|||
assert b._cn_models["flux-union-pro"] is p1.controlnet
|
||||
|
||||
|
||||
def test_controlnet_pipe_blocks_flagged_remote_repo(monkeypatch):
|
||||
# A bare owner/name ControlNet is accepted by resolve_controlnet without the base
|
||||
# trust gate, so the load path must run the Hub malware preflight: a flagged remote
|
||||
# repo must raise BEFORE from_pretrained downloads/deserializes it.
|
||||
import threading
|
||||
|
||||
import utils.security
|
||||
from core.inference.diffusion import DiffusionBackend
|
||||
|
||||
loaded = {"called": False}
|
||||
|
||||
class _TrapModel(_FakeCNModel):
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
path,
|
||||
torch_dtype = None,
|
||||
token = None,
|
||||
):
|
||||
loaded["called"] = True
|
||||
return super().from_pretrained(path, torch_dtype = torch_dtype, token = token)
|
||||
|
||||
mod = _fake_diffusers()
|
||||
mod.FluxControlNetModel = _TrapModel
|
||||
monkeypatch.setitem(sys.modules, "diffusers", mod)
|
||||
monkeypatch.setattr(
|
||||
utils.security,
|
||||
"evaluate_file_security",
|
||||
lambda name, hf_token = None, **kw: types.SimpleNamespace(
|
||||
blocked = True, reason = "Hugging Face security scan flagged unsafe files: evil.bin"
|
||||
),
|
||||
)
|
||||
b = DiffusionBackend()
|
||||
st = _state()
|
||||
b._state = st
|
||||
resolved = dc.ResolvedControlNet("evil/cn", "evil/cn", is_local = False)
|
||||
with pytest.raises(ValueError, match = "security scan flagged"):
|
||||
b._controlnet_pipe(st, resolved, threading.Event())
|
||||
assert loaded["called"] is False
|
||||
|
||||
|
||||
def test_controlnet_pipe_skips_scan_for_local_dir(monkeypatch, tmp_path):
|
||||
# A local dir the user picked has no Hub scan; the preflight must not block it even
|
||||
# if the (unused) scan stub would say blocked.
|
||||
import threading
|
||||
|
||||
import utils.security
|
||||
from core.inference.diffusion import DiffusionBackend
|
||||
|
||||
monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
|
||||
monkeypatch.setattr(
|
||||
utils.security,
|
||||
"evaluate_file_security",
|
||||
lambda name, hf_token = None, **kw: types.SimpleNamespace(blocked = True, reason = "x"),
|
||||
)
|
||||
b = DiffusionBackend()
|
||||
st = _state()
|
||||
b._state = st
|
||||
resolved = dc.ResolvedControlNet("my-cn", str(tmp_path), is_local = True)
|
||||
p = b._controlnet_pipe(st, resolved, threading.Event())
|
||||
assert isinstance(p, _FakeCNPipe)
|
||||
|
||||
|
||||
def test_controlnet_pipe_rejects_family_without_classes():
|
||||
import threading
|
||||
|
||||
|
|
@ -260,3 +337,19 @@ def test_controlnet_pipe_rejects_family_without_classes():
|
|||
st.family.controlnet_pipeline_class = None
|
||||
with pytest.raises(ValueError, match = "not supported"):
|
||||
b._controlnet_pipe(st, dc.ResolvedControlNet("x", "y", False), threading.Event())
|
||||
|
||||
|
||||
def test_controlnet_pipe_not_cached_after_unload_race(monkeypatch):
|
||||
# An unload that lands while from_pipe is assembling must not let the wrapper
|
||||
# repopulate the cache around the torn-down base pipe.
|
||||
import threading
|
||||
|
||||
from core.inference.diffusion import DiffusionBackend
|
||||
|
||||
monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
|
||||
b = DiffusionBackend()
|
||||
st = _state() # never committed to b._state: the load is already gone
|
||||
resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False)
|
||||
with pytest.raises(RuntimeError, match = "cancelled"):
|
||||
b._controlnet_pipe(st, resolved, threading.Event())
|
||||
assert b._cn_pipes == {}
|
||||
|
|
|
|||
|
|
@ -177,7 +177,8 @@ def test_delete_image_cleans_sidecar_and_thumb(client, ds_root):
|
|||
folder.mkdir()
|
||||
_write_png(folder / "x.png")
|
||||
(folder / "x.txt").write_text("cap", encoding = "utf-8")
|
||||
# Generate a thumbnail so we can assert it is cleaned up too.
|
||||
# Generate a thumbnail so we can assert it is cleaned up too. Thumbs are keyed on
|
||||
# the full filename (stem + extension) to avoid same-stem collisions across formats.
|
||||
client.get("/api/train/diffusion/dataset/d/image/x.png?thumb=32")
|
||||
# Thumb cache key includes the extension (x.png_32.jpg), so png and jpg
|
||||
# siblings can't collide.
|
||||
|
|
@ -190,6 +191,19 @@ def test_delete_image_cleans_sidecar_and_thumb(client, ds_root):
|
|||
assert not list((folder / ".thumbs").glob("x.png_*.jpg"))
|
||||
|
||||
|
||||
def test_thumb_cache_key_distinguishes_same_stem_extensions(client, ds_root):
|
||||
# sample.png and sample.jpg share a stem; each must get its OWN thumbnail cache
|
||||
# file, so the labeling grid never serves one image's thumbnail for the other.
|
||||
folder = ds_root / "d"
|
||||
folder.mkdir()
|
||||
Image.new("RGB", (8, 8), (10, 20, 30)).save(folder / "sample.png", format = "PNG")
|
||||
Image.new("RGB", (8, 8), (200, 210, 220)).save(folder / "sample.jpg", format = "JPEG")
|
||||
client.get("/api/train/diffusion/dataset/d/image/sample.png?thumb=32")
|
||||
client.get("/api/train/diffusion/dataset/d/image/sample.jpg?thumb=32")
|
||||
thumbs = sorted(p.name for p in (folder / ".thumbs").glob("*.jpg"))
|
||||
assert thumbs == ["sample.jpg_32.jpg", "sample.png_32.jpg"]
|
||||
|
||||
|
||||
# ── traversal / validation ───────────────────────────────────────────────────
|
||||
def test_dataset_name_traversal_rejected_over_http(client, ds_root):
|
||||
# A name that fails the folder-name validator returns 400, never touches disk.
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ import pytest
|
|||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from auth.authentication import get_current_subject
|
||||
from auth.authentication import authenticated_via_api_key, get_current_subject
|
||||
from core.training.diffusion_training_service import DiffusionTrainingService
|
||||
from routes.training import router as training_router
|
||||
|
||||
|
|
@ -265,11 +265,22 @@ def client(monkeypatch):
|
|||
|
||||
monkeypatch.setattr(tr, "get_training_backend", lambda: _FakeLLMBackend(active = False))
|
||||
monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None)
|
||||
# The dataset preflight runs the trainer's discovery against _BODY's fake
|
||||
# data_dir; stub it here so wiring tests pass, and let the dedicated preflight
|
||||
# tests below re-point it at a real tmp dataset.
|
||||
monkeypatch.setattr(
|
||||
"core.training.diffusion_train_common.discover_image_caption_pairs",
|
||||
lambda data_dir, **kw: [("img.png", "caption")],
|
||||
)
|
||||
app = FastAPI()
|
||||
app.include_router(training_router, prefix = "/api/train")
|
||||
app.dependency_overrides[get_current_subject] = lambda: "test-user"
|
||||
# Default to session (UI) auth: the API-key inference-in-flight guard is a no-op there,
|
||||
# so the wiring tests below behave as before. The guard test flips this override.
|
||||
app.dependency_overrides[authenticated_via_api_key] = lambda: False
|
||||
c = TestClient(app)
|
||||
c._fake = fake # type: ignore[attr-defined]
|
||||
c._app = app # type: ignore[attr-defined]
|
||||
return c
|
||||
|
||||
|
||||
|
|
@ -358,6 +369,37 @@ def test_route_start_conflict_maps_to_409(client):
|
|||
assert r.status_code == 409
|
||||
|
||||
|
||||
def test_route_start_over_api_with_inference_in_flight_is_409(client, monkeypatch):
|
||||
# An API-key client must not start diffusion training (which frees VRAM by unloading
|
||||
# chat) while an inference request is streaming; it should 409 instead of killing it.
|
||||
client._app.dependency_overrides[authenticated_via_api_key] = lambda: True
|
||||
monkeypatch.setattr(
|
||||
"core.inference.llama_keepwarm.other_inference_request_count",
|
||||
lambda current_request_counted = False: 1,
|
||||
)
|
||||
freed = {"called": False}
|
||||
import routes.training as tr
|
||||
|
||||
monkeypatch.setattr(
|
||||
tr, "_free_gpu_for_diffusion_training", lambda: freed.__setitem__("called", True)
|
||||
)
|
||||
r = client.post("/api/train/diffusion/start", json = _BODY)
|
||||
assert r.status_code == 409
|
||||
# The guard must run BEFORE any GPU is freed, so the live inference stream survives.
|
||||
assert freed["called"] is False
|
||||
|
||||
|
||||
def test_route_start_over_api_without_inference_proceeds(client, monkeypatch):
|
||||
# Same API-key path but no inference in flight: the start proceeds normally.
|
||||
client._app.dependency_overrides[authenticated_via_api_key] = lambda: True
|
||||
monkeypatch.setattr(
|
||||
"core.inference.llama_keepwarm.other_inference_request_count",
|
||||
lambda current_request_counted = False: 0,
|
||||
)
|
||||
r = client.post("/api/train/diffusion/start", json = _BODY)
|
||||
assert r.status_code == 200
|
||||
|
||||
|
||||
def test_route_status_and_stop(client):
|
||||
client.post("/api/train/diffusion/start", json = _BODY)
|
||||
s = client.get("/api/train/diffusion/status")
|
||||
|
|
|
|||
|
|
@ -2282,6 +2282,14 @@ export function HubModelPicker({
|
|||
}
|
||||
|
||||
if (section === "recommended") {
|
||||
// Curated safetensors rows render ABOVE the recommended rows (and call
|
||||
// getOptionProps), so their keys must lead here or they fall back to the
|
||||
// duplicate ...-option-missing id and drop out of arrow-key navigation.
|
||||
keys.push(
|
||||
...curatedSafetensorsRows.map((m) =>
|
||||
makeModelOptionKey("curated-safetensors", m.id),
|
||||
),
|
||||
);
|
||||
keys.push(
|
||||
...recommendedRows.map((r) => makeModelOptionKey("recommended", r.id)),
|
||||
);
|
||||
|
|
@ -2291,6 +2299,7 @@ export function HubModelPicker({
|
|||
}, [
|
||||
cachedReady,
|
||||
chatOnly,
|
||||
curatedSafetensorsRows,
|
||||
sortedCustomFolderModels,
|
||||
customFoldersCollapsed,
|
||||
downloadedCollapsed,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue