# 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 threading import time from dataclasses import dataclass, field from pathlib import Path 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). """ if fam.name != "flux.2-klein": return fam.base_repo lower = (repo_id or "").lower() is_9b = "9b" in lower is_base = "base" in lower 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 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. ``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 for fam in _FAMILIES: if fam.name in needle: return fam for alias in fam.aliases: if alias and alias in needle: 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 the # pipe call inside generate_image. `_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). The # locks are taken in order load -> state so a generation in # flight cannot deadlock the next load. self._lock = threading.Lock() self._load_lock = threading.Lock() self._family: Optional[DiffusionFamily] = None self._repo_id: Optional[str] = None self._gguf_path: Optional[str] = None self._base_repo: Optional[str] = None self._device: Optional[str] = None self._dtype: Optional[str] = None self._loaded_at: Optional[float] = None self._loading: bool = False self._last_error: 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) -> dict[str, Any]: # 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. gguf_basename = Path(self._gguf_path).name if self._gguf_path else None return { "is_loaded": self.is_loaded, "is_loading": self._loading, "repo_id": self._repo_id, "family": self._family.name if self._family else None, "pipeline_class": self._family.pipeline_class if self._family else None, "base_repo": self._base_repo, "gguf_filename": gguf_basename, "device": self._device, "dtype": self._dtype, "loaded_at": self._loaded_at, "last_error": self._last_error, "supported_families": supported_families(), } 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). """ import torch if torch.cuda.is_available(): return "cuda", torch.bfloat16 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; the previous pipeline (if any) stays loaded so a failed swap does not leave Studio in an unusable state. """ from huggingface_hub import hf_hub_download import diffusers import torch fam = detect_family(repo_id, override_family = family_override) if fam is None: raise RuntimeError( f"Could not infer a diffusion family for '{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() # _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. with self._load_lock: with self._lock: self._loading = True self._last_error = None try: # Unload any chat model that is holding GPU memory so the # diffusion load does not OOM on a < 24 GB GPU. Best # effort: if the llama-cpp backend module is absent (eg # tests, headless tooling) we just continue. _release_chat_backend_for_diffusion() 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. caller-supplied base_repo wins # 2. if no GGUF file was requested the user is loading a # full diffusers repo; use repo_id directly so we do # not silently substitute the family default # 3. otherwise use the family + repo_id heuristic so a # 9B GGUF picks the 9B base, not the 4B fallback if base_repo: effective_base = base_repo elif 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 pass base_repo to override the full-repo " "load target." ) effective_base = repo_id else: effective_base = _smart_base_repo(fam, repo_id) logger.info( "Loading diffusion model %s (family=%s, device=%s, dtype=%s, base=%s)", repo_id, fam.name, device, dtype, 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." ) local_gguf_path = hf_hub_download( repo_id = repo_id, filename = gguf_filename, token = hf_token, ) quant_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype) transformer = transformer_cls.from_single_file( local_gguf_path, quantization_config = quant_config, torch_dtype = dtype, ) 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 # Release the previous pipeline BEFORE allocating the # new one so peak VRAM stays at one model's worth, not # two. This matters on 16-24 GB consumer GPUs where the # combined footprint would OOM the from_pretrained call. old = self._pipe if old is not None: with self._lock: self._pipe = None _release(old) old = None pipe = pipeline_cls.from_pretrained(effective_base, **pipe_kwargs) if enable_model_cpu_offload and device == "cuda": pipe.enable_model_cpu_offload() else: pipe.to(device) with self._lock: self._pipe = pipe self._family = fam self._repo_id = repo_id self._gguf_path = local_gguf_path self._base_repo = effective_base self._device = device self._dtype = str(dtype).replace("torch.", "") self._loaded_at = time.time() # ``old`` was released above before the new allocation; # nothing left to free here. return self.status() except Exception as exc: with self._lock: self._last_error = str(exc) logger.exception("Diffusion load failed for %s", repo_id) raise RuntimeError(f"Failed to load diffusion model: {exc}") from exc finally: with self._lock: self._loading = False def unload_model(self) -> dict[str, Any]: # Take the load lock too 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}. with self._load_lock: with self._lock: old = self._pipe self._pipe = None self._family = None self._repo_id = None self._gguf_path = None self._base_repo = None self._device = None self._dtype = None self._loaded_at = None _release(old) 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. The mutex is held for the entire call: diffusion pipelines are not thread-safe, and overlapping ``__call__``s on a shared pipeline frequently corrupt their internal scheduler state. """ 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" generator = None if seed is not None: # Match the device of the pipeline so determinism holds # across reload cycles. For CPU offload, the noise still # has to live on the device the diffusion forward runs on. 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 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 _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 _release_chat_backend_for_diffusion() -> 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. Best effort: if a chat backend module is not importable (CI, isolated tests, custom builds) or fails on the unload, we log and continue; the diffusion load can still try and surface its own OOM. """ # 1. GGUF chat backend (llama-server subprocess). try: from routes.inference import get_llama_cpp_backend # type: ignore backend = get_llama_cpp_backend() if getattr(backend, "is_loaded", False): logger.info("Unloading llama-server before diffusion load") backend.unload_model() except Exception as exc: logger.debug("llama-server unload skipped: %s", exc) # 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. The orchestrator's unload_model takes a # model_name; passing it without args raised TypeError and was # swallowed, leaving the chat model resident. try: from core.inference import get_inference_backend # type: ignore backend = get_inference_backend() active_model_name = getattr(backend, "active_model_name", None) if active_model_name: logger.info( "Unloading safetensors chat backend '%s' before diffusion load", active_model_name, ) backend.unload_model(active_model_name) except Exception as exc: logger.debug("safetensors unload skipped: %s", exc) def _release(obj: Any) -> None: """Best-effort GPU-memory release for a pipeline being swapped out.""" if obj is None: return try: del obj except Exception: pass gc.collect() try: import torch if torch.cuda.is_available(): torch.cuda.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))