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.

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();

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 {

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);
 

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}
       />
     </>
   );

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"