From dfa9777fe3fdb03fe477e1c024deb7e4100ef141 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 09:58:01 +0000 Subject: [PATCH] Studio Images: add the FLUX.2-dev model family Loading unsloth/FLUX.2-dev-GGUF failed because detect_family knew only the Qwen3-based FLUX.2-klein, so FLUX.2-dev (the full, Mistral-based Flux2Pipeline) resolved to nothing and the load errored. Add a flux.2-dev family: Flux2Pipeline + Flux2Transformer2DModel over the black-forest-labs/FLUX.2-dev base repo (gated, reachable with an HF token), with its FLUX.2 32-channel VAE and Mistral text encoder wired for the sd-cli path from the open Comfy-Org/flux2-dev mirror. text-to-image only: diffusers 0.38 ships no Flux2 img2img / inpaint pipeline for dev. Frontend gets sensible dev defaults (28 steps, guidance 4), distinct from klein's turbo defaults. Verified live: GGUF load resolves the family + gated base repo and generates a real 1024x1024 image on GPU. --- .../core/inference/diffusion_families.py | 18 ++++++++++++++++++ studio/backend/core/inference/sd_cpp_args.py | 1 + studio/backend/tests/test_diffusion_backend.py | 9 +++++++-- .../src/features/images/images-page.tsx | 2 ++ 4 files changed, 28 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index a5ef4460c2..4a8bd1dd8a 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -138,6 +138,24 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( ("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 diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index ea259b9000..5b6407b970 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -40,6 +40,7 @@ from core.inference.diffusion_memory import ( _TE_FLAGS_BY_FAMILY: dict[str, tuple[str, ...]] = { "z-image": ("--llm",), "flux.2-klein": ("--llm",), + "flux.2-dev": ("--llm",), "qwen-image": ("--qwen2vl",), "flux.1": ("--clip_l", "--t5xxl"), } diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 3eb4cd34b9..931f22920b 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -45,8 +45,13 @@ def test_detect_family_from_repo_id(): assert klein.cfg_kwarg == "guidance_scale" # Both klein sizes share the one family (base repo resolved per-variant). assert detect_family("unsloth/FLUX.2-klein-9B-GGUF").name == "flux.2-klein" - # Only klein is wired up; the Mistral-based FLUX.2-dev base repo is gated. - assert detect_family("unsloth/FLUX.2-dev-GGUF") is None + # FLUX.2-dev is the Mistral-based Flux2Pipeline, a distinct family from klein; its + # gated base repo is reachable with an HF token. It must not collide with klein. + dev = detect_family("unsloth/FLUX.2-dev-GGUF") + assert dev.name == "flux.2-dev" + assert dev.pipeline_class == "Flux2Pipeline" + assert dev.base_repo == "black-forest-labs/FLUX.2-dev" + assert detect_family("black-forest-labs/FLUX.2-dev").name == "flux.2-dev" # 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" diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 04485702fd..0bddf05c45 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -194,6 +194,8 @@ const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }> { match: "kontext", steps: 28, guidance: 2.5 }, { match: "flux.1", steps: 28, guidance: 3.5 }, { match: "flux.2-klein", steps: 4, guidance: 0 }, + // FLUX.2-dev is the full (non-distilled) model: more steps + real guidance, unlike klein. + { match: "flux.2-dev", steps: 28, guidance: 4 }, { match: "qwen-image", steps: 20, guidance: 4 }, { match: "z-image", steps: 20, guidance: 4 }, ];