/add-training-support

name: add-training-support
description: Wire an existing ecosystem into the LoRA training system so it appears as a trainable base model in the training form. Adds the base-model entry, schema enums, orchestrator validation, feature flag, and per-ecosystem default params. Use when a model family (already in basemodel.constants.ts) needs training support — e.g. AI-Toolkit ecosystems like Anima, HiDream, Boogu. Always checks @civitai/client for the ecosystem's training input type first.

Add Training Support

Wires an existing ecosystem into the **training** form (the "Train a LoRA" flow), distinct from the generation form. After this, the ecosystem shows up as a selectable base model in Training Step 1 and submits a valid `training` step to the orchestrator.

Training is a **separate subsystem from generation**. The generation handlers/graphs in `src/server/services/orchestrator/ecosystems/*` and `src/shared/data-graph/generation/*` are NOT part of this — do not touch them. The files below are the training path.

When to use

• A new model family was just added via `add-ecosystem` and now needs to be trainable
• Re-enabling training for an ecosystem that was commented out
• The orchestrator/`@civitai/client` gained a new `*TrainingInput` type

Prerequisites

The ecosystem, base model, family and license must already exist in [basemodel.constants.ts](src/shared/constants/basemodel.constants.ts) (`ECO.<Name>`, an `EcosystemRecord`, a `BaseModelRecord`). If any are missing, run **add-ecosystem** first.

Almost every ecosystem added recently is **AI-Toolkit-based** (engine `'ai-toolkit'`, mandatory). The worked reference is **`anima`** — it is the simplest complete example (image, AI-Toolkit-only, no `modelVariant`). The single highest-leverage move when adding a new ecosystem: **grep every file below for `anima` and add a parallel entry.**

Step 0 — Check `@civitai/client` for the training input type

**Always** confirm the SDK ships the ecosystem's training type before wiring. It tells you the exact field shape, whether `modelVariant` is required, and any fixed fields.

grep -niE "<Name>AiToolkitTrainingInput|ecosystem: '<name>'" \
  node_modules/@civitai/client/dist/generated/types.gen.d.ts

Note from the type:

• **`ecosystem` literal** — the exact string the orchestrator expects (e.g. `'boogu'`).
• **`modelVariant`** — present (e.g. flux1 `'dev'|'schnell'`, wan `'2.1'|'2.2'`) or absent.
• **Fixed/constrained fields** — e.g. Boogu declares `batchSize?: null | number` "Fixed at 1

for this ecosystem". Honor these in the UI param bounds.

• **`readonly` outputs** (`defaultSteps`, `storageBuzzPerEpoch`, `maxBatchSize`, …) — server-

computed; never send them.

If the type is missing in the installed version, check the latest (`npm view @civitai/client versions --json | tail`) and bump with `pnpm add @civitai/client@<v>`. If still absent, the orchestrator dispatch (`training.orch.ts`) already casts `ecosystem as any`, so it works — but prefer a typed SDK.

Step 1 — Gather the ecosystem's training defaults

Get these from the user, the model card, or an orchestrator `whatif` sample request:

• **Default steps** (the AI-Toolkit primary length knob — drives pricing)
• **lr** (unet LR), whether the **text encoder** is trained (usually disabled)
• **networkDim / networkAlpha**, **lrScheduler**, **optimizerType**, **noiseOffset**
• **batch size** bounds (many AI-Toolkit image ecosystems are fixed at 1)
• **resolution** (image ecosystems are typically 1024)
• The **base-model AIR**. If the model isn't on civitai yet, use the HF repository URN from the

sample as a placeholder and leave a comment to swap in `urn:air:<eco>:checkpoint:civitai:<modelId>@<versionId>` once uploaded. For AI-Toolkit-only ecosystems this AIR is NOT sent to the orchestrator (it resolves the base model from the ecosystem); it's used for UI display / `getTrainingFields.getModel`.

Pick a stable **base-model key** (the `trainingModelInfo` key, e.g. `boogu`) and a **baseType** (the `TrainingBaseModelType`, e.g. `'boogu'`). They can differ (SD has `sd_1_5`/`anime`/… keys all mapping to baseType `sd15`), but for a single-checkpoint ecosystem keep them the same.

Step 2 — Confirm the plan with the user

Adding training support for: <Name>  (baseType: <baseType>, key: <key>)
Engine: ai-toolkit (mandatory)   Variant: <none | enum>
Defaults: steps <n>, lr <n>, dim/alpha <n>/<n>, scheduler <x>, batch <n> (max <n>), res <n>
AIR: <placeholder|civitai urn>   Feature flag: <name>Training / <name>-training (Flipt)

Wait for confirmation, then make all edits in one pass.

Step 3 — Edits

3a. `src/utils/training.ts` (the core — ~7 spots)

1. **`trainingBaseModelTypesImage`** (or `…Video` / `…Audio`) — add `'<baseType>'`. 2. **`aiToolkitStepDefault`** — add a branch returning the ecosystem's default steps if it differs from 2000. 3. **`aiToolkitBatchMax`** — add a branch only if max batch > 1 (default returns 1). 4. **`trainingModelInfo`** — add the `<key>: { label, pretty, type: '<baseType>', description, air, baseModel: '<BaseModelName>', isNew: true, aiToolkit: { ecosystem: '<eco>' [, modelVariant] } }` entry. `baseModel` MUST match the `BaseModelRecord.name` / ecosystem `key` in basemodel.constants.ts exactly (a mismatch produces a malformed AIR on training completion and fails the post-train scan with 400 — see the inline comment on the `hidream_o1` entry). 5. **`baseTypeToEcosystem`** — add `<baseType>: '<eco>'`. 6. **`isAiToolkitSupported`** `supportedTypes` — add `'<baseType>'`. 7. **`isAiToolkitMandatory`** `mandatoryTypes` — add `'<baseType>'` (for AI-Toolkit-only ecosystems). This auto-enables sample-prompt requirements and AI-Toolkit gating. 8. **`getDefaultEngine`** — add `if (baseType === '<baseType>') return 'ai-toolkit';`.

3b. `src/server/schema/model-version.schema.ts`

• Add `export const trainingDetailsBaseModels<Name> = ['<key>'] as const;`