From 68e9b38220ccbb8def77df39fdc3f5824c05f120 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Tue, 26 May 2026 08:38:49 -0600 Subject: [PATCH] Add ability to trigger a save from the ui which whill make the trainer save on the next step --- .../sd_trainer/DiffusionTrainer.py | 33 +++++++++++++++++++ ui/prisma/schema.prisma | 1 + ui/src/app/api/jobs/[jobID]/save_now/route.ts | 19 +++++++++++ ui/src/components/JobActionBar.tsx | 26 ++++++++++++--- ui/src/utils/jobs.ts | 16 +++++++++ version.py | 2 +- 6 files changed, 91 insertions(+), 6 deletions(-) create mode 100644 ui/src/app/api/jobs/[jobID]/save_now/route.ts diff --git a/extensions_built_in/sd_trainer/DiffusionTrainer.py b/extensions_built_in/sd_trainer/DiffusionTrainer.py index 252da00..51c081b 100644 --- a/extensions_built_in/sd_trainer/DiffusionTrainer.py +++ b/extensions_built_in/sd_trainer/DiffusionTrainer.py @@ -8,6 +8,8 @@ from typing import Literal, Optional import threading import time import signal +from toolkit.basic import flush +from toolkit.print import print_acc AITK_Status = Literal["running", "stopped", "error", "completed"] @@ -172,6 +174,36 @@ class DiffusionTrainer(SDTrainer): self.is_stopping = True raise Exception("Job returning to queue") + def should_save(self): + if not self.is_ui_trainer: + return False + def _check_save(): + with self._db_connect() as conn: + cursor = conn.cursor() + cursor.execute( + "SELECT save_now FROM Job WHERE id = ?", (self.job_id,)) + save_now = cursor.fetchone() + return False if save_now is None else save_now[0] == 1 + + return self._retry_db_operation(_check_save) + + def maybe_save(self): + if not self.is_ui_trainer: + return + if self.should_save(): + self.update_db_key("save_now", 0) + if self.progress_bar is not None: + self.progress_bar.pause() + print_acc(f"\nSaving at step {self.step_num}") + # clear any grads + self.optimizer.zero_grad() + self.save(self.step_num) + self.ensure_params_requires_grad() + flush() + if self.progress_bar is not None: + self.progress_bar.unpause() + self.save(self.step_num) + async def _update_key(self, key, value): if not self.accelerator.is_main_process: return @@ -287,6 +319,7 @@ class DiffusionTrainer(SDTrainer): if self.is_ui_trainer: self.update_step() self.maybe_stop() + self.maybe_save() def hook_before_model_load(self): super().hook_before_model_load() diff --git a/ui/prisma/schema.prisma b/ui/prisma/schema.prisma index f183e58..bb10e9a 100644 --- a/ui/prisma/schema.prisma +++ b/ui/prisma/schema.prisma @@ -38,6 +38,7 @@ model Job { pid Int? job_type String @default("train") // 'train', 'caption' job_ref String? // can be used for anything for special jobs, like dataset path for caption jobs + save_now Boolean @default(false) // if true, the job will be saved on the next step @@index([status]) @@index([gpu_ids]) diff --git a/ui/src/app/api/jobs/[jobID]/save_now/route.ts b/ui/src/app/api/jobs/[jobID]/save_now/route.ts new file mode 100644 index 0000000..0a6e54b --- /dev/null +++ b/ui/src/app/api/jobs/[jobID]/save_now/route.ts @@ -0,0 +1,19 @@ +import { NextRequest, NextResponse } from 'next/server'; +import { PrismaClient } from '@prisma/client'; + +const prisma = new PrismaClient(); + +export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) { + const { jobID } = await params; + + const job = await prisma.job.update({ + where: { id: jobID }, + data: { + save_now: true, + }, + }); + + console.log(`Job ${jobID} marked to save on next step`); + + return NextResponse.json(job); +} diff --git a/ui/src/components/JobActionBar.tsx b/ui/src/components/JobActionBar.tsx index 4520a04..8971e0b 100644 --- a/ui/src/components/JobActionBar.tsx +++ b/ui/src/components/JobActionBar.tsx @@ -1,9 +1,9 @@ import Link from 'next/link'; -import { Eye, Trash2, Pen, Play, Pause, Cog, X } from 'lucide-react'; +import { Eye, Trash2, Pen, Play, Pause, Cog, X, Copy, Save, OctagonX } from 'lucide-react'; import { Button } from '@headlessui/react'; import { openConfirm } from '@/components/ConfirmModal'; import { Job } from '@prisma/client'; -import { startJob, stopJob, deleteJob, getAvaliableJobActions, markJobAsStopped } from '@/utils/jobs'; +import { startJob, stopJob, deleteJob, getAvaliableJobActions, markJobAsStopped, saveJobNow } from '@/utils/jobs'; import { startQueue } from '@/utils/queue'; import { Menu, MenuButton, MenuItem, MenuItems } from '@headlessui/react'; import { redirect } from 'next/navigation'; @@ -140,20 +140,35 @@ export default function JobActionBar({ - + {job.job_type === 'train' && ( + Clone Job )} + {canStop && ( + +
{ + await saveJobNow(job.id); + if (onRefresh) onRefresh(); + }} + > + + Save Next Step +
+ + )}
{ let message = `Are you sure you want to mark this job as stopped? This will set the job status to 'stopped' if the status is hung. Only do this if you are 100% sure the job is stopped. This will NOT stop the job.`; openConfirm({ @@ -168,6 +183,7 @@ export default function JobActionBar({ }); }} > + Mark as Stopped
diff --git a/ui/src/utils/jobs.ts b/ui/src/utils/jobs.ts index 43d0bdc..a5db97b 100644 --- a/ui/src/utils/jobs.ts +++ b/ui/src/utils/jobs.ts @@ -50,6 +50,22 @@ export const deleteJob = (jobID: string) => { }); }; +export const saveJobNow = (jobID: string) => { + return new Promise((resolve, reject) => { + apiClient + .get(`/api/jobs/${jobID}/save_now`) + .then(res => res.data) + .then(data => { + console.log('Job set to save on next step:', data); + resolve(); + }) + .catch(error => { + console.error('Error setting job to save on next step:', error); + reject(error); + }); + }); +}; + export const markJobAsStopped = (jobID: string) => { return new Promise((resolve, reject) => { apiClient diff --git a/version.py b/version.py index 611ec14..00eb0b2 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.9.13" +VERSION = "0.9.14"