From 5d80d30168d76c42f88c43027e5daa102c07ea6d Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 24 Jun 2026 21:42:57 -0300 Subject: [PATCH] Support Qwen-Image and FLUX.1 image models alongside Z-Image --- studio/backend/core/inference/diffusion.py | 55 ++++++++++-- .../core/inference/diffusion_families.py | 41 +++++++-- .../backend/tests/test_diffusion_backend.py | 27 +++++- .../assistant-ui/model-selector/pickers.tsx | 5 +- .../assistant-ui/model-selector/types.ts | 3 + .../src/features/images/images-page.tsx | 87 +++++++++++-------- 6 files changed, 162 insertions(+), 56 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 7e7ce2acfd..e7a0d5fa49 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -24,6 +24,7 @@ from loggers import get_logger from utils.hardware import clear_gpu_cache from .diffusion_families import ( + DiffusionFamily, detect_family, resolve_base_repo, resolve_local_gguf_child, @@ -135,7 +136,7 @@ class DiffusionBackend: raise ValueError( f"Could not infer a diffusion family for '{repo_id}'. Pass family_override (z-image)." ) - base = resolve_base_repo(fam, base_repo) + base = _resolve_base_repo(repo_id, base_repo, fam, hf_token) with self._lock: # Allow starting over a previously-failed load, but not over a live one. @@ -254,7 +255,7 @@ class DiffusionBackend: raise ValueError( f"Could not infer a diffusion family for '{repo_id}'. Pass family_override (z-image)." ) - base = resolve_base_repo(fam, base_repo) + base = _resolve_base_repo(repo_id, base_repo, fam, hf_token) device, dtype = self._pick_device_and_dtype() with self._lock: @@ -306,8 +307,11 @@ class DiffusionBackend: negative_prompt: Optional[str] = None, width: int = 1024, height: int = 1024, - steps: int = 9, # Z-Image-Turbo: 9 steps = 8 DiT forwards (official default). - guidance: float = 0.0, # Turbo is distilled CFG-free; guidance must be 0. + # Fallbacks for a caller that passes nothing; the route always sends the + # per-model values the UI seeds (few steps / no CFG for distilled models, + # more steps / real CFG for full ones). + steps: int = 9, + guidance: float = 0.0, seed: Optional[int] = None, batch_size: int = 1, ) -> dict[str, Any]: @@ -332,13 +336,19 @@ class DiffusionBackend: "width": width, "height": height, "num_inference_steps": steps, - "guidance_scale": guidance, + # Most pipelines take guidance via "guidance_scale"; Qwen-Image + # uses "true_cfg_scale" (its distilled guidance is off). + state.family.cfg_kwarg: guidance, "generator": generator, # Generate the whole batch in one forward pass (VRAM-heavy). All # share this call's seed, drawn sequentially from one generator. "num_images_per_prompt": batch_size, } - if negative_prompt: + # Pipelines vary in which kwargs they accept (a distilled pipeline may + # take neither a negative prompt nor a step callback), so only pass + # those where the signature has them. + call_params = inspect.signature(state.pipe.__call__).parameters + if negative_prompt and "negative_prompt" in call_params: kwargs["negative_prompt"] = negative_prompt gen = _GenState(total_steps = steps) @@ -351,9 +361,7 @@ class DiffusionBackend: gen.eta_seconds = _estimate_eta(gen.total_steps, gen.step, gen.first_step_at, now) return callback_kwargs - # Not every pipeline accepts the callback; only pass it where supported - # so the step counter never breaks generation. - if "callback_on_step_end" in inspect.signature(state.pipe.__call__).parameters: + if "callback_on_step_end" in call_params: kwargs["callback_on_step_end"] = _on_step self._gen = gen @@ -417,6 +425,35 @@ class DiffusionBackend: } +def _resolve_base_repo( + repo_id: str, base_repo: Optional[str], fam: DiffusionFamily, hf_token: Optional[str] +) -> str: + """The companion diffusers repo: caller's base, else the GGUF repo's own + ``base_model`` tag, else the family fallback. Shared by both load paths so a + direct ``load_pipeline`` call resolves the variant base the same way.""" + return resolve_base_repo(fam, (base_repo or "").strip() or _hf_base_model(repo_id, hf_token)) + + +def _hf_base_model(repo_id: str, hf_token: Optional[str]) -> Optional[str]: + """The diffusers base repo from a GGUF repo's ``base_model`` tag, or None. + + Lets one family entry cover every variant (Turbo/full, schnell/dev, the + 2512 Qwen revision). Skipped for local paths; None on any lookup failure. + """ + if Path(repo_id).expanduser().exists(): + return None + try: + from huggingface_hub import HfApi + + meta = HfApi().model_info(repo_id, token = hf_token).cardData or {} + except Exception: # noqa: BLE001 — best-effort; fall back to the family default + return None + base = meta.get("base_model") + if isinstance(base, list): + base = base[0] if base else None + return base if isinstance(base, str) and base.strip() else None + + def _base_file_downloaded(rfilename: str) -> bool: """True for base-repo files ``from_pretrained`` actually fetches. diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index b9c2d5170c..44676f189b 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -25,30 +25,55 @@ class DiffusionFamily: pipeline_class: str transformer_class: str base_repo: str + # Pipeline kwarg carrying the guidance value. Most use "guidance_scale"; + # Qwen-Image's distilled guidance is off, so its real CFG is "true_cfg_scale". + cfg_kwarg: str = "guidance_scale" # Extra lowercased substrings (besides ``name``) that map a repo id here. aliases: tuple[str, ...] = field(default_factory = tuple) -# MVP: Z-Image-Turbo only. Its single-file GGUF is the transformer (Lumina2 DiT); -# the VAE / text-encoder / scheduler come from the base diffusers repo. More -# families (FLUX, Qwen-Image) can be appended here when we ship them. +# Keyed by architecture, not per model variant: a checkpoint's specific base repo +# is read from its HF base_model tag at load time, so one entry covers Turbo/full, +# schnell/dev, etc. base_repo here is only a fallback. Only archs whose diffusers +# transformer supports from_single_file load here (ERNIE-Image does not, yet); +# FLUX.2 is also excluded — its Mistral text-encoder chat template is incompatible +# with the pinned transformers. _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"), + ), + DiffusionFamily( + name = "qwen-image", + pipeline_class = "QwenImagePipeline", + transformer_class = "QwenImageTransformer2DModel", + base_repo = "Qwen/Qwen-Image", + cfg_kwarg = "true_cfg_scale", + aliases = ("qwen_image", "qwenimage"), + ), DiffusionFamily( name = "z-image", pipeline_class = "ZImagePipeline", transformer_class = "ZImageTransformer2DModel", base_repo = "Tongyi-MAI/Z-Image-Turbo", - aliases = ("z-image-turbo", "zimage", "z_image"), + aliases = ("zimage", "z_image"), ), ) +# Editing / inpaint checkpoints share an arch keyword but need a different +# pipeline and an input image, which this text-to-image backend doesn't drive. +_EDIT_KEYWORDS = ("edit", "kontext", "inpaint") + def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[DiffusionFamily]: """Resolve a ``DiffusionFamily`` from a repo id, or an explicit override. ``override`` matches a family ``name`` or alias exactly; otherwise the repo - id is scanned for the first family whose name/alias appears in it. Returns - ``None`` when nothing matches so the caller can raise a clean error. + id is scanned for the first family whose name/alias appears in it. Image + editing checkpoints are rejected (None) since this backend is text-to-image. """ if override: key = override.strip().lower() @@ -57,6 +82,8 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff return fam return None needle = repo_id.lower() + if any(kw in needle for kw in _EDIT_KEYWORDS): + return None for fam in _FAMILIES: if fam.name in needle or any(alias in needle for alias in fam.aliases): return fam @@ -64,7 +91,7 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str: - """The companion diffusers repo: caller-supplied if given, else the family default.""" + """The companion diffusers repo: caller-supplied if given, else the family fallback.""" base = (base_repo or "").strip() return base or fam.base_repo diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index d01dce967a..c1f039b2f9 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -30,10 +30,20 @@ from core.inference.diffusion_families import ( def test_detect_family_from_repo_id(): + # Detection is by architecture; Turbo/full and schnell/dev map to one family. assert detect_family("unsloth/Z-Image-Turbo-GGUF").name == "z-image" assert detect_family("unsloth/Z-Image-GGUF").name == "z-image" + assert detect_family("unsloth/Qwen-Image-2512-GGUF").name == "qwen-image" + assert detect_family("unsloth/FLUX.1-schnell-GGUF").name == "flux.1" + # FLUX.2's text encoder isn't compatible yet, so it's not a supported family. + assert detect_family("unsloth/FLUX.2-klein-4B-GGUF") is None + # Qwen-Image guides via true_cfg_scale, not guidance_scale. + assert detect_family("unsloth/Qwen-Image-2512-GGUF").cfg_kwarg == "true_cfg_scale" + assert detect_family("unsloth/Z-Image-GGUF").cfg_kwarg == "guidance_scale" + # Image-editing checkpoints are rejected (text-to-image backend only). + assert detect_family("unsloth/Qwen-Image-Edit-2511-GGUF") is None + assert detect_family("unsloth/FLUX.1-Kontext-dev-GGUF") is None assert detect_family("meta-llama/Llama-3-8B") is None - assert detect_family("unsloth/FLUX.2-klein-4B-GGUF") is None # not in the MVP family table def test_detect_family_override(): @@ -226,6 +236,21 @@ def test_generate_without_load_raises(fake_runtime): backend.generate(prompt = "x") +def test_resolve_base_repo_prefers_caller_then_hf_tag_then_fallback(monkeypatch): + from core.inference import diffusion + from core.inference.diffusion_families import detect_family + + fam = detect_family("unsloth/Qwen-Image-2512-GGUF") + monkeypatch.setattr(diffusion, "_hf_base_model", lambda repo, tok: "Qwen/Qwen-Image-2512") + # Caller's explicit base wins and the HF tag is not consulted. + assert diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", "my/base", fam, None) == "my/base" + # No caller base: the repo's base_model tag (the variant base) is used. + assert diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", None, fam, None) == "Qwen/Qwen-Image-2512" + # No caller base and no tag: the family fallback. + monkeypatch.setattr(diffusion, "_hf_base_model", lambda repo, tok: None) + assert diffusion._resolve_base_repo("unsloth/Qwen-Image-2512-GGUF", " ", fam, None) == fam.base_repo + + def test_load_without_gguf_raises(): backend = DiffusionBackend() with pytest.raises(ValueError): diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 47e48d727d..b0fc7177c4 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -697,7 +697,7 @@ function GgufVariantExpander({ ); const handleVariantClick = useCallback( - (quant: string, downloaded?: boolean, sizeBytes?: number) => { + (quant: string, filename: string, downloaded?: boolean, sizeBytes?: number) => { // Only seed the staged context for picks whose weights are already on // disk. The staging effect short-circuits on a known contextLength // (pendingHasContext) before starting the download, so attaching it to an @@ -708,6 +708,7 @@ function GgufVariantExpander({ source: sourceOverride ?? (isLocalPath ? "local" : "hub"), isLora: false, ggufVariant: quant, + ggufFilename: filename, isDownloaded: isLocalPath ? true : downloaded, expectedBytes: sizeBytes, contextLength: isAvailable ? nativeContext : undefined, @@ -879,7 +880,7 @@ function GgufVariantExpander({ type="button" {...variantList.getOptionProps(variantOptionKey, false)} onClick={() => - handleVariantClick(v.quant, v.downloaded, v.size_bytes) + handleVariantClick(v.quant, v.filename, v.downloaded, v.size_bytes) } className={cn( "flex min-w-0 flex-1 items-center justify-between gap-2 rounded-full px-2 py-1 text-left text-sm transition-colors hover:bg-[#ececec] focus-visible:bg-[#ececec] focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring/45 dark:hover:bg-[var(--sidebar-accent)] dark:focus-visible:bg-[var(--sidebar-accent)]", diff --git a/studio/frontend/src/components/assistant-ui/model-selector/types.ts b/studio/frontend/src/components/assistant-ui/model-selector/types.ts index 6a86515267..98523f0a2d 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/types.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/types.ts @@ -29,6 +29,9 @@ export interface ModelSelectorChangeMeta { source: "hub" | "lora" | "exported" | "local" | "external"; isLora: boolean; ggufVariant?: string; + /** Exact GGUF filename for the picked quant (filenames don't always follow the + * repo name, e.g. FLUX.1-schnell -> flux1-schnell-*.gguf). */ + ggufFilename?: string; isDownloaded?: boolean; expectedBytes?: number; /** Native GGUF context, threaded so a staged pick can seed the slider. */ diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index ab0e3b2723..c1f6e160db 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -29,7 +29,6 @@ import { import { Slider } from "@/components/ui/slider"; import { Spinner } from "@/components/ui/spinner"; import { Textarea } from "@/components/ui/textarea"; -import { SectionCard } from "@/components/section-card"; import { InfoHint } from "@/components/ui/info-hint"; import { ModelSelector } from "@/components/assistant-ui/model-selector"; import { IMAGE_GEN_TASKS } from "@/components/assistant-ui/model-selector/pickers"; @@ -58,20 +57,42 @@ import { unloadDiffusionModel, } from "./api"; -// MVP: a single curated diffusion GGUF. The chat ModelSelector lists/picks its -// quants; the base diffusers repo (VAE/text-encoder) is resolved server-side. -const MODEL = { - repo_id: "unsloth/Z-Image-Turbo-GGUF", - label: "Z-Image-Turbo", - family: "z-image", - // Hub-canonical filename pattern: z-image-turbo-.gguf - ggufFor: (quant: string) => `z-image-turbo-${quant}.gguf`, -}; - +// Curated diffusion GGUFs the picker recommends. The backend resolves each one's +// pipeline + base diffusers repo from its repo id, so the rail just lists them; +// the chat ModelSelector also surfaces any other on-device image GGUF. +const txt2img = (id: string, name: string): ModelOption => ({ + id, + name, + description: "Text-to-image · GGUF", + isGguf: true, +}); const MODELS: ModelOption[] = [ - { id: MODEL.repo_id, name: MODEL.label, description: "Text-to-image · GGUF", isGguf: true }, + txt2img("unsloth/Z-Image-Turbo-GGUF", "Z-Image-Turbo"), + txt2img("unsloth/Z-Image-GGUF", "Z-Image"), + txt2img("unsloth/Qwen-Image-2512-GGUF", "Qwen-Image 2512"), + txt2img("unsloth/Qwen-Image-GGUF", "Qwen-Image"), + txt2img("unsloth/FLUX.1-schnell-GGUF", "FLUX.1 schnell"), + txt2img("unsloth/FLUX.1-dev-GGUF", "FLUX.1 dev"), ]; +// Per-model generation defaults (steps + guidance), matched by repo-id substring, +// most specific first. Distilled "turbo/schnell" models want few steps and little +// guidance; the full "dev" models want more steps and real CFG. +const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }> = [ + { match: "z-image-turbo", steps: 9, guidance: 0 }, + { match: "flux.1-schnell", steps: 4, guidance: 0 }, + { match: "flux.1", steps: 28, guidance: 3.5 }, + { match: "qwen-image", steps: 20, guidance: 4 }, + { match: "z-image", steps: 20, guidance: 4 }, +]; + +function defaultsFor(repoId: string): { steps: number; guidance: number } { + const id = repoId.toLowerCase(); + // Fallback (a curated entry covers every model in MODELS) is the distilled + // few-step / no-CFG shape, hit only for an unrecognised on-device image GGUF. + return MODEL_DEFAULTS.find((d) => id.includes(d.match)) ?? { steps: 9, guidance: 0 }; +} + // Common aspect ratios (landscape; Flip gives the portrait mirror). Picking one // locks the W:H proportion; the sliders set the size. const ASPECT_RATIOS: Record = { @@ -181,7 +202,7 @@ function loadToastDescription(p: DiffusionLoadProgress) { return ( { + async (repoId: string, ggufFilename: string) => { // Cancel any prior poll loop so two can't run at once. if (pollTimer.current) clearTimeout(pollTimer.current); setBusy("loading"); @@ -587,11 +608,8 @@ export function ImagesPage() { loadToastId.current = toast(null, loadToastArgs(IDLE_PROGRESS)); try { // Returns immediately — the load runs in the background; we poll for it. - await loadDiffusionModel({ - model_path: MODEL.repo_id, - gguf_filename: ggufFilename, - family_override: MODEL.family, - }); + // The backend infers the family + base diffusers repo from the repo id. + await loadDiffusionModel({ model_path: repoId, gguf_filename: ggufFilename }); } catch (err) { dismissLoadToast(); toast.error(err instanceof Error ? err.message : "Failed to start load"); @@ -604,16 +622,16 @@ export function ImagesPage() { [pollLoadProgress, refreshStatus, dismissLoadToast], ); - // The chat picker emits (modelId, picked quant); load its matching GGUF. + // The chat picker emits (modelId, picked quant + its exact filename); load it, + // and seed the inputs with that model's defaults. const handleModelSelect = useCallback( (id: string, meta: ModelSelectorChangeMeta) => { - if (!meta.ggufVariant) return; // a non-quant pick; ignore - if (id !== MODEL.repo_id) { - toast.error("Only Z-Image-Turbo is supported right now."); - return; - } + if (!meta.ggufVariant || !meta.ggufFilename) return; // not a quant pick setQuant(meta.ggufVariant); - void handleLoad(MODEL.ggufFor(meta.ggufVariant)); + const d = defaultsFor(id); + setSteps(d.steps); + setGuidance(d.guidance); + void handleLoad(id, meta.ggufFilename); }, [handleLoad], ); @@ -725,13 +743,9 @@ export function ImagesPage() { {/* ── Controls rail + preview canvas. Padding mirrors the other tabs (Export, Data Recipes): px-5 / sm:px-9, with a roomy bottom. ── */}
- } - title="Generate" - description="Prompt and settings" - accent="indigo" - className="w-[340px] shrink-0 gap-4 overflow-y-auto" - > + {/* The controls rail. Plain card (the gray surface) with no header — + the prompt + Generate button make the panel self-explanatory. */} +