Reworked ui sample image modal to show more information and function a lot better.

ui/next.config.ts

@@ -1,6 +1,9 @@
 import type { NextConfig } from 'next';
 
 const nextConfig: NextConfig = {
+  devIndicators: {
+    buildActivity: false,
+  },
   typescript: {
     // Remove this. Build fails because of route types
     ignoreBuildErrors: true,

ui/src/app/layout.tsx

@@ -4,7 +4,6 @@ import './globals.css';
 import Sidebar from '@/components/Sidebar';
 import { ThemeProvider } from '@/components/ThemeProvider';
 import ConfirmModal from '@/components/ConfirmModal';
-import SampleImageModal from '@/components/SampleImageModal';
 import { Suspense } from 'react';
 import AuthWrapper from '@/components/AuthWrapper';
 import DocModal from '@/components/DocModal';
@@ -40,7 +39,6 @@ export default function RootLayout({ children }: { children: React.ReactNode })
         </ThemeProvider>
         <ConfirmModal />
         <DocModal />
-        <SampleImageModal />
       </body>
     </html>
   );

ui/src/components/SampleImageCard.tsx

@@ -1,5 +1,4 @@
 import React, { useRef, useEffect, useState, ReactNode } from 'react';
-import { sampleImageModalState } from '@/components/SampleImageModal';
 import { isVideo } from '@/utils/basic';
 
 interface SampleImageCardProps {
@@ -10,6 +9,7 @@ interface SampleImageCardProps {
   children?: ReactNode;
   className?: string;
   onDelete?: () => void;
+  onClick?: () => void;
 }
 
 const SampleImageCard: React.FC<SampleImageCardProps> = ({
@@ -19,6 +19,7 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
   sampleImages,
   children,
   className = '',
+  onClick = () => {},
 }) => {
   const cardRef = useRef<HTMLDivElement>(null);
   const [isVisible, setIsVisible] = useState<boolean>(false);
@@ -56,7 +57,7 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
         ref={cardRef}
         className="relative w-full cursor-pointer"
         style={{ paddingBottom: '100%' }} // Make it square
-        onClick={() => sampleImageModalState.set({ imgPath: imageUrl, numSamples, sampleImages })}
+        onClick={onClick}
       >
         <div className="absolute inset-0 rounded-t-lg shadow-md">
           {isVisible && (

ui/src/components/SampleImageModal.tsx

@@ -1,190 +0,0 @@
-'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 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]);
-
-  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;
-    if (imgInfo.promptIdx === 0) return;
-    console.log('Arrow Left pressed');
-    // go to previous sample
-    const currentIdx = imageModal.sampleImages.findIndex(img => img === imageModal.imgPath);
-    if (currentIdx === -1) return;
-    const minIdx = currentIdx - imgInfo.promptIdx;
-    const nextIdx = currentIdx - 1;
-    if (nextIdx < minIdx) 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 stepMinIdx = currentIdx - imgInfo.promptIdx;
-    const maxIdx = stepMinIdx + imageModal.numSamples - 1;
-    const nextIdx = currentIdx + 1;
-    if (nextIdx > maxIdx) return;
-    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, imgInfo]);
-
-  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/SampleImageViewer.tsx

@@ -0,0 +1,217 @@
+'use client';
+import { useState, useEffect, useMemo, useCallback } from 'react';
+import { createPortal } from 'react-dom';
+import { Dialog, DialogBackdrop, DialogPanel } from '@headlessui/react';
+import { SampleConfig, SampleItem } from '@/types';
+
+interface Props {
+  imgPath: string | null; // current image path
+  numSamples: number; // number of samples per row
+  sampleImages: string[]; // all sample images
+  sampleConfig: SampleConfig | null;
+  onChange: (nextPath: string | null) => void; // parent setter
+}
+
+export default function SampleImageViewer({ imgPath, numSamples, sampleImages, sampleConfig, onChange }: Props) {
+  const [mounted, setMounted] = useState(false);
+  const [isOpen, setIsOpen] = useState(Boolean(imgPath));
+
+  useEffect(() => setMounted(true), []);
+
+  // open/close based on external value
+  useEffect(() => {
+    setIsOpen(Boolean(imgPath));
+  }, [imgPath]);
+
+  // after close, clear parent state post-transition
+  useEffect(() => {
+    if (!isOpen && imgPath) {
+      const t = setTimeout(() => onChange(null), 300);
+      return () => clearTimeout(t);
+    }
+  }, [isOpen, imgPath, onChange]);
+
+  const onCancel = useCallback(() => setIsOpen(false), []);
+
+  const imgInfo = useMemo(() => {
+    const ii = { filename: '', step: 0, promptIdx: 0 };
+    if (imgPath) {
+      const filename = imgPath.split('/').pop();
+      if (!filename) return ii;
+      ii.filename = filename;
+      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;
+  }, [imgPath]);
+
+  const setImageAtIndex = useCallback(
+    (idx: number) => {
+      if (idx < 0 || idx >= sampleImages.length) return;
+      onChange(sampleImages[idx]);
+    },
+    [sampleImages, numSamples, onChange],
+  );
+
+  const currentIndex = useMemo(() => {
+    if (!imgPath) return -1;
+    return sampleImages.findIndex(img => img === imgPath);
+  }, [imgPath, sampleImages]);
+
+  const handleArrowUp = useCallback(() => {
+    if (currentIndex === -1) return;
+    setImageAtIndex(currentIndex - numSamples);
+  }, [numSamples, currentIndex, setImageAtIndex]);
+
+  const handleArrowDown = useCallback(() => {
+    if (currentIndex === -1) return;
+    setImageAtIndex(currentIndex + numSamples);
+  }, [numSamples, currentIndex, setImageAtIndex]);
+
+  const handleArrowLeft = useCallback(() => {
+    if (currentIndex === -1) return;
+    if (imgInfo.promptIdx === 0) return;
+    const minIdx = currentIndex - imgInfo.promptIdx;
+    const nextIdx = currentIndex - 1;
+    if (nextIdx < minIdx) return;
+    setImageAtIndex(nextIdx);
+  }, [sampleImages, currentIndex, imgInfo.promptIdx, setImageAtIndex]);
+
+  const handleArrowRight = useCallback(() => {
+    if (currentIndex === -1) return;
+    const stepMinIdx = currentIndex - imgInfo.promptIdx;
+    const maxIdx = stepMinIdx + numSamples - 1;
+    const nextIdx = currentIndex + 1;
+    if (nextIdx > maxIdx) return;
+    setImageAtIndex(nextIdx);
+  }, [sampleImages, currentIndex, imgInfo.promptIdx, setImageAtIndex]);
+
+  const sampleItem = useMemo<SampleItem | null>(() => {
+    if (!sampleConfig) return null;
+    if (imgInfo.promptIdx < 0) return null;
+    if (imgInfo.promptIdx >= sampleConfig.samples.length) return null;
+    return sampleConfig.samples[imgInfo.promptIdx];
+  }, [sampleConfig, imgInfo.promptIdx]);
+
+  const controlImages = useMemo<string[]>(() => {
+    if (!imgPath) return [];
+    let controlImageArr: string[] = [];
+    if (sampleItem?.ctrl_img) {
+      // can be a an array of paths, or a single path
+      if (Array.isArray(sampleItem.ctrl_img)) {
+        controlImageArr = sampleItem.ctrl_img;
+      } else {
+        controlImageArr = [sampleItem.ctrl_img];
+      }
+    } else if (sampleItem?.ctrl_img_1) {
+      controlImageArr.push(sampleItem.ctrl_img_1);
+    }
+    if (sampleItem?.ctrl_img_2) {
+      controlImageArr.push(sampleItem.ctrl_img_2);
+    }
+    if (sampleItem?.ctrl_img_3) {
+      controlImageArr.push(sampleItem.ctrl_img_3);
+    }
+    // filter out nulls
+    controlImageArr = controlImageArr.filter(ci => ci !== null && ci !== undefined && ci !== '');
+    return controlImageArr;
+  }, [sampleItem, imgPath]);
+
+  // keyboard events while open
+  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, onCancel, handleArrowUp, handleArrowDown, handleArrowLeft, handleArrowRight]);
+
+  if (!mounted) return null;
+
+  return createPortal(
+    <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 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"
+          >
+            <div className="overflow-hidden flex items-center justify-center">
+              {imgPath && (
+                <img
+                  src={`/api/img/${encodeURIComponent(imgPath)}`}
+                  alt="Sample Image"
+                  className="w-auto h-auto max-w-[95vw] max-h-[82vh] object-contain"
+                />
+              )}
+            </div>
+            {/* # make full width */}
+            <div className="bg-gray-950 text-sm flex justify-between items-center px-4 py-2">
+              <div className="flex-1 relative h-10 min-w-0">
+                {sampleItem?.prompt && (
+                  <div className="absolute inset-0 grid place-items-center overflow-auto mr-4">
+                    <div className="w-full">
+                      <span className="text-gray-400 mr-1">Prompt:</span>
+                      <span className="whitespace-pre-wrap break-words">{sampleItem.prompt}</span>
+                    </div>
+                  </div>
+                )}
+              </div>
+              {controlImages.length > 0 && (
+                <div key={imgPath} className="flex space-x-2 mr-4">
+                  {controlImages.map((ci, idx) => (
+                    <img
+                      key={idx}
+                      src={`/api/img/${encodeURIComponent(ci)}`}
+                      alt={`Control ${idx + 1}`}
+                      className="max-h-12 max-w-12 object-contain bg-black border border-gray-700 rounded"
+                    />
+                  ))}
+                </div>
+              )}
+
+              <div className="text-xs">
+                <div>
+                  <span className="text-gray-400">Step:</span> {imgInfo.step.toLocaleString()}
+                </div>
+                <div>
+                  <span className="text-gray-400">Sample #:</span> {imgInfo.promptIdx + 1}
+                </div>
+              </div>
+            </div>
+          </DialogPanel>
+        </div>
+      </div>
+    </Dialog>,
+    document.body,
+  );
+}

ui/src/components/SampleImages.tsx

@@ -8,6 +8,7 @@ import { Button } from '@headlessui/react';
 import { FaDownload } from 'react-icons/fa';
 import { apiClient } from '@/utils/api';
 import classNames from 'classnames';
+import SampleImageViewer from './SampleImageViewer';
 
 interface SampleImagesMenuProps {
   job?: Job | null;
@@ -66,15 +67,14 @@ interface SampleImagesProps {
 
 export default function SampleImages({ job }: SampleImagesProps) {
   const { sampleImages, status, refreshSampleImages } = useSampleImages(job.id, 5000);
+  const [selectedSamplePath, setSelectedSamplePath] = useState<string | null>(null);
   const numSamples = useMemo(() => {
     if (job?.job_config) {
       const jobConfig = JSON.parse(job.job_config) as JobConfig;
       const sampleConfig = jobConfig.config.process[0].sample;
-      if (sampleConfig.prompts) {
-        return sampleConfig.prompts.length;
-      } else {
-        return sampleConfig.samples.length;
-      }
+      const numPrompts = sampleConfig.prompts ? sampleConfig.prompts.length : 0;
+      const numSamples = sampleConfig.samples.length;
+      return Math.max(numPrompts, numSamples, 1);
     }
     return 10;
   }, [job]);
@@ -223,6 +223,14 @@ export default function SampleImages({ job }: SampleImagesProps) {
     }
   }, [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>
       <div className="pb-4">
@@ -236,11 +244,19 @@ export default function SampleImages({ job }: SampleImagesProps) {
                 numSamples={numSamples}
                 sampleImages={sampleImages}
                 alt="Sample Image"
+                onClick={() => setSelectedSamplePath(sample)}
               />
             ))}
           </div>
         )}
       </div>
+      <SampleImageViewer
+        imgPath={selectedSamplePath}
+        numSamples={numSamples}
+        sampleImages={sampleImages}
+        onChange={setPath => setSelectedSamplePath(setPath)}
+        sampleConfig={sampleConfig}
+      />
     </div>
   );
 }