From 5a58e9df1347be4bc0ead322a5ef2b2919b31e9f Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 17 Jul 2026 10:49:04 +0000
Subject: [PATCH] Add FLUX.1 Krea dev to the image model catalog

Krea's guidance-distilled FLUX.1-dev finetune keeps the exact dev layout, so it
runs under the existing flux.1 family unchanged. Wire it up end to end:

- Catalog group with the gated official bf16 pipeline and the open QuantStack
  GGUF quants; the gated artifact is skipped on auto-routing when undownloaded.
- Trust the official repo for non-GGUF from_pretrained loads, next to the other
  black-forest-labs bases.
- Generation defaults: 28 steps at guidance 4.5 per the model card. The generic
  "krea" defaults key (Krea-2-Turbo's 8-step no-CFG recipe) used to swallow the
  id, which would have produced garbage output; the new flux.1-krea key precedes
  it on both the backend table and the images page table.
- The flux.1 prequant checkpoints are schnell-based; the loader's baked
  base_model_id validation refuses them for the Krea-dev base, so int8/fp8
  requests dense-quantize instead (covered by existing prequant tests).
---
 studio/backend/core/inference/diffusion.py    |  3 ++
 .../core/inference/diffusion_families.py      |  4 +++
 .../tests/test_diffusion_more_families.py     | 32 +++++++++++++++++++
 .../model-selector/model-catalog.check.ts     | 13 ++++++++
 .../model-selector/model-catalog.ts           | 13 ++++++++
 .../src/features/images/images-page.tsx       |  4 +++
 6 files changed, 69 insertions(+)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 8366e01502..02ab08f4bf 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -285,6 +285,9 @@ _TRUSTED_NON_GGUF_REPOS = frozenset(
         "black-forest-labs/flux.1-dev",
         "black-forest-labs/flux.1-schnell",
         "black-forest-labs/flux.1-kontext-dev",
+        # Krea's guidance-distilled FLUX.1-dev finetune: same arch/layout as dev (FluxPipeline,
+        # CLIP+T5+ae), gated like dev. Detected as the flux.1 family via the "flux.1" token.
+        "black-forest-labs/flux.1-krea-dev",
         "tongyi-mai/z-image-turbo",
         "qwen/qwen-image",
         "qwen/qwen-image-2512",
diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py
index 2fcbbff539..1dffd37e74 100644
--- a/studio/backend/core/inference/diffusion_families.py
+++ b/studio/backend/core/inference/diffusion_families.py
@@ -454,6 +454,10 @@ def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
 # 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.
diff --git a/studio/backend/tests/test_diffusion_more_families.py b/studio/backend/tests/test_diffusion_more_families.py
index b2ffcd093e..9bda972c5e 100644
--- a/studio/backend/tests/test_diffusion_more_families.py
+++ b/studio/backend/tests/test_diffusion_more_families.py
@@ -53,6 +53,38 @@ def test_ideogram4_repos_are_trusted_non_gguf():
     assert not _is_trusted_diffusion_repo("ideogram-ai/some-future-repo")
 
 
+# ── FLUX.1 Krea dev (flux.1 family variant) ──────────────────────────────────
+@pytest.mark.parametrize(
+    "repo_id",
+    [
+        "black-forest-labs/FLUX.1-Krea-dev",
+        "QuantStack/FLUX.1-Krea-dev-GGUF",
+        # A local GGUF pick where the family keyword lives in the filename.
+        "QuantStack/FLUX.1-Krea-dev-GGUF/flux1-krea-dev-Q4_K_M.gguf",
+    ],
+)
+def test_detect_family_flux1_krea_dev(repo_id):
+    # Krea's FLUX.1-dev finetune keeps the exact dev layout, so it must resolve to the
+    # existing flux.1 family (FluxPipeline), never to krea-2 (a different arch).
+    fam = detect_family(repo_id)
+    assert fam is not None and fam.name == "flux.1"
+    assert fam.pipeline_class == "FluxPipeline"
+
+
+def test_flux1_krea_dev_is_trusted_non_gguf():
+    # The gated official pipeline loads via from_pretrained -> needs the allowlist.
+    assert _is_trusted_diffusion_repo("black-forest-labs/FLUX.1-Krea-dev")
+
+
+def test_flux1_krea_dev_generation_defaults():
+    # Model-card recipe: 28 steps at guidance 4.5. The generic "krea" key (Krea-2-Turbo's
+    # 8-step no-CFG shape) must NOT swallow it, and the krea-2 defaults must stay intact.
+    assert default_generation_params("black-forest-labs/FLUX.1-Krea-dev") == (28, 4.5)
+    assert default_generation_params("QuantStack/FLUX.1-Krea-dev-GGUF") == (28, 4.5)
+    assert default_generation_params("krea/Krea-2-Turbo") == (8, 0.0)
+    assert default_generation_params("krea/Krea-2-Raw") == (52, 3.5)
+
+
 def test_ideogram4_generation_defaults():
     # Model-card settings: 48 steps, guidance 7 (the backend keeps the pipeline's
     # recommended tapered schedule when the request matches exactly).
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts
index 9db3dae9e5..d7727491a2 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts
+++ b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts
@@ -347,6 +347,19 @@ assert.equal(
   }).repoId,
   "black-forest-labs/FLUX.1-dev",
 );
