diff --git a/ui/package-lock.json b/ui/package-lock.json index a86a7bac..2ff2e530 100644 --- a/ui/package-lock.json +++ b/ui/package-lock.json @@ -9,6 +9,7 @@ "version": "0.1.0", "dependencies": { "@headlessui/react": "^2.2.0", + "@monaco-editor/react": "^4.7.0", "@prisma/client": "^6.3.1", "axios": "^1.7.9", "classnames": "^2.5.1", @@ -21,7 +22,8 @@ "react-dropzone": "^14.3.5", "react-global-hooks": "^1.3.5", "react-icons": "^5.5.0", - "sqlite3": "^5.1.7" + "sqlite3": "^5.1.7", + "yaml": "^2.7.0" }, "devDependencies": { "@types/node": "^20", @@ -535,6 +537,27 @@ "@jridgewell/sourcemap-codec": "^1.4.14" } }, + "node_modules/@monaco-editor/loader": { + "version": "1.5.0", + "resolved": "https://registry.npmjs.org/@monaco-editor/loader/-/loader-1.5.0.tgz", + "integrity": "sha512-hKoGSM+7aAc7eRTRjpqAZucPmoNOC4UUbknb/VNoTkEIkCPhqV8LfbsgM1webRM7S/z21eHEx9Fkwx8Z/C/+Xw==", + "dependencies": { + "state-local": "^1.0.6" + } + }, + "node_modules/@monaco-editor/react": { + "version": "4.7.0", + "resolved": "https://registry.npmjs.org/@monaco-editor/react/-/react-4.7.0.tgz", + "integrity": "sha512-cyzXQCtO47ydzxpQtCGSQGOC8Gk3ZUeBXFAxD+CWXYFo5OqZyZUonFl0DwUlTyAfRHntBfw2p3w4s9R6oe1eCA==", + "dependencies": { + "@monaco-editor/loader": "^1.5.0" + }, + "peerDependencies": { + "monaco-editor": ">= 0.25.0 < 1", + "react": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0", + "react-dom": "^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0" + } + }, "node_modules/@next/env": { "version": "15.1.7", "resolved": "https://registry.npmjs.org/@next/env/-/env-15.1.7.tgz", @@ -2751,6 +2774,12 @@ "integrity": "sha512-gKLcREMhtuZRwRAfqP3RFW+TK4JqApVBtOIftVgjuABpAtpxhPGaDcfvbhNvD0B8iD1oUr/txX35NjcaY6Ns/A==", "license": "MIT" }, + "node_modules/monaco-editor": { + "version": "0.52.2", + "resolved": "https://registry.npmjs.org/monaco-editor/-/monaco-editor-0.52.2.tgz", + "integrity": "sha512-GEQWEZmfkOGLdd3XK8ryrfWz3AIP8YymVXiPHEdewrUq7mh0qrKrfHLNCXcbB6sTnMLnOZ3ztSiKcciFUkIJwQ==", + "peer": true + }, "node_modules/ms": { "version": "2.1.3", "resolved": "https://registry.npmjs.org/ms/-/ms-2.1.3.tgz", @@ -3947,6 +3976,11 @@ "node": ">=8" } }, + "node_modules/state-local": { + "version": "1.0.7", + "resolved": "https://registry.npmjs.org/state-local/-/state-local-1.0.7.tgz", + "integrity": "sha512-HTEHMNieakEnoe33shBYcZ7NX83ACUjCu8c40iOGEZsngj9zRnkqS9j1pqQPXwobB0ZcVTk27REb7COQ0UR59w==" + }, "node_modules/streamsearch": { "version": "1.1.0", "resolved": "https://registry.npmjs.org/streamsearch/-/streamsearch-1.1.0.tgz", @@ -4504,7 +4538,6 @@ "version": "2.7.0", "resolved": "https://registry.npmjs.org/yaml/-/yaml-2.7.0.tgz", "integrity": "sha512-+hSoy/QHluxmC9kCIJyL/uyFmLmc+e5CFR5Wa+bpIhIj85LVb9ZH2nVnqrHoSvKogwODv0ClqZkmiSSaIH5LTA==", - "dev": true, "bin": { "yaml": "bin.mjs" }, diff --git a/ui/package.json b/ui/package.json index 1e524bba..b08d2e58 100644 --- a/ui/package.json +++ b/ui/package.json @@ -12,6 +12,7 @@ }, "dependencies": { "@headlessui/react": "^2.2.0", + "@monaco-editor/react": "^4.7.0", "@prisma/client": "^6.3.1", "axios": "^1.7.9", "classnames": "^2.5.1", @@ -24,7 +25,8 @@ "react-dropzone": "^14.3.5", "react-global-hooks": "^1.3.5", "react-icons": "^5.5.0", - "sqlite3": "^5.1.7" + "sqlite3": "^5.1.7", + "yaml": "^2.7.0" }, "devDependencies": { "@types/node": "^20", diff --git a/ui/src/app/jobs/new/jobConfig.ts b/ui/src/app/jobs/new/jobConfig.ts index c03fd18c..0f2a4311 100644 --- a/ui/src/app/jobs/new/jobConfig.ts +++ b/ui/src/app/jobs/new/jobConfig.ts @@ -22,13 +22,13 @@ export const defaultJobConfig: JobConfig = { type: 'ui_trainer', training_folder: 'output', sqlite_db_path: './aitk_db.db', - device: 'cuda:0', + device: 'cuda', trigger_word: null, performance_log_every: 10, network: { type: 'lora', - linear: 16, - linear_alpha: 16, + linear: 32, + linear_alpha: 32, lokr_full_rank: true, lokr_factor: -1 }, diff --git a/ui/src/app/jobs/new/page.tsx b/ui/src/app/jobs/new/page.tsx index 871d63c2..65474cdc 100644 --- a/ui/src/app/jobs/new/page.tsx +++ b/ui/src/app/jobs/new/page.tsx @@ -17,6 +17,8 @@ import path from 'path'; import { TopBar, MainContent } from '@/components/layout'; import { Button } from '@headlessui/react'; import { FaChevronLeft } from 'react-icons/fa'; +import SimpleJob from './SimpleJob'; +import AdvancedJob from './AdvancedJob'; const isDev = process.env.NODE_ENV === 'development'; @@ -29,12 +31,11 @@ export default function TrainingForm() { const { gpuList, isGPUInfoLoaded } = useGPUInfo(); const { datasets, status: datasetFetchStatus } = useDatasetList(); const [datasetOptions, setDatasetOptions] = useState<{ value: string; label: string }[]>([]); + const [showAdvancedView, setShowAdvancedView] = useState(false); const [jobConfig, setJobConfig] = useNestedState(objectCopy(defaultJobConfig)); const [status, setStatus] = useState<'idle' | 'saving' | 'success' | 'error'>('idle'); - const isVideoModel = isVideoModelFromArch(jobConfig.config.process[0].model.arch); - useEffect(() => { if (!isSettingsLoaded) return; if (datasetFetchStatus !== 'success') return; @@ -130,6 +131,27 @@ export default function TrainingForm() {

{runId ? 'Edit Training Job' : 'New Training Job'}

+ {showAdvancedView && ( + <> +
+ setGpuIDs(value)} + options={gpuList.map((gpu: any) => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))} + /> +
+
+ + )} + +
+ +
- -
-
- - setJobConfig(value, 'config.name')} - placeholder="Enter training name" - disabled={runId !== null} - required - /> - setGpuIDs(value)} - options={gpuList.map(gpu => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))} - /> - { - if (value?.trim() === '') { - value = null; - } - setJobConfig(value, 'config.process[0].trigger_word'); - }} - placeholder="" - required - /> - - {/* Model Configuration Section */} - - { - // see if model changed - const currentModel = options.model.find( - model => model.name_or_path === jobConfig.config.process[0].model.name_or_path, - ); - if (!currentModel || currentModel.name_or_path === value) { - // model has not changed - return; - } - // revert defaults from previous model - for (const key in currentModel.defaults) { - setJobConfig(currentModel.defaults[key][1], key); - } - // set new model - setJobConfig(value, 'config.process[0].model.name_or_path'); - // update the defaults when a model is selected - const model = options.model.find(model => model.name_or_path === value); - if (model?.defaults) { - for (const key in model.defaults) { - setJobConfig(model.defaults[key][0], key); - } - } - }} - options={ - options.model - .map(model => { - if (model.dev_only && !isDev) { - return null; - } - return { - value: model.name_or_path, - label: model.name_or_path, - }; - }) - .filter(x => x) as { value: string; label: string }[] - } - /> - { - const currentArch = modelArchs.find( - a => a.name === jobConfig.config.process[0].model.arch, - ); - if (!currentArch || currentArch.name === value) { - return; - } - // set new model - setJobConfig(value, 'config.process[0].model.arch'); - }} - options={ - modelArchs - .map(model => { - return { - value: model.name, - label: model.label, - }; - }) - .filter(x => x) as { value: string; label: string }[] - } - /> - -
- setJobConfig(value, 'config.process[0].model.quantize')} - /> - setJobConfig(value, 'config.process[0].model.quantize_te')} - /> -
-
-
- - setJobConfig(value, 'config.process[0].network.type')} - options={[ - { value: 'lora', label: 'LoRA' }, - { value: 'lokr', label: 'LoKr' }, - ]} - /> - {jobConfig.config.process[0].network?.type == 'lokr' && ( - setJobConfig(parseInt(value), 'config.process[0].network.lokr_factor')} - options={[ - { value: '-1', label: 'Auto' }, - { value: '4', label: '4' }, - { value: '8', label: '8' }, - { value: '16', label: '16' }, - { value: '32', label: '32' }, - ]} - /> - )} - {jobConfig.config.process[0].network?.type == 'lora' && ( - { - console.log('onChange', value); - setJobConfig(value, 'config.process[0].network.linear'); - setJobConfig(value, 'config.process[0].network.linear_alpha'); - }} - placeholder="eg. 16" - min={0} - max={1024} - required - /> - )} - - - setJobConfig(value, 'config.process[0].save.dtype')} - options={[ - { value: 'bf16', label: 'BF16' }, - { value: 'fp16', label: 'FP16' }, - { value: 'fp32', label: 'FP32' }, - ]} - /> - setJobConfig(value, 'config.process[0].save.save_every')} - placeholder="eg. 250" - min={1} - required - /> - setJobConfig(value, 'config.process[0].save.max_step_saves_to_keep')} - placeholder="eg. 4" - min={1} - required - /> - -
-
- -
-
- setJobConfig(value, 'config.process[0].train.batch_size')} - placeholder="eg. 4" - min={1} - required - /> - setJobConfig(value, 'config.process[0].train.gradient_accumulation')} - placeholder="eg. 1" - min={1} - required - /> - setJobConfig(value, 'config.process[0].train.steps')} - placeholder="eg. 2000" - min={1} - required - /> -
-
- setJobConfig(value, 'config.process[0].train.optimizer')} - options={[ - { value: 'adamw8bit', label: 'AdamW8Bit' }, - { value: 'adafactor', label: 'Adafactor' }, - ]} - /> - setJobConfig(value, 'config.process[0].train.lr')} - placeholder="eg. 0.0001" - min={0} - required - /> - setJobConfig(value, 'config.process[0].train.optimizer_params.weight_decay')} - placeholder="eg. 0.0001" - min={0} - required - /> -
-
- setJobConfig(value, 'config.process[0].train.timestep_type')} - options={[ - { value: 'sigmoid', label: 'Sigmoid' }, - { value: 'linear', label: 'Linear' }, - { value: 'flux_shift', label: 'Flux Shift' }, - ]} - /> - setJobConfig(value, 'config.process[0].train.content_or_style')} - options={[ - { value: 'balanced', label: 'Balanced' }, - { value: 'content', label: 'High Noise' }, - { value: 'style', label: 'Low Noise' }, - ]} - /> - setJobConfig(value, 'config.process[0].train.noise_scheduler')} - options={[ - { value: 'flowmatch', label: 'FlowMatch' }, - { value: 'ddpm', label: 'DDPM' }, - ]} - /> -
-
- - setJobConfig(value, 'config.process[0].train.ema_config.use_ema')} - /> - - setJobConfig(value, 'config.process[0].train.ema_config?.ema_decay')} - placeholder="eg. 0.99" - min={0} - /> - -
- setJobConfig(value, 'config.process[0].train.unload_text_encoder')} - /> -
-
-
-
- - setJobConfig(value, 'config.process[0].train.diff_output_preservation')} - /> - - - setJobConfig(value, 'config.process[0].train.diff_output_preservation_multiplier') - } - placeholder="eg. 1.0" - min={0} - /> - setJobConfig(value, 'config.process[0].train.diff_output_preservation_class')} - placeholder="eg. woman" - /> -
-
-
-
-
- - <> - {jobConfig.config.process[0].datasets.map((dataset, i) => ( -
- -

Dataset {i + 1}

-
-
- setJobConfig(value, `config.process[0].datasets[${i}].folder_path`)} - options={datasetOptions} - /> - setJobConfig(value, `config.process[0].datasets[${i}].network_weight`)} - placeholder="eg. 1.0" - /> -
-
- setJobConfig(value, `config.process[0].datasets[${i}].default_caption`)} - placeholder="eg. A photo of a cat" - /> - - setJobConfig(value, `config.process[0].datasets[${i}].caption_dropout_rate`) - } - placeholder="eg. 0.05" - min={0} - required - /> -
-
- - - setJobConfig(value, `config.process[0].datasets[${i}].cache_latents_to_disk`) - } - /> - setJobConfig(value, `config.process[0].datasets[${i}].is_reg`)} - /> - -
-
- -
- {[ - [256, 512, 768], - [1024, 1280, 1536], - ].map(resGroup => ( -
- {resGroup.map(res => ( - { - const resolutions = dataset.resolution.includes(res) - ? dataset.resolution.filter(r => r !== res) - : [...dataset.resolution, res]; - setJobConfig(resolutions, `config.process[0].datasets[${i}].resolution`); - }} - /> - ))} -
- ))} -
-
-
-
-
- ))} - - -
-
-
- -
-
- setJobConfig(value, 'config.process[0].sample.sample_every')} - placeholder="eg. 250" - min={1} - required - /> - setJobConfig(value, 'config.process[0].sample.sampler')} - options={[ - { value: 'flowmatch', label: 'FlowMatch' }, - { value: 'ddpm', label: 'DDPM' }, - ]} - /> -
-
- setJobConfig(value, 'config.process[0].sample.guidance_scale')} - placeholder="eg. 1.0" - min={0} - required - /> - setJobConfig(value, 'config.process[0].sample.sample_steps')} - placeholder="eg. 1" - className="pt-2" - min={1} - required - /> -
-
- setJobConfig(value, 'config.process[0].sample.width')} - placeholder="eg. 1024" - min={0} - required - /> - setJobConfig(value, 'config.process[0].sample.height')} - placeholder="eg. 1024" - className="pt-2" - min={0} - required - /> -
+ {showAdvancedView ? ( +
+ +
+ ) : ( + + -
- setJobConfig(value, 'config.process[0].sample.seed')} - placeholder="eg. 0" - min={0} - required - /> - setJobConfig(value, 'config.process[0].sample.walk_seed')} - /> -
- { isVideoModel && ( -
- setJobConfig(value, 'config.process[0].sample.num_frames')} - placeholder="eg. 0" - min={0} - required - /> - setJobConfig(value, 'config.process[0].sample.fps')} - placeholder="eg. 0" - min={0} - required - /> -
- )} -
- - {jobConfig.config.process[0].sample.prompts.map((prompt, i) => ( -
-
- setJobConfig(value, `config.process[0].sample.prompts[${i}]`)} - placeholder="Enter prompt" - required - /> -
-
- -
-
- ))} - -
-
-
- - {status === 'success' &&

Training saved successfully!

} - {status === 'error' &&

Error saving training. Please try again.

} -
-
-
+
+ + )} ); } diff --git a/ui/src/hooks/useSettings.tsx b/ui/src/hooks/useSettings.tsx index 15b5175e..7d17bbd5 100644 --- a/ui/src/hooks/useSettings.tsx +++ b/ui/src/hooks/useSettings.tsx @@ -2,6 +2,12 @@ import { useEffect, useState } from 'react'; +export interface Settings { + HF_TOKEN: string; + TRAINING_FOLDER: string; + DATASETS_FOLDER: string; +} + export default function useSettings() { const [settings, setSettings] = useState({ HF_TOKEN: '',