From 2d974219bf36aade162eb4ea2552efb10677b594 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 11:16:25 +0000 Subject: [PATCH] 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 || "",