Preview samples full screen and use arrow keys to navigate them

ui/src/app/layout.tsx

@@ -4,6 +4,7 @@ import './globals.css';
 import Sidebar from '@/components/Sidebar';
 import { ThemeProvider } from '@/components/ThemeProvider';
 import ConfirmModal from '@/components/ConfirmModal';
+import SampleImageModal from '@/components/SampleImageModal';
 
 const inter = Inter({ subsets: ['latin'] });
 
@@ -23,6 +24,7 @@ export default function RootLayout({ children }: { children: React.ReactNode })
           </div>
         </ThemeProvider>
         <ConfirmModal />
+        <SampleImageModal />
       </body>
     </html>
   );

ui/src/components/SampleImageCard.tsx

@@ -1,14 +1,17 @@
 import React, { useRef, useEffect, useState, ReactNode } from 'react';
+import { sampleImageModalState } from '@/components/SampleImageModal';
 
 interface SampleImageCardProps {
   imageUrl: string;
   alt: string;
+  numSamples: number;
+  sampleImages: string[];
   children?: ReactNode;
   className?: string;
   onDelete?: () => void;
 }
 
-const SampleImageCard: React.FC<SampleImageCardProps> = ({ imageUrl, alt, children, className = '' }) => {
+const SampleImageCard: React.FC<SampleImageCardProps> = ({ imageUrl, alt, numSamples, sampleImages, children, className = '' }) => {
   const cardRef = useRef<HTMLDivElement>(null);
   const [isVisible, setIsVisible] = useState<boolean>(false);
   const [loaded, setLoaded] = useState<boolean>(false);
@@ -43,8 +46,9 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({ imageUrl, alt, childr
       {/* Square image container */}
       <div
         ref={cardRef}
-        className="relative w-full"
+        className="relative w-full cursor-pointer"
         style={{ paddingBottom: '100%' }} // Make it square
+        onClick={() => sampleImageModalState.set({ imgPath: imageUrl, numSamples, sampleImages })}
       >
         <div className="absolute inset-0 rounded-t-lg shadow-md">
           {isVisible && (

ui/src/components/SampleImageModal.tsx

@@ -0,0 +1,185 @@
+'use client';
+import { useState, useEffect, useMemo } from 'react';
+import { createGlobalState } from 'react-global-hooks';
+import { Dialog, DialogBackdrop, DialogPanel } from '@headlessui/react';
+
+export interface SampleImageModalState {
+  imgPath: string;
+  numSamples: number;
+  sampleImages: string[];
+}
+
+export const sampleImageModalState = createGlobalState<SampleImageModalState | null>(null);
+
+export const openSampleImage = (sampleImageProps: SampleImageModalState) => {
+  sampleImageModalState.set(sampleImageProps);
+};
+
+export default function SampleImageModal() {
+  const [imageModal, setImageModal] = sampleImageModalState.use();
+  const [isOpen, setIsOpen] = useState(false);
+
+  useEffect(() => {
+    if (imageModal) {
+      setIsOpen(true);
+    }
+  }, [imageModal]);
+
+  useEffect(() => {
+    if (!isOpen) {
+      // use timeout to allow the dialog to close before resetting the state
+      setTimeout(() => {
+        setImageModal(null);
+      }, 500);
+    }
+  }, [isOpen]);
+
+  const onCancel = () => {
+    setIsOpen(false);
+  };
+
+  const handleArrowUp = () => {
+    if (!imageModal) return;
+    console.log('Arrow Up pressed');
+    // Change image to same sample but up one step
+    const currentIdx = imageModal.sampleImages.findIndex(img => img === imageModal.imgPath);
+    if (currentIdx === -1) return;
+    const nextIdx = currentIdx - imageModal.numSamples;
+    if (nextIdx < 0) return;
+    openSampleImage({
+      imgPath: imageModal.sampleImages[nextIdx],
+      numSamples: imageModal.numSamples,
+      sampleImages: imageModal.sampleImages,
+    });
+  };
+
+  const handleArrowDown = () => {
+    if (!imageModal) return;
+    console.log('Arrow Down pressed');
+    // Change image to same sample but down one step
+    const currentIdx = imageModal.sampleImages.findIndex(img => img === imageModal.imgPath);
+    if (currentIdx === -1) return;
+    const nextIdx = currentIdx + imageModal.numSamples;
+    if (nextIdx >= imageModal.sampleImages.length) return;
+    openSampleImage({
+      imgPath: imageModal.sampleImages[nextIdx],
+      numSamples: imageModal.numSamples,
+      sampleImages: imageModal.sampleImages,
+    });
+  };
+
+  const handleArrowLeft = () => {
+    if (!imageModal) return;
+    console.log('Arrow Left pressed');
+    // go to previous sample
+    const currentIdx = imageModal.sampleImages.findIndex(img => img === imageModal.imgPath);
+    if (currentIdx === -1) return;
+    const nextIdx = currentIdx - 1;
+    if (nextIdx < 0) return;
+    openSampleImage({
+      imgPath: imageModal.sampleImages[nextIdx],
+      numSamples: imageModal.numSamples,
+      sampleImages: imageModal.sampleImages,
+    });
+  };
+
+  const handleArrowRight = () => {
+    if (!imageModal) return;
+    console.log('Arrow Right pressed');
+    // go to next sample
+    const currentIdx = imageModal.sampleImages.findIndex(img => img === imageModal.imgPath);
+    if (currentIdx === -1) return;
+    const nextIdx = currentIdx + 1;
+    if (nextIdx >= imageModal.sampleImages.length) return;
+    openSampleImage({
+      imgPath: imageModal.sampleImages[nextIdx],
+      numSamples: imageModal.numSamples,
+      sampleImages: imageModal.sampleImages,
+    });
+  };
+
+  // Handle keyboard events
+  useEffect(() => {
+    const handleKeyDown = (event: KeyboardEvent) => {
+      if (!isOpen) return;
+
+      switch (event.key) {
+        case 'Escape':
+          onCancel();
+          break;
+        case 'ArrowUp':
+          handleArrowUp();
+          break;
+        case 'ArrowDown':
+          handleArrowDown();
+          break;
+        case 'ArrowLeft':
+          handleArrowLeft();
+          break;
+        case 'ArrowRight':
+          handleArrowRight();
+          break;
+        default:
+          break;
+      }
+    };
+
+    window.addEventListener('keydown', handleKeyDown);
+
+    return () => {
+      window.removeEventListener('keydown', handleKeyDown);
+    };
+  }, [isOpen, imageModal]);
+
+  const imgInfo = useMemo(() => {
+    const ii = {
+      filename: '',
+      step: 0,
+      promptIdx: 0,
+    };
+    if (imageModal?.imgPath) {
+      const filename = imageModal.imgPath.split('/').pop();
+      if (!filename) return ii;
+      // filename is <timestep>__<zero_pad_step>_<prompt_idx>.<ext>
+      ii.filename = filename as string;
+      const parts = filename
+        .split('.')[0]
+        .split('_')
+        .filter(p => p !== '');
+      if (parts.length === 3) {
+        ii.step = parseInt(parts[1]);
+        ii.promptIdx = parseInt(parts[2]);
+      }
+    }
+    return ii;
+  }, [imageModal]);
+
+  return (
+    <Dialog open={isOpen} onClose={onCancel} className="relative z-10">
+      <DialogBackdrop
+        transition
+        className="fixed inset-0 bg-gray-900/75 transition-opacity data-closed:opacity-0 data-enter:duration-300 data-enter:ease-out data-leave:duration-200 data-leave:ease-in"
+      />
+
+      <div className="fixed inset-0 z-10 w-screen overflow-y-auto">
+        <div className="flex min-h-full items-center justify-center p-4 text-center">
+          <DialogPanel
+            transition
+            className="relative transform overflow-hidden 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"
+          >
+            <div className="flex justify-center items-center">
+              {imageModal?.imgPath && (
+                <img
+                  src={`/api/img/${encodeURIComponent(imageModal.imgPath)}`}
+                  alt="Sample Image"
+                  className="max-w-full max-h-[calc(95vh-2rem)] object-contain"
+                />
+              )}
+            </div>
+            <div className="bg-gray-950 text-center text-sm p-2">step: {imgInfo.step}</div>
+          </DialogPanel>
+        </div>
+      </div>
+    </Dialog>
+  );
+}

ui/src/components/SampleImages.tsx

@@ -11,13 +11,13 @@ interface SampleImagesProps {
 export default function SampleImages({ job }: SampleImagesProps) {
   const { sampleImages, status, refreshSampleImages } = useSampleImages(job.id, 5000);
   const numSamples = useMemo(() => {
-      if (job?.job_config) {
-        const jobConfig = JSON.parse(job.job_config) as JobConfig;
-        const sampleConfig = jobConfig.config.process[0].sample;
-        return sampleConfig.prompts.length;
-      }
-      return 10;
-    }, [job]);
+    if (job?.job_config) {
+      const jobConfig = JSON.parse(job.job_config) as JobConfig;
+      const sampleConfig = jobConfig.config.process[0].sample;
+      return sampleConfig.prompts.length;
+    }
+    return 10;
+  }, [job]);
 
   // Use direct Tailwind class without string interpolation
   // This way Tailwind can properly generate the class
@@ -73,13 +73,19 @@ export default function SampleImages({ job }: SampleImagesProps) {
 
   return (
     <div>
-      <div className='pb-4'>
+      <div className="pb-4">
         {status === 'loading' && sampleImages.length === 0 && <p>Loading...</p>}
         {status === 'error' && <p>Error fetching sample images</p>}
         {sampleImages && (
           <div className={`grid ${gridColsClass} gap-1`}>
             {sampleImages.map((sample: string) => (
-              <SampleImageCard key={sample} imageUrl={sample} alt="Sample Image" />
+              <SampleImageCard
+                key={sample}
+                imageUrl={sample}
+                numSamples={numSamples}
+                sampleImages={sampleImages}
+                alt="Sample Image"
+              />
             ))}
           </div>
         )}