diff --git a/.gitignore b/.gitignore index 0f02393b..04c233ac 100644 --- a/.gitignore +++ b/.gitignore @@ -180,4 +180,5 @@ cython_debug/ .DS_Store ._.DS_Store aitk_db.db -/notes.md \ No newline at end of file +/notes.md +/data \ No newline at end of file diff --git a/README.md b/README.md index db05500f..a57aff60 100644 --- a/README.md +++ b/README.md @@ -425,6 +425,9 @@ Everything else should work the same including layer targeting. Only larger updates are listed here. There are usually smaller daily updated that are omitted. +### Jul 17, 2025 +- Make it easy to add control images to the samples in the ui + ### Jul 11, 2025 - Added better video config settings to the UI for video models. - Added Wan I2V training to the UI diff --git a/ui/package-lock.json b/ui/package-lock.json index 6bebedce..f4ed097b 100644 --- a/ui/package-lock.json +++ b/ui/package-lock.json @@ -24,6 +24,7 @@ "react-icons": "^5.5.0", "react-select": "^5.10.1", "sqlite3": "^5.1.7", + "uuid": "^11.1.0", "yaml": "^2.7.0" }, "devDependencies": { @@ -5370,6 +5371,18 @@ "resolved": "https://registry.npmjs.org/util-deprecate/-/util-deprecate-1.0.2.tgz", "integrity": "sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==" }, + "node_modules/uuid": { + "version": "11.1.0", + "resolved": "https://registry.npmjs.org/uuid/-/uuid-11.1.0.tgz", + "integrity": "sha512-0/A9rDy9P7cJ+8w1c9WD9V//9Wj15Ce2MPz8Ri6032usz+NfePxx5AcN3bN+r6ZL6jEo066/yNYB3tn4pQEx+A==", + "funding": [ + "https://github.com/sponsors/broofa", + "https://github.com/sponsors/ctavan" + ], + "bin": { + "uuid": "dist/esm/bin/uuid" + } + }, "node_modules/v8-compile-cache-lib": { "version": "3.0.1", "resolved": "https://registry.npmjs.org/v8-compile-cache-lib/-/v8-compile-cache-lib-3.0.1.tgz", diff --git a/ui/package.json b/ui/package.json index de04bd2b..594aac6d 100644 --- a/ui/package.json +++ b/ui/package.json @@ -28,6 +28,7 @@ "react-icons": "^5.5.0", "react-select": "^5.10.1", "sqlite3": "^5.1.7", + "uuid": "^11.1.0", "yaml": "^2.7.0" }, "devDependencies": { diff --git a/ui/src/app/api/img/[...imagePath]/route.ts b/ui/src/app/api/img/[...imagePath]/route.ts index 2a67586f..80fc7272 100644 --- a/ui/src/app/api/img/[...imagePath]/route.ts +++ b/ui/src/app/api/img/[...imagePath]/route.ts @@ -2,7 +2,7 @@ import { NextRequest, NextResponse } from 'next/server'; import fs from 'fs'; import path from 'path'; -import { getDatasetsRoot, getTrainingFolder } from '@/server/settings'; +import { getDatasetsRoot, getTrainingFolder, getDataRoot } from '@/server/settings'; export async function GET(request: NextRequest, { params }: { params: { imagePath: string } }) { const { imagePath } = await params; @@ -13,8 +13,9 @@ export async function GET(request: NextRequest, { params }: { params: { imagePat // Get allowed directories const datasetRoot = await getDatasetsRoot(); const trainingRoot = await getTrainingFolder(); + const dataRoot = await getDataRoot(); - const allowedDirs = [datasetRoot, trainingRoot]; + const allowedDirs = [datasetRoot, trainingRoot, dataRoot]; // Security check: Ensure path is in allowed directory const isAllowed = allowedDirs.some(allowedDir => filepath.startsWith(allowedDir)) && !filepath.includes('..'); diff --git a/ui/src/app/api/img/upload/route.ts b/ui/src/app/api/img/upload/route.ts new file mode 100644 index 00000000..56615bd0 --- /dev/null +++ b/ui/src/app/api/img/upload/route.ts @@ -0,0 +1,58 @@ +// src/app/api/datasets/upload/route.ts +import { NextRequest, NextResponse } from 'next/server'; +import { writeFile, mkdir } from 'fs/promises'; +import { join } from 'path'; +import { getDataRoot } from '@/server/settings'; +import {v4 as uuidv4} from 'uuid'; + +export async function POST(request: NextRequest) { + try { + const dataRoot = await getDataRoot(); + if (!dataRoot) { + return NextResponse.json({ error: 'Data root path not found' }, { status: 500 }); + } + const imgRoot = join(dataRoot, 'images'); + + + const formData = await request.formData(); + const files = formData.getAll('files'); + + if (!files || files.length === 0) { + return NextResponse.json({ error: 'No files provided' }, { status: 400 }); + } + + // make it recursive if it doesn't exist + await mkdir(imgRoot, { recursive: true }); + const savedFiles = await Promise.all( + files.map(async (file: any) => { + const bytes = await file.arrayBuffer(); + const buffer = Buffer.from(bytes); + + const extension = file.name.split('.').pop() || 'jpg'; + + // Clean filename and ensure it's unique + const fileName = `${uuidv4()}`; // Use UUID for unique file names + const filePath = join(imgRoot, `${fileName}.${extension}`); + + await writeFile(filePath, buffer); + return filePath; + }), + ); + + return NextResponse.json({ + message: 'Files uploaded successfully', + files: savedFiles, + }); + } catch (error) { + console.error('Upload error:', error); + return NextResponse.json({ error: 'Error uploading files' }, { status: 500 }); + } +} + +// Increase payload size limit (default is 4mb) +export const config = { + api: { + bodyParser: false, + responseLimit: '50mb', + }, +}; diff --git a/ui/src/app/jobs/new/AdvancedJob.tsx b/ui/src/app/jobs/new/AdvancedJob.tsx index 42365158..6a0d4388 100644 --- a/ui/src/app/jobs/new/AdvancedJob.tsx +++ b/ui/src/app/jobs/new/AdvancedJob.tsx @@ -5,6 +5,7 @@ import YAML from 'yaml'; import Editor, { OnMount } from '@monaco-editor/react'; import type { editor } from 'monaco-editor'; import { Settings } from '@/hooks/useSettings'; +import { migrateJobConfig } from './jobConfig'; type Props = { jobConfig: JobConfig; @@ -115,6 +116,7 @@ export default function AdvancedJob({ jobConfig, setJobConfig, settings }: Props } catch (e) { console.warn(e); } + migrateJobConfig(parsed); setJobConfig(parsed); } } catch (e) { diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index 157bcc6c..52f6454a 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/ui/src/app/jobs/new/SimpleJob.tsx @@ -7,6 +7,7 @@ import { objectCopy } from '@/utils/basic'; import { TextInput, SelectInput, Checkbox, FormGroup, NumberInput } from '@/components/formInputs'; import Card from '@/components/Card'; import { X } from 'lucide-react'; +import AddSingleImageModal, { openAddImageModal } from '@/components/AddSingleImageModal'; type Props = { jobConfig: JobConfig; @@ -116,6 +117,17 @@ export default function SimpleJob({ return newDataset; }); setJobConfig(datasets, 'config.process[0].datasets'); + + // update samples + const hasSampleCtrlImg = newArch?.additionalSections?.includes('sample.ctrl_img') || false; + const samples = jobConfig.config.process[0].sample.samples.map(sample => { + const newSample = objectCopy(sample); + if (!hasSampleCtrlImg) { + delete newSample.ctrl_img; // remove ctrl_img if not applicable + } + return newSample; + }); + setJobConfig(samples, 'config.process[0].sample.samples'); }} options={groupedModelOptions} /> @@ -648,32 +660,58 @@ export default function SimpleJob({ - - {modelArch?.additionalSections?.includes('sample.ctrl_img') && ( -
-

