Add support for ACE-Step 1.5 and ACE-Step 1.5 XL. Also added dataset captioning through the UI. (#785)

* Base ace step 1.5 xl added. Generating, still wip on training and ui

* Base training code done

* Fix some issues with caching text embeddings. Update sample cards to show audio

* Fix issue with quantizing ace step

* Add album artwork to samples with waveform.

* Cleanup logs

* Add album art endpoint to speed up album art loading

* Made an make video with artwork script

* Make ui handle basic audio models. Make multi line adjustments to the editor and better syntax hilighting.

* Add prompt tagging system for special tagged models.

* prompt tagging processing for ui working.

* Moved default samples to a special file so we can add more when needed and they can be adjusted for a specific model

* Add a captioner job with music captioner that is prepped for use with the ui

* Add basit ui setup for captioning modal and handeling captioning jobs

* Starting captioning job from ui working. Still better management for it.

* Better filtering of job options in the job view for captioning jobs

* Added qwen3 vl as a captioner for images

* Have an indicator when a dataset is being captioned.

* Adjust the way caption jobs look in the queue

* Fix a few issues. Adjust defaults.

* Version bump

* Added ace step to the readme.

ui/src/components/AudioPlayer.tsx

@@ -1,6 +1,7 @@
 'use client';
 
 import React, { useEffect, useMemo, useRef, useState } from 'react';
+import { apiClient } from '@/utils/api';
 
 type AudioPlayerProps = {
   src: string;
@@ -38,214 +39,17 @@ function broadcastExclusivePlay(token: string) {
 }
 
 /**
- * ID3 helpers (v2.2/v2.3/v2.4):
- * - robust album art extraction: APIC (v2.3/2.4) + PIC (v2.2)
- * - basic text frames: title/artist/album
- * - handles tag-level unsynchronisation
- *
- * Requires fetch() byte access; if CORS blocks, it will fall back gracefully.
+ * Build the server-side album-art URL from the audio src.
+ * The audio src is `/api/img/{encodedPath}` — we extract the path
+ * and point to `/api/audio/art/{encodedPath}` instead.
  */
-type Id3Meta = {
-  title?: string;
-  artist?: string;
-  album?: string;
-  pictureUrl?: string; // object URL
-};
-
-function synchsafeToInt(b0: number, b1: number, b2: number, b3: number) {
-  return ((b0 & 0x7f) << 21) | ((b1 & 0x7f) << 14) | ((b2 & 0x7f) << 7) | (b3 & 0x7f);
-}
-
-function deUnsync(bytes: Uint8Array) {
-  const out: number[] = [];
-  for (let i = 0; i < bytes.length; i++) {
-    const cur = bytes[i];
-    out.push(cur);
-    if (cur === 0xff && i + 1 < bytes.length && bytes[i + 1] === 0x00) i += 1;
+function albumArtUrlFromSrc(src: string): string {
+  const prefix = '/api/img/';
+  if (src.startsWith(prefix)) {
+    return `/api/audio/art/${src.slice(prefix.length)}`;
   }
-  return new Uint8Array(out);
-}
-
-function decodeText(encoding: number, bytes: Uint8Array) {
-  let end = bytes.length;
-  while (end > 0 && bytes[end - 1] === 0) end--;
-  const b = bytes.slice(0, end);
-
-  try {
-    if (encoding === 0) return new TextDecoder('latin1').decode(b);
-    if (encoding === 1) return new TextDecoder('utf-16').decode(b);
-    if (encoding === 2) return new TextDecoder('utf-16be').decode(b);
-    if (encoding === 3) return new TextDecoder('utf-8').decode(b);
-  } catch {
-    // ignore
-  }
-  return new TextDecoder('latin1').decode(b);
-}
-
-function readNullTerminated(bytes: Uint8Array, start: number, encoding: number) {
-  if (encoding === 1 || encoding === 2) {
-    let i = start;
-    while (i + 1 < bytes.length && !(bytes[i] === 0 && bytes[i + 1] === 0)) i += 2;
-    const textBytes = bytes.slice(start, i);
-    return { text: decodeText(encoding, textBytes), next: i + 2 };
-  } else {
-    let i = start;
-    while (i < bytes.length && bytes[i] !== 0) i++;
-    const textBytes = bytes.slice(start, i);
-    return { text: decodeText(encoding, textBytes), next: i + 1 };
-  }
-}
-
-async function fetchBytes(src: string, start: number, endInclusive: number) {
-  const wantLen = endInclusive - start + 1;
-
-  try {
-    const r = await fetch(src, { headers: { Range: `bytes=${start}-${endInclusive}` } });
-    if (!r.ok) throw new Error('range not ok');
-    const buf = await r.arrayBuffer();
-    return new Uint8Array(buf);
-  } catch {
-    const r = await fetch(src);
-    if (!r.ok) throw new Error('fetch not ok');
-    const buf = await r.arrayBuffer();
-    const u8 = new Uint8Array(buf);
-    if (start === 0 && u8.length >= wantLen) return u8.slice(0, wantLen);
-    return u8.slice(start, Math.min(u8.length, endInclusive + 1));
-  }
-}
-
-async function extractId3MetaAndArt(src: string, maxTagBytes = 4_000_000): Promise<Id3Meta> {
-  const head = await fetchBytes(src, 0, 64 * 1024 - 1).catch(() => null);
-  if (!head || head.length < 10) return {};
-  if (head[0] !== 0x49 || head[1] !== 0x44 || head[2] !== 0x33) return {};
-
-  const verMajor = head[3]; // 2,3,4
-  const flags = head[5];
-  const tagSize = synchsafeToInt(head[6], head[7], head[8], head[9]);
-  const tagEnd = 10 + tagSize;
-
-  const need = Math.min(tagEnd, maxTagBytes);
-  let tagBytes = head;
-  if (head.length < need) {
-    const more = await fetchBytes(src, 0, need - 1).catch(() => null);
-    if (!more) return {};
-    tagBytes = more;
-  } else {
-    tagBytes = head.slice(0, need);
-  }
-
-  const tagUnsync = (flags & 0x80) !== 0;
-  const tagDataRaw = tagBytes.slice(10, Math.min(tagBytes.length, tagEnd));
-  const tagData = tagUnsync ? deUnsync(tagDataRaw) : tagDataRaw;
-
-  let offset = 0;
-
-  // Extended header (v2.3/v2.4)
-  if (verMajor === 3 || verMajor === 4) {
-    const hasExt = (flags & 0x40) !== 0;
-    if (hasExt && tagData.length >= 4) {
-      let extSize = 0;
-      if (verMajor === 4) extSize = synchsafeToInt(tagData[0], tagData[1], tagData[2], tagData[3]);
-      else extSize = (tagData[0] << 24) | (tagData[1] << 16) | (tagData[2] << 8) | tagData[3];
-      offset += 4 + Math.max(0, extSize);
-    }
-  }
-
-  const meta: Id3Meta = {};
-  const setIfEmpty = (k: keyof Id3Meta, v?: string) => {
-    if (!v) return;
-    if (!meta[k]) meta[k] = v;
-  };
-
-  while (offset < tagData.length) {
-    if (tagData[offset] === 0x00) break;
-
-    if (verMajor === 2) {
-      if (offset + 6 > tagData.length) break;
-      const id = new TextDecoder('latin1').decode(tagData.slice(offset, offset + 3));
-      const size = (tagData[offset + 3] << 16) | (tagData[offset + 4] << 8) | tagData[offset + 5];
-      offset += 6;
-      if (!id.trim() || size <= 0 || offset + size > tagData.length) break;
-
-      const frame = tagData.slice(offset, offset + size);
-
-      if (id === 'TT2' || id === 'TP1' || id === 'TAL') {
-        const enc = frame[0];
-        const txt = decodeText(enc, frame.slice(1));
-        if (id === 'TT2') setIfEmpty('title', txt);
-        if (id === 'TP1') setIfEmpty('artist', txt);
-        if (id === 'TAL') setIfEmpty('album', txt);
-      }
-
-      if (id === 'PIC' && frame.length > 6) {
-        const enc = frame[0];
-        const fmt = new TextDecoder('latin1').decode(frame.slice(1, 4)).toLowerCase();
-        const imgType = fmt === 'png' ? 'image/png' : 'image/jpeg';
-        let p = 5;
-        const desc = readNullTerminated(frame, p, enc);
-        p = desc.next;
-        if (p < frame.length) {
-          const img = frame.slice(p);
-          if (img.length > 64) {
-            const blob = new Blob([img], { type: imgType });
-            meta.pictureUrl = URL.createObjectURL(blob);
-          }
-        }
-      }
-
-      offset += size;
-    } else {
-      if (offset + 10 > tagData.length) break;
-
-      const id = new TextDecoder('latin1').decode(tagData.slice(offset, offset + 4));
-      let size = 0;
-      if (verMajor === 4)
-        size = synchsafeToInt(tagData[offset + 4], tagData[offset + 5], tagData[offset + 6], tagData[offset + 7]);
-      else
-        size =
-          (tagData[offset + 4] << 24) | (tagData[offset + 5] << 16) | (tagData[offset + 6] << 8) | tagData[offset + 7];
-
-      const flag2 = tagData[offset + 9];
-      offset += 10;
-
-      if (!id.trim() || size <= 0 || offset + size > tagData.length) break;
-
-      let frame = tagData.slice(offset, offset + size);
-
-      const frameUnsync = verMajor === 4 && (flag2 & 0x02) !== 0;
-      if (frameUnsync) frame = deUnsync(frame);
-
-      if (id === 'TIT2' || id === 'TPE1' || id === 'TALB') {
-        const enc = frame[0];
-        const txt = decodeText(enc, frame.slice(1));
-        if (id === 'TIT2') setIfEmpty('title', txt);
-        if (id === 'TPE1') setIfEmpty('artist', txt);
-        if (id === 'TALB') setIfEmpty('album', txt);
-      }
-
-      if (id === 'APIC' && frame.length > 10) {
-        const enc = frame[0];
-        const mimeZ = readNullTerminated(frame, 1, 0);
-        const mime = mimeZ.text || 'image/jpeg';
-        let p = mimeZ.next;
-        if (p < frame.length) p += 1; // pictureType
-        const desc = readNullTerminated(frame, p, enc);
-        p = desc.next;
-
-        if (p < frame.length) {
-          const img = frame.slice(p);
-          if (img.length > 64) {
-            const blob = new Blob([img], { type: mime });
-            if (!meta.pictureUrl) meta.pictureUrl = URL.createObjectURL(blob);
-          }
-        }
-      }
-
-      offset += size;
-    }
-  }
-
-  return meta;
+  // Fallback: assume src is already an encoded path
+  return `/api/audio/art/${encodeURIComponent(src)}`;
 }
 
 export default function AudioPlayer({
@@ -278,12 +82,8 @@ export default function AudioPlayer({
   const [dragging, setDragging] = useState(false);
   const [dragValue, setDragValue] = useState(0);
 
-  // Meta + artwork
-  const [metaTitle, setMetaTitle] = useState<string | null>(null);
-  const [metaArtist, setMetaArtist] = useState<string | null>(null);
-  const [metaAlbum, setMetaAlbum] = useState<string | null>(null);
+  // Album art: served by /api/audio/art endpoint (fast, server-side extraction)
   const [albumArtUrl, setAlbumArtUrl] = useState<string | null>(null);
-  const albumArtBlobUrlRef = useRef<string | null>(null);
 
   // WebAudio analyser
   const audioCtxRef = useRef<AudioContext | null>(null);
@@ -318,57 +118,23 @@ export default function AudioPlayer({
     c.height = Math.max(1, Math.floor(cssH * dpr));
   }, [size]);
 
-  // Extract meta + art
+  // Set album art URL from server endpoint
   useEffect(() => {
+    const artUrl = albumArtUrlFromSrc(src);
     let cancelled = false;
-
-    (async () => {
-      if (albumArtBlobUrlRef.current) {
-        URL.revokeObjectURL(albumArtBlobUrlRef.current);
-        albumArtBlobUrlRef.current = null;
-      }
-      setAlbumArtUrl(null);
-      setMetaTitle(null);
-      setMetaArtist(null);
-      setMetaAlbum(null);
-
-      try {
-        const meta = await extractId3MetaAndArt(src);
-        if (cancelled) return;
-
-        if (meta.title) setMetaTitle(meta.title);
-        if (meta.artist) setMetaArtist(meta.artist);
-        if (meta.album) setMetaAlbum(meta.album);
-
-        if (meta.pictureUrl) {
-          albumArtBlobUrlRef.current = meta.pictureUrl;
-          setAlbumArtUrl(meta.pictureUrl);
-        } else if (defaultAlbumArtUrl) {
-          setAlbumArtUrl(defaultAlbumArtUrl);
-        }
-      } catch {
-        if (!cancelled && defaultAlbumArtUrl) setAlbumArtUrl(defaultAlbumArtUrl);
-      }
-    })();
-
+    apiClient
+      .head(artUrl)
+      .then(() => {
+        if (!cancelled) setAlbumArtUrl(artUrl);
+      })
+      .catch(() => {
+        if (!cancelled) setAlbumArtUrl(defaultAlbumArtUrl ?? null);
+      });
     return () => {
       cancelled = true;
     };
   }, [src, defaultAlbumArtUrl]);
 
-  // Cleanup artwork blob URL on unmount
-  useEffect(() => {
-    return () => {
-      if (albumArtBlobUrlRef.current) {
-        URL.revokeObjectURL(albumArtBlobUrlRef.current);
-        albumArtBlobUrlRef.current = null;
-      }
-    };
-  }, []);
-
-  const effectiveTitle = metaTitle || title;
-  const effectiveSubtitle = metaArtist || subtitle || metaAlbum || '';
-
   const progress = useMemo(() => {
     const cur = dragging ? dragValue : t;
     if (!duration) return 0;
@@ -677,11 +443,11 @@ export default function AudioPlayer({
             className="truncate text-gray-200"
             style={{ fontSize: titleSize, lineHeight: 1.1, letterSpacing: '0.01em' }}
           >
-            {effectiveTitle}
+            {title}
           </div>
-          {effectiveSubtitle ? (
+          {subtitle ? (
             <div className="mt-1 truncate text-gray-400" style={{ fontSize: subSize, lineHeight: 1.15 }}>
-              {effectiveSubtitle}
+              {subtitle}
             </div>
           ) : null}
           {err ? (

ui/src/components/AutoCaptionButton.tsx

@@ -0,0 +1,39 @@
+import React, {use, useEffect} from 'react';
+import { Button } from '@headlessui/react';
+import { CaptionDatasetModal, openCaptionDatasetModal } from '@/components/CaptionDatasetModal';
+import useJobByRef from '@/hooks/useJobByRef';
+import Link from 'next/link';
+import { Loader2 } from 'lucide-react';
+
+type AutoCaptionButtonProps = {
+  datasetPath: string;
+  setIsAutoCaptioning?: (isAutoCaptioning: boolean) => void;
+};
+
+export default function AutoCaptionButton({ datasetPath, setIsAutoCaptioning }: AutoCaptionButtonProps) {
+  const { job, status, refreshJob } = useJobByRef(datasetPath, 5000);
+  useEffect(() => {
+    if (setIsAutoCaptioning) {
+      setIsAutoCaptioning(!!(job && (job.status === 'running')));
+    }
+  }, [job, setIsAutoCaptioning]);
+  
+  if (job && (job.status === 'running' || job.status === 'queued')) {
+    return (
+      <Link href={`/jobs/${job.id}`} className="text-white bg-gray-400 px-3 py-1 rounded-md mr-2 inline-flex items-center gap-1.5">
+        <Loader2 className="w-4 h-4 animate-spin" />
+        Auto Captioning...
+      </Link>
+    );
+  }
+  return (
+    <Button
+      className="text-white bg-blue-600 px-3 py-1 rounded-md mr-2"
+      onClick={() => openCaptionDatasetModal(datasetPath, () => {
+        refreshJob();
+      })}
+    >
+      Auto Caption
+    </Button>
+  );
+}

ui/src/components/CaptionDatasetModal.tsx

@@ -0,0 +1,273 @@
+import React, { useState, useEffect, useRef } from 'react';
+import { Modal } from '@/components/Modal';
+import { createGlobalState } from 'react-global-hooks';
+import { useFromNull } from '@/hooks/useFromNull';
+import {
+  Checkbox,
+  CreatableSelectInput,
+  FormGroup,
+  SelectInput,
+  TextAreaInput,
+  TextInput,
+} from '@/components/formInputs';
+import { CaptionJobConfig } from '@/types';
+import { defaultCaptionJobConfig, handleCaptionerTypeChange } from '@/helpers/captionJobConfig';
+import { objectCopy } from '@/utils/basic';
+import { useNestedState } from '@/utils/hooks';
+import {
+  captionerTypes,
+  defaultQtype,
+  groupedCaptionerTypes,
+  maxNewTokensOptions,
+  maxResOptions,
+  quantizationOptions,
+} from '@/helpers/captionOptions';
+import { isMac } from '@/helpers/basic';
+import useGPUInfo from '@/hooks/useGPUInfo';
+import { apiClient } from '@/utils/api';
+import { v4 as uuidv4 } from 'uuid';
+import { startJob } from '@/utils/jobs';
+import { startQueue } from '@/utils/queue';
+
+export interface CaptionDatasetModalState {
+  datasetPath: string;
+  onClose?: () => void;
+}
+
+export const captionDatasetModalState = createGlobalState<CaptionDatasetModalState | null>(null);
+
+export const openCaptionDatasetModal = (datasetPath: string, onClose?: () => void) => {
+  captionDatasetModalState.set({ datasetPath, onClose });
+};
+
+export const CaptionDatasetModal: React.FC = () => {
+  const [modalInfo, setModalInfo] = captionDatasetModalState.use();
+  const [jobConfig, setJobConfig] = useNestedState<CaptionJobConfig>(objectCopy(defaultCaptionJobConfig));
+  const [gpuIDs, setGpuIDs] = useState<string | null>(null);
+  const { gpuList, isGPUInfoLoaded } = useGPUInfo();
+  const open = modalInfo !== null;
+  const isSavingRef = useRef(false);
+  const showGPUSelect = !isMac();
+
+  useFromNull(() => {
+    // reset the state
+    setJobConfig(objectCopy(defaultCaptionJobConfig));
+    // set the path_to_caption
+    if (modalInfo?.datasetPath) {
+      setJobConfig(modalInfo.datasetPath, 'config.process[0].caption.path_to_caption');
+    }
+  }, [modalInfo]);
+
+  useEffect(() => {
+    if (isGPUInfoLoaded) {
+      if (gpuIDs === null && gpuList.length > 0) {
+        setGpuIDs(`${gpuList[0].index}`);
+      }
+    }
+  }, [gpuList, isGPUInfoLoaded]);
+
+  const handleClose = () => {
+    if (modalInfo?.onClose) {
+      modalInfo.onClose();
+    }
+    setModalInfo(null);
+  };
+
+  const selectedCaptionOption = captionerTypes.find(option => option.name === jobConfig.config.process[0].type);
+
+  const saveJob = async () => {
+    if (isSavingRef.current) return;
+    if (!modalInfo?.datasetPath) {
+      alert('Dataset path is missing. Please try again.');
+      return;
+    }
+    isSavingRef.current = true;
+
+    apiClient
+      .post('/api/jobs', {
+        id: null,
+        name: uuidv4(),
+        gpu_ids: gpuIDs,
+        job_config: jobConfig,
+        job_type: 'caption',
+        job_ref: modalInfo.datasetPath,
+      })
+      .then(async res => {
+        const jobId = res.data.id;
+        await startJob(jobId);
+        // start the queue as well
+        await startQueue(gpuIDs || '');
+        isSavingRef.current = false;
+        handleClose();
+      })
+      .catch(error => {
+        if (error.response?.status === 409) {
+          alert('A caption job for this dataset already exists. Please check your jobs list.');
+        } else {
+          alert('Failed to save job. Please try again.');
+        }
+        console.log('Error saving training:', error);
+        isSavingRef.current = false;
+      });
+  };
+
+  return (
+    <Modal isOpen={open} onClose={handleClose} title="Caption Dataset" size="lg">
+      <div className="space-y-4 text-gray-200">
+        <form
+          onSubmit={e => {
+            e.preventDefault();
+            saveJob();
+          }}
+        >
+          <div className="text-sm text-gray-400">
+            <div className="grid grid-cols-1 md:grid-cols-2 gap-4 mt-4">
+              <div>
+                <SelectInput
+                  label="Captioner Type"
+                  value={jobConfig.config.process[0].type}
+                  onChange={value => {
+                    handleCaptionerTypeChange(jobConfig.config.process[0].type, value, jobConfig, setJobConfig);
+                  }}
+                  options={groupedCaptionerTypes}
+                />
+              </div>
+              {showGPUSelect && (
+                <div>
+                  <SelectInput
+                    label="GPU ID"
+                    value={`${gpuIDs}`}
+                    onChange={value => setGpuIDs(value)}
+                    options={gpuList.map((gpu: any) => ({ value: `${gpu.index}`, label: `GPU #${gpu.index}` }))}
+                  />
+                </div>
+              )}
+            </div>
+            <div className="mt-4">
+              <CreatableSelectInput
+                label="Name or Path"
+                value={jobConfig.config.process[0].caption.model_name_or_path}
+                docKey="config.process[0].caption.model_name_or_path"
+                onChange={(value: string | null) => {
+                  if (value?.trim() === '') {
+                    value = null;
+                  }
+                  setJobConfig(value, 'config.process[0].caption.model_name_or_path');
+                }}
+                placeholder=""
+                options={selectedCaptionOption?.name_or_path_options || []}
+                required
+              />
+            </div>
+            {selectedCaptionOption?.additionalSections?.includes('caption.model_name_or_path2') && (
+              <div className="mt-4">
+                <CreatableSelectInput
+                  label="Name or Path 2"
+                  value={jobConfig.config.process[0].caption.model_name_or_path2 || ''}
+                  onChange={(value: string | null) => {
+                    if (value?.trim() === '') {
+                      value = null;
+                    }
+                    setJobConfig(value, 'config.process[0].caption.model_name_or_path2');
+                  }}
+                  placeholder=""
+                  options={selectedCaptionOption?.name_or_path2_options || []}
+                />
+              </div>
+            )}
+            <div className="grid grid-cols-1 md:grid-cols-2 gap-4 mt-4">
+              <div>
+                <SelectInput
+                  label="Quantize"
+                  value={jobConfig.config.process[0].caption.quantize ? jobConfig.config.process[0].caption.qtype : ''}
+                  onChange={value => {
+                    if (value === '') {
+                      setJobConfig(false, 'config.process[0].caption.quantize');
+                      value = defaultQtype;
+                    } else {
+                      setJobConfig(true, 'config.process[0].caption.quantize');
+                    }
+                    setJobConfig(value, 'config.process[0].caption.qtype');
+                  }}
+                  options={quantizationOptions}
+                />
+                {selectedCaptionOption?.additionalSections?.includes('caption.max_res') && (
+                  <div className="mt-4">
+                    <SelectInput
+                      label="Max Resolution"
+                      value={`${jobConfig.config.process[0].caption.max_res || ''}`}
+                      onChange={value => {
+                        const intVal = parseInt(value);
+                        if (!isNaN(intVal)) {
+                          setJobConfig(intVal, 'config.process[0].caption.max_res');
+                        }
+                      }}
+                      options={maxResOptions}
+                    />
+                  </div>
+                )}
+                {selectedCaptionOption?.additionalSections?.includes('caption.max_new_tokens') && (
+                  <div className="mt-4">
+                    <SelectInput
+                      label="Max New Tokens"
+                      value={`${jobConfig.config.process[0].caption.max_new_tokens || ''}`}
+                      onChange={value => {
+                        const intVal = parseInt(value);
+                        if (!isNaN(intVal)) {
+                          setJobConfig(intVal, 'config.process[0].caption.max_new_tokens');
+                        }
+                      }}
+                      options={maxNewTokensOptions}
+                    />
+                  </div>
+                )}
+              </div>
+              <div>
+                <FormGroup label="Options">
+                  <Checkbox
+                    label="Low VRAM"
+                    checked={jobConfig.config.process[0].caption.low_vram}
+                    onChange={value => setJobConfig(value, 'config.process[0].caption.low_vram')}
+                  />
+                  <Checkbox
+                    label="Recaption"
+                    checked={jobConfig.config.process[0].caption.recaption}
+                    onChange={value => setJobConfig(value, 'config.process[0].caption.recaption')}
+                  />
+                </FormGroup>
+              </div>
+            </div>
+            {selectedCaptionOption?.additionalSections?.includes('caption.caption_prompt') && (
+              <div className="mt-4">
+                <TextAreaInput
+                  label="Caption Prompt"
+                  value={jobConfig.config.process[0].caption.caption_prompt || ''}
+                  onChange={value => {
+                    setJobConfig(value, 'config.process[0].caption.caption_prompt');
+                  }}
+                  placeholder="Enter caption prompt"
+                />
+              </div>
+            )}
+          </div>
+
+          <div className="mt-6 flex justify-end space-x-3">
+            <button
+              type="button"
+              className="rounded-md bg-gray-700 px-4 py-2 text-gray-200 hover:bg-gray-600 focus:outline-none focus:ring-2 focus:ring-gray-500"
+              onClick={handleClose}
+            >
+              Cancel
+            </button>
+            <button
+              type="submit"
+              className="rounded-md bg-blue-600 px-4 py-2 text-white hover:bg-blue-700 focus:outline-none focus:ring-2 focus:ring-blue-500"
+            >
+              Add to Queue
+            </button>
+          </div>
+        </form>
+      </div>
+    </Modal>
+  );
+};

ui/src/components/DatasetImageCard.tsx

@@ -9,6 +9,7 @@ import { isVideo, isAudio } from '@/utils/basic';
 interface DatasetImageCardProps {
   imageUrl: string;
   alt: string;
+  isAutoCaptioning: boolean;
   children?: ReactNode;
   className?: string;
   onDelete?: () => void;
@@ -17,6 +18,7 @@ interface DatasetImageCardProps {
 const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
   imageUrl,
   alt,
+  isAutoCaptioning,
   children,
   className = '',
   onDelete = () => {},
@@ -80,6 +82,16 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
     }
   }, [inViewport, isVisible]);
 
+  // Poll for caption updates every 5 seconds while auto-captioning
+  useEffect(() => {
+    if (!isAutoCaptioning || !inViewport || !isVisible) return;
+    const interval = setInterval(() => {
+      // Reset so fetchCaption will re-fetch
+      setIsCaptionLoaded(false);
+    }, 5000);
+    return () => clearInterval(interval);
+  }, [isAutoCaptioning, inViewport, isVisible]);
+
   useEffect(() => {
     // Create intersection observer to check viewport visibility
     const observer = new IntersectionObserver(
@@ -129,6 +141,8 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
 
   const isCaptionCurrent = caption.trim() === savedCaption;
 
+  const [showAudioPlayer, setShowAudioPlayer] = useState(true);
+
   const isItAVideo = isVideo(imageUrl);
   const isItAudio = isAudio(imageUrl);
   const isItImage = !isItAVideo && !isItAudio;
@@ -154,7 +168,22 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
                   controls
                 />
               )}
-              {isItAudio && (
+              {isItAudio && !showAudioPlayer && (
+                <div
+                  className="w-full h-full cursor-pointer flex items-center justify-center bg-gray-900"
+                  onClick={() => setShowAudioPlayer(true)}
+                >
+                  <img
+                    src={`/api/audio/art/${encodeURIComponent(imageUrl)}`}
+                    alt={alt}
+                    className="w-full h-full object-contain"
+                    onError={e => {
+                      (e.target as HTMLImageElement).style.display = 'none';
+                    }}
+                  />
+                </div>
+              )}
+              {isItAudio && showAudioPlayer && (
                 <AudioPlayer
                   src={`/api/img/${encodeURIComponent(imageUrl)}`}
                   title={imageUrl.replace(/^.*[\\/]/, '')}
@@ -221,9 +250,12 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
             onBlur={saveCaption}
           >
             <textarea
-              className="w-full bg-transparent resize-none outline-none focus:ring-0 focus:outline-none"
+              className={classNames("w-full bg-transparent resize-none outline-none focus:ring-0 focus:outline-none", {
+                'opacity-50 cursor-not-allowed': isAutoCaptioning,
+              })}
               value={caption}
               rows={3}
+              readOnly={isAutoCaptioning}
               onChange={e => setCaption(e.target.value)}
               onKeyDown={handleKeyDown}
             />

ui/src/components/JobActionBar.tsx

@@ -123,11 +123,16 @@ export default function JobActionBar({
           <Cog />
         </MenuButton>
         <MenuItems anchor="bottom" className="bg-gray-900 border border-gray-700 rounded shadow-lg w-48 px-2 py-2 mt-4">
-          <MenuItem>
-            <Link href={`/jobs/new?cloneId=${job.id}`} className="cursor-pointer px-4 py-1 hover:bg-gray-800 rounded block">
-              Clone Job
-            </Link>
-          </MenuItem>
+          {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"
+              >
+                Clone Job
+              </Link>
+            </MenuItem>
+          )}
           <MenuItem>
             <div
               className="cursor-pointer px-4 py-1 hover:bg-gray-800 rounded"

ui/src/components/JobConfigViewer.tsx

@@ -21,11 +21,23 @@ const yamlConfig: YAML.DocumentOptions &
   directives: true,
 };
 
+function toYaml(obj: any): string {
+  const doc = new YAML.Document(obj, yamlConfig);
+  YAML.visit(doc, {
+    Scalar(_key, node) {
+      if (typeof node.value === 'string' && node.value.includes('\n')) {
+        node.type = YAML.Scalar.BLOCK_LITERAL;
+      }
+    },
+  });
+  return doc.toString(yamlConfig);
+}
+
 export default function JobConfigViewer({ job }: Props) {
   const [editorValue, setEditorValue] = useState<string>('');
   useEffect(() => {
     if (job?.job_config) {
-      const yamlContent = YAML.stringify(JSON.parse(job.job_config), yamlConfig);
+      const yamlContent = toYaml(JSON.parse(job.job_config));
       setEditorValue(yamlContent);
     }
   }, [job]);

ui/src/components/JobOverview.tsx

@@ -82,6 +82,8 @@ export default function JobOverview({ job }: JobOverviewProps) {
     }
   };
 
+  const jobType = job?.job_type || 'unknown';
+
   let status = job.status;
   if (isStopping) {
     status = 'stopping';
@@ -100,17 +102,19 @@ export default function JobOverview({ job }: JobOverviewProps) {
 
         <div className="p-4 space-y-6 flex flex-col flex-grow">
           {/* Progress Bar */}
-          <div className="space-y-2">
-            <div className="flex items-center justify-between text-sm">
-              <span className="text-gray-400">Progress</span>
-              <span className="text-gray-200">
-                Step {job.step} of {totalSteps}
-              </span>
+          {job.job_type === 'train' && (
+            <div className="space-y-2">
+              <div className="flex items-center justify-between text-sm">
+                <span className="text-gray-400">Progress</span>
+                <span className="text-gray-200">
+                  Step {job.step} of {totalSteps}
+                </span>
+              </div>
+              <div className="w-full bg-gray-800 rounded-full h-2">
+                <div className="h-2 rounded-full bg-blue-500 transition-all" style={{ width: `${progress}%` }} />
+              </div>
             </div>
-            <div className="w-full bg-gray-800 rounded-full h-2">
-              <div className="h-2 rounded-full bg-blue-500 transition-all" style={{ width: `${progress}%` }} />
-            </div>
-          </div>
+          )}
 
           {/* Job Info Grid */}
           <div className="grid gap-4 grid-cols-1 md:grid-cols-3">
@@ -164,9 +168,11 @@ export default function JobOverview({ job }: JobOverviewProps) {
       <div className="col-span-1">
         <div>{isCPUInfoLoaded && cpuInfo && <CPUWidget cpu={cpuInfo} />}</div>
         <div className="mt-4">{isGPUInfoLoaded && gpuList.length > 0 && <GPUWidget gpu={gpuList[0]} />}</div>
-        <div className="mt-4">
-          <FilesWidget jobID={job.id} />
-        </div>
+        {jobType === 'train' && (
+          <div className="mt-4">
+            <FilesWidget jobID={job.id} />
+          </div>
+        )}
       </div>
     </div>
   );

ui/src/components/JobsTable.tsx

@@ -14,10 +14,11 @@ import useGPUInfo from '@/hooks/useGPUInfo';
 interface JobsTableProps {
   autoStartQueue?: boolean;
   onlyActive?: boolean;
+  job_type?: string | null;
 }
 
-export default function JobsTable({ onlyActive = false }: JobsTableProps) {
-  const { jobs, status, refreshJobs } = useJobsList(onlyActive, 5000);
+export default function JobsTable({ onlyActive = false, job_type = null }: JobsTableProps) {
+  const { jobs, status, refreshJobs } = useJobsList({ onlyActive, reloadInterval: 5000, job_type });
   const { queues, status: queueStatus, refreshQueues } = useQueueList();
   const { gpuList, isGPUInfoLoaded } = useGPUInfo();
 
@@ -30,21 +31,37 @@ export default function JobsTable({ onlyActive = false }: JobsTableProps) {
     {
       title: 'Name',
       key: 'name',
-      render: row => (
-        <Link href={`/jobs/${row.id}`} className="font-medium whitespace-nowrap">
-          {['running', 'stopping'].includes(row.status) ? (
-            <CgSpinner className="inline animate-spin mr-2 text-blue-400" />
-          ) : null}
-          {row.name}
-        </Link>
-      ),
+      render: row => {
+        let title = row.name;
+        // if (row.job_type === 'train') title = `Train: ${title}`;
+        if (row.job_type === 'caption') {
+          let splits = row.job_ref.split(/[/\\]/);
+          const datasetPath = `${splits[splits.length - 1]}`;
+          title = (
+            <>
+              <small className="opacity-50">CAPTION: </small> {datasetPath}
+            </>
+          );
+        }
+        return (
+          <Link href={`/jobs/${row.id}`} className="font-medium whitespace-nowrap">
+            {['running', 'stopping'].includes(row.status) ? (
+              <CgSpinner className="inline animate-spin mr-2 text-blue-400" />
+            ) : null}
+            {title}
+          </Link>
+        );
+      },
     },
     {
       title: 'Steps',
       key: 'steps',
       render: row => {
         const jobConfig: JobConfig = JSON.parse(row.job_config);
-        const totalSteps = jobConfig.config.process[0].train.steps;
+        if (row.job_type !== 'train') {
+          return <></>;
+        }
+        const totalSteps = jobConfig.config.process[0].train?.steps;
 
         return (
           <div>

ui/src/components/SampleImageCard.tsx

@@ -1,5 +1,5 @@
 import React, { useRef, useEffect, useState, ReactNode } from 'react';
-import { isVideo } from '@/utils/basic';
+import { isVideo, isAudio } from '@/utils/basic';
 
 interface SampleImageCardProps {
   imageUrl: string;
@@ -63,7 +63,18 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
       <div ref={cardRef} className="relative w-full cursor-pointer" style={{ paddingBottom: '100%' }} onClick={onClick}>
         <div className="absolute inset-0 rounded-t-lg shadow-md">
           {isVisible ? (
-            isVideo(imageUrl) ? (
+            isAudio(imageUrl) ? (
+              <div className="w-full h-full flex items-center justify-center bg-gray-900">
+                <img
+                  src={`/api/audio/art/${encodeURIComponent(imageUrl)}`}
+                  alt={alt}
+                  className="w-full h-full object-cover"
+                  onError={e => {
+                    (e.target as HTMLImageElement).style.display = 'none';
+                  }}
+                />
+              </div>
+            ) : isVideo(imageUrl) ? (
               <video
                 ref={videoRef}
                 src={`/api/img/${encodeURIComponent(imageUrl)}`}

ui/src/components/SampleImageViewer.tsx

@@ -7,7 +7,8 @@ import { Cog } from 'lucide-react';
 import { Menu, MenuButton, MenuItem, MenuItems } from '@headlessui/react';
 import { openConfirm } from './ConfirmModal';
 import { apiClient } from '@/utils/api';
-import { isVideo } from '@/utils/basic';
+import { isVideo, isAudio } from '@/utils/basic';
+import AudioPlayer from './AudioPlayer';
 
 interface Props {
   imgPath: string | null; // current image path
@@ -214,7 +215,15 @@ export default function SampleImageViewer({
           >
             <div className="overflow-hidden flex items-center justify-center">
               {displayedImgPath &&
-                (isVideo(displayedImgPath) ? (
+                (isAudio(displayedImgPath) ? (
+                  <div className="w-[500px] h-[500px] max-w-[95vw] max-h-[82vh]">
+                    <AudioPlayer
+                      src={`/api/img/${encodeURIComponent(displayedImgPath)}`}
+                      title={displayedImgPath.replace(/^.*[\\/]/, '')}
+                      autoPlay
+                    />
+                  </div>
+                ) : isVideo(displayedImgPath) ? (
                   <video
                     src={`/api/img/${encodeURIComponent(displayedImgPath)}`}
                     className="w-auto h-auto max-w-[95vw] max-h-[82vh] object-contain"
@@ -282,7 +291,7 @@ export default function SampleImageViewer({
                 </div>
               </div>
             </div>
-            <div className="absolute top-2 right-2 bg-gray-900 rounded-full p-1 leading-[0px] opacity-50 hover:opacity-100">
+            <div className="absolute top-2 right-2 bg-gray-900 rounded-full p-1 leading-[0px] opacity-50 hover:opacity-100 z-20">
               <Menu>
                 <MenuButton>
                   <Cog />

ui/src/components/Sidebar.tsx

@@ -8,7 +8,7 @@ const Sidebar = () => {
   const navigation = [
     { name: 'Dashboard', href: '/dashboard', icon: Home },
     { name: 'New Job', href: '/jobs/new', icon: Plus },
-    { name: 'Training Queue', href: '/jobs', icon: BrainCircuit },
+    { name: 'Queue', href: '/jobs', icon: BrainCircuit },
     { name: 'Datasets', href: '/datasets', icon: Images },
     { name: 'Settings', href: '/settings', icon: Settings },
   ];

ui/src/components/formInputs.tsx

@@ -67,6 +67,49 @@ export const TextInput = forwardRef<HTMLInputElement, TextInputProps>((props: Te
 // 👇 Helpful for debugging
 TextInput.displayName = 'TextInput';
 
+export interface TextAreaInputProps extends InputProps {
+  value: string;
+  onChange: (value: string) => void;
+  disabled?: boolean;
+  rows?: number;
+}
+
+export const TextAreaInput = forwardRef<HTMLTextAreaElement, TextAreaInputProps>((props: TextAreaInputProps, ref) => {
+  const { label, value, onChange, placeholder, required, disabled, rows = 4, className, docKey = null } = props;
+  let { doc } = props;
+  if (!doc && docKey) {
+    doc = getDoc(docKey);
+  }
+  return (
+    <div className={classNames(className)}>
+      {label && (
+        <label className={labelClasses}>
+          {label}{' '}
+          {doc && (
+            <div className="inline-block ml-1 text-xs text-gray-500 cursor-pointer" onClick={() => openDoc(doc)}>
+              <CircleHelp className="inline-block w-4 h-4 cursor-pointer" />
+            </div>
+          )}
+        </label>
+      )}
+      <textarea
+        ref={ref}
+        value={value}
+        onChange={e => {
+          if (!disabled) onChange(e.target.value);
+        }}
+        className={`${inputClasses} ${disabled ? 'opacity-30 cursor-not-allowed' : ''}`}
+        placeholder={placeholder}
+        required={required}
+        disabled={disabled}
+        rows={rows}
+      />
+    </div>
+  );
+});
+
+TextAreaInput.displayName = 'TextAreaInput';
+
 export interface NumberInputProps extends InputProps {
   value: number | null;
   onChange: (value: number | null) => void;
@@ -200,6 +243,115 @@ export const SelectInput = (props: SelectInputProps) => {
   );
 };
 
+export interface CreatableSelectInputProps extends InputProps {
+  value: string;
+  disabled?: boolean;
+  onChange: (value: string) => void;
+  options: GroupedSelectOption[] | SelectOption[];
+}
+
+const CUSTOM_SELECT_VALUE = '__custom__';
+
+export const CreatableSelectInput = (props: CreatableSelectInputProps) => {
+  const { label, value, onChange, options, docKey = null } = props;
+  let { doc } = props;
+  if (!doc && docKey) {
+    doc = getDoc(docKey);
+  }
+
+  // Check if current value matches any predefined option
+  let isInOptions = false;
+  if (options && options.length > 0) {
+    if ('options' in options[0]) {
+      isInOptions = (options as GroupedSelectOption[]).flatMap(g => g.options).some(opt => opt.value === value);
+    } else {
+      isInOptions = (options as SelectOption[]).some(opt => opt.value === value);
+    }
+  }
+
+  const [isCustom, setIsCustom] = React.useState(!isInOptions && !!value);
+
+  // Build select options with "Custom" at the top
+  const customOption: SelectOption = { value: CUSTOM_SELECT_VALUE, label: 'Custom' };
+  const selectOptions = React.useMemo(() => {
+    if (options && options.length > 0 && 'options' in options[0]) {
+      return [{ label: '', options: [customOption] }, ...(options as GroupedSelectOption[])];
+    }
+    return [customOption, ...(options as SelectOption[])];
+  }, [options]);
+
+  const selectedOption = isCustom
+    ? customOption
+    : (() => {
+        if (!options || options.length === 0) return undefined;
+        if ('options' in options[0]) {
+          return (options as GroupedSelectOption[]).flatMap(g => g.options).find(opt => opt.value === value);
+        }
+        return (options as SelectOption[]).find(opt => opt.value === value);
+      })();
+
+  return (
+    <div
+      className={classNames(props.className, {
+        'opacity-30 cursor-not-allowed': props.disabled,
+      })}
+    >
+      {label && (
+        <label className={labelClasses}>
+          {label}{' '}
+          {doc && (
+            <div className="inline-block ml-1 text-xs text-gray-500 cursor-pointer" onClick={() => openDoc(doc)}>
+              <CircleHelp className="inline-block w-4 h-4 cursor-pointer" />
+            </div>
+          )}
+        </label>
+      )}
+      <div className="flex gap-2">
+        <div className={isCustom ? 'w-1/3' : 'w-full'}>
+          <Select
+            value={selectedOption}
+            options={selectOptions}
+            isDisabled={props.disabled}
+            className="aitk-react-select-container"
+            classNamePrefix="aitk-react-select"
+            formatOptionLabel={(option: unknown) => {
+              const opt = option as SelectOption;
+              return opt.value === CUSTOM_SELECT_VALUE ? (
+                <span className="opacity-50 italic">~ Custom ~</span>
+              ) : (
+                opt.label
+              );
+            }}
+            onChange={selected => {
+              if (selected) {
+                const val = (selected as { value: string }).value;
+                if (val === CUSTOM_SELECT_VALUE) {
+                  setIsCustom(true);
+                  onChange('');
+                } else {
+                  setIsCustom(false);
+                  onChange(val);
+                }
+              }
+            }}
+          />
+        </div>
+        {isCustom && (
+          <input
+            type="text"
+            value={value}
+            onChange={e => onChange(e.target.value)}
+            className={`${inputClasses} w-2/3`}
+            placeholder={props.placeholder ?? 'Enter custom value'}
+            disabled={props.disabled}
+            autoFocus
+          />
+        )}
+      </div>
+    </div>
+  );
+};
+
 export interface CheckboxProps {
   label?: string | React.ReactNode;
   checked: boolean;