Add compiling option for captioning. Added progress bar for captioning on active job widget. Track steps on caption. Added reverse proxy iframe that plugins can use to add ui functionality.

extensions_built_in/captioner/BaseCaptioner.py

@@ -43,6 +43,7 @@ class CaptionConfig:
         self.caption_prompt = kwargs.get(
             "caption_prompt", "Describe this image in detail."
         )
+        self.compile = kwargs.get("compile", False)
 
 
 class BaseCaptioner(BaseExtensionProcess):
@@ -82,6 +83,7 @@ class BaseCaptioner(BaseExtensionProcess):
         self.model2 = None
         self.processor2 = None
         self.file_paths = []
+        self.step_num = 0
         self.device_torch = torch.device(self.caption_config.device)
         self.torch_dtype = get_torch_dtype(self.caption_config.dtype)
 
@@ -91,8 +93,11 @@ class BaseCaptioner(BaseExtensionProcess):
             self.start_stop_watcher()
             self.update_status("running", "Loading Model")
             self.load_model()
+            self.maybe_compile_models()
             self.update_status("running", "Looking for files")
             self.find_files()
+            self.update_db_key("total_steps", len(self.file_paths))
+            self.update_step()
             self.update_status("running", f"Captioning {len(self.file_paths)} files")
             self.run_caption_loop()
             self.update_status("completed", "Captioning completed")
@@ -117,6 +122,9 @@ class BaseCaptioner(BaseExtensionProcess):
             except Exception as e:
                 print(f"Error captioning file {file_path}: {e}")
                 continue
+            finally:
+                self.step_num += 1
+                self.update_step()
 
     def load_pil_image(self, file_path: str, max_res: Optional[int] = None) -> Image:
         image = Image.open(file_path).convert("RGB")
@@ -183,6 +191,30 @@ class BaseCaptioner(BaseExtensionProcess):
     def load_model(self):
         raise NotImplementedError("Model loading not implemented for this captioner")
 