Control Images

- To use control images on samples, add --ctrl_img to the prompts below. -
- Example: make this a cartoon --ctrl_img /path/to/image.png -
- )} - {jobConfig.config.process[0].sample.prompts.map((prompt, i) => ( -
+ +
+ + {jobConfig.config.process[0].sample.samples.map((sample, i) => ( +
+
- setJobConfig(value, `config.process[0].sample.prompts[${i}]`)} - placeholder="Enter prompt" - required - /> +
+
+ setJobConfig(value, `config.process[0].sample.samples[${i}].prompt`)} + placeholder="Enter prompt" + required + /> +
+ + {modelArch?.additionalSections?.includes('sample.ctrl_img') && ( +
{ + openAddImageModal(imagePath => { + console.log('Selected image path:', imagePath); + if (!imagePath) return; + setJobConfig(imagePath, `config.process[0].sample.samples[${i}].ctrl_img`); + }); + }} + > + {!sample.ctrl_img && ( +
Add Control Image
+ )} +
+ )} +
+
- ))} - - +
+ ))} +
{status === 'success' &&

Training saved successfully!

} {status === 'error' &&

Error saving training. Please try again.

} + ); } diff --git a/ui/src/app/jobs/new/jobConfig.ts b/ui/src/app/jobs/new/jobConfig.ts index 94c0de6a..a50358a5 100644 --- a/ui/src/app/jobs/new/jobConfig.ts +++ b/ui/src/app/jobs/new/jobConfig.ts @@ -90,17 +90,37 @@ export const defaultJobConfig: JobConfig = { sample_every: 250, width: 1024, height: 1024, - prompts: [ - 'woman with red hair, playing chess at the park, bomb going off in the background', - 'a woman holding a coffee cup, in a beanie, sitting at a cafe', - 'a horse is a DJ at a night club, fish eye lens, smoke machine, lazer lights, holding a martini', - 'a man showing off his cool new t shirt at the beach, a shark is jumping out of the water in the background', - 'a bear building a log cabin in the snow covered mountains', - 'woman playing the guitar, on stage, singing a song, laser lights, punk rocker', - 'hipster man with a beard, building a chair, in a wood shop', - 'photo of a man, white background, medium shot, modeling clothing, studio lighting, white backdrop', - "a man holding a sign that says, 'this is a sign'", - 'a bulldog, in a post apocalyptic world, with a shotgun, in a leather jacket, in a desert, with a motorcycle', + samples: [ + { + 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', + }, + { + 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, @@ -118,3 +138,23 @@ export const defaultJobConfig: JobConfig = { version: '1.0', }, }; + +export const migrateJobConfig = (jobConfig: JobConfig): JobConfig => { + // upgrade prompt strings to samples + if ( + jobConfig?.config?.process && + jobConfig.config.process[0]?.sample && + Array.isArray(jobConfig.config.process[0].sample.prompts) && + jobConfig.config.process[0].sample.prompts.length > 0 + ) { + let newSamples = []; + for (const prompt of jobConfig.config.process[0].sample.prompts) { + newSamples.push({ + prompt: prompt, + }); + } + jobConfig.config.process[0].sample.samples = newSamples; + delete jobConfig.config.process[0].sample.prompts; + } + return jobConfig; +}; diff --git a/ui/src/app/jobs/new/page.tsx b/ui/src/app/jobs/new/page.tsx index 5c8d027b..fb2b8546 100644 --- a/ui/src/app/jobs/new/page.tsx +++ b/ui/src/app/jobs/new/page.tsx @@ -2,11 +2,11 @@ import { useEffect, useState } from 'react'; import { useSearchParams, useRouter } from 'next/navigation'; -import { defaultJobConfig, defaultDatasetConfig } from './jobConfig'; +import { defaultJobConfig, defaultDatasetConfig, migrateJobConfig } from './jobConfig'; import { JobConfig } from '@/types'; import { objectCopy } from '@/utils/basic'; import { useNestedState } from '@/utils/hooks'; -import { SelectInput} from '@/components/formInputs'; +import { SelectInput } from '@/components/formInputs'; import useSettings from '@/hooks/useSettings'; import useGPUInfo from '@/hooks/useGPUInfo'; import useDatasetList from '@/hooks/useDatasetList'; @@ -61,7 +61,7 @@ export default function TrainingForm() { .then(data => { console.log('Training:', data); setGpuIDs(data.gpu_ids); - setJobConfig(JSON.parse(data.job_config)); + setJobConfig(migrateJobConfig(JSON.parse(data.job_config))); }) .catch(error => console.error('Error fetching training:', error)); } @@ -181,11 +181,13 @@ export default function TrainingForm() { ) : ( - - Advanced job detected. Please switch to advanced view to continue. - - }> + + Advanced job detected. Please switch to advanced view to continue. + + } + > ); -} \ No newline at end of file +} diff --git a/ui/src/components/AddSingleImageModal.tsx b/ui/src/components/AddSingleImageModal.tsx new file mode 100644 index 00000000..ba32ef9d --- /dev/null +++ b/ui/src/components/AddSingleImageModal.tsx @@ -0,0 +1,141 @@ +'use client'; +import { createGlobalState } from 'react-global-hooks'; +import { Dialog, DialogBackdrop, DialogPanel, DialogTitle } from '@headlessui/react'; +import { FaUpload } from 'react-icons/fa'; +import { useCallback, useState } from 'react'; +import { useDropzone } from 'react-dropzone'; +import { apiClient } from '@/utils/api'; + +export interface AddSingleImageModalState { + + onComplete?: (imagePath: string|null) => void; +} + +export const addSingleImageModalState = createGlobalState(null); + +export const openAddImageModal = (onComplete: (imagePath: string|null) => void) => { + addSingleImageModalState.set({onComplete }); +}; + +export default function AddSingleImageModal() { + const [addSingleImageModalInfo, setAddSingleImageModalInfo] = addSingleImageModalState.use(); + const [uploadProgress, setUploadProgress] = useState(0); + const [isUploading, setIsUploading] = useState(false); + const open = addSingleImageModalInfo !== null; + + const onCancel = () => { + if (!isUploading) { + setAddSingleImageModalInfo(null); + } + }; + + const onDone = (imagePath: string|null) => { + if (addSingleImageModalInfo?.onComplete && !isUploading) { + addSingleImageModalInfo.onComplete(imagePath); + setAddSingleImageModalInfo(null); + } + }; + + const onDrop = useCallback( + async (acceptedFiles: File[]) => { + if (acceptedFiles.length === 0) return; + + setIsUploading(true); + setUploadProgress(0); + + const formData = new FormData(); + acceptedFiles.forEach(file => { + formData.append('files', file); + }); + + try { + const resp = await apiClient.post(`/api/img/upload`, formData, { + headers: { + 'Content-Type': 'multipart/form-data', + }, + onUploadProgress: progressEvent => { + const percentCompleted = Math.round((progressEvent.loaded * 100) / (progressEvent.total || 100)); + setUploadProgress(percentCompleted); + }, + timeout: 0, // Disable timeout + }); + console.log('Upload successful:', resp.data); + + onDone(resp.data.files[0] || null); + } catch (error) { + console.error('Upload failed:', error); + } finally { + setIsUploading(false); + setUploadProgress(0); + } + }, + [addSingleImageModalInfo], + ); + + const { getRootProps, getInputProps, isDragActive } = useDropzone({ + onDrop, + accept: { + 'image/*': ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp'], + }, + multiple: false, + }); + + return ( + + + +
+
+ +
+
+ + Add Control Image + +
+
+ + +

+ {isDragActive ? 'Drop the image here...' : 'Drag & drop an image here, or click to select one'} +

+
+ {isUploading && ( +
+
+
+
+

Uploading... {uploadProgress}%

+
+ )} +
+
+
+
+ +
+ +
+
+ + ); +} diff --git a/ui/src/components/SampleImages.tsx b/ui/src/components/SampleImages.tsx index c0653001..6f134ad4 100644 --- a/ui/src/components/SampleImages.tsx +++ b/ui/src/components/SampleImages.tsx @@ -14,7 +14,11 @@ export default function SampleImages({ job }: SampleImagesProps) { if (job?.job_config) { const jobConfig = JSON.parse(job.job_config) as JobConfig; const sampleConfig = jobConfig.config.process[0].sample; - return sampleConfig.prompts.length; + if (sampleConfig.prompts) { + return sampleConfig.prompts.length; + } else { + return sampleConfig.samples.length; + } } return 10; }, [job]); diff --git a/ui/src/paths.ts b/ui/src/paths.ts index 9b11ea95..92311f43 100644 --- a/ui/src/paths.ts +++ b/ui/src/paths.ts @@ -2,3 +2,4 @@ import path from 'path'; export const TOOLKIT_ROOT = path.resolve('@', '..', '..'); export const defaultTrainFolder = path.join(TOOLKIT_ROOT, 'output'); export const defaultDatasetsFolder = path.join(TOOLKIT_ROOT, 'datasets'); +export const defaultDataRoot = path.join(TOOLKIT_ROOT, 'data'); diff --git a/ui/src/server/settings.ts b/ui/src/server/settings.ts index efebc8cc..9c6abe68 100644 --- a/ui/src/server/settings.ts +++ b/ui/src/server/settings.ts @@ -1,5 +1,5 @@ import { PrismaClient } from '@prisma/client'; -import { defaultDatasetsFolder } from '@/paths'; +import { defaultDatasetsFolder, defaultDataRoot } from '@/paths'; import { defaultTrainFolder } from '@/paths'; import NodeCache from 'node-cache'; @@ -66,3 +66,22 @@ export const getHFToken = async () => { myCache.set(key, token); return token; }; + +export const getDataRoot = async () => { + const key = 'DATA_ROOT'; + let dataRoot = myCache.get(key) as string; + if (dataRoot) { + return dataRoot; + } + let row = await prisma.settings.findFirst({ + where: { + key: key, + }, + }); + dataRoot = defaultDataRoot; + if (row?.value && row.value !== '') { + dataRoot = row.value; + } + myCache.set(key, dataRoot); + return dataRoot; +}; diff --git a/ui/src/types.ts b/ui/src/types.ts index 02e89b5c..19d64cf1 100644 --- a/ui/src/types.ts +++ b/ui/src/types.ts @@ -133,12 +133,27 @@ export interface ModelConfig { model_kwargs: { [key: string]: any }; } +export interface SampleItem { + prompt: string; + width?: number + height?: number; + neg?: string; + seed?: number; + guidance_scale?: number; + sample_steps?: number; + fps?: number; + num_frames?: number; + ctrl_img?: string | null; + ctrl_idx?: number; +} + export interface SampleConfig { sampler: string; sample_every: number; width: number; height: number; - prompts: string[]; + prompts?: string[]; + samples: SampleItem[]; neg: string; seed: number; walk_seed: boolean; diff --git a/version.py b/version.py index aeeea4a5..9ca78d05 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.3.7" \ No newline at end of file +VERSION = "0.3.8" \ No newline at end of file