Auto-caption QoL change, build errors cleared.

.gitignore

@@ -184,4 +184,6 @@ cython_debug/
 aitk_db.db
 /notes.md
 /data
-.claude
+.claude
+launch.bat
+instructions.md

extensions_built_in/captioner/BaseCaptioner.py

@@ -9,6 +9,7 @@
 import signal
 import concurrent.futures
 from PIL import Image
+from jobs.exceptions import JobReturnedToQueueException, JobStoppedException
 
 import torch
 from jobs.process import BaseExtensionProcess
@@ -314,11 +315,11 @@ def maybe_stop(self):
         if self.should_stop():
             self._run_async_operation(self._update_status("stopped", "Job stopped"))
             self.is_stopping = True
-            raise Exception("Job stopped")
+            raise JobStoppedException("Job stopped")
         if self.should_return_to_queue():
             self._run_async_operation(self._update_status("queued", "Job queued"))
             self.is_stopping = True
-            raise Exception("Job returning to queue")
+            raise JobReturnedToQueueException("Job returning to queue")
 
     async def _update_key(self, key, value):
         def _do_update():

extensions_built_in/sd_trainer/DiffusionTrainer.py

@@ -8,6 +8,7 @@
 import threading
 import time
 import signal
+from jobs.exceptions import JobReturnedToQueueException, JobStoppedException
 from toolkit.basic import flush
 from toolkit.print import print_acc
 
@@ -167,12 +168,12 @@ def maybe_stop(self):
             self._run_async_operation(
                 self._update_status("stopped", "Job stopped"))
             self.is_stopping = True
-            raise Exception("Job stopped")
+            raise JobStoppedException("Job stopped")
         if self.should_return_to_queue():
             self._run_async_operation(
                 self._update_status("queued", "Job queued"))
             self.is_stopping = True
-            raise Exception("Job returning to queue")
+            raise JobReturnedToQueueException("Job returning to queue")
 
     def should_save(self):
         if not self.is_ui_trainer:
@@ -238,6 +239,53 @@ def update_db_key(self, key, value):
         if self.accelerator.is_main_process and self.is_ui_trainer:
             self._run_async_operation(self._update_key(key, value))
 
+    def _queue_job_to_bottom(self, info: Optional[str] = None):
+        if not self.is_ui_trainer:
+            return
+
+        def _queue_job():
+            with self._db_connect() as conn:
+                cursor = conn.cursor()
+                cursor.execute("BEGIN IMMEDIATE")
+                try:
+                    cursor.execute("SELECT COALESCE(MAX(queue_position), 0) FROM Job")
+                    highest_queue_position = cursor.fetchone()
+                    next_queue_position = (
+                        0 if highest_queue_position is None or highest_queue_position[0] is None else highest_queue_position[0]
+                    ) + 1000
+                    cursor.execute(
+                        "UPDATE Job SET queue_position = ?, return_to_queue = ?, info = ? WHERE id = ?",
+                        (next_queue_position, 1, info or "Job queued", self.job_id),
+                    )
+                finally:
+                    cursor.execute("COMMIT")
+
+        self._retry_db_operation(_queue_job)
+
+    def post_save_hook(self, save_path):
+        super(DiffusionTrainer, self).post_save_hook(save_path)
+        if not self.is_ui_trainer:
+            return
+
+        rolling_pause = getattr(self.save_config, "rolling_pause", 0) or 0
+        if rolling_pause <= 0 or not self.save_config.save_every:
+            return
+        if self.step_num <= 0 or self.step_num % self.save_config.save_every != 0:
+            return
+
+        checkpoint_count = self.step_num // self.save_config.save_every
+        if checkpoint_count <= 0 or checkpoint_count % rolling_pause != 0:
+            return
+
+        if self.accelerator.is_main_process:
+            print_acc(
+                f"Rolling pause triggered after checkpoint {checkpoint_count}; returning job to the back of the queue"
+            )
+            self._queue_job_to_bottom(f"Rolling pause after checkpoint {checkpoint_count}")
+
+        self.accelerator.wait_for_everyone()
+        self.maybe_stop()
+
     async def _update_status(self, status: AITK_Status, info: Optional[str] = None):
         if not self.accelerator.is_main_process or not self.is_ui_trainer:
             return

extensions_built_in/sd_trainer/UITrainer.py

@@ -8,6 +8,7 @@
 import threading
 import time
 import signal
+from jobs.exceptions import JobReturnedToQueueException, JobStoppedException
 
 AITK_Status = Literal["running", "stopped", "error", "completed"]
 
@@ -132,12 +133,12 @@ def maybe_stop(self):
             self._run_async_operation(
                 self._update_status("stopped", "Job stopped"))
             self.is_stopping = True
-            raise Exception("Job stopped")
+            raise JobStoppedException("Job stopped")
         if self.should_return_to_queue():
             self._run_async_operation(
                 self._update_status("queued", "Job queued"))
             self.is_stopping = True
-            raise Exception("Job returning to queue")
+            raise JobReturnedToQueueException("Job returning to queue")
 
     async def _update_key(self, key, value):
         if not self.accelerator.is_main_process:

