diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index efb8491c..5411f78a 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/ui/src/app/jobs/new/SimpleJob.tsx @@ -332,19 +332,13 @@ export default function SimpleJob({ const gateUrl = modelArch.gateUrl as string; openDoc({ title: `Notes - ${modelArch.label}`, - description: ( -
- {modelArch.modelNotes} -
- ), + description:
{modelArch.modelNotes}
, }); }} className="w-full flex items-center gap-2 rounded-md bg-blue-950/60 border border-blue-800 px-3 py-2 text-sm text-blue-200 hover:bg-blue-900/60 text-left" > - - Model notes - + Model notes )} @@ -910,6 +904,39 @@ export default function SimpleJob({ )} )} + + <> + { + if (value) { + setJobConfig(true, 'config.process[0].train.do_guidance_loss'); + if (!jobConfig.config.process[0].train.guidance_loss_target) { + setJobConfig(3.0, 'config.process[0].train.guidance_loss_target'); + } + } else { + setJobConfig(undefined, 'config.process[0].train.do_guidance_loss'); + setJobConfig(undefined, 'config.process[0].train.guidance_loss_target'); + } + }} + /> + {jobConfig.config.process[0].train.do_guidance_loss && ( + <> + setJobConfig(value, 'config.process[0].train.guidance_loss_target')} + placeholder="eg. 3.0" + min={0} + /> + + )} + + diff --git a/ui/src/app/jobs/new/options.tsx b/ui/src/app/jobs/new/options.tsx index 29d31b04..90446759 100644 --- a/ui/src/app/jobs/new/options.tsx +++ b/ui/src/app/jobs/new/options.tsx @@ -717,6 +717,8 @@ export const modelArchs: ModelArch[] = [ 'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'], 'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'], 'config.process[0].train.cache_text_embeddings': [true, false], + 'config.process[0].train.do_guidance_loss': [true, undefined], + 'config.process[0].train.guidance_loss_target': [3.0, undefined], 'config.process[0].network.linear': [16, defaultLinearRank], 'config.process[0].network.linear_alpha': [16, defaultLinearRank], 'config.process[0].sample.num_frames': [107, 1], diff --git a/ui/src/docs.tsx b/ui/src/docs.tsx index e178bbb7..2542b376 100644 --- a/ui/src/docs.tsx +++ b/ui/src/docs.tsx @@ -351,6 +351,14 @@ const docs: { [key: string]: ConfigDoc } = { ), }, + 'train.guidance_loss_target': { + title: 'Guidance Loss Target', + description: ( + <> + For contrastive guidance loss, this is the target CGF to amplify predictions to. + + ), + }, }; export const getDoc = (key: string | null | undefined): ConfigDoc | null => { diff --git a/ui/src/types.ts b/ui/src/types.ts index 8f24f080..554dbc08 100644 --- a/ui/src/types.ts +++ b/ui/src/types.ts @@ -164,6 +164,8 @@ export interface TrainConfig { audio_loss_multiplier?: number; max_loss?: number | null; validation_config?: ValidationConfig; + do_guidance_loss?: boolean; + guidance_loss_target?: number; } export interface QuantizeKwargsConfig {