# 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: RuntimeError messages for client-recoverable generate states. The # /images/generate route matches these EXACTLY for a 409 (vs a 500), so both engines raise them # verbatim -- 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 guidance. Most use "guidance_scale"; Qwen-Image's real CFG is # "true_cfg_scale" (its distilled guidance is off). cfg_kwarg: str = "guidance_scale" # The pipe attribute holding the denoiser: DiT families ``pipe.transformer`` (default), U-Net # families (SDXL) ``pipe.unet``. Read wherever the backend touches the denoiser generically. denoiser_attr: str = "transformer" # True when a single-file ``.safetensors`` is the WHOLE pipeline (SDXL), so the loader calls # ``pipeline_class.from_single_file`` directly. DiT families leave this False (transformer-only). single_file_is_pipeline: bool = False # True for families whose pipeline needs MULTIPLE denoisers no single file carries (Ideogram 4: # conditional + unconditional_transformer), so only a full ``pipeline`` load is valid; # validate_load_request rejects single-file / GGUF kinds up front. pipeline_only: bool = False # Optional diffusers pipeline classes for image-conditioned workflows, built around the resident # modules via ``Pipeline.from_pipe`` (no reload). None = unsupported (UI gates it off). img2img_pipeline_class: Optional[str] = None inpaint_pipeline_class: Optional[str] = None # ControlNet pipeline + model classes: the backend loads the model via from_pretrained and # builds the pipeline via ``from_pipe(base, controlnet=model)`` (no reload). None on both = # no support (UI gates it off). controlnet_pipeline_class: Optional[str] = None controlnet_model_class: Optional[str] = None # True when the inpaint pipeline keeps the canvas size, so it can also drive outpaint. False for # FLUX.2 (it scales >1MP inputs to ~1MP, shrinking an outpaint canvas) -> Inpaint but not Extend. inpaint_preserves_size: bool = True # True for instruction-editing families (Qwen-Image-Edit / FLUX Kontext): the OWN pipeline IS the # edit pipeline (image + instruction, no plain text-to-image), used directly (no from_pipe). # ``base_repo`` supplies the VAE / text-encoder / processor / scheduler for the GGUF transformer. edit: bool = False # True for families whose text-to-image pipeline ALSO accepts reference image(s) (FLUX.2's # ``image`` arg). Unlike ``edit`` they still do plain text-to-image; unlike img2img the # conditioning is reference-based (no ``strength``, output size from width/height). Used directly. 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 (-> inf/NaN -> black image); the backend # promotes a resolved float16 to float32 for these. fp16_incompatible: bool = False # False only for a family whose denoiser block doesn't compile cleanly with regional # torch.compile. Consulted on the GGUF path too; all current families compile, so this stays True. supports_torch_compile: bool = True # Optional pre-quantized transformer checkpoints as (scheme, repo_id) pairs. When the fast quant # path resolves a scheme with a hosted checkpoint, the loader fetches the already-quantized # weights instead of the dense bf16 (lower load VRAM + smaller download). Empty -> unchanged. prequant_repos: tuple[tuple[str, str], ...] = field(default_factory = tuple) # Hosted checkpoints for NON-DEFAULT bases of the family, as (base_repo, scheme, repo_id) # triples with base_repo lowercased. One family entry covers several published variants # (flux.1: schnell/dev/Krea-dev) whose weights differ, so each variant needs its own baked # checkpoint; the loader's base_model_id validation correctly refuses the default entry for # them. Resolution prefers an exact variant match, then falls back to ``prequant_repos``. prequant_variant_repos: tuple[tuple[str, str, str], ...] = field(default_factory = tuple) # Hosted PRE-CAST text-encoder checkpoints as (scheme, component, repo_id) triples # (component is the pipeline attribute, e.g. "text_encoder"). Serves the layerwise-fp8 # storage scheme only: the cast is a deterministic transform, so the stored artifact is # bit-identical to dense-load-then-cast while skipping the multi-GB dense TE download # (see diffusion_te_prequant.py). Empty -> the TE loads dense and casts as before. te_prequant_repos: tuple[tuple[str, str, str], ...] = field(default_factory = tuple) # Native (sd.cpp) single-file assets, used only on the no-GPU sd.cpp engine. The transformer GGUF # is shared with diffusers; sd-cli also needs a single-file VAE + text encoder(s) (the base repo # ships those sharded). Each is a (repo_id, filename); ``sd_cpp_text_encoders`` carries a trailing # SdCppModelFiles field name (clip_l / t5xxl / llm / qwen2vl / clip_g) for the sd-cli flag. Empty # -> no native mapping (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 FLUX.2, None 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 so the native output matches the model's supported # invocation (e.g. Qwen-Image needs euler + flow-shift 3). None leaves sd-cli defaults. 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). Opt-in per family # (each arch needs its own loop); the training-start path refuses a non-trainable family up front. trainable: bool = False # Recommended base repos to train FROM, most-preferred first (e.g. a QLoRA prequant repo, then # bf16). Surfaced by the Train UI. train_base_repos: tuple[str, ...] = field(default_factory = tuple) # When set, deploying a LoRA trained on this family loads THIS repo instead of the trained-on # checkpoint (Krea: train on Raw, preview on Turbo). Both sides must be the same precision so the # swap never enlarges the load. Unset elsewhere. deploy_base_repo: Optional[str] = None # Keyed by architecture, not per 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 # a fallback). Only archs whose diffusers transformer supports from_single_file load here. _FAMILIES: tuple[DiffusionFamily, ...] = ( DiffusionFamily( name = "flux.1", pipeline_class = "FluxPipeline", transformer_class = "FluxTransformer2DModel", base_repo = "black-forest-labs/FLUX.1-schnell", # Hosted pre-quantized DiT checkpoints (gate-validated vs same-seed bf16). The loader # verifies the checkpoint's baked base_model_id against the repo actually being loaded, # so a non-default base (e.g. FLUX.1-dev under this family) safely falls back to the # dense-quantize path instead of loading schnell weights. prequant_repos = ( ("int8", "unsloth/FLUX.1-schnell-FP8"), ("fp8", "unsloth/FLUX.1-schnell-FP8"), ), # Gate-validated checkpoints baked from the dev / Krea-dev weights (same arch, different # weights): without these entries the default schnell checkpoint is refused for those # bases and every int8/fp8 load pays the dense download + on-the-fly quantise. prequant_variant_repos = ( ("black-forest-labs/flux.1-dev", "int8", "unsloth/FLUX.1-dev-FP8"), ("black-forest-labs/flux.1-dev", "fp8", "unsloth/FLUX.1-dev-FP8"), ("black-forest-labs/flux.1-krea-dev", "int8", "unsloth/FLUX.1-Krea-dev-FP8"), ("black-forest-labs/flux.1-krea-dev", "fp8", "unsloth/FLUX.1-Krea-dev-FP8"), ), # Pre-cast T5-XXL (9.52 -> 5.90 GB; CLIP-L stays dense, 0.25 GB). One artifact # serves schnell/dev/Krea-dev: the T5 shards are byte-identical across all three # (verified sha256, see diffusion_te_prequant._TE_EQUIVALENT_BASES). te_prequant_repos = (("fp8", "text_encoder_2", "unsloth/FLUX.1-schnell-FP8"),), aliases = ("flux1", "flux-1"), # LoRA training targets FLUX.1-dev via the DiT trainer (QLoRA nf4); the dev repo is gated. 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 Flux2KleinPipeline (Qwen3 encoder), not the Mistral Flux2Pipeline; must # precede a generic flux match. The Mistral Flux2Pipeline is the 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", prequant_repos = ( ("int8", "unsloth/FLUX.2-klein-4B-FP8"), ("fp8", "unsloth/FLUX.2-klein-4B-FP8"), ), aliases = ("flux2-klein",), # LoRA training via the DiT trainer (QLoRA nf4 by default); klein-4B is not gated. trainable = True, train_base_repos = ("black-forest-labs/FLUX.2-klein-4B",), # Flux2KleinPipeline takes reference image(s) via `image`, so it exposes a "reference" # workflow atop text-to-image. It has an inpaint pipeline (no img2img) -> inpaint + extend. reference = True, inpaint_pipeline_class = "Flux2KleinInpaintPipeline", # FLUX.2 scales >1MP inputs to ~1MP, so outpaint can't grow. inpaint_preserves_size = False, # FLUX.2's 32-channel AE needs the latent-format override; the single-file VAE ships in # Comfy-Org/flux2-dev (klein-4B has only 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: full (non-distilled) FLUX.2 on the Mistral Flux2Pipeline (distinct from klein), so # its own entry. Base repo is gated. Text-to-image only (no Flux2 img2img/inpaint in diffusers # 0.38). VAE + Mistral encoder come from the open Comfy-Org/flux2-dev mirror for sd-cli. DiffusionFamily( name = "flux.2-dev", pipeline_class = "Flux2Pipeline", transformer_class = "Flux2Transformer2DModel", base_repo = "black-forest-labs/FLUX.2-dev", prequant_repos = ( ("int8", "unsloth/FLUX.2-dev-FP8"), ("fp8", "unsloth/FLUX.2-dev-FP8"), ), # Pre-cast Mistral-Small-24B conditioner (bf16 ~48 GB dense, ~24.7 GB pre-cast). te_prequant_repos = (("fp8", "text_encoder", "unsloth/FLUX.2-dev-FP8"),), aliases = ("flux2-dev", "flux2dev"), # LoRA training via the DiT trainer (QLoRA nf4 by default); the base repo is gated, so # training requires an HF token with the FLUX.2-dev license accepted. trainable = True, train_base_repos = ("black-forest-labs/FLUX.2-dev",), 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( # FLUX instruction editing: FluxKontextPipeline takes an image + edit instruction; the GGUF # transformer is standard FluxTransformer2DModel. Specific aliases first so detect_family # prefers this over "flux.1" and un-rejects the "kontext" keyword. 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( # Qwen instruction editing: the 2511 checkpoint ships as QwenImageEditPlusPipeline # (multi-image); the GGUF transformer is standard QwenImageTransformer2DModel. Specific # aliases first so detect_family prefers this over "qwen-image". 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", # int8 only: fp8 is family-denied (_FAMILY_SCHEME_DENY) so a repo entry would be dead. prequant_repos = (("int8", "unsloth/Qwen-Image-FP8"),), # Pre-cast Qwen2.5-VL-7B (bf16 ~16.6 GB dense, ~8.8 GB pre-cast). The DiT fp8 denial # is a transformer-scheme rule; the layerwise TE cast is unaffected. te_prequant_repos = (("fp8", "text_encoder", "unsloth/Qwen-Image-FP8"),), 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"), # Qwen2.5-VL as a Q4_K_M GGUF keeps the CPU RAM win (bf16 encoder is ~15 GB). sd-cli's # --qwen2vl aliases --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", prequant_repos = ( ("int8", "unsloth/Z-Image-Turbo-FP8"), ("fp8", "unsloth/Z-Image-Turbo-FP8"), ), # Pre-cast Qwen3-4B TE (8.04 -> 4.41 GB). NOT shared with flux.2-klein-4B: klein's # TE retrained layer 35's MLP (up/down_proj maxdiff 0.86 vs this checkpoint). te_prequant_repos = (("fp8", "text_encoder", "unsloth/Z-Image-Turbo-FP8"),), aliases = ("zimage", "z_image"), # LoRA training via the DiT trainer (bf16); defaults to the prequant nf4 repo for QLoRA. 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"), ), ), # Krea 2 (diffusers >= 0.39): a ~12B single-stream DiT with a Qwen3-VL-4B encoder and the # Qwen-Image VAE. Loaded per-component (diffusion_krea2.py) because the repo ships # transformers-5.x configs. No GGUF/sd.cpp mapping yet. DiffusionFamily( name = "krea-2", pipeline_class = "Krea2Pipeline", transformer_class = "Krea2Transformer2DModel", base_repo = "krea/Krea-2-Turbo", prequant_repos = ( ("int8", "unsloth/Krea-2-Turbo-FP8"), ("fp8", "unsloth/Krea-2-Turbo-FP8"), ), # Pre-cast Qwen3-VL-4B TE (8.88 -> 4.83 GB); handed into load_krea2_pipeline # directly (constructor assembly never sees pipe_kwargs). te_prequant_repos = (("fp8", "text_encoder", "unsloth/Krea-2-Turbo-FP8"),), aliases = ("krea2",), # LoRA training via the DiT trainer (no prequant repo yet, so nf4 quantizes on the fly). # Krea's guidance: train on the undistilled Raw, run adapters on Turbo, so Raw is the # default training base and Turbo the inference/base repo. trainable = True, train_base_repos = ("krea/Krea-2-Raw", "krea/Krea-2-Turbo"), # Adapters trained on Raw run on Turbo; deploy previews them there (same bf16 precision). deploy_base_repo = "krea/Krea-2-Turbo", # Exported bf16-only; fp16 unvalidated upstream, so keep the fp16 fallback off like z-image. fp16_incompatible = True, ), # Lumina Image 2.0: a 2.6B single-stream DiT with a Gemma2-2B encoder and a standard # 16-channel AutoencoderKL, all transformers-4.x-compatible, so the generic # from_pretrained pipeline path loads it. No GGUF/sd.cpp mapping exists upstream. # NOT aliased to bare "lumina": Lumina-Next checkpoints are a different arch # (LuminaText2ImgPipeline) and must stay unknown rather than crash mid-load. DiffusionFamily( name = "lumina-2", pipeline_class = "Lumina2Pipeline", transformer_class = "Lumina2Transformer2DModel", base_repo = "Alpha-VLLM/Lumina-Image-2.0", # Gate-validated hosted checkpoints (28/28 pairs each; LPIPS mean 0.146 int8 / # 0.116 fp8 vs same-seed bf16). prequant_repos = ( ("int8", "unsloth/Lumina-Image-2.0-FP8"), ("fp8", "unsloth/Lumina-Image-2.0-FP8"), ), # Pre-cast Gemma2-2B TE. The Hub stores it fp32 (10.46 GB), so the 3.20 GB # artifact is a 3.3x download cut even though the model is small. te_prequant_repos = (("fp8", "text_encoder", "unsloth/Lumina-Image-2.0-FP8"),), aliases = ("lumina-image-2.0", "lumina-image-2", "lumina2"), # Published and validated bf16-only upstream; keep the fp16 fallback off like z-image. fp16_incompatible = True, ), # HunyuanImage 2.1 (diffusers >= 0.39): a 17B dual-stream DiT with a Qwen2.5-VL text # encoder, a ByT5 glyph encoder, and the 32x-compression HunyuanImage VAE. The community # mirror also ships guider/ocr_guider components (AdaptiveProjectedMixGuidance), which # 0.39 loads natively, so the generic from_pretrained pipeline path covers the whole # stack. 2K-native (the card recipe renders 2048x2048); classifier-free guidance runs # inside the repo's guider at its baked scale, and the call's own guidance knob is # distilled_guidance_scale (there is no guidance_scale kwarg). Distinct from # HunyuanImage-3.0, which stays excluded above: 2.1 has a real diffusers pipeline. DiffusionFamily( name = "hunyuanimage-2.1", # Hosted checkpoints, verified bit-identical to on-the-fly quantize (the family's # guider pipeline is not run-to-run deterministic, so same-seed LPIPS vs bf16 blends # trajectory divergence with harness noise; per-case hard checks pass and the drift is # compositional, reviewed visually). prequant_repos = ( ("int8", "unsloth/HunyuanImage-2.1-FP8"), ("fp8", "unsloth/HunyuanImage-2.1-FP8"), ), # The Qwen2.5-VL TE is byte-identical to Qwen-Image's (verified sha256, see # _TE_EQUIVALENT_BASES), so the family reuses the Qwen-Image artifact: zero new # hosting, 16.58 -> 8.84 GB download. ByT5 (text_encoder_2) stays dense. te_prequant_repos = (("fp8", "text_encoder", "unsloth/Qwen-Image-FP8"),), pipeline_class = "HunyuanImagePipeline", transformer_class = "HunyuanImageTransformer2DModel", base_repo = "hunyuanvideo-community/HunyuanImage-2.1-Diffusers", cfg_kwarg = "distilled_guidance_scale", aliases = ("hunyuanimage-2.1-diffusers", "hunyuanimage2.1"), # Exported bf16-only; keep the fp16 fallback off like z-image / krea-2. fp16_incompatible = True, ), # HiDream-I1: a 17B MoE DiT (16 double + 32 single layers, 4 routed experts) with FOUR text # encoders. The repos ship CLIP-L/CLIP-G/T5-XXL but NOT the Llama-3.1-8B text_encoder_4 their # model_index names: the loader assembles it from the open unsloth mirror # (diffusion_hidream.py). Full / Dev / Fast share the arch, so one family covers all three # (per-variant step/guidance defaults below). city96 publishes a GGUF but the GGUF path would # need the same TE4 assembly for tiny demand, so no GGUF artifact is wired yet. DiffusionFamily( name = "hidream-i1", # Hosted checkpoints: 28/28 per-case gate pairs per scheme (LPIPS suite means 0.291 # int8 / 0.278 fp8, the 50-step trajectory band); int8 verified bit-identical to # on-the-fly quantize across all 1615 state dict tensors. prequant_repos = ( ("int8", "unsloth/HiDream-I1-Full-FP8"), ("fp8", "unsloth/HiDream-I1-Full-FP8"), ), # Pre-cast Llama-3.1-8B TE4 (16.1 GB bf16 -> 8.1 GB). The generic TE pass only covers # text_encoder.._3, so TE4 engages via hidream_te4_kwargs, not te_prequant_pipe_kwargs. te_prequant_repos = (("fp8", "text_encoder_4", "unsloth/HiDream-I1-Full-FP8"),), pipeline_class = "HiDreamImagePipeline", transformer_class = "HiDreamImageTransformer2DModel", base_repo = "HiDream-ai/HiDream-I1-Full", aliases = ("hidream", "hidream-i1-full", "hidream-i1-dev", "hidream-i1-fast"), # Exported bf16-only; keep the fp16 fallback off like the other modern DiTs. fp16_incompatible = True, ), # Ideogram 4 (diffusers >= 0.39): a 34-layer DiT PAIR (conditional + unconditional_transformer # for dual-branch CFG, both ~9B, so memory planning counts two DiTs) with a Qwen3-VL encoder. # No bf16 checkpoint: ideogram-4-fp8 (raw float8, upcast on load) is the highest-precision # artifact and the family base; the -nf4 repos carry bnb-4bit quantization_configs. All gated. # No GGUF/sd.cpp mapping. CFG quirk: the pipeline takes guidance_scale OR a per-step # guidance_schedule (see the loader's IDEOGRAM4 branch). DiffusionFamily( name = "ideogram-4", pipeline_class = "Ideogram4Pipeline", transformer_class = "Ideogram4Transformer2DModel", base_repo = "ideogram-ai/ideogram-4-fp8", aliases = ("ideogram4", "ideogram-v4", "ideogram"), # Two DiTs assembled per-component, so no transformer-only single-file / GGUF load. pipeline_only = True, ), # SDXL is the one U-Net family: the denoiser is ``pipe.unet`` and a single-file ``.safetensors`` # is the WHOLE pipeline, so it sets ``denoiser_attr="unet"`` + ``single_file_is_pipeline=True`` # and loads via the pipeline class. img2img / inpaint / ControlNet are the standard SDXL # pipelines via from_pipe. No GGUF/single-file transformer path and no sd.cpp mapping. 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 uses the U-Net LoRA trainer. 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) # The family whose CFG uses a guidance_scale/guidance_schedule pair (the loader special-cases the # call). Named here so the two modules can't drift. IDEOGRAM4_FAMILY_NAME = "ideogram-4" # The family whose generate call carries the card's CFG-truncation ratio (the loader # special-cases the call). Named here so the two modules can't drift. LUMINA2_FAMILY_NAME = "lumina-2" # Models Studio deliberately does NOT support, reason surfaced verbatim in the load error (vs the # generic unknown-family message). Keyed by a lowercase repo-id substring. The bar is a diffusers # pipeline: HunyuanImage-3.0 is an 80B MoE needing AutoModelForCausalLM + trust_remote_code (RCE # out of the question). _EXCLUDED_MODELS: tuple[tuple[str, str], ...] = ( ( # "-3" scoped so a future HunyuanImage 2.x with a diffusers pipeline falls through normally. "hunyuanimage-3", "HunyuanImage-3.0 has no diffusers pipeline (it is an 80B autoregressive MoE " "that requires trust_remote_code), so Studio does not support it.", ), ) def excluded_model_reason(repo_id: str) -> Optional[str]: """The stated reason ``repo_id`` is unsupported, or None when it is simply unknown.""" needle = (repo_id or "").lower() for token, reason in _EXCLUDED_MODELS: if _token_in_needle(token, needle): return reason return None # Editing / inpaint checkpoints share an arch keyword but need a different pipeline + input image. # "layered" rejects Qwen-Image-Layered (its transformer expects an extra addition_t_cond input # the standard pipeline never supplies, so it crashes at the first denoise). Fails the load fast. _EDIT_KEYWORDS = ("edit", "kontext", "inpaint", "layered") def _token_in_needle(token: str, needle: str) -> bool: """True when ``token`` appears in ``needle`` as a whole segment (delimited by ``- _ . / \\`` or a boundary), not a raw substring, so 'qwen-image-edit' matches '...-2511' but 'kontext' doesn't match 'kontextual'.""" 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 '...qwen-image-edit-2511...' matches 'qwen-image-edit', not 'qwen-image').""" 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``/alias exactly; otherwise the most-specific family whose name/alias is a substring of the repo id wins. Supported editing families match here; unsupported editing/inpaint/layered checkpoints sharing only an arch keyword are rejected (None). """ 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 family (qwen-image) swallow a variant it can't run # (qwen-image-LAYERED): if the id carries a reject keyword the matched family doesn't # declare, reject. Scope the check to the LAST path component so a parent folder named # `edit` doesn't reject a valid file (the repo_id/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 (registry order).""" 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 local ``.gguf`` pick where the family keyword lives only in the filename. Only a fallback, so remote picks and overrides behave exactly as ``detect_family``. Shared by both engines.""" 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 (the OpenAI # /v1/images/generations endpoint has no step/guidance knobs). Matched by substring, most specific # first -- same values as the UI's MODEL_DEFAULTS table (images-page.tsx); keep in sync. _GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = ( ("z-image-turbo", 9, 0.0), # FLUX.1 Krea dev is a FLUX.1-dev finetune (flux.1 family), NOT a Krea-2: its card runs # 28 steps at guidance 4.5. Must precede the generic "krea" key below, which would # otherwise hand it Krea-2-Turbo's 8-step no-CFG recipe. ("flux.1-krea", 28, 4.5), # Krea 2 Raw (undistilled): 52 steps / guidance 3.5. Must precede the generic "krea" key. ("krea-2-raw", 52, 3.5), # Krea 2 Turbo (distilled): 8 steps, no CFG. "krea" then covers Turbo and other krea ids but Raw. ("krea", 8, 0.0), ("flux.1-schnell", 4, 0.0), ("kontext", 28, 2.5), # editing: before the generic flux.1 ("flux.1", 28, 3.5), ("flux.2-klein", 4, 0.0), ("flux.2-dev", 28, 4.0), # full (non-distilled) ("qwen-image", 20, 4.0), ("z-image", 20, 4.0), # Lumina Image 2.0 model-card: 50 steps, guidance 4 (plus cfg_trunc_ratio 0.25, which the # loader passes itself; see LUMINA2_FAMILY_NAME). ("lumina", 50, 4.0), # HunyuanImage 2.1 model-card: 50 steps; the guidance value feeds the call's # distilled_guidance_scale (default 3.25), while real CFG runs inside the repo guiders. ("hunyuanimage", 50, 3.25), # HiDream-I1 upstream inference.py: Full 50 steps / guidance 5; the distilled Dev (28) and # Fast (16) run guidance-free. Specific keys precede the generic "hidream" (Full + fallback). ("hidream-i1-dev", 28, 0.0), ("hidream-i1-fast", 16, 0.0), ("hidream", 50, 5.0), # Ideogram 4 model-card: 48 steps, guidance 7 (its schedule tapers the last 3 steps to 3.0; # the loader keeps that taper when the request matches these defaults exactly). ("ideogram", 48, 7.0), # SDXL: Turbo distilled; base wants ~30 steps + CFG ~7. "sdxl-turbo" precedes "sdxl". ("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 naming a known model wins (repo id, then resolved base repo), so a local-path load still resolves via its base repo. Keys matched as substrings, most specific first.""" 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, base_repo: Optional[str] = None ) -> Optional[str]: """The hosted pre-quantized transformer repo for ``scheme`` in this family, or None. ``base_repo`` (when known) selects a variant-specific checkpoint first: a checkpoint is baked from ONE base's weights and the loader refuses it for any other base, so a variant without its own entry still returns the family default (harmless: the base_model_id validation then falls back to dense-quantise, exactly as before this table existed).""" base = (base_repo or "").strip().lower() if base: for entry_base, entry_scheme, repo_id in fam.prequant_variant_repos: if entry_base == base and entry_scheme == scheme: return repo_id 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 sd.cpp needs; without it the no-GPU route falls back to diffusers.""" 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`` (user-supplied) to a file under ``repo_root``, rejecting absolute paths and ``..`` escapes.""" 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 lexical guards miss a symlink escape). 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