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 {