Add LTX-2 Support (#644)

* WIP, adding support for LTX2

* Training on images working

* Fix loading comfy models

* Handle converting and deconverting lora so it matches original format

* Reworked ui to habdle ltx and propert dataset default overwriting.

* Update the way lokr saves to it is more compatable with comfy

* Audio loading and synchronization/resampling is working

* Add audio to training. Does it work? Maybe, still testing.

* Fixed fps default issue for sound

* Have ui set fps for accurate audio mapping on ltx

* Added audio procession options to the ui for ltx

* Clean up requirements

ui/src/app/api/img/delete/route.ts

@@ -15,7 +15,7 @@ export async function POST(request: Request) {
     }
 
     // make sure it is an image
-    if (!/\.(jpg|jpeg|png|bmp|gif|tiff|webp)$/i.test(imgPath.toLowerCase())) {
+    if (!/\.(jpg|jpeg|png|bmp|gif|tiff|webp|mp4)$/i.test(imgPath.toLowerCase())) {
       return NextResponse.json({ error: 'Not an image' }, { status: 400 });
     }
 

ui/src/app/api/jobs/[jobID]/samples/route.ts

@@ -29,7 +29,7 @@ export async function GET(request: NextRequest, { params }: { params: { jobID: s
   const samples = fs
     .readdirSync(samplesFolder)
     .filter(file => {
-      return file.endsWith('.png') || file.endsWith('.jpg') || file.endsWith('.jpeg') || file.endsWith('.webp');
+      return file.endsWith('.png') || file.endsWith('.jpg') || file.endsWith('.jpeg') || file.endsWith('.webp') || file.endsWith('.mp4');
     })
     .map(file => {
       return path.join(samplesFolder, file);

ui/src/app/jobs/new/SimpleJob.tsx

@@ -862,6 +862,48 @@ export default function SimpleJob({
                             docKey="datasets.do_i2v"
                           />
                         )}
+                        {modelArch?.additionalSections?.includes('datasets.do_audio') && (
+                          <Checkbox
+                            label="Do Audio"
+                            checked={dataset.do_audio || false}
+                            onChange={value => {
+                              if (!value) {
+                                setJobConfig(undefined, `config.process[0].datasets[${i}].do_audio`);
+                              } else {
+                                setJobConfig(value, `config.process[0].datasets[${i}].do_audio`);
+                              }
+                            }}
+                            docKey="datasets.do_audio"
+                          />
+                        )}
+                        {modelArch?.additionalSections?.includes('datasets.audio_normalize') && (
+                          <Checkbox
+                            label="Audio Normalize"
+                            checked={dataset.audio_normalize || false}
+                            onChange={value => {
+                              if (!value) {
+                                setJobConfig(undefined, `config.process[0].datasets[${i}].audio_normalize`);
+                              } else {
+                                setJobConfig(value, `config.process[0].datasets[${i}].audio_normalize`);
+                              }
+                            }}
+                            docKey="datasets.audio_normalize"
+                          />
+                        )}
+                        {modelArch?.additionalSections?.includes('datasets.audio_preserve_pitch') && (
+                          <Checkbox
+                            label="Audio Preserve Pitch"
+                            checked={dataset.audio_preserve_pitch || false}
+                            onChange={value => {
+                              if (!value) {
+                                setJobConfig(undefined, `config.process[0].datasets[${i}].audio_preserve_pitch`);
+                              } else {
+                                setJobConfig(value, `config.process[0].datasets[${i}].audio_preserve_pitch`);
+                              }
+                            }}
+                            docKey="datasets.audio_preserve_pitch"
+                          />
+                        )}
                       </FormGroup>
                       <FormGroup label="Flipping" docKey={'datasets.flip'} className="mt-2">
                         <Checkbox

ui/src/app/jobs/new/jobConfig.ts

@@ -14,7 +14,6 @@ export const defaultDatasetConfig: DatasetConfig = {
   controls: [],
   shrink_video_to_frames: true,
   num_frames: 1,
-  do_i2v: true,
   flip_x: false,
   flip_y: false,
 };

ui/src/app/jobs/new/options.ts

@@ -17,6 +17,9 @@ type AdditionalSections =
   | 'datasets.control_path'
   | 'datasets.multi_control_paths'
   | 'datasets.do_i2v'
+  | 'datasets.do_audio'
+  | 'datasets.audio_normalize'
+  | 'datasets.audio_preserve_pitch'
   | 'sample.ctrl_img'
   | 'sample.multi_ctrl_imgs'
   | 'datasets.num_frames'
@@ -288,6 +291,7 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].sample.width': [768, 1024],
       'config.process[0].sample.height': [768, 1024],
       'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+      'config.process[0].datasets[x].do_i2v': [true, undefined],
     },
     disableSections: ['network.conv'],
     additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.low_vram', 'datasets.do_i2v'],
@@ -601,6 +605,31 @@ export const modelArchs: ModelArch[] = [
     disableSections: ['network.conv'],
     additionalSections: ['model.low_vram', 'model.layer_offloading'],
   },
+  {
+    name: 'ltx2',
+    label: 'LTX-2',
+    group: 'video',
+    isVideoModel: true,
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Lightricks/LTX-2', defaultNameOrPath],
+      'config.process[0].model.quantize': [true, false],
+      'config.process[0].model.quantize_te': [true, false],
+      'config.process[0].model.low_vram': [true, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].sample.num_frames': [121, 1],
+      'config.process[0].sample.fps': [24, 1],
+      'config.process[0].sample.width': [768, 1024],
+      'config.process[0].sample.height': [768, 1024],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+      'config.process[0].datasets[x].do_i2v': [false, undefined],
+      'config.process[0].datasets[x].do_audio': [true, undefined],
+      'config.process[0].datasets[x].fps': [24, undefined],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['datasets.num_frames', 'model.layer_offloading', 'model.low_vram', 'datasets.do_audio', 'datasets.audio_normalize', 'datasets.audio_preserve_pitch'],
+  },
 ].sort((a, b) => {
   // Sort by label, case-insensitive
   return a.label.localeCompare(b.label, undefined, { sensitivity: 'base' });

ui/src/app/jobs/new/utils.ts

@@ -2,6 +2,25 @@ import { GroupedSelectOption, JobConfig, SelectOption } from '@/types';
 import { modelArchs, ModelArch } from './options';
 import { objectCopy } from '@/utils/basic';
 
+const expandDatasetDefaults = (
+  defaults: { [key: string]: any },
+  numDatasets: number,
+): { [key: string]: any } => {
+  // expands the defaults for datasets[x] to datasets[0], datasets[1], etc.
+  const expandedDefaults: { [key: string]: any } = { ...defaults };
+  for (const key in defaults) {
+    if (key.includes('datasets[x].')) {
+      for (let i = 0; i < numDatasets; i++) {
+        const datasetKey = key.replace('datasets[x].', `datasets[${i}].`);
+        const v = defaults[key];
+        expandedDefaults[datasetKey] = Array.isArray(v) ? [...v] : objectCopy(v);
+      }
+      delete expandedDefaults[key];
+    }
+  }
+  return expandedDefaults;
+};
+
 export const handleModelArchChange = (
   currentArchName: string,
   newArchName: string,
@@ -39,16 +58,11 @@ export const handleModelArchChange = (
     }
   }
 
-  // revert defaults from previous model
-  for (const key in currentArch.defaults) {
-    setJobConfig(currentArch.defaults[key][1], key);
-  }
+  const numDatasets = jobConfig.config.process[0].datasets.length;
+
+  let currentDefaults = expandDatasetDefaults(currentArch.defaults || {}, numDatasets);
+  let newDefaults = expandDatasetDefaults(newArch?.defaults || {}, numDatasets);
 
-  if (newArch?.defaults) {
-    for (const key in newArch.defaults) {
-      setJobConfig(newArch.defaults[key][0], key);
-    }
-  }
   // set new model
   setJobConfig(newArchName, 'config.process[0].model.arch');
 
@@ -79,27 +93,27 @@ export const handleModelArchChange = (
       if (newDataset.control_path_1 && newDataset.control_path_1 !== '') {
         newDataset.control_path = newDataset.control_path_1;
       }
-      if (newDataset.control_path_1) {
+      if ('control_path_1' in newDataset) {
         delete newDataset.control_path_1;
       }
-      if (newDataset.control_path_2) {
+      if ('control_path_2' in newDataset) {
         delete newDataset.control_path_2;
       }
-      if (newDataset.control_path_3) {
+      if ('control_path_3' in newDataset) {
         delete newDataset.control_path_3;
       }
     } else {
       // does not have control images
-      if (newDataset.control_path) {
+      if ('control_path' in newDataset) {
         delete newDataset.control_path;
       }
-      if (newDataset.control_path_1) {
+      if ('control_path_1' in newDataset) {
         delete newDataset.control_path_1;
       }
-      if (newDataset.control_path_2) {
+      if ('control_path_2' in newDataset) {
         delete newDataset.control_path_2;
       }
-      if (newDataset.control_path_3) {
+      if ('control_path_3' in newDataset) {
         delete newDataset.control_path_3;
       }
     }
@@ -120,4 +134,13 @@ export const handleModelArchChange = (
     return newSample;
   });
   setJobConfig(samples, 'config.process[0].sample.samples');
+
+  // revert defaults from previous model
+  for (const key in currentDefaults) {
+    setJobConfig(currentDefaults[key][1], key);
+  }
+
+  for (const key in newDefaults) {
+    setJobConfig(newDefaults[key][0], key);
+  }
 };