jobs/exceptions.py

@@ -0,0 +1,10 @@
+class JobControlFlowException(Exception):
+    pass
+
+
+class JobStoppedException(JobControlFlowException):
+    pass
+
+
+class JobReturnedToQueueException(JobControlFlowException):
+    pass

run.py

@@ -38,6 +38,7 @@
 from toolkit.job import get_job
 from toolkit.accelerator import get_accelerator
 from toolkit.print import print_acc, setup_log_to_file
+from jobs.exceptions import JobReturnedToQueueException, JobStoppedException
 
 accelerator = get_accelerator()
 
@@ -111,13 +112,23 @@ def main():
             job.run()
             job.cleanup()
             jobs_completed += 1
+        except (JobReturnedToQueueException, JobStoppedException) as e:
+            print_acc(str(e))
+            try:
+                job.process[0].on_error(e)
+            except Exception as e2:
+                print_acc(f"Error running cleanup: {e2}")
+            finally:
+                job.cleanup()
         except Exception as e:
             print_acc(f"Error running job: {e}")
             jobs_failed += 1
             try:
                 job.process[0].on_error(e)
             except Exception as e2:
                 print_acc(f"Error running on_error: {e2}")
+            finally:
+                job.cleanup()
             if not args.recover:
                 print_end_message(jobs_completed, jobs_failed)
                 raise e
@@ -126,6 +137,8 @@ def main():
                 job.process[0].on_error(e)
             except Exception as e2:
                 print_acc(f"Error running on_error: {e2}")
+            finally:
+                job.cleanup()
             if not args.recover:
                 print_end_message(jobs_completed, jobs_failed)
                 raise e

run_modal.py

@@ -86,6 +86,7 @@
 
 import argparse
 from toolkit.job import get_job
+from jobs.exceptions import JobReturnedToQueueException, JobStoppedException
 
 def print_end_message(jobs_completed, jobs_failed):
     failure_string = f"{jobs_failed} failure{'' if jobs_failed == 1 else 's'}" if jobs_failed > 0 else ""
@@ -133,9 +134,23 @@ def main(config_file_list_str: str, recover: bool = False, name: str = None):
 
             job.cleanup()
             jobs_completed += 1
+        except (JobReturnedToQueueException, JobStoppedException) as e:
+            print(str(e))
+            try:
+                job.process[0].on_error(e)
+            except Exception as e2:
+                print(f"Error running cleanup: {e2}")
+            finally:
+                job.cleanup()
         except Exception as e:
             print(f"Error running job: {e}")
             jobs_failed += 1
+            try:
+                job.process[0].on_error(e)
+            except Exception as e2:
+                print(f"Error running on_error: {e2}")
+            finally:
+                job.cleanup()
             if not recover:
                 print_end_message(jobs_completed, jobs_failed)
                 raise e

toolkit/config_modules.py

@@ -25,6 +25,7 @@ def __init__(self, **kwargs):
         self.save_every: int = kwargs.get('save_every', 1000)
         self.dtype: str = kwargs.get('dtype', 'float16')
         self.max_step_saves_to_keep: int = kwargs.get('max_step_saves_to_keep', 5)
+        self.rolling_pause: int = kwargs.get('rolling_pause', 0)
         self.save_format: SaveFormat = kwargs.get('save_format', 'safetensors')
         if self.save_format not in ['safetensors', 'diffusers']:
             raise ValueError(f"save_format must be safetensors or diffusers, got {self.save_format}")

ui/cron/actions/startJob.ts

@@ -165,6 +165,7 @@ export default async function startJob(jobID: string) {
     data: {
       status: 'running',
       stop: false,
+      return_to_queue: false,
       info: 'Starting job...',
     },
   });

ui/next.config.ts

@@ -11,19 +11,12 @@ const nextConfig: NextConfig = {
     NEXT_PUBLIC_APP_VERSION: appVersion,
   },
   serverExternalPackages: ['macstats', 'osx-temperature-sensor'],