+// FLUX.1 Krea dev: gated BF16 skipped when undownloaded -> the open QuantStack GGUF; the
+// GGUF repo id also resolves to the group (cross-owner via the artifact list).
+const kreaDevRoute = groupForRepoId("black-forest-labs/FLUX.1-Krea-dev", IMAGE_CATALOG);
+assert.ok(kreaDevRoute);
+assert.equal(
+  pickDefaultArtifact(kreaDevRoute, { gpuGb: 80, systemRamGb: 128, isDownloaded: notDownloaded })
+    .repoId,
+  "QuantStack/FLUX.1-Krea-dev-GGUF",
+);
+assert.equal(
+  groupForRepoId("QuantStack/FLUX.1-Krea-dev-GGUF", IMAGE_CATALOG),
+  kreaDevRoute,
+);
 // FLUX.1-schnell is Apache-2.0 (not gated): its BF16 IS auto-routed on a GPU that fits it.
 const fluxSchnellRoute = groupForRepoId("unsloth/FLUX.1-schnell", IMAGE_CATALOG);
 assert.ok(fluxSchnellRoute);
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts
index ee6decc2c1..34451e3897 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts
+++ b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts
@@ -198,6 +198,19 @@ export const IMAGE_CATALOG: CatalogGroup[] = [
       gguf("unsloth/FLUX.1-dev-GGUF"),
     ],
   },
+  {
+    // Krea's guidance-distilled FLUX.1-dev finetune ("opinionated aesthetics"): same
+    // arch/layout as FLUX.1-dev, so it runs under the existing flux.1 family. The base
+    // repo is gated like dev; QuantStack publishes the open GGUF quants.
+    canonicalId: "black-forest-labs/FLUX.1-Krea-dev",
+    displayName: "FLUX.1 Krea dev",
+    description: "Text-to-image",
+    scope: "image",
+    artifacts: [
+      bf16Pipeline("black-forest-labs/FLUX.1-Krea-dev", 32, { gated: true }),
+      gguf("QuantStack/FLUX.1-Krea-dev-GGUF"),
+    ],
+  },
   {
     canonicalId: "unsloth/FLUX.2-klein-4B",
     displayName: "FLUX.2 klein 4B",
diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx
index 88aa810ef1..7b6e8ce71a 100644
--- a/studio/frontend/src/features/images/images-page.tsx
+++ b/studio/frontend/src/features/images/images-page.tsx
@@ -155,6 +155,10 @@ const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }>
   { match: "flux.1-schnell", steps: 4, guidance: 0 },
   // Kontext (editing) before the generic flux.1: ~28 steps, lower guidance (~2.5).
   { match: "kontext", steps: 28, guidance: 2.5 },
+  // Krea's FLUX.1-dev finetune runs its card recipe (28 steps, guidance 4.5); before
+  // the generic flux.1 key. It never hits the krea-2 keys above ("krea-2" is not a
+  // substring of "flux.1-krea-dev").
+  { match: "flux.1-krea", steps: 28, guidance: 4.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.