From fca32d5a5b86db80be9468b26ef4d17e73a05b4b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 3 Jul 2026 13:25:54 +0000
Subject: [PATCH 01/17] Add Krea 2 Turbo: diffusion family, inference loader,
 LoRA training (diffusers 0.39)

Inference:
- krea-2 DiffusionFamily (Krea2Pipeline / Krea2Transformer2DModel, base
  krea/Krea-2-Turbo, bf16 only, no GGUF/sd.cpp mapping yet)
- Per-component pipeline loader (core/inference/diffusion_krea2.py): the krea
  repo is exported with transformers 5.2, so the tokenizer config
  (extra_special_tokens as a list, no slow-tokenizer vocab files) and the
  text encoder rope settings (rope_parameters vs rope_scaling) need explicit
  compat on the 4.x line; values are copied verbatim and equal the 4.x
  Qwen3-VL defaults, so the math is unchanged. from_pretrained also
  type-checks the tokenizer against the declared slow class, so the pipeline
  is assembled through its constructor with the model_index init config
  (is_distilled carries Turbo's fixed mu=1.15 schedule)
- Trust allowlist entry, curated picker entry + 8 step / cfg 0 defaults,
  int8 exclusion token for the M=1 Krea2TimestepEmbedding projection

Training:
- krea-2 _FamilySpec in the DiT trainer: phased conditioning/transformer
  load through the compat loader, shared Qwen-Image VAE latent path,
  fixed-512 text embeds (static shapes, plain concat collate), inline 2x2
  latent packing + shared position grid, the authors' recommended LoRA
  target set and rank/alpha 32, lr 3e-4, 512px presets
- GPU smokes on B200: nf4 2.9 steps/s at 11.5 GB, bf16 3.4 steps/s at
  30.1 GB, bf16 + regional compile 5.2 steps/s; adapter round-trip
  generation verified
---
 studio/backend/core/inference/diffusion.py    |  31 ++-
 .../core/inference/diffusion_families.py      |  19 ++
 .../backend/core/inference/diffusion_krea2.py | 137 ++++++++++++
 .../inference/diffusion_transformer_quant.py  |   4 +
 .../core/training/diffusion_dit_trainer.py    | 133 +++++++++++
 .../core/training/diffusion_train_common.py   |   9 +-
 .../backend/tests/test_diffusion_backend.py   |  10 +-
 .../tests/test_diffusion_dit_trainer.py       |  11 +-
 studio/backend/tests/test_diffusion_krea2.py  | 209 ++++++++++++++++++
 .../src/features/images/images-page.tsx       |   6 +
 10 files changed, 554 insertions(+), 15 deletions(-)
 create mode 100644 studio/backend/core/inference/diffusion_krea2.py
 create mode 100644 studio/backend/tests/test_diffusion_krea2.py

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 28eafe6829..3418d6f23b 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -42,6 +42,7 @@ from .diffusion_device import (
     diffusion_device_target_from_torch_device,
     resolve_diffusion_device_target,
 )
+from .diffusion_krea2 import KREA2_FAMILY_NAME, load_krea2_pipeline
 from .diffusion_memory import (
     OFFLOAD_NONE,
     apply_memory_plan,
@@ -203,6 +204,9 @@ _TRUSTED_NON_GGUF_REPOS = frozenset(
         "black-forest-labs/flux.1-dev",
         "tongyi-mai/z-image-turbo",
         "qwen/qwen-image",
+        # Krea 2 Turbo: official vendor repo, safetensors-only, no remote code. Loaded
+        # per-component via core/inference/diffusion_krea2.py (no GGUF variant yet).
+        "krea/krea-2-turbo",
     }
 )
 
@@ -1004,10 +1008,16 @@ class DiffusionBackend:
                         # (transformer + VAE + text encoders + scheduler) from the repo
                         # and re-applies any embedded quantization_config (e.g. bnb-4bit),
                         # so a pre-quantized pipeline reloads quantized with no extra config.
-                        pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype}
-                        if hf_token:
-                            pipe_kwargs["token"] = hf_token
-                        pipe = pipeline_cls.from_pretrained(repo_id, **pipe_kwargs)
+                        if fam.name == KREA2_FAMILY_NAME:
+                            # The krea repo ships transformers-5.x style configs the 4.x
+                            # line cannot parse; assemble the pipeline per-component
+                            # (see diffusion_krea2.py for the exact compat story).
+                            pipe = load_krea2_pipeline(repo_id, dtype, hf_token = hf_token)
+                        else:
+                            pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype}
+                            if hf_token:
+                                pipe_kwargs["token"] = hf_token
+                            pipe = pipeline_cls.from_pretrained(repo_id, **pipe_kwargs)
                     elif kind == "single_file" and fam.single_file_is_pipeline:
                         # A single-file SDXL-style checkpoint is the WHOLE pipeline
                         # (U-Net + VAE + both text encoders), not a transformer-only file,
@@ -1040,10 +1050,15 @@ class DiffusionBackend:
                             single_file_path, **sf_kwargs
                         )
 
-                        pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer}
-                        if hf_token:
-                            pipe_kwargs["token"] = hf_token
-                        pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
+                        if fam.name == KREA2_FAMILY_NAME:
+                            pipe = load_krea2_pipeline(
+                                base, dtype, hf_token = hf_token, transformer = transformer
+                            )
+                        else:
+                            pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer}
+                            if hf_token:
+                                pipe_kwargs["token"] = hf_token
+                            pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
 
                 # Resolve the effective speed mode: GGUF models default to the
                 # near-lossless `default` profile (compile is ~2.2x and sits below
diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py
index 775be1c962..65eefcbae9 100644
--- a/studio/backend/core/inference/diffusion_families.py
+++ b/studio/backend/core/inference/diffusion_families.py
@@ -282,6 +282,25 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
             ("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors", "llm"),
         ),
     ),
+    # Krea 2 (diffusers >= 0.39): a ~12B single-stream flow-matching DiT with a
+    # Qwen3-VL-4B text encoder (12 tapped hidden layers fused in-transformer) and the
+    # Qwen-Image VAE. Loaded per-component through core/inference/diffusion_krea2.py
+    # because the krea repo ships transformers-5.x style configs (see that module).
+    # No GGUF/sd.cpp mapping yet, so the no-GPU route falls back to diffusers.
+    DiffusionFamily(
+        name = "krea-2",
+        pipeline_class = "Krea2Pipeline",
+        transformer_class = "Krea2Transformer2DModel",
+        base_repo = "krea/Krea-2-Turbo",
+        aliases = ("krea2",),
+        # LoRA training via the DiT trainer (no prequant repo yet, so nf4 quantizes the
+        # 12B transformer on the fly under the default precision).
+        trainable = True,
+        train_base_repos = ("krea/Krea-2-Turbo",),
+        # The checkpoint is exported bf16-only (the model card pins bfloat16); fp16 is
+        # unvalidated upstream, so keep the fp16 fallback off like z-image.
+        fp16_incompatible = True,
+    ),
     # 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.
diff --git a/studio/backend/core/inference/diffusion_krea2.py b/studio/backend/core/inference/diffusion_krea2.py
new file mode 100644
index 0000000000..312a4c8aa2
--- /dev/null
+++ b/studio/backend/core/inference/diffusion_krea2.py
@@ -0,0 +1,137 @@
+"""Krea 2 pipeline loader: assembles ``Krea2Pipeline`` from per-component loads.
+
+Why not ``Krea2Pipeline.from_pretrained``: the ``krea/Krea-2-Turbo`` repo was exported
+with transformers 5.2, and two of its configs use 5.x-only conventions that the 4.x
+line cannot parse:
+
+- ``tokenizer/tokenizer_config.json`` declares ``Qwen2Tokenizer`` (slow -- 5.x unified
+  slow/fast under the plain name) but ships only ``tokenizer.json``. 4.x's slow class
+  needs vocab.json/merges.txt (absent), and its fast class trips over
+  ``extra_special_tokens`` stored as a LIST (4.x expects a dict). Loading the fast
+  class with an explicit ``extra_special_tokens = {}`` override is id-identical: every
+  listed token is already registered as an added special token inside tokenizer.json,
+  and the pipeline templates prompts manually (it never uses a chat template).
+- ``text_encoder/config.json`` keeps the rope settings under ``rope_parameters`` (the
+  5.x name). 4.x reads ``rope_scaling`` + a top-level ``rope_theta`` and crashes on the
+  missing key (``NoneType.get``). The values are copied across verbatim -- and they
+  equal 4.x's Qwen3-VL defaults (theta 5e6, mrope_section [24, 20, 20], interleaved
+  mrope applied unconditionally), so the rotary embedding is numerically identical.
+  The state dict itself round-trips 1:1 (checkpoint keys == 4.x module keys).
+
+``from_pretrained`` additionally type-checks a passed ``tokenizer`` against the
+declared SLOW class (a fast tokenizer does not subclass it), so the pipeline is built
+through its constructor instead, forwarding the ``is_distilled`` /
+``text_encoder_select_layers`` / ``patch_size`` init config from model_index.json --
+Turbo's fixed mu=1.15 timestep shift rides on ``is_distilled = True``, so dropping it
+would silently degrade the schedule.
+
+Both workarounds are self-disabling on a transformers 5.x runtime: the plain tokenizer
+load succeeds (no fallback taken) and ``rope_scaling`` parses non-None (no patch).
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Optional
+
+from loggers import get_logger
+
+logger = get_logger(__name__)
+
+KREA2_FAMILY_NAME = "krea-2"
+
+
+def load_krea2_tokenizer(repo_id: str, hf_token: Optional[str] = None):
+    """The Krea 2 tokenizer, tolerating the repo's transformers-5.x tokenizer config."""
+    from transformers import AutoTokenizer
+
+    kwargs: dict[str, Any] = {"subfolder": "tokenizer"}
+    if hf_token:
+        kwargs["token"] = hf_token
+    try:
+        return AutoTokenizer.from_pretrained(repo_id, **kwargs)
+    except Exception as exc:  # noqa: BLE001 -- 4.x config-parse failure, retry with override
+        logger.info("diffusion.krea2 tokenizer compat fallback: %s", exc)
+        return AutoTokenizer.from_pretrained(repo_id, extra_special_tokens = {}, **kwargs)
+
+
+def remap_rope_parameters(text_config) -> None:
+    """Copy 5.x ``rope_parameters`` onto the 4.x ``rope_scaling`` / ``rope_theta`` slots
+    in place. A no-op when ``rope_scaling`` already parsed non-None (a 5.x runtime) or
+    the config carries no ``rope_parameters`` dict."""
+    rope_parameters = getattr(text_config, "rope_parameters", None)
+    if getattr(text_config, "rope_scaling", None) is None and isinstance(rope_parameters, dict):
+        text_config.rope_scaling = {
+            k: v for k, v in rope_parameters.items() if k != "rope_theta"
+        }
+        if "rope_theta" in rope_parameters:
+            text_config.rope_theta = rope_parameters["rope_theta"]
+
+
+def load_krea2_text_encoder(repo_id: str, dtype, hf_token: Optional[str] = None):
+    """The Qwen3-VL text encoder, remapping 5.x ``rope_parameters`` for a 4.x runtime."""
+    from transformers import AutoConfig, Qwen3VLModel
+
+    kwargs: dict[str, Any] = {"subfolder": "text_encoder"}
+    if hf_token:
+        kwargs["token"] = hf_token
+    config = AutoConfig.from_pretrained(repo_id, **kwargs)
+    remap_rope_parameters(getattr(config, "text_config", config))
+    return Qwen3VLModel.from_pretrained(repo_id, config = config, dtype = dtype, **kwargs)
+
+
+def _load_model_index(repo_id: str, hf_token: Optional[str] = None) -> dict[str, Any]:
+    """model_index.json as a dict, from a local path or the Hub cache."""
+    try:
+        local = Path(repo_id).expanduser() / "model_index.json"
+        if local.is_file():
+            return json.loads(local.read_text())
+    except OSError:
+        pass
+    from huggingface_hub import hf_hub_download
+
+    path = hf_hub_download(repo_id, "model_index.json", token = hf_token or None)
+    return json.loads(Path(path).read_text())
+
+
+def load_krea2_pipeline(
+    repo_id: str,
+    dtype,
+    hf_token: Optional[str] = None,
+    transformer = None,
+    with_transformer: bool = True,
+):
+    """A ready ``Krea2Pipeline`` for ``repo_id`` (still on CPU; caller places it).
+
+    ``transformer`` lets the single-file/quant paths hand in a prebuilt denoiser;
+    ``with_transformer = False`` skips the (26 GB) denoiser entirely for a
+    conditioning-only pipeline (the trainer's phased load). The remaining components
+    (VAE, text encoder, tokenizer, scheduler) come from the repo.
+    """
+    import diffusers
+
+    token = hf_token or None
+    tokenizer = load_krea2_tokenizer(repo_id, hf_token = token)
+    text_encoder = load_krea2_text_encoder(repo_id, dtype, hf_token = token)
+    scheduler = diffusers.FlowMatchEulerDiscreteScheduler.from_pretrained(
+        repo_id, subfolder = "scheduler", token = token
+    )
+    vae = diffusers.AutoencoderKLQwenImage.from_pretrained(
+        repo_id, subfolder = "vae", torch_dtype = dtype, token = token
+    )
+    if transformer is None and with_transformer:
+        transformer = diffusers.Krea2Transformer2DModel.from_pretrained(
+            repo_id, subfolder = "transformer", torch_dtype = dtype, token = token
+        )
+    model_index = _load_model_index(repo_id, hf_token = token)
+    return diffusers.Krea2Pipeline(
+        scheduler = scheduler,
+        vae = vae,
+        text_encoder = text_encoder,
+        tokenizer = tokenizer,
+        transformer = transformer,
+        text_encoder_select_layers = model_index.get("text_encoder_select_layers"),
+        is_distilled = bool(model_index.get("is_distilled", False)),
+        patch_size = int(model_index.get("patch_size", 2)),
+    )
diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py
index 9e2e56a2a5..7dcfc95a73 100644
--- a/studio/backend/core/inference/diffusion_transformer_quant.py
+++ b/studio/backend/core/inference/diffusion_transformer_quant.py
@@ -60,6 +60,10 @@ _INT8_EXCLUDE_NAME_TOKENS = (
     "guidance_embed",
     "time_text_embed",  # Flux/Qwen time_text_embed.* (pooled-text + timestep); NOT context_embedder
     "pooled",
+    # Krea 2's Krea2TimestepEmbedding ("time_embed.linear_2", 6144->6144 at M = batch);
+    # its other M=1 projection ("time_mod_proj") is already caught by "_mod", and
+    # img_in / final_layer.linear / text_fusion.projector fall under min_features.
+    "time_embed",
 )
 
 
diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py
index 5776f87652..5e6ae46bf4 100644
--- a/studio/backend/core/training/diffusion_dit_trainer.py
+++ b/studio/backend/core/training/diffusion_dit_trainer.py
@@ -65,6 +65,26 @@ _FLUX_TARGETS = (
 )
 _QWEN_TARGETS = _FLUX_TARGETS
 _ZIMAGE_TARGETS = ("to_q", "to_k", "to_v", "to_out.0")
+# The Krea 2 authors' recommended default target set (their DreamBooth reference script):
+# attention + SwiGLU + the text-fusion projector + the conditioning embedders. For long
+# runs they suggest narrowing to the attention layers so prompt adherence doesn't drop.
+_KREA2_TARGETS = (
+    "img_in",
+    "final_layer.linear",
+    "to_q",
+    "to_k",
+    "to_v",
+    "to_out.0",
+    "to_gate",
+    "ff.up",
+    "ff.down",
+    "text_fusion.projector",
+    "txt_in.linear_1",
+    "txt_in.linear_2",
+    "time_embed.linear_1",
+    "time_embed.linear_2",
+    "time_mod_proj",
+)
 
 
 @dataclass
@@ -647,6 +667,105 @@ def _zimage_save(pipe_cls, out_dir, transformer_lora_layers):
     )
 
 
