diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py new file mode 100644 index 0000000000..538a152c35 --- /dev/null +++ b/studio/backend/core/inference/diffusion.py @@ -0,0 +1,612 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Local diffusion (text-to-image) backend. + +A torch-only singleton: it dequantises a single-file GGUF on-device via +``GGUFQuantizationConfig`` and pulls the rest of the pipeline (VAE, text +encoders, scheduler) from the matching base repo. torch/diffusers are imported +lazily so this stays importable in a no-torch runtime. ``begin_load`` runs on a +background thread; poll ``load_progress`` for the download bar. GPU-handoff +policy lives in the arbiter the routes call, not here. +""" + +from __future__ import annotations + +import inspect +import threading +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional + +from loggers import get_logger +from utils.hardware import clear_gpu_cache + +from .diffusion_families import ( + DiffusionFamily, + detect_family, + resolve_base_repo, + resolve_local_gguf_child, +) + +logger = get_logger(__name__) + + +@dataclass(frozen = True) +class _LoadState: + """Everything about the currently-loaded pipeline, swapped as one unit.""" + + pipe: Any + family: Any + repo_id: str + base_repo: str + device: str + dtype: str + cpu_offload: bool + + +@dataclass +class _LoadingState: + """An in-flight background load, polled for download progress.""" + + repo_id: str + base_repo: str + expected_bytes: int = 0 + error: Optional[str] = None + + +@dataclass +class _GenState: + """An in-flight generation, updated per denoising step for the progress bar.""" + + total_steps: int + step: int = 0 + # Set when the first step finishes; the ETA rate is measured from there so the + # slower first step (warmup) doesn't skew it. + first_step_at: float = 0.0 + # Computed once per step (in the callback) so it's stable between polls. + eta_seconds: Optional[float] = None + + +def _estimate_eta(total_steps: int, step: int, first_step_at: float, now: float) -> Optional[float]: + """Seconds remaining, from the average step time measured after the first step. + None until at least one step has elapsed since the first.""" + steps_since_first = step - 1 + if not first_step_at or steps_since_first <= 0: + return None + per_step = (now - first_step_at) / steps_since_first + return max(0.0, (total_steps - step) * per_step) + + +class DiffusionBackend: + """Holds at most one loaded diffusers pipeline. All mutations are serialised.""" + + def __init__(self) -> None: + # One lock serialises load / generate / unload — a generate must not run + # while the pipeline is being swapped out from under it. status() and + # load_progress() read the single state references without the lock, so + # polling never blocks a slow load. + self._lock = threading.Lock() + self._state: Optional[_LoadState] = None + self._loading: Optional[_LoadingState] = None + # Bumped on every begin_load and unload so a worker whose load was + # superseded (a new load) or cancelled (unload, incl. an arbiter eviction) + # neither commits its pipeline nor stamps progress onto the current load. + self._load_token = 0 + # Set by unload() to abort an in-flight download (which runs without the + # lock, like the chat backend), so an eviction/unload can preempt a slow + # load instead of blocking on the lock for the whole download. + self._cancel_event = threading.Event() + # The callback mutates this and generate_progress() reads it, both without + # the lock (generate holds it for the whole call), so polling stays live. + self._gen: Optional[_GenState] = None + + @property + def is_loaded(self) -> bool: + return self._state is not None + + def _pick_device_and_dtype(self) -> tuple[str, Any]: + import torch + + if torch.cuda.is_available(): + # BF16 needs Ampere+ (compute capability >= 8); pre-Ampere cards + # (Turing/Volta/Pascal) only emulate it, so use FP16 there. (Checked by + # capability, not torch.cuda.is_bf16_supported(), which returns True via + # emulation on those cards and would still pick BF16.) + dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16 + return "cuda", dtype + # Intel XPU enables the Images page (CHAT_ONLY=False in hardware.py), so + # without this branch an Arc user would silently run diffusion on CPU. + if hasattr(torch, "xpu") and torch.xpu.is_available(): + return "xpu", torch.bfloat16 + mps = getattr(torch.backends, "mps", None) + if mps is not None and mps.is_available(): + return "mps", torch.float16 + return "cpu", torch.float32 + + def _resolve_gguf_path(self, repo_id: str, gguf_filename: str, hf_token: Optional[str]) -> str: + local_root = Path(repo_id).expanduser() + if local_root.exists(): + return str(resolve_local_gguf_child(local_root, gguf_filename)) + from huggingface_hub import hf_hub_download + + return hf_hub_download(repo_id, gguf_filename, token = hf_token) + + def _prefetch_files( + self, + repo_id: str, + gguf_filename: Optional[str], + base: str, + base_files: list[str], + hf_token: Optional[str], + ) -> None: + """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")``.""" + from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback + + # GGUF transformer (hub repos only; a local path is already on disk). + if gguf_filename and not Path(repo_id).expanduser().exists(): + hf_hub_download_with_xet_fallback( + repo_id, gguf_filename, hf_token, cancel_event = self._cancel_event + ) + # Base repo (VAE / text-encoder / scheduler); list comes from the estimate. + for rfilename in base_files: + if self._cancel_event.is_set(): + raise RuntimeError("Cancelled") + hf_hub_download_with_xet_fallback( + base, rfilename, hf_token, cancel_event = self._cancel_event + ) + + # ── Background load + progress ───────────────────────────────────────── + + def validate_load( + self, + repo_id: str, + *, + gguf_filename: Optional[str] = None, + family_override: Optional[str] = None, + ) -> DiffusionFamily: + """Cheap, pure-synchronous pre-flight: returns the resolved family or + raises ValueError. No I/O, so a caller can reject a bad request before + taking the GPU (begin_load and load_pipeline re-check on the load path).""" + if not gguf_filename: + raise ValueError( + "gguf_filename is required: this backend loads single-file GGUF checkpoints only." + ) + fam = detect_family(repo_id, family_override) + if fam is None: + raise ValueError( + f"'{repo_id}' isn't a supported image-generation model. " + f"Supported: Z-Image, Qwen-Image, FLUX.1, FLUX.2-klein." + ) + return fam + + def begin_load( + self, + repo_id: str, + *, + gguf_filename: Optional[str] = None, + base_repo: Optional[str] = None, + family_override: Optional[str] = None, + hf_token: Optional[str] = None, + cpu_offload: bool = False, + ) -> dict[str, Any]: + """Validate, then run the (slow) load on a daemon thread. Returns at once.""" + fam = self.validate_load( + repo_id, gguf_filename = gguf_filename, family_override = family_override + ) + + with self._lock: + # Allow starting over a previously-failed load, but not over a live one. + if self._loading is not None and self._loading.error is None: + raise RuntimeError("A diffusion load is already in progress.") + self._load_token += 1 + token = self._load_token + # Best-effort download preemption only; the token (not this event) is + # the real guard that a superseded worker can't commit its pipeline. + self._cancel_event.clear() + # Seed with the family fallback; the worker resolves the real base + # (a network lookup) and updates this, so begin_load never blocks. + self._loading = _LoadingState(repo_id = repo_id, base_repo = fam.base_repo) + + threading.Thread( + target = self._run_load, + kwargs = dict( + repo_id = repo_id, + gguf_filename = gguf_filename, + base_repo = base_repo, + family_override = family_override, + hf_token = hf_token, + cpu_offload = cpu_offload, + _load_token = token, + ), + daemon = True, + ).start() + return self.status() + + def _run_load(self, **kwargs: Any) -> None: + token = kwargs.get("_load_token") + try: + # Resolve the base repo and estimate sizes on this thread (both network + # calls) so begin_load returns instantly; the bar shows raw bytes until + # the total lands. This is the only writer of _loading's fields here. + fam = detect_family(kwargs["repo_id"], kwargs.get("family_override")) + base = _resolve_base_repo( + kwargs["repo_id"], kwargs.get("base_repo"), fam, kwargs.get("hf_token") + ) + kwargs["base_repo"] = base + expected, base_files = self._estimate_download_bytes( + kwargs["repo_id"], kwargs.get("gguf_filename"), base, kwargs.get("hf_token") + ) + with self._lock: + # Stamp progress only if this load is still current; a superseding + # load (or unload) has its own token and its own _LoadingState. + if self._load_token == token and self._loading is not None: + self._loading.base_repo = base + 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["repo_id"], + kwargs.get("gguf_filename"), + base, + base_files, + kwargs.get("hf_token"), + ) + self.load_pipeline(**kwargs) + with self._lock: + # Only clear the marker if this load is still the current one; a + # newer begin_load (or an unload) has its own token. + if self._load_token == token: + self._loading = None + except Exception as exc: # noqa: BLE001 — surfaced to the client via load_progress + # A cancelled/superseded load raised below; don't log it as a failure + # or stamp its error onto whatever load is current now. + if self._load_token != token: + return + logger.error("diffusion.load_failed: %s", exc) + # Redact native paths: this error is surfaced verbatim via the + # load-progress poll, and Studio can run as a shared server. + from utils.native_path_leases import redact_native_paths + + with self._lock: + if self._load_token == token and self._loading is not None: + self._loading.error = redact_native_paths(str(exc)) + + def load_progress(self) -> dict[str, Any]: + """Phase + downloaded/total bytes for the in-flight load (cache-scan based).""" + loading = self._loading + if loading is not None and loading.error: + return _progress("error", error = loading.error) + if loading is None: + return _progress("ready" if self._state is not None else None) + + downloaded = self._cache_bytes(loading.repo_id) + self._cache_bytes(loading.base_repo) + expected = loading.expected_bytes + # Downloads done but pipeline still dequantising / moving to GPU. The cache + # scan can slightly exceed the estimate (extra cached quants, blob padding), + # so clamp the reported bytes/fraction so the bar never overshoots 100%. + if expected > 0 and downloaded >= expected * 0.999: + return _progress("finalizing", min(downloaded, expected), expected, 1.0) + fraction = min(downloaded / expected, 1.0) if expected > 0 else 0.0 + return _progress("downloading", downloaded, expected, fraction) + + @staticmethod + def _estimate_download_bytes( + repo_id: str, gguf_filename: Optional[str], base_repo: str, hf_token: Optional[str] + ) -> tuple[int, list[str]]: + """Total download size for the progress bar, plus the base-repo files to + fetch (the prefetch reuses this list, so the base is listed only once).""" + from huggingface_hub import HfApi + + api = HfApi() + total = 0 + base_files: list[str] = [] + try: + if gguf_filename: + info = api.model_info(repo_id, files_metadata = True, token = hf_token) + total += sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename) + base_info = api.model_info(base_repo, files_metadata = True, token = hf_token) + for s in base_info.siblings: + if _base_file_downloaded(s.rfilename): + base_files.append(s.rfilename) + total += s.size or 0 + except Exception as exc: # noqa: BLE001 — estimate is best-effort + logger.warning("diffusion.size_estimate_failed: %s", exc) + return total, base_files + + @staticmethod + def _cache_bytes(repo_id: str) -> int: + from huggingface_hub import constants + + blobs = Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}" / "blobs" + total = 0 + try: + for entry in blobs.iterdir(): + try: + total += entry.stat().st_size + except OSError: + continue # broken symlink / unreadable + except OSError: + return 0 # repo not in cache yet + return total + + # ── Synchronous load / generate / unload ─────────────────────────────── + + def load_pipeline( + self, + repo_id: str, + *, + gguf_filename: Optional[str] = None, + base_repo: Optional[str] = None, + family_override: Optional[str] = None, + hf_token: Optional[str] = None, + cpu_offload: bool = False, + _load_token: Optional[int] = None, + ) -> dict[str, Any]: + import diffusers + + fam = self.validate_load( + repo_id, gguf_filename = gguf_filename, family_override = family_override + ) + base = _resolve_base_repo(repo_id, base_repo, fam, hf_token) + device, dtype = self._pick_device_and_dtype() + + with self._lock: + # Bail before the (slow, VRAM-heavy) build if an unload/eviction or a + # newer load superseded this one while we were resolving/downloading — + # otherwise an evicted load would resurrect a pipeline into VRAM. + if _load_token is not None and _load_token != self._load_token: + raise RuntimeError("Diffusion load was cancelled.") + + # Free the old pipeline before allocating the new one so two + # checkpoints never sit in VRAM at once. + self._unload_locked() + + # Dequantise the GGUF transformer on-device; the VAE / text-encoder / + # scheduler come from the base diffusers repo (the GGUF is transformer-only). + gguf_path = self._resolve_gguf_path(repo_id, gguf_filename, hf_token) + transformer_cls = getattr(diffusers, fam.transformer_class) + transformer = transformer_cls.from_single_file( + gguf_path, + quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype), + torch_dtype = dtype, + config = base, + subfolder = "transformer", + # Forward the token: the config is fetched from the (possibly gated) + # base repo before from_pretrained gets a chance to authenticate. + token = hf_token, + ) + + pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer} + if hf_token: + pipe_kwargs["token"] = hf_token + pipeline_cls = getattr(diffusers, fam.pipeline_class) + pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) + + use_offload = bool(cpu_offload and device == "cuda") + if use_offload: + pipe.enable_model_cpu_offload() + else: + pipe.to(device) + + self._state = _LoadState( + pipe = pipe, + family = fam, + repo_id = repo_id, + base_repo = base, + device = device, + dtype = str(dtype).replace("torch.", ""), + cpu_offload = use_offload, + ) + + logger.info("diffusion.loaded: repo=%s base=%s device=%s", repo_id, base, device) + return self.status() + + def generate( + self, + *, + prompt: str, + negative_prompt: Optional[str] = None, + width: int = 1024, + height: int = 1024, + # Fallbacks for a caller that passes nothing; the route always sends the + # per-model values the UI seeds (few steps / no CFG for distilled models, + # more steps / real CFG for full ones). + steps: int = 9, + guidance: float = 0.0, + seed: Optional[int] = None, + batch_size: int = 1, + ) -> dict[str, Any]: + import torch + with self._lock: + state = self._state + if state is None: + raise RuntimeError("No diffusion model is loaded.") + + generator = torch.Generator(device = state.device) + if seed is None: + # Draw a fresh random seed but keep it within JS's safe-integer + # range (< 2**53), so the reported seed round-trips through JSON + # and actually reproduces the image (a raw 64-bit seed would lose + # precision in the browser and the recipe couldn't be replayed). + seed = generator.seed() & ((1 << 53) - 1) + else: + seed = int(seed) + generator.manual_seed(seed) + + kwargs: dict[str, Any] = { + "prompt": prompt, + "width": width, + "height": height, + "num_inference_steps": steps, + # Most pipelines take guidance via "guidance_scale"; Qwen-Image + # uses "true_cfg_scale" (its distilled guidance is off). + state.family.cfg_kwarg: guidance, + "generator": generator, + # Generate the whole batch in one forward pass (VRAM-heavy). All + # share this call's seed, drawn sequentially from one generator. + "num_images_per_prompt": batch_size, + } + # Pipelines vary in which kwargs they accept (a distilled pipeline may + # take neither a negative prompt nor a step callback), so only pass + # those where the signature has them. + call_params = inspect.signature(state.pipe.__call__).parameters + if negative_prompt and "negative_prompt" in call_params: + kwargs["negative_prompt"] = negative_prompt + + gen = _GenState(total_steps = steps) + + def _on_step(pipe, step_index, timestep, callback_kwargs): + now = time.time() + gen.step = step_index + 1 + if gen.first_step_at == 0.0: + gen.first_step_at = now + gen.eta_seconds = _estimate_eta(gen.total_steps, gen.step, gen.first_step_at, now) + return callback_kwargs + + if "callback_on_step_end" in call_params: + kwargs["callback_on_step_end"] = _on_step + + self._gen = gen + try: + images = state.pipe(**kwargs).images + finally: + self._gen = None + # Return the PIL images (not yet encoded): the route embeds each + # image's recipe and persists it via the gallery. + return {"images": list(images), "seed": int(seed), "repo_id": state.repo_id} + + def generate_progress(self) -> dict[str, Any]: + """Live per-step progress for an in-flight generation (lock-free read).""" + gen = self._gen + if gen is None or gen.total_steps <= 0: + return { + "active": False, + "step": 0, + "total_steps": 0, + "fraction": 0.0, + "eta_seconds": None, + } + return { + "active": True, + "step": gen.step, + "total_steps": gen.total_steps, + "fraction": gen.step / gen.total_steps, # step is 1..total, never over 1.0 + "eta_seconds": gen.eta_seconds, + } + + def unload(self) -> dict[str, Any]: + # Abort an in-flight download (it runs without the lock and checks this), + # so unload/an eviction returns promptly instead of waiting it out. + self._cancel_event.set() + with self._lock: + self._unload_locked() + # Cancel any in-flight load (its worker checks this token before + # committing) and drop the marker so the next load starts clean. + self._load_token += 1 + self._loading = None + return self.status() + + def _unload_locked(self) -> None: + state = self._state + if state is None: + return + self._state = None + del state + clear_gpu_cache() + + def status(self) -> dict[str, Any]: + state = self._state + if state is None: + return { + "loaded": False, + "repo_id": None, + "family": None, + "base_repo": None, + "device": None, + "dtype": None, + "cpu_offload": False, + } + return { + "loaded": True, + "repo_id": state.repo_id, + "family": state.family.name, + "base_repo": state.base_repo, + "device": state.device, + "dtype": state.dtype, + "cpu_offload": state.cpu_offload, + } + + +def _resolve_base_repo( + repo_id: str, base_repo: Optional[str], fam: DiffusionFamily, hf_token: Optional[str] +) -> str: + """The companion diffusers repo: caller's base, else the GGUF repo's own + ``base_model`` tag, else the family fallback. Shared by both load paths so a + direct ``load_pipeline`` call resolves the variant base the same way.""" + return resolve_base_repo(fam, (base_repo or "").strip() or _hf_base_model(repo_id, hf_token)) + + +def _hf_base_model(repo_id: str, hf_token: Optional[str]) -> Optional[str]: + """The diffusers base repo from a GGUF repo's ``base_model`` tag, or None. + + Lets one family entry cover every variant (Turbo/full, schnell/dev, the + 2512 Qwen revision). Skipped for local paths; None on any lookup failure. + """ + if Path(repo_id).expanduser().exists(): + return None + try: + from huggingface_hub import HfApi + meta = HfApi().model_info(repo_id, token = hf_token).cardData or {} + except Exception: # noqa: BLE001 — best-effort; fall back to the family default + return None + base = meta.get("base_model") + if isinstance(base, list): + base = base[0] if base else None + return base if isinstance(base, str) and base.strip() else None + + +def _base_file_downloaded(rfilename: str) -> bool: + """True for base-repo files ``from_pretrained`` actually fetches. + + The transformer is supplied by the GGUF, and repo docs (``assets/``, the + top-level README/PDF/images) are never downloaded — counting them would peg + the progress estimate above what lands on disk, so the bar would sit short of + 100% for the whole pipeline-load phase instead of advancing to "finalizing". + """ + if rfilename.startswith("transformer/"): + return False + if "/" not in rfilename: # top-level: only the pipeline manifest is fetched + return rfilename == "model_index.json" + return not rfilename.startswith("assets/") + + +def _progress( + phase: Optional[str], + bytes_downloaded: int = 0, + bytes_total: int = 0, + fraction: float = 0.0, + *, + error: Optional[str] = None, +) -> dict[str, Any]: + return { + "phase": phase, + "bytes_downloaded": bytes_downloaded, + "bytes_total": bytes_total, + "fraction": fraction, + "error": error, + } + + +_diffusion_backend: Optional[DiffusionBackend] = None + + +def get_diffusion_backend() -> DiffusionBackend: + global _diffusion_backend + if _diffusion_backend is None: + _diffusion_backend = DiffusionBackend() + return _diffusion_backend diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py new file mode 100644 index 0000000000..450fa6f63a --- /dev/null +++ b/studio/backend/core/inference/diffusion_families.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pure helpers for diffusion model identification. + +No torch/diffusers imports here: everything in this module is a pure function of +its string/path arguments so it can be unit-tested without the heavy runtime. + +A diffusion checkpoint published as a single-file GGUF only carries the +transformer weights; the matching VAE / text encoders / scheduler come from a +companion ``diffusers`` base repo. ``DiffusionFamily`` maps a checkpoint to the +diffusers classes and base repo needed to assemble the full pipeline. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from pathlib import Path, PurePosixPath +from typing import Optional + + +@dataclass(frozen = True) +class DiffusionFamily: + name: str + pipeline_class: str + transformer_class: str + base_repo: str + # Pipeline kwarg carrying the guidance value. Most use "guidance_scale"; + # Qwen-Image's distilled guidance is off, so its real CFG is "true_cfg_scale". + cfg_kwarg: str = "guidance_scale" + # Extra lowercased substrings (besides ``name``) that map a repo id here. + aliases: tuple[str, ...] = field(default_factory = tuple) + + +# Keyed by architecture, not per model variant: a checkpoint's specific base repo +# is read from its HF base_model tag at load time, so one entry covers Turbo/full, +# schnell/dev, etc. base_repo here is only a fallback. Only archs whose diffusers +# transformer supports from_single_file load here (ERNIE-Image does not, yet). +# FLUX.2-dev and FLUX.2-klein-9B are left out only because their base diffusers +# repos are gated; the open klein-4B base stands in for the klein family below. +_FAMILIES: tuple[DiffusionFamily, ...] = ( + DiffusionFamily( + name = "flux.1", + pipeline_class = "FluxPipeline", + transformer_class = "FluxTransformer2DModel", + base_repo = "black-forest-labs/FLUX.1-schnell", + aliases = ("flux1", "flux-1"), + ), + # FLUX.2-klein is a distinct pipeline (Flux2KleinPipeline) with a Qwen3 text + # encoder, not the Mistral-based Flux2Pipeline; it must precede a generic + # flux match. The base Flux2Pipeline (FLUX.2-dev) is gated, so it's omitted. + DiffusionFamily( + name = "flux.2-klein", + pipeline_class = "Flux2KleinPipeline", + transformer_class = "Flux2Transformer2DModel", + base_repo = "black-forest-labs/FLUX.2-klein-4B", + aliases = ("flux2-klein",), + ), + DiffusionFamily( + name = "qwen-image", + pipeline_class = "QwenImagePipeline", + transformer_class = "QwenImageTransformer2DModel", + base_repo = "Qwen/Qwen-Image", + cfg_kwarg = "true_cfg_scale", + aliases = ("qwen_image", "qwenimage"), + ), + DiffusionFamily( + name = "z-image", + pipeline_class = "ZImagePipeline", + transformer_class = "ZImageTransformer2DModel", + base_repo = "Tongyi-MAI/Z-Image-Turbo", + aliases = ("zimage", "z_image"), + ), +) + +# Editing / inpaint checkpoints share an arch keyword but need a different +# pipeline and an input image, which this text-to-image backend doesn't drive. +_EDIT_KEYWORDS = ("edit", "kontext", "inpaint", "inpainting") + + +def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[DiffusionFamily]: + """Resolve a ``DiffusionFamily`` from a repo id, or an explicit override. + + ``override`` matches a family ``name`` or alias exactly; otherwise the repo + id is scanned for the first family whose name/alias appears in it. Image + editing checkpoints are rejected (None) since this backend is text-to-image. + """ + if override: + key = override.strip().lower() + for fam in _FAMILIES: + if key == fam.name or key in fam.aliases: + return fam + return None + needle = repo_id.lower() + # Match edit keywords as whole id segments, not raw substrings, so a normal + # text-to-image repo like ".../some-image-edition" isn't misread as an editing + # checkpoint. Qwen-Image-Edit / FLUX.1-Kontext still match (edit/kontext are + # whole tokens there). Split on both path separators so a Windows local path + # is segmented too. + segments = set(re.split(r"[-_./\\]+", needle)) + if any(kw in segments for kw in _EDIT_KEYWORDS): + return None + for fam in _FAMILIES: + if fam.name in needle or any(alias in needle for alias in fam.aliases): + return fam + return None + + +def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str: + """The companion diffusers repo: caller-supplied if given, else the family fallback.""" + base = (base_repo or "").strip() + return base or fam.base_repo + + +def resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path: + """Resolve ``gguf_filename`` to a file under ``repo_root``, rejecting escapes. + + ``gguf_filename`` is user-supplied, so an absolute path or a ``..`` segment + could otherwise read outside the local repo directory. + """ + if ( + Path(gguf_filename).is_absolute() + or PurePosixPath(gguf_filename).is_absolute() + or gguf_filename.startswith(("/", "\\")) + or "\\" in gguf_filename + ): + raise ValueError("gguf_filename must be a relative path inside the repo.") + rel = PurePosixPath(gguf_filename) + if any(part in ("", ".", "..") for part in rel.parts): + raise ValueError("gguf_filename must not contain '', '.', or '..' segments.") + # Resolve symlinks before the containment check: the guards above stop + # lexical escapes, but a symlink inside the repo could still point outside it. + repo_real = repo_root.resolve() + child = repo_root.joinpath(*rel.parts).resolve() + if child != repo_real and repo_real not in child.parents: + raise ValueError("gguf_filename must resolve to a file inside the repo.") + if not child.exists(): + raise FileNotFoundError(f"'{gguf_filename}' not found under {repo_root}.") + return child diff --git a/studio/backend/core/inference/gpu_arbiter.py b/studio/backend/core/inference/gpu_arbiter.py new file mode 100644 index 0000000000..daa0ef4e94 --- /dev/null +++ b/studio/backend/core/inference/gpu_arbiter.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Single-GPU arbiter for Studio's two heavy GPU consumers. + +The chat backends (llama-server + the Unsloth subprocess) and the diffusion +backend share one GPU. Before either takes the GPU it calls ``acquire_for(owner)``, +which evicts the current *other* owner so two large models never sit in VRAM at +once. The arbiter only sequences ownership; the actual freeing is delegated to +each backend's existing teardown. + +Eviction runs under the arbiter lock, so an ownership transfer is atomic with +respect to other acquires. +""" + +from __future__ import annotations + +import threading +from typing import Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +CHAT = "chat" +DIFFUSION = "diffusion" + +_lock = threading.Lock() +_owner: Optional[str] = None + + +def _evict_chat() -> None: + import time + + from core.inference import get_inference_backend + from routes.inference import get_llama_cpp_backend + + llama = get_llama_cpp_backend() + # is_active (process exists), not is_loaded (process exists AND healthy): a + # chat model still starting up holds/keeps allocating VRAM but isn't healthy + # yet, so gating on is_loaded would skip it and let the load race the + # diffusion pipeline. unload_model() sets _cancel_event and kills the process. + if llama.is_active: + llama.unload_model() + orchestrator = get_inference_backend() + if orchestrator.active_model_name: + orchestrator.unload_model(orchestrator.active_model_name) + # Kill the subprocess too, not just the model: its base CUDA context holds + # VRAM the diffusion pipeline needs. + orchestrator._shutdown_subprocess(timeout = 5.0) + # The driver reclaims the killed process's VRAM asynchronously; wait for free + # memory to settle before the diffusion pipeline allocates, mirroring the chat + # reload path — otherwise a warm chat→diffusion handoff can transiently OOM. + llama._wait_for_vram_settle(since_kill = time.monotonic()) + + +def _evict_diffusion() -> None: + from core.inference.diffusion import get_diffusion_backend + get_diffusion_backend().unload() + + +# Patchable in tests via monkeypatch.setitem. +_EVICTORS = {CHAT: _evict_chat, DIFFUSION: _evict_diffusion} + + +def acquire_for(owner: str) -> None: + """Make ``owner`` the sole GPU owner, evicting the other if it holds it.""" + global _owner + if owner not in _EVICTORS: + raise ValueError(f"unknown GPU owner: {owner!r}") + with _lock: + if _owner is not None and _owner != owner: + logger.info("gpu_arbiter: evicting %s for %s", _owner, owner) + _EVICTORS[_owner]() + _owner = owner + + +def release(owner: str) -> None: + """Drop ``owner``'s claim (no-op if it isn't the current owner).""" + global _owner + with _lock: + if _owner == owner: + _owner = None + + +def current_owner() -> Optional[str]: + return _owner diff --git a/studio/backend/core/inference/image_gallery.py b/studio/backend/core/inference/image_gallery.py new file mode 100644 index 0000000000..66016f807a --- /dev/null +++ b/studio/backend/core/inference/image_gallery.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Disk-backed persistence for generated images. + +Each image is a PNG under ``studio_root()/images`` with its full generation +recipe embedded as PNG text chunks: a structured ``unsloth`` JSON blob (the +source of truth the gallery reads back) plus an Automatic1111-style +``parameters`` string for interop with other tools. Because the recipe lives +inside the file, a downloaded PNG carries its own settings. + +The gallery is intentionally dumb storage: the route owns the metadata schema +and passes a plain dict; this module only writes/reads/sorts files. +""" + +from __future__ import annotations + +import base64 +import json +import re +import uuid +from pathlib import Path +from typing import Any, Optional + +from loggers import get_logger +from utils.paths import ensure_dir, studio_root + +logger = get_logger(__name__) + +# PNG text-chunk key holding our structured recipe JSON. +_META_KEY = "unsloth" +# Image ids are file stems; restrict to filename-safe chars so a crafted id +# can't escape the gallery directory. +_ID_RE = re.compile(r"^[A-Za-z0-9_-]{1,128}$") + + +def gallery_dir() -> Path: + return ensure_dir(studio_root() / "images") + + +def _params_text(meta: dict[str, Any]) -> str: + """Automatic1111-style ``parameters`` string for cross-tool interop.""" + lines = [str(meta.get("prompt", ""))] + negative = meta.get("negative_prompt") + if negative: + lines.append(f"Negative prompt: {negative}") + lines.append( + f"Steps: {meta.get('steps')}, CFG scale: {meta.get('guidance')}, " + f"Seed: {meta.get('seed')}, Size: {meta.get('width')}x{meta.get('height')}, " + f"Model: {meta.get('model', '')}" + ) + return "\n".join(lines) + + +def _png_bytes(image: Any, meta: dict[str, Any]) -> bytes: + import io + + from PIL.PngImagePlugin import PngInfo + + info = PngInfo() + info.add_text(_META_KEY, json.dumps(meta)) + info.add_text("parameters", _params_text(meta)) + buf = io.BytesIO() + image.save(buf, format = "PNG", pnginfo = info) + return buf.getvalue() + + +def save(image: Any, meta: dict[str, Any]) -> dict[str, Any]: + """Persist a PIL image with its recipe embedded; return the gallery record.""" + image_id = uuid.uuid4().hex + (gallery_dir() / f"{image_id}.png").write_bytes(_png_bytes(image, meta)) + return _record(image_id, meta) + + +def _record(image_id: str, meta: dict[str, Any]) -> dict[str, Any]: + return { + **meta, + "id": image_id, + "url": f"/api/inference/images/gallery/{image_id}/file", + } + + +def image_path(image_id: str) -> Optional[Path]: + """Resolve an id to its on-disk PNG, or None if missing / unsafe.""" + if not _ID_RE.match(image_id): + return None + path = gallery_dir() / f"{image_id}.png" + # Defence in depth: confirm the resolved path is still inside the gallery. + try: + path.resolve().relative_to(gallery_dir().resolve()) + except ValueError: + return None + return path if path.is_file() else None + + +def image_b64(image_id: str) -> Optional[str]: + path = image_path(image_id) + if path is None: + return None + return base64.b64encode(path.read_bytes()).decode("ascii") + + +# Recipe keys a gallery record must carry (the required GalleryImage fields, minus +# id/url which _record adds). A PNG missing any is treated as foreign and skipped, +# so a hand-dropped or older-schema file can't 500 the whole listing. +_REQUIRED_META = ("prompt", "width", "height", "steps", "guidance", "seed", "created_at") + + +def _read_meta(path: Path) -> Optional[dict[str, Any]]: + from PIL import Image + + try: + with Image.open(path) as im: + raw = im.text.get(_META_KEY) # type: ignore[attr-defined] + except Exception: + return None + if not raw: + return None + try: + meta = json.loads(raw) + except (ValueError, TypeError): + return None + if not isinstance(meta, dict) or any(k not in meta for k in _REQUIRED_META): + return None + return meta + + +def _mtime(path: Path) -> float: + try: + return path.stat().st_mtime + except OSError: + return 0.0 + + +def list_images(limit: Optional[int] = None, offset: int = 0) -> list[dict[str, Any]]: + """A newest-first window of images for infinite scroll. + + Ordered by file mtime (a cheap stat, ~= generation order) so a months-old + gallery isn't opened in full just to sort it; only the window's recipes are + read. limit=None returns everything from ``offset`` on.""" + try: + paths = list(gallery_dir().glob("*.png")) + except OSError: + return [] + paths.sort(key = _mtime, reverse = True) + window = paths[offset:] if limit is None else paths[offset : offset + limit] + records = [] + for path in window: + meta = _read_meta(path) + if meta is None: # not one of ours (no recipe chunk) — skip + continue + records.append(_record(path.stem, meta)) + return records + + +def delete(image_id: str) -> bool: + path = image_path(image_id) + if path is None: + return False + try: + path.unlink() + return True + except OSError as exc: + logger.warning("image_gallery.delete_failed: %s", exc) + return False + + +def clear() -> int: + """Delete every gallery PNG; return how many were removed.""" + removed = 0 + try: + paths = list(gallery_dir().glob("*.png")) + except OSError: + return 0 + for path in paths: + try: + path.unlink() + removed += 1 + except OSError: + continue + return removed diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 4a3162b09e..9f175a0045 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1680,3 +1680,118 @@ class AnthropicMessagesResponse(BaseModel): stop_reason: Optional[str] = None stop_sequence: Optional[str] = None usage: AnthropicUsage = Field(default_factory = AnthropicUsage) + + +# ── Diffusion (local text-to-image) ── + + +class DiffusionLoadRequest(BaseModel): + """Request to load a local diffusion (text-to-image) checkpoint.""" + + model_path: str = Field(..., description = "Diffusion GGUF repo id or local path") + gguf_filename: str = Field( + ..., description = "The chosen single-file GGUF quant inside model_path" + ) + base_repo: Optional[str] = Field( + None, description = "Companion diffusers repo for VAE/text-encoders (default: family base)" + ) + family_override: Optional[str] = Field( + None, description = "Force a family when it can't be inferred from the repo id" + ) + hf_token: Optional[str] = Field(None, description = "HuggingFace token for gated repos") + cpu_offload: bool = Field(False, description = "Enable model CPU offload to fit low-VRAM cards") + + +class DiffusionGenerateRequest(BaseModel): + """Request to generate one image from the loaded diffusion model.""" + + prompt: str = Field(..., min_length = 1, description = "Text prompt") + negative_prompt: Optional[str] = Field( + None, description = "What to avoid (if the model supports it)" + ) + width: int = Field(1024, ge = 256, le = 2048, description = "Image width in pixels (multiple of 16)") + height: int = Field( + 1024, ge = 256, le = 2048, description = "Image height in pixels (multiple of 16)" + ) + steps: int = Field(9, ge = 1, le = 100, description = "Number of denoising steps") + guidance: float = Field(0.0, ge = 0.0, le = 20.0, description = "Classifier-free guidance scale") + seed: Optional[int] = Field( + None, ge = 0, le = 2**64 - 1, description = "Seed for reproducibility (random if omitted)" + ) + batch_size: int = Field( + 1, ge = 1, le = 32, description = "Images generated in one forward pass (VRAM-heavy)" + ) + + @field_validator("width", "height") + @classmethod + def _multiple_of_16(cls, value: int) -> int: + # Z-Image requires dimensions divisible by 16 (8x VAE downsample + 2x + # patch). Non-multiples crash deep in the pipeline, so reject them here + # for a clean 422 instead of a cryptic 500. + if value % 16 != 0: + raise ValueError("must be a multiple of 16") + return value + + +class GalleryImage(BaseModel): + """A persisted image's full generation recipe (embedded in the PNG too).""" + + id: str = Field(..., description = "Stable id (the on-disk filename stem)") + url: str = Field(..., description = "Relative URL to fetch the PNG bytes") + prompt: str = Field(..., description = "Prompt used") + negative_prompt: Optional[str] = Field(None, description = "Negative prompt, if any") + width: int = Field(..., description = "Image width") + height: int = Field(..., description = "Image height") + steps: int = Field(..., description = "Denoising steps") + guidance: float = Field(..., description = "Guidance scale") + seed: int = Field(..., description = "Seed used") + batch_index: int = Field(0, description = "Position within its batch (0-based)") + model: Optional[str] = Field(None, description = "Model repo id that produced it") + created_at: float = Field(..., description = "Creation time (epoch seconds)") + + +class DiffusionGenerateResponse(BaseModel): + """The persisted gallery records for one generation call (a batch).""" + + images: list[GalleryImage] = Field(..., description = "Saved records, one per image in the batch") + + +class GalleryListResponse(BaseModel): + """A newest-first page of persisted images, for infinite scroll.""" + + images: list[GalleryImage] = Field(default_factory = list) + has_more: bool = Field(False, description = "Whether older images remain past this page") + + +class DiffusionGenerateProgressResponse(BaseModel): + """Live per-step progress for an in-flight generation.""" + + active: bool = Field(False, description = "Whether a generation is running") + step: int = Field(0, description = "Denoising steps completed so far") + total_steps: int = Field(0, description = "Total denoising steps for this run") + fraction: float = Field(0.0, description = "step / total_steps, clamped to [0,1]") + eta_seconds: Optional[float] = Field(None, description = "Estimated seconds remaining") + + +class DiffusionLoadProgressResponse(BaseModel): + """Download/finalize progress for an in-flight diffusion load.""" + + phase: Optional[Literal["downloading", "finalizing", "ready", "error"]] = Field( + None, description = "Load phase; null when idle" + ) + bytes_downloaded: int = Field(0, description = "Bytes present in the HF cache so far") + bytes_total: int = Field(0, description = "Estimated total bytes to download (0 = unknown)") + fraction: float = Field(0.0, description = "bytes_downloaded / bytes_total, clamped to [0,1]") + error: Optional[str] = Field(None, description = "Failure message when phase is 'error'") + + +class DiffusionStatusResponse(BaseModel): + """Current diffusion backend state.""" + + loaded: bool = Field(False, description = "Whether a diffusion model is loaded") + repo_id: Optional[str] = Field(None, description = "Loaded repo id or local path") + family: Optional[str] = Field(None, description = "Detected diffusion family") + base_repo: Optional[str] = Field(None, description = "Companion diffusers base repo") + device: Optional[str] = Field(None, description = "Device the pipeline is on") + dtype: Optional[str] = Field(None, description = "Compute dtype") + cpu_offload: bool = Field(False, description = "Whether CPU offload is engaged") diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 20dea5ec12..1c7ce86c2c 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -183,6 +183,12 @@ class LocalModelInfo(BaseModel): None, description = "Unix timestamp of latest observed update", ) + task: Optional[str] = Field( + None, + description = "HF pipeline task inferred from a GGUF's architecture " + "('text-to-image' for diffusion, 'text-generation' otherwise). Lets the " + "Images picker show only diffusion GGUFs.", + ) class LocalModelListResponse(BaseModel): diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 1b7f7a668c..b75f114c21 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -18,6 +18,10 @@ cryptography>=42.0.0 boto3>=1.34.0 # optional: S3 dataset loading httpx>=0.27.0 fastmcp>=3.0.2 +# Local image generation (Images page): diffusers' GGUFQuantizationConfig needs +# the gguf reader. diffusers itself is torch-bound and comes via the ML stack; +# gguf is torch-free so it lives here, where it's installed in every mode. +gguf # RAG (knowledge bases, hybrid retrieval). sentence-transformers lives in # extras-no-deps.txt; these add the lexical+dense store and document parsing. sqlite-vec==0.1.9 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index caf2a262a3..718814ced0 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1049,6 +1049,14 @@ from models.inference import ( LoadRequest, UnloadRequest, GenerateRequest, + DiffusionLoadRequest, + DiffusionGenerateRequest, + DiffusionGenerateResponse, + DiffusionGenerateProgressResponse, + DiffusionStatusResponse, + DiffusionLoadProgressResponse, + GalleryImage, + GalleryListResponse, LoadResponse, LoadProgressResponse, UnloadResponse, @@ -2510,6 +2518,14 @@ async def load_model( user_override = request.chat_template_override, ) + # Reclaim the GPU from the diffusion (Images) backend before any chat + # load path — including the already-loaded fast path below, so chat + # ownership is asserted without depending on chat/diffusion exclusivity + # holding. No-op when diffusion isn't loaded. + from core.inference.gpu_arbiter import acquire_for, CHAT + + await asyncio.to_thread(acquire_for, CHAT) + # ── Already-loaded check: skip reload if the exact model is active ── backend = get_inference_backend() llama_backend = get_llama_cpp_backend() @@ -10150,3 +10166,231 @@ async def _openai_passthrough_non_streaming( # redundant parse + re-serialize round-trip. return Response(content = resp.content, media_type = "application/json") return JSONResponse(content = data) + + +# ────────────────────────────────────────────────────────────────────────── +# Diffusion (local text-to-image) +# +# Studio-only routes (studio_router is not mounted under /v1). The diffusion +# backend runs in-process and is synchronous, so the blocking load/generate/ +# unload calls are offloaded with asyncio.to_thread to keep the event loop free. +# This is the single error boundary: backend methods raise, we map to HTTP here. +# ────────────────────────────────────────────────────────────────────────── + + +def _guard_diffusion_load_against_training() -> None: + """Refuse loading an image model while a training run is active. Unlike chat, + a diffusion pipeline's VRAM can't be cheaply estimated before the load, so the + load is refused outright rather than fit-checked. No-op when training is + inactive or its state can't be read. Raises HTTP 409.""" + from core.training import get_training_backend + + try: + if not get_training_backend().is_training_active(): + return + except Exception as e: + logger.warning("Could not check training state for image-load guard: %s", e) + return + raise HTTPException( + status_code = 409, + detail = ( + "Can't load an image model while training is running: the diffusion " + "pipeline would compete with the training run for GPU memory. Training " + "was left untouched. Try again after training finishes." + ), + ) + + +@studio_router.post("/images/load", response_model = DiffusionStatusResponse) +async def load_diffusion_model( + request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject) +): + from core.inference.diffusion import get_diffusion_backend + from core.inference.gpu_arbiter import acquire_for, DIFFUSION + from utils.native_path_leases import redact_native_paths + + backend = get_diffusion_backend() + try: + # Reject an unsupported/malformed request BEFORE taking the GPU, so a + # bad image-model pick doesn't evict the user's loaded chat model only + # to 400. validate_load is pure (no I/O), so it runs inline. + backend.validate_load( + request.model_path, + gguf_filename = request.gguf_filename, + family_override = request.family_override, + ) + # Refuse while training is running: a multi-GB diffusion pipeline would + # compete with the training subprocess for VRAM. The chat path does the + # same via _guard_chat_load_against_training; this is its image sibling. + _guard_diffusion_load_against_training() + # Take the GPU from the chat backend, then kick the (slow) load onto a + # background thread and return at once — the client polls images/load-progress. + await asyncio.to_thread(acquire_for, DIFFUSION) + status_dict = await asyncio.to_thread( + backend.begin_load, + request.model_path, + gguf_filename = request.gguf_filename, + base_repo = request.base_repo, + family_override = request.family_override, + hf_token = request.hf_token, + cpu_offload = request.cpu_offload, + ) + return DiffusionStatusResponse(**status_dict) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc))) + except RuntimeError as exc: + # A load is already in progress. + raise HTTPException(status_code = 409, detail = str(exc)) + + +@studio_router.post("/images/generate", response_model = DiffusionGenerateResponse) +async def generate_diffusion_image( + request: DiffusionGenerateRequest, current_subject: str = Depends(get_current_subject) +): + from core.inference import image_gallery + from core.inference.diffusion import get_diffusion_backend + + backend = get_diffusion_backend() + try: + result = await asyncio.to_thread( + backend.generate, + prompt = request.prompt, + negative_prompt = request.negative_prompt, + width = request.width, + height = request.height, + steps = request.steps, + guidance = request.guidance, + seed = request.seed, + batch_size = request.batch_size, + ) + except RuntimeError as exc: + if not backend.is_loaded: + # The only genuine client-state 409: nothing is loaded to generate with. + raise HTTPException(status_code = 409, detail = "No diffusion model is loaded.") + # A pipeline RuntimeError (CUDA OOM, shape/device) is a server failure; fall + # through to the sanitized 500 instead of echoing raw exception text (which + # would 409 an OOM as retryable and leak VRAM totals / tensor shapes). + logger.error("diffusion.generate_failed: %s", exc) + raise HTTPException(status_code = 500, detail = "Image generation failed.") + except Exception as exc: + logger.error("diffusion.generate_failed: %s", exc) + raise HTTPException(status_code = 500, detail = "Image generation failed.") + + # Persist each image with its full recipe embedded. The whole batch shares + # one seed (drawn sequentially from one generator) and one timestamp (the + # images are generated together), so their gallery order is stable on reload. + created_at = time.time() + + def _persist() -> list[dict]: + records = [] + for index, image in enumerate(result["images"]): + records.append( + image_gallery.save( + image, + { + "prompt": request.prompt, + "negative_prompt": request.negative_prompt, + "width": request.width, + "height": request.height, + "steps": request.steps, + "guidance": request.guidance, + "seed": result["seed"], + # Position within the batch: images here share a seed + timestamp, + # so the export filename needs this to stay unique. + "batch_index": index, + "model": result.get("repo_id"), + "created_at": created_at, + }, + ) + ) + return records + + try: + records = await asyncio.to_thread(_persist) + except Exception as exc: + logger.error("diffusion.persist_failed: %s", exc) + raise HTTPException(status_code = 500, detail = "Failed to save the generated image.") + + return DiffusionGenerateResponse(images = [GalleryImage(**r) for r in records]) + + +@studio_router.get("/images/gallery", response_model = GalleryListResponse) +async def list_gallery_images( + limit: int = 50, + offset: int = 0, + current_subject: str = Depends(get_current_subject), +): + from core.inference import image_gallery + + limit = max(1, min(limit, 200)) + offset = max(0, offset) + # Fetch one extra to learn whether more remain, without a second scan. + records = await asyncio.to_thread(image_gallery.list_images, limit + 1, offset) + has_more = len(records) > limit + return GalleryListResponse( + images = [GalleryImage(**r) for r in records[:limit]], + has_more = has_more, + ) + + +@studio_router.get("/images/gallery/{image_id}/file") +async def get_gallery_image_file( + image_id: str, current_subject: str = Depends(get_current_subject) +): + from core.inference import image_gallery + + path = await asyncio.to_thread(image_gallery.image_path, image_id) + if path is None: + raise HTTPException(status_code = 404, detail = "Image not found.") + data = await asyncio.to_thread(path.read_bytes) + # Immutable content (id is unique per image), so let the browser cache it. + return Response( + content = data, + media_type = "image/png", + headers = {"Cache-Control": "private, max-age=31536000, immutable"}, + ) + + +@studio_router.delete("/images/gallery/{image_id}") +async def delete_gallery_image(image_id: str, current_subject: str = Depends(get_current_subject)): + from core.inference import image_gallery + + deleted = await asyncio.to_thread(image_gallery.delete, image_id) + if not deleted: + raise HTTPException(status_code = 404, detail = "Image not found.") + return {"deleted": True} + + +@studio_router.delete("/images/gallery") +async def clear_gallery_images(current_subject: str = Depends(get_current_subject)): + from core.inference import image_gallery + removed = await asyncio.to_thread(image_gallery.clear) + return {"removed": removed} + + +@studio_router.post("/images/unload", response_model = DiffusionStatusResponse) +async def unload_diffusion_model(current_subject: str = Depends(get_current_subject)): + from core.inference.diffusion import get_diffusion_backend + from core.inference.gpu_arbiter import release, DIFFUSION + + status_dict = await asyncio.to_thread(get_diffusion_backend().unload) + release(DIFFUSION) + return DiffusionStatusResponse(**status_dict) + + +@studio_router.get("/images/status", response_model = DiffusionStatusResponse) +async def diffusion_status(current_subject: str = Depends(get_current_subject)): + from core.inference.diffusion import get_diffusion_backend + return DiffusionStatusResponse(**get_diffusion_backend().status()) + + +@studio_router.get("/images/load-progress", response_model = DiffusionLoadProgressResponse) +async def diffusion_load_progress(current_subject: str = Depends(get_current_subject)): + from core.inference.diffusion import get_diffusion_backend + return DiffusionLoadProgressResponse(**get_diffusion_backend().load_progress()) + + +@studio_router.get("/images/generate-progress", response_model = DiffusionGenerateProgressResponse) +async def diffusion_generate_progress(current_subject: str = Depends(get_current_subject)): + from core.inference.diffusion import get_diffusion_backend + return DiffusionGenerateProgressResponse(**get_diffusion_backend().generate_progress()) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index e22f65751c..d5fa9f2664 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -859,6 +859,10 @@ async def list_local_models( try: models = collect_local_models(models_root) + # Tag each GGUF with its task so the Images picker can filter to diffusion. + models = [ + m.model_copy(update = {"task": _local_model_task(m.path, m.model_format)}) for m in models + ] return LocalModelListResponse( models_dir = str(models_root), @@ -3072,6 +3076,79 @@ def _repo_gguf_last_modified(repo_info) -> float: return latest +# GGUF general.architecture values that denote a diffusion (image) model; +# everything else is treated as a text model. Lets the Images picker show only +# image GGUFs in its On Device list. +_DIFFUSION_GGUF_ARCHS = frozenset( + { + "flux", + "flux2", + "sd1", + "sd2", + "sd3", + "sdxl", + "stable_diffusion", + "lumina2", + "qwen_image", + "qwenimage", + "auraflow", + "pixart", + "hunyuan_video", + "wan", + } +) + + +def _gguf_architecture(path: str) -> Optional[str]: + """The GGUF ``general.architecture``, or None. Delegates to the shared, + bounds-checked header reader (cached by path/mtime/size).""" + from utils.models.gguf_metadata import read_gguf_general_metadata + + arch = (read_gguf_general_metadata(path) or {}).get("general.architecture") + return arch.strip() if isinstance(arch, str) and arch.strip() else None + + +def _arch_to_task(arch: Optional[str]) -> Optional[str]: + if arch is None: + return None + return "text-to-image" if arch.lower() in _DIFFUSION_GGUF_ARCHS else "text-generation" + + +def _repo_gguf_task(repo_info) -> Optional[str]: + """HF pipeline task of a cached GGUF repo, from its architecture: + 'text-to-image' for diffusion archs, else 'text-generation' (None if unreadable).""" + try: + for path in _iter_gguf_paths(Path(repo_info.repo_path)): + if _is_mmproj_filename(path.name): + continue + task = _arch_to_task(_gguf_architecture(str(path))) + if task is not None: + return task + except Exception: + pass + return None + + +def _local_model_task(path: str, model_format: Optional[str]) -> Optional[str]: + """Same classification for a local model: read its GGUF architecture. The + path may be the .gguf file itself or a folder containing one.""" + if model_format != "gguf": + return None + try: + p = Path(path) + if p.suffix.lower() == ".gguf" and p.is_file(): + return _arch_to_task(_gguf_architecture(str(p))) + for f in _iter_gguf_paths(p): + if _is_mmproj_filename(f.name): + continue + task = _arch_to_task(_gguf_architecture(str(f))) + if task is not None: + return task + except Exception: + pass + return None + + @router.get("/cached-gguf") async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): """List GGUF repos downloaded to HF cache, legacy Unsloth cache, and HF default cache.""" @@ -3099,6 +3176,7 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): "size_bytes": total_size, "cache_path": str(repo_info.repo_path), "has_vision": _repo_has_mmproj(repo_info), + "task": _repo_gguf_task(repo_info), } # Keep the newest timestamp across duplicate caches; # attach only when known so absent rows sort as oldest. @@ -3123,6 +3201,20 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)): return {"cached": []} +def _repo_is_diffusers(repo_info) -> bool: + """A diffusers pipeline repo (e.g. a Z-Image / FLUX base) carries a top-level + model_index.json. These render images rather than chat, so the chat picker + hides them — mirroring how cached diffusion GGUFs are classified by arch.""" + try: + for rev in repo_info.revisions: + for f in rev.files: + if f.file_name == "model_index.json" or f.file_name.endswith("/model_index.json"): + return True + except Exception: + pass + return False + + @router.get("/cached-models") async def list_cached_models(current_subject: str = Depends(get_current_subject)): """List non-GGUF model repos downloaded to HF cache, legacy Unsloth cache, and HF default cache.""" @@ -3169,6 +3261,7 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject) row = { "repo_id": repo_id, "size_bytes": total_size, + "task": "text-to-image" if _repo_is_diffusers(repo_info) else None, } # Keep the newest timestamp across duplicate caches; # attach only when known so absent rows sort as oldest. @@ -3238,6 +3331,24 @@ async def delete_cached_model( except Exception: pass + # Also refuse if the diffusion (Images) backend has this repo loaded; its + # delete guard is otherwise chat-only, so its GGUF could be removed from + # under a live pipeline. Repo-level match, like the chat guards above. + try: + from core.inference.diffusion import get_diffusion_backend + diffusion_status = get_diffusion_backend().status() + if diffusion_status.get("loaded") and diffusion_status.get("repo_id"): + loaded_id = str(diffusion_status["repo_id"]).lower() + if loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()): + raise HTTPException( + status_code = 400, + detail = "Unload the model before deleting", + ) + except HTTPException: + raise + except Exception: + pass + try: cache_scans = _all_hf_cache_scans() diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 3818fe9f73..4056493af9 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -362,6 +362,25 @@ async def start_training( except Exception as e: logger.warning("Could not shut down export subprocess: %s", e) + try: + # A resident or in-flight diffusion (Images) pipeline also holds + # GPU memory the training run needs, and it can't be cheaply sized, + # so tear it down unconditionally like the export subprocess above + # (the chat block below fit-checks; diffusion can't). unload() is a + # no-op when nothing is loaded and also preempts an in-flight load; + # release the arbiter so it doesn't think the gone pipeline owns + # the GPU. Must precede the chat block, which early-returns. + from core.inference import gpu_arbiter + from core.inference.diffusion import get_diffusion_backend + + diffusion = get_diffusion_backend() + if diffusion.is_loaded: + logger.info("Unloading diffusion (Images) model to free GPU memory for training") + diffusion.unload() + gpu_arbiter.release(gpu_arbiter.DIFFUSION) + except Exception as e: + logger.warning("Could not unload diffusion model for training: %s", e) + try: from routes.training_vram import ( can_keep_chat_during_training, diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index b2ead305ba..b8ffb5a3e6 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -83,6 +83,7 @@ def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf(monk "size_bytes": 5_000, "cache_path": str(repo.repo_path), "has_vision": False, + "task": None, } ] @@ -105,6 +106,7 @@ def test_list_cached_gguf_matches_extension_case_insensitively(monkeypatch, tmp_ "size_bytes": 7_000, "cache_path": str(repo.repo_path), "has_vision": False, + "task": None, } ] @@ -200,6 +202,7 @@ def test_list_cached_gguf_keeps_largest_duplicate_repo_across_scans(monkeypatch, "size_bytes": 6_000, "cache_path": str(larger.repo_path), "has_vision": False, + "task": None, } ] @@ -230,6 +233,7 @@ def test_list_cached_gguf_dedupes_shared_blobs_across_revisions(monkeypatch, tmp "size_bytes": 5_000, "cache_path": str(repo.repo_path), "has_vision": False, + "task": None, } ] @@ -280,6 +284,7 @@ def test_list_cached_gguf_includes_mixed_repo_with_gguf_and_safetensors(monkeypa "size_bytes": 5_000, "cache_path": str(mixed.repo_path), "has_vision": False, + "task": None, } ] @@ -307,6 +312,7 @@ def test_list_cached_gguf_handles_none_size_on_disk(monkeypatch, tmp_path): "size_bytes": 5_000, "cache_path": str(partial.repo_path), "has_vision": False, + "task": None, } ] @@ -343,6 +349,7 @@ def test_list_cached_gguf_skips_malformed_repo_without_wiping_response(monkeypat "size_bytes": 5_000, "cache_path": str(healthy.repo_path), "has_vision": False, + "task": None, } ] @@ -390,7 +397,35 @@ def test_list_cached_models_includes_repo_with_only_mmproj_gguf(monkeypatch, tmp result = asyncio.run(models_route.list_cached_models(current_subject = "test-user")) - assert result["cached"] == [{"repo_id": "Org/MmprojAux", "size_bytes": 15_000}] + assert result["cached"] == [{"repo_id": "Org/MmprojAux", "size_bytes": 15_000, "task": None}] + + +def test_list_cached_models_tags_diffusers_pipeline_as_text_to_image(monkeypatch, tmp_path): + """A cached diffusers pipeline repo (model_index.json present) is tagged + text-to-image so the chat picker hides it, while a plain checkpoint isn't.""" + diffusion = _repo( + "Tongyi-MAI/Z-Image-Turbo", + [_file("model_index.json", 1_000), _file("text_encoder/model.safetensors", 9_000)], + tmp_path / "models--Tongyi-MAI--Z-Image-Turbo", + ) + checkpoint = _repo( + "unsloth/Llama-3.2-1B-Instruct", + [_file("config.json", 1_000), _file("model.safetensors", 9_000)], + tmp_path / "models--unsloth--Llama-3.2-1B-Instruct", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [diffusion, checkpoint])], + ) + + result = asyncio.run(models_route.list_cached_models(current_subject = "test-user")) + by_repo = {c["repo_id"]: c["task"] for c in result["cached"]} + assert by_repo == { + "Tongyi-MAI/Z-Image-Turbo": "text-to-image", + "unsloth/Llama-3.2-1B-Instruct": None, + } def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeypatch, tmp_path): @@ -419,6 +454,7 @@ def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeyp "size_bytes": 5_000, "cache_path": str(vision_repo.repo_path), "has_vision": True, + "task": None, } ] @@ -473,6 +509,7 @@ def test_all_hf_cache_scans_survives_inaccessible_aux_cache(monkeypatch, tmp_pat "size_bytes": 5_000, "cache_path": str(tmp_path / "active"), "has_vision": False, + "task": None, } ] @@ -669,3 +706,39 @@ def test_gguf_download_progress_counts_quant_subdir(monkeypatch, tmp_path): assert result["downloaded_bytes"] == 20_000 assert result["progress"] == 1.0 + + +def test_delete_cached_refuses_diffusion_loaded_repo(monkeypatch): + # The cached-delete guard refuses deleting a repo the diffusion (Images) + # backend has loaded, mirroring the chat guard, so its GGUF can't be removed + # from under a live pipeline. + from fastapi import HTTPException + import core.inference.diffusion as diffusion_mod + import routes.inference as routes_inference + + # Chat and orchestrator report nothing loaded; only diffusion holds the repo. + # delete_cached_model resolves get_inference_backend from the models module + # namespace, so patch it there (not on core.inference) to isolate that guard. + monkeypatch.setattr( + routes_inference, "get_llama_cpp_backend", + lambda: SimpleNamespace(is_loaded = False, model_identifier = None), + ) + monkeypatch.setattr( + models_route, "get_inference_backend", + lambda: SimpleNamespace(active_model_name = None), + ) + monkeypatch.setattr( + diffusion_mod, "get_diffusion_backend", + lambda: SimpleNamespace(status = lambda: {"loaded": True, "repo_id": "org/Z-Image-GGUF"}), + ) + + try: + asyncio.run( + models_route.delete_cached_model( + repo_id = "org/Z-Image-GGUF", variant = None, current_subject = "u", + ) + ) + assert False, "expected HTTPException refusing the delete" + except HTTPException as e: + assert e.status_code == 400 + assert "Unload the model before deleting" in e.detail diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py new file mode 100644 index 0000000000..5fc5d23bbe --- /dev/null +++ b/studio/backend/tests/test_diffusion_backend.py @@ -0,0 +1,529 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""CPU-only unit tests for the diffusion backend. + +The family helpers are pure functions, tested directly. The backend lifecycle is +exercised with ``torch`` / ``diffusers`` stubbed via ``sys.modules`` so no real +GPU, weights, or network access is needed (sub-second, CI-friendly). +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from core.inference.diffusion import ( + DiffusionBackend, + _base_file_downloaded, +) +from core.inference.diffusion_families import ( + detect_family, + resolve_base_repo, + resolve_local_gguf_child, +) + + +# Pure family helpers + + +def test_detect_family_from_repo_id(): + # Detection is by architecture; Turbo/full and schnell/dev map to one family. + assert detect_family("unsloth/Z-Image-Turbo-GGUF").name == "z-image" + assert detect_family("unsloth/Z-Image-GGUF").name == "z-image" + assert detect_family("unsloth/Qwen-Image-2512-GGUF").name == "qwen-image" + assert detect_family("unsloth/FLUX.1-schnell-GGUF").name == "flux.1" + # FLUX.2-klein is its own pipeline (Qwen3 TE), distinct from FLUX.1. + klein = detect_family("unsloth/FLUX.2-klein-4B-GGUF") + assert klein.name == "flux.2-klein" + assert klein.pipeline_class == "Flux2KleinPipeline" + assert klein.cfg_kwarg == "guidance_scale" + # Both klein sizes share the one family (base repo resolved per-variant). + assert detect_family("unsloth/FLUX.2-klein-9B-GGUF").name == "flux.2-klein" + # Only klein is wired up; the Mistral-based FLUX.2-dev base repo is gated. + assert detect_family("unsloth/FLUX.2-dev-GGUF") is None + # Qwen-Image guides via true_cfg_scale, not guidance_scale. + assert detect_family("unsloth/Qwen-Image-2512-GGUF").cfg_kwarg == "true_cfg_scale" + assert detect_family("unsloth/Z-Image-GGUF").cfg_kwarg == "guidance_scale" + # Image-editing checkpoints are rejected (text-to-image backend only): the + # edit keyword is matched as a whole id segment, so an "edit" that's only a + # substring of a normal word ("Edition") still loads. + assert detect_family("unsloth/Qwen-Image-Edit-2511-GGUF") is None + assert detect_family("unsloth/FLUX.1-Kontext-dev-GGUF") is None + assert detect_family("unsloth/Qwen-Image-Inpainting-GGUF") is None + assert detect_family("unsloth/Z-Image-Edition-GGUF").name == "z-image" + assert detect_family("meta-llama/Llama-3-8B") is None + + +def test_detect_family_override(): + assert detect_family("local/path", override = "z-image").name == "z-image" + assert detect_family("local/path", override = "zimage").name == "z-image" + assert detect_family("local/path", override = "not-a-family") is None + + +def test_resolve_base_repo(): + fam = detect_family("x", override = "z-image") + assert resolve_base_repo(fam, None) == fam.base_repo + assert resolve_base_repo(fam, " ") == fam.base_repo + assert resolve_base_repo(fam, "custom/base") == "custom/base" + + +def test_resolve_local_gguf_child(tmp_path): + (tmp_path / "model.gguf").write_bytes(b"x") + assert resolve_local_gguf_child(tmp_path, "model.gguf") == (tmp_path / "model.gguf").resolve() + with pytest.raises(ValueError): + resolve_local_gguf_child(tmp_path, "/etc/passwd") + with pytest.raises(ValueError): + resolve_local_gguf_child(tmp_path, "../secret.gguf") + with pytest.raises(ValueError): + resolve_local_gguf_child(tmp_path, "..\\secret.gguf") + with pytest.raises(FileNotFoundError): + resolve_local_gguf_child(tmp_path, "missing.gguf") + + +def test_resolve_local_gguf_child_blocks_symlink_escape(tmp_path): + outside = tmp_path / "outside.gguf" + outside.write_bytes(b"secret") + repo = tmp_path / "repo" + repo.mkdir() + try: + (repo / "model.gguf").symlink_to(outside) + except (OSError, NotImplementedError): + pytest.skip("symlinks not supported on this platform") + with pytest.raises(ValueError): + resolve_local_gguf_child(repo, "model.gguf") + + +# Stubbed runtime for backend lifecycle + + +class _FakeDtype: + def __init__(self, name: str) -> None: + self._name = name + + def __repr__(self) -> str: + return f"torch.{self._name}" + + __str__ = __repr__ + + +class _FakeGenerator: + def __init__(self, device = None) -> None: + self.device = device + self.manual = None + + def seed(self) -> int: + return 4242 + + def manual_seed(self, value: int): + self.manual = value + return self + + +class _FakeImage: + """Stand-in for a generated PIL image (the route persists it; here we only + count how many come back).""" + + +class _FakePipe: + def __init__(self) -> None: + self.moved_to = None + self.offloaded = False + self.last_kwargs = None + + def to(self, device): + self.moved_to = device + return self + + def enable_model_cpu_offload(self) -> None: + self.offloaded = True + + # Explicit signature (not just **kwargs) so generate()'s signature-gated + # guards for negative_prompt / callback_on_step_end actually take effect — + # a **kwargs-only fake would make `"negative_prompt" in signature` always False. + def __call__( + self, + *, + prompt = None, + negative_prompt = None, + callback_on_step_end = None, + guidance_scale = None, + true_cfg_scale = None, + **kwargs, + ): + self.last_kwargs = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "callback_on_step_end": callback_on_step_end, + "guidance_scale": guidance_scale, + "true_cfg_scale": true_cfg_scale, + **kwargs, + } + n = kwargs.get("num_images_per_prompt", 1) + return types.SimpleNamespace(images = [_FakeImage() for _ in range(n)]) + + +class _FakePipeline: + last: dict = {} + + @classmethod + def from_pretrained(cls, base, **kwargs): + _FakePipeline.last = {"base": base, **kwargs} + return _FakePipe() + + +class _FakeTransformer: + last: dict = {} + + @classmethod + def from_single_file(cls, path, **kwargs): + _FakeTransformer.last = {"path": path, **kwargs} + return object() + + +@pytest.fixture +def fake_runtime(monkeypatch): + torch = types.ModuleType("torch") + torch.bfloat16 = _FakeDtype("bfloat16") + torch.float16 = _FakeDtype("float16") + torch.float32 = _FakeDtype("float32") + torch.Generator = _FakeGenerator + torch.cuda = types.SimpleNamespace(is_available = lambda: False) + torch.backends = types.SimpleNamespace(mps = None) + + diffusers = types.ModuleType("diffusers") + diffusers.GGUFQuantizationConfig = lambda compute_dtype = None: ("quant", compute_dtype) + diffusers.ZImagePipeline = _FakePipeline + diffusers.ZImageTransformer2DModel = _FakeTransformer + # Qwen-Image too, so the true_cfg_scale cfg-kwarg path is exercisable. + diffusers.QwenImagePipeline = _FakePipeline + diffusers.QwenImageTransformer2DModel = _FakeTransformer + + monkeypatch.setitem(sys.modules, "torch", torch) + monkeypatch.setitem(sys.modules, "diffusers", diffusers) + # The backend imports clear_gpu_cache by reference; no-op it so unload doesn't + # run real hardware detection against the stubbed torch. + monkeypatch.setattr("core.inference.diffusion.clear_gpu_cache", lambda: None) + _FakePipeline.last = {} + _FakeTransformer.last = {} + yield + + +def test_load_generate_unload_gguf(fake_runtime, tmp_path): + (tmp_path / "model.gguf").write_bytes(b"weights") + backend = DiffusionBackend() + + status = backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "base/repo", + family_override = "z-image", + hf_token = "hf_secret", + ) + assert status["loaded"] is True + assert status["family"] == "z-image" + assert status["base_repo"] == "base/repo" + assert status["device"] == "cpu" + assert status["dtype"] == "float32" + assert status["cpu_offload"] is False + # Transformer built from the local GGUF, pipeline assembled from the base repo. + assert _FakeTransformer.last["path"] == str((tmp_path / "model.gguf").resolve()) + assert _FakeTransformer.last["subfolder"] == "transformer" + # The token reaches the (possibly gated) base config fetch and the pipeline. + assert _FakeTransformer.last["token"] == "hf_secret" + assert _FakePipeline.last["base"] == "base/repo" + assert "transformer" in _FakePipeline.last + + gen = backend.generate( + prompt = "a sloth", negative_prompt = "blurry", width = 512, height = 512, steps = 4, guidance = 3.0 + ) + assert gen["seed"] == 4242 # random seed reported back + assert gen["repo_id"] == str(tmp_path) # echoed so the route can record the model + assert len(gen["images"]) == 1 # PIL images handed to the route for persistence + # z-image guides via guidance_scale (not true_cfg_scale); the signature-gated + # negative_prompt and per-step callback both reach the pipeline call. + call = backend._state.pipe.last_kwargs + assert call["guidance_scale"] == 3.0 and call["true_cfg_scale"] is None + assert call["negative_prompt"] == "blurry" + assert callable(call["callback_on_step_end"]) + + gen2 = backend.generate(prompt = "again", seed = 99) + assert gen2["seed"] == 99 + + # batch_size produces that many images in one call, all sharing the seed. + batch = backend.generate(prompt = "batch", seed = 7, batch_size = 3) + assert len(batch["images"]) == 3 and batch["seed"] == 7 + + assert backend.unload()["loaded"] is False + assert backend.is_loaded is False + + +def test_cpu_offload_ignored_off_cuda(fake_runtime, tmp_path): + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + status = backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + family_override = "z-image", + base_repo = "base/repo", + cpu_offload = True, + ) + # No CUDA in the stub, so offload is not engaged. + assert status["cpu_offload"] is False + + +def test_generate_without_load_raises(fake_runtime): + backend = DiffusionBackend() + with pytest.raises(RuntimeError): + backend.generate(prompt = "x") + + +def test_resolve_base_repo_prefers_caller_then_hf_tag_then_fallback(monkeypatch): + from core.inference import diffusion + from core.inference.diffusion_families import detect_family + + fam = detect_family("unsloth/Qwen-Image-2512-GGUF") + monkeypatch.setattr(diffusion, "_hf_base_model", lambda repo, tok: "Qwen/Qwen-Image-2512") + # Caller's explicit base wins and the HF tag is not consulted. + assert ( + diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", "my/base", fam, None) + == "my/base" + ) + # No caller base: the repo's base_model tag (the variant base) is used. + assert ( + diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", None, fam, None) + == "Qwen/Qwen-Image-2512" + ) + # No caller base and no tag: the family fallback. + monkeypatch.setattr(diffusion, "_hf_base_model", lambda repo, tok: None) + assert ( + diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", " ", fam, None) + == fam.base_repo + ) + + +def test_load_without_gguf_raises(): + backend = DiffusionBackend() + with pytest.raises(ValueError): + backend.load_pipeline("unsloth/Z-Image-Turbo-GGUF") # no gguf_filename + + +def test_load_unknown_family_raises(): + backend = DiffusionBackend() + # User-facing message: names the supported models, no internal-API jargon. + with pytest.raises(ValueError, match = "isn't a supported image-generation model"): + backend.load_pipeline("some/unrecognised-repo", gguf_filename = "x.gguf") + + +# load_progress state machine (no threads / network / real cache) + +from core.inference.diffusion import _LoadingState, _LoadState # noqa: E402 + + +def test_load_progress_idle_and_ready(): + backend = DiffusionBackend() + assert backend.load_progress()["phase"] is None + backend._state = _LoadState(object(), None, "r", "b", "cpu", "float32", False) + assert backend.load_progress()["phase"] == "ready" + + +def test_load_progress_error(): + backend = DiffusionBackend() + backend._loading = _LoadingState(repo_id = "r", base_repo = "b", error = "boom") + p = backend.load_progress() + assert p["phase"] == "error" and p["error"] == "boom" + + +def test_load_progress_downloading_then_finalizing(monkeypatch): + backend = DiffusionBackend() + backend._loading = _LoadingState(repo_id = "r", base_repo = "b", expected_bytes = 1000) + + monkeypatch.setattr(DiffusionBackend, "_cache_bytes", staticmethod(lambda repo: 150)) + p = backend.load_progress() + assert p["phase"] == "downloading" + assert p["bytes_downloaded"] == 300 # summed across repo + base + assert abs(p["fraction"] - 0.3) < 1e-9 + + monkeypatch.setattr(DiffusionBackend, "_cache_bytes", staticmethod(lambda repo: 500)) + assert backend.load_progress()["phase"] == "finalizing" # 1000/1000 + + +def test_base_file_downloaded_excludes_undownloaded(): + # Counted: the pipeline manifest + component subfolders from_pretrained fetches. + assert _base_file_downloaded("model_index.json") + assert _base_file_downloaded("text_encoder/model-00001-of-00003.safetensors") + assert _base_file_downloaded("vae/diffusion_pytorch_model.safetensors") + # Excluded: the GGUF supplies the transformer; docs/assets and top-level files + # are never downloaded, so counting them would peg the bar short of 100%. + assert not _base_file_downloaded( + "transformer/diffusion_pytorch_model-00001-of-00003.safetensors" + ) + assert not _base_file_downloaded("assets/Z-Image-Gallery.pdf") + assert not _base_file_downloaded("README.md") + assert not _base_file_downloaded(".gitattributes") + + +def test_load_progress_fraction_clamped(monkeypatch): + # The cache scan can exceed the estimate (e.g. a second cached quant); the + # reported fraction must still clamp to 1.0 rather than overshoot. + backend = DiffusionBackend() + backend._loading = _LoadingState(repo_id = "r", base_repo = "b", expected_bytes = 1000) + monkeypatch.setattr(DiffusionBackend, "_cache_bytes", staticmethod(lambda repo: 900)) + p = backend.load_progress() # summed 1800 > expected 1000 + assert p["phase"] == "finalizing" + assert p["fraction"] == 1.0 + assert p["bytes_downloaded"] == 1000 # clamped to the estimate + + +def test_estimate_eta(): + from core.inference.diffusion import _estimate_eta + + # No rate yet until a step has elapsed since the first. + assert _estimate_eta(8, 1, first_step_at = 100.0, now = 100.0) is None + assert _estimate_eta(8, 0, first_step_at = 0.0, now = 100.0) is None + # 3 steps in 3s since the first ⇒ 1s/step ⇒ 4 steps left ⇒ ~4s. + assert _estimate_eta(8, 4, first_step_at = 100.0, now = 103.0) == 4.0 + # Last step ⇒ 0 remaining. + assert _estimate_eta(8, 8, first_step_at = 100.0, now = 107.0) == 0.0 + + +def test_generate_qwen_uses_true_cfg_scale(fake_runtime, tmp_path): + (tmp_path / "model.gguf").write_bytes(b"weights") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "Qwen/Qwen-Image", + family_override = "qwen-image", + ) + backend.generate(prompt = "a sloth", guidance = 4.0) + # Qwen-Image's distilled guidance is off; the real CFG must land on true_cfg_scale. + call = backend._state.pipe.last_kwargs + assert call["true_cfg_scale"] == 4.0 and call["guidance_scale"] is None + + +def test_begin_load_rejects_concurrent(monkeypatch): + backend = DiffusionBackend() + # The worker resolves the base + downloads, both over the network; stub them + # so the test is offline. + monkeypatch.setattr("core.inference.diffusion._hf_base_model", lambda *a, **k: None) + monkeypatch.setattr(DiffusionBackend, "_prefetch_files", lambda self, *a, **k: None) + monkeypatch.setattr( + DiffusionBackend, "_estimate_download_bytes", staticmethod(lambda *a, **k: (0, [])) + ) + # Block the spawned worker so the load stays "in progress". + monkeypatch.setattr( + DiffusionBackend, "load_pipeline", lambda self, **k: __import__("time").sleep(0.2) + ) + backend.begin_load("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "z-image-turbo-Q4_K_S.gguf") + with pytest.raises(RuntimeError): + backend.begin_load("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "z-image-turbo-Q4_K_S.gguf") + + +def test_unload_cancels_in_flight_load(fake_runtime): + # An unload (or an arbiter eviction, which calls unload) while a load's worker + # is still resolving/downloading must cancel it: load_pipeline sees the bumped + # token and aborts, so the evicted load never resurrects a pipeline into VRAM. + backend = DiffusionBackend() + fam = detect_family("unsloth/Z-Image-Turbo-GGUF") + token = 7 + backend._load_token = token + with pytest.raises(RuntimeError, match = "cancelled"): + # Simulate the worker reaching load_pipeline after unload bumped the token. + backend._load_token = token + 1 + backend.load_pipeline( + "unsloth/Z-Image-Turbo-GGUF", + gguf_filename = "z-image-turbo-Q4_K_S.gguf", + base_repo = fam.base_repo, + _load_token = token, + ) + + +def test_pick_dtype_bf16_only_on_ampere(fake_runtime, monkeypatch): + # BF16 only on Ampere+ (cc >= 8); pre-Ampere cards must fall back to FP16. + torch = sys.modules["torch"] + backend = DiffusionBackend() + monkeypatch.setattr(torch.cuda, "is_available", lambda: True, raising = False) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (8, 0), raising = False) + assert backend._pick_device_and_dtype() == ("cuda", torch.bfloat16) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (7, 5), raising = False) + assert backend._pick_device_and_dtype() == ("cuda", torch.float16) + + +def test_unload_sets_cancel_event(fake_runtime): + # unload signals an in-flight download (which runs without the lock) to abort. + backend = DiffusionBackend() + assert not backend._cancel_event.is_set() + backend.unload() + assert backend._cancel_event.is_set() + + +def test_prefetch_aborts_when_cancelled(tmp_path): + # A prefetch interrupted by unload (cancel event set) raises rather than + # downloading the whole base, so the load can be preempted mid-download. + backend = DiffusionBackend() + backend._cancel_event.set() + # Local gguf path so the transformer download is skipped; the base loop hits + # the cancel check on its first file (no network). + (tmp_path / "model.gguf").write_bytes(b"x") + with pytest.raises(RuntimeError, match = "Cancelled"): + backend._prefetch_files( + str(tmp_path), + "model.gguf", + "Tongyi-MAI/Z-Image-Turbo", + ["vae/diffusion_pytorch_model.safetensors"], + None, + ) + + +def test_prefetch_downloads_gguf_and_base(monkeypatch, tmp_path): + backend = DiffusionBackend() + calls: list = [] + monkeypatch.setattr( + "utils.hf_xet_fallback.hf_hub_download_with_xet_fallback", + lambda repo, fn, tok, **k: (calls.append((repo, fn)), f"/cache/{fn}")[1], + ) + # Hub repo: the GGUF transformer and each base file are fetched. + backend._prefetch_files( + "unsloth/Z-Image-Turbo-GGUF", + "model.gguf", + "base/repo", + ["vae/x.safetensors", "text_encoder/y.safetensors"], + "hf_tok", + ) + assert ("unsloth/Z-Image-Turbo-GGUF", "model.gguf") in calls + assert ("base/repo", "vae/x.safetensors") in calls + assert ("base/repo", "text_encoder/y.safetensors") in calls + # Local GGUF path: the transformer download is skipped, base still fetched. + calls.clear() + (tmp_path / "model.gguf").write_bytes(b"x") + backend._prefetch_files(str(tmp_path), "model.gguf", "base/repo", ["vae/x.safetensors"], None) + assert all(repo != str(tmp_path) for repo, _ in calls) + assert ("base/repo", "vae/x.safetensors") in calls + + +def test_run_load_does_not_stamp_superseded_progress(fake_runtime, monkeypatch): + # A worker whose load is superseded mid-resolve must not stamp its progress + # (base_repo / expected_bytes) onto the new load's _LoadingState. + backend = DiffusionBackend() + backend._loading = _LoadingState(repo_id = "unsloth/Z-Image-Turbo-GGUF", base_repo = "seed") + backend._load_token = 5 + monkeypatch.setattr("core.inference.diffusion._hf_base_model", lambda *a, **k: None) + + def supersede_then_estimate(*a, **k): + backend._load_token = 6 # a newer begin_load bumped the token mid-resolve + return (99999, []) + + monkeypatch.setattr( + DiffusionBackend, "_estimate_download_bytes", staticmethod(supersede_then_estimate) + ) + monkeypatch.setattr(DiffusionBackend, "_prefetch_files", lambda self, *a, **k: None) + monkeypatch.setattr(DiffusionBackend, "load_pipeline", lambda self, **k: None) + + backend._run_load( + repo_id = "unsloth/Z-Image-Turbo-GGUF", gguf_filename = "m.gguf", base_repo = None, _load_token = 5 + ) + assert backend._loading.expected_bytes == 0 + assert backend._loading.base_repo == "seed" diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py new file mode 100644 index 0000000000..878de0d477 --- /dev/null +++ b/studio/backend/tests/test_diffusion_routes.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""FastAPI round-trip tests for the diffusion image routes. + +The diffusion backend is replaced with a lightweight fake, so these exercise the +route wiring, validation (422), error mapping, and response shapes without torch, +diffusers, weights, or a GPU. +""" + +from __future__ import annotations + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import core.inference.diffusion as diffusion_module +import core.inference.gpu_arbiter as gpu_arbiter +import core.inference.image_gallery as gallery_module +from auth.authentication import get_current_subject +from routes.inference import studio_router + + +class _FakeBackend: + def __init__(self) -> None: + self.loaded = False + + @property + def is_loaded(self) -> bool: + return self.loaded + + def validate_load(self, model_path, **kwargs): + # The route runs this cheap pre-flight before taking the GPU. The fake + # accepts anything; tests that exercise rejection patch it to raise. + return None + + def begin_load(self, model_path, **kwargs): + # The real backend loads on a thread; the fake completes instantly. + self.loaded = True + return { + "loaded": True, + "repo_id": model_path, + "family": "z-image", + "base_repo": kwargs.get("base_repo") or "base/repo", + "device": "cpu", + "dtype": "float32", + "cpu_offload": False, + } + + def load_progress(self): + return { + "phase": "ready" if self.loaded else None, + "bytes_downloaded": 0, + "bytes_total": 0, + "fraction": 1.0 if self.loaded else 0.0, + "error": None, + } + + def generate( + self, + *, + seed = None, + batch_size = 1, + **kwargs, + ): + if not self.loaded: + raise RuntimeError("No diffusion model is loaded.") + # The real backend returns the PIL images; the route persists them. The + # fake returns sentinels since image_gallery is stubbed in the fixture. + return { + "images": [object() for _ in range(batch_size)], + "seed": seed if seed is not None else 4242, + "repo_id": "x/z-image", + } + + def unload(self): + self.loaded = False + return _unloaded_status() + + def status(self): + return {**_unloaded_status(), "loaded": self.loaded} + + +def _unloaded_status(): + return { + "loaded": False, + "repo_id": None, + "family": None, + "base_repo": None, + "device": None, + "dtype": None, + "cpu_offload": False, + } + + +@pytest.fixture +def client(monkeypatch, tmp_path): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + # Isolate from the real GPU arbiter: reset ownership and stub the evictors so + # the load route's acquire_for() never touches live backend singletons. + monkeypatch.setattr(gpu_arbiter, "_owner", None) + monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: None) + monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.DIFFUSION, lambda: None) + + # In-memory gallery backed by tmp files, so routes exercise persistence wiring + # without PIL/real disk under studio_root. + store: dict[str, dict] = {} + + def _save(image, meta): + image_id = f"img{len(store)}" + (tmp_path / f"{image_id}.png").write_bytes(b"PNG") + record = {**meta, "id": image_id, "url": f"/api/inference/images/gallery/{image_id}/file"} + store[image_id] = record + return record + + def _clear(): + n = len(store) + store.clear() + return n + + monkeypatch.setattr(gallery_module, "save", _save) + monkeypatch.setattr(gallery_module, "image_b64", lambda i: "QUJD" if i in store else None) + + def _list_images(limit = None, offset = 0): + ordered = sorted(store.values(), key = lambda r: r.get("created_at", 0.0), reverse = True) + return ordered[offset:] if limit is None else ordered[offset : offset + limit] + + monkeypatch.setattr(gallery_module, "list_images", _list_images) + monkeypatch.setattr( + gallery_module, + "image_path", + lambda i: (tmp_path / f"{i}.png") if i in store else None, + ) + monkeypatch.setattr(gallery_module, "delete", lambda i: store.pop(i, None) is not None) + monkeypatch.setattr(gallery_module, "clear", _clear) + + app = FastAPI() + app.include_router(studio_router, prefix = "/api/inference") + app.dependency_overrides[get_current_subject] = lambda: "test-user" + return TestClient(app) + + +def test_load_generate_status_unload_roundtrip(client): + loaded = client.post( + "/api/inference/images/load", + json = { + "model_path": "unsloth/Z-Image-Turbo-GGUF", + "gguf_filename": "z-image-turbo-Q4_K_S.gguf", + "base_repo": "base/repo", + }, + ) + assert loaded.status_code == 200 + body = loaded.json() + assert body["loaded"] is True and body["family"] == "z-image" + + assert client.get("/api/inference/images/status").json()["loaded"] is True + + gen = client.post("/api/inference/images/generate", json = {"prompt": "a sloth", "seed": 7}) + assert gen.status_code == 200 + # One persisted record carrying the full recipe back. + images = gen.json()["images"] + assert len(images) == 1 + img = images[0] + assert img["seed"] == 7 and img["prompt"] == "a sloth" and img["id"] + + # The image is now listable, fetchable, and deletable. + listed = client.get("/api/inference/images/gallery").json()["images"] + assert [i["id"] for i in listed] == [img["id"]] + assert client.get(img["url"]).status_code == 200 + assert client.delete(img["url"].removesuffix("/file")).status_code == 200 + assert client.get("/api/inference/images/gallery").json()["images"] == [] + + unloaded = client.post("/api/inference/images/unload") + assert unloaded.status_code == 200 and unloaded.json()["loaded"] is False + assert client.get("/api/inference/images/status").json()["loaded"] is False + + +def test_generate_batch_size_persists_each_image(client): + client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) + resp = client.post( + "/api/inference/images/generate", + json = {"prompt": "p", "batch_size": 3, "seed": 5}, + ) + assert resp.status_code == 200 + images = resp.json()["images"] + assert len(images) == 3 + assert all(i["seed"] == 5 for i in images) # the batch shares one seed + assert len({i["id"] for i in images}) == 3 # but each is a distinct record + assert len(client.get("/api/inference/images/gallery").json()["images"]) == 3 + + +def test_gallery_pagination(client): + client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) + client.post("/api/inference/images/generate", json = {"prompt": "p", "batch_size": 5, "seed": 1}) + page1 = client.get("/api/inference/images/gallery?limit=2&offset=0").json() + assert len(page1["images"]) == 2 and page1["has_more"] is True + last = client.get("/api/inference/images/gallery?limit=2&offset=4").json() + assert len(last["images"]) == 1 and last["has_more"] is False + + +def test_generate_rejects_non_multiple_of_16(client): + client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) + # Odd, and a multiple of 8 that isn't a multiple of 16: both rejected, since + # Z-Image requires dimensions divisible by 16. + for bad in (1001, 1000): + resp = client.post("/api/inference/images/generate", json = {"prompt": "p", "width": bad}) + assert resp.status_code == 422, bad + # A multiple of 16 is accepted. + ok = client.post("/api/inference/images/generate", json = {"prompt": "p", "width": 1024}) + assert ok.status_code == 200 + + +def test_load_requires_gguf_filename(client): + # gguf_filename is now mandatory — a load without it is a 422. + resp = client.post("/api/inference/images/load", json = {"model_path": "x/z-image"}) + assert resp.status_code == 422 + + +def test_generate_without_load_returns_409(client): + resp = client.post("/api/inference/images/generate", json = {"prompt": "p"}) + assert resp.status_code == 409 + + +def test_generate_pipeline_error_returns_sanitized_500(client, monkeypatch): + # A loaded model that fails mid-pipeline (CUDA OOM, a RuntimeError) is a server + # failure: 500 with a generic message, not a 409 echoing the raw exception. + backend = diffusion_module.get_diffusion_backend() + backend.loaded = True + + def _oom(**kwargs): + raise RuntimeError("CUDA out of memory. Tried to allocate 20.00 GiB (24.00 GiB total)") + + monkeypatch.setattr(backend, "generate", _oom) + resp = client.post("/api/inference/images/generate", json = {"prompt": "p"}) + assert resp.status_code == 500 + assert resp.json()["detail"] == "Image generation failed." + assert "CUDA" not in resp.json()["detail"] + + +def test_load_unknown_family_returns_400(client, monkeypatch): + def _raise(*a, **k): + raise ValueError("'x/y' isn't a supported image-generation model. Supported: Z-Image.") + + backend = _FakeBackend() + # Validation runs in the pre-flight (before the GPU is taken), so that is + # where an unsupported model is rejected now. + backend.validate_load = _raise + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + resp = client.post( + "/api/inference/images/load", json = {"model_path": "x/y", "gguf_filename": "q.gguf"} + ) + assert resp.status_code == 400 + assert "isn't a supported image-generation model" in resp.json()["detail"] + + +def test_load_validation_failure_does_not_evict_chat(client, monkeypatch): + # A rejected image-model pick must not tear down the user's loaded chat model: + # validation runs before acquire_for, so chat keeps the GPU on a 400. + monkeypatch.setattr(gpu_arbiter, "_owner", gpu_arbiter.CHAT) + evicted = [] + monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: evicted.append(True)) + + backend = _FakeBackend() + + def _raise(*a, **k): + raise ValueError("'x/y' isn't a supported image-generation model.") + + backend.validate_load = _raise + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + resp = client.post( + "/api/inference/images/load", json = {"model_path": "x/y", "gguf_filename": "q.gguf"} + ) + assert resp.status_code == 400 + assert evicted == [] # chat backend was never evicted + assert gpu_arbiter.current_owner() == gpu_arbiter.CHAT + + +def test_load_refused_during_training_does_not_evict_chat(client, monkeypatch): + # An image load while training is active is refused (409) before the GPU is + # taken, so the training run and the loaded chat model are both untouched. + import core.training as core_training + + monkeypatch.setattr(gpu_arbiter, "_owner", gpu_arbiter.CHAT) + evicted = [] + monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: evicted.append(True)) + + class _Training: + def is_training_active(self): + return True + + monkeypatch.setattr(core_training, "get_training_backend", lambda: _Training()) + + resp = client.post( + "/api/inference/images/load", + json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}, + ) + assert resp.status_code == 409 + assert "training" in resp.json()["detail"].lower() + assert evicted == [] # chat backend was never evicted + assert gpu_arbiter.current_owner() == gpu_arbiter.CHAT + + +def test_load_progress_route(client): + # Before load: idle. + idle = client.get("/api/inference/images/load-progress") + assert idle.status_code == 200 and idle.json()["phase"] is None + # After load: the fake reports ready. + client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) + ready = client.get("/api/inference/images/load-progress") + assert ready.json()["phase"] == "ready" + + +def test_routes_require_auth(): + # No dependency override: the auth dependency must reject the request. + app = FastAPI() + app.include_router(studio_router, prefix = "/api/inference") + unauth = TestClient(app) + assert unauth.get("/api/inference/images/status").status_code in (401, 403) diff --git a/studio/backend/tests/test_gpu_arbiter.py b/studio/backend/tests/test_gpu_arbiter.py new file mode 100644 index 0000000000..0487830def --- /dev/null +++ b/studio/backend/tests/test_gpu_arbiter.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the single-GPU arbiter. + +The real evictors (which tear down live backends) are replaced with recorders, so +these verify only the ownership/eviction sequencing — no torch, GPU, or subprocess. +""" + +from __future__ import annotations + +import pytest + +import core.inference.gpu_arbiter as arb + + +@pytest.fixture +def calls(monkeypatch): + recorded: list[str] = [] + monkeypatch.setattr(arb, "_owner", None) + monkeypatch.setitem(arb._EVICTORS, arb.CHAT, lambda: recorded.append("evict-chat")) + monkeypatch.setitem(arb._EVICTORS, arb.DIFFUSION, lambda: recorded.append("evict-diffusion")) + return recorded + + +def test_first_acquire_evicts_nothing(calls): + arb.acquire_for(arb.CHAT) + assert calls == [] + assert arb.current_owner() == arb.CHAT + + +def test_diffusion_load_evicts_chat(calls): + arb.acquire_for(arb.CHAT) + arb.acquire_for(arb.DIFFUSION) + assert calls == ["evict-chat"] + assert arb.current_owner() == arb.DIFFUSION + + +def test_chat_load_evicts_diffusion(calls): + arb.acquire_for(arb.DIFFUSION) + arb.acquire_for(arb.CHAT) + assert calls == ["evict-diffusion"] + assert arb.current_owner() == arb.CHAT + + +def test_reacquiring_same_owner_does_not_evict(calls): + arb.acquire_for(arb.CHAT) + arb.acquire_for(arb.CHAT) + assert calls == [] + assert arb.current_owner() == arb.CHAT + + +def test_release_clears_owner(calls): + arb.acquire_for(arb.DIFFUSION) + arb.release(arb.DIFFUSION) + assert arb.current_owner() is None + # A subsequent chat load then has nothing to evict. + arb.acquire_for(arb.CHAT) + assert calls == [] + + +def test_release_by_non_owner_is_noop(calls): + arb.acquire_for(arb.CHAT) + arb.release(arb.DIFFUSION) + assert arb.current_owner() == arb.CHAT + + +def test_unknown_owner_raises(calls): + with pytest.raises(ValueError): + arb.acquire_for("gpu") + + +def test_evict_chat_unloads_a_still_loading_chat_backend(monkeypatch): + # A chat model still starting up is is_active (process exists) but not yet + # is_loaded (healthy). Eviction must still unload it, or the load would keep + # allocating VRAM after the GPU was handed to diffusion. + import core.inference as core_inference + import routes.inference as routes_inference + + unloaded: list[bool] = [] + + class _FakeLlama: + is_active = True + is_loaded = False # still loading: skipped if eviction gates on is_loaded + + def unload_model(self): + unloaded.append(True) + + def _wait_for_vram_settle(self, *, since_kill): + pass + + class _FakeOrchestrator: + active_model_name = None + + def unload_model(self, name): + pass + + def _shutdown_subprocess(self, timeout = 5.0): + pass + + monkeypatch.setattr(routes_inference, "get_llama_cpp_backend", lambda: _FakeLlama()) + monkeypatch.setattr(core_inference, "get_inference_backend", lambda: _FakeOrchestrator()) + + arb._evict_chat() + + assert unloaded == [True] # still-loading chat backend was unloaded, not skipped diff --git a/studio/backend/tests/test_image_gallery.py b/studio/backend/tests/test_image_gallery.py new file mode 100644 index 0000000000..883d8b8793 --- /dev/null +++ b/studio/backend/tests/test_image_gallery.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the disk-backed image gallery: PNG-embedded recipe round-trips, +listing order, safe id handling, and delete/clear.""" + +from __future__ import annotations + +import base64 +import io +import os + +import pytest + +import core.inference.image_gallery as gallery + +PIL = pytest.importorskip("PIL") +from PIL import Image # noqa: E402 + + +@pytest.fixture(autouse = True) +def _tmp_gallery(monkeypatch, tmp_path): + # Point the gallery at a throwaway root instead of ~/.unsloth/studio. + monkeypatch.setattr(gallery, "studio_root", lambda: tmp_path) + + +def _img(color = (10, 20, 30)): + return Image.new("RGB", (16, 16), color) + + +def _meta(**over): + base = { + "prompt": "a sloth", + "negative_prompt": None, + "width": 1024, + "height": 1024, + "steps": 9, + "guidance": 0.0, + "seed": 7, + "model": "unsloth/Z-Image-Turbo-GGUF", + "created_at": 100.0, + } + base.update(over) + return base + + +def test_save_embeds_recipe_and_round_trips(): + record = gallery.save(_img(), _meta()) + assert record["id"] and record["url"].endswith(f"{record['id']}/file") + + # The recipe is embedded in the PNG itself (portable), not just in a sidecar. + raw = base64.b64decode(gallery.image_b64(record["id"])) + with Image.open(io.BytesIO(raw)) as im: + assert im.text["unsloth"] + assert "Negative prompt" not in im.text["parameters"] # none given + assert "Steps: 9" in im.text["parameters"] + + listed = gallery.list_images() + assert len(listed) == 1 + assert listed[0]["prompt"] == "a sloth" and listed[0]["seed"] == 7 + + +def _save_with_mtime(prompt: str, t: float) -> dict: + record = gallery.save(_img(), _meta(prompt = prompt, created_at = t)) + # Listing orders by mtime; set it explicitly so a tight test loop can't tie it. + os.utime(gallery.gallery_dir() / f"{record['id']}.png", (t, t)) + return record + + +def test_list_is_newest_first(): + old = _save_with_mtime("old", 100.0) + new = _save_with_mtime("new", 200.0) + assert [r["id"] for r in gallery.list_images()] == [new["id"], old["id"]] + + +def test_list_paginates_with_limit_offset(): + # 5 images, newest (t=4) first. + for i in range(5): + _save_with_mtime(f"p{i}", float(i)) + page1 = gallery.list_images(limit = 2, offset = 0) + page2 = gallery.list_images(limit = 2, offset = 2) + assert [r["prompt"] for r in page1] == ["p4", "p3"] + assert [r["prompt"] for r in page2] == ["p2", "p1"] + # limit=None still returns everything from the offset. + assert len(gallery.list_images()) == 5 + assert len(gallery.list_images(offset = 4)) == 1 + + +def test_negative_prompt_recorded_in_parameters(): + record = gallery.save(_img(), _meta(negative_prompt = "blurry")) + raw = base64.b64decode(gallery.image_b64(record["id"])) + with Image.open(io.BytesIO(raw)) as im: + assert "Negative prompt: blurry" in im.text["parameters"] + + +def test_delete_and_clear(): + a = gallery.save(_img(), _meta(prompt = "a")) + gallery.save(_img(), _meta(prompt = "b")) + assert gallery.delete(a["id"]) is True + assert gallery.delete(a["id"]) is False # already gone + assert len(gallery.list_images()) == 1 + assert gallery.clear() == 1 + assert gallery.list_images() == [] + + +def test_image_path_rejects_unsafe_ids(): + # Traversal / bad chars never resolve to a path. + assert gallery.image_path("../../etc/passwd") is None + assert gallery.image_path("a/b") is None + assert gallery.image_path("missing") is None + + +def test_list_skips_foreign_pngs(tmp_path): + # A PNG without our recipe chunk (user dropped a file) is ignored. + foreign = gallery.gallery_dir() / "foreign.png" + _img().save(foreign, format = "PNG") + gallery.save(_img(), _meta(prompt = "ours")) + listed = gallery.list_images() + assert [r["prompt"] for r in listed] == ["ours"] + + +def test_list_skips_recipe_missing_required_fields(tmp_path): + # A PNG carrying our chunk but an incomplete/older-schema recipe (no seed etc.) + # must be skipped, not crash the whole listing when the route builds GalleryImage. + import json + + from PIL.PngImagePlugin import PngInfo + + info = PngInfo() + info.add_text("unsloth", json.dumps({"prompt": "partial"})) # missing width/seed/... + _img().save(gallery.gallery_dir() / "partial.png", format = "PNG", pnginfo = info) + gallery.save(_img(), _meta(prompt = "ours")) + listed = gallery.list_images() + assert [r["prompt"] for r in listed] == ["ours"] diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index 5c18e637e2..9272dfab2b 100644 --- a/studio/frontend/src/app/router.tsx +++ b/studio/frontend/src/app/router.tsx @@ -10,6 +10,7 @@ import { Route as dataRecipesRoute } from "./routes/data-recipes"; import { Route as dataRecipeRoute } from "./routes/data-recipes.$recipeId"; import { Route as chatRoute } from "./routes/chat"; import { Route as exportRoute } from "./routes/export"; +import { Route as imagesRoute } from "./routes/images"; import { Route as gridTestRoute } from "./routes/grid-test"; import { Route as indexRoute } from "./routes/index"; import { Route as loginRoute } from "./routes/login"; @@ -32,6 +33,7 @@ const routeTree = rootRoute.addChildren([ chatRoute, projectsRoute, exportRoute, + imagesRoute, dataRecipesRoute, dataRecipeRoute, ]); diff --git a/studio/frontend/src/app/routes/images.tsx b/studio/frontend/src/app/routes/images.tsx new file mode 100644 index 0000000000..1761612140 --- /dev/null +++ b/studio/frontend/src/app/routes/images.tsx @@ -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, +}); diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 06f2701a16..0905940609 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -74,6 +74,7 @@ import { PinOffIcon, PlusSignIcon, PowerIcon, + PaintBrush02Icon, PencilEdit02Icon, LayoutAlignLeftIcon, Settings02Icon, @@ -1213,6 +1214,18 @@ export function AppSidebar() { closeMobileIfOpen(); }} /> + { + if (chatOnly) return; + navigate({ to: "/images" }); + closeMobileIfOpen(); + }} + /> {/* Train has a labelled section when expanded; plain icon here only when collapsed. */} s.isChatOnly()); @@ -460,6 +465,7 @@ function ModelSelectorContent({ deleteDisabled={deleteDisabled} section={effectiveHubSection} onEject={hasSelection && onEject ? onEject : undefined} + task={task} sectionToggle={ ); diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 386f233c7b..a33f4b57f0 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -40,6 +40,7 @@ import { useHubInfiniteScroll } from "@/features/hub/hooks/use-hub-infinite-scro import { type HfModelResult, type HfSortKey, + type HfTaskFilter, useHubModelSearch, } from "@/features/hub/hooks/use-hub-model-search"; import { useOnlineStatus } from "@/features/hub/hooks/use-online-status"; @@ -696,7 +697,7 @@ function GgufVariantExpander({ ); const handleVariantClick = useCallback( - (quant: string, downloaded?: boolean, sizeBytes?: number) => { + (quant: string, filename: string, downloaded?: boolean, sizeBytes?: number) => { // Only seed the staged context for picks whose weights are already on // disk. The staging effect short-circuits on a known contextLength // (pendingHasContext) before starting the download, so attaching it to an @@ -707,6 +708,7 @@ function GgufVariantExpander({ source: sourceOverride ?? (isLocalPath ? "local" : "hub"), isLora: false, ggufVariant: quant, + ggufFilename: filename, isDownloaded: isLocalPath ? true : downloaded, expectedBytes: sizeBytes, contextLength: isAvailable ? nativeContext : undefined, @@ -878,7 +880,7 @@ function GgufVariantExpander({ type="button" {...variantList.getOptionProps(variantOptionKey, false)} onClick={() => - handleVariantClick(v.quant, v.downloaded, v.size_bytes) + handleVariantClick(v.quant, v.filename, v.downloaded, v.size_bytes) } className={cn( "flex min-w-0 flex-1 items-center justify-between gap-2 rounded-full px-2 py-1 text-left text-sm transition-colors hover:bg-[#ececec] focus-visible:bg-[#ececec] focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring/45 dark:hover:bg-[var(--sidebar-accent)] dark:focus-visible:bg-[var(--sidebar-accent)]", @@ -966,6 +968,51 @@ function isGgufRepo(id: string, hintedIsGguf?: boolean): boolean { return Boolean(hintedIsGguf) || hasGgufSuffix(id); } +// True when a repo's inferred task is within the picker's task filter (or no +// filter is set). Unknown task (null) only passes when there's no filter. +function taskMatchesFilter(repoTask: string | null | undefined, filter: HfTaskFilter): boolean { + if (!filter) return true; + const wanted = Array.isArray(filter) ? filter : [filter]; + return repoTask != null && (wanted as readonly string[]).includes(repoTask); +} + +// Image-generation pipeline tasks: handled by the Images page, never loadable as +// chat models. The backend reports "text-to-image" for diffusion-arch GGUFs. The +// Images page reuses this as its picker's `task` filter, so it lives here. +export const IMAGE_GEN_TASKS = [ + "text-to-image", + "image-to-image", + "image-text-to-image", +] as const; + +// Editing/inpaint checkpoints are tagged image-to-image but need an input image, +// which the text-to-image backend rejects (mirrors its _EDIT_KEYWORDS). Hidden by +// id so they don't show in the Images picker only to 400 on load. Keeping the +// image-to-image task itself is required: some supported models (FLUX.2-klein) +// carry that tag too. +const IMAGE_EDIT_KEYWORDS = ["edit", "kontext", "inpaint", "inpainting"] as const; +function isImageEditModel(repoId: string | null | undefined): boolean { + if (!repoId) return false; + // Whole-segment match (not substring) so a normal model like "...-edition" + // isn't hidden; mirrors the backend detect_family segment check. Split on + // both path separators so a Windows local path is segmented too. + const segments = new Set(repoId.toLowerCase().split(/[-_./\\]+/)); + return IMAGE_EDIT_KEYWORDS.some((kw) => segments.has(kw)); +} + +// Gate an on-device model by the picker's task scope. With a filter (the Images +// page) keep only matching, non-editing tasks; with no filter (chat) drop +// image-generation models so a downloaded diffusion GGUF doesn't show up as a +// loadable chat model. +function passesTaskGate( + repoTask: string | null | undefined, + repoId: string | null | undefined, + filter: HfTaskFilter, +): boolean { + if (filter) return taskMatchesFilter(repoTask, filter) && !isImageEditModel(repoId); + return !(repoTask != null && (IMAGE_GEN_TASKS as readonly string[]).includes(repoTask)); +} + // Module-level caches so re-mounting the popover shows results instantly let _cachedGgufCache: CachedGgufRepo[] = []; let _cachedModelsCache: CachedModelRepo[] = []; @@ -1147,6 +1194,7 @@ export function HubModelPicker({ section = "downloaded", sectionToggle, onEject, + task, }: { models: ModelOption[]; /** Fine-tuned models, shown as a section in the On Device view. */ @@ -1166,6 +1214,9 @@ export function HubModelPicker({ sectionToggle?: ReactNode; /** Eject the loaded model. Rendered as the last list row when set. */ onEject?: () => void; + /** Restrict Hub results to a pipeline task (e.g. text-to-image for the + * Images page). Undefined = all tasks (chat default). */ + task?: HfTaskFilter; }) { const gpu = useGpuInfo(); // Last-loaded timestamps power the "Recent" sort (vs "Downloaded" = file date). @@ -1196,6 +1247,7 @@ export function HubModelPicker({ hasMore, } = useHubModelSearch(debouncedQuery, { ownerScope: "unsloth", + task, sortBy: recommendedSortBy, sortDirection: "desc", pinUnslothFirst: true, @@ -1208,6 +1260,7 @@ export function HubModelPicker({ }); const recommendedSearch = useHubModelSearch("", { ownerScope: "unsloth", + task, sortBy: recommendedSortBy, sortDirection: "desc", pinUnslothFirst: true, @@ -1524,18 +1577,28 @@ export function HubModelPicker({ const isMac = deviceType === "mac"; // Drop models Studio can't run for chat (diffusion / image / video / etc.) - // using the Hub's classifier on the tags the listing already carries. + // using the Hub's classifier on the tags the listing already carries. When the + // picker is scoped to a task (e.g. the Images page asks for text-to-image), + // models matching that task are exactly what we want — keep them even though + // the chat classifier marks image tasks "unsupported". const isChatSupported = useCallback( - (r: HfModelResult) => - classifyUnslothSupport({ - modelId: r.id, - pipelineTag: r.pipelineTag, - tags: r.tags, - libraryName: r.libraryName, - quantMethod: r.quantMethod, - deviceType, - }).status !== "unsupported", - [deviceType], + (r: HfModelResult) => { + // Image tab (task set): only task-matching, non-editing results. Anything + // else (e.g. a chat GGUF surfaced by a typed query) is dropped rather than + // falling through to the chat classifier and appearing as loadable. + if (task) return taskMatchesFilter(r.pipelineTag, task) && !isImageEditModel(r.id); + return ( + classifyUnslothSupport({ + modelId: r.id, + pipelineTag: r.pipelineTag, + tags: r.tags, + libraryName: r.libraryName, + quantMethod: r.quantMethod, + deviceType, + }).status !== "unsupported" + ); + }, + [deviceType, task], ); const recommendedIds = useMemo(() => { @@ -1576,6 +1639,10 @@ export function HubModelPicker({ .filter((r) => !isMobileVariant(r.id)); // Drop models Studio can't run for chat (diffusion / image / video / etc.). rows = rows.filter(isChatSupported); + // Images (task set) loads single-file GGUF only, so never surface non-GGUF + // rows regardless of the format dropdown (mirrors hfIds and the empty + // Recommended view); selecting a non-GGUF row is a silent no-op. + if (task) rows = rows.filter((r) => r.isGguf); // With no explicit format, show the device-recommended formats (GGUF, plus // MLX on Mac). When the user picks a format, honor it instead so Safetensors // is not dropped by the recommendation default. @@ -1615,6 +1682,7 @@ export function HubModelPicker({ isMac, gpu, isChatSupported, + task, ]); // Per-row meta + VRAM badge from the recommended listing's own metadata. @@ -1698,14 +1766,31 @@ export function HubModelPicker({ return map; }, [results, recommendedSearch.results]); - // Ordered by the On Device dropdown (recent/download date/size/name). + // Ordered by the On Device dropdown (recent/download date/size/name). The task + // gate keeps the Images picker to diffusion GGUFs and, conversely, hides those + // diffusion GGUFs from the chat picker (where they aren't loadable models). const sortedCachedGguf = useMemo( - () => sortCachedRepos(cachedGguf, downloadedSort, loadTimes), - [cachedGguf, downloadedSort, loadTimes], + () => + sortCachedRepos( + cachedGguf.filter((c) => passesTaskGate(c.task, c.repo_id, task)), + downloadedSort, + loadTimes, + ), + [cachedGguf, downloadedSort, loadTimes, task], ); + // Non-GGUF (safetensors) cached repos aren't single-file diffusion GGUFs, so + // hide them entirely when a task filter is active (the Images picker). In chat, + // drop cached diffusers pipeline repos (a Z-Image / FLUX base) the same way. const sortedCachedModels = useMemo( - () => sortCachedRepos(cachedModels, downloadedSort, loadTimes), - [cachedModels, downloadedSort, loadTimes], + () => + task + ? [] + : sortCachedRepos( + cachedModels.filter((c) => passesTaskGate(c.task, c.repo_id, task)), + downloadedSort, + loadTimes, + ), + [cachedModels, downloadedSort, loadTimes, task], ); // Each local section's search is scoped to its own models (matched by name). const localQuery = normalizeForSearch(debouncedQuery.trim()); @@ -1714,18 +1799,23 @@ export function HubModelPicker({ normalizeForSearch( `${m.model_id ?? ""} ${m.display_name} ${m.id}`, ).includes(localQuery); + // The Images page wants diffusion GGUFs only; chat wants everything but those. + // passesTaskGate handles both directions off the GGUF architecture the backend + // reports as each model's task. const sortedLmStudio = useMemo( () => sortLocalModels( lmStudioModels.filter( (m) => - localModelMatchesFormat(m, formatFilter) && matchesLocalQuery(m), + passesTaskGate(m.task, m.model_id ?? m.id, task) && + localModelMatchesFormat(m, formatFilter) && + matchesLocalQuery(m), ), downloadedSort, loadTimes, ), // eslint-disable-next-line react-hooks/exhaustive-deps - [lmStudioModels, downloadedSort, formatFilter, loadTimes, localQuery], + [lmStudioModels, downloadedSort, formatFilter, loadTimes, localQuery, task], ); // Local ./models entries. Chat-only Studio runs GGUF (any host) and MLX (Mac // only), so raw checkpoints there are hidden (mirrors the cached non-GGUF @@ -1735,6 +1825,7 @@ export function HubModelPicker({ sortLocalModels( localDirModels.filter( (m) => + passesTaskGate(m.task, m.model_id ?? m.id, task) && (!chatOnly || localModelIsGguf(m) || (isMac && localModelIsMlx(m))) && @@ -1753,6 +1844,7 @@ export function HubModelPicker({ loadTimes, localQuery, chatOnly, + task, ], ); const sortedCustomFolderModels = useMemo( @@ -1760,18 +1852,22 @@ export function HubModelPicker({ sortLocalModels( customFolderModels.filter( (m) => - localModelMatchesFormat(m, formatFilter) && matchesLocalQuery(m), + passesTaskGate(m.task, m.model_id ?? m.id, task) && + localModelMatchesFormat(m, formatFilter) && + matchesLocalQuery(m), ), customSort, loadTimes, ), // eslint-disable-next-line react-hooks/exhaustive-deps - [customFolderModels, customSort, formatFilter, loadTimes, localQuery], + [customFolderModels, customSort, formatFilter, loadTimes, localQuery, task], ); // Fine-tuned models for the On Device "Fine-tuned" section: flat, query- - // filtered, newest first. + // filtered, newest first. Hidden under a task filter (no fine-tuning of e.g. + // image models). const fineTunedRows = useMemo(() => { + if (task) return []; const needle = normalizeForSearch(debouncedQuery.trim()); return loraModels .filter((m) => { @@ -1787,7 +1883,7 @@ export function HubModelPicker({ if (aTime !== bTime) return bTime - aTime; return a.name.localeCompare(b.name); }); - }, [loraModels, debouncedQuery]); + }, [loraModels, debouncedQuery, task]); // While searching, filter Downloaded by the query instead of hiding it, so a // downloaded model the user is searching for stays visible. @@ -1883,10 +1979,14 @@ export function HubModelPicker({ .filter((id) => !isHiddenModelId(id)) .filter((id) => id.toLowerCase().startsWith("unsloth/")) .filter((id) => !recommendedSet.has(id)) - // Chat-only keeps runnable formats: GGUF anywhere, plus MLX/safetensors - // on Mac (matches the empty Recommended view so search stays consistent). - .filter( - (id) => !chatOnly || isRecommendableFormat(id, isKnownGgufRepo(id), isMac), + // Images (task set) loads single-file GGUF only, so don't surface non-GGUF + // text-to-image rows the page can't load (mirrors the Recommended view). + // Otherwise chat-only keeps runnable formats: GGUF anywhere, plus MLX/ + // safetensors on Mac. + .filter((id) => + task + ? isKnownGgufRepo(id) + : !chatOnly || isRecommendableFormat(id, isKnownGgufRepo(id), isMac), ) .filter((id) => !/-FP8[-.]|FP8-Dynamic/i.test(id)) .filter((id) => @@ -1902,6 +2002,7 @@ export function HubModelPicker({ isChatSupported, formatFilter, isMac, + task, ]); const hubOptionKeys = useMemo(() => { @@ -2567,6 +2668,7 @@ export function HubModelPicker({ ) : null} + {!task && ( + + +
+

