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

@@ -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",
     }
 )
 

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,

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)

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)."
+    ),
 }
 
 

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