ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/backend/core/inference/diffusion_families.py
Daniel Han 6a8b0b47e7 Fix review findings on image generation: failed-load VRAM, API defaults, preflights 
- Free reserved VRAM in the diffusion load worker's failure path: a load-time OOM
  never commits _state and the next load's _unload_locked early-returns, so nothing
  else reclaimed the half-built pipeline's memory
- Use a monotonic clock for the denoise ETA rate
- Sync _GENERATION_DEFAULTS with the UI table: kontext, flux.2-dev, sdxl-turbo and
  SDXL base rows so /v1/images/generations stops falling back to 9 steps / CFG 0
- 400 (not sanitized 500) when /v1/images/generations hits an edit-only model
- Fail fast on pre-Ampere CUDA in the DiT trainer instead of dying in model load
- Run the trainer trust gate in the diffusion training route before freeing GPU
  residents so an untrusted base cannot tear down loaded chat/Images models
- Protect native sd.cpp companion VAE/text-encoder repos from cache deletion while
  a load is downloading them
- Exempt the task-scoped Images picker from the chat-only GGUF/MLX format gate so
  local diffusers pipelines stay selectable on no-GPU hosts
2026-07-05 01:00:47 +00:00

504 lines
27 KiB
Python
Raw Permalink Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Pure helpers for diffusion model identification.
No torch/diffusers imports here: everything in this module is a pure function of
its string/path arguments so it can be unit-tested without the heavy runtime.
A diffusion checkpoint published as a single-file GGUF only carries the
transformer weights; the matching VAE / text encoders / scheduler come from a
companion ``diffusers`` base repo. ``DiffusionFamily`` maps a checkpoint to the
diffusers classes and base repo needed to assemble the full pipeline.
"""
from __future__ import annotations
import re
from dataclasses import dataclass, field
from pathlib import Path, PurePosixPath
from typing import Optional
# Runtime->route contract: the RuntimeError messages a backend raises for
# client-recoverable generate states. The /images/generate route matches these
# EXACTLY to return 409 (vs a sanitized 500 for real failures), so both engines
# must raise them verbatim -- keep them named here, not as scattered literals.
DIFFUSION_NOT_LOADED_MSG = "No diffusion model is loaded."
DIFFUSION_CANCELLED_MSG = "Diffusion generation was cancelled."
@dataclass(frozen = True)
class DiffusionFamily:
name: str
pipeline_class: str
transformer_class: str
base_repo: str
# Pipeline kwarg carrying the guidance value. Most use "guidance_scale";
# Qwen-Image's distilled guidance is off, so its real CFG is "true_cfg_scale".
cfg_kwarg: str = "guidance_scale"
# The pipe attribute holding the denoiser module. DiT families expose it as
# ``pipe.transformer`` (the default); U-Net families (SDXL) as ``pipe.unet``.
# Read wherever the backend touches the denoiser generically (VAE dtype
# alignment, optimisation guards), so a U-Net family works without assuming a
# ``transformer`` attribute exists.
denoiser_attr: str = "transformer"
# True when a single-file ``.safetensors`` checkpoint is the WHOLE pipeline
# (U-Net + VAE + text encoders), not a transformer-only file. SDXL ships this
# way, so the loader calls ``pipeline_class.from_single_file`` on it directly
# rather than ``transformer_class.from_single_file`` + a companion base repo.
# DiT families leave this False (their single file is transformer-only).
single_file_is_pipeline: bool = False
# Optional diffusers pipeline classes for image-conditioned workflows. The backend
# builds these around the ALREADY-loaded transformer/VAE/text-encoder via
# ``Pipeline.from_pipe`` (no extra weights, no reload), so a family only needs the
# class name here to gain the workflow. None = the family does not support it (the
# UI gates the workflow off). The base text-to-image pipeline is ``pipeline_class``.
img2img_pipeline_class: Optional[str] = None
inpaint_pipeline_class: Optional[str] = None
# ControlNet: the diffusers ControlNet pipeline + model classes for this family. The backend
# loads the (small) ControlNet model via from_pretrained and builds the pipeline via
# ``Pipeline.from_pipe(base, controlnet=model)`` around the resident modules (no reload),
# then passes the control image + conditioning scale at generate time. None on both = the
# family has no diffusers ControlNet support and the UI gates the workflow off.
controlnet_pipeline_class: Optional[str] = None
controlnet_model_class: Optional[str] = None
# True when the inpaint pipeline keeps the input canvas size, so it can also drive
# outpaint (extend), where the padded canvas is LARGER than the original. False for
# FLUX.2 (its pipelines scale any >1MP input down to ~1MP, which shrinks an outpaint
# canvas back and defeats the extend). Such families get Inpaint but not Extend.
inpaint_preserves_size: bool = True
# True for instruction-editing families (Qwen-Image-Edit / FLUX Kontext): the model's
# OWN pipeline (``pipeline_class``) is the edit pipeline -- it takes an input image plus
# a text instruction and has no plain text-to-image mode. So these expose only the
# "edit" workflow, require an input image at generate time, and the loaded pipe is used
# directly (no from_pipe). ``base_repo`` here is the matching diffusers repo that
# supplies the VAE / text-encoder / processor / scheduler for the GGUF transformer.
edit: bool = False
# True for families whose OWN text-to-image pipeline ALSO accepts reference image(s)
# (FLUX.2: Flux2KleinPipeline takes an optional ``image`` arg). Unlike ``edit`` these
# families still do plain text-to-image (no image), and unlike img2img the conditioning
# is reference-based, not a denoise blend: there is no ``strength`` and the output size
# comes from the requested width/height, not the reference's size. The loaded pipe is
# used directly (no from_pipe). Exposes a "reference" workflow alongside "txt2img".
reference: bool = False
# Extra lowercased substrings (besides ``name``) that map a repo id here.
aliases: tuple[str, ...] = field(default_factory = tuple)
# True for families whose activations overflow float16's finite range
# (~6.5e4) and produce inf -> NaN latents -> a black image. The backend
# promotes a resolved float16 to float32 for these at load time.
fp16_incompatible: bool = False
# Set False only for a family whose denoiser block does not compile cleanly with
# regional torch.compile. Now consulted on the GGUF path too (compile runs on the
# GGUF transformer); all current families compile, so this stays True.
supports_torch_compile: bool = True
# Optional pre-quantized transformer checkpoints, as (scheme, repo_id) pairs (a
# hashable mapping). When the fast transformer_quant path resolves a scheme with a
# hosted checkpoint, the loader fetches the already-quantized weights instead of
# materialising the dense bf16 transformer on the GPU (much lower load VRAM + a
# smaller download). Empty until checkpoints are hosted -> behaviour is unchanged.
prequant_repos: tuple[tuple[str, str], ...] = field(default_factory = tuple)
# Native (stable-diffusion.cpp) single-file assets, used only when the no-GPU
# sd.cpp engine is selected (CPU / Apple). The transformer GGUF is shared with
# the diffusers path; sd-cli additionally needs a single-file VAE and text
# encoder(s), because the diffusers base repo ships those sharded and sd-cli
# cannot read that layout. Each asset is a hashable (repo_id, filename) the
# backend fetches with hf_hub_download. ``sd_cpp_text_encoders`` carries a
# trailing SdCppModelFiles field name (clip_l / t5xxl / llm / qwen2vl / clip_g)
# so the backend maps each file onto the right sd-cli flag. Empty -> the family
# has no native mapping and the sd.cpp route falls back to diffusers.
sd_cpp_vae: Optional[tuple[str, str]] = None
# VAE latent-format override for sd-cli (--vae-format): "flux2" for the FLUX.2
# autoencoder, None (auto) otherwise.
sd_cpp_vae_format: Optional[str] = None
sd_cpp_text_encoders: tuple[tuple[str, str, str], ...] = field(default_factory = tuple)
# Family-specific sd-cli sampler settings, applied on the native path so the
# output matches the model's supported invocation (e.g. Qwen-Image needs
# --sampling-method euler --flow-shift 3 per stable-diffusion.cpp's docs). None
# leaves sd-cli's defaults (correct for the distilled flux/z-image families).
sd_cpp_sampling_method: Optional[str] = None
sd_cpp_flow_shift: Optional[float] = None
# True when Studio can TRAIN a LoRA on this family (a trainer is registered for it in
# core.training). Loadable-for-inference is the default; training is opt-in per family
# because each architecture needs its own training loop. The diffusion training start
# path resolves the base model's family and refuses a non-trainable one up front.
trainable: bool = False
# Recommended base repos to train FROM, most-preferred first (e.g. a QLoRA-friendly
# prequant repo, then a bf16 repo). Surfaced by the Train UI as the base-model choices.
train_base_repos: tuple[str, ...] = field(default_factory = tuple)
# Keyed by architecture, not per model variant: a checkpoint's specific base repo
# is read from its HF base_model tag at load time, so one entry covers Turbo/full,
# schnell/dev, etc. base_repo here is only a fallback. Only archs whose diffusers
# transformer supports from_single_file load here (ERNIE-Image does not, yet; LTX
# video models are out of scope). FLUX.2-klein-9B shares the klein family (its base
# repo is resolved per-variant), and FLUX.2-dev has its own family below.
_FAMILIES: tuple[DiffusionFamily, ...] = (
DiffusionFamily(
name = "flux.1",
pipeline_class = "FluxPipeline",
transformer_class = "FluxTransformer2DModel",
base_repo = "black-forest-labs/FLUX.1-schnell",
aliases = ("flux1", "flux-1"),
# LoRA training targets the guidance-distilled FLUX.1-dev via the DiT trainer
# (QLoRA nf4). The dev repo is gated on the Hub, so a user HF token is required.
trainable = True,
train_base_repos = ("black-forest-labs/FLUX.1-dev",),
img2img_pipeline_class = "FluxImg2ImgPipeline",
inpaint_pipeline_class = "FluxInpaintPipeline",
controlnet_pipeline_class = "FluxControlNetPipeline",
controlnet_model_class = "FluxControlNetModel",
sd_cpp_vae = ("black-forest-labs/FLUX.1-schnell", "ae.safetensors"),
sd_cpp_text_encoders = (
("comfyanonymous/flux_text_encoders", "clip_l.safetensors", "clip_l"),
("comfyanonymous/flux_text_encoders", "t5xxl_fp16.safetensors", "t5xxl"),
),
),
# FLUX.2-klein is a distinct pipeline (Flux2KleinPipeline) with a Qwen3 text
# encoder, not the Mistral-based Flux2Pipeline; it must precede a generic flux
# match. The Mistral-based Flux2Pipeline is the separate flux.2-dev family below.
DiffusionFamily(
name = "flux.2-klein",
pipeline_class = "Flux2KleinPipeline",
transformer_class = "Flux2Transformer2DModel",
base_repo = "black-forest-labs/FLUX.2-klein-4B",
aliases = ("flux2-klein",),
# Flux2KleinPipeline natively accepts reference image(s) via its `image` arg, so it
# exposes a "reference" workflow on top of plain text-to-image. It has a dedicated
# inpaint pipeline too (no img2img one), so it also gets inpaint + extend (outpaint).
reference = True,
inpaint_pipeline_class = "Flux2KleinInpaintPipeline",
# FLUX.2 scales >1MP inputs down to ~1MP, so outpaint (a larger canvas) can't grow.
inpaint_preserves_size = False,
# FLUX.2 uses a distinct 32-channel autoencoder; sd-cli needs the latent
# format override. The single-file VAE ships in Comfy-Org/flux2-dev (the
# klein-4B repo only has a sharded diffusers VAE). Shares Qwen3-4B with z-image.
sd_cpp_vae = ("Comfy-Org/flux2-dev", "split_files/vae/flux2-vae.safetensors"),
sd_cpp_vae_format = "flux2",
sd_cpp_text_encoders = (
("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors", "llm"),
),
),
# FLUX.2-dev is the full (non-distilled) FLUX.2. It uses the Mistral-based
# Flux2Pipeline, distinct from klein's Qwen3-based Flux2KleinPipeline, so it needs
# its own entry. Its base diffusers repo is gated (gated=auto) but reachable with an
# HF token. text-to-image only: diffusers 0.38 ships no Flux2 img2img / inpaint
# pipeline for dev. VAE + Mistral text encoder come from the open Comfy-Org/flux2-dev
# mirror for the sd-cli path (shares the FLUX.2 32-channel AE with klein).
DiffusionFamily(
name = "flux.2-dev",
pipeline_class = "Flux2Pipeline",
transformer_class = "Flux2Transformer2DModel",
base_repo = "black-forest-labs/FLUX.2-dev",
aliases = ("flux2-dev", "flux2dev"),
sd_cpp_vae = ("Comfy-Org/flux2-dev", "split_files/vae/flux2-vae.safetensors"),
sd_cpp_vae_format = "flux2",
sd_cpp_text_encoders = (
(
"Comfy-Org/flux2-dev",
"split_files/text_encoders/mistral_3_small_flux2_bf16.safetensors",
"llm",
),
),
),
DiffusionFamily(
# Instruction editing with FLUX. FluxKontextPipeline takes an input image + an edit
# instruction; the GGUF transformer is the standard FluxTransformer2DModel, with the
# T5/CLIP text encoders + VAE from the base diffusers repo. cfg defaults to
# guidance_scale (FLUX). Most-specific aliases first so detect_family prefers this
# over the plain "flux.1" family and un-rejects the "kontext" keyword for it.
name = "flux.1-kontext",
pipeline_class = "FluxKontextPipeline",
transformer_class = "FluxTransformer2DModel",
base_repo = "black-forest-labs/FLUX.1-Kontext-dev",
aliases = ("flux.1-kontext-dev", "flux1-kontext", "flux-kontext", "kontext"),
edit = True,
),
DiffusionFamily(
# Instruction editing (image-in + text-instruction-out). The 2511 checkpoint ships
# as QwenImageEditPlusPipeline (multi-image-capable); the GGUF transformer is the
# standard QwenImageTransformer2DModel, with the VAE / Qwen2.5-VL text-encoder /
# image processor / scheduler coming from the base diffusers repo. Most-specific
# aliases first so detect_family prefers this over the plain "qwen-image" family.
name = "qwen-image-edit",
pipeline_class = "QwenImageEditPlusPipeline",
transformer_class = "QwenImageTransformer2DModel",
base_repo = "Qwen/Qwen-Image-Edit-2511",
cfg_kwarg = "true_cfg_scale",
aliases = (
"qwen-image-edit-2511",
"qwen-image-edit-2509",
"qwen-image-edit",
"qwen_image_edit",
"qwenimageedit",
),
edit = True,
),
DiffusionFamily(
name = "qwen-image",
pipeline_class = "QwenImagePipeline",
transformer_class = "QwenImageTransformer2DModel",
base_repo = "Qwen/Qwen-Image",
cfg_kwarg = "true_cfg_scale",
aliases = ("qwen_image", "qwenimage"),
# LoRA training via the DiT trainer, defaulting to the prequant nf4 repo (QLoRA).
trainable = True,
train_base_repos = ("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", "Qwen/Qwen-Image"),
img2img_pipeline_class = "QwenImageImg2ImgPipeline",
inpaint_pipeline_class = "QwenImageInpaintPipeline",
controlnet_pipeline_class = "QwenImageControlNetPipeline",
controlnet_model_class = "QwenImageControlNetModel",
sd_cpp_vae = ("Comfy-Org/Qwen-Image_ComfyUI", "split_files/vae/qwen_image_vae.safetensors"),
# The Qwen2.5-VL text encoder as a Q4_K_M GGUF keeps the CPU RAM win (the
# bf16 safetensors encoder is ~15 GB). sd-cli's --qwen2vl is an alias of --llm.
sd_cpp_text_encoders = (
(
"unsloth/Qwen2.5-VL-7B-Instruct-GGUF",
"Qwen2.5-VL-7B-Instruct-Q4_K_M.gguf",
"qwen2vl",
),
),
# Qwen-Image's supported sd.cpp invocation (docs/qwen_image.md).
sd_cpp_sampling_method = "euler",
sd_cpp_flow_shift = 3.0,
),
DiffusionFamily(
name = "z-image",
pipeline_class = "ZImagePipeline",
transformer_class = "ZImageTransformer2DModel",
base_repo = "Tongyi-MAI/Z-Image-Turbo",
aliases = ("zimage", "z_image"),
# LoRA training via the DiT trainer (bf16 only). Defaults to the prequant nf4 repo
# for QLoRA; the bf16 Tongyi-MAI base is the alternative.
trainable = True,
train_base_repos = ("unsloth/Z-Image-Turbo-unsloth-bnb-4bit", "Tongyi-MAI/Z-Image-Turbo"),
img2img_pipeline_class = "ZImageImg2ImgPipeline",
inpaint_pipeline_class = "ZImageInpaintPipeline",
# Z-Image's MLP down-projections peak near 9e5, which overflows float16.
fp16_incompatible = True,
sd_cpp_vae = ("Comfy-Org/z_image_turbo", "split_files/vae/ae.safetensors"),
sd_cpp_text_encoders = (
("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors", "llm"),
),
),
# SDXL is the one U-Net family here: the denoiser is ``pipe.unet``
# (UNet2DConditionModel), not a DiT ``pipe.transformer``, and a single-file
# ``.safetensors`` is the WHOLE pipeline rather than a transformer-only file.
# So it declares ``denoiser_attr = "unet"`` + ``single_file_is_pipeline = True``
# and loads via the pipeline class (from_pretrained for a repo, from_single_file
# for a single .safetensors). The base repo supplies both CLIP text encoders,
# the VAE and the scheduler on the pipeline path. img2img / inpaint / ControlNet
# are the standard SDXL pipelines, built around the resident modules via
# from_pipe like every other family. There is no GGUF/single-file transformer
# path for SDXL (the whole checkpoint is one file), and no native sd.cpp mapping
# yet, so the no-GPU route falls back to diffusers.
DiffusionFamily(
name = "sdxl",
pipeline_class = "StableDiffusionXLPipeline",
transformer_class = "UNet2DConditionModel",
base_repo = "stabilityai/stable-diffusion-xl-base-1.0",
aliases = ("stable-diffusion-xl", "sd-xl", "sd_xl", "sdxl-turbo", "sdxl-base"),
denoiser_attr = "unet",
single_file_is_pipeline = True,
img2img_pipeline_class = "StableDiffusionXLImg2ImgPipeline",
inpaint_pipeline_class = "StableDiffusionXLInpaintPipeline",
controlnet_pipeline_class = "StableDiffusionXLControlNetPipeline",
controlnet_model_class = "ControlNetModel",
# SDXL is the one family with a shipped LoRA trainer today (the U-Net trainer).
# DiT families (flux.1 / qwen-image / z-image) become trainable in a follow-up.
trainable = True,
train_base_repos = (
"stabilityai/stable-diffusion-xl-base-1.0",
"stabilityai/sdxl-turbo",
),
),
)
def trainable_family_names() -> tuple[str, ...]:
"""Names of families Studio can train a LoRA on, in registry order."""
return tuple(fam.name for fam in _FAMILIES if fam.trainable)
# Editing / inpaint checkpoints share an arch keyword but need a different
# pipeline and an input image, which this text-to-image backend doesn't drive.
# "layered" rejects Qwen-Image-Layered: its transformer sets additional_t_cond=True
# and expects an extra addition_t_cond input that the standard QwenImagePipeline
# never supplies, so it loads but crashes at the first denoise step. Rejecting it
# here fails the load fast with a clear message and hides it from the picker.
_EDIT_KEYWORDS = ("edit", "kontext", "inpaint", "layered")
def _token_in_needle(token: str, needle: str) -> bool:
"""True when ``token`` appears in ``needle`` as a whole path/name segment, i.e.
delimited by a separator (``- _ . / \\``) or a string boundary, not merely as a
raw substring. This keeps multi-part tokens matching where they should
('qwen-image-edit' in 'qwen-image-edit-2511') while preventing a short token from
matching inside an unrelated word ('kontext' must not match 'kontextual', 'edit'
must not match 'edition')."""
return re.search(r"(?:^|[-_./\\])" + re.escape(token) + r"(?:$|[-_./\\])", needle) is not None
def _best_family_match(needle: str) -> Optional[DiffusionFamily]:
"""The family whose name/alias is the LONGEST whole-segment token of ``needle``.
Longest = most specific, so an edit checkpoint ('...qwen-image-edit-2511...')
matches the 'qwen-image-edit' family rather than the generic 'qwen-image' one.
Segment matching (not raw substring) stops a short alias like 'kontext' from
hijacking an unrelated path such as '.../kontextual/z-image-...gguf'."""
best: Optional[tuple[DiffusionFamily, int]] = None
for fam in _FAMILIES:
for token in (fam.name, *fam.aliases):
if _token_in_needle(token, needle) and (best is None or len(token) > best[1]):
best = (fam, len(token))
return best[0] if best else None
def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[DiffusionFamily]:
"""Resolve a ``DiffusionFamily`` from a repo id, or an explicit override.
``override`` matches a family ``name`` or alias exactly. Otherwise the most-specific
family whose name/alias is a substring of the repo id wins. Supported editing families
(Qwen-Image-Edit) match here; unsupported editing/inpaint/layered checkpoints that only
share a base family's arch keyword are still rejected (None), because they need a
different pipeline + input this backend's base text-to-image path doesn't drive.
"""
if override:
key = override.strip().lower()
for fam in _FAMILIES:
if key == fam.name or key in fam.aliases:
return fam
return None
needle = repo_id.lower()
match = _best_family_match(needle)
if match is not None:
# Don't let a generic base family (e.g. qwen-image) swallow a variant it can't run
# (qwen-image-LAYERED, ...-Inpaint): if the id still carries a reject keyword the
# matched family does not itself declare, reject so the load fails fast + clearly.
# Scope the keyword check to the LAST path component (the model id or
# filename), not arbitrary parent directories: a valid file selected as
# repo_id `/models/edit` + filename `Z-Image-Turbo-Q4.gguf` must not be
# rejected because a parent folder happens to be named `edit`. The
# combined `repo_id/gguf_filename` fallback passes the filename last.
basename = re.split(r"[/\\]+", needle)[-1]
matched_tokens = (match.name, *match.aliases)
if any(
_token_in_needle(kw, basename) and not any(kw in tok for tok in matched_tokens)
for kw in _EDIT_KEYWORDS
):
return None
return match
return None
def supported_family_names() -> tuple[str, ...]:
"""Family names accepted as ``family_override`` and shown in the unknown-model
error. Kept in registry order so the message lists what the backend can load."""
return tuple(fam.name for fam in _FAMILIES)
def detect_family_for_pick(
repo_id: str,
gguf_filename: Optional[str] = None,
override: Optional[str] = None,
) -> Optional[DiffusionFamily]:
"""``detect_family``, falling back to the combined path/filename for a direct
local ``.gguf`` pick. The frontend splits such a pick into (parent dir, basename),
so the family keyword can live only in the filename (e.g.
``/models/z-image-turbo-Q4_K_M.gguf``) while the parent directory carries none;
scan the combined string too when the directory alone is undetectable. Only a
fallback, so remote ``org/name`` picks and explicit overrides behave exactly as
``detect_family``. Shared by both engines so validation and load can't diverge."""
fam = detect_family(repo_id, override)
if fam is None and gguf_filename and not override:
fam = detect_family(f"{repo_id}/{gguf_filename}", override)
return fam
def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
"""The companion diffusers repo: caller-supplied if given, else the family fallback."""
base = (base_repo or "").strip()
return base or fam.base_repo
# Default (steps, guidance) per model for callers that can't pass them — namely
# the OpenAI /v1/images/generations endpoint, whose spec has no step/guidance
# knobs. Distilled "turbo/schnell" models want few steps and no CFG; the full
# "dev" models want more steps and real CFG. Matched by substring, most specific
# first — the same scheme and values as the UI's MODEL_DEFAULTS table
# (studio/frontend/src/features/images/images-page.tsx); keep the two in sync.
_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = (
("z-image-turbo", 9, 0.0),
("flux.1-schnell", 4, 0.0),
# Kontext (editing) before the generic flux.1: ~28 steps, lower guidance (~2.5).
("kontext", 28, 2.5),
("flux.1", 28, 3.5),
("flux.2-klein", 4, 0.0),
# FLUX.2-dev is the full (non-distilled) model: more steps + real guidance.
("flux.2-dev", 28, 4.0),
("qwen-image", 20, 4.0),
("z-image", 20, 4.0),
# SDXL: Turbo is distilled (few steps, no CFG); base/full SDXL wants ~30 steps and
# real CFG (~7). "sdxl-turbo" must precede the generic "sdxl" substring match.
("sdxl-turbo", 3, 0.0),
("stable-diffusion-xl", 30, 7.0),
("sdxl", 30, 7.0),
)
# Unrecognised model: distilled few-step / no-CFG shape, matching the UI fallback.
_GENERATION_DEFAULT_FALLBACK = (9, 0.0)
def default_generation_params(*identifiers: Optional[str]) -> tuple[int, float]:
"""Default ``(steps, guidance)`` for a loaded model. The first identifier that
names a known model wins (the repo id, then the resolved base repo), so a
local-path load — whose repo id is just a filesystem path that may not name
the model — still resolves via its base repo. Within an identifier, keys are
matched as substrings, most specific first (the same scheme as the UI)."""
for identifier in identifiers:
needle = (identifier or "").lower()
for key, steps, guidance in _GENERATION_DEFAULTS:
if key in needle:
return steps, guidance
return _GENERATION_DEFAULT_FALLBACK
def family_prequant_repo(fam: DiffusionFamily, scheme: str) -> Optional[str]:
"""The hosted pre-quantized transformer repo for ``scheme`` in this family, or None."""
for entry_scheme, repo_id in fam.prequant_repos:
if entry_scheme == scheme:
return repo_id
return None
def family_sd_cpp_supported(fam: DiffusionFamily) -> bool:
"""True when the family has the single-file VAE + text-encoder mapping the
native sd.cpp engine needs. A family without it can only run on diffusers, so
the no-GPU route falls back rather than routing to sd-cli."""
return bool(fam.sd_cpp_vae and fam.sd_cpp_text_encoders)
def resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path:
"""Resolve ``gguf_filename`` to a file under ``repo_root``, rejecting escapes.
``gguf_filename`` is user-supplied, so an absolute path or a ``..`` segment
could otherwise read outside the local repo directory.
"""
if (
Path(gguf_filename).is_absolute()
or PurePosixPath(gguf_filename).is_absolute()
or gguf_filename.startswith(("/", "\\"))
or "\\" in gguf_filename
):
raise ValueError("gguf_filename must be a relative path inside the repo.")
rel = PurePosixPath(gguf_filename)
if any(part in ("", ".", "..") for part in rel.parts):
raise ValueError("gguf_filename must not contain '', '.', or '..' segments.")
# Resolve symlinks before the containment check: the guards above stop
# lexical escapes, but a symlink inside the repo could still point outside it.
repo_real = repo_root.resolve()
child = repo_root.joinpath(*rel.parts).resolve()
if child != repo_real and repo_real not in child.parents:
raise ValueError("gguf_filename must resolve to a file inside the repo.")
if not child.is_file():
raise FileNotFoundError(f"'{gguf_filename}' is not a file under {repo_root}.")
return child
Reference in a new issue View git blame Copy permalink