update

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

@@ -1,8 +1,7 @@
-import { JobConfig, DatasetConfig } from '@/types';
+import { JobConfig, DatasetConfig, SliderConfig } from '@/types';
 
 export const defaultDatasetConfig: DatasetConfig = {
   folder_path: '/path/to/images/folder',
-  control_path: null,
   mask_path: null,
   mask_min_value: 0.1,
   default_caption: '',
@@ -20,13 +19,22 @@ export const defaultDatasetConfig: DatasetConfig = {
   flip_y: false,
 };
 
+export const defaultSliderConfig: SliderConfig = {
+  guidance_strength: 3.0,
+  anchor_strength: 1.0,
+  positive_prompt: 'person who is happy',
+  negative_prompt: 'person who is sad',
+  target_class: 'person',
+  anchor_class: '',
+};
+
 export const defaultJobConfig: JobConfig = {
   job: 'extension',
   config: {
     name: 'my_first_lora_v1',
     process: [
       {
-        type: 'ui_trainer',
+        type: 'diffusion_trainer',
         training_folder: 'output',
         sqlite_db_path: './aitk_db.db',
         device: 'cuda',
@@ -53,53 +61,55 @@ export const defaultJobConfig: JobConfig = {
         },
         datasets: [defaultDatasetConfig],
         train: {
-          batch_size: 4,
+          batch_size: 1,
           bypass_guidance_embedding: true,
-          steps: 1200,
+          steps: 3000,
           gradient_accumulation: 1,
           train_unet: true,
           train_text_encoder: false,
           gradient_checkpointing: true,
           noise_scheduler: 'flowmatch',
           optimizer: 'adamw8bit',
-          timestep_type: 'weighted',
+          timestep_type: 'sigmoid',
           content_or_style: 'balanced',
           optimizer_params: {
             weight_decay: 1e-4,
           },
           unload_text_encoder: false,
           cache_text_embeddings: false,
-          lr: 0.0003,
+          lr: 0.0001,
           ema_config: {
             use_ema: false,
             ema_decay: 0.99,
           },
-          skip_first_sample: true,
-          disable_sampling: true,
+          skip_first_sample: false,
+          force_first_sample: false,
+          disable_sampling: false,
           dtype: 'bf16',
           diff_output_preservation: false,
           diff_output_preservation_multiplier: 1.0,
           diff_output_preservation_class: 'person',
           switch_boundary_every: 1,
+          loss_type: 'mse',
         },
         model: {
-          name_or_path: 'Qwen/Qwen-Image',
+          name_or_path: 'ostris/Flex.1-alpha',
           quantize: true,
           qtype: 'qfloat8',
           quantize_te: true,
           qtype_te: 'qfloat8',
-          arch: 'qwen_image',
+          arch: 'flex1',
           low_vram: false,
           model_kwargs: {},
         },
         sample: {
           sampler: 'flowmatch',
-          sample_every: 1500,
+          sample_every: 250,
           width: 1024,
           height: 1024,
           samples: [
             {
-              prompt: 'woman with red hair, playing chess at the park, bomb going off in the background'
+              prompt: 'woman with red hair, playing chess at the park, bomb going off in the background',
             },
             {
               prompt: 'a woman holding a coffee cup, in a beanie, sitting at a cafe',
@@ -107,6 +117,30 @@ export const defaultJobConfig: JobConfig = {
             {
               prompt: 'a horse is a DJ at a night club, fish eye lens, smoke machine, lazer lights, holding a martini',
             },
+            {
+              prompt:
+                'a man showing off his cool new t shirt at the beach, a shark is jumping out of the water in the background',
+            },
+            {
+              prompt: 'a bear building a log cabin in the snow covered mountains',
+            },
+            {
+              prompt: 'woman playing the guitar, on stage, singing a song, laser lights, punk rocker',
+            },
+            {
+              prompt: 'hipster man with a beard, building a chair, in a wood shop',
+            },
+            {
+              prompt:
+                'photo of a man, white background, medium shot, modeling clothing, studio lighting, white backdrop',
+            },
+            {
+              prompt: "a man holding a sign that says, 'this is a sign'",
+            },
+            {
+              prompt:
+                'a bulldog, in a post apocalyptic world, with a shotgun, in a leather jacket, in a desert, with a motorcycle',
+            },
           ],
           neg: '',
           seed: 42,
@@ -142,5 +176,16 @@ export const migrateJobConfig = (jobConfig: JobConfig): JobConfig => {
     jobConfig.config.process[0].sample.samples = newSamples;
     delete jobConfig.config.process[0].sample.prompts;
   }
+
+  // upgrade job from ui_trainer to diffusion_trainer
+  if (jobConfig?.config?.process && jobConfig.config.process[0]?.type === 'ui_trainer') {
+    jobConfig.config.process[0].type = 'diffusion_trainer';
+  }
+
+  if ('auto_memory' in jobConfig.config.process[0].model) {
+    jobConfig.config.process[0].model.layer_offloading = (jobConfig.config.process[0].model.auto_memory ||
+      false) as boolean;
+    delete jobConfig.config.process[0].model.auto_memory;
+  }
   return jobConfig;
 };

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

@@ -1,4 +1,5 @@
-import { GroupedSelectOption, SelectOption } from '@/types';
+import { GroupedSelectOption, SelectOption, JobConfig } from '@/types';
+import { defaultSliderConfig } from './jobConfig';
 
 type Control = 'depth' | 'line' | 'pose' | 'inpaint';
 
@@ -8,6 +9,7 @@ type DisableableSections =
   | 'network.conv'
   | 'trigger_word'
   | 'train.diff_output_preservation'
+  | 'train.blank_prompt_preservation'
   | 'train.unload_text_encoder'
   | 'slider';
 
@@ -19,7 +21,10 @@ type AdditionalSections =
   | 'sample.multi_ctrl_imgs'
   | 'datasets.num_frames'
   | 'model.multistage'
-  | 'model.low_vram';
+  | 'model.layer_offloading'
+  | 'model.low_vram'
+  | 'model.qie.match_target_res'
+  | 'model.assistant_lora_path';
 type ModelGroup = 'image' | 'instruction' | 'video';
 
 export interface ModelArch {
@@ -221,7 +226,7 @@ export const modelArchs: ModelArch[] = [
       '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': [false, false],
+      'config.process[0].model.low_vram': [true, false],
       'config.process[0].train.timestep_type': ['linear', 'sigmoid'],
       'config.process[0].model.model_kwargs': [
         {
@@ -232,7 +237,7 @@ export const modelArchs: ModelArch[] = [
       ],
     },
     disableSections: ['network.conv'],
-    additionalSections: ['datasets.num_frames', 'model.low_vram', 'model.multistage'],
+    additionalSections: ['datasets.num_frames', 'model.low_vram', 'model.multistage', 'model.layer_offloading'],
     accuracyRecoveryAdapters: {
       // '3 bit with ARA': 'uint3|ostris/accuracy_recovery_adapters/wan22_14b_t2i_torchao_uint3.safetensors',
       '4 bit with ARA': 'uint4|ostris/accuracy_recovery_adapters/wan22_14b_t2i_torchao_uint4.safetensors',
@@ -252,7 +257,7 @@ export const modelArchs: ModelArch[] = [
       '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': [false, false],
+      'config.process[0].model.low_vram': [true, false],
       'config.process[0].train.timestep_type': ['linear', 'sigmoid'],
       'config.process[0].model.model_kwargs': [
         {
@@ -263,7 +268,13 @@ export const modelArchs: ModelArch[] = [
       ],
     },
     disableSections: ['network.conv'],
-    additionalSections: ['sample.ctrl_img', 'datasets.num_frames', 'model.low_vram', 'model.multistage'],
+    additionalSections: [
+      'sample.ctrl_img',
+      'datasets.num_frames',
+      'model.low_vram',
+      'model.multistage',
+      'model.layer_offloading',
+    ],
     accuracyRecoveryAdapters: {
       '4 bit with ARA': 'uint4|ostris/accuracy_recovery_adapters/wan22_14b_i2v_torchao_uint4.safetensors',
     },
@@ -278,7 +289,7 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].model.name_or_path': ['Wan-AI/Wan2.2-TI2V-5B-Diffusers', defaultNameOrPath],
       'config.process[0].model.quantize': [true, false],
       'config.process[0].model.quantize_te': [true, false],
-      'config.process[0].model.low_vram': [false, 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],
@@ -313,7 +324,7 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].model.name_or_path': ['Qwen/Qwen-Image', defaultNameOrPath],
       'config.process[0].model.quantize': [true, false],
       'config.process[0].model.quantize_te': [true, false],
-      'config.process[0].model.low_vram': [false, 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].train.timestep_type': ['weighted', 'sigmoid'],
@@ -322,7 +333,7 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].train.lr': [0.0003, 0.0001],
     },
     disableSections: ['network.conv'],
-    additionalSections: ['model.low_vram'],
+    additionalSections: ['model.low_vram', 'model.layer_offloading'],
     accuracyRecoveryAdapters: {
       '3 bit with ARA': 'uint3|ostris/accuracy_recovery_adapters/qwen_image_torchao_uint3.safetensors',
     },
@@ -336,7 +347,7 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].model.name_or_path': ['Qwen/Qwen-Image-Edit', defaultNameOrPath],
       'config.process[0].model.quantize': [true, false],
       'config.process[0].model.quantize_te': [true, false],
-      'config.process[0].model.low_vram': [false, 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].train.timestep_type': ['weighted', 'sigmoid'],
@@ -345,7 +356,7 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].train.lr': [0.0003, 0.0001],
     },
     disableSections: ['network.conv'],
-    additionalSections: ['datasets.control_path', 'sample.ctrl_img', 'model.low_vram'],
+    additionalSections: ['datasets.control_path', 'sample.ctrl_img', 'model.low_vram', 'model.layer_offloading'],
     accuracyRecoveryAdapters: {
       '3 bit with ARA': 'uint3|ostris/accuracy_recovery_adapters/qwen_image_edit_torchao_uint3.safetensors',
     },
@@ -359,7 +370,7 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].model.name_or_path': ['Qwen/Qwen-Image-Edit-2509', defaultNameOrPath],
       'config.process[0].model.quantize': [true, false],
       'config.process[0].model.quantize_te': [true, false],
-      'config.process[0].model.low_vram': [false, false],
+      'config.process[0].model.low_vram': [true, false],
       'config.process[0].train.unload_text_encoder': [false, false],
       'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
       'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
@@ -367,9 +378,21 @@ export const modelArchs: ModelArch[] = [
       'config.process[0].model.qtype': ['qfloat8', 'qfloat8'],
       'config.process[0].train.batch_size': [4, 1],
       'config.process[0].train.lr': [0.0003, 0.0001],
+      'config.process[0].model.model_kwargs': [
+        {
+          match_target_res: false,
+        },
+        {},
+      ],
     },
     disableSections: ['network.conv', 'train.unload_text_encoder'],
-    additionalSections: ['datasets.control_path', 'sample.ctrl_img', 'model.low_vram'],
+    additionalSections: [
+      'datasets.multi_control_paths',
+      'sample.multi_ctrl_imgs',
+      'model.low_vram',
+      'model.layer_offloading',
+      'model.qie.match_target_res',
+    ],
     accuracyRecoveryAdapters: {
       '3 bit with ARA': 'uint3|ostris/accuracy_recovery_adapters/qwen_image_edit_2509_torchao_uint3.safetensors',
     },
@@ -460,6 +483,84 @@ export const modelArchs: ModelArch[] = [
     disableSections: ['network.conv'],
     additionalSections: ['datasets.control_path', 'sample.ctrl_img'],
   },
+  {
+    name: 'flux2',
+    label: 'FLUX.2',
+    group: 'image',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['black-forest-labs/FLUX.2-dev', 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].train.unload_text_encoder': [false, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+      'config.process[0].model.qtype': ['qfloat8', 'qfloat8'],
+      'config.process[0].model.model_kwargs': [
+        {
+          match_target_res: false,
+        },
+        {},
+      ],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: [
+      'datasets.multi_control_paths',
+      'sample.multi_ctrl_imgs',
+      'model.low_vram',
+      'model.layer_offloading',
+      'model.qie.match_target_res',
+    ],
+  },
+  {
+    name: 'zimage:turbo',
+    label: 'Z-Image Turbo (w/ Training Adapter)',
+    group: 'image',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['Tongyi-MAI/Z-Image-Turbo', 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].train.unload_text_encoder': [false, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+      'config.process[0].model.qtype': ['qfloat8', 'qfloat8'],
+      'config.process[0].model.assistant_lora_path': [
+        'ostris/zimage_turbo_training_adapter/zimage_turbo_training_adapter_v2.safetensors',
+        undefined,
+      ],
+      'config.process[0].sample.guidance_scale': [1, 4],
+      'config.process[0].sample.sample_steps': [8, 25],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['model.low_vram', 'model.layer_offloading', 'model.assistant_lora_path'],
+  },
+  {
+    name: 'zimage:deturbo',
+    label: 'Z-Image De-Turbo (De-Distilled)',
+    group: 'image',
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      'config.process[0].model.name_or_path': ['ostris/Z-Image-De-Turbo', defaultNameOrPath],
+      'config.process[0].model.extras_name_or_path': ['Tongyi-MAI/Z-Image-Turbo', undefined],
+      '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].train.unload_text_encoder': [false, false],
+      'config.process[0].sample.sampler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.noise_scheduler': ['flowmatch', 'flowmatch'],
+      'config.process[0].train.timestep_type': ['weighted', 'sigmoid'],
+      'config.process[0].model.qtype': ['qfloat8', 'qfloat8'],
+      'config.process[0].sample.guidance_scale': [3, 4],
+      'config.process[0].sample.sample_steps': [25, 25],
+    },
+    disableSections: ['network.conv'],
+    additionalSections: ['model.low_vram', 'model.layer_offloading'],
+  },
 ].sort((a, b) => {
   // Sort by label, case-insensitive
   return a.label.localeCompare(b.label, undefined, { sensitivity: 'base' });
@@ -481,7 +582,6 @@ export const groupedModelOptions: GroupedSelectOption[] = modelArchs.reduce((acc
 export const quantizationOptions: SelectOption[] = [
   { value: '', label: '- NONE -' },
   { value: 'qfloat8', label: 'float8 (default)' },
-  { value: 'uint8', label: '8 bit' },
   { value: 'uint7', label: '7 bit' },
   { value: 'uint6', label: '6 bit' },
   { value: 'uint5', label: '5 bit' },
@@ -491,3 +591,33 @@ export const quantizationOptions: SelectOption[] = [
 ];
 
 export const defaultQtype = 'qfloat8';
+
+interface JobTypeOption extends SelectOption {
+  disableSections?: DisableableSections[];
+  processSections?: string[];
+  onActivate?: (config: JobConfig) => JobConfig;
+  onDeactivate?: (config: JobConfig) => JobConfig;
+}
+
+export const jobTypeOptions: JobTypeOption[] = [
+  {
+    value: 'diffusion_trainer',
+    label: 'LoRA Trainer',
+    disableSections: ['slider'],
+  },
+  {
+    value: 'concept_slider',
+    label: 'Concept Slider',
+    disableSections: ['trigger_word', 'train.diff_output_preservation'],
+    onActivate: (config: JobConfig) => {
+      // add default slider config
+      config.config.process[0].slider = { ...defaultSliderConfig };
+      return config;
+    },
+    onDeactivate: (config: JobConfig) => {
+      // remove slider config
+      delete config.config.process[0].slider;
+      return config;
+    },
+  },
+];

ui/src/types.ts

@@ -39,6 +39,16 @@ export interface GpuInfo {
   fan: GpuFan;
 }
 
+export interface CpuInfo {
+  name: string;
+  cores: number;
+  temperature: number;
+  totalMemory: number;
+  freeMemory: number;
+  availableMemory: number;
+  currentLoad: number;
+}
+
 export interface GPUApiResponse {
   hasNvidiaSmi: boolean;
   gpus: GpuInfo[];
@@ -83,12 +93,15 @@ export interface DatasetConfig {
   cache_latents_to_disk?: boolean;
   resolution: number[];
   controls: string[];
-  control_path: string | null;
+  control_path?: string | null;
   num_frames: number;
   shrink_video_to_frames: boolean;
   do_i2v: boolean;
   flip_x: boolean;
   flip_y: boolean;
+  control_path_1?: string | null;
+  control_path_2?: string | null;
+  control_path_3?: string | null;
 }
 
 export interface EMAConfig {
@@ -117,11 +130,17 @@ export interface TrainConfig {
     weight_decay: number;
   };
   skip_first_sample: boolean;
+  force_first_sample: boolean;
   disable_sampling: boolean;
   diff_output_preservation: boolean;
   diff_output_preservation_multiplier: number;
   diff_output_preservation_class: string;
+  blank_prompt_preservation?: boolean;
+  blank_prompt_preservation_multiplier?: number;
   switch_boundary_every: number;
+  loss_type: 'mse' | 'mae' | 'wavelet' | 'stepped';
+  do_differential_guidance?: boolean;
+  differential_guidance_scale?: number;
 }
 
 export interface QuantizeKwargsConfig {
@@ -130,6 +149,7 @@ export interface QuantizeKwargsConfig {
 
 export interface ModelConfig {
   name_or_path: string;
+  extras_name_or_path?: string;
   quantize: boolean;
   quantize_te: boolean;
   qtype: string;
@@ -138,11 +158,15 @@ export interface ModelConfig {
   arch: string;
   low_vram: boolean;
   model_kwargs: { [key: string]: any };
+  layer_offloading?: boolean;
+  layer_offloading_transformer_percent?: number;
+  layer_offloading_text_encoder_percent?: number;
+  assistant_lora_path?: string;
 }
 
 export interface SampleItem {
   prompt: string;
-  width?: number
+  width?: number;
   height?: number;
   neg?: string;
   seed?: number;
@@ -152,6 +176,10 @@ export interface SampleItem {
   num_frames?: number;
   ctrl_img?: string | null;
   ctrl_idx?: number;
+  network_multiplier?: number;
+  ctrl_img_1?: string | null;
+  ctrl_img_2?: string | null;
+  ctrl_img_3?: string | null;
 }
 
 export interface SampleConfig {
@@ -170,14 +198,24 @@ export interface SampleConfig {
   fps: number;
 }
 
+export interface SliderConfig {
+  guidance_strength?: number;
+  anchor_strength?: number;
+  positive_prompt?: string;
+  negative_prompt?: string;
+  target_class?: string;
+  anchor_class?: string | null;
+}
+
 export interface ProcessConfig {
-  type: 'ui_trainer';
+  type: string;
   sqlite_db_path?: string;
   training_folder: string;
   performance_log_every: number;
   trigger_word: string | null;
   device: string;
   network?: NetworkConfig;
+  slider?: SliderConfig;
   save: SaveConfig;
   datasets: DatasetConfig[];
   train: TrainConfig;
@@ -210,7 +248,7 @@ export interface JobConfig {
 }
 
 export interface ConfigDoc {
-  title: string;
+  title: string | React.ReactNode;
   description: React.ReactNode;
 }
 
@@ -256,3 +294,5 @@ export interface JobUpsertPayload {
   info?: string;
   speed_string?: string;
 }
+
+export type JobStatus = 'queued' | 'running' | 'stopping' | 'stopped' | 'completed' | 'error';