327 lines
10 KiB
TypeScript
327 lines
10 KiB
TypeScript
import { useMemo, useState, useRef, useCallback } from 'react';
|
|
import { Virtuoso, VirtuosoHandle } from 'react-virtuoso';
|
|
import useSampleImages from '@/hooks/useSampleImages';
|
|
import SampleImageCard from './SampleImageCard';
|
|
import { Job } from '@prisma/client';
|
|
import { JobConfig } from '@/types';
|
|
import { LuImageOff, LuLoader, LuBan } from 'react-icons/lu';
|
|
import { Button } from '@headlessui/react';
|
|
import { FaDownload } from 'react-icons/fa';
|
|
import { apiClient } from '@/utils/api';
|
|
import classNames from 'classnames';
|
|
import { FaCaretDown, FaCaretUp } from 'react-icons/fa';
|
|
import SampleImageViewer from './SampleImageViewer';
|
|
|
|
interface SampleImagesMenuProps {
|
|
job?: Job | null;
|
|
}
|
|
|
|
export const SampleImagesMenu = ({ job }: SampleImagesMenuProps) => {
|
|
const [isZipping, setIsZipping] = useState(false);
|
|
|
|
const downloadZip = async () => {
|
|
if (isZipping) return;
|
|
setIsZipping(true);
|
|
|
|
try {
|
|
const res = await apiClient.post('/api/zip', {
|
|
zipTarget: 'samples',
|
|
jobName: job?.name,
|
|
});
|
|
|
|
const zipPath = res.data.zipPath; // e.g. /mnt/Train2/out/ui/.../samples.zip
|
|
if (!zipPath) throw new Error('No zipPath in response');
|
|
|
|
const downloadPath = `/api/files/${encodeURIComponent(zipPath)}`;
|
|
const a = document.createElement('a');
|
|
a.href = downloadPath;
|
|
// optional: suggest filename (browser may ignore if server sets Content-Disposition)
|
|
a.download = 'samples.zip';
|
|
document.body.appendChild(a);
|
|
a.click();
|
|
a.remove();
|
|
} catch (err) {
|
|
console.error('Error downloading zip:', err);
|
|
} finally {
|
|
setIsZipping(false);
|
|
}
|
|
};
|
|
return (
|
|
<Button
|
|
onClick={downloadZip}
|
|
className={classNames(
|
|
`flex-1 sm:flex-initial justify-center px-2 sm:px-4 py-1 h-8 hover:bg-gray-200 dark:hover:bg-gray-700 flex items-center`,
|
|
{
|
|
'opacity-50 cursor-not-allowed': isZipping,
|
|
},
|
|
)}
|
|
>
|
|
{isZipping ? (
|
|
<LuLoader className="animate-spin inline-block sm:mr-2" />
|
|
) : (
|
|
<FaDownload className="inline-block sm:mr-2" />
|
|
)}
|
|
<span className="hidden sm:inline">{isZipping ? 'Preparing' : 'Download'}</span>
|
|
</Button>
|
|
);
|
|
};
|
|
|
|
interface SampleImagesProps {
|
|
job: Job;
|
|
}
|
|
|
|
export default function SampleImages({ job }: SampleImagesProps) {
|
|
const { sampleImages, status, refreshSampleImages } = useSampleImages(job.id, 5000);
|
|
const [selectedSamplePath, setSelectedSamplePath] = useState<string | null>(null);
|
|
const [scrollParent, setScrollParent] = useState<HTMLDivElement | null>(null);
|
|
const scrollParentCallback = useCallback((el: HTMLDivElement | null) => setScrollParent(el), []);
|
|
const virtuosoRef = useRef<VirtuosoHandle>(null);
|
|
const numSamples = useMemo(() => {
|
|
if (job?.job_config) {
|
|
const jobConfig = JSON.parse(job.job_config) as JobConfig;
|
|
const sampleConfig = jobConfig.config.process[0].sample;
|
|
const numPrompts = sampleConfig.prompts ? sampleConfig.prompts.length : 0;
|
|
const numSamples = sampleConfig.samples.length;
|
|
return Math.max(numPrompts, numSamples, 1);
|
|
}
|
|
return 10;
|
|
}, [job]);
|
|
|
|
// Group samples into rows of `numSamples` for the virtualized list — one row per sample iteration.
|
|
const rows = useMemo(() => {
|
|
const out: string[][] = [];
|
|
for (let i = 0; i < sampleImages.length; i += numSamples) {
|
|
out.push(sampleImages.slice(i, i + numSamples));
|
|
}
|
|
return out;
|
|
}, [sampleImages, numSamples]);
|
|
|
|
const scrollToBottom = () => {
|
|
virtuosoRef.current?.scrollToIndex({ index: 'LAST', align: 'end' });
|
|
};
|
|
|
|
const scrollToTop = () => {
|
|
virtuosoRef.current?.scrollToIndex({ index: 0, align: 'start' });
|
|
};
|
|
|
|
const PageInfoContent = useMemo(() => {
|
|
let icon = null;
|
|
let text = '';
|
|
let subtitle = '';
|
|
let showIt = false;
|
|
let bgColor = '';
|
|
let textColor = '';
|
|
let iconColor = '';
|
|
|
|
if (sampleImages.length > 0) return null;
|
|
|
|
if (status == 'loading') {
|
|
icon = <LuLoader className="animate-spin w-8 h-8" />;
|
|
text = 'Loading Samples';
|
|
subtitle = 'Please wait while we fetch your samples...';
|
|
showIt = true;
|
|
bgColor = 'bg-gray-50 dark:bg-gray-800/50';
|
|
textColor = 'text-gray-900 dark:text-gray-100';
|
|
iconColor = 'text-gray-500 dark:text-gray-400';
|
|
}
|
|
if (status == 'error') {
|
|
icon = <LuBan className="w-8 h-8" />;
|
|
text = 'Error Loading Samples';
|
|
subtitle = 'There was a problem fetching the samples.';
|
|
showIt = true;
|
|
bgColor = 'bg-red-50 dark:bg-red-950/20';
|
|
textColor = 'text-red-900 dark:text-red-100';
|
|
iconColor = 'text-red-600 dark:text-red-400';
|
|
}
|
|
if (status == 'success' && sampleImages.length === 0) {
|
|
icon = <LuImageOff className="w-8 h-8" />;
|
|
text = 'No Samples Found';
|
|
subtitle = 'No samples have been generated yet';
|
|
showIt = true;
|
|
bgColor = 'bg-gray-50 dark:bg-gray-800/50';
|
|
textColor = 'text-gray-900 dark:text-gray-100';
|
|
iconColor = 'text-gray-500 dark:text-gray-400';
|
|
}
|
|
|
|
if (!showIt) return null;
|
|
|
|
return (
|
|
<div
|
|
className={`mt-10 flex flex-col items-center justify-center py-16 px-8 rounded-xl border-2 border-gray-700 border-dashed ${bgColor} ${textColor} mx-auto max-w-md text-center`}
|
|
>
|
|
<div className={`${iconColor} mb-4`}>{icon}</div>
|
|
<h3 className="text-lg font-semibold mb-2">{text}</h3>
|
|
<p className="text-sm opacity-75 leading-relaxed">{subtitle}</p>
|
|
</div>
|
|
);
|
|
}, [status, sampleImages.length]);
|
|
|
|
// Use direct Tailwind class without string interpolation
|
|
// This way Tailwind can properly generate the class
|
|
// I hate this, but it's the only way to make it work
|
|
const gridColsClass = useMemo(() => {
|
|
const cols = Math.min(numSamples, 40);
|
|
|
|
switch (cols) {
|
|
case 1:
|
|
case 2:
|
|
case 3:
|
|
return 'grid-cols-3';
|
|
case 4:
|
|
return 'grid-cols-4';
|
|
case 5:
|
|
return 'grid-cols-5';
|
|
case 6:
|
|
return 'grid-cols-6';
|
|
case 7:
|
|
return 'grid-cols-7';
|
|
case 8:
|
|
return 'grid-cols-8';
|
|
case 9:
|
|
return 'grid-cols-9';
|
|
case 10:
|
|
return 'grid-cols-10';
|
|
case 11:
|
|
return 'grid-cols-11';
|
|
case 12:
|
|
return 'grid-cols-12';
|
|
case 13:
|
|
return 'grid-cols-13';
|
|
case 14:
|
|
return 'grid-cols-14';
|
|
case 15:
|
|
return 'grid-cols-15';
|
|
case 16:
|
|
return 'grid-cols-16';
|
|
case 17:
|
|
return 'grid-cols-17';
|
|
case 18:
|
|
return 'grid-cols-18';
|
|
case 19:
|
|
return 'grid-cols-19';
|
|
case 20:
|
|
return 'grid-cols-20';
|
|
case 21:
|
|
return 'grid-cols-21';
|
|
case 22:
|
|
return 'grid-cols-22';
|
|
case 23:
|
|
return 'grid-cols-23';
|
|
case 24:
|
|
return 'grid-cols-24';
|
|
case 25:
|
|
return 'grid-cols-25';
|
|
case 26:
|
|
return 'grid-cols-26';
|
|
case 27:
|
|
return 'grid-cols-27';
|
|
case 28:
|
|
return 'grid-cols-28';
|
|
case 29:
|
|
return 'grid-cols-29';
|
|
case 30:
|
|
return 'grid-cols-30';
|
|
case 31:
|
|
return 'grid-cols-31';
|
|
case 32:
|
|
return 'grid-cols-32';
|
|
case 33:
|
|
return 'grid-cols-33';
|
|
case 34:
|
|
return 'grid-cols-34';
|
|
case 35:
|
|
return 'grid-cols-35';
|
|
case 36:
|
|
return 'grid-cols-36';
|
|
case 37:
|
|
return 'grid-cols-37';
|
|
case 38:
|
|
return 'grid-cols-38';
|
|
case 39:
|
|
return 'grid-cols-39';
|
|
case 40:
|
|
return 'grid-cols-40';
|
|
default:
|
|
return 'grid-cols-3';
|
|
}
|
|
}, [numSamples]);
|
|
|
|
const sampleConfig = useMemo(() => {
|
|
if (job?.job_config) {
|
|
const jobConfig = JSON.parse(job.job_config) as JobConfig;
|
|
return jobConfig.config.process[0].sample;
|
|
}
|
|
return null;
|
|
}, [job]);
|
|
|
|
return (
|
|
<div ref={scrollParentCallback} className="absolute top-[80px] left-0 right-0 bottom-0 overflow-y-auto">
|
|
<div className="pb-4">
|
|
{PageInfoContent}
|
|
{sampleImages && rows.length > 0 && scrollParent && (
|
|
<Virtuoso
|
|
ref={virtuosoRef}
|
|
customScrollParent={scrollParent}
|
|
totalCount={rows.length}
|
|
initialTopMostItemIndex={rows.length - 1}
|
|
followOutput="auto"
|
|
increaseViewportBy={400}
|
|
computeItemKey={index => rows[index]?.[0] ?? index}
|
|
itemContent={index => {
|
|
const row = rows[index];
|
|
if (!row) return null;
|
|
|
|
// Only pad the final row when numSamples < MIN_COLS and the row is short.
|
|
const MIN_COLS = 3;
|
|
const shouldPad = numSamples < MIN_COLS && row.length < MIN_COLS;
|
|
const padsNeeded = shouldPad ? MIN_COLS - row.length : 0;
|
|
|
|
return (
|
|
// pb-1 recreates the vertical gap between rows that the original single CSS grid provided via `gap-1`.
|
|
<div className={`grid ${gridColsClass} gap-1 pb-1`}>
|
|
{row.map(sample => (
|
|
<SampleImageCard
|
|
key={sample}
|
|
imageUrl={sample}
|
|
numSamples={numSamples}
|
|
sampleImages={sampleImages}
|
|
alt="Sample Image"
|
|
onClick={() => setSelectedSamplePath(sample)}
|
|
observerRoot={scrollParent}
|
|
/>
|
|
))}
|
|
{Array.from({ length: padsNeeded }).map((_, i) => (
|
|
<div key={`pad-${index}-${i}`} className="invisible" />
|
|
))}
|
|
</div>
|
|
);
|
|
}}
|
|
/>
|
|
)}
|
|
</div>
|
|
<SampleImageViewer
|
|
imgPath={selectedSamplePath}
|
|
numSamples={numSamples}
|
|
sampleImages={sampleImages}
|
|
onChange={setPath => setSelectedSamplePath(setPath)}
|
|
sampleConfig={sampleConfig}
|
|
refreshSampleImages={refreshSampleImages}
|
|
/>
|
|
<div
|
|
className="hidden md:flex fixed top-20 mt-4 right-6 w-10 h-10 rounded-full bg-gray-900 shadow-lg items-center justify-center text-white opacity-80 hover:opacity-100 cursor-pointer"
|
|
onClick={scrollToTop}
|
|
title="Scroll to Top"
|
|
>
|
|
<FaCaretUp className="text-gray-500 dark:text-gray-400" />
|
|
</div>
|
|
<div
|
|
className="hidden md:flex fixed bottom-5 right-6 w-10 h-10 rounded-full bg-gray-900 shadow-lg items-center justify-center text-white opacity-80 hover:opacity-100 cursor-pointer"
|
|
onClick={scrollToBottom}
|
|
title="Scroll to Bottom"
|
|
>
|
|
<FaCaretDown className="text-gray-500 dark:text-gray-400" />
|
|
</div>
|
|
</div>
|
|
);
|
|
}
|