+    def maybe_compile_models(self):
+        if not self.caption_config.compile:
+            return
+        import importlib.util
+
+        if importlib.util.find_spec("triton") is None:
+            print(
+                "[AITK] compile requested but triton is not installed, skipping compilation."
+            )
+            return
+        try:
+            # compilation happens lazily on first forward, so fall back to
+            # eager there too if the backend fails (e.g. broken triton install)
+            torch._dynamo.config.suppress_errors = True
+            for model in [self.model, self.model2]:
+                if model is not None and isinstance(model, torch.nn.Module):
+                    # dynamic=True avoids recompiling for every new image/token shape
+                    model.compile(dynamic=True)
+            print(
+                "[AITK] Model compilation enabled. The first few items will be slow while the model compiles."
+            )
+        except Exception as e:
+            print(f"[AITK] Failed to compile model, continuing without compile: {e}")
+
     def start_stop_watcher(self, interval_sec: float = 5.0):
         """
         Start a daemon thread that periodically checks should_stop()

ui/next.config.ts

@@ -11,6 +11,14 @@ const nextConfig: NextConfig = {
     NEXT_PUBLIC_APP_VERSION: appVersion,
   },
   serverExternalPackages: ['macstats', 'osx-temperature-sensor'],
+  async rewrites() {
+    return [
+      {
+        source: '/proxy-8866/:path*',
+        destination: 'http://localhost:8866/:path*',
+      },
+    ];
+  },
   webpack: (config, { isServer }) => {
     if (isServer) {
       config.externals.push('osx-temperature-sensor', 'macstats');

ui/prisma/schema.prisma

@@ -32,6 +32,7 @@ model Job {
   stop            Boolean  @default(false)
   return_to_queue Boolean  @default(false) // same as stop, but will be set to 'queued' when stopped
   step            Int      @default(0)
+  total_steps     Int?
   info            String   @default("")
   speed_string    String   @default("")
   queue_position  Int      @default(0)

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

@@ -0,0 +1,47 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { PrismaClient } from '@prisma/client';
+import path from 'path';
+import fs from 'fs';
+import { getTrainingFolder } from '@/server/settings';
+
+const prisma = new PrismaClient();
+
+export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
+  const { jobID } = await params;
+
+  const job = await prisma.job.findUnique({
+    where: { id: jobID },
+  });
+
+  if (!job) {
+    return NextResponse.json({ error: 'Job not found' }, { status: 404 });
+  }
+
+  const trainingFolder = await getTrainingFolder();
+  const jobFolder = path.join(trainingFolder, job.name);
+  const pluginPath = path.join(jobFolder, 'plugin.html');
+
+  if (!fs.existsSync(pluginPath)) {
+    return NextResponse.json({ exists: false, html: null });
+  }
+
+  // lightweight existence check used to decide if the Plugin tab should show
+  if (request.nextUrl.searchParams.get('check') === '1') {
+    return NextResponse.json({ exists: true, html: null });
+  }
+
+  // serve the raw html so it can be loaded directly as an iframe src
+  let html = '';
+  try {
+    html = fs.readFileSync(pluginPath, 'utf-8');
+  } catch (error) {
+    console.error('Error reading plugin file:', error);
+    return NextResponse.json({ error: 'Error reading plugin file' }, { status: 500 });
+  }
+  return new NextResponse(html, {
+    headers: {
+      'Content-Type': 'text/html; charset=utf-8',
+      'Cache-Control': 'no-store',
+    },
+  });
+}

ui/src/app/jobs/[jobID]/page.tsx

@@ -1,8 +1,8 @@
 'use client';
 
-import { useState, use } from 'react';
+import { useState, useEffect, use } from 'react';
 import { FaChevronLeft } from 'react-icons/fa';
-import { MdDashboard, MdImage, MdShowChart, MdCode } from 'react-icons/md';
+import { MdDashboard, MdImage, MdShowChart, MdCode, MdExtension } from 'react-icons/md';
 import { Button } from '@headlessui/react';
 import { TopBar, MainContent } from '@/components/layout';
 import useJob from '@/hooks/useJob';
@@ -12,9 +12,11 @@ import { redirect } from 'next/navigation';
 import JobActionBar from '@/components/JobActionBar';
 import JobConfigViewer from '@/components/JobConfigViewer';
 import JobLossGraph from '@/components/JobLossGraph';
+import JobPlugin from '@/components/JobPlugin';
 import { Job } from '@prisma/client';
+import { apiClient } from '@/utils/api';
 
-type PageKey = 'overview' | 'samples' | 'config' | 'loss_log';
+type PageKey = 'overview' | 'samples' | 'config' | 'loss_log' | 'plugin';
 
 interface Page {
   name: string;
@@ -58,6 +60,13 @@ const pages: Page[] = [
     component: JobConfigViewer,
     mainCss: 'pt-[80px] px-0 pb-0',
   },
+  {
+    name: 'Plugin',
+    value: 'plugin',
+    icon: MdExtension,
+    component: JobPlugin,
+    mainCss: 'pt-[80px] px-0 pb-0',
+  },
 ];
 
 export default function JobPage({ params }: { params: { jobID: string } }) {
@@ -65,6 +74,21 @@ export default function JobPage({ params }: { params: { jobID: string } }) {
   const jobID = usableParams.jobID;
   const { job, status, refreshJob } = useJob(jobID, 5000);
   const [pageKey, setPageKey] = useState<PageKey>('overview');
+  const [hasPlugin, setHasPlugin] = useState(false);
+
+  // poll for plugin.html in the job folder; show the Plugin tab if it exists
+  useEffect(() => {
+    const checkPlugin = () => {
+      apiClient
+        .get(`/api/jobs/${jobID}/plugin?check=1`)
+        .then(res => res.data)
+        .then(data => setHasPlugin(!!data.exists))
+        .catch(() => {});
+    };
+    checkPlugin();
+    const interval = setInterval(checkPlugin, 5000);
+    return () => clearInterval(interval);
+  }, [jobID]);
 
   const page = pages.find(p => p.value === pageKey);
 
@@ -117,6 +141,9 @@ export default function JobPage({ params }: { params: { jobID: string } }) {
           if (page.jobTypes && !page.jobTypes.includes(jobType)) {
             return null;
           }
+          if (page.value === 'plugin' && !hasPlugin) {
+            return null;
+          }
           return (
             <Button
               key={page.value}

ui/src/components/ActiveJobWidget.tsx

@@ -3,7 +3,7 @@
 import Link from 'next/link';
 import { CgSpinner } from 'react-icons/cg';
 import useJobsList from '@/hooks/useJobsList';
-import { JobConfig } from '@/types';
+import { getTotalSteps } from '@/utils/jobs';
 
 export default function ActiveJobWidget() {
   const { jobs } = useJobsList({ onlyActive: true, reloadInterval: 5000 });
@@ -17,13 +17,11 @@ export default function ActiveJobWidget() {
           {jobs.map(job => {
             let totalSteps: number | undefined;
             try {
-              const cfg: JobConfig = JSON.parse(job.job_config);
-              totalSteps = cfg.config.process[0].train?.steps;
+              totalSteps = getTotalSteps(job);
             } catch {
               totalSteps = undefined;
             }
-            const isTrain = job.job_type === 'train';
-            const pct = isTrain && totalSteps ? Math.min(100, (job.step / totalSteps) * 100) : 0;
+            const pct = totalSteps ? Math.min(100, (job.step / totalSteps) * 100) : 0;
             const isRunning = job.status === 'running' || job.status === 'stopping';
 
             let label = job.name;
@@ -47,7 +45,7 @@ export default function ActiveJobWidget() {
                     {isRunning && <CgSpinner className="animate-spin text-blue-400 flex-shrink-0" />}
                     <span className="text-xs text-gray-100 truncate min-w-0 flex-1">{label}</span>
                   </div>
-                  {isTrain && totalSteps ? (
+                  {totalSteps ? (
                     <div className="mt-1.5">
                       <div className="bg-gray-700 rounded-full h-1">
                         <div className="bg-blue-500 h-1 rounded-full transition-all" style={{ width: `${pct}%` }} />

ui/src/components/CaptionSimpleJob.tsx

@@ -178,6 +178,11 @@ const CaptionSimpleJob: React.FC<Props> = ({ jobConfig, setJobConfig, gpuIDs, se
               checked={jobConfig.config.process[0].caption.recaption}
               onChange={value => setJobConfig(value, 'config.process[0].caption.recaption')}
             />
+            <Checkbox
+              label="Compile Models"
+              checked={jobConfig.config.process[0].caption.compile || false}
+              onChange={value => setJobConfig(value, 'config.process[0].caption.compile')}
+            />
           </FormGroup>
         </div>
       </div>

ui/src/components/JobOverview.tsx

@@ -102,7 +102,7 @@ export default function JobOverview({ job }: JobOverviewProps) {
 
         <div className="p-4 space-y-6 flex flex-col flex-grow">
           {/* Progress Bar */}
-          {job.job_type === 'train' && (
+          {totalSteps > 0 && (
             <div className="space-y-2">
               <div className="flex items-center justify-between text-sm">
                 <span className="text-gray-400">Progress</span>

ui/src/components/JobPlugin.tsx

@@ -0,0 +1,43 @@
+'use client';
+
+import { useEffect, useState } from 'react';
+import { Job } from '@prisma/client';
+import { apiClient } from '@/utils/api';
+
+interface JobPluginProps {
+  job: Job;
+}
+
+export default function JobPlugin({ job }: JobPluginProps) {
+  const [html, setHtml] = useState<string | null>(null);
+  const [status, setStatus] = useState<'loading' | 'success' | 'error'>('loading');
+
+  // An iframe src request cannot send the Authorization header the API requires,
+  // so fetch the html through apiClient and load it via srcDoc instead.
+  useEffect(() => {
+    apiClient
+      .get(`/api/jobs/${job.id}/plugin`, { responseType: 'text' })
+      .then(res => {
+        setHtml(res.data);
+        setStatus('success');
+      })
+      .catch(error => {
+        console.error('Error fetching plugin:', error);
+        setStatus('error');
+      });
+  }, [job.id]);
+
+  if (status === 'loading') {
+    return <p className="p-4">Loading plugin...</p>;
+  }
+  if (status === 'error' || html == null) {
+    return <p className="p-4">Error loading plugin</p>;
+  }
+
+  // Sandboxed iframe keeps the plugin's scripts and styles isolated from the app.
+  // allow-scripts lets its <script> tags run, but it cannot touch the parent page,
+  // cookies, or storage. The plugin sizes itself to 100% of its container, so the
+  // iframe must have an explicit size.
+  // 80px = top bar (48px) + tab bar (32px)
+  return <iframe srcDoc={html} sandbox="allow-scripts" className="block w-full h-[calc(100dvh-80px)] border-0" />;
+}

ui/src/components/JobsTable.tsx

@@ -2,7 +2,7 @@ import { useMemo, useState } from 'react';
 import useJobsList from '@/hooks/useJobsList';
 import Link from 'next/link';
 import UniversalTable, { TableColumn } from '@/components/UniversalTable';
-import { GpuInfo, JobConfig } from '@/types';
+import { GpuInfo } from '@/types';
 import JobActionBar from './JobActionBar';
 import { Job, Queue } from '@prisma/client';
 import useQueueList from '@/hooks/useQueueList';
@@ -11,7 +11,7 @@ import { startQueue, stopQueue } from '@/utils/queue';
 import { CgSpinner } from 'react-icons/cg';
 import useGPUInfo from '@/hooks/useGPUInfo';
 import { openConfirm } from '@/components/ConfirmModal';
-import { deleteJob, stopJob } from '@/utils/jobs';
+import { deleteJob, getTotalSteps, stopJob } from '@/utils/jobs';
 import { Trash2 } from 'lucide-react';
 
 interface JobsTableProps {
@@ -150,11 +150,10 @@ export default function JobsTable({ onlyActive = false, job_type = null }: JobsT
       title: 'Steps',
       key: 'steps',
       render: row => {
-        const jobConfig: JobConfig = JSON.parse(row.job_config);
-        if (row.job_type !== 'train') {
+        const totalSteps = getTotalSteps(row);
+        if (!totalSteps) {
           return <></>;
         }
-        const totalSteps = jobConfig.config.process[0].train?.steps;
 
         return (
           <div>

ui/src/helpers/captionJobConfig.ts

@@ -21,6 +21,7 @@ export const defaultCaptionJobConfig: CaptionJobConfig = {
           extensions: ['mp3', 'wav', 'flac', 'ogg'],
           path_to_caption: '',
           recaption: false,
+          compile: false,
           caption_extension: 'txt',
         },
       },

ui/src/types.ts

@@ -268,6 +268,7 @@ export interface CaptionProcessConfig {
     extensions: string[];
     path_to_caption: string;
     recaption: boolean;
+    compile?: boolean;
     caption_prompt?: string;
     max_res?: number;
     max_new_tokens?: number;

ui/src/utils/jobs.ts

@@ -107,6 +107,9 @@ export const getNumberOfSamples = (job: Job) => {
 };
 
 export const getTotalSteps = (job: Job) => {
+  if (job.total_steps != null) {
+    return job.total_steps;
+  }
   const jobConfig = getJobConfig(job);
   return jobConfig.config.process[0].train?.steps || 0;
 };