Added ability to set the caption extention in dataset viewer, captioner, and trainer so one dataset can have multiple caption styles in different files with different extensions. Added dataset caption template for a blank ideogram 4 formatted template.

2026-06-06 08:32:24 -06:00
parent 10cdeb394e
commit 41157b460c
19 changed files with 219 additions and 270 deletions
--- a/ui/src/app/api/caption/get/route.ts
+++ b/ui/src/app/api/caption/get/route.ts
@@ -22,15 +22,16 @@ export async function POST(request: NextRequest) {
    return new NextResponse(null, { status: 499 });
  }

-  const { imgPath } = body;
+  const { imgPath, ext } = body;
  console.log('Received POST request for caption:', imgPath);
  try {
    // Decode the path
    const filepath = imgPath;
    console.log('Decoded image path:', filepath);

-    // caption name is the filepath without extension but with .txt
-    const captionPath = filepath.replace(/\.[^/.]+$/, '') + '.txt';
+    // caption name is the filepath without extension but with the caption extension (default txt)
+    const captionExt = ((ext || 'txt') as string).replace(/^\.+/, '').trim() || 'txt';
+    const captionPath = filepath.replace(/\.[^/.]+$/, '') + '.' + captionExt;

    // Get allowed directories
    const allowedDir = await getDatasetsRoot();
--- a/ui/src/app/api/caption/getBatch/route.ts
+++ b/ui/src/app/api/caption/getBatch/route.ts
@@ -21,11 +21,12 @@ export async function POST(request: NextRequest) {
    return new NextResponse(null, { status: 499 });
  }

-  const { imgPaths } = body as { imgPaths?: string[] };
+  const { imgPaths, ext } = body as { imgPaths?: string[]; ext?: string };
  if (!Array.isArray(imgPaths)) {
    return NextResponse.json({ error: 'imgPaths must be an array' }, { status: 400 });
  }

+  const captionExt = ((ext || 'txt') as string).replace(/^\.+/, '').trim() || 'txt';
  const allowedDir = await getDatasetsRoot();
  const captions: Record<string, string> = {};

@@ -33,7 +34,7 @@ export async function POST(request: NextRequest) {
    if (typeof imgPath !== 'string') continue;
    if (!isUnderRoot(imgPath, allowedDir)) continue;

-    const captionPath = imgPath.replace(/\.[^/.]+$/, '') + '.txt';
+    const captionPath = imgPath.replace(/\.[^/.]+$/, '') + '.' + captionExt;
    try {
      captions[imgPath] = fs.existsSync(captionPath) ? fs.readFileSync(captionPath, 'utf-8') : '';
    } catch {
--- a/ui/src/app/api/img/caption/route.ts
+++ b/ui/src/app/api/img/caption/route.ts
@@ -5,7 +5,7 @@ import { getDatasetsRoot } from '@/server/settings';
 export async function POST(request: Request) {
  try {
    const body = await request.json();
-    const { imgPath, caption } = body;
+    const { imgPath, caption, ext } = body;
    let datasetsPath = await getDatasetsRoot();
    // make sure the dataset path is in the image path
    if (!imgPath.startsWith(datasetsPath)) {
@@ -17,8 +17,9 @@ export async function POST(request: Request) {
      return NextResponse.json({ error: 'Image does not exist' }, { status: 404 });
    }

-    // check for caption
-    const captionPath = imgPath.replace(/\.[^/.]+$/, '') + '.txt';
+    // check for caption (default extension txt)
+    const captionExt = ((ext || 'txt') as string).replace(/^\.+/, '').trim() || 'txt';
+    const captionPath = imgPath.replace(/\.[^/.]+$/, '') + '.' + captionExt;
    // save caption to file
    fs.writeFileSync(captionPath, caption);

--- a/ui/src/app/datasets/[datasetName]/page.tsx
+++ b/ui/src/app/datasets/[datasetName]/page.tsx
@@ -13,6 +13,7 @@ import { apiClient } from '@/utils/api';
 import useSettings from '@/hooks/useSettings';
 import { pathJoin } from '@/utils/basic';
 import AutoCaptionButton from '@/components/AutoCaptionButton';
+import { CreatableSelectInput } from '@/components/formInputs';

 export default function DatasetPage({ params }: { params: { datasetName: string } }) {
  const [imgList, setImgList] = useState<{ img_path: string }[]>([]);
@@ -22,6 +23,7 @@ export default function DatasetPage({ params }: { params: { datasetName: string
  const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle');
  const { settings, isSettingsLoaded } = useSettings();
  const [selectedImgPath, setSelectedImgPath] = useState<string | null>(null);
+  const [captionExt, setCaptionExt] = useState<string>('txt');
  const [captionRefreshKeys, setCaptionRefreshKeys] = useState<Record<string, number>>({});
  const [scrollParent, setScrollParent] = useState<HTMLDivElement | null>(null);
  const scrollParentCallback = useCallback((el: HTMLDivElement | null) => setScrollParent(el), []);
@@ -118,9 +120,23 @@ export default function DatasetPage({ params }: { params: { datasetName: string
        </div>
        <div className="flex-1"></div>
        <div className="flex-shrink-0 flex items-center gap-1 sm:gap-2">
+          <div className="flex items-center gap-1">
+            <label className="text-xs text-gray-400 hidden sm:inline whitespace-nowrap">Caption ext</label>
+            <CreatableSelectInput
+              className="w-44"
+              value={captionExt}
+              onChange={value => setCaptionExt(value)}
+              options={[
+                { value: 'txt', label: 'txt' },
+                { value: 'json', label: 'json' },
+                { value: 'caption', label: 'caption' },
+              ]}
+            />
+          </div>
          <AutoCaptionButton
            datasetPath={`${pathJoin(settings.DATASETS_FOLDER, datasetName)}`}
            setIsAutoCaptioning={setIsAutoCaptioning}
+            captionExt={captionExt}
          />
          <Button
            className="text-white bg-slate-600 px-2 sm:px-3 py-1 rounded-md text-sm sm:text-base whitespace-nowrap"
@@ -151,6 +167,7 @@ export default function DatasetPage({ params }: { params: { datasetName: string
                  onImageClick={() => setSelectedImgPath(img.img_path)}
                  captionRefreshKey={captionRefreshKeys[img.img_path] || 0}
                  observerRoot={scrollParent}
+                  captionExt={captionExt}
                />
              );
            }}
@@ -165,6 +182,7 @@ export default function DatasetPage({ params }: { params: { datasetName: string
        onChange={setSelectedImgPath}
        refreshImages={() => refreshImageList(datasetName)}
        onCaptionSaved={path => setCaptionRefreshKeys(prev => ({ ...prev, [path]: (prev[path] || 0) + 1 }))}
+        captionExt={captionExt}
      />
    </>
  );
--- a/ui/src/app/jobs/new/SimpleJob.tsx
+++ b/ui/src/app/jobs/new/SimpleJob.tsx
@@ -20,6 +20,7 @@ import {
  FormGroup,
  NumberInput,
  SliderInput,
+  CreatableSelectInput,
 } from '@/components/formInputs';
 import Card from '@/components/Card';
 import { X, Copy, Wand2, SquareDashed } from 'lucide-react';
@@ -959,6 +960,18 @@ export default function SimpleJob({
                        min={0}
                        required
                      />
+                      <CreatableSelectInput
+                        label="Caption Extension"
+                        className="pt-2"
+                        value={dataset.caption_ext || 'txt'}
+                        onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].caption_ext`)}
+                        options={[
+                          { value: 'txt', label: 'txt' },
+                          { value: 'json', label: 'json' },
+                          { value: 'caption', label: 'caption' },
+                        ]}
+                      />
+
                      {modelArch?.additionalSections?.includes('datasets.num_frames') && !dataset.auto_frame_count && (
                        <NumberInput
                          label="Num Frames"
--- a/ui/src/components/AutoCaptionButton.tsx
+++ b/ui/src/components/AutoCaptionButton.tsx
@@ -8,9 +8,10 @@ import { Loader2 } from 'lucide-react';
 type AutoCaptionButtonProps = {
  datasetPath: string;
  setIsAutoCaptioning?: (isAutoCaptioning: boolean) => void;
+  captionExt?: string;
 };

-export default function AutoCaptionButton({ datasetPath, setIsAutoCaptioning }: AutoCaptionButtonProps) {
+export default function AutoCaptionButton({ datasetPath, setIsAutoCaptioning, captionExt }: AutoCaptionButtonProps) {
  const { job, status, refreshJob } = useJobByRef(datasetPath, 5000);
  useEffect(() => {
    if (setIsAutoCaptioning) {
@@ -34,9 +35,13 @@ export default function AutoCaptionButton({ datasetPath, setIsAutoCaptioning }:
    <Button
      className="text-white bg-blue-600 px-2 sm:px-3 py-1 rounded-md mr-1 sm:mr-2 text-sm sm:text-base whitespace-nowrap"
      onClick={() =>
-        openCaptionDatasetModal(datasetPath, () => {
-          refreshJob();
-        })
+        openCaptionDatasetModal(
+          datasetPath,
+          () => {
+            refreshJob();
+          },
+          { defaultCaptionExt: captionExt },
+        )
      }
    >
      <span className="hidden sm:inline">Auto Caption</span>
--- a/ui/src/components/CaptionDatasetModal.tsx
+++ b/ui/src/components/CaptionDatasetModal.tsx
@@ -22,6 +22,7 @@ export interface CaptionDatasetModalState {
  datasetPath: string;
  jobId?: string | null;
  cloneId?: string | null;
+  defaultCaptionExt?: string;
  onClose?: () => void;
 }

@@ -30,13 +31,14 @@ export const captionDatasetModalState = createGlobalState<CaptionDatasetModalSta
 export const openCaptionDatasetModal = (
  datasetPath: string,
  onClose?: () => void,
-  options?: { jobId?: string | null; cloneId?: string | null },
+  options?: { jobId?: string | null; cloneId?: string | null; defaultCaptionExt?: string },
 ) => {
  captionDatasetModalState.set({
    datasetPath,
    onClose,
    jobId: options?.jobId ?? null,
    cloneId: options?.cloneId ?? null,
+    defaultCaptionExt: options?.defaultCaptionExt,
  });
 };

@@ -64,6 +66,11 @@ export const CaptionDatasetModal: React.FC = () => {
    if (modalInfo?.datasetPath) {
      setJobConfig(modalInfo.datasetPath, 'config.process[0].caption.path_to_caption');
    }
+    // default the caption extension to the current header selection. Editing it
+    // here only changes this job, not the header (separate state).
+    if (modalInfo?.defaultCaptionExt) {
+      setJobConfig(modalInfo.defaultCaptionExt, 'config.process[0].caption.caption_extension');
+    }
  }, [modalInfo]);

  // clone existing caption job
--- a/ui/src/components/CaptionSimpleJob.tsx
+++ b/ui/src/components/CaptionSimpleJob.tsx
@@ -121,6 +121,20 @@ const CaptionSimpleJob: React.FC<Props> = ({ jobConfig, setJobConfig, gpuIDs, se
            }}
            options={quantizationOptions}
          />
+          <div className="mt-4">
+            <CreatableSelectInput
+              label="Caption Extension"
+              value={jobConfig.config.process[0].caption.caption_extension || 'txt'}
+              onChange={value => {
+                setJobConfig(value, 'config.process[0].caption.caption_extension');
+              }}
+              options={[
+                { value: 'txt', label: 'txt' },
+                { value: 'json', label: 'json' },
+                { value: 'caption', label: 'caption' },
+              ]}
+            />
+          </div>
          {additionalSections.includes('caption.max_res') && (
            <div className="mt-4">
              <SelectInput
--- a/ui/src/components/DatasetImageCard.tsx
+++ b/ui/src/components/DatasetImageCard.tsx
@@ -18,6 +18,7 @@ interface DatasetImageCardProps {
  captionRefreshKey?: number;
  observerRoot?: Element | null;
  rootMargin?: string;
+  captionExt?: string;
 }

 const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
@@ -31,6 +32,7 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
  captionRefreshKey = 0,
  observerRoot = null,
  rootMargin = '200px 0px',
+  captionExt = 'txt',
 }) => {
  const [loaded, setLoaded] = useState<boolean>(false);
  const [showAudioPlayer, setShowAudioPlayer] = useState(true);
@@ -110,6 +112,7 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
  const { caption: fetchedCaption, isLoaded: isCaptionLoaded } = useCaptionBatch(
    isVisible ? imageUrl : null,
    combinedRefreshKey,
+    captionExt,
  );

  const [caption, setCaption] = useState<string>('');
@@ -138,10 +141,10 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
      return;
    }
    apiClient
-      .post('/api/img/caption', { imgPath: imageUrl, caption: trimmedCaption })
+      .post('/api/img/caption', { imgPath: imageUrl, caption: trimmedCaption, ext: captionExt })
      .then(() => {
        setSavedCaption(trimmedCaption);
-        setCachedCaption(imageUrl, trimmedCaption);
+        setCachedCaption(imageUrl, trimmedCaption, captionExt);
        dirtyRef.current = false;
      })
      .catch(error => {
@@ -150,19 +153,19 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
  };

  // Save any pending edit if the card unmounts (e.g. scrolled out of the virtualized window).
-  const latestRef = useRef({ caption, savedCaption, imageUrl });
+  const latestRef = useRef({ caption, savedCaption, imageUrl, captionExt });
  useEffect(() => {
-    latestRef.current = { caption, savedCaption, imageUrl };
+    latestRef.current = { caption, savedCaption, imageUrl, captionExt };
  });
  useEffect(() => {
    return () => {
      if (!dirtyRef.current) return;
-      const { caption: c, savedCaption: s, imageUrl: url } = latestRef.current;
+      const { caption: c, savedCaption: s, imageUrl: url, captionExt: ext } = latestRef.current;
      const trimmed = c.trim();
      if (trimmed === s) return;
      apiClient
-        .post('/api/img/caption', { imgPath: url, caption: trimmed })
-        .then(() => setCachedCaption(url, trimmed))
+        .post('/api/img/caption', { imgPath: url, caption: trimmed, ext })
+        .then(() => setCachedCaption(url, trimmed, ext))
        .catch(err => console.error('Error saving caption on unmount:', err));
    };
  }, []);
--- a/ui/src/components/DatasetImageViewer.tsx
+++ b/ui/src/components/DatasetImageViewer.tsx
@@ -12,6 +12,7 @@ import AudioPlayer from './AudioPlayer';
 import { TransformWrapper, TransformComponent } from 'react-zoom-pan-pinch';
 import { BoundingBoxEditor, parseBoundingBoxes, extractBoxes } from './BoundingBoxOverlay';
 import IdeogramCaptionSidebar, { isIdeogramCaption } from './IdeogramCaptionSidebar';
+import datasetTemplates from '@/helpers/datasetTemplates';

 function safeParse(text: string): any {
  try {
@@ -27,9 +28,17 @@ interface Props {
  onChange: (nextPath: string | null) => void; // parent setter
  refreshImages?: () => void;
  onCaptionSaved?: (imgPath: string, caption: string) => void;
+  captionExt?: string;
 }

-export default function DatasetImageViewer({ imgPath, imageList, onChange, refreshImages, onCaptionSaved }: Props) {
+export default function DatasetImageViewer({
+  imgPath,
+  imageList,
+  onChange,
+  refreshImages,
+  onCaptionSaved,
+  captionExt = 'txt',
+}: Props) {
  const [mounted, setMounted] = useState(false);
  const [isOpen, setIsOpen] = useState(Boolean(imgPath));
  const [caption, setCaption] = useState<string>('');
@@ -98,7 +107,7 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre
      const trimmed = value.trim();
      if (trimmed === prevSaved.trim()) return;
      apiClient
-        .post('/api/img/caption', { imgPath: path, caption: trimmed })
+        .post('/api/img/caption', { imgPath: path, caption: trimmed, ext: captionExt })
        .then(() => {
          if (currentImgPathRef.current === path) {
            setSavedCaption(trimmed);
@@ -109,7 +118,7 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre
          console.error('Error saving caption:', error);
        });
    },
-    [onCaptionSaved],
+    [onCaptionSaved, captionExt],
  );

  // Stable handle to the latest saveCaptionForPath so the fetch effect doesn't
@@ -151,7 +160,11 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre
      // transformResponse identity: keep the caption as a raw string. Axios's
      // default parses any JSON-looking body into an object (our bbox captions
      // are JSON), which would render as "[object Object]".
-      .post('/api/caption/get', { imgPath }, { signal: controller.signal, transformResponse: [d => d] })
+      .post(
+        '/api/caption/get',
+        { imgPath, ext: captionExt },
+        { signal: controller.signal, transformResponse: [d => d] },
+      )
      .then(res => res.data)
      .then(data => {
        if (controller.signal.aborted) return;
@@ -169,7 +182,7 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre
    return () => {
      controller.abort();
    };
-  }, [imgPath]);
+  }, [imgPath, captionExt]);

  // Save any pending caption when the viewer fully unmounts
  useEffect(() => {
@@ -509,6 +522,23 @@ export default function DatasetImageViewer({ imgPath, imageList, onChange, refre
                  {currentIndex >= 0 ? `${currentIndex + 1} / ${imageList.length}` : ''}
                </div>
              </div>
+              {isCaptionLoaded && caption.trim() === '' && (
+                <select
+                  className="w-full bg-gray-900 border border-gray-700 text-gray-100 text-sm rounded p-2 outline-none focus:ring-0 focus:outline-none"
+                  value=""
+                  onChange={e => {
+                    const template = datasetTemplates[e.target.value];
+                    if (template) setCaption(template.trim());
+                  }}
+                >
+                  <option value="">Templates...</option>
+                  {Object.keys(datasetTemplates).map(key => (
+                    <option key={key} value={key}>
+                      {key}
+                    </option>
+                  ))}
+                </select>
+              )}
              {isIdeogram ? (
                <IdeogramCaptionSidebar
                  caption={caption}
--- a/ui/src/components/formInputs.tsx
+++ b/ui/src/components/formInputs.tsx
@@ -363,13 +363,15 @@ export const CreatableSelectInput = (props: CreatableSelectInputProps) => {
        </label>
      )}
      <div className="flex gap-2">
-        <div className={isCustom ? 'w-1/3' : 'w-full'}>
+        <div className={isCustom ? 'w-20 shrink-0' : 'w-full'}>
          <Select
            value={selectedOption}
            options={selectOptions}
            isDisabled={props.disabled}
            className="aitk-react-select-container"
            classNamePrefix="aitk-react-select"
+            menuPosition="fixed"
+            menuPlacement="auto"
            formatOptionLabel={(option: unknown) => {
              const opt = option as SelectOption;
              return opt.value === CUSTOM_SELECT_VALUE ? (
@@ -397,7 +399,7 @@ export const CreatableSelectInput = (props: CreatableSelectInputProps) => {
            type="text"
            value={value}
            onChange={e => onChange(e.target.value)}
-            className={`${inputClasses} w-2/3`}
+            className={`${inputClasses} flex-1 min-w-0`}
            placeholder={props.placeholder ?? 'Enter custom value'}
            disabled={props.disabled}
            autoFocus
--- a/ui/src/helpers/captionJobConfig.ts
+++ b/ui/src/helpers/captionJobConfig.ts
@@ -21,6 +21,7 @@ export const defaultCaptionJobConfig: CaptionJobConfig = {
          extensions: ['mp3', 'wav', 'flac', 'ogg'],
          path_to_caption: '',
          recaption: false,
+          caption_extension: 'txt',
        },
      },
    ],
--- a/ui/src/helpers/datasetTemplates.ts
+++ b/ui/src/helpers/datasetTemplates.ts
@@ -0,0 +1,21 @@
+const datasetTemplates: { [key: string]: string } = {
+    ideogram4: `
+{
+  "high_level_description": "",
+  "style_description": {
+    "aesthetics": "",
+    "lighting": "",
+    "photo": "",
+    "medium": "",
+    "color_palette": []
+  },
+  "compositional_deconstruction": {
+    "background": "",
+    "elements": [
+    ]
+  }
+}
+`
+};
+
+export default datasetTemplates;
--- a/ui/src/hooks/useCaptionBatch.tsx
+++ b/ui/src/hooks/useCaptionBatch.tsx
@@ -4,13 +4,24 @@ import { apiClient } from '@/utils/api';

 // Module-level batcher: many cards mount at once when the virtualized grid scrolls;
 // instead of N HTTP requests, queue paths and flush them as a single batch.
+// Entries are keyed by extension + path so different caption extensions don't
+// collide in the cache or the pending batch.
 type Resolver = { resolve: (caption: string) => void; reject: (err: unknown) => void };
-const pending = new Map<string, Resolver[]>();
+type Pending = { path: string; ext: string; resolvers: Resolver[] };
+const pending = new Map<string, Pending>();
 const cache = new Map<string, string>();
 let flushTimer: ReturnType<typeof setTimeout> | null = null;
 const FLUSH_DELAY_MS = 30;
 const MAX_BATCH = 200;

+function normExt(ext: string | undefined): string {
+  return (ext || 'txt').replace(/^\.+/, '').trim() || 'txt';
+}
+
+function keyFor(path: string, ext: string): string {
+  return `${ext}\n${path}`;
+}
+
 function scheduleFlush() {
  if (flushTimer) return;
  flushTimer = setTimeout(flush, FLUSH_DELAY_MS);
@@ -20,55 +31,69 @@ async function flush() {
  flushTimer = null;
  if (pending.size === 0) return;

-  // Drain up to MAX_BATCH paths; if more arrived, reschedule.
-  const paths: string[] = [];
-  for (const path of pending.keys()) {
-    paths.push(path);
-    if (paths.length >= MAX_BATCH) break;
+  // Drain up to MAX_BATCH entries; if more arrived, reschedule.
+  const keys: string[] = [];
+  for (const key of pending.keys()) {
+    keys.push(key);
+    if (keys.length >= MAX_BATCH) break;
  }
-  const batchResolvers = paths.map(p => ({ path: p, resolvers: pending.get(p)! }));
-  for (const p of paths) pending.delete(p);
+  const drained = keys.map(k => pending.get(k)!);
+  for (const k of keys) pending.delete(k);

-  try {
-    const res = await apiClient.post('/api/caption/getBatch', { imgPaths: paths });
-    const captions: Record<string, string> = res.data?.captions ?? {};
-    for (const { path, resolvers } of batchResolvers) {
-      const value = captions[path] ?? '';
-      cache.set(path, value);
-      for (const r of resolvers) r.resolve(value);
-    }
-  } catch (err) {
-    for (const { resolvers } of batchResolvers) {
-      for (const r of resolvers) r.reject(err);
-    }
+  // Group by extension; each extension is a separate batch request.
+  const byExt = new Map<string, Pending[]>();
+  for (const entry of drained) {
+    const group = byExt.get(entry.ext);
+    if (group) group.push(entry);
+    else byExt.set(entry.ext, [entry]);
  }

+  await Promise.all(
+    Array.from(byExt.entries()).map(async ([ext, entries]) => {
+      const paths = entries.map(e => e.path);
+      try {
+        const res = await apiClient.post('/api/caption/getBatch', { imgPaths: paths, ext });
+        const captions: Record<string, string> = res.data?.captions ?? {};
+        for (const { path, ext: e, resolvers } of entries) {
+          const value = captions[path] ?? '';
+          cache.set(keyFor(path, e), value);
+          for (const r of resolvers) r.resolve(value);
+        }
+      } catch (err) {
+        for (const { resolvers } of entries) {
+          for (const r of resolvers) r.reject(err);
+        }
+      }
+    }),
+  );
+
  if (pending.size > 0) scheduleFlush();
 }

-function requestCaption(path: string, signal?: AbortSignal): Promise<string> {
+function requestCaption(path: string, ext: string, signal?: AbortSignal): Promise<string> {
  return new Promise((resolve, reject) => {
    if (signal?.aborted) {
      reject(new DOMException('Aborted', 'AbortError'));
      return;
    }
+    const key = keyFor(path, ext);
    const resolver: Resolver = { resolve, reject };
-    const list = pending.get(path);
-    if (list) {
-      list.push(resolver);
+    const entry = pending.get(key);
+    if (entry) {
+      entry.resolvers.push(resolver);
    } else {
-      pending.set(path, [resolver]);
+      pending.set(key, { path, ext, resolvers: [resolver] });
    }
    if (signal) {
      const onAbort = () => {
        // Remove this resolver from the pending batch. If no other card is
-        // still waiting on the same path, drop the path entirely so the next
+        // still waiting on the same key, drop the entry entirely so the next
        // batch doesn't include it.
-        const arr = pending.get(path);
-        if (arr) {
-          const idx = arr.indexOf(resolver);
-          if (idx >= 0) arr.splice(idx, 1);
-          if (arr.length === 0) pending.delete(path);
+        const e = pending.get(key);
+        if (e) {
+          const idx = e.resolvers.indexOf(resolver);
+          if (idx >= 0) e.resolvers.splice(idx, 1);
+          if (e.resolvers.length === 0) pending.delete(key);
        }
        reject(new DOMException('Aborted', 'AbortError'));
      };
@@ -78,19 +103,20 @@ function requestCaption(path: string, signal?: AbortSignal): Promise<string> {
  });
 }

-export function invalidateCaption(path: string) {
-  cache.delete(path);
+export function invalidateCaption(path: string, ext?: string) {
+  cache.delete(keyFor(path, normExt(ext)));
 }

-export function setCachedCaption(path: string, caption: string) {
-  cache.set(path, caption);
+export function setCachedCaption(path: string, caption: string, ext?: string) {
+  cache.set(keyFor(path, normExt(ext)), caption);
 }

 // Fetches caption for a path, using the module-level batcher + cache.
 // `refreshKey` busts the cache (e.g. after external edits or auto-captioning poll).
-export default function useCaptionBatch(imgPath: string | null, refreshKey: number = 0) {
-  const [caption, setCaption] = useState<string>(() => (imgPath ? (cache.get(imgPath) ?? '') : ''));
-  const [isLoaded, setIsLoaded] = useState<boolean>(() => Boolean(imgPath && cache.has(imgPath)));
+export default function useCaptionBatch(imgPath: string | null, refreshKey: number = 0, ext: string = 'txt') {
+  const captionExt = normExt(ext);
+  const [caption, setCaption] = useState<string>(() => (imgPath ? (cache.get(keyFor(imgPath, captionExt)) ?? '') : ''));
+  const [isLoaded, setIsLoaded] = useState<boolean>(() => Boolean(imgPath && cache.has(keyFor(imgPath, captionExt))));
  const lastPathRef = useRef<string | null>(null);

  useEffect(() => {
@@ -100,9 +126,9 @@ export default function useCaptionBatch(imgPath: string | null, refreshKey: numb
      return;
    }

-    if (refreshKey > 0) invalidateCaption(imgPath);
+    if (refreshKey > 0) invalidateCaption(imgPath, captionExt);

-    const cached = cache.get(imgPath);
+    const cached = cache.get(keyFor(imgPath, captionExt));
    if (cached !== undefined) {
      setCaption(cached);
      setIsLoaded(true);
@@ -114,7 +140,7 @@ export default function useCaptionBatch(imgPath: string | null, refreshKey: numb
    const controller = new AbortController();
    lastPathRef.current = imgPath;
    setIsLoaded(false);
-    requestCaption(imgPath, controller.signal)
+    requestCaption(imgPath, captionExt, controller.signal)
      .then(value => {
        if (cancelled || lastPathRef.current !== imgPath) return;
        setCaption(value);
@@ -130,7 +156,7 @@ export default function useCaptionBatch(imgPath: string | null, refreshKey: numb
      cancelled = true;
      controller.abort();
    };
-  }, [imgPath, refreshKey]);
+  }, [imgPath, refreshKey, captionExt]);

  return { caption, isLoaded, setCaption };
 }
--- a/ui/src/types.ts
+++ b/ui/src/types.ts
@@ -272,6 +272,7 @@ export interface CaptionProcessConfig {
    max_res?: number;
    max_new_tokens?: number;
    fixed_caption?: string;
+    caption_extension?: string;
  }
 }