+# ── Krea 2 ────────────────────────────────────────────────────────────────────
+def _krea2_load_conditioners(cfg, device, weight_dtype):
+    # The krea repo ships transformers-5.x style configs the pinned 4.x line cannot
+    # parse, so the conditioning pipeline is assembled per-component (with the tokenizer
+    # and rope compat) instead of from_pretrained(transformer = None); see
+    # core/inference/diffusion_krea2.py for the exact compat story.
+    import torch
+    from core.inference.diffusion_krea2 import load_krea2_pipeline
+
+    pipe = load_krea2_pipeline(
+        cfg.base_model, torch.bfloat16, hf_token = cfg.hf_token, with_transformer = False
+    )
+    pipe.vae.to(device, dtype = torch.float32)
+    return pipe, pipe.vae
+
+
+def _krea2_load_transformer(cfg, device, weight_dtype, base_precision):
+    # The transformer subfolder is diffusers-format (no transformers compat needed).
+    # There is no prequant repo yet, so nf4 quantizes the 12B transformer on the fly.
+    from diffusers import Krea2Transformer2DModel
+    return _load_dit_transformer(Krea2Transformer2DModel, cfg, device, base_precision)
+
+
+def _krea2_encode_prompts(pipe, captions, device):
+    import torch
+
+    _encoders_to_device(pipe, device)
+    out = []
+    with torch.no_grad():
+        for cap in captions:
+            # encode_prompt pads/truncates to the fixed max_sequence_length, so every
+            # embed is [1, 512, num_text_layers, 2560] with a [1, 512] validity mask --
+            # static shapes (the padding sits mid-template, BEFORE the assistant suffix,
+            # matching how the model was sampled at training time).
+            pe, mask = pipe.encode_prompt(
+                prompt = cap,
+                device = device,
+                num_images_per_prompt = 1,
+                max_sequence_length = 512,
+            )
+            out.append((pe.cpu(), mask.cpu()))
+    return out
+
+
+# Krea 2 conditions on the Qwen-Image VAE (AutoencoderKLQwenImage) with the same
+# per-channel latents_mean / latents_std normalisation, so latent encoding is shared.
+_krea2_encode_latents = _qwen_encode_latents
+_krea2_encode_latent_stats = _qwen_encode_latent_stats
+
+
+def _krea2_collate(
+    entries,
+    device,
+    weight_dtype,
+    pad_to = None,
+):
+    import torch
+
+    # Fixed-length embeds (see _krea2_encode_prompts), so collation is a plain concat
+    # with the mask riding along; ``pad_to`` is moot because the shapes are static.
+    pe_b = torch.cat([e[0] for e in entries]).to(device = device, dtype = weight_dtype)
+    mask_b = torch.cat([e[1] for e in entries]).to(device)
+    return (pe_b, mask_b)
+
+
+def _krea2_forward(transformer, noisy, timesteps, sigmas, embeds_batch, cfg, device, weight_dtype):
+    from diffusers import Krea2Pipeline
+
+    pe, mask = embeds_batch
+    # [B,16,1,H,W] -> [B, (H/2)*(W/2), 64] 2x2 patches. Krea2Pipeline._pack_latents /
+    # _unpack_latents are instance methods (they read self.patch_size), so the packing is
+    # inlined here exactly like the reference DreamBooth script (patch_size = 2).
+    bsz, c, _f, h, w = noisy.shape
+    packed = noisy.reshape(bsz, c, h // 2, 2, w // 2, 2)
+    packed = packed.permute(0, 2, 4, 1, 3, 5).reshape(bsz, (h // 2) * (w // 2), c * 4)
+    # Text tokens sit at the rotary origin, so one shared position grid serves the batch.
+    position_ids = Krea2Pipeline.prepare_position_ids(pe.shape[1], h // 2, w // 2, device)
+    pred = transformer(
+        hidden_states = packed,
+        encoder_hidden_states = pe,
+        timestep = timesteps / 1000,
+        position_ids = position_ids,
+        encoder_attention_mask = mask,
+        return_dict = False,
+    )[0]
+    pred = pred.view(bsz, h // 2, w // 2, c, 2, 2)
+    pred = pred.permute(0, 3, 1, 4, 2, 5)
+    return pred.reshape(bsz, c, 1, h, w)
+
+
+def _krea2_save(pipe_cls, out_dir, transformer_lora_layers):
+    from diffusers import Krea2Pipeline
+    Krea2Pipeline.save_lora_weights(
+        save_directory = out_dir,
+        transformer_lora_layers = transformer_lora_layers,
+        weight_name = DEFAULT_LORA_FILENAME,
+    )
+
+
 _SPECS: dict[str, _FamilySpec] = {
     "flux.1": _FamilySpec(
         family = "flux.1",
@@ -690,6 +809,20 @@ _SPECS: dict[str, _FamilySpec] = {
         forward = _zimage_forward,
         save = _zimage_save,
     ),
+    "krea-2": _FamilySpec(
+        family = "krea-2",
+        lora_targets = _KREA2_TARGETS,
+        force_bf16 = True,
+        dense_bf16_gb = 26.3,
+        load_conditioners = _krea2_load_conditioners,
+        load_transformer = _krea2_load_transformer,
+        encode_prompts = _krea2_encode_prompts,
+        encode_latents = _krea2_encode_latents,
+        encode_latent_stats = _krea2_encode_latent_stats,
+        collate = _krea2_collate,
+        forward = _krea2_forward,
+        save = _krea2_save,
+    ),
 }
 
 
diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py
index 56b1020cd5..5cba60b5d4 100644
--- a/studio/backend/core/training/diffusion_train_common.py
+++ b/studio/backend/core/training/diffusion_train_common.py
@@ -159,7 +159,7 @@ def get_trainer(family: str) -> Callable[..., str]:
     if key == "sdxl":
         from core.training.diffusion_lora_trainer import run_diffusion_lora_training
         return run_diffusion_lora_training
-    if key in ("flux.1", "qwen-image", "z-image"):
+    if key in ("flux.1", "qwen-image", "z-image", "krea-2"):
         from core.training.diffusion_dit_trainer import run_dit_lora_training
         return run_dit_lora_training
     raise ValueError(f"No trainer is registered for family {family!r}.")
@@ -173,6 +173,9 @@ FAMILY_TRAIN_DEFAULTS: dict[str, dict[str, Any]] = {
     "flux.1": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 512},
     "qwen-image": {"lora_rank": 16, "learning_rate": 5e-5, "resolution": 512},
     "z-image": {"lora_rank": 16, "learning_rate": 1e-4, "resolution": 768},
+    # The Krea 2 authors' recommended starting point (their DreamBooth script defaults):
+    # rank/alpha 32, lr 3e-4, 512px.
+    "krea-2": {"lora_rank": 32, "learning_rate": 3e-4, "resolution": 512},
 }
 
 
@@ -188,6 +191,7 @@ _FAMILY_LABELS = {
     "flux.1": "FLUX.1-dev",
     "qwen-image": "Qwen-Image",
     "z-image": "Z-Image",
+    "krea-2": "Krea 2",
 }
 _FAMILY_VRAM_NOTES = {
     "sdxl": "Trains on ~12 GB+ (bf16 LoRA). The lightest, fastest option.",
@@ -197,6 +201,7 @@ _FAMILY_VRAM_NOTES = {
     ),
     "qwen-image": "20B model, QLoRA (nf4) by default (~24 GB+). The heaviest option.",
     "z-image": "6B model, QLoRA (nf4) by default (~12 GB+). bf16 only.",
+    "krea-2": "12B model, QLoRA (nf4) by default (~18 GB+). bf16 only.",
 }
 
 
@@ -215,7 +220,7 @@ def family_train_infos() -> list[dict[str, Any]]:
         repos = list(fam.train_base_repos) or [fam.base_repo]
         # base_precision / compile apply to the DiT trainer only; SDXL keeps its
         # mixed_precision lever, so the UI hides the selector for it.
-        is_dit = name in ("flux.1", "qwen-image", "z-image")
+        is_dit = name in ("flux.1", "qwen-image", "z-image", "krea-2")
         infos.append(
             {
                 "name": name,
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index d36e62f3cf..d2f0ea63fd 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -82,6 +82,14 @@ def test_detect_family_from_repo_id():
     assert detect_family("unsloth/FLUX.1-dev-GGUF").name == "flux.1"
     # A plain Qwen-Image checkpoint must still resolve to the base family, not edit.
     assert detect_family("unsloth/Qwen-Image-2512-GGUF").name == "qwen-image"
+    # Krea 2 (diffusers >= 0.39): bf16-only single-stream DiT, no GGUF/sd.cpp mapping.
+    krea2 = detect_family("krea/Krea-2-Turbo")
+    assert krea2.name == "krea-2"
+    assert krea2.pipeline_class == "Krea2Pipeline"
+    assert krea2.transformer_class == "Krea2Transformer2DModel"
+    assert krea2.cfg_kwarg == "guidance_scale"
+    assert krea2.fp16_incompatible is True
+    assert krea2.sd_cpp_text_encoders == ()
     assert detect_family("meta-llama/Llama-3-8B") is None
 
 
@@ -122,7 +130,7 @@ def test_detect_family_override():
 def test_supported_family_names():
     names = supported_family_names()
     # The unknown-model error lists these, so the key families must be present.
-    for expected in ("flux.1", "flux.2-klein", "flux.2-dev", "qwen-image", "z-image"):
+    for expected in ("flux.1", "flux.2-klein", "flux.2-dev", "qwen-image", "z-image", "krea-2"):
         assert expected in names
     # Every listed name is a valid family_override (round-trips through detect_family).
     for name in names:
diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py
index ba54480003..f4e9c804ca 100644
--- a/studio/backend/tests/test_diffusion_dit_trainer.py
+++ b/studio/backend/tests/test_diffusion_dit_trainer.py
@@ -21,15 +21,18 @@ from core.training.diffusion_dit_trainer import (
 from core.training.diffusion_train_common import DiffusionLoraConfig, family_train_infos
 
 
-def test_specs_cover_the_three_dit_families():
-    assert set(_SPECS) == {"flux.1", "qwen-image", "z-image"}
-    # FLUX / Qwen share the added-kv attention target set; Z-Image is single-stream.
+def test_specs_cover_the_dit_families():
+    assert set(_SPECS) == {"flux.1", "qwen-image", "z-image", "krea-2"}
+    # FLUX / Qwen share the added-kv attention target set; Z-Image and Krea 2 are
+    # single-stream.
     assert "add_q_proj" in _SPECS["flux.1"].lora_targets
     assert "add_q_proj" in _SPECS["qwen-image"].lora_targets
     assert "add_q_proj" not in _SPECS["z-image"].lora_targets
-    # Z-Image and Qwen are bf16-only.
+    assert "add_q_proj" not in _SPECS["krea-2"].lora_targets
+    # Z-Image, Qwen and Krea 2 are bf16-only.
     assert _SPECS["z-image"].force_bf16 is True
     assert _SPECS["qwen-image"].force_bf16 is True
+    assert _SPECS["krea-2"].force_bf16 is True
 
 
 @pytest.mark.parametrize(
diff --git a/studio/backend/tests/test_diffusion_krea2.py b/studio/backend/tests/test_diffusion_krea2.py
new file mode 100644
index 0000000000..b360ae43f9
--- /dev/null
+++ b/studio/backend/tests/test_diffusion_krea2.py
@@ -0,0 +1,209 @@
+"""Unit tests for the Krea 2 per-component pipeline loader (CPU-only, no network)."""
+
+from __future__ import annotations
+
+import json
+import sys
+from types import SimpleNamespace
+
+from core.inference.diffusion_krea2 import (
+    KREA2_FAMILY_NAME,
+    _load_model_index,
+    load_krea2_pipeline,
+    remap_rope_parameters,
+)
+
+
+# ── rope_parameters (transformers 5.x) -> rope_scaling (4.x) remap ──────────
+
+
+def test_remap_rope_parameters_copies_5x_values():
+    cfg = SimpleNamespace(
+        rope_scaling = None,
+        rope_theta = 1000000.0,
+        rope_parameters = {
+            "mrope_interleaved": True,
+            "mrope_section": [24, 20, 20],
+            "rope_theta": 5000000,
+            "rope_type": "default",
+        },
+    )
+    remap_rope_parameters(cfg)
+    # rope_theta is hoisted to the top-level slot, the rest lands in rope_scaling.
+    assert cfg.rope_theta == 5000000
+    assert cfg.rope_scaling == {
+        "mrope_interleaved": True,
+        "mrope_section": [24, 20, 20],
+        "rope_type": "default",
+    }
+
+
+def test_remap_rope_parameters_noop_on_5x_runtime_or_plain_4x_config():
+    # rope_scaling already parsed (5.x runtime exposing the alias): untouched.
+    parsed = {"rope_type": "default", "mrope_section": [1, 2, 3]}
+    cfg = SimpleNamespace(rope_scaling = parsed, rope_theta = 7.0, rope_parameters = {"x": 1})
+    remap_rope_parameters(cfg)
+    assert cfg.rope_scaling is parsed
+    assert cfg.rope_theta == 7.0
+    # No rope_parameters at all (a plain 4.x-exported config): untouched.
+    cfg = SimpleNamespace(rope_scaling = None, rope_theta = 7.0)
+    remap_rope_parameters(cfg)
+    assert cfg.rope_scaling is None
+
+
+# ── model_index.json resolution ──────────────────────────────────────────────
+
+
+def test_load_model_index_from_local_path(tmp_path):
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"is_distilled": True, "patch_size": 2})
+    )
+    assert _load_model_index(str(tmp_path)) == {"is_distilled": True, "patch_size": 2}
+
+
+# ── pipeline assembly threads the model_index init config ────────────────────
+
+
+def test_load_krea2_pipeline_threads_init_config(monkeypatch, tmp_path):
+    (tmp_path / "model_index.json").write_text(
+        json.dumps(
+            {
+                "is_distilled": True,
+                "patch_size": 2,
+                "text_encoder_select_layers": [2, 5, 8],
+            }
+        )
+    )
+
+    captured: dict = {}
+
+    class _FromPretrained:
+        def __init__(self, tag):
+            self.tag = tag
+
+        def from_pretrained(self, repo_id, **kwargs):
+            captured.setdefault("components", {})[self.tag] = (repo_id, kwargs)
+            return SimpleNamespace(tag = self.tag)
+
+    def _pipeline_ctor(**kwargs):
+        captured["pipeline"] = kwargs
+        return SimpleNamespace(**kwargs)
+
+    fake_diffusers = SimpleNamespace(
+        FlowMatchEulerDiscreteScheduler = _FromPretrained("scheduler"),
+        AutoencoderKLQwenImage = _FromPretrained("vae"),
+        Krea2Transformer2DModel = _FromPretrained("transformer"),
+        Krea2Pipeline = _pipeline_ctor,
+    )
+    monkeypatch.setitem(sys.modules, "diffusers", fake_diffusers)
+    monkeypatch.setattr(
+        "core.inference.diffusion_krea2.load_krea2_tokenizer",
+        lambda repo_id, hf_token = None: SimpleNamespace(tag = "tokenizer"),
+    )
+    monkeypatch.setattr(
+        "core.inference.diffusion_krea2.load_krea2_text_encoder",
+        lambda repo_id, dtype, hf_token = None: SimpleNamespace(tag = "text_encoder"),
+    )
+
+    pipe = load_krea2_pipeline(str(tmp_path), "bf16")
+
+    # Turbo's fixed-mu schedule rides on is_distilled; dropping any of these would
+    # silently degrade generations, so the ctor kwargs are asserted exactly.
+    assert captured["pipeline"]["is_distilled"] is True
+    assert captured["pipeline"]["patch_size"] == 2
+    assert captured["pipeline"]["text_encoder_select_layers"] == [2, 5, 8]
+    assert pipe.transformer.tag == "transformer"
+    # A prebuilt transformer (single-file/quant path) must be used as-is.
+    prebuilt = SimpleNamespace(tag = "prebuilt")
+    pipe = load_krea2_pipeline(str(tmp_path), "bf16", transformer = prebuilt)
+    assert pipe.transformer is prebuilt
+
+
+# ── registry / trust / int8 exclusion wiring ─────────────────────────────────
+
+
+def test_krea2_family_wiring():
+    from core.inference.diffusion import _is_trusted_diffusion_repo
+    from core.inference.diffusion_families import detect_family, family_sd_cpp_supported
+    from core.inference.diffusion_transformer_quant import TQ_INT8, exclude_tokens_for_scheme
+
+    fam = detect_family("krea/Krea-2-Turbo")
+    assert fam is not None and fam.name == KREA2_FAMILY_NAME
+    # The vendor repo is non-GGUF allowlisted; no sd.cpp mapping -> diffusers fallback.
+    assert _is_trusted_diffusion_repo("krea/Krea-2-Turbo")
+    assert not family_sd_cpp_supported(fam)
+    # Krea2TimestepEmbedding runs at M = batch; int8 (torch._int_mm, M > 16) must skip it.
+    assert "time_embed" in exclude_tokens_for_scheme(TQ_INT8)
+
+
+# ── training wiring ──────────────────────────────────────────────────────────
+
+
+def test_krea2_training_registry():
+    from core.inference.diffusion_families import trainable_family_names
+    from core.training.diffusion_train_common import (
+        family_train_infos,
+        get_trainer,
+        train_defaults,
+    )
+    from core.training.diffusion_dit_trainer import run_dit_lora_training
+
+    assert "krea-2" in trainable_family_names()
+    assert get_trainer("krea-2") is run_dit_lora_training
+    # The Krea 2 authors' recommended starting point (their reference script defaults).
+    assert train_defaults("krea-2") == {
+        "lora_rank": 32,
+        "learning_rate": 3e-4,
+        "resolution": 512,
+    }
+    info = {i["name"]: i for i in family_train_infos()}["krea-2"]
+    assert info["default_base"] == "krea/Krea-2-Turbo"
+    assert info["supports_compile"] is True
+
+
+def test_krea2_spec_registered_with_authors_targets():
+    from core.training.diffusion_dit_trainer import _KREA2_TARGETS, _SPECS
+
+    spec = _SPECS["krea-2"]
+    assert spec.force_bf16 is True
+    assert spec.lora_targets == _KREA2_TARGETS
+    # The authors' full recommended set: attention + SwiGLU + text fusion + embedders.
+    for t in ("to_q", "to_gate", "ff.up", "text_fusion.projector", "time_mod_proj"):
+        assert t in _KREA2_TARGETS
+
+
+def test_krea2_collate_and_forward_roundtrip():
+    import torch
+    from core.training.diffusion_dit_trainer import _SPECS
+
+    spec = _SPECS["krea-2"]
+    # Two fixed-length embed entries collate to a plain concat with the mask batched.
+    entries = [
+        (torch.randn(1, 8, 12, 16), torch.ones(1, 8, dtype = torch.int64)),
+        (torch.randn(1, 8, 12, 16), torch.ones(1, 8, dtype = torch.int64)),
+    ]
+    pe_b, mask_b = spec.collate(entries, "cpu", torch.float32)
+    assert pe_b.shape == (2, 8, 12, 16)
+    assert mask_b.shape == (2, 8)
+
+    captured = {}
+
+    class _FakeTransformer:
+        def __call__(self, **kwargs):
+            captured.update(kwargs)
+            # Echo the packed sequence: unpack(pack(x)) == x proves the inlined
+            # packing mirrors Krea2Pipeline exactly (they are mutual inverses).
+            return (kwargs["hidden_states"],)
+
+    noisy = torch.randn(2, 16, 1, 8, 8)
+    timesteps = torch.tensor([250.0, 750.0])
+    pred = spec.forward(
+        _FakeTransformer(), noisy, timesteps, None, (pe_b, mask_b), None, "cpu", torch.float32
+    )
+    assert torch.equal(pred, noisy)
+    # [B, (H/2)*(W/2), C*4] patches, one shared [(txt+img), 3] position grid, and the
+    # [0, 1] timestep convention.
+    assert captured["hidden_states"].shape == (2, 16, 64)
+    assert captured["position_ids"].shape == (8 + 16, 3)
+    assert torch.allclose(captured["timestep"], torch.tensor([0.25, 0.75]))
+    assert captured["encoder_attention_mask"] is mask_b
diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx
index bff5f93e9d..e4e2e6f299 100644
--- a/studio/frontend/src/features/images/images-page.tsx
+++ b/studio/frontend/src/features/images/images-page.tsx
@@ -98,6 +98,8 @@ const editGguf = (id: string, name: string): ModelOption => ({
 type SafetensorsSpec = { kind: "pipeline" | "single_file"; filename?: string };
 const SAFETENSORS_MODELS: Record<string, SafetensorsSpec> = {
   "unsloth/Z-Image-Turbo-unsloth-bnb-4bit": { kind: "pipeline" },
+  // Krea 2 Turbo: official vendor repo (bf16 pipeline), on the backend allowlist.
+  "krea/Krea-2-Turbo": { kind: "pipeline" },
   "unsloth/Qwen-Image-2512-unsloth-bnb-4bit": { kind: "pipeline" },
   "unsloth/Qwen-Image-2512-FP8": {
     kind: "single_file",
@@ -132,6 +134,7 @@ const MODELS: ModelOption[] = [
     "Z-Image-Turbo (bnb-4bit)",
     "Safetensors · bnb-4bit",
   ),
+  safetensors("krea/Krea-2-Turbo", "Krea 2 Turbo", "Safetensors · bf16"),
   safetensors(
     "unsloth/Qwen-Image-2512-unsloth-bnb-4bit",
     "Qwen-Image 2512 (bnb-4bit)",
@@ -208,6 +211,9 @@ const DEFAULT_GEN = { steps: 9, guidance: 0 };
 
 const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }> = [
   { match: "z-image-turbo", steps: 9, guidance: 0 },
+  // Krea 2 Turbo is distilled (TDM): 8 steps, no CFG (the base/midtrain checkpoints
+  // would want ~28 steps + CFG 4.5, but only Turbo is curated today).
+  { match: "krea-2", steps: 8, guidance: 0 },
   { 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 },

From e850ba79f999a3166658502b214e79b9bce80555 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 3 Jul 2026 13:34:23 +0000
Subject: [PATCH 02/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/inference/diffusion_krea2.py | 10 ++++++----
 studio/backend/tests/test_diffusion_krea2.py     |  4 +---
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/studio/backend/core/inference/diffusion_krea2.py b/studio/backend/core/inference/diffusion_krea2.py
index 312a4c8aa2..65619ccd47 100644
--- a/studio/backend/core/inference/diffusion_krea2.py
+++ b/studio/backend/core/inference/diffusion_krea2.py
@@ -62,14 +62,16 @@ def remap_rope_parameters(text_config) -> None:
     the config carries no ``rope_parameters`` dict."""
     rope_parameters = getattr(text_config, "rope_parameters", None)
     if getattr(text_config, "rope_scaling", None) is None and isinstance(rope_parameters, dict):
-        text_config.rope_scaling = {
-            k: v for k, v in rope_parameters.items() if k != "rope_theta"
-        }
+        text_config.rope_scaling = {k: v for k, v in rope_parameters.items() if k != "rope_theta"}
         if "rope_theta" in rope_parameters:
             text_config.rope_theta = rope_parameters["rope_theta"]
 
 
-def load_krea2_text_encoder(repo_id: str, dtype, hf_token: Optional[str] = None):
+def load_krea2_text_encoder(
+    repo_id: str,
+    dtype,
+    hf_token: Optional[str] = None,
+):
     """The Qwen3-VL text encoder, remapping 5.x ``rope_parameters`` for a 4.x runtime."""
     from transformers import AutoConfig, Qwen3VLModel
 
diff --git a/studio/backend/tests/test_diffusion_krea2.py b/studio/backend/tests/test_diffusion_krea2.py
index b360ae43f9..4cface12cd 100644
--- a/studio/backend/tests/test_diffusion_krea2.py
+++ b/studio/backend/tests/test_diffusion_krea2.py
@@ -55,9 +55,7 @@ def test_remap_rope_parameters_noop_on_5x_runtime_or_plain_4x_config():
 
 
 def test_load_model_index_from_local_path(tmp_path):
-    (tmp_path / "model_index.json").write_text(
-        json.dumps({"is_distilled": True, "patch_size": 2})
-    )
+    (tmp_path / "model_index.json").write_text(json.dumps({"is_distilled": True, "patch_size": 2}))
     assert _load_model_index(str(tmp_path)) == {"is_distilled": True, "patch_size": 2}
 
 

From c7477ea2369eeeaab58c6c6370a99ea8971ca43a Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 3 Jul 2026 14:50:35 +0000
Subject: [PATCH 03/17] Add the Krea 2 per-arch eager fusion
 (Krea2TransformerBlock addcmul patch)

Fuses the block's two inline modulations (1 + scale) * norm(x) + shift and two
gated residuals x + gate * out to torch.addcmul, matching the existing qwen /
z-image / flux fusions (compile-safe, 1-ULP more accurate, body-drift guarded).
Stock-vs-patched equivalence test included; install count is now 7.
---
 .../core/inference/diffusion_arch_patches.py  | 51 ++++++++++++++++++-
 .../tests/test_diffusion_arch_patches.py      | 30 ++++++++++-
 2 files changed, 79 insertions(+), 2 deletions(-)

diff --git a/studio/backend/core/inference/diffusion_arch_patches.py b/studio/backend/core/inference/diffusion_arch_patches.py
index e079e09603..f73085197e 100644
--- a/studio/backend/core/inference/diffusion_arch_patches.py
+++ b/studio/backend/core/inference/diffusion_arch_patches.py
@@ -27,13 +27,14 @@ source-body check (``_body_has``) confirms the exact lines we rewrite are still
 future diffusers that changed the block body simply leaves the block UNPATCHED (correctness
 first) rather than running a stale copy. Kill-switch: ``UNSLOTH_DIFFUSION_ARCH_PATCHES=0``.
 
-Implemented for all four families (extend by adding entries to ``_SPECS``):
+Implemented for all five families (extend by adding entries to ``_SPECS``):
   * qwen-image  ``QwenImageTransformerBlock._modulate`` -- modulation addcmul (all 4 sites).
   * z-image     ``ZImageTransformerBlock.forward``      -- the 2 gated-residual addcmuls.
   * flux.1      ``FluxTransformerBlock.forward`` (inline norm2 modulation + 4 gated residuals)
                 + ``FluxSingleTransformerBlock.forward`` (residual + gate*proj_out).
   * flux.2-klein ``Flux2TransformerBlock.forward`` (4 inline modulations + 4 gated residuals)
                 + ``Flux2SingleTransformerBlock.forward`` (inline modulation + gated residual).
+  * krea-2      ``Krea2TransformerBlock.forward`` (2 inline modulations + 2 gated residuals).
 """
 
 from __future__ import annotations
@@ -468,6 +469,53 @@ def _spec_flux2_single():
     return (cls, "forward", _flux2_single_forward)
 
 
+# =====================================================================================
+# krea-2: Krea2TransformerBlock.forward  (2 inline modulations + 2 gated residuals)
+# =====================================================================================
+def _krea2_block_forward(
+    self,
+    hidden_states,
+    temb,
+    image_rotary_emb,
+    attention_mask = None,
+):
+    """diffusers 0.39 ``Krea2TransformerBlock.forward`` with the two inline modulations
+    ``(1 + scale) * norm(x) + shift`` and the two gated residuals ``x + gate * out`` each
+    fused to one ``torch.addcmul``."""
+    # temb: (B, 1, 6 * hidden_size), shared across all blocks; each block only learns an
+    # additive table.
+    modulation = temb.unflatten(-1, (6, -1)) + self.scale_shift_table
+    prescale, preshift, pregate, postscale, postshift, postgate = modulation.unbind(-2)
+
+    norm1 = self.norm1(hidden_states)
+    attn_out = self.attn(
+        torch.addcmul(preshift, norm1, 1 + prescale),
+        attention_mask = attention_mask,
+        image_rotary_emb = image_rotary_emb,
+    )
+    hidden_states = torch.addcmul(hidden_states, pregate, attn_out)
+    norm2 = self.norm2(hidden_states)
+    ff_out = self.ff(torch.addcmul(postshift, norm2, 1 + postscale))
+    return torch.addcmul(hidden_states, postgate, ff_out)
+
+
+def _spec_krea2_forward():
+    try:
+        from diffusers.models.transformers.transformer_krea2 import Krea2TransformerBlock as cls
+    except Exception:  # noqa: BLE001
+        return None
+    orig = getattr(cls, "forward", None)
+    if orig is None or not _body_has(
+        orig,
+        "(1.0 + prescale) * self.norm1(hidden_states) + preshift",
+        "hidden_states = hidden_states + pregate * attn_out",
+        "(1.0 + postscale) * self.norm2(hidden_states) + postshift",
+        "hidden_states = hidden_states + postgate * ff_out",
+    ):
+        return None
+    return (cls, "forward", _krea2_block_forward)
+
+
 # =====================================================================================
 # registry + lifecycle
 # =====================================================================================
@@ -480,6 +528,7 @@ _SPECS: tuple[Callable[[], Optional[tuple]], ...] = (
     _spec_flux_single,
     _spec_flux2_double,
     _spec_flux2_single,
+    _spec_krea2_forward,
 )
 
 # (cls, attr) pairs we successfully patched, for an exact reverse.
diff --git a/studio/backend/tests/test_diffusion_arch_patches.py b/studio/backend/tests/test_diffusion_arch_patches.py
index 77276a0fb1..02e975c5a8 100644
--- a/studio/backend/tests/test_diffusion_arch_patches.py
+++ b/studio/backend/tests/test_diffusion_arch_patches.py
@@ -236,6 +236,34 @@ def test_flux2_single_forward_matches_stock():
     _close_any(got, ref)
 
 
+def test_krea2_forward_matches_stock():
+    from diffusers.models.transformers.transformer_krea2 import (
+        Krea2RotaryPosEmbed,
+        Krea2TransformerBlock,
+    )
+    from diffusers.pipelines.krea2.pipeline_krea2 import Krea2Pipeline
+
+    torch.manual_seed(4)
+    blk = Krea2TransformerBlock(
+        hidden_size = D, intermediate_size = 2 * D, num_heads = H, num_kv_heads = H // 2,
+        norm_eps = 1e-6,
+    ).eval()
+    # Give the zero-init modulation table real values so all six scale/shift/gate
+    # branches contribute to the output.
+    with torch.no_grad():
+        blk.scale_shift_table.normal_()
+    # A [text + 2x2 image grid] sequence with the real rotary embed (axes sum to head_dim).
+    position_ids = Krea2Pipeline.prepare_position_ids(4, 2, 2, torch.device("cpu"))
+    rope = Krea2RotaryPosEmbed(theta = 10000, axes_dim = [D // H // 2, D // H // 4, D // H // 4])
+    image_rotary_emb = rope(position_ids)
+    hs = torch.randn(B, position_ids.shape[0], D)
+    tm = torch.randn(B, 1, 6 * D)
+    with torch.inference_mode():
+        ref = Krea2TransformerBlock.forward(blk, hs, tm, image_rotary_emb).clone()
+        got = ap._krea2_block_forward(blk, hs, tm, image_rotary_emb)
+    torch.testing.assert_close(got, ref, atol = 1e-5, rtol = 1e-4)
+
+
 # ── lifecycle ───────────────────────────────────────────────────────────────────
 
 
@@ -246,7 +274,7 @@ def test_install_idempotent_and_reversible():
     q_orig, z_orig = Q._modulate, Z.forward
     n1 = ap.install_arch_patches()
     n2 = ap.install_arch_patches()  # idempotent
-    assert n1 == 6 and n2 == n1  # qwen + z-image + flux.1 x2 + flux.2 x2
+    assert n1 == 7 and n2 == n1  # qwen + z-image + flux.1 x2 + flux.2 x2 + krea-2
     assert Q._modulate is not q_orig and Z.forward is not z_orig
     assert ap.is_installed()
 

From 11a00d0238c854ffff673a26828a4bf38f64b526 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 3 Jul 2026 14:55:53 +0000
Subject: [PATCH 04/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/tests/test_diffusion_arch_patches.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/studio/backend/tests/test_diffusion_arch_patches.py b/studio/backend/tests/test_diffusion_arch_patches.py
index 02e975c5a8..ceeb047669 100644
--- a/studio/backend/tests/test_diffusion_arch_patches.py
+++ b/studio/backend/tests/test_diffusion_arch_patches.py
@@ -245,7 +245,10 @@ def test_krea2_forward_matches_stock():
 
     torch.manual_seed(4)
     blk = Krea2TransformerBlock(
-        hidden_size = D, intermediate_size = 2 * D, num_heads = H, num_kv_heads = H // 2,
+        hidden_size = D,
+        intermediate_size = 2 * D,
+        num_heads = H,
+        num_kv_heads = H // 2,
         norm_eps = 1e-6,
     ).eval()
     # Give the zero-init modulation table real values so all six scale/shift/gate

From 98d6c29ae1511ab2773410219d183fb5ac83b5ae Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 4 Jul 2026 01:07:14 +0000
Subject: [PATCH 05/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/inference/sd_cpp_backend.py |  5 ++++-
 studio/backend/tests/test_sd_cpp_backend.py     | 14 +++++++++++---
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py
index 7632551661..d8cb21d409 100644
--- a/studio/backend/core/inference/sd_cpp_backend.py
+++ b/studio/backend/core/inference/sd_cpp_backend.py
@@ -893,7 +893,10 @@ class SdCppDiffusionBackend:
                 lora_stage = Path(server_lora_dir) / f"gen_{os.urandom(6).hex()}"
                 materialized = diffusion_lora.materialize_native_dir(lora_resolved, lora_stage)
                 lora_payload = [
-                    {"path": f"{lora_stage.name}/{Path(m.path).name}", "multiplier": float(m.weight)}
+                    {
+                        "path": f"{lora_stage.name}/{Path(m.path).name}",
+                        "multiplier": float(m.weight),
+                    }
                     for m in materialized
                 ]
         try:
diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py
index 0b52ea6e14..be2fc22a4b 100644
--- a/studio/backend/tests/test_sd_cpp_backend.py
+++ b/studio/backend/tests/test_sd_cpp_backend.py
@@ -691,7 +691,11 @@ def _fake_materialize(resolved, dest):
     return out
 
 
-def _patch_lora(monkeypatch, resolved, supported = True):
+def _patch_lora(
+    monkeypatch,
+    resolved,
+    supported = True,
+):
     from core.inference import diffusion_lora as dl
 
     monkeypatch.setattr(dl, "supports_lora", lambda **k: supported)
@@ -706,7 +710,9 @@ def test_generate_oneshot_applies_loras_via_prompt_tags(monkeypatch):
 
     eng = _FakeEngine()
     b = _loaded_backend(engine = eng)  # mode = "oneshot"
-    _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)])
+    _patch_lora(
+        monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)]
+    )
     b.generate(prompt = "a fox", steps = 4, seed = 1, loras = [("id1", 0.8)])
     _, params, _, _ = eng.calls[0]
     assert params.lora_dir is not None and params.lora_apply_mode == "auto"
@@ -725,7 +731,9 @@ def test_generate_server_stages_loras_and_sends_structured_field(monkeypatch, tm
     servers: list = []
     _run_server_load(monkeypatch, b, servers)
     servers[0].lora_dir = str(tmp_path)
-    _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)])
+    _patch_lora(
+        monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)]
+    )
     b.generate(prompt = "x", steps = 4, seed = 1, batch_size = 1, loras = [("id1", 0.7)])
     payload = servers[0].payloads[0]
     assert "lora" in payload and len(payload["lora"]) == 1

From bbc840eb33f3b77e988630ed373105d60d1adb12 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 4 Jul 2026 01:31:41 +0000
Subject: [PATCH 06/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/training/diffusion_dit_trainer.py     | 8 +++++++-
 .../backend/core/training/diffusion_training_service.py   | 4 +++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py
index 1a63e5c014..7370866203 100644
--- a/studio/backend/core/training/diffusion_dit_trainer.py
+++ b/studio/backend/core/training/diffusion_dit_trainer.py
@@ -300,7 +300,13 @@ def _apply_fp8_training(transformer, on_event) -> bool:
 
 
 def _pick_auto_precision(
-    prequant, device, free_gb, dense_gb, capability, has_fp8, has_torchao = True
+    prequant,
+    device,
+    free_gb,
+    dense_gb,
+    capability,
+    has_fp8,
+    has_torchao = True,
 ) -> str:
     """Pure policy for base_precision="auto": nf4 for a prequant base or no CUDA; else the
     fastest dense mode whose weights + headroom (activations, optimizer, cache) fit the
diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py
index b33b061f1d..67f64f14ad 100644
--- a/studio/backend/core/training/diffusion_training_service.py
+++ b/studio/backend/core/training/diffusion_training_service.py
@@ -450,7 +450,9 @@ class DiffusionTrainingService:
                     if "learning_rate" in ev
                     else s["learning_rate"]
                 )
-                grad_norm = _finite_or_none(ev["grad_norm"]) if "grad_norm" in ev else s["grad_norm"]
+                grad_norm = (
+                    _finite_or_none(ev["grad_norm"]) if "grad_norm" in ev else s["grad_norm"]
+                )
                 s.update(
                     status = "running",
                     step = ev.get("step", s["step"]),

From dc290bdf71b28e1b70d26ec42a0c462961d675aa Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 4 Jul 2026 03:23:30 +0000
Subject: [PATCH 07/17] Train Krea 2 LoRAs on the undistilled Raw checkpoint by
 default

Krea's release guidance is to train on Krea-2-Raw and run adapters on
Turbo. Raw now leads the krea-2 training bases (Turbo stays available),
both vendor repos are trust-listed, and load_krea2_pipeline fails fast
with an upgrade hint on diffusers older than 0.39 instead of a bare
AttributeError mid-load
---
 studio/backend/core/inference/diffusion.py    |  5 ++++-
 .../core/inference/diffusion_families.py      |  7 +++++--
 .../backend/core/inference/diffusion_krea2.py |  9 +++++++++
 .../core/training/diffusion_train_common.py   |  5 ++++-
 studio/backend/tests/test_diffusion_krea2.py  | 20 +++++++++++++++++--
 5 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 3418d6f23b..72398e47ce 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -204,9 +204,12 @@ _TRUSTED_NON_GGUF_REPOS = frozenset(
         "black-forest-labs/flux.1-dev",
         "tongyi-mai/z-image-turbo",
         "qwen/qwen-image",
-        # Krea 2 Turbo: official vendor repo, safetensors-only, no remote code. Loaded
+        # Krea 2: official vendor repos, safetensors-only, no remote code. Loaded
         # per-component via core/inference/diffusion_krea2.py (no GGUF variant yet).
+        # Turbo is the inference model; Raw is the undistilled base Krea recommends
+        # training LoRAs on (train on Raw, run adapters on Turbo).
         "krea/krea-2-turbo",
+        "krea/krea-2-raw",
     }
 )
 
diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py
index 62b9baa511..ed2dc1e9c1 100644
--- a/studio/backend/core/inference/diffusion_families.py
+++ b/studio/backend/core/inference/diffusion_families.py
@@ -294,9 +294,12 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
         base_repo = "krea/Krea-2-Turbo",
         aliases = ("krea2",),
         # LoRA training via the DiT trainer (no prequant repo yet, so nf4 quantizes the
-        # 12B transformer on the fly under the default precision).
+        # 12B transformer on the fly under the default precision). Krea's release guidance
+        # is explicit: train LoRAs on the undistilled Raw checkpoint and apply them on
+        # Turbo for inference, so Raw is the default training base and Turbo stays the
+        # inference/base repo.
         trainable = True,
-        train_base_repos = ("krea/Krea-2-Turbo",),
+        train_base_repos = ("krea/Krea-2-Raw", "krea/Krea-2-Turbo"),
         # The checkpoint is exported bf16-only (the model card pins bfloat16); fp16 is
         # unvalidated upstream, so keep the fp16 fallback off like z-image.
         fp16_incompatible = True,
diff --git a/studio/backend/core/inference/diffusion_krea2.py b/studio/backend/core/inference/diffusion_krea2.py
index 65619ccd47..4475559cdb 100644
--- a/studio/backend/core/inference/diffusion_krea2.py
+++ b/studio/backend/core/inference/diffusion_krea2.py
@@ -113,6 +113,15 @@ def load_krea2_pipeline(
     """
     import diffusers
 
+    # diffusers gained Krea2Pipeline in 0.39; on an older install the getattr chain below
+    # would die with a bare AttributeError mid-load, so fail first with the actionable fix.
+    if not hasattr(diffusers, "Krea2Pipeline"):
+        raise RuntimeError(
+            f"Krea 2 needs diffusers >= 0.39.0 (Krea2Pipeline); this environment has "
+            f"diffusers {getattr(diffusers, '__version__', 'unknown')}. "
+            f"Upgrade with: pip install -U diffusers"
+        )
+
     token = hf_token or None
     tokenizer = load_krea2_tokenizer(repo_id, hf_token = token)
     text_encoder = load_krea2_text_encoder(repo_id, dtype, hf_token = token)
diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py
index 4f531641fe..6a587174b7 100644
--- a/studio/backend/core/training/diffusion_train_common.py
+++ b/studio/backend/core/training/diffusion_train_common.py
@@ -202,7 +202,10 @@ _FAMILY_VRAM_NOTES = {
     ),
     "qwen-image": "20B model, QLoRA (nf4) by default (~24 GB+). The heaviest option.",
     "z-image": "6B model, QLoRA (nf4) by default (~12 GB+). bf16 only.",
-    "krea-2": "12B model, QLoRA (nf4) by default (~18 GB+). bf16 only.",
+    "krea-2": (
+        "12B model, QLoRA (nf4) by default (~18 GB+). bf16 only. Trains on the "
+        "undistilled Krea-2-Raw (Krea's guidance: train on Raw, run adapters on Turbo)."
+    ),
 }
 
 
diff --git a/studio/backend/tests/test_diffusion_krea2.py b/studio/backend/tests/test_diffusion_krea2.py
index 4cface12cd..f2bfcac2f0 100644
--- a/studio/backend/tests/test_diffusion_krea2.py
+++ b/studio/backend/tests/test_diffusion_krea2.py
@@ -120,6 +120,17 @@ def test_load_krea2_pipeline_threads_init_config(monkeypatch, tmp_path):
 # ── registry / trust / int8 exclusion wiring ─────────────────────────────────
 
 
+def test_load_krea2_pipeline_requires_krea_capable_diffusers(monkeypatch):
+    # On diffusers < 0.39 (no Krea2Pipeline) the loader must fail fast with the upgrade
+    # hint instead of dying with a bare AttributeError mid-load.
+    import pytest
+
+    fake = SimpleNamespace(__version__ = "0.38.0")
+    monkeypatch.setitem(sys.modules, "diffusers", fake)
+    with pytest.raises(RuntimeError, match = "0.39"):
+        load_krea2_pipeline("krea/Krea-2-Turbo", "bf16")
+
+
 def test_krea2_family_wiring():
     from core.inference.diffusion import _is_trusted_diffusion_repo
     from core.inference.diffusion_families import detect_family, family_sd_cpp_supported
@@ -127,8 +138,10 @@ def test_krea2_family_wiring():
 
     fam = detect_family("krea/Krea-2-Turbo")
     assert fam is not None and fam.name == KREA2_FAMILY_NAME
-    # The vendor repo is non-GGUF allowlisted; no sd.cpp mapping -> diffusers fallback.
+    # Both vendor repos are non-GGUF allowlisted (Turbo for inference, Raw for training);
+    # no sd.cpp mapping -> diffusers fallback.
     assert _is_trusted_diffusion_repo("krea/Krea-2-Turbo")
+    assert _is_trusted_diffusion_repo("krea/Krea-2-Raw")
     assert not family_sd_cpp_supported(fam)
     # Krea2TimestepEmbedding runs at M = batch; int8 (torch._int_mm, M > 16) must skip it.
     assert "time_embed" in exclude_tokens_for_scheme(TQ_INT8)
@@ -155,7 +168,10 @@ def test_krea2_training_registry():
         "resolution": 512,
     }
     info = {i["name"]: i for i in family_train_infos()}["krea-2"]
-    assert info["default_base"] == "krea/Krea-2-Turbo"
+    # Krea's guidance: train LoRAs on the undistilled Raw model, run them on Turbo, so
+    # Raw leads the training bases while Turbo stays available.
+    assert info["default_base"] == "krea/Krea-2-Raw"
+    assert info["base_repos"] == ["krea/Krea-2-Raw", "krea/Krea-2-Turbo"]
     assert info["supports_compile"] is True
 
 

From 48e1cbc242893d2be5dc21b9d9e62e3f00c0d7a9 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 4 Jul 2026 03:31:20 +0000
Subject: [PATCH 08/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/routes/models.py                    |  4 +---
 studio/backend/tests/test_diffusion_dataset_api.py |  6 ++++--
 studio/backend/tests/test_diffusion_dit_trainer.py |  5 ++++-
 studio/backend/tests/test_diffusion_training.py    |  6 ++++--
 studio/backend/tests/test_local_model_format.py    | 13 +++++++++----
 5 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py
index 54bc490052..d4501517c8 100644
--- a/studio/backend/routes/models.py
+++ b/studio/backend/routes/models.py
@@ -911,9 +911,7 @@ async def list_local_models(
         models = collect_local_models(models_root)
         # Tag each model with its task so the Images picker can filter to diffusion
         # (GGUF by architecture; local diffusers checkpoints by pipeline / family).
-        models = [
-            m.model_copy(update = {"task": _local_model_task(m)}) for m in models
-        ]
+        models = [m.model_copy(update = {"task": _local_model_task(m)}) for m in models]
 
         return LocalModelListResponse(
             models_dir = str(models_root),
diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py
index e1c2c0a186..b10dd04f4c 100644
--- a/studio/backend/tests/test_diffusion_dataset_api.py
+++ b/studio/backend/tests/test_diffusion_dataset_api.py
@@ -64,8 +64,10 @@ def test_list_images_caption_precedence(client, ds_root):
     # a.png -> sidecar (an explicit edit beats the metadata row), b.png -> metadata-only,
     # c.png -> none.
     (folder / "metadata.jsonl").write_text(
-        json.dumps({"file_name": "a.png", "text": "from metadata"}) + "\n"
-        + json.dumps({"file_name": "b.png", "text": "from metadata"}) + "\n",
+        json.dumps({"file_name": "a.png", "text": "from metadata"})
+        + "\n"
+        + json.dumps({"file_name": "b.png", "text": "from metadata"})
+        + "\n",
         encoding = "utf-8",
     )
     (folder / "a.txt").write_text("edited sidecar", encoding = "utf-8")
diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py
index a727f0f13b..c40a0575a3 100644
--- a/studio/backend/tests/test_diffusion_dit_trainer.py
+++ b/studio/backend/tests/test_diffusion_dit_trainer.py
@@ -61,7 +61,10 @@ def test_select_lora_targets_explicit_override_wins():
         base_model = "black-forest-labs/FLUX.1-dev", data_dir = "d", output_dir = "o"
     ).normalized()
     assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS
-    assert _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets) == _FLUX_TARGETS
+    assert (
+        _select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets)
+        == _FLUX_TARGETS
+    )
 
 
 @pytest.mark.parametrize(
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index 1efd19f18f..517186a196 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/studio/backend/tests/test_diffusion_training.py
@@ -547,8 +547,10 @@ def test_diffusion_info_counts_metadata_captions(client, dataset_roots):
     (folder / "c.png").write_bytes(b"x")
     # a.png + b.png via metadata; a.png also has a sidecar (must count once); c.png none.
     (folder / "metadata.jsonl").write_text(
-        json.dumps({"file_name": "a.png", "text": "cap a"}) + "\n"
-        + json.dumps({"file_name": "b.png", "text": "cap b"}) + "\n",
+        json.dumps({"file_name": "a.png", "text": "cap a"})
+        + "\n"
+        + json.dumps({"file_name": "b.png", "text": "cap b"})
+        + "\n",
         encoding = "utf-8",
     )
     (folder / "a.txt").write_text("edited a", encoding = "utf-8")
diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py
index dd888eb50d..f9184774d6 100644
--- a/studio/backend/tests/test_local_model_format.py
+++ b/studio/backend/tests/test_local_model_format.py
@@ -121,7 +121,14 @@ def test_scan_models_dir_classifies_root_gguf_with_config(tmp_path):
 from models.models import LocalModelInfo  # noqa: E402
 
 
-def _local(path, *, model_format = None, model_id = None, display_name = "m", id = "m"):
+def _local(
+    path,
+    *,
+    model_format = None,
+    model_id = None,
+    display_name = "m",
+    id = "m",
+):
     return LocalModelInfo(
         id = id,
         display_name = display_name,
@@ -147,9 +154,7 @@ def test_local_task_tags_diffusers_by_family_id(tmp_path):
     d = tmp_path / "flux-checkpoint"
     _touch(d / "flux1-dev.safetensors")
     assert (
-        models_route._local_model_task(
-            _local(d, model_id = "black-forest-labs/FLUX.1-dev")
-        )
+        models_route._local_model_task(_local(d, model_id = "black-forest-labs/FLUX.1-dev"))
         == "text-to-image"
     )
 

From ff59b435338080fac476b8a73afff8f20ae226db Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 4 Jul 2026 04:31:09 +0000
Subject: [PATCH 09/17] Fail clearly when a local Krea 2 dir lacks
 model_index.json

Falling through to hf_hub_download with a filesystem path as the repo id
raised an opaque HFValidationError; a local dir without the file now
raises FileNotFoundError naming the directory
---
 studio/backend/core/inference/diffusion_krea2.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/studio/backend/core/inference/diffusion_krea2.py b/studio/backend/core/inference/diffusion_krea2.py
index 4475559cdb..86a1d0413c 100644
--- a/studio/backend/core/inference/diffusion_krea2.py
+++ b/studio/backend/core/inference/diffusion_krea2.py
@@ -85,12 +85,20 @@ def load_krea2_text_encoder(
 
 def _load_model_index(repo_id: str, hf_token: Optional[str] = None) -> dict[str, Any]:
     """model_index.json as a dict, from a local path or the Hub cache."""
+    is_local_dir = False
     try:
-        local = Path(repo_id).expanduser() / "model_index.json"
+        root = Path(repo_id).expanduser()
+        is_local_dir = root.is_dir()
+        local = root / "model_index.json"
         if local.is_file():
             return json.loads(local.read_text())
     except OSError:
         pass
+    if is_local_dir:
+        # A local checkpoint dir without the file must fail clearly here: falling through
+        # to hf_hub_download with a filesystem path as the repo id would die with an
+        # opaque HFValidationError instead.
+        raise FileNotFoundError(f"model_index.json not found in local model dir {repo_id}")
     from huggingface_hub import hf_hub_download
 
     path = hf_hub_download(repo_id, "model_index.json", token = hf_token or None)

From 6ecbb2b8a35c649d11db67cbe002a7ee1b95794f Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 4 Jul 2026 04:41:52 +0000
Subject: [PATCH 10/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/training/diffusion_train_common.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py
index 60741f8e8d..c4fe3c06be 100644
--- a/studio/backend/core/training/diffusion_train_common.py
+++ b/studio/backend/core/training/diffusion_train_common.py
@@ -493,7 +493,9 @@ def discover_image_caption_pairs(
         # 2. metadata row keyed by file name (basename or the relative path; as_posix so a
         #    Windows backslash path still matches the jsonl's forward-slash keys).
         if caption is None:
-            caption = meta_caption.get(img.name) or meta_caption.get(img.relative_to(root).as_posix())
+            caption = meta_caption.get(img.name) or meta_caption.get(
+                img.relative_to(root).as_posix()
+            )
         # 3. dreambooth instance prompt.
         if caption is None and instance_prompt:
             caption = instance_prompt

From 466f86c45aed8312e957481682405493a5fac568 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 4 Jul 2026 07:26:56 +0000
Subject: [PATCH 11/17] Skip the krea-2 forward roundtrip test on hosts without
 diffusers

spec.forward imports Krea2Pipeline for prepare_position_ids, so the test needs a
real diffusers install; the backend CI matrix runs without one and failed on the
import. Same importorskip guard the sigmas gather test already uses.
---
 studio/backend/tests/test_diffusion_krea2.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/studio/backend/tests/test_diffusion_krea2.py b/studio/backend/tests/test_diffusion_krea2.py
index f2bfcac2f0..82649c1a6e 100644
--- a/studio/backend/tests/test_diffusion_krea2.py
+++ b/studio/backend/tests/test_diffusion_krea2.py
@@ -6,6 +6,8 @@ import json
 import sys
 from types import SimpleNamespace
 
+import pytest
+
 from core.inference.diffusion_krea2 import (
     KREA2_FAMILY_NAME,
     _load_model_index,
@@ -187,6 +189,9 @@ def test_krea2_spec_registered_with_authors_targets():
 
 
 def test_krea2_collate_and_forward_roundtrip():
+    # spec.forward imports Krea2Pipeline (prepare_position_ids), so this needs a real
+    # diffusers install; CI hosts run the backend suite without one.
+    pytest.importorskip("diffusers")
     import torch
     from core.training.diffusion_dit_trainer import _SPECS
 

From 20a3650108a7f54debfb94ca6c49f95b13dbfa6c Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 5 Jul 2026 02:12:24 +0000
Subject: [PATCH 12/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/inference/sd_cpp_backend.py | 4 +---
 studio/backend/tests/test_diffusion_backend.py  | 4 +---
 2 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py
index 78425833d6..75572b500a 100644
--- a/studio/backend/core/inference/sd_cpp_backend.py
+++ b/studio/backend/core/inference/sd_cpp_backend.py
@@ -698,9 +698,7 @@ class SdCppDiffusionBackend:
             loading = self._loading
             if loading is None or loading.error is not None:
                 return ()
-            return tuple(
-                r for r in (loading.repo_id, loading.base_repo, *loading.asset_repos) if r
-            )
+            return tuple(r for r in (loading.repo_id, loading.base_repo, *loading.asset_repos) if r)
 
     # ── Generate ───────────────────────────────────────────────────────────
 
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index 8341bf46f6..5bf371f633 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -1414,9 +1414,7 @@ def test_bad_mode_strings_fail_before_eviction(fake_runtime):
         {"text_encoder_quant": "fp3"},
     ):
         with pytest.raises(ValueError):
-            backend.load_pipeline(
-                "unsloth/Z-Image-GGUF", gguf_filename = "m.gguf", **kwargs
-            )
+            backend.load_pipeline("unsloth/Z-Image-GGUF", gguf_filename = "m.gguf", **kwargs)
         assert backend._state is not None
 
 

From 237d197b5d9977e9f017be8022cb35bbacaa8fba Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 5 Jul 2026 04:58:34 +0000
Subject: [PATCH 13/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/tests/test_diffusion_controlnet.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py
index 3d4f7c98cf..50bf7159b0 100644
--- a/studio/backend/tests/test_diffusion_controlnet.py
+++ b/studio/backend/tests/test_diffusion_controlnet.py
@@ -263,6 +263,7 @@ def test_controlnet_pipe_rejects_family_without_classes():
     with pytest.raises(ValueError, match = "not supported"):
         b._controlnet_pipe(st, dc.ResolvedControlNet("x", "y", False), threading.Event())
 
+
 def test_controlnet_pipe_not_cached_after_unload_race(monkeypatch):
     # An unload that lands while from_pipe is assembling must not let the wrapper
     # repopulate the cache around the torn-down base pipe.

From d0f7dad7ec24feab48cdc32617c01f1e1a9e1e29 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 5 Jul 2026 07:58:48 +0000
Subject: [PATCH 14/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/training/diffusion_train_common.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py
index 3b536e6d81..1239ddde0d 100644
--- a/studio/backend/core/training/diffusion_train_common.py
+++ b/studio/backend/core/training/diffusion_train_common.py
@@ -592,7 +592,7 @@ def _plan_cache_variants(
 # fallback. Over this budget the default falls back to per-step VAE encoding. A fixed
 # constant (rather than a psutil RAM fraction) keeps the gate dependency-free and identical
 # across hosts; it is deliberately conservative, well under a typical training host's RAM.
-_LATENT_CACHE_BUDGET_BYTES = 4 * 1024 ** 3  # 4 GiB
+_LATENT_CACHE_BUDGET_BYTES = 4 * 1024**3  # 4 GiB
 
 # Returned by the cache builders when the estimated cache exceeds the budget: the caller
 # keeps the VAE resident and encodes each step's latents in-loop. A distinct sentinel from
@@ -609,7 +609,9 @@ def _latent_cache_forced() -> bool:
 
 
 def _latent_cache_over_budget(
-    per_variant_bytes: int, total_variants: int, budget_bytes: Optional[int] = None
+    per_variant_bytes: int,
+    total_variants: int,
+    budget_bytes: Optional[int] = None,
 ) -> bool:
     """True when a cache of ``total_variants`` entries, each two fp32 tensors totalling
     ``per_variant_bytes``, is estimated to exceed ``budget_bytes``. ``per_variant_bytes`` is

From 293b770c78340829047f3167dd36b1b7a4f7abff Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 5 Jul 2026 11:43:19 +0000
Subject: [PATCH 15/17] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/inference/diffusion.py        | 1 -
 studio/backend/tests/test_diffusion_controlnet.py | 8 ++++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index a810bd22e1..7ce31d5ecf 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -1564,7 +1564,6 @@ class DiffusionBackend:
             # local dir the user picked has no Hub scan and is exempt (fail-open there).
             if not getattr(resolved_cn, "is_local", False):
                 from utils.security import evaluate_file_security
-
                 _cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None)
                 if _cn_fs.blocked:
                     raise ValueError(_cn_fs.reason)
diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py
index b7ae037197..bcb7393457 100644
--- a/studio/backend/tests/test_diffusion_controlnet.py
+++ b/studio/backend/tests/test_diffusion_controlnet.py
@@ -234,7 +234,6 @@ def _state():
 def _allow_cn_security(monkeypatch):
     """Stub the Hub malware preflight to allow the load (hermetic, no network)."""
     import utils.security
-
     monkeypatch.setattr(
         utils.security,
         "evaluate_file_security",
@@ -277,7 +276,12 @@ def test_controlnet_pipe_blocks_flagged_remote_repo(monkeypatch):
 
     class _TrapModel(_FakeCNModel):
         @classmethod
-        def from_pretrained(cls, path, torch_dtype = None, token = None):
+        def from_pretrained(
+            cls,
+            path,
+            torch_dtype = None,
+            token = None,
+        ):
             loaded["called"] = True
             return super().from_pretrained(path, torch_dtype = torch_dtype, token = token)
 

From 2d974219bf36aade162eb4ea2552efb10677b594 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 6 Jul 2026 11:16:25 +0000
Subject: [PATCH 16/17] Deploy Krea adapters on Turbo and use its distilled
 recipe over the API

- DiffusionFamily gains deploy_base_repo (krea/Krea-2-Turbo): deploying a LoRA
  trained on Raw now previews it on Turbo, not the non-distilled Raw checkpoint.
  Scoped to a same-precision override so it never turns an nf4 train base into a
  larger bf16 deploy load; exposed through family_train_infos -> the Train UI's
  onDeployClick / historical-run deploy resolve the deploy base.
- _GENERATION_DEFAULTS gains a Krea entry (8 steps, 0 CFG) so the OpenAI
  /v1/images/generations route matches the Create UI's documented distilled recipe
  instead of falling through to the generic (9, 0.0).
---
 .../core/inference/diffusion_families.py      | 14 ++++++++++
 .../core/training/diffusion_train_common.py   |  2 ++
 studio/backend/models/training.py             |  4 +++
 studio/backend/tests/test_diffusion_krea2.py  | 17 +++++++++++-
 studio/frontend/src/features/images/api.ts    |  4 +++
 .../images/train/diffusion-train-panel.tsx    | 27 ++++++++++++++-----
 6 files changed, 61 insertions(+), 7 deletions(-)

diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py
index a06c66a48d..53792feac2 100644
--- a/studio/backend/core/inference/diffusion_families.py
+++ b/studio/backend/core/inference/diffusion_families.py
@@ -126,6 +126,13 @@ class DiffusionFamily:
     # 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)
+    # When set, deploying a LoRA trained on this family loads THIS repo instead of the
+    # checkpoint it was trained on -- for families whose release guidance is to train on one
+    # checkpoint but run adapters on another (Krea: train on Raw, preview on Turbo). Both
+    # sides must be the same precision so the swap never enlarges the load (unlike the
+    # nf4 -> bf16 gap that would risk an OOM on deploy). Unset elsewhere, so every other
+    # family deploys on the base it was trained on.
+    deploy_base_repo: Optional[str] = None
 
 
 # Keyed by architecture, not per model variant: a checkpoint's specific base repo
@@ -300,6 +307,9 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
         # inference/base repo.
         trainable = True,
         train_base_repos = ("krea/Krea-2-Raw", "krea/Krea-2-Turbo"),
+        # Per Krea's guidance, adapters trained on Raw are meant to run on Turbo; deploy
+        # previews them on Turbo (same bf16 precision, so the swap never enlarges the load).
+        deploy_base_repo = "krea/Krea-2-Turbo",
         # The checkpoint is exported bf16-only (the model card pins bfloat16); fp16 is
         # unvalidated upstream, so keep the fp16 fallback off like z-image.
         fp16_incompatible = True,
@@ -451,6 +461,10 @@ def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
 # (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),
+    # Krea 2 Turbo is distilled (TDM): 8 steps, no CFG -- matching the Create UI seed, so
+    # the OpenAI /v1/images/generations route uses the documented recipe instead of falling
+    # through to the generic (9, 0.0). "krea" collides with no other model id.
+    ("krea", 8, 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),
diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py
index 1239ddde0d..97684163af 100644
--- a/studio/backend/core/training/diffusion_train_common.py
+++ b/studio/backend/core/training/diffusion_train_common.py
@@ -281,6 +281,8 @@ def family_train_infos() -> list[dict[str, Any]]:
                 "precision_modes": dit_modes if is_dit else [],
                 "recommended_precision": dit_recommended if is_dit else "nf4",
                 "supports_compile": is_dit,
+                # Krea trains on Raw but previews adapters on Turbo; None elsewhere.
+                "deploy_base": fam.deploy_base_repo,
             }
         )
     return infos
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 95a44c4b19..d6d5241072 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -870,6 +870,10 @@ class DiffusionTrainableFamily(BaseModel):
     precision_modes: List[str] = Field(default_factory = list)
     recommended_precision: str = "nf4"
     supports_compile: bool = False
+    # When set, deploying a LoRA trained on this family previews it on this repo instead of
+    # the training base (Krea trains on Raw but runs adapters on Turbo). Null for families
+    # that deploy on the base they trained on.
+    deploy_base: Optional[str] = None
 
 
 class DiffusionTrainingInfoResponse(BaseModel):
diff --git a/studio/backend/tests/test_diffusion_krea2.py b/studio/backend/tests/test_diffusion_krea2.py
index 82649c1a6e..b17be7e38a 100644
--- a/studio/backend/tests/test_diffusion_krea2.py
+++ b/studio/backend/tests/test_diffusion_krea2.py
@@ -135,7 +135,11 @@ def test_load_krea2_pipeline_requires_krea_capable_diffusers(monkeypatch):
 
 def test_krea2_family_wiring():
     from core.inference.diffusion import _is_trusted_diffusion_repo
-    from core.inference.diffusion_families import detect_family, family_sd_cpp_supported
+    from core.inference.diffusion_families import (
+        default_generation_params,
+        detect_family,
+        family_sd_cpp_supported,
+    )
     from core.inference.diffusion_transformer_quant import TQ_INT8, exclude_tokens_for_scheme
 
     fam = detect_family("krea/Krea-2-Turbo")
@@ -147,6 +151,13 @@ def test_krea2_family_wiring():
     assert not family_sd_cpp_supported(fam)
     # Krea2TimestepEmbedding runs at M = batch; int8 (torch._int_mm, M > 16) must skip it.
     assert "time_embed" in exclude_tokens_for_scheme(TQ_INT8)
+    # Adapters train on Raw but run on Turbo, so the family carries a deploy override.
+    assert fam.deploy_base_repo == "krea/Krea-2-Turbo"
+    # The OpenAI /v1/images/generations route reads (steps, guidance) from this table; Krea
+    # Turbo is distilled (8 steps, no CFG), matching the Create UI seed instead of the
+    # generic (9, 0.0) fallback.
+    assert default_generation_params("krea/Krea-2-Turbo") == (8, 0.0)
+    assert default_generation_params("krea/Krea-2-Raw") == (8, 0.0)
 
 
 # ── training wiring ──────────────────────────────────────────────────────────
@@ -175,6 +186,10 @@ def test_krea2_training_registry():
     assert info["default_base"] == "krea/Krea-2-Raw"
     assert info["base_repos"] == ["krea/Krea-2-Raw", "krea/Krea-2-Turbo"]
     assert info["supports_compile"] is True
+    # Deploy previews the adapter on Turbo, not the Raw checkpoint it trained on, so the UI
+    # loads the distilled inference recipe; other families leave this None.
+    assert info["deploy_base"] == "krea/Krea-2-Turbo"
+    assert {i["name"]: i for i in family_train_infos()}["flux.1"]["deploy_base"] is None
 
 
 def test_krea2_spec_registered_with_authors_targets():
diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts
index eb8ecd785e..610e5ec935 100644
--- a/studio/frontend/src/features/images/api.ts
+++ b/studio/frontend/src/features/images/api.ts
@@ -443,6 +443,10 @@ export interface DiffusionTrainableFamily {
   recommended_precision?: string;
   // Whether the family's transformer can be torch.compile'd (gates the Speed > Compile row).
   supports_compile?: boolean;
+  // When set, deploying a LoRA trained on this family previews it on this repo instead of
+  // the checkpoint it was trained on (Krea trains on Raw but runs adapters on Turbo). Null
+  // for families that deploy on the base they trained on.
+  deploy_base?: string | null;
 }
 
 // Where diffusion training reads/writes on this Studio, plus usable dataset folders.
diff --git a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx
index 463ca86dc4..2bb2824d57 100644
--- a/studio/frontend/src/features/images/train/diffusion-train-panel.tsx
+++ b/studio/frontend/src/features/images/train/diffusion-train-panel.tsx
@@ -723,23 +723,38 @@ export function DiffusionTrainPanel({
     [poll],
   );
 
+  // Resolve the repo an adapter should be PREVIEWED on. Krea (and any family that trains on
+  // one checkpoint but runs adapters on another) declares a deploy_base: preview the adapter
+  // there instead of the training checkpoint, so the default Krea train-on-Raw flow does not
+  // load the adapter on Raw's non-distilled recipe. Only a recognised training base is
+  // overridden; a custom repo the user typed is respected as-is.
+  const deployBaseFor = useCallback(
+    (trainedBase: string, famName: string): string => {
+      const rec = info?.families?.find((f) => f.name === famName);
+      if (rec?.deploy_base && rec.base_repos.includes(trainedBase)) return rec.deploy_base;
+      return trainedBase;
+    },
+    [info?.families],
+  );
+
   const onDeployClick = useCallback(() => {
     if (!status?.catalog_path) {
       toast.error("The trained adapter is not available yet.");
       return;
     }
-    const baseRepo = status.base_model || (effectiveBase === CUSTOM_BASE ? customBase : effectiveBase);
-    if (!baseRepo) {
+    const trainedBase = status.base_model || (effectiveBase === CUSTOM_BASE ? customBase : effectiveBase);
+    if (!trainedBase) {
       toast.error("Could not determine the base model to load for this adapter.");
       return;
     }
+    const famName = status.family || family?.name || "";
     onDeploy?.({
-      baseRepo,
-      family: status.family || family?.name || "",
+      baseRepo: deployBaseFor(trainedBase, famName),
+      family: famName,
       catalogPath: status.catalog_path,
       trigger: instancePrompt.trim(),
     });
-  }, [status, baseChoice, customBase, family, instancePrompt, onDeploy]);
+  }, [status, effectiveBase, customBase, family, instancePrompt, onDeploy, deployBaseFor]);
 
   const numberField = (
     label: string,
@@ -1217,7 +1232,7 @@ export function DiffusionTrainPanel({
                     size="sm"
                     onClick={() =>
                       onDeploy?.({
-                        baseRepo: viewRun.base_model || "",
+                        baseRepo: deployBaseFor(viewRun.base_model || "", viewRun.family || ""),
                         family: viewRun.family || "",
                         catalogPath: viewRun.catalog_path || "",
                         trigger: viewRun.instance_prompt || "",

From 9ab68825ee3d4356a05e7f8056fc865afef06ebe Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 6 Jul 2026 14:00:39 +0000
Subject: [PATCH 17/17] Give Krea-2-Raw its undistilled 52-step recipe instead
 of the distilled default

Krea-2-Raw is in _TRUSTED_NON_GGUF_REPOS, so it is inference-loadable, but the generic
"krea" generation-defaults key matched it too and applied Turbo's distilled 8-step / no-CFG
recipe, producing degraded output on the undistilled base. Add a more specific krea-2-raw key
(52 steps, guidance 3.5 per the model card) ahead of the generic one, in both the backend
table and the frontend MODEL_DEFAULTS so the Studio UI and the OpenAI images route agree.
---
 studio/backend/core/inference/diffusion_families.py | 10 +++++++---
 studio/backend/tests/test_diffusion_krea2.py        |  5 +++--
 studio/frontend/src/features/images/images-page.tsx |  8 ++++++--
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py
index 53792feac2..a6cf3d7b03 100644
--- a/studio/backend/core/inference/diffusion_families.py
+++ b/studio/backend/core/inference/diffusion_families.py
@@ -461,9 +461,13 @@ def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
 # (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),
-    # Krea 2 Turbo is distilled (TDM): 8 steps, no CFG -- matching the Create UI seed, so
-    # the OpenAI /v1/images/generations route uses the documented recipe instead of falling
-    # through to the generic (9, 0.0). "krea" collides with no other model id.
+    # Krea 2 Raw is the undistilled base (both Turbo and Raw are in _TRUSTED_NON_GGUF_REPOS, so
+    # either is inference-loadable): its model card runs 52 steps at guidance 3.5. It must precede
+    # the generic "krea" key, or the distilled recipe below would degrade a Raw load to garbage.
+    ("krea-2-raw", 52, 3.5),
+    # Krea 2 Turbo is distilled (TDM): 8 steps, no CFG -- matching the Create UI seed, so the
+    # OpenAI /v1/images/generations route uses the documented recipe instead of falling through
+    # to the generic (9, 0.0). "krea" then covers Turbo and any other krea id but Raw (above).
     ("krea", 8, 0.0),
     ("flux.1-schnell", 4, 0.0),
     # Kontext (editing) before the generic flux.1: ~28 steps, lower guidance (~2.5).
diff --git a/studio/backend/tests/test_diffusion_krea2.py b/studio/backend/tests/test_diffusion_krea2.py
index b17be7e38a..86ff8fdcb0 100644
--- a/studio/backend/tests/test_diffusion_krea2.py
+++ b/studio/backend/tests/test_diffusion_krea2.py
@@ -155,9 +155,10 @@ def test_krea2_family_wiring():
     assert fam.deploy_base_repo == "krea/Krea-2-Turbo"
     # The OpenAI /v1/images/generations route reads (steps, guidance) from this table; Krea
     # Turbo is distilled (8 steps, no CFG), matching the Create UI seed instead of the
-    # generic (9, 0.0) fallback.
+    # generic (9, 0.0) fallback. Raw is the undistilled base (also inference-loadable) and runs
+    # its full 52-step / CFG 3.5 recipe, so its more specific key must win over the "krea" one.
     assert default_generation_params("krea/Krea-2-Turbo") == (8, 0.0)
-    assert default_generation_params("krea/Krea-2-Raw") == (8, 0.0)
+    assert default_generation_params("krea/Krea-2-Raw") == (52, 3.5)
 
 
 # ── training wiring ──────────────────────────────────────────────────────────
diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx
index 92379f1ddb..800ba0a0d9 100644
--- a/studio/frontend/src/features/images/images-page.tsx
+++ b/studio/frontend/src/features/images/images-page.tsx
@@ -211,8 +211,12 @@ const DEFAULT_GEN = { steps: 9, guidance: 0 };
 
 const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }> = [
   { match: "z-image-turbo", steps: 9, guidance: 0 },
-  // Krea 2 Turbo is distilled (TDM): 8 steps, no CFG (the base/midtrain checkpoints
-  // would want ~28 steps + CFG 4.5, but only Turbo is curated today).
+  // Krea 2 Raw is the undistilled base (also inference-loadable): its card runs 52 steps at
+  // guidance 3.5, so it must precede the distilled "krea-2" key below or a Raw load would run
+  // the 8-step recipe and produce garbage.
+  { match: "krea-2-raw", steps: 52, guidance: 3.5 },
+  // Krea 2 Turbo is distilled (TDM): 8 steps, no CFG. "krea-2" then covers Turbo (and any
+  // other krea id) but Raw, which is matched more specifically above.
   { match: "krea-2", steps: 8, guidance: 0 },
   { match: "flux.1-schnell", steps: 4, guidance: 0 },
   // Kontext (editing) before the generic flux.1: ~28 steps, lower guidance (~2.5).