github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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.

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-06 08:32:24 -06:00
parent 10cdeb394e
commit 41157b460c
19 changed files with 219 additions and 270 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
ui/src/components/DatasetImageCard.tsx
View File

@@ -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));
};
}, []);