Generation settings

+

{formatTimestamp(image.created_at)}

+
+
+ + {image.negative_prompt ? ( + + ) : null} + {image.model ? : null} + + + + +
+
+ +
+ + + ); +} + +type Busy = "loading" | "unloading" | "generating" | null; + +export function ImagesPage() { + const [quant, setQuant] = useState(galleryCache.quant); + const [prompt, setPrompt] = useState( + "a tiny ginger sloth coding in a sunlit treehouse, photorealistic", + ); + const [negativePrompt, setNegativePrompt] = useState(""); + // width/height are the source of truth; `aspect` locks their proportion + // ("custom" = free) and `portrait` tracks orientation, so Flip keeps the lock. + const [width, setWidth] = useState(1024); + const [height, setHeight] = useState(1024); + const [aspect, setAspect] = useState("1:1"); + const [portrait, setPortrait] = useState(false); + // Z-Image-Turbo official defaults: 9 steps (= 8 DiT forwards), guidance 0 + // (distilled CFG-free; a negative prompt is ignored at this guidance). + const [steps, setSteps] = useState(DEFAULT_GEN.steps); + const [guidance, setGuidance] = useState(DEFAULT_GEN.guidance); + const [seed, setSeed] = useState(""); + // Batch size = images per forward pass (VRAM-heavy); count = sequential loops. + const [batchSize, setBatchSize] = useState(1); + const [count, setCount] = useState(1); + + const [busy, setBusy] = useState(null); + // {done, total} while a multi-run generation is in flight (for the button). + // Number of runs finished in the current multi-run generation (null = idle). + // The total is just `count`, so it isn't stored separately. + const [genDone, setGenDone] = useState(null); + // Live per-step progress (step / total + ETA) polled during generation. + const [genStep, setGenStep] = useState(null); + const genPollTimer = useRef | null>(null); + const [status, setStatus] = useState(null); + // Records come from the backend (durable); srcById maps each id to its object + // URL (loaded images) or data URL (the one just generated). + const [images, setImages] = useState(() => galleryCache.images); + const [hasMore, setHasMore] = useState(() => galleryCache.hasMore); + const [selectedId, setSelectedId] = useState(() => galleryCache.selectedId); + const [srcById, setSrcById] = useState>(() => + Object.fromEntries(galleryCache.srcById), + ); + // Guards a "load more" so a fast scroll can't fire several at once. + const loadingMore = useRef(false); + const pollTimer = useRef | null>(null); + // The persistent load toast's id, so each poll updates it in place (chat-style). + const loadToastId = useRef(null); + // Last load-progress signature shown, so a tick that moved nothing skips the toast. + const lastLoadSig = useRef(null); + + const dismissLoadToast = useCallback(() => { + if (loadToastId.current != null) toast.dismiss(loadToastId.current); + loadToastId.current = null; + }, []); + + // Mirror to the module cache so a tab switch re-renders instantly. + useEffect(() => { + galleryCache.images = images; + galleryCache.hasMore = hasMore; + galleryCache.selectedId = selectedId; + galleryCache.quant = quant; + }, [images, hasMore, selectedId, quant]); + + const selected = useMemo( + () => images.find((i) => i.id === selectedId) ?? images[0] ?? null, + [images, selectedId], + ); + const selectedSrc = selected ? srcById[selected.id] : undefined; + + // Fetch (once) the object URL for a record's PNG; cached across remounts. + const ensureSrc = useCallback(async (image: GalleryImage) => { + if (galleryCache.srcById.has(image.id) || galleryCache.inflight.has(image.id)) return; + galleryCache.inflight.add(image.id); + try { + const url = await fetchGalleryObjectUrl(image.url); + galleryCache.srcById.set(image.id, url); + setSrcById((prev) => ({ ...prev, [image.id]: url })); + } catch { + // Leave it without a src; the tile shows a placeholder. + } finally { + galleryCache.inflight.delete(image.id); + } + }, []); + + const loadGallery = useCallback(async () => { + try { + const page = await getGallery(0, PAGE_SIZE); + galleryCache.images = page.images; + galleryCache.hasMore = page.has_more; + setImages(page.images); + setHasMore(page.has_more); + page.images.forEach((image) => void ensureSrc(image)); + } catch { + // Best-effort: a failed gallery load shouldn't block the page. + } + }, [ensureSrc]); + + // Load the next older page. offset = how many we've loaded so far. A newly + // generated image gets a newer mtime, so it sorts to the front on the backend + // too — the offset keeps pointing at the same older images. + const loadMore = useCallback(async () => { + if (loadingMore.current || !galleryCache.hasMore) return; + loadingMore.current = true; + try { + const page = await getGallery(galleryCache.images.length, PAGE_SIZE); + setImages((prev) => { + const seen = new Set(prev.map((i) => i.id)); + const next = [...prev, ...page.images.filter((i) => !seen.has(i.id))]; + galleryCache.images = next; + return next; + }); + galleryCache.hasMore = page.has_more; + setHasMore(page.has_more); + page.images.forEach((image) => void ensureSrc(image)); + } catch { + // transient; the user can scroll again to retry + } finally { + loadingMore.current = false; + } + }, [ensureSrc]); + + useEffect(() => { + void loadGallery(); + }, [loadGallery]); + + const handleDelete = useCallback(async (id: string) => { + try { + await deleteGalleryImage(id); + } catch (err) { + toast.error(err instanceof Error ? err.message : "Failed to delete image"); + return; + } + const url = galleryCache.srcById.get(id); + if (url?.startsWith("blob:")) URL.revokeObjectURL(url); + galleryCache.srcById.delete(id); + setSrcById((prev) => { + const next = { ...prev }; + delete next[id]; + return next; + }); + setImages((prev) => prev.filter((i) => i.id !== id)); + setSelectedId((cur) => (cur === id ? null : cur)); + }, []); + + // Load an image's recipe back into the form inputs. + const restoreSettings = useCallback((image: GalleryImage) => { + setPrompt(image.prompt); + setNegativePrompt(image.negative_prompt ?? ""); + setSteps(image.steps); + setGuidance(image.guidance); + setSeed(String(image.seed)); + setWidth(image.width); + setHeight(image.height); + const m = matchAspect(image.width, image.height); + setAspect(m.key); + setPortrait(m.portrait); + toast.success("Settings restored to inputs"); + }, []); + + // A locked ratio keeps the paired dimension in step while dragging; "custom" + // frees both; Flip swaps W/H. ratioHW is the h/w for ratio [a,b] (long:short): + // landscape h = w*b/a, portrait h = w*a/b. + const ratioHW = (a: number, b: number) => (portrait ? a / b : b / a); + const changeAspect = (key: string) => { + setAspect(key); + if (key === "custom") return; + const [a, b] = ASPECT_RATIOS[key]; + setHeight(snapDim(width * ratioHW(a, b))); + }; + const changeWidth = (v: number) => { + setWidth(v); + if (aspect === "custom") return; + const [a, b] = ASPECT_RATIOS[aspect]; + setHeight(snapDim(v * ratioHW(a, b))); + }; + const changeHeight = (v: number) => { + setHeight(v); + if (aspect === "custom") return; + const [a, b] = ASPECT_RATIOS[aspect]; + setWidth(snapDim(v / ratioHW(a, b))); + }; + const flipDimensions = () => { + setWidth(height); + setHeight(width); + setPortrait((p) => !p); + }; + + const refreshStatus = useCallback(async () => { + try { + setStatus(await getDiffusionStatus()); + } catch { + // Status is best-effort; a failed poll shouldn't surface an error toast. + } + }, []); + + // Poll load-progress until the background load reaches "ready" or "error", + // updating the persistent toast in place each tick. + const pollLoadProgress = useCallback(async () => { + try { + const p = await getDiffusionLoadProgress(); + if (p.phase === "ready") { + dismissLoadToast(); + setStatus(await getDiffusionStatus()); + toast.success("Model loaded"); + setBusy(null); + return; + } + if (p.phase === "error") { + dismissLoadToast(); + toast.error(p.error || "Failed to load model"); + setBusy(null); + // A failed load may have freed a previously-loaded model, so resync to + // the real backend state (the synchronous failure path does the same). + void refreshStatus(); + return; + } + if (p.phase === null) { + // No load in flight and nothing loaded: the load was cancelled or + // evicted (e.g. a chat load took the GPU) and the backend cleared its + // state. Terminal — otherwise this loop would spin forever and leave + // busy stuck on "loading", deadening the picker and Generate button. + dismissLoadToast(); + setBusy(null); + void refreshStatus(); + return; + } + // Include bytes_total: the estimate lands as a 0→real jump while phase and + // bytes_downloaded hold, and the toast shows the percentage off that total. + const sig = `${p.phase}:${p.bytes_downloaded}:${p.bytes_total}`; + if (loadToastId.current != null && sig !== lastLoadSig.current) { + lastLoadSig.current = sig; + toast(null, loadToastArgs(p, loadToastId.current)); + } + } catch { + // Transient poll failure: keep trying. + } + pollTimer.current = setTimeout(() => void pollLoadProgress(), 1000); + }, [dismissLoadToast, refreshStatus]); + + useEffect(() => { + void (async () => { + await refreshStatus(); + // A load runs on the backend as a daemon thread that survives navigation. + // On (re)mount, resume tracking one that's still in flight so the page + // shows progress and updates on completion, instead of a stale view that + // never polls (no toast, and no refresh when the load finishes). + try { + const p = await getDiffusionLoadProgress(); + if (p.phase === "downloading" || p.phase === "finalizing") { + setBusy("loading"); + dismissLoadToast(); + lastLoadSig.current = null; + loadToastId.current = toast(null, loadToastArgs(p)); + void pollLoadProgress(); + } + } catch { + // Resume is best-effort; a failed probe just leaves the idle view. + } + })(); + // Stop polling if the page unmounts mid-load / mid-generate, and dismiss the + // load toast — its poll loop is gone, so it would otherwise hang forever + // (duration: Infinity) on whatever page the user navigated to. + return () => { + if (pollTimer.current) clearTimeout(pollTimer.current); + if (genPollTimer.current) clearInterval(genPollTimer.current); + dismissLoadToast(); + }; + }, [refreshStatus, dismissLoadToast, pollLoadProgress]); + + const handleLoad = useCallback( + async (repoId: string, ggufFilename: string) => { + // Cancel any prior poll loop so two can't run at once. + if (pollTimer.current) clearTimeout(pollTimer.current); + setBusy("loading"); + // Show the chat-style toast immediately; the poll updates it by id. + dismissLoadToast(); + lastLoadSig.current = null; + loadToastId.current = toast(null, loadToastArgs(IDLE_PROGRESS)); + try { + // Returns immediately — the load runs in the background; we poll for it. + // The backend infers the family + base diffusers repo from the repo id. + // Forward the saved HF token so gated bases (FLUX dev/klein) can download. + await loadDiffusionModel({ + model_path: repoId, + gguf_filename: ggufFilename, + hf_token: hfApiToken(getHfToken()), + }); + } catch (err) { + dismissLoadToast(); + toast.error(err instanceof Error ? err.message : "Failed to start load"); + setBusy(null); + void refreshStatus(); + return; + } + void pollLoadProgress(); + }, + [pollLoadProgress, refreshStatus, dismissLoadToast], + ); + + // The chat picker emits (modelId, picked quant + its exact filename); load it, + // and seed the inputs with that model's defaults. + const handleModelSelect = useCallback( + (id: string, meta: ModelSelectorChangeMeta) => { + // Ignore picks while a load/generation/unload is in flight: starting a + // replacement load now would tear down the live poll/toast and reset + // busy, while the backend rejects the second load with a 409. + if (busy !== null) return; + if (!meta.ggufVariant || !meta.ggufFilename) return; // not a quant pick + setQuant(meta.ggufVariant); + const d = defaultsFor(id); + setSteps(d.steps); + setGuidance(d.guidance); + void handleLoad(id, meta.ggufFilename); + }, + [busy, handleLoad], + ); + + const handleUnload = useCallback(async () => { + // Ejecting cancels any in-flight replacement load on the backend, so tear + // down its client-side tracking too: the load poll reschedules on phase + // null and the persistent toast never resolves, so both would otherwise + // leak forever after the unload. + if (pollTimer.current) clearTimeout(pollTimer.current); + pollTimer.current = null; + dismissLoadToast(); + lastLoadSig.current = null; + setBusy("unloading"); + try { + setStatus(await unloadDiffusionModel()); + setQuant(null); + } catch (err) { + toast.error(err instanceof Error ? err.message : "Failed to unload model"); + void refreshStatus(); + } finally { + setBusy(null); + } + }, [refreshStatus, dismissLoadToast]); + + const handleGenerate = useCallback(async () => { + if (!prompt.trim()) { + toast.error("Prompt is empty"); + return; + } + // Resolve a base seed up front. With an explicit seed the run is fully + // reproducible; with a random one we still pick a concrete base now so each + // sequential image gets a distinct, reproducible seed (base + i). + let baseSeed: number; + if (seed.trim()) { + const n = Number(seed); + if (!Number.isInteger(n) || n < 0 || n > Number.MAX_SAFE_INTEGER) { + toast.error("Seed must be a non-negative integer"); + return; + } + baseSeed = n; + } else { + baseSeed = Math.floor(Math.random() * 2 ** 32); + } + + // Snap custom dims to the model's grid so a half-typed value can't 422. + const w = snapDim(width); + const h = snapDim(height); + + // A large run count (generate all night) is a legitimate choice, so there's + // no upper cap; just floor at 1 and ignore non-numeric input (the number box + // can yield NaN), which would otherwise make the loop a silent no-op. + const runs = Number.isFinite(count) && count >= 1 ? Math.floor(count) : 1; + if (runs !== count) setCount(runs); + + setBusy("generating"); + setGenDone(0); + setGenStep(null); + // Poll the backend's per-step progress across the whole run (all sequential + // generations), so the bar tracks the live denoising steps. + genPollTimer.current = setInterval(async () => { + try { + const p = await getGenerateProgress(); + // Skip the state update (and re-render) when nothing the bar shows moved. + setGenStep((prev) => { + if (!p.active) return null; + if (prev && prev.step === p.step && prev.eta_seconds === p.eta_seconds) return prev; + return p; + }); + } catch { + // transient; keep polling + } + }, 300); + try { + for (let i = 0; i < runs; i++) { + const res = await generateDiffusionImage({ + prompt: prompt.trim(), + // Only send a negative prompt when guidance uses it, so the recipe + // doesn't record one the model ignored. + negative_prompt: guidance > 0 ? negativePrompt.trim() || undefined : undefined, + width: w, + height: h, + steps, + guidance, + seed: baseSeed + i, + batch_size: batchSize, + }); + // Prepend this run's records (newest first) and load their blobs. + setImages((prev) => [...res.images, ...prev]); + if (res.images[0]) setSelectedId(res.images[0].id); + res.images.forEach((image) => void ensureSrc(image)); + setGenDone(i + 1); + } + } catch (err) { + toast.error(err instanceof Error ? err.message : "Image generation failed"); + } finally { + if (genPollTimer.current) clearInterval(genPollTimer.current); + genPollTimer.current = null; + setBusy(null); + setGenDone(null); + setGenStep(null); + } + }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, ensureSrc]); + + return ( +
+ {/* ── Top: the model selector, kept at the chat tab's exact position so the + shared element matches. The load progress shows in a chat-style toast, + not here. ── */} +
+ +
+ + {/* ── Controls rail + preview canvas. Padding mirrors the other tabs + (Export, Data Recipes): px-5 / sm:px-9, with a roomy bottom. ── */} +
+ {/* The controls rail. Plain card (the gray surface) with no header — + the prompt + Generate button make the panel self-explanatory. */} +
+ +