-  webpack: (config, { isServer }) => {
+   webpack: (config, { isServer }) => {
     if (isServer) {
       config.externals.push('osx-temperature-sensor', 'macstats');
     }
     return config;
   },
-  devIndicators: {
-    buildActivity: false,
-  },
-  typescript: {
-    // Remove this. Build fails because of route types
-    ignoreBuildErrors: true,
-  },
   experimental: {
     serverActions: {
       bodySizeLimit: '100gb',

ui/src/app/api/audio/art/[...audioPath]/route.ts

@@ -125,10 +125,10 @@ function extractArtFromTag(buf: Buffer): ArtResult {
   return null;
 }
 
-export async function GET(request: NextRequest, { params }: { params: { audioPath: string } }) {
+export async function GET(request: NextRequest, { params }: { params: Promise<{ audioPath: string[] }> }) {
   const { audioPath } = await params;
   try {
-    const filepath = decodeURIComponent(audioPath);
+    const filepath = decodeURIComponent(audioPath.join('/'));
 
     // Security check
     const datasetRoot = await getDatasetsRoot();

ui/src/app/api/datasets/fileCounts/route.ts

@@ -0,0 +1,77 @@
+import { NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot } from '@/server/settings';
+
+const imageExtensions = new Set(['.jpg', '.jpeg', '.png', '.bmp', '.webp']);
+const audioExtensions = new Set(['.mp3', '.wav', '.flac', '.ogg']);
+
+export async function POST(request: Request) {
+  const datasetsRoot = await getDatasetsRoot();
+  const body = await request.json();
+  const datasetPath = body?.datasetPath;
+
+  if (typeof datasetPath !== 'string' || datasetPath.trim() === '') {
+    return NextResponse.json({ error: 'datasetPath is required' }, { status: 400 });
+  }
+
+  const normalizedRoot = path.resolve(datasetsRoot);
+  const normalizedDatasetPath = path.resolve(datasetPath);
+  const relativePath = path.relative(normalizedRoot, normalizedDatasetPath);
+  const isOutsideRoot = relativePath.startsWith('..') || path.isAbsolute(relativePath);
+
+  if (isOutsideRoot) {
+    return NextResponse.json({ error: 'Dataset path must be inside the datasets folder' }, { status: 400 });
+  }
+
+  if (!fs.existsSync(normalizedDatasetPath)) {
+    return NextResponse.json({ error: 'Dataset folder not found' }, { status: 404 });
+  }
+
+  try {
+    const counts = countDatasetFiles(normalizedDatasetPath);
+    return NextResponse.json(counts);
+  } catch (error) {
+    console.error('Error counting dataset files:', error);
+    return NextResponse.json({ error: 'Failed to count dataset files' }, { status: 500 });
+  }
+}
+
+function countDatasetFiles(dir: string) {
+  let imageFileCount = 0;
+  let audioFileCount = 0;
+  const entries = fs.readdirSync(dir, { withFileTypes: true });
+
+  for (const entry of entries) {
+    if (entry.name.startsWith('.')) {
+      continue;
+    }
+
+    const entryPath = path.join(dir, entry.name);
+
+    if (entry.isDirectory()) {
+      if (entry.name === '_controls') {
+        continue;
+      }
+
+      const nestedCounts = countDatasetFiles(entryPath);
+      imageFileCount += nestedCounts.imageFileCount;
+      audioFileCount += nestedCounts.audioFileCount;
+      continue;
+    }
+
+    if (!entry.isFile()) {
+      continue;
+    }
+
+    const extension = path.extname(entry.name).toLowerCase();
+    if (imageExtensions.has(extension)) {
+      imageFileCount += 1;
+    }
+    if (audioExtensions.has(extension)) {
+      audioFileCount += 1;
+    }
+  }
+
+  return { imageFileCount, audioFileCount };
+}

ui/src/app/api/files/[...filePath]/route.ts

@@ -4,11 +4,11 @@ import fs from 'fs';
 import path from 'path';
 import { getDatasetsRoot, getTrainingFolder } from '@/server/settings';
 
-export async function GET(request: NextRequest, { params }: { params: { filePath: string } }) {
+export async function GET(request: NextRequest, { params }: { params: Promise<{ filePath: string[] }> }) {
   const { filePath } = await params;
   try {
     // Decode the path
-    const decodedFilePath = decodeURIComponent(filePath);
+    const decodedFilePath = decodeURIComponent(filePath.join('/'));
 
     // Get allowed directories
     const datasetRoot = await getDatasetsRoot();

ui/src/app/api/img/[...imagePath]/route.ts

@@ -29,11 +29,11 @@ const contentTypeMap: { [key: string]: string } = {
   '.ogg': 'audio/ogg',
 };
 
-export async function GET(request: NextRequest, { params }: { params: { imagePath: string } }) {
+export async function GET(request: NextRequest, { params }: { params: Promise<{ imagePath: string[] }> }) {
   const { imagePath } = await params;
   try {
     // Decode the path
-    const filepath = decodeURIComponent(imagePath);
+    const filepath = decodeURIComponent(imagePath.join('/'));
 
     // Get allowed directories
     const datasetRoot = await getDatasetsRoot();

ui/src/app/api/jobs/[jobID]/delete/route.ts

@@ -6,7 +6,7 @@ import fs from 'fs';
 
 const prisma = new PrismaClient();
 
-export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
+export async function GET(request: NextRequest, { params }: { params: Promise<{ jobID: string }> }) {
   const { jobID } = await params;
 
   const job = await prisma.job.findUnique({

ui/src/app/api/jobs/[jobID]/files/route.ts

@@ -6,7 +6,7 @@ import { getTrainingFolder } from '@/server/settings';
 
 const prisma = new PrismaClient();
 
-export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
+export async function GET(request: NextRequest, { params }: { params: Promise<{ jobID: string }> }) {
   const { jobID } = await params;
 
   const job = await prisma.job.findUnique({

ui/src/app/api/jobs/[jobID]/log/route.ts

@@ -6,7 +6,7 @@ import { getTrainingFolder } from '@/server/settings';
 
 const prisma = new PrismaClient();
 
-export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
+export async function GET(request: NextRequest, { params }: { params: Promise<{ jobID: string }> }) {
   const { jobID } = await params;
 
   const job = await prisma.job.findUnique({