Coordinate diffusion with training VRAM and guard its cache deletes

# Conflicts:
#	studio/backend/routes/models.py

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

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

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

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

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")

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):

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

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())

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()
 

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,

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

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"

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)

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

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"]

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,
 ]);

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,
+});

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();
                 }}
               />
+              <NavItem
+                icon={PaintBrush02Icon}
+                label={t("shell.navigation.images")}
+                active={pathname === "/images" || pathname.startsWith("/images/")}
+                disabled={chatOnly}
+                tooltip={trainExportDisabledHint}
+                onClick={() => {
+                  if (chatOnly) return;
+                  navigate({ to: "/images" });
+                  closeMobileIfOpen();
+                }}
+              />
               {/* Train has a labelled section when expanded; plain icon here only when collapsed. */}
               <NavItem
                 icon={TestTubeOutlineIcon}

studio/frontend/src/components/assistant-ui/model-selector.tsx

@@ -33,6 +33,7 @@ import {
   useState,
 } from "react";
 import { Input } from "../ui/input";
+import type { HfTaskFilter } from "@/features/hub/hooks/use-hub-model-search";
 import { HubModelPicker, hasDownloadedModels } from "./model-selector/pickers";
 import { PillTabs } from "./model-selector/pill-tabs";
 import {
@@ -136,6 +137,8 @@ interface ModelSelectorProps {
   triggerDataTour?: string;
   contentDataTour?: string;
   showCloudIndicator?: boolean;
+  /** Restrict the Hub tab to a pipeline task (e.g. text-to-image). */
+  task?: HfTaskFilter;
 }
 
 function ModelSelectorTrigger({
@@ -276,6 +279,7 @@ function ModelSelectorContent({
   deleteDisabled,
   className,
   dataTour,
+  task,
 }: {
   open: boolean;
   models: ModelOption[];
@@ -291,6 +295,7 @@ function ModelSelectorContent({
   deleteDisabled?: boolean;
   className?: string;
   dataTour?: string;
+  task?: HfTaskFilter;
 }) {
   const hasSelection = Boolean(value);
   const chatOnly = usePlatformStore((s) => s.isChatOnly());
@@ -460,6 +465,7 @@ function ModelSelectorContent({
             deleteDisabled={deleteDisabled}
             section={effectiveHubSection}
             onEject={hasSelection && onEject ? onEject : undefined}
+            task={task}
             sectionToggle={
               <PillTabs
                 ariaLabel="Hub section"
@@ -539,6 +545,7 @@ export function ModelSelector({
   triggerDataTour,
   contentDataTour,
   showCloudIndicator = false,
+  task,
 }: ModelSelectorProps) {
   const [uncontrolledOpen, setUncontrolledOpen] = useState(false);
   const open = controlledOpen ?? uncontrolledOpen;
@@ -655,11 +662,15 @@ export function ModelSelector({
         onEject={onEject ? handleEject : undefined}
         onFoldersChange={onFoldersChange}
         onPickLocalModel={onPickLocalModel ? handlePickLocalModel : undefined}
-        onBrowseHub={handleBrowseHub}
+        // The image tab (the only caller passing `task`) is a self-contained
+        // curated + on-device picker, so it omits the "Search Hub" button that
+        // navigates to the general Hub page.
+        onBrowseHub={task ? undefined : handleBrowseHub}
         onModelsChange={onModelsChange}
         deleteDisabled={deleteDisabled}
         className={contentClassName}
         dataTour={contentDataTour}
+        task={task}
       />
     </Popover>
   );

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({
                             </TooltipContent>
                           </Tooltip>
                         ) : null}
+                        {!task && (
                         <Tooltip delayDuration={0}>
                           <TooltipTrigger asChild={true}>
                             <button
@@ -2588,6 +2690,7 @@ export function HubModelPicker({
                             Go to fine-tuned models
                           </TooltipContent>
                         </Tooltip>
+                        )}
                         <Tooltip delayDuration={0}>
                           <TooltipTrigger asChild={true}>
                             <button
@@ -2646,8 +2749,9 @@ export function HubModelPicker({
 
               {/* Fine-tuned models: a section above Custom Folders. Always shown on
               On Device so the train shortcut always has a target, with an empty
-              state when none exist. */}
-              {section === "downloaded" ? (
+              state when none exist. Hidden under a task filter (e.g. Images —
+              we don't fine-tune image models). */}
+              {section === "downloaded" && !task ? (
                 <>
                   <div
                     ref={fineTunedSectionRef}

studio/frontend/src/components/assistant-ui/model-selector/types.ts

@@ -29,6 +29,9 @@ export interface ModelSelectorChangeMeta {
   source: "hub" | "lora" | "exported" | "local" | "external";
   isLora: boolean;
   ggufVariant?: string;
+  /** Exact GGUF filename for the picked quant (filenames don't always follow the
+   *  repo name, e.g. FLUX.1-schnell -> flux1-schnell-*.gguf). */
+  ggufFilename?: string;
   isDownloaded?: boolean;
   expectedBytes?: number;
   /** Native GGUF context, threaded so a staged pick can seed the slider. */

studio/frontend/src/features/chat/api/chat-api.ts

@@ -203,6 +203,10 @@ export interface CachedGgufRepo {
   /** True when the repo ships an mmproj adapter (image inputs). Optional for
    * older-backend compatibility. */
   has_vision?: boolean;
+  /** HF pipeline task inferred from the GGUF architecture ("text-to-image" for
+   * diffusion, "text-generation" otherwise). Lets the Images picker show only
+   * diffusion GGUFs. Optional for older-backend compatibility. */
+  task?: string | null;
 }
 
 export async function getGgufDownloadProgress(
@@ -285,6 +289,10 @@ export interface LocalModelInfo {
   // classify scanned folders whose name lacks a -GGUF suffix.
   model_format?: string | null;
   updated_at?: number | null;
+  // HF pipeline task inferred from the GGUF architecture, so the Images picker
+  // can filter local models to diffusion ("text-to-image"). Optional for
+  // older-backend compatibility.
+  task?: string | null;
 }
 
 interface LocalModelListResponse {
@@ -311,6 +319,9 @@ export interface CachedModelRepo {
   /** Epoch seconds of the newest downloaded weight file; sorts Downloaded
    * newest-first. Optional for older-backend compatibility. */
   last_modified?: number;
+  /** HF pipeline task: "text-to-image" for a cached diffusers pipeline repo
+   * (model_index.json present), so the chat picker can hide it. Absent = chat. */
+  task?: string | null;
 }
 
 export async function listCachedModels(): Promise<CachedModelRepo[]> {

studio/frontend/src/features/chat/components/model-load-status.tsx

@@ -4,12 +4,15 @@
 import { Progress } from "@/components/ui/progress";
 import { Spinner } from "@/components/ui/spinner";
 import { Button } from "@/components/ui/button";
+import { cn } from "@/lib/utils";
 
 type ModelLoadDescriptionProps = {
   title?: string | null;
   message?: string | null;
   progressPercent?: number | null;
   progressLabel?: string | null;
+  // Extra classes for the root row (e.g. a titleless caller dropping min-h-12).
+  className?: string;
 };
 
 function clampProgress(value: number): number {
@@ -39,6 +42,7 @@ export function ModelLoadDescription({
   message,
   progressPercent,
   progressLabel,
+  className,
 }: ModelLoadDescriptionProps) {
   const hasProgress = typeof progressPercent === "number";
   // Split once at the top so the JSX below stays flat (no IIFE).
@@ -46,7 +50,7 @@ export function ModelLoadDescription({
     splitProgressLabel(progressLabel);
 
   return (
-    <div className="relative flex min-h-12 w-full items-stretch gap-2">
+    <div className={cn("relative flex min-h-12 w-full items-stretch gap-2", className)}>
       <div className="flex h-full shrink-0 items-center self-center">
         <Spinner className="size-3.5 text-muted-foreground" />
       </div>

studio/frontend/src/features/images/api.ts

@@ -0,0 +1,145 @@
+// 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 { authFetch } from "@/features/auth";
+import { readFastApiError } from "@/lib/format-fastapi-error";
+
+export interface DiffusionStatus {
+  loaded: boolean;
+  repo_id: string | null;
+  family: string | null;
+  base_repo: string | null;
+  device: string | null;
+  dtype: string | null;
+  cpu_offload: boolean;
+}
+
+export interface DiffusionGenerateProgress {
+  active: boolean;
+  step: number;
+  total_steps: number;
+  fraction: number;
+  eta_seconds: number | null;
+}
+
+export interface DiffusionLoadProgress {
+  phase: "downloading" | "finalizing" | "ready" | "error" | null;
+  bytes_downloaded: number;
+  bytes_total: number;
+  fraction: number;
+  error: string | null;
+}
+
+export interface DiffusionLoadRequest {
+  model_path: string;
+  gguf_filename: string;
+  base_repo?: string;
+  family_override?: string;
+  hf_token?: string;
+  cpu_offload?: boolean;
+}
+
+export interface DiffusionGenerateRequest {
+  prompt: string;
+  negative_prompt?: string;
+  width?: number;
+  height?: number;
+  steps?: number;
+  guidance?: number;
+  seed?: number;
+  batch_size?: number;
+}
+
+// A persisted image's full generation recipe (also embedded in the PNG).
+export interface GalleryImage {
+  id: string;
+  url: string;
+  prompt: string;
+  negative_prompt: string | null;
+  width: number;
+  height: number;
+  steps: number;
+  guidance: number;
+  seed: number;
+  batch_index: number;
+  model: string | null;
+  created_at: number;
+}
+
+export interface DiffusionGenerateResponse {
+  images: GalleryImage[];
+}
+
+async function parseJson<T>(response: Response): Promise<T> {
+  if (!response.ok) {
+    throw new Error(await readFastApiError(response));
+  }
+  return (await response.json()) as T;
+}
+
+export async function getDiffusionStatus(): Promise<DiffusionStatus> {
+  return parseJson(await authFetch("/api/inference/images/status"));
+}
+
+export async function getDiffusionLoadProgress(): Promise<DiffusionLoadProgress> {
+  return parseJson(await authFetch("/api/inference/images/load-progress"));
+}
+
+export async function getGenerateProgress(): Promise<DiffusionGenerateProgress> {
+  return parseJson(await authFetch("/api/inference/images/generate-progress"));
+}
+
+export async function loadDiffusionModel(body: DiffusionLoadRequest): Promise<DiffusionStatus> {
+  return parseJson(
+    await authFetch("/api/inference/images/load", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify(body),
+    }),
+  );
+}
+
+export async function generateDiffusionImage(
+  body: DiffusionGenerateRequest,
+): Promise<DiffusionGenerateResponse> {
+  return parseJson(
+    await authFetch("/api/inference/images/generate", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify(body),
+    }),
+  );
+}
+
+export async function unloadDiffusionModel(): Promise<DiffusionStatus> {
+  return parseJson(await authFetch("/api/inference/images/unload", { method: "POST" }));
+}
+
+export interface GalleryPage {
+  images: GalleryImage[];
+  has_more: boolean;
+}
+
+export async function getGallery(offset = 0, limit = 50): Promise<GalleryPage> {
+  return parseJson(
+    await authFetch(`/api/inference/images/gallery?offset=${offset}&limit=${limit}`),
+  );
+}
+
+export async function deleteGalleryImage(id: string): Promise<void> {
+  const res = await authFetch(`/api/inference/images/gallery/${id}`, { method: "DELETE" });
+  if (!res.ok) throw new Error(await readFastApiError(res));
+}
+
+export async function clearGallery(): Promise<void> {
+  const res = await authFetch("/api/inference/images/gallery", { method: "DELETE" });
+  if (!res.ok) throw new Error(await readFastApiError(res));
+}
+
+/** Fetch a gallery PNG (auth-protected, so it can't be a plain <img src>) and
+ *  wrap it in an object URL. Callers must revoke the URL when done. */
+export async function fetchGalleryObjectUrl(url: string): Promise<string> {
+  const res = await authFetch(url);
+  if (!res.ok) throw new Error(await readFastApiError(res));
+  return URL.createObjectURL(await res.blob());
+}

studio/frontend/src/features/images/index.ts

@@ -0,0 +1,4 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+export { ImagesPage } from "./images-page";

studio/frontend/src/i18n/locales/en.ts

@@ -39,6 +39,7 @@ export const en = {
       hub: "Hub",
       train: "Train",
       recipes: "Recipes",
+      images: "Images",
       export: "Export",
       recents: "Recents",
       settings: "Settings",

studio/frontend/src/i18n/locales/zh-CN.ts

@@ -39,6 +39,7 @@ export const zhCN = {
       hub: "Hub",
       train: "训练",
       recipes: "配方",
+      images: "图像",
       export: "导出",
       recents: "最近",
       settings: "设置",