Add ability to trigger a save from the ui which whill make the trainer save on the next step

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()

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])

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);
+}

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({
         <MenuButton className={'ml-1 sm:ml-2'}>
           <Cog className={iconSizeClass} />
         </MenuButton>
-        <MenuItems anchor="bottom" className="bg-gray-900 border border-gray-700 rounded shadow-lg w-48 px-2 py-2 mt-4">
+        <MenuItems anchor="bottom" className="bg-gray-900 border border-gray-700 rounded shadow-lg w-52 px-2 py-2 mt-4">
           {job.job_type === 'train' && (
             <MenuItem>
               <Link
                 href={`/jobs/new?cloneId=${job.id}`}
-                className="cursor-pointer px-4 py-1 hover:bg-gray-800 rounded block"
+                className="cursor-pointer px-4 py-1 hover:bg-gray-800 rounded flex items-center gap-2"
               >
+                <Copy className="w-4 h-4" />
                 Clone Job
               </Link>
             </MenuItem>
           )}
+          {canStop && (
+            <MenuItem>
+              <div
+                className="cursor-pointer px-4 py-1 hover:bg-gray-800 rounded flex items-center gap-2"
+                onClick={async () => {
+                  await saveJobNow(job.id);
+                  if (onRefresh) onRefresh();
+                }}
+              >
+                <Save className="w-4 h-4" />
+                Save Next Step
+              </div>
+            </MenuItem>
+          )}
           <MenuItem>
             <div
-              className="cursor-pointer px-4 py-1 hover:bg-gray-800 rounded"
+              className="cursor-pointer px-4 py-1 hover:bg-gray-800 rounded flex items-center gap-2"
               onClick={() => {
                 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({
                 });
               }}
             >
+              <OctagonX className="w-4 h-4" />
               Mark as Stopped
             </div>
           </MenuItem>

ui/src/utils/jobs.ts

@@ -50,6 +50,22 @@ export const deleteJob = (jobID: string) => {
   });
 };
 
+export const saveJobNow = (jobID: string) => {
+  return new Promise<void>((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<void>((resolve, reject) => {
     apiClient

version.py

@@ -1 +1 @@
-VERSION = "0.9.13"
+VERSION = "0.9.14"