# SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """Diffusion image generation backend. Loads Hugging Face diffusion checkpoints in either the standard ``diffusers`` layout or the single-file GGUF layout published under ``unsloth/*-GGUF`` (Flux 2, Flux 2 Klein, Qwen-Image, SD3, SDXL, ...). GGUF files are dynamically dequantised on-device via ``diffusers.GGUFQuantizationConfig``, then the rest of the pipeline (VAE, text encoders, scheduler) is pulled from the matching ``diffusers`` repo so end users only ever need one local file plus the metadata repo. The module is intentionally torch-only: it never spawns a subprocess and shares the active CUDA / MPS device with the rest of Studio. The cost of not having a separate process is that loading a diffusion model and a GGUF chat model at the same time can OOM on consumer GPUs; the routes layer must therefore swap between the two as needed (the orchestrator unloads llama-server before any diffusion load on hosts with < 24 GB). The class deliberately exposes a small, llama-cpp-style surface: load_model(repo_id, ...) generate_image(prompt, ...) -> PIL.Image unload_model() status() -> dict so the route layer at ``studio/backend/routes/inference.py`` can mirror the existing llama-server lifecycle (probe + load + generate + unload) without learning a second API. """ from __future__ import annotations import asyncio import gc import io import re import threading import time from dataclasses import dataclass, field from pathlib import Path, PurePosixPath from typing import Any, Optional from loggers import get_logger logger = get_logger(__name__) # ─── Pipeline registry ──────────────────────────────────────────────── # # Keep this list narrow on purpose: only ship the small text-to-image # families with first-class GGUF coverage on the Hub. Anything else is # either video (LTX*, Wan) or research-grade (Sana, SD3.5) and can be # added once it has a working GGUF release plus a smoke test. # # Each entry maps a substring of the loaded repo id (case-insensitive) # to the (pipeline_class_name, transformer_class_name, default base # repo for missing pieces). ``base_repo`` is what we pass to # ``Pipeline.from_pretrained`` to pick up the VAE + text encoders when # the user gave us a GGUF-only repo. The base_repo is documented to the # user via ``status()`` so they understand why a second download fires. @dataclass(frozen = True) class DiffusionFamily: name: str pipeline_class: str transformer_class: str base_repo: str # Optional: list of HF "trigger" substrings besides ``name`` that map # to this family (e.g. "flux1-dev" plus "flux.1-dev"). Lowercased. aliases: tuple[str, ...] = field(default_factory = tuple) _FAMILIES: tuple[DiffusionFamily, ...] = ( # The "9b" alias is checked first so a "flux-2-klein-9b" GGUF picks # the 9B base instead of the 4B one when the user does not pass an # explicit base_repo. Apache 2.0 is preferred as the auto-default for # the 4B path because BFL's 9B base is gated. DiffusionFamily( name = "flux.2-klein", pipeline_class = "Flux2KleinPipeline", transformer_class = "Flux2Transformer2DModel", # Default for klein when no explicit base_repo: Apache-2.0 4B Base. # The frontend curated picker always passes base_repo explicitly, # so this default only fires for "custom HF repo" mode. base_repo = "black-forest-labs/FLUX.2-klein-base-4B", aliases = ("flux2-klein", "flux-2-klein", "flux.2.klein"), ), DiffusionFamily( name = "flux.2", pipeline_class = "Flux2Pipeline", transformer_class = "Flux2Transformer2DModel", base_repo = "black-forest-labs/FLUX.2-dev", aliases = ("flux2-dev", "flux-2-dev", "flux.2.dev"), ), DiffusionFamily( name = "flux.1", pipeline_class = "FluxPipeline", transformer_class = "FluxTransformer2DModel", base_repo = "black-forest-labs/FLUX.1-dev", aliases = ("flux1-dev", "flux-1-dev", "flux.1.dev", "flux-dev"), ), DiffusionFamily( name = "qwen-image", pipeline_class = "QwenImagePipeline", transformer_class = "QwenImageTransformer2DModel", base_repo = "Qwen/Qwen-Image", aliases = ("qwenimage", "qwen_image"), ), DiffusionFamily( name = "stable-diffusion-3", pipeline_class = "StableDiffusion3Pipeline", transformer_class = "SD3Transformer2DModel", base_repo = "stabilityai/stable-diffusion-3-medium-diffusers", # Intentionally NOT including "sd3.5" / "stable-diffusion-3.5" # here: the SD3.5 family uses a different transformer config and # base repo than SD3 Medium, and silently pairing SD3.5 GGUFs # with the Medium base produces a misleading load. Add a # dedicated SD3.5 family with its own base_repo when we ship # smoke coverage for it. aliases = ("sd3-medium", "stable-diffusion-3-medium"), ), # SDXL: full diffusers path only (no GGUF). SDXL uses a UNet (not a # transformer) and wiring UNet2DConditionModel.from_single_file + # GGUF is a separate code path the rest of this module does not # exercise. The family is intentionally NOT in _FAMILIES so the # frontend status panel does not advertise GGUF support we do not # implement; callers wanting SDXL full repos can still do so by # passing the diffusers repo with no gguf_filename and # family_override = "stable-diffusion-xl" via the route, which uses # the lookup in _FULL_REPO_FAMILIES. ) # Families available via family_override on the routes layer when the # user is loading a full diffusers checkpoint (no GGUF). Kept separate # from _FAMILIES so the GGUF-only status panel does not over-advertise. _FULL_REPO_FAMILIES: tuple[DiffusionFamily, ...] = ( DiffusionFamily( name = "stable-diffusion-xl", pipeline_class = "StableDiffusionXLPipeline", transformer_class = "", base_repo = "stabilityai/stable-diffusion-xl-base-1.0", aliases = ("sdxl",), ), ) def _smart_base_repo(fam: DiffusionFamily, repo_id: str) -> str: """Pick the best matching base diffusers repo for a given GGUF repo when the caller did not pass an explicit base_repo. Currently only specialises the flux.2-klein family: a repo name containing "9b" gets the 9B base, "base-4b" / "base-9b" map to the Base variants, everything else falls back to the family default (Apache 2.0 4B Base). Only the LAST segment of the repo id / path is inspected so a namespace or parent directory like ``baseorg/...`` or ``/home/me/.cache/base/...`` does not falsely select the Base variant (round 12 review #9). Splits on BOTH ``/`` and ``\\`` so Windows local paths like ``C:\\Users\\me\\base\\FLUX.2-klein-4B`` do not get scored as "base" via the parent directory either (round 13 P2 #13). """ if fam.name != "flux.2-klein": return fam.base_repo cleaned = (repo_id or "").rstrip("/\\") last_segment = re.split(r"[\\/]+", cleaned)[-1].lower() if cleaned else "" is_9b = "9b" in last_segment is_base = "base" in last_segment if is_9b and is_base: return "black-forest-labs/FLUX.2-klein-base-9B" if is_9b: return "black-forest-labs/FLUX.2-klein-9B" if is_base: return "black-forest-labs/FLUX.2-klein-base-4B" # Distilled 4B is the default for any flux-2-klein GGUF that does # not advertise 9B or "base". return "black-forest-labs/FLUX.2-klein-4B" def _expand_existing_local_path(value: str) -> str: """Expand ``~`` in ``value`` when the expanded path exists locally. Round 14 P2 #11: the GGUF local path branch already calls ``Path(repo_id).expanduser()``, but the full-diffusers-repo and base-companion-repo paths passed the literal ``~/...`` straight into ``from_pretrained``, which treated it as a Hub id and tried to download. Keep behaviour identical for Hub ids (no leading ``~`` -> return as-is) and for non-existent expansions (the diffusers loader will surface its own ``not found`` error). """ if not value or not isinstance(value, str) or not value.startswith("~"): return value candidate = Path(value).expanduser() if candidate.exists(): return str(candidate) return value def _preflight_diffusers_subfolder_config( repo: str, subfolder: str, hf_token: Optional[str], ) -> None: """Round 21 P2 #6: also probe ``{subfolder}/config.json``. The full-repo preflight at ``_preflight_full_diffusers_repo`` only proves ``model_index.json`` exists. For GGUF loads the follow-up ``from_single_file(..., config=effective_base, subfolder="transformer")`` still needs a matching ``transformer/config.json`` on the base companion. Without this second probe a base that has model_index.json but no transformer config would still unload chat before the load failed. """ if not repo or not subfolder: return try: local = Path(repo).expanduser() except (OSError, ValueError): local = None if local is not None and local.exists(): config_path = local / subfolder / "config.json" if not config_path.is_file(): raise RuntimeError( f"Diffusion repo '{_display_repo_id(repo)}' is missing " f"{subfolder}/config.json." ) return if (local is not None and local.is_absolute()) or repo.startswith("~"): # Local-only path that does not exist -- _preflight_full_diffusers_repo # already raised for the absent directory, so reaching here means the # caller is loading a Hub id that just looks like a path. Fall through # to the network probe. pass try: from huggingface_hub import hf_hub_download as _hf_hub_download except Exception: return try: _hf_hub_download( repo_id = repo, filename = "config.json", subfolder = subfolder, token = hf_token, ) except Exception as exc: raise RuntimeError( f"Could not access diffusion repo '{_display_repo_id(repo)}' " f"{subfolder}/config.json before unloading the current model." ) from exc def _preflight_full_diffusers_repo(repo: str, hf_token: Optional[str]) -> None: """Prove a full diffusers repo is accessible before any unloads. Round 19 P1 #3: the GGUF path's ``hf_hub_download(gguf_filename)`` above this function fails fast on a bad / private / gated / typo'd repo before we touch the chat backend. The full diffusers path used to skip that round-trip and only discover the issue inside ``from_pretrained`` AFTER the user's chat model was already unloaded. Add the same one-file probe (``model_index.json`` is the diffusers manifest; every diffusers repo has one). Local paths are checked structurally so we do not hit the network for a missing on-disk directory; both branches raise RuntimeError so the surrounding load_model bails out before the chat unload. The display label is collapsed via ``_display_repo_id`` so an absolute filesystem path in the error message does not leak the operator's layout (see round 17 P2 #9). """ if not repo: return try: local = Path(repo).expanduser() except (OSError, ValueError): local = None if local is not None and local.exists(): if not local.is_dir(): raise RuntimeError( f"Diffusion repo '{_display_repo_id(repo)}' is not a directory." ) if not (local / "model_index.json").is_file(): raise RuntimeError( f"Diffusion repo '{_display_repo_id(repo)}' is missing " "model_index.json." ) return if (local is not None and local.is_absolute()) or repo.startswith("~"): raise RuntimeError( f"Local diffusion repo '{_display_repo_id(repo)}' does not exist." ) try: from huggingface_hub import hf_hub_download as _hf_hub_download except Exception: # diffusers is installed but huggingface_hub is missing -- let # the downstream loader produce the canonical error. return try: _hf_hub_download( repo_id = repo, filename = "model_index.json", token = hf_token, ) except Exception as exc: raise RuntimeError( f"Could not access diffusion repo '{_display_repo_id(repo)}' " "before unloading the current model." ) from exc def _display_repo_id(value: Any) -> Any: """Return a public-facing label for a repo_id / base_repo. For Hub-style identifiers (``owner/repo``) the value passes through unchanged so the Images panel and result figcaption stay informative. Absolute local paths (``/home/me/exports/...`` or ``C:\\Users\\...``) collapse to the leaf name so ``/images/status`` does not leak the user's filesystem layout to other authenticated browser sessions (round 15 P2 #6). HF tokens are scrubbed defensively in case they slipped past the request-side validator. """ if not isinstance(value, str) or not value: return value try: candidate = Path(value).expanduser() if candidate.is_absolute() or candidate.exists(): return candidate.name or value except (OSError, ValueError): pass return _redact_hf_tokens(value) _HF_TOKEN_RE = re.compile(r"hf_[A-Za-z0-9]{20,}") def _redact_hf_tokens(value: Any) -> Any: """Scrub embedded ``hf_xxxxxxxx`` tokens out of a string before logging. Round 14 P2 #9: callers can wrap an authenticated URL (``https://hf_token@huggingface.co/...``) into ``repo_id`` / ``base_repo`` / paths; the token would otherwise reach structured-log sinks via the load-info / load-failure log lines. Non-strings are returned unchanged so the helper is safe to sprinkle through ``logger.info`` / ``logger.error`` argument lists. """ if not isinstance(value, str): return value return _HF_TOKEN_RE.sub("", value) def _resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path: """Resolve a GGUF filename inside a local repo directory safely. Returns the resolved absolute path or raises ``RuntimeError`` if: - ``gguf_filename`` is absolute (``/etc/passwd``) or contains a Windows separator (``..\\..\\secret.gguf``); - the parts contain ``""`` / ``.`` / ``..`` (``../other.gguf``); - the resolved candidate escapes ``repo_root`` after symlinks / ``..`` collapse; - the resolved candidate is not a regular file. This is the only path that bridges a user-supplied ``gguf_filename`` string into ``Path``s the loader opens, so confining it to the chosen repo here protects the delete-ownership guards downstream (round 13 P1 #2). ``hf_hub_download`` already enforces the same invariant for Hub repos. """ # ``Path("/etc/passwd").is_absolute()`` is False on Windows (POSIX # absolute paths read as drive-relative), so check both pathlib # flavours plus a leading separator so the rejection is portable. if ( Path(gguf_filename).is_absolute() or PurePosixPath(gguf_filename).is_absolute() or gguf_filename.startswith(("/", "\\")) or "\\" in gguf_filename ): raise RuntimeError("gguf_filename must be a relative file path inside repo_id.") rel = PurePosixPath(gguf_filename) if any(part in ("", ".", "..") for part in rel.parts): raise RuntimeError( "gguf_filename must not contain empty, '.', or '..' segments." ) root = repo_root.expanduser().resolve(strict = True) try: candidate = (root / Path(*rel.parts)).resolve(strict = True) except (OSError, FileNotFoundError) as exc: # strict=True raises FileNotFoundError on a missing leaf or # parent component, and OSError on a malformed Windows path # (e.g. drive letters injected through the user-supplied # string). Either way the candidate does not exist inside the # chosen repo, which is exactly the "file not in repo" failure # mode the caller cares about. raise RuntimeError( f"Local repo path '{repo_root}' does not contain '{gguf_filename}'." ) from exc try: candidate.relative_to(root) except ValueError as exc: raise RuntimeError( "gguf_filename must stay inside the local repo_id directory." ) from exc if not candidate.is_file(): raise RuntimeError( f"Local repo path '{repo_root}' does not contain '{gguf_filename}'." ) return candidate # Negative substrings that disqualify a candidate family even when its # name appears as a substring of the repo id. Prevents # "stable-diffusion-3" matching SD3.5 and "qwen-image" matching # Qwen-Image-Edit. Each entry maps a family name to substrings that # must NOT appear anywhere in the repo id. _FAMILY_EXCLUDE: dict[str, tuple[str, ...]] = { "stable-diffusion-3": ( "3.5", "3-5", "3_5", "stable-diffusion-3.5", "stable_diffusion_3_5", ), # All underscore / hyphen spellings that appear in Hub repo ids for # the *-Edit family must exclude Qwen-Image, otherwise # ``unsloth/qwen_image_edit-GGUF`` matches the Qwen-Image base. "qwen-image": ( "qwen-image-edit", "qwenimage-edit", "qwen_image_edit", "qwenimageedit", ), } def detect_family( repo_id: str, *, override_family: Optional[str] = None ) -> Optional[DiffusionFamily]: """Return the diffusion family matching ``repo_id``. Matching is substring-based and case-insensitive, with a small deny list (``_FAMILY_EXCLUDE``) for known false positives such as SD3.5 (would otherwise match SD3 Medium) and Qwen-Image-Edit (would otherwise match Qwen-Image). ``override_family`` bypasses substring matching and looks up by ``DiffusionFamily.name`` or (when explicitly asked) by ``_FULL_REPO_FAMILIES.name``. Returns ``None`` when no family applies so callers can surface a clear "unsupported model" error rather than guessing wrong. """ if override_family: wanted = override_family.strip().lower() for fam in _FAMILIES + _FULL_REPO_FAMILIES: if fam.name == wanted: return fam return None needle = (repo_id or "").lower() if not needle: return None # Round 17 P2 #10: if repo_id is an absolute local path, the # whole path goes into ``needle`` and the _FAMILY_EXCLUDE deny # lists match against parent-directory names too. That means # ``/home/me/qwen-image-edit-cache/flux-2-klein-4b`` would be # excluded from the Flux family because the parent contains # ``qwen-image-edit``. Reduce to the leaf when the candidate # looks like a filesystem path so excludes only consider the # model directory itself. if "/" in needle or "\\" in needle: try: candidate = Path(repo_id).expanduser() if candidate.is_absolute() or candidate.exists(): leaf = candidate.name if leaf: needle = leaf.lower() except (OSError, ValueError): pass # Normalise mixed separator spellings (``Qwen_Image-Edit-GGUF``, # ``Qwen-Image_Edit-GGUF``, ``Qwen.Image.Edit-GGUF``) and the # compact concatenation (``QwenImageEdit-GGUF``) so the # _FAMILY_EXCLUDE deny lists do not need every permutation of # ``-``, ``_``, ``.`` and run-together spellings to keep # Qwen-Image-Edit out of the base Qwen-Image family (round 14 # P2 #8). needle_norm = re.sub(r"[^a-z0-9]+", "-", needle).strip("-") needle_compact = re.sub(r"[^a-z0-9]+", "", needle) # Per-token compact strings let ``unsloth/Flux2Klein-GGUF`` match # the ``flux2klein`` alias: the whole-needle compact is # ``unslothflux2kleingguf`` and the regex boundary check rejects # the embedded match, but the token ``Flux2Klein`` (between the # ``/`` and the ``-``) compacts to exactly ``flux2klein`` (round # 16 P2 #9). needle_compact_tokens = { re.sub(r"[^a-z0-9]+", "", token) for token in re.split(r"[^a-z0-9]+", needle) if token } def _matches_family_token(term: str) -> bool: """Token-boundary match on the normalised needle. Prevents ``owner/flux.20-model`` from matching ``flux.2`` because ``flux.20`` does not have a separator after ``flux-2`` (round 15 P2 #8). Compact spellings (``flux2klein``) match only when they appear as a complete repo-name token, not as a substring of a longer token (round 16 P2 #9).""" term_norm = re.sub(r"[^a-z0-9]+", "-", term.lower()).strip("-") if not term_norm: return False if re.search(rf"(^|-){re.escape(term_norm)}($|-)", needle_norm): return True term_compact = re.sub(r"[^a-z0-9]+", "", term.lower()) if not term_compact: return False return term_compact in needle_compact_tokens or term_compact == needle_compact # Scan _FAMILIES first (GGUF-supported), then _FULL_REPO_FAMILIES # so a repo like ``stabilityai/stable-diffusion-xl-base-1.0`` is # auto-detected as SDXL instead of returning None. for fam in _FAMILIES + _FULL_REPO_FAMILIES: excludes = _FAMILY_EXCLUDE.get(fam.name, ()) if any( e in needle or re.sub(r"[^a-z0-9]+", "-", e).strip("-") in needle_norm or re.sub(r"[^a-z0-9]+", "", e) in needle_compact for e in excludes ): continue if _matches_family_token(fam.name): return fam for alias in fam.aliases: if alias and _matches_family_token(alias): return fam return None def supported_families() -> list[dict[str, str]]: """Public-facing list of families for ``/api/inference/images/status``.""" return [ { "name": fam.name, "pipeline_class": fam.pipeline_class, "base_repo": fam.base_repo, } for fam in _FAMILIES ] # ─── Backend ────────────────────────────────────────────────────────── class DiffusionBackend: """Singleton-style diffusion backend. One pipeline at a time; ``load_model`` swaps the previous one out. Generation is mutex'd so concurrent requests serialise rather than racing GPU memory. """ def __init__(self) -> None: self._pipe: Any = None # `_lock` protects mutations to the small state fields and is # the only lock taken by status(). It is intentionally NOT held # for the long pipeline forward pass: holding it for the whole # generate would block status() polls (frontend at 1 Hz) and # any concurrent unload requests for minutes at a time. # # `_load_lock` serialises the entire load_model call so two # concurrent /images/load requests cannot both reach # pipeline_cls.from_pretrained at the same time (which would # double-spend VRAM and corrupt _pipe). # # `_generate_lock` serialises pipeline __call__ since diffusers # pipelines are not thread-safe; overlapping forwards on the # shared pipe corrupt internal scheduler state. # # Lock order is load -> state and generate -> state (never # state -> load/generate) so a forward in flight cannot # deadlock the next load or a status poll. self._lock = threading.Lock() self._load_lock = threading.Lock() self._generate_lock = threading.Lock() self._family: Optional[DiffusionFamily] = None self._repo_id: Optional[str] = None self._gguf_path: Optional[str] = None # Original ``gguf_filename`` the caller passed in, preserved # so delete guards can compare against subdirectory variants # like ``BF16/model.gguf`` or ``Q4_K_M/model.gguf`` instead # of the collapsed basename (round 14 P1 #4). The basename # alone (``model.gguf``) loses the quant directory and lets # /delete-cached unlink the wrong file. self._gguf_filename: Optional[str] = None self._base_repo: Optional[str] = None self._device: Optional[str] = None self._dtype: Optional[str] = None # True when ``enable_model_cpu_offload()`` was applied on the # loaded pipeline. Diffusers' offload moves the active # submodule between CPU and GPU on each step, so a CUDA # ``torch.Generator`` mismatches the CPU-resident embeddings # and generation crashes mid-forward (round 14 P1 #6). When # this is True, seeded generation has to use a CPU generator # regardless of self._device. self._cpu_offload_enabled: bool = False self._loaded_at: Optional[float] = None self._loading: bool = False self._last_error: Optional[str] = None # `_pending_*` fields advertise the target of an in-flight load # so cache- and finetuned-delete guards can refuse to rmtree a # repo while it is being downloaded / read. They are set under # _lock at the start of load_model and cleared on success or # in the finally block. The route layer reads them via # status() under _lock. self._pending_repo_id: Optional[str] = None self._pending_base_repo: Optional[str] = None self._pending_gguf_filename: Optional[str] = None # ── lifecycle ───────────────────────────────────────────────── @property def is_loaded(self) -> bool: return self._pipe is not None @property def repo_id(self) -> Optional[str]: return self._repo_id def status(self, *, include_internal: bool = False) -> dict[str, Any]: # Take _lock so the snapshot cannot observe a torn state where # _pipe was already swapped but _family/_repo_id haven't been # updated yet (or vice versa). Frontend polling at 1 Hz would # otherwise render impossible "loaded but no repo_id" states. # Only echo the GGUF basename; full absolute path leaks the # local HF cache layout (and the system username on default # POSIX layouts) to any authenticated Studio session. # # Round 16 P1 #5: the guard-facing ``active_*`` / ``pending_*`` # fields hold the EXACT raw path (so /delete-cached can match # an HF snapshot mmap) but are NOT safe to surface to the # browser. Callers that need the raw path (route-internal # delete guards) pass ``include_internal=True``; the public # ``/api/inference/images/status`` route always uses the # public payload. with self._lock: # UI-facing collapsed basename. Full local path leaks the # HF cache layout + system username; the original caller- # supplied filename (e.g. ``BF16/model.gguf``) is kept # separately as ``active_gguf_filename`` for delete # guards. gguf_basename = Path(self._gguf_path).name if self._gguf_path else None # Expose BOTH the resident pipeline's id AND the pending # load target. Delete guards must check both: when model A # is already loaded and a swap to model B is in flight, # only checking one would let the user rmtree whichever # repo the guard ignored. UI-facing ``repo_id`` / # ``base_repo`` / ``gguf_filename`` still prefer pending # during a swap so the panel shows the load target the # user just clicked. active_repo = self._repo_id active_base = self._base_repo active_gguf = self._gguf_filename pending_repo = self._pending_repo_id if self._loading else None pending_base = self._pending_base_repo if self._loading else None pending_gguf = self._pending_gguf_filename if self._loading else None # When a swap is in flight, the UI-facing repo_id / # base_repo / gguf_filename advertise the PENDING model # but ``self._family`` still points at the previously # loaded pipeline. Reporting them together produces a # repo/family pair that never existed (round 11 #6). # Null the family / pipeline_class while a swap is in # flight; the frontend can fall back to "unknown". ui_family = self._family.name if self._family else None ui_pipeline_class = self._family.pipeline_class if self._family else None if pending_repo and pending_repo != active_repo: ui_family = None ui_pipeline_class = None # UI-facing ``gguf_filename`` collapses to the basename # so the Images panel does not surface internal cache / # variant directory names. Guard-facing ``active_*`` / # ``pending_*`` retain the full caller-supplied filename # so /delete-cached can compare against subdirectory # variants like ``BF16/model.gguf`` (round 14 P1 #4-5). ui_gguf = pending_gguf or active_gguf ui_gguf_basename = Path(ui_gguf).name if ui_gguf else None # UI-facing ``repo_id`` / ``base_repo`` collapse absolute # local paths to their leaf name so ``/images/status`` # does not leak the user's filesystem layout to other # authenticated browser sessions (round 15 P2 #6). The # guard-facing ``active_*`` / ``pending_*`` fields below # preserve the exact value so delete guards still match # against the snapshot path. payload: dict[str, Any] = { "is_loaded": self._pipe is not None, "is_loading": self._loading, "repo_id": _display_repo_id(pending_repo or active_repo), "family": ui_family, "pipeline_class": ui_pipeline_class, "base_repo": _display_repo_id(pending_base or active_base), "gguf_filename": ui_gguf_basename, "device": self._device, "dtype": self._dtype, "loaded_at": self._loaded_at, "last_error": self._last_error, "supported_families": supported_families(), } if include_internal: # Guard-facing fields: every repo / path / GGUF # filename the backend owns RIGHT NOW. Delete routes # iterate both, paired so the variant-filename check # is compared against the SAME repo that owns it # (round 13 P1 #3-5). Round 16 P1 #5: never returned # by the public /images/status route. payload.update( { "active_repo_id": active_repo, "active_base_repo": active_base, "active_gguf_filename": active_gguf, "pending_repo_id": pending_repo, "pending_base_repo": pending_base, "pending_gguf_filename": pending_gguf, } ) return payload def _pick_device_and_dtype(self) -> tuple[str, "Any"]: """Pick (device, dtype) for the current host. CUDA-first because that is the only path our diffusion GGUFs are validated on. On macOS we use MPS in float16 to keep the pipeline on the Metal GPU. CPU is allowed only as a last resort because running FLUX on CPU is unusably slow (> 10 minutes per image). BF16 is gated on ``torch.cuda.is_bf16_supported`` because the Pascal / Turing class (sm_60 / sm_70 / sm_75) reports ``is_available() == True`` but lacks BF16 ALUs; FLUX kernels then fail inside ``from_pretrained`` or at the first denoise step. Those cards still work on FP16, so fall back rather than refuse to load. """ import torch if torch.cuda.is_available(): bf16_ok = False try: bf16_ok = bool(torch.cuda.is_bf16_supported()) except Exception: bf16_ok = False return "cuda", torch.bfloat16 if bf16_ok else torch.float16 if ( hasattr(torch, "backends") and getattr(torch.backends, "mps", None) and torch.backends.mps.is_available() ): return "mps", torch.float16 return "cpu", torch.float32 def load_model( self, repo_id: str, *, gguf_filename: Optional[str] = None, base_repo: Optional[str] = None, hf_token: Optional[str] = None, family_override: Optional[str] = None, enable_model_cpu_offload: bool = True, ) -> dict[str, Any]: """Load a diffusion model. ``repo_id`` is the Hugging Face repo id of either a GGUF-only repo (e.g. ``unsloth/FLUX.2-klein-4B-GGUF``) or a full diffusers repo (e.g. ``black-forest-labs/FLUX.2-klein``). When the repo contains a GGUF, ``gguf_filename`` picks which quant to load; otherwise diffusers' standard config-driven load runs. ``base_repo`` overrides the auto-detected diffusers base used for VAE / text encoders. ``family_override`` short-circuits the substring matcher when an exotic repo name confuses it. Raises ``RuntimeError`` on failure with a user-facing message. On a failed swap the previous pipeline is also released to keep peak VRAM bounded; status() reports is_loaded=false with last_error set so the caller can react. """ # Surface a friendly load error when the no-torch / partial # install path is active: the user clicked Load on the Images # page but the runtime never installed torch + diffusers (round # 13 P2 #12). Without this wrapper the import surfaces as a # raw ``ModuleNotFoundError`` -> 500 instead of a 400 the UI # can display. try: from huggingface_hub import hf_hub_download import diffusers import torch except ModuleNotFoundError as exc: missing = exc.name or str(exc) raise RuntimeError( "Diffusion image generation requires the torch / diffusers " f"runtime. Missing dependency: {missing}. Install the Studio " "torch runtime (re-run setup.sh / install.ps1) before " "loading an image model." ) from exc # Round 30 P1 #11: also preflight transformers BEFORE any # destructive unload. Diffusers can expose stub pipeline # classes when transformers is missing or broken, so the load # would otherwise tear down chat first and fail later inside # from_pretrained. Use find_spec (no module execution) so test # environments that stub these modules still pass the preflight # without us actually importing them. # Round 34: accelerate is pulled in transitively by every # supported transformers install path (it is a hard runtime # dep of transformers' PyTorch backend), so a separate # find_spec("accelerate") guard is redundant in practice and # broke the CI test matrix where the test env ships # transformers without accelerate. The offload code path # (``enable_model_cpu_offload`` / ``device_map="auto"``) # will surface a clean ModuleNotFoundError if a user somehow # arrives at an offload-needed load without it. import importlib.util as _ilu if _ilu.find_spec("transformers") is None: raise RuntimeError( "Diffusion image generation requires the Studio torch " "runtime. Missing dependency: transformers. Install the " "Studio torch runtime (re-run setup.sh / install.ps1) " "before loading an image model." ) fam = detect_family(repo_id, override_family = family_override) if fam is None: # Round 22 P2 #4: route the repo label through # ``_display_repo_id`` so a local absolute path that did # not match any family does not leak the operator's # filesystem layout via the error message / last_error # / 400 response body. raise RuntimeError( f"Could not infer a diffusion family for '{_display_repo_id(repo_id)}'. " "Pass family_override = 'flux.2-klein' / 'flux.2' / " "'flux.1' / 'qwen-image' / 'stable-diffusion-3' / " "'stable-diffusion-xl' to disambiguate." ) device, dtype = self._pick_device_and_dtype() # Round 32 P1 #3: track whether the backend-side # helper-busy check published a "diffusion-backend" pending # entry so the outer finally clears the matching publish # exactly once. Set inside the try below right after the # snapshot succeeds. backend_pending_published = False # _load_lock serialises the entire load so two concurrent calls # cannot both kick off a multi-GB download + GPU upload at once. # The second caller waits behind the first and then loads on top # of the now-populated state via the normal swap path. # _generate_lock is also taken so we do not start swapping the # pipeline (release old + allocate new) while a previous # generation is still iterating denoising steps; releasing the # pipe out from under an in-flight forward corrupts scheduler # state. Order: _load_lock -> _generate_lock -> _lock so a # forward (which only takes _generate_lock + briefly _lock) # cannot block a queued load forever. with self._load_lock, self._generate_lock: with self._lock: self._loading = True self._last_error = None # Publish the pending target so cache / finetuned # delete guards can see what is mid-download even # before _repo_id / _base_repo are populated on # success. self._pending_repo_id = repo_id self._pending_base_repo = base_repo # Store the caller's full ``gguf_filename`` (e.g. # ``BF16/model.gguf``) so the variant-aware delete # guards have the subdirectory info. The UI side of # status() still collapses to the basename for display. self._pending_gguf_filename = gguf_filename if gguf_filename else None try: pipeline_cls = getattr(diffusers, fam.pipeline_class, None) if pipeline_cls is None: raise RuntimeError( f"diffusers {diffusers.__version__} has no " f"{fam.pipeline_class}; upgrade diffusers and retry." ) transformer_cls = ( getattr(diffusers, fam.transformer_class, None) if fam.transformer_class else None ) # Resolution rules for the "what repo to call # from_pretrained on" question: # 1. no GGUF file -> caller is loading a full # diffusers repo; use repo_id directly so we do # not silently substitute the family default # AND ignore any base_repo input (it is only # meaningful as a GGUF companion override). The # old order let ``base_repo`` swap a fine-tuned # ``owner/my-flux.1-finetune`` for # ``black-forest-labs/FLUX.1-dev`` while status # still advertised the user's repo (round 13 # P2 #10). # 2. otherwise prefer caller-supplied base_repo for # the missing VAE / text encoder components. # 3. otherwise use the family + repo_id heuristic so # a 9B GGUF picks the 9B base, not the 4B fallback. if not gguf_filename: # Guard: a repo that ends in "-GGUF" (the unsloth # convention) is GGUF-only and will 500 on # from_pretrained; surface a clear error instead of # letting diffusers raise a confusing model-index # failure deep in the loader. if repo_id.lower().endswith("-gguf"): raise RuntimeError( f"'{repo_id}' looks like a GGUF-only repo. " "Either provide gguf_filename to pick a quant, " "or load a full diffusers repo (base_repo only " "applies when picking a GGUF quant)." ) # ``~/models/my-flux`` must be expanded so # diffusers' from_pretrained does not pass the # literal tilde through to ``os.path.isdir`` and # fall back to the Hub (round 14 P2 #11). effective_base = _expand_existing_local_path(repo_id) with self._lock: self._pending_base_repo = effective_base elif base_repo: effective_base = _expand_existing_local_path(base_repo) # Refresh pending so delete guards see the actual # base, not just caller-supplied None. with self._lock: self._pending_base_repo = effective_base else: effective_base = _smart_base_repo(fam, repo_id) with self._lock: self._pending_base_repo = effective_base # ``repo_id`` / ``effective_base`` are user-supplied # strings that can embed an ``hf_xxxxx`` token via a # URL-style path (``https://hf_token@huggingface.co/...``). # Scrub them BEFORE the logger formats the line so the # token never reaches structured-log sinks (round 14 # P2 #9). # Round 23 P2 #11: ``_redact_hf_tokens`` only scrubs # ``hf_xxxxx`` substrings, so an absolute local # path like ``/home/alice/private/FLUX.2-klein-GGUF`` # used to land in this log line verbatim. Route # through ``_display_repo_id`` so the leaf is # logged when the value is a filesystem path, with # the token-redaction step inside that helper as a # belt-and-braces defence. logger.info( "Loading diffusion model %s (family=%s, device=%s, dtype=%s, base=%s)", _display_repo_id(repo_id), fam.name, device, dtype, _display_repo_id(effective_base), ) transformer = None local_gguf_path: Optional[str] = None if gguf_filename: if transformer_cls is None: raise RuntimeError( f"Family {fam.name} does not have a GGUF transformer " "path wired in this build; load the full repo instead." ) # DiffusionLoadRequest.repo_id is documented to # accept either a Hub repo id OR a local path # (Studio export, downloaded HF snapshot, etc.). # We accept BOTH absolute and relative local # directories: Studio exports surface as relative # paths like ``exports/my-flux`` and earlier # versions only accepted absolute paths, falling # through to ``hf_hub_download`` which then # raised HFValidationError on the relative path # (round 13 P1 #2). For local paths we route the # gguf_filename through ``_resolve_local_gguf_child`` # so traversal (``../secret.gguf``) and absolute # filename escapes (``/etc/passwd``) are rejected # BEFORE the file is opened, which also keeps the # delete-ownership guards aligned with what was # actually loaded. repo_id_path = Path(repo_id).expanduser() if repo_id_path.is_dir(): local_gguf_path = str( _resolve_local_gguf_child(repo_id_path, gguf_filename) ) else: local_gguf_path = hf_hub_download( repo_id = repo_id, filename = gguf_filename, token = hf_token, ) # Round 20 P1 #1: every load mode (full diffusers # repo, GGUF + explicit base_repo, GGUF + auto-picked # base_repo) feeds ``effective_base`` into # ``from_pretrained`` further down. The round 19 # preflight only ran for the first two, so an # auto-picked GGUF companion that turned out to be # gated / private / missing still unloaded chat # before the load failed. Always preflight # ``effective_base`` so a bad companion repo is # caught BEFORE chat / export are released. _preflight_full_diffusers_repo(effective_base, hf_token) # Round 21 P2 #6: the GGUF transformer path also # consumes ``effective_base`` via # ``from_single_file(config=effective_base, # subfolder="transformer")``. A base that has # ``model_index.json`` but lacks # ``transformer/config.json`` would pass the # round-19 preflight and only fail AFTER the chat # unload. Run the subfolder probe too so the # second cheap failure mode is also caught early. if gguf_filename and fam.transformer_class: _preflight_diffusers_subfolder_config( effective_base, "transformer", hf_token, ) # Round 20 P1 #2: ``diffusers.GGUFQuantizationConfig`` # imports the ``gguf`` package lazily at construction # time. Partial Studio installs (``diffusers`` present, # ``gguf`` not) used to discover that AFTER the chat / # export release calls. Build the quant config up # front so the missing-dependency surface raises # while the user's chat model is still resident. quant_config = None if gguf_filename: try: quant_config = diffusers.GGUFQuantizationConfig( compute_dtype = dtype ) except ModuleNotFoundError as exc: missing = exc.name or str(exc) raise RuntimeError( "Diffusion GGUF loading requires the gguf " "runtime package. Missing dependency: " f"{missing}. Re-run Studio setup before " "loading an image GGUF." ) from exc # All cheap failure points (bad gguf_filename, missing # pipeline / transformer class, gated download token, # transient Hub error on the GGUF download) have now # been validated. Anything past this line allocates # GPU memory, so: # 1. Verify training is idle and the export job (if # any) is also idle. ``_release_other_gpu_owners # _for_diffusion`` RAISES on conflict, so it must # run BEFORE we unload chat (round 16 P1 #2): a # route precheck -> worker race could otherwise # drop the user's chat model only to bail out # because training started in between, and a # direct ``DiffusionBackend.load_model`` caller # that did not run the route prechecks would also # leave chat unloaded for nothing. # 2. Release the chat backend (llama-server + the # safetensors orchestrator) now that we know the # load can actually proceed. # 3. Release any *previous* diffusion pipeline so the # new transformer / new from_pretrained does not # race the old pipe for VRAM. Switching between # FLUX.2 klein 4B and 9B on a 16-24 GB GPU OOMs # otherwise: from_single_file allocates the new # transformer while the old pipeline still owns # its weights. # 4. THEN call from_single_file / from_pretrained. # Round 29 P1 #1: do ALL cheap conflict checks BEFORE # any destructive unload, so a training/export conflict # caught inside _release_other_gpu_owners_for_diffusion # does NOT leave the user with no chat model after we # already unloaded it. The helper-busy check is # split out of _release_chat_backend_for_diffusion; # _release_other_gpu_owners_for_diffusion raises # RuntimeError early when training/export is active # without touching the chat backend. # Round 32 P1 #3: publish a backend-side pending # entry under the helper-advisor start lock so a # direct / test / future caller of this method is # symmetric with the route layer's # _raise_if_helper_advisor_busy("diffusion"). The # route's "diffusion" tag and this "diffusion- # backend" tag refcount independently; both # contribute to public_load_pending(). backend_pending_published = _raise_if_helper_advisor_busy_for_diffusion( publish_pending = True, ) _release_other_gpu_owners_for_diffusion() _release_chat_backend_for_diffusion(check_helper_advisor = False) old = self._pipe if old is not None: with self._lock: # Clear ALL metadata together so a failed swap # cannot leave status() reporting the previous # repo / family / base_repo on top of an empty # pipe. The except block below will restore # last_error so the caller knows what happened. self._pipe = None self._family = None self._repo_id = None self._gguf_path = None self._gguf_filename = None self._base_repo = None self._device = None self._dtype = None self._cpu_offload_enabled = False self._loaded_at = None _release(old) old = None # Now that both the attribute and the local # have been nulled, the pipeline is unreachable; # ask the CUDA allocator to release its slabs so # the next from_pretrained does not OOM behind # an already-freed-but-cached arena. _drain_cuda_cache() if gguf_filename: # ``quant_config`` was already constructed above # (round 20 P1 #2 pre-release fail-fast). # Diffusers-format GGUFs (FLUX.2 klein / Qwen-Image / # SD3) need the matching base repo's component config # at config=, subfolder="transformer". # Older city96-style GGUFs ignore those kwargs. The # token is also passed because gated GGUF repos # require it both at download and at config read time. single_file_kwargs: dict[str, Any] = { "quantization_config": quant_config, "torch_dtype": dtype, "config": effective_base, "subfolder": "transformer", } if hf_token: single_file_kwargs["token"] = hf_token transformer = transformer_cls.from_single_file( local_gguf_path, **single_file_kwargs, ) pipe_kwargs: dict[str, Any] = { "torch_dtype": dtype, # use_safetensors=True refuses pickle-backed .bin # weights at load time. Diffusers will fall back to # safetensors variants on repos that publish both, # and hard-error on repos that only ship .bin (which # is the threat model we want to block since pickle # files can execute arbitrary code in this process). "use_safetensors": True, } if transformer is not None: pipe_kwargs["transformer"] = transformer if hf_token: pipe_kwargs["token"] = hf_token pipe = None cpu_offload_enabled = bool( enable_model_cpu_offload and device == "cuda" ) try: pipe = pipeline_cls.from_pretrained(effective_base, **pipe_kwargs) # Device placement / offload can ALSO raise after # from_pretrained succeeded (OOM at the .to(device) # copy, accelerate offload hook misconfigured, etc.). # If we let the exception escape now, the local # ``pipe`` lives on the traceback frame until the # caller drops it, holding multi-GB of VRAM behind # the next load attempt. Explicitly release both # pipe and transformer in the same try (round 13 # P2 #11). if cpu_offload_enabled: pipe.enable_model_cpu_offload() else: pipe.to(device) except Exception: if pipe is not None: _release(pipe) pipe = None if transformer is not None: _release(transformer) transformer = None _drain_cuda_cache() raise with self._lock: self._pipe = pipe self._family = fam self._repo_id = repo_id self._gguf_path = local_gguf_path # Preserve the full caller-supplied filename, not # just the basename, so per-variant delete guards # see ``BF16/model.gguf`` (round 14 P1 #4). self._gguf_filename = gguf_filename if gguf_filename else None self._base_repo = effective_base self._device = device self._dtype = str(dtype).replace("torch.", "") self._cpu_offload_enabled = cpu_offload_enabled self._loaded_at = time.time() # Clear loading + pending here, BEFORE returning, # so the response payload reports the resident # pipeline cleanly (is_loading=false, no pending_*). # The ``finally`` block below is idempotent and # still clears on error / early raise paths. self._loading = False self._pending_repo_id = None self._pending_base_repo = None self._pending_gguf_filename = None return self.status() except Exception as exc: # Scrub hf_token and pipe_kwargs from frame locals BEFORE # logger.exception() captures them. Rich tracebacks and # some structlog formatters render frame locals, which # would otherwise echo the raw hf_... token into logs # and any error reporting sink the user has wired up. # ALSO scrub the exception message itself: huggingface_hub # / diffusers can include the bearer token verbatim in # 401 / 403 messages, which would propagate through # ``_last_error`` (rendered in status()) and the # user-facing RuntimeError (rendered in route responses). scrub_token = hf_token hf_token = None # noqa: F841 pipe_kwargs = None # noqa: F841 single_file_kwargs = None # noqa: F841 exc_msg = str(exc) if scrub_token: exc_msg = exc_msg.replace(scrub_token, "") # Hugging Face tokens are prefixed ``hf_``; replace any # leftover ``hf_...`` substrings to catch tokens we did # not store as ``scrub_token`` (e.g. cached tokens that # huggingface_hub picked up on its own). import re exc_msg = re.sub(r"hf_[A-Za-z0-9]{20,}", "", exc_msg) # Round 17 P2 #9: diffusers / safetensors raise errors # like ``FileNotFoundError: /home/alice/models/foo.gguf`` # or ``OSError: Error while loading state dict from # C:\\Users\\bob\\repos\\flux``. These messages flow # into ``_last_error`` (rendered by status() to every # authenticated browser tab) and the user-facing # RuntimeError, which would leak the operator's # filesystem layout to other sessions. Collapse the # known repo / base / gguf paths to their leaf name # using the same convention as _display_repo_id(). def _collapse_local(msg: str, candidate: Optional[str]) -> str: if not candidate or not isinstance(candidate, str): return msg try: p = Path(candidate).expanduser() except (OSError, ValueError): return msg leaf = p.name or candidate needles: set[str] = set() # Round 20 P2 #6: a relative candidate like # ``exports/my-flux`` used to collapse only the # exact ``exports/my-flux`` substring, but # downstream libraries (diffusers / safetensors) # resolve and emit ``/mnt/disks/.../exports/my-flux/...`` # absolute strings that leaked the operator's # filesystem layout. Also scrub the resolved # absolute form so the leaf is the only path # fragment that survives. try: if p.exists(): needles.add(str(p.resolve())) elif p.is_absolute(): needles.add(str(p)) except (OSError, ValueError): pass if "/" in candidate or "\\" in candidate: needles.add(candidate) # Replace longest first so a parent-directory # substring does not blank out the leaf-only # context the user needs. for needle in sorted( (n for n in needles if n and n != leaf), key = len, reverse = True, ): msg = msg.replace(needle, leaf) return msg # ``effective_base`` and ``gguf_filename`` are local # to the try block above and may be unbound if the # exception fired before assignment (e.g. the GGUF # repo / filename validation raises before # ``effective_base`` is computed). ``locals().get`` # keeps the scrub a no-op in that case. # Round 18 P2 #9: also scrub ``local_gguf_path``. The # GGUF quant is loaded via # ``transformer_cls.from_single_file(local_gguf_path)``, # and diffusers / safetensors errors include the # resolved absolute HF cache path # (``/home/alice/.cache/huggingface/hub/.../flux.gguf``). # Without this the cache path would leak into # ``_last_error`` (and therefore status() / log lines). _locals = locals() exc_msg = _collapse_local(exc_msg, repo_id) exc_msg = _collapse_local(exc_msg, _locals.get("effective_base")) exc_msg = _collapse_local(exc_msg, _locals.get("gguf_filename")) exc_msg = _collapse_local(exc_msg, _locals.get("local_gguf_path")) with self._lock: self._last_error = exc_msg # ``logger.exception`` would emit the raw exception # (including any unredacted ``hf_...`` token inside # the message OR traceback locals on rich loggers). # Use ``logger.error`` with the already-scrubbed # message and exc_info=False so the bearer token # cannot leak through structured logging sinks. # Round 23 P2 #12: same fix as the start-of-load # log above. ``_redact_hf_tokens`` alone left # absolute local repo paths in this failure line. logger.error( "Diffusion load failed for %s: %s", _display_repo_id(repo_id), exc_msg, ) raise RuntimeError( f"Failed to load diffusion model: {exc_msg}" ) from exc finally: with self._lock: self._loading = False # Clear pending so status() falls back to publishing # the resident pipeline (or nothing, on a failed # swap). Keeping pending alive after the load # finishes would falsely block deletes forever. self._pending_repo_id = None self._pending_base_repo = None self._pending_gguf_filename = None # Round 32 P1 #3: clear the backend-side public-load # pending publish if it was set. Skipped when the # helper-busy snapshot raised (no publish to clear) # so the counter stays in sync with publishes. if backend_pending_published: _clear_diffusion_backend_pending() def unload_model(self) -> dict[str, Any]: # Take the load lock and the generate lock so unload cannot: # * race with an in-flight load_model and have the load # thread overwrite the cleared state after we already # returned {"is_loaded": false}. # * return is_loaded=false while a forward pass is still # iterating denoising steps on the soon-to-be-freed pipe. # The generate forward only holds _generate_lock (briefly # _lock), so acquiring _generate_lock here blocks until any # in-flight generation completes. with self._load_lock, self._generate_lock: with self._lock: old = self._pipe self._pipe = None self._family = None self._repo_id = None self._gguf_path = None self._gguf_filename = None self._base_repo = None self._device = None self._dtype = None self._cpu_offload_enabled = False self._loaded_at = None _release(old) old = None # noqa: F841 _drain_cuda_cache() return {"is_loaded": False} # ── generation ──────────────────────────────────────────────── def generate_image( self, *, prompt: str, negative_prompt: Optional[str] = None, num_inference_steps: int = 24, guidance_scale: float = 3.5, width: int = 1024, height: int = 1024, seed: Optional[int] = None, ) -> "Any": """Generate a single PIL image and return it. Concurrent generations are serialised by ``_generate_lock`` so diffusion pipelines (not thread-safe; overlapping ``__call__``s corrupt internal scheduler state) only ever run one at a time. The state ``_lock`` is taken only to snapshot ``_pipe`` / ``_device`` and immediately released: holding it for the whole forward pass blocked ``status()`` polls and concurrent unload requests for the entire (minutes-long) generation, which made the UI feel frozen. """ # Take _generate_lock FIRST so a concurrent unload/load that # observes us holding it will queue behind this generation # (and `unload_model` then waits its turn before clearing # state). Snapshotting `self._pipe` outside the lock and then # taking the lock let a load/unload race in between, so the # forward could run against a freed or swapped pipeline. with self._generate_lock: return self._generate_image_unlocked( prompt = prompt, negative_prompt = negative_prompt, num_inference_steps = num_inference_steps, guidance_scale = guidance_scale, width = width, height = height, seed = seed, ) def _generate_image_unlocked( self, *, prompt: str, negative_prompt: Optional[str] = None, num_inference_steps: int = 24, guidance_scale: float = 3.5, width: int = 1024, height: int = 1024, seed: Optional[int] = None, ) -> "Any": """Inner body of ``generate_image`` that ASSUMES the caller already holds ``_generate_lock``. Lets ``generate_image_with_metadata`` snapshot metadata under the same lock without deadlocking on a non-reentrant ``threading.Lock`` (round 13 P2 #9).""" if not prompt or not prompt.strip(): raise ValueError("prompt is empty") if num_inference_steps < 1 or num_inference_steps > 200: raise ValueError("num_inference_steps must be in [1, 200]") if width <= 0 or height <= 0 or width > 2048 or height > 2048: raise ValueError("width and height must be in (0, 2048]") # Snap to a multiple of 8: Flux / SD pipelines require it and a # silent crash deep in the VAE is much worse than a clear error # message up front. if width % 8 or height % 8: raise ValueError("width and height must be multiples of 8") import torch with self._lock: if self._pipe is None: raise RuntimeError("No diffusion model is loaded.") pipe = self._pipe device = self._device or "cpu" cpu_offload_enabled = self._cpu_offload_enabled generator = None if seed is not None: # Match the device of the pipeline so determinism holds # across reload cycles. When CPU offload is enabled # (the default on CUDA hosts), diffusers shuttles each # submodule between CPU and GPU on every step. A CUDA # torch.Generator then mismatches the CPU-resident # embeddings at the start of the forward and the run # crashes (round 14 P1 #6). Use a CPU generator in that # case; numerical determinism for the same seed is # preserved because the seed feeds an int rather than a # device-local RNG state. if cpu_offload_enabled: gen_device = "cpu" else: gen_device = ( "cuda" if device == "cuda" and torch.cuda.is_available() else "cpu" ) generator = torch.Generator(device = gen_device).manual_seed(int(seed)) call_kwargs: dict[str, Any] = { "prompt": prompt, "num_inference_steps": int(num_inference_steps), "guidance_scale": float(guidance_scale), "width": int(width), "height": int(height), } # FLUX.2 / FLUX.2 klein pipelines do NOT accept # negative_prompt and 500 if you pass it in. Inspect the # signature and only forward when supported; warn otherwise # so the UI can disable the field for incompatible families. if negative_prompt is not None and negative_prompt.strip(): if _pipe_accepts_kwarg(pipe, "negative_prompt"): call_kwargs["negative_prompt"] = negative_prompt # QwenImagePipeline and FluxPipeline treat # guidance_scale as distilled CFG and use # true_cfg_scale as the real classifier-free # guidance knob; the negative prompt is only # effective when true_cfg_scale > 1. Forward the # user-supplied guidance_scale through both so the # negative prompt actually steers generation. if _pipe_accepts_kwarg(pipe, "true_cfg_scale"): call_kwargs["true_cfg_scale"] = float(guidance_scale) else: logger.info( "Dropping negative_prompt: %s does not accept it", type(pipe).__name__, ) if generator is not None: call_kwargs["generator"] = generator out = pipe(**call_kwargs) images = getattr(out, "images", None) or [] if not images: raise RuntimeError("Diffusion pipeline returned no images.") return images[0] def generate_image_with_metadata( self, **kwargs: Any, ) -> tuple[Any, dict[str, Any]]: """Generate a single image AND snapshot its identifying metadata. Returns ``(pil_image, {"model": , "family": })`` where the metadata reflects the pipeline that produced the image. Snapshotted under ``_generate_lock + _lock`` so a queued unload / load that promotes a different pipeline cannot replace ``self._repo_id`` / ``self._family`` between the forward returning and the route reading status (round 13 P2 #9). The route uses these values directly in the response instead of re-calling ``status()``. """ with self._generate_lock: image = self._generate_image_unlocked(**kwargs) with self._lock: # Round 16 P1 #6: route ``model`` through # _display_repo_id so a generation response for a # locally-loaded model cannot echo back an absolute # filesystem path to the browser. meta = { "model": _display_repo_id(self._repo_id), "family": self._family.name if self._family else None, } return image, meta def _pipe_accepts_kwarg(pipe: Any, name: str) -> bool: """True if ``pipe.__call__`` advertises a kwarg called ``name``. Cheap inspect-based probe so we do not have to maintain a manual list of which pipeline classes accept negative_prompt. Returns False on any introspection error so callers stay on the safe path. """ import inspect try: sig = inspect.signature(pipe.__call__) except (TypeError, ValueError): return False if name in sig.parameters: return True return any(p.kind is inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()) def encode_png_base64(pil_image: "Any") -> str: """Encode a PIL image to base64-encoded PNG.""" import base64 buf = io.BytesIO() pil_image.save(buf, format = "PNG", optimize = True) return base64.b64encode(buf.getvalue()).decode("ascii") # ─── Helpers ────────────────────────────────────────────────────────── def _raise_if_helper_advisor_busy_for_diffusion( *, publish_pending: bool = False, ) -> bool: """Round 29 P1 #1: split the helper-busy check out of _release_chat_backend_for_diffusion so the diffusion load can check ALL conflicts (helper, training, export) BEFORE doing ANY destructive unloads. Otherwise a route-precheck race or a direct backend call would unload the user's chat while training was active, then 409 with the user holding no model at all. Round 32 P1 #3: when ``publish_pending=True`` also takes ``_HELPER_ADVISOR_START_LOCK`` and publishes a ``diffusion-backend`` public-load pending entry so a concurrent AI Assist helper / advisor start that wins the start lock sees the pending public owner and refuses VRAM. The route layer publishes its own ``diffusion`` tag (refcount semantics, so the two publishes coexist without erasing each other). Returns True when a pending entry was actually published so the caller can pair it with ``_clear_diffusion_backend_pending`` in finally. Direct callers (tests, scripts) opt in with ``publish_pending= True`` to get the same atomic check + publish the route gets. The ``check_helper_advisor`` callback in ``_release_chat_backend_for_diffusion`` keeps the default False so legacy callers do not double-publish or leak pending entries. """ try: from utils.datasets.llm_assist import ( _HELPER_ADVISOR_START_LOCK, _publish_public_load_pending, helper_advisor_busy, ) except Exception: return False with _HELPER_ADVISOR_START_LOCK: if helper_advisor_busy(): raise RuntimeError( "AI Assist (helper / advisor GGUF) is still using the GPU. " "Wait for it to finish before loading a diffusion image model." ) if publish_pending: _publish_public_load_pending("diffusion-backend") return True return False def _clear_diffusion_backend_pending() -> None: """Round 32 P1 #3: paired clear for ``_raise_if_helper_advisor_busy_for_diffusion(publish_pending=True)``. Safe to call when the helpers module is unavailable (no-op).""" try: from utils.datasets.llm_assist import _release_public_load_pending except Exception: return try: _release_public_load_pending("diffusion-backend") except Exception: pass def _release_chat_backend_for_diffusion(*, check_helper_advisor: bool = True) -> None: """Unload any running chat backend before a diffusion load. Diffusion pipelines on FLUX-class models can eat 12-24 GB of VRAM, and the chat backends (llama-server for GGUF, the safetensors Inference orchestrator for HF / Unsloth) typically hold onto their loaded weights until told to drop them. Asking both to release their weights first means a typical 24 GB consumer GPU can host one chat model OR one diffusion model without manual unload steps. A missing chat backend module is a silent no-op (fresh install / no GGUF use). An unload that ACTUALLY fails (raises or leaves the backend resident) raises ``RuntimeError`` so the surrounding diffusion ``load_model`` bails out instead of double-owning VRAM (round 17 P1 #2). """ # Round 27 P1 #2 / round 29 P1 #1: helper / advisor GGUF loads # run on a PRIVATE LlamaCppBackend so the global llama check below # cannot see them. The actual busy check now lives in # _raise_if_helper_advisor_busy_for_diffusion so the caller can do # ALL conflict checks BEFORE any destructive unload. Kept here as # a default-on safety net for callers that did not run the # standalone check. if check_helper_advisor: _raise_if_helper_advisor_busy_for_diffusion() # 1. GGUF chat backend (llama-server subprocess). We unload when # EITHER is_loaded is True (resident model) OR is_active is # True (mid-download / startup) OR loading_model_identifier is # populated (HF GGUF download in progress, before is_active / # is_loaded flip). The last case is what round 13 P1 #8 flagged. try: from routes.inference import get_llama_cpp_backend # type: ignore except Exception as exc: logger.debug("llama-server unavailable before diffusion load: %s", exc) else: backend = get_llama_cpp_backend() is_loaded = bool(getattr(backend, "is_loaded", False)) is_active = bool(getattr(backend, "is_active", False)) is_loading = bool(getattr(backend, "loading_model_identifier", None)) if is_loaded or is_active or is_loading: logger.info( "Unloading llama-server (loaded=%s active=%s loading=%s) before diffusion load", is_loaded, is_active, is_loading, ) try: ok = backend.unload_model() except Exception as exc: raise RuntimeError( "Could not unload the existing GGUF chat model before " "loading a diffusion image model." ) from exc # Round 28 P1 #12: a cancelled pending GGUF download takes # up to a few seconds to clear loading_model_identifier in # its finally block. Wait briefly so the same retryable # cancel path used by the unload route does not 503 us. deadline = time.monotonic() + 5.0 while ( getattr(backend, "loading_model_identifier", None) and time.monotonic() < deadline ): time.sleep(0.1) # Round 18 P1 #4: also reject when ``loading_model_identifier`` # is still set after the unload call. Without this, a GGUF # download / startup that was already in flight before the # diffusion handoff (and which never flipped is_active to # True before the unload landed) keeps allocating into VRAM # while diffusion proceeds, double-owning the GPU. if ( ok is False or getattr(backend, "is_loaded", False) or getattr(backend, "is_active", False) or getattr(backend, "loading_model_identifier", None) ): raise RuntimeError( "The existing GGUF chat model is still active or loading " "after unload; retry before loading a diffusion image model." ) # 2. Safetensors / HF chat backend (the InferenceOrchestrator that # serves FastVisionModel / FastLanguageModel weights). When this # backend has a model resident on the same GPU, a diffusion load # will OOM the same way. We also flush any loading_models set so # a chat load that is mid-download cannot race the diffusion # allocation. try: from core.inference import get_inference_backend # type: ignore except Exception as exc: logger.debug("safetensors unavailable before diffusion load: %s", exc) return backend = get_inference_backend() active_model_name = getattr(backend, "active_model_name", None) loading_models = set(getattr(backend, "loading_models", set()) or set()) owned_names = {name for name in ({active_model_name} | loading_models) if name} def _require_unload(model_name: str) -> None: try: ok = backend.unload_model(model_name) except Exception as exc: raise RuntimeError( f"Could not unload safetensors chat model '{model_name}' " "before loading a diffusion image model." ) from exc if ok is False: raise RuntimeError( f"Safetensors backend refused to unload '{model_name}' " "before loading a diffusion image model." ) # Round 19 P1 #2: per-name post-state check. ``unload_model`` # returning ``True`` does not guarantee the orchestrator # actually dropped the weights; the worker may have responded # while still holding them, or a concurrent ``load`` may have # repopulated the tracker. Verify the specific name is gone # so the surrounding diffusion load bails out instead of # silently double-owning VRAM. active_after = getattr(backend, "active_model_name", None) loading_after = set(getattr(backend, "loading_models", set()) or set()) if active_after == model_name or model_name in loading_after: raise RuntimeError( f"Safetensors chat model '{model_name}' is still active " "or loading after unload; retry before loading a diffusion image model." ) if active_model_name: logger.info( "Unloading safetensors chat backend '%s' before diffusion load", active_model_name, ) _require_unload(active_model_name) for loading in loading_models: if loading == active_model_name: continue logger.info( "Unloading in-flight safetensors chat load '%s' before diffusion", loading, ) _require_unload(loading) # Round 21 P1 #5: final sweep without the owned_names filter. # A concurrent ``/load`` that appeared AFTER the initial # snapshot was previously ignored, so a chat model that started # loading during the diffusion handoff slipped through and # raced the diffusion allocation for VRAM. Treat ANY surviving # active / loading entry as a failure so the surrounding # load_model raises and the caller retries. remaining_loading = set(getattr(backend, "loading_models", set()) or set()) remaining_active = getattr(backend, "active_model_name", None) if remaining_loading or remaining_active: raise RuntimeError( "A safetensors chat model is still active or loading " "after unload; retry before loading a diffusion image model." ) def _release_other_gpu_owners_for_diffusion() -> None: """Best-effort: shut down export subprocess + active training before a diffusion load. Both can hold multi-GB of VRAM and would OOM the diffusion allocation on consumer GPUs.""" # Export resident checkpoint. We tear down a SETTLED export # (current_checkpoint populated AND is_export_active() False) # because that means the export ran to completion and the user # can re-load the result. An in-flight export job # (is_export_active() True) is NEVER touched here: terminating # it would corrupt the user's partial output artifact. # # The route layer also rejects /images/load with HTTP 409 via # _raise_if_export_active when is_export_active() is True. This # helper repeats the local check anyway so that direct backend # callers (tests, scripts, future routes that forget the # higher-level guard) cannot still kill an active export. # Training-active check runs FIRST so direct backend callers # (tests, scripts, future routes) cannot bypass the route layer's # 409 by calling ``load_model`` directly while a training run is # active (round 15 P1 #3). The route layer's # ``_raise_if_training_active`` still runs ahead of the load to # surface the conflict as 409; this helper re-raises so direct # callers see the same RuntimeError the export-active path raises. try: from core.training import get_training_backend # type: ignore except Exception as exc: logger.debug("training module not importable: %s", exc) else: try: training_active = bool(get_training_backend().is_training_active()) except Exception as exc: # Unverifiable status -> fail closed (might be active). raise RuntimeError( "Could not verify training status before loading a " "diffusion image model." ) from exc if training_active: raise RuntimeError( "Training is currently active. Stop the training run " "before loading a diffusion image model." ) try: from core.export import get_export_backend # type: ignore except Exception as exc: logger.debug("export module not importable: %s", exc) return # Round 18 P1 #6: ``get_export_backend()`` raising used to be a # silent ``return`` so direct ``DiffusionBackend.load_model`` # callers could proceed toward GPU allocation without being able # to verify export ownership. Fail closed instead, matching the # route-level helper which already maps "Could not verify" / # "Could not access" failures to HTTP 503. try: exp = get_export_backend() except Exception as exc: raise RuntimeError( "Could not verify export status before loading a " "diffusion image model." ) from exc is_export_active_fn = getattr(exp, "is_export_active", None) if is_export_active_fn is not None: try: export_is_active = bool(is_export_active_fn()) except Exception as exc: # Round 16 P2 #8: distinguish unverifiable status from # active export. The previous "treat as active" mapping # surfaced as a misleading 409 conflict; raise a # "Could not verify" RuntimeError so the route layer # maps it to 503 (retryable) instead. raise RuntimeError( "Could not verify export status before loading a " "diffusion image model." ) from exc if export_is_active: # Round 14 P2 #10: the prior behaviour logged a warning # and continued, so direct ``DiffusionBackend.load_model`` # callers (tests, scripts) silently bypassed the route # layer's 409. Hard-refuse instead so any code path that # reaches this helper while an export is active sees the # same failure mode the route returns. raise RuntimeError( "An export job is currently active. Stop the export " "job before loading a diffusion image model." ) if getattr(exp, "current_checkpoint", None): # Round 18 P1 #2: a wedged ``_shutdown_subprocess`` used to log # at debug level and continue, so direct backend callers could # allocate diffusion VRAM on top of an export checkpoint that # still owned the GPU. Mirror the route-level helper and raise # so the surrounding ``load_model`` bails out with a clean # RuntimeError that the route layer maps to HTTP 503. try: logger.info("Shutting down idle export subprocess before diffusion load") exp._shutdown_subprocess() except Exception as exc: raise RuntimeError( "Could not unload the idle export checkpoint before " "loading a diffusion image model." ) from exc exp.current_checkpoint = None exp.is_vision = False exp.is_peft = False # Note: active training is *not* stopped here. The route layer # (`_raise_if_training_active` in routes/inference.py) refuses # /images/load with HTTP 409 before this helper runs, so reaching # this point with training still active would only happen in # programmatic backend calls (tests, scripts). Silently terminating # someone's training run when the diffusion load might still fail # is worse than letting the load OOM and surfacing it explicitly. def _release(obj: Any) -> None: """Best-effort GPU-memory release for a pipeline being swapped out. Only drops the local reference (which the caller has already nulled in its own scope) and runs ``gc.collect()`` so __del__ fires. Does NOT call ``torch.cuda.empty_cache()`` here because when the caller still holds the actual reference in a local / attribute, ``empty_cache()`` would run before __del__ released the weights and would not actually free GPU memory. Use ``_drain_cuda_cache()`` AFTER the last reference has been nulled. """ if obj is None: return try: del obj except Exception: pass gc.collect() def _drain_cuda_cache() -> None: """Hand freed weights back to the active accelerator's allocator. Call this AFTER every reference to the freed object has been dropped (caller's local + attribute) and a ``gc.collect()`` has fired __del__. Calling earlier would empty an already-pinned cache and not actually release the memory. Handles CUDA *and* MPS (Apple Silicon) so a diffusion swap on macOS actually returns VRAM to the Metal allocator. """ try: import torch if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception: pass try: import torch mps_backend = getattr(getattr(torch, "backends", None), "mps", None) if mps_backend is not None and mps_backend.is_available(): mps_module = getattr(torch, "mps", None) empty_cache = ( getattr(mps_module, "empty_cache", None) if mps_module else None ) if empty_cache is not None: empty_cache() except Exception: pass # ─── Module-level singleton ─────────────────────────────────────────── _singleton: Optional[DiffusionBackend] = None _singleton_lock = threading.Lock() def get_diffusion_backend() -> DiffusionBackend: """Return the process-wide diffusion backend (lazy-instantiated).""" global _singleton if _singleton is None: with _singleton_lock: if _singleton is None: _singleton = DiffusionBackend() return _singleton async def async_generate( backend: DiffusionBackend, **kwargs: Any, ) -> "Any": """Run ``generate_image`` in the default executor so route handlers do not block the event loop for the 5-30 s a diffusion step takes.""" loop = asyncio.get_event_loop() return await loop.run_in_executor(None, lambda: backend.generate_image(**kwargs)) async def async_generate_with_metadata( backend: DiffusionBackend, **kwargs: Any, ) -> tuple[Any, dict[str, Any]]: """Run ``generate_image_with_metadata`` in the default executor. Used by the /images/generate route so the response model / family fields reflect the pipeline that actually produced the image, even if an unload races the route between the forward returning and the response being assembled (round 13 P2 #9).""" loop = asyncio.get_event_loop() return await loop.run_in_executor( None, lambda: backend.generate_image_with_metadata(**kwargs), )