From dc290bdf71b28e1b70d26ec42a0c462961d675aa Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 03:23:30 +0000 Subject: [PATCH] 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