diff --git a/ui/src/components/SampleImageViewer.tsx b/ui/src/components/SampleImageViewer.tsx index e8804a9a..69a141f9 100644 --- a/ui/src/components/SampleImageViewer.tsx +++ b/ui/src/components/SampleImageViewer.tsx @@ -28,6 +28,7 @@ export default function SampleImageViewer({ }: Props) { const [mounted, setMounted] = useState(false); const [isOpen, setIsOpen] = useState(Boolean(imgPath)); + const [showingControlIdx, setShowingControlIdx] = useState(null); useEffect(() => setMounted(true), []); @@ -44,7 +45,10 @@ export default function SampleImageViewer({ } }, [isOpen, imgPath, onChange]); - const onCancel = useCallback(() => setIsOpen(false), []); + const onCancel = useCallback(() => { + setShowingControlIdx(null); + setIsOpen(false); + }, []); const imgInfo = useMemo(() => { // handle windows C:\\Apps\\AI-Toolkit\\AI-Toolkit\\output\\LoRA-Name\\samples\\1763563000704__000004000_0.jpg @@ -80,6 +84,7 @@ export default function SampleImageViewer({ const setImageAtIndex = useCallback( (idx: number) => { if (idx < 0 || idx >= sampleImages.length) return; + setShowingControlIdx(null); onChange(sampleImages[idx]); }, [sampleImages, numSamples, onChange], @@ -158,6 +163,13 @@ export default function SampleImageViewer({ return sampleConfig?.seed ?? '?'; }, [sampleItem, sampleConfig]); + const displayedImgPath = useMemo(() => { + if (showingControlIdx !== null && controlImages[showingControlIdx]) { + return controlImages[showingControlIdx]; + } + return imgPath; + }, [showingControlIdx, controlImages, imgPath]); + // keyboard events while open useEffect(() => { const handleKeyDown = (event: KeyboardEvent) => { @@ -201,10 +213,10 @@ export default function SampleImageViewer({ className="relative transform rounded-lg bg-gray-800 text-left shadow-xl transition-all data-closed:translate-y-4 data-closed:opacity-0 data-enter:duration-300 data-enter:ease-out data-leave:duration-200 data-leave:ease-in max-w-[95%] max-h-[95vh] data-closed:sm:translate-y-0 data-closed:sm:scale-95 flex flex-col overflow-hidden" >
- {imgPath && - (isVideo(imgPath) ? ( + {displayedImgPath && + (isVideo(displayedImgPath) ? (
{controlImages.length > 0 && (
+ {showingControlIdx !== null && ( + Main setShowingControlIdx(null)} + title="Main image" + /> + )} {controlImages.map((ci, idx) => ( {`Control setShowingControlIdx(idx)} + title={`Control image ${idx + 1}`} /> ))}