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).
This commit is contained in:
Daniel Han
parent
099357cf40
commit
2d974219bf
6 changed files with 61 additions and 7 deletions
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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 || "",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue