github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Added training to the ui. Still testing, but everything seems to be working.

Browse Source
This commit is contained in:
Jaret Burkett
2025-08-16 05:51:37 -06:00
parent ca7bfa414b
commit 8ea2cf00f6
12 changed files with 268 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
ui/src/app/jobs/new/SimpleJob.tsx
View File

@@ -40,10 +40,25 @@ export default function SimpleJob({
const isVideoModel = !!(modelArch?.group === 'video');
let topBarClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 xl:grid-cols-5 gap-6';
const numTopCards = useMemo(() => {
let count = 4; // job settings, model config, target config, save config
if (modelArch?.additionalSections?.includes('model.multistage')) {
count += 1; // add multistage card
}
if (!modelArch?.disableSections?.includes('model.quantize')) {
count += 1; // add quantization card
}
return count;
}, [modelArch]);
if (modelArch?.disableSections?.includes('model.quantize')) {
topBarClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 xl:grid-cols-4 gap-6';
let topBarClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-4 xl:grid-cols-4 gap-6';
if (numTopCards == 5) {
topBarClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 xl:grid-cols-5 gap-6';
}
if (numTopCards == 6) {
topBarClass = 'grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 xl:grid-cols-3 2xl:grid-cols-6 gap-6';
}
const transformerQuantizationOptions: GroupedSelectOption[] | SelectOption[] = useMemo(() => {
@@ -91,7 +106,7 @@ export default function SimpleJob({
<>
<form onSubmit={handleSubmit} className="space-y-8">
<div className={topBarClass}>
<Card title="Job Settings">
<Card title="Job">
<TextInput
label="Training Name"
value={jobConfig.config.name}
@@ -124,7 +139,7 @@ export default function SimpleJob({
</Card>
{/* Model Configuration Section */}
<Card title="Model Configuration">
<Card title="Model">
<SelectInput
label="Model Architecture"
value={jobConfig.config.process[0].model.arch}
@@ -239,7 +254,32 @@ export default function SimpleJob({
/>
</Card>
)}
<Card title="Target Configuration">
{modelArch?.additionalSections?.includes('model.multistage') && (
<Card title="Multistage">
<FormGroup label="Stages to Train" docKey={'model.multistage'}>
<Checkbox
label="High Noise"
checked={jobConfig.config.process[0].model.model_kwargs?.train_high_noise || false}
onChange={value => setJobConfig(value, 'config.process[0].model.model_kwargs.train_high_noise')}
/>
<Checkbox
label="Low Noise"
checked={jobConfig.config.process[0].model.model_kwargs?.train_low_noise || false}
onChange={value => setJobConfig(value, 'config.process[0].model.model_kwargs.train_low_noise')}
/>
</FormGroup>
<NumberInput
label="Switch Every"
value={jobConfig.config.process[0].train.switch_boundary_every}
onChange={value => setJobConfig(value, 'config.process[0].train.switch_boundary_every')}
placeholder="eg. 1"
docKey={'train.switch_boundary_every'}
min={1}
required
/>
</Card>
)}
<Card title="Target">
<SelectInput
label="Target Type"
value={jobConfig.config.process[0].network?.type ?? 'lora'}
@@ -295,7 +335,7 @@ export default function SimpleJob({
</>
)}
</Card>
<Card title="Save Configuration">
<Card title="Save">
<SelectInput
label="Data Type"
value={jobConfig.config.process[0].save.dtype}
@@ -325,7 +365,7 @@ export default function SimpleJob({
</Card>
</div>
<div>
<Card title="Training Configuration">
<Card title="Training">
<div className="grid grid-cols-1 md:grid-cols-3 lg:grid-cols-5 gap-6">
<div>
<NumberInput
@@ -645,7 +685,7 @@ export default function SimpleJob({
</Card>
</div>
<div>
<Card title="Sample Configuration">
<Card title="Sample">
<div
className={
isVideoModel

1
ui/src/app/jobs/new/jobConfig.ts
View File

@@ -78,6 +78,7 @@ export const defaultJobConfig: JobConfig = {
diff_output_preservation: false,
diff_output_preservation_multiplier: 1.0,
diff_output_preservation_class: 'person',
switch_boundary_every: 1,
},
model: {
name_or_path: 'ostris/Flex.1-alpha',

46
ui/src/app/jobs/new/options.ts
View File

@@ -3,7 +3,13 @@ import { GroupedSelectOption, SelectOption } from '@/types';
type Control = 'depth' | 'line' | 'pose' | 'inpaint';
type DisableableSections = 'model.quantize' | 'train.timestep_type' | 'network.conv';
type AdditionalSections = 'datasets.control_path' | 'datasets.do_i2v' | 'sample.ctrl_img' | 'datasets.num_frames' | 'model.low_vram';
type AdditionalSections =
| 'datasets.control_path'
| 'datasets.do_i2v'
| 'sample.ctrl_img'
| 'datasets.num_frames'
| 'model.multistage'
| 'model.low_vram';
type ModelGroup = 'image' | 'video';
export interface ModelArch {
@@ -121,7 +127,7 @@ export const modelArchs: ModelArch[] = [
'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
'config.process[0].sample.num_frames': [41, 1],
'config.process[0].sample.fps': [15, 1],
'config.process[0].sample.fps': [16, 1],
},
disableSections: ['network.conv'],
additionalSections: ['datasets.num_frames', 'model.low_vram'],
@@ -139,7 +145,7 @@ export const modelArchs: ModelArch[] = [
'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
'config.process[0].sample.num_frames': [41, 1],
'config.process[0].sample.fps': [15, 1],
'config.process[0].sample.fps': [16, 1],
'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
},
disableSections: ['network.conv'],
@@ -158,7 +164,7 @@ export const modelArchs: ModelArch[] = [
'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
'config.process[0].sample.num_frames': [41, 1],
'config.process[0].sample.fps': [15, 1],
'config.process[0].sample.fps': [16, 1],
'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
},
disableSections: ['network.conv'],
@@ -177,11 +183,41 @@ export const modelArchs: ModelArch[] = [
'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
'config.process[0].sample.num_frames': [41, 1],
'config.process[0].sample.fps': [15, 1],
'config.process[0].sample.fps': [16, 1],
},
disableSections: ['network.conv'],
additionalSections: ['datasets.num_frames', 'model.low_vram'],
},
{
name: 'wan22_14b:t2v',
label: 'Wan 2.2 (14B)',
group: 'video',
isVideoModel: true,
defaults: {
// default updates when [selected, unselected] in the UI
'config.process[0].model.name_or_path': ['ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16', defaultNameOrPath],
'config.process[0].model.quantize': [true, false],
'config.process[0].model.quantize_te': [true, false],
'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
'config.process[0].sample.num_frames': [41, 1],
'config.process[0].sample.fps': [16, 1],
'config.process[0].model.low_vram': [true, false],
'config.process[0].train.timestep_type': ['linear', 'sigmoid'],
'config.process[0].model.model_kwargs': [
{
train_high_noise: true,
train_low_noise: true,
},
{},
],
},
disableSections: ['network.conv'],
additionalSections: ['datasets.num_frames', 'model.low_vram', 'model.multistage'],
// accuracyRecoveryAdapters: {
// '3 bit with ARA': 'uint3|ostris/accuracy_recovery_adapters/wan22_14b_t2i_torchao_uint3.safetensors',
// },
},
{
name: 'wan22_5b',
label: 'Wan 2.2 TI2V (5B)',

4
ui/src/components/formInputs.tsx
View File

@@ -283,7 +283,7 @@ export const FormGroup: React.FC<FormGroupProps> = props => {
return (
<div className={classNames(className)}>
{label && (
<label className={labelClasses}>
<label className={classNames(labelClasses, 'mb-2')}>
{label}{' '}
{doc && (
<div className="inline-block ml-1 text-xs text-gray-500 cursor-pointer" onClick={() => openDoc(doc)}>
@@ -292,7 +292,7 @@ export const FormGroup: React.FC<FormGroupProps> = props => {
)}
</label>
)}
<div className="px-4 space-y-2">{children}</div>
<div className="space-y-2">{children}</div>
</div>
);
};

30
ui/src/docs.tsx
View File

@@ -111,6 +111,36 @@ const docs: { [key: string]: ConfigDoc } = {
</>
),
},
'model.multistage': {
title: 'Stages to Train',
description: (
<>
Some models have multi stage networks that are trained and used separately in the denoising process. Most
common, is to have 2 stages. One for high noise and one for low noise. You can choose to train both stages at
once or train them separately. If trained at the same time, The trainer will alternate between training each
model every so many steps and will output 2 different LoRAs. If you choose to train only one stage, the
trainer will only train that stage and output a single LoRA.
</>
),
},
'train.switch_boundary_every': {
title: 'Switch Boundary Every',
description: (
<>
When training a model with multiple stages, this setting controls how often the trainer will switch between
training each stage.
<br />
<br />
For low vram settings, the model not being trained will be unloaded from the gpu to save memory. This takes some
time to do, so it is recommended to alternate less often when using low vram. A setting like 10 or 20 is
recommended for low vram settings.
<br />
<br />
The swap happens at the batch level, meaning it will swap between a gradient accumulation steps. To train both
stages in a single step, set them to switch every 1 step and set gradient accumulation to 2.
</>
),
},
};
export const getDoc = (key: string | null | undefined): ConfigDoc | null => {

1
ui/src/types.ts
View File

@@ -119,6 +119,7 @@ export interface TrainConfig {
diff_output_preservation: boolean;
diff_output_preservation_multiplier: number;
diff_output_preservation_class: string;
switch_boundary_every: number;
}
export interface QuantizeKwargsConfig {