From 0f580f0663279ace30317c8bf90ae736064957cb Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 12 Jun 2026 10:26:52 -0600 Subject: [PATCH] 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. --- .../captioner/BaseCaptioner.py | 32 +++++++++++++ ui/next.config.ts | 8 ++++ ui/prisma/schema.prisma | 1 + ui/src/app/api/jobs/[jobID]/plugin/route.ts | 47 +++++++++++++++++++ ui/src/app/jobs/[jobID]/page.tsx | 33 +++++++++++-- ui/src/components/ActiveJobWidget.tsx | 10 ++-- ui/src/components/CaptionSimpleJob.tsx | 5 ++ ui/src/components/JobOverview.tsx | 2 +- ui/src/components/JobPlugin.tsx | 43 +++++++++++++++++ ui/src/components/JobsTable.tsx | 9 ++-- ui/src/helpers/captionJobConfig.ts | 1 + ui/src/types.ts | 1 + ui/src/utils/jobs.ts | 3 ++ 13 files changed, 180 insertions(+), 15 deletions(-) create mode 100644 ui/src/app/api/jobs/[jobID]/plugin/route.ts create mode 100644 ui/src/components/JobPlugin.tsx diff --git a/extensions_built_in/captioner/BaseCaptioner.py b/extensions_built_in/captioner/BaseCaptioner.py index cd61a8b..cb63002 100644 --- a/extensions_built_in/captioner/BaseCaptioner.py +++ b/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() diff --git a/ui/next.config.ts b/ui/next.config.ts index 2502cc7..244f141 100644 --- a/ui/next.config.ts +++ b/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'); diff --git a/ui/prisma/schema.prisma b/ui/prisma/schema.prisma index bb10e9a..79c90fb 100644 --- a/ui/prisma/schema.prisma +++ b/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) diff --git a/ui/src/app/api/jobs/[jobID]/plugin/route.ts b/ui/src/app/api/jobs/[jobID]/plugin/route.ts new file mode 100644 index 0000000..83f09e5 --- /dev/null +++ b/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', + }, + }); +} diff --git a/ui/src/app/jobs/[jobID]/page.tsx b/ui/src/app/jobs/[jobID]/page.tsx index 07a0526..7f5c2a8 100644 --- a/ui/src/app/jobs/[jobID]/page.tsx +++ b/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('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 (