From 8042cbe9d29f6811894adefdb8d3f971befb286f Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Mon, 25 May 2026 07:22:24 -0600
Subject: [PATCH] Added virtulization for sample images to handle huge number
 of samples more efficientyly.

---
 ui/src/components/DatasetImageCard.tsx |  6 +-
 ui/src/components/SampleImageCard.tsx  |  8 ++-
 ui/src/components/SampleImages.tsx     | 99 +++++++++++++-------------
 3 files changed, 62 insertions(+), 51 deletions(-)

diff --git a/ui/src/components/DatasetImageCard.tsx b/ui/src/components/DatasetImageCard.tsx
index bedb42b2..a931c99e 100644
--- a/ui/src/components/DatasetImageCard.tsx
+++ b/ui/src/components/DatasetImageCard.tsx
@@ -115,7 +115,11 @@ const DatasetImageCard: React.FC<DatasetImageCardProps> = ({
   return (
     <div className={`flex flex-col ${className}`}>
       <div className="relative w-full" style={{ paddingBottom: '100%' }}>
-        <div className="absolute inset-0 rounded-t-lg shadow-md">
+        <div
+          className={classNames('absolute inset-0 rounded-t-lg shadow-md bg-gray-900', {
+            'animate-pulse': isItImage && !loaded,
+          })}
+        >
           {isItAVideo && (
             <video
               src={`/api/img/${encodeURIComponent(imageUrl)}`}
diff --git a/ui/src/components/SampleImageCard.tsx b/ui/src/components/SampleImageCard.tsx
index baf9497a..a801f69e 100644
--- a/ui/src/components/SampleImageCard.tsx
+++ b/ui/src/components/SampleImageCard.tsx
@@ -58,10 +58,16 @@ const SampleImageCard: React.FC<SampleImageCardProps> = ({
 
   const handleLoad = () => setLoaded(true);
 
+  const isImageType = !isAudio(imageUrl) && !isVideo(imageUrl);
+
   return (
     <div className={`flex flex-col ${className}`}>
       <div ref={cardRef} className="relative w-full cursor-pointer" style={{ paddingBottom: '100%' }} onClick={onClick}>
-        <div className="absolute inset-0 rounded-t-lg shadow-md">
+        <div
+          className={`absolute inset-0 rounded-t-lg shadow-md bg-gray-900 ${
+            isVisible && isImageType && !loaded ? 'animate-pulse' : ''
+          }`}
+        >
           {isVisible ? (
             isAudio(imageUrl) ? (
               <div className="w-full h-full flex items-center justify-center bg-gray-900">
diff --git a/ui/src/components/SampleImages.tsx b/ui/src/components/SampleImages.tsx
index e3a61ae3..7a6ebf3a 100644
--- a/ui/src/components/SampleImages.tsx
+++ b/ui/src/components/SampleImages.tsx
@@ -1,4 +1,5 @@
-import { useMemo, useState, useRef, useEffect } from 'react';
+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';
@@ -72,8 +73,9 @@ interface SampleImagesProps {
 export default function SampleImages({ job }: SampleImagesProps) {
   const { sampleImages, status, refreshSampleImages } = useSampleImages(job.id, 5000);
   const [selectedSamplePath, setSelectedSamplePath] = useState<string | null>(null);
-  const containerRef = useRef<HTMLDivElement>(null);
-  const didFirstScroll = useRef(false);
+  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;
@@ -85,16 +87,21 @@ export default function SampleImages({ job }: SampleImagesProps) {
     return 10;
   }, [job]);
 
-  const scrollToBottom = () => {
-    if (containerRef.current) {
-      containerRef.current.scrollTo({ top: containerRef.current.scrollHeight, behavior: 'instant' });
+  // 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 = () => {
-    if (containerRef.current) {
-      containerRef.current.scrollTo({ top: 0, behavior: 'instant' });
-    }
+    virtuosoRef.current?.scrollToIndex({ index: 0, align: 'start' });
   };
 
   const PageInfoContent = useMemo(() => {
@@ -247,55 +254,49 @@ export default function SampleImages({ job }: SampleImagesProps) {
     return null;
   }, [job]);
 
-  // scroll to bottom on first load of samples
-  useEffect(() => {
-    if (status === 'success' && sampleImages.length > 0 && !didFirstScroll.current) {
-      didFirstScroll.current = true;
-      setTimeout(() => {
-        scrollToBottom();
-      }, 100);
-    }
-  }, [status, sampleImages.length]);
-
   return (
-    <div ref={containerRef} className="absolute top-[80px] left-0 right-0 bottom-0 overflow-y-auto">
+    <div ref={scrollParentCallback} className="absolute top-[80px] left-0 right-0 bottom-0 overflow-y-auto">
       <div className="pb-4">
         {PageInfoContent}
-        {sampleImages && (
-          <div className={`grid ${gridColsClass} gap-1`}>
-            {sampleImages.map((sample: string, idx: number) => {
-              // Compute current group (groups are size = numSamples)
-              const groupIndex = Math.floor(idx / numSamples);
-              const groupStart = groupIndex * numSamples;
-              const groupEnd = Math.min(groupStart + numSamples, sampleImages.length);
-              const groupSize = groupEnd - groupStart;
-              const isEndOfGroup = idx === groupEnd - 1;
+        {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 enforce a MIN of 3 when the group's planned width is < 3
+              // Only pad the final row when numSamples < MIN_COLS and the row is short.
               const MIN_COLS = 3;
-              const shouldPad = numSamples < MIN_COLS && groupSize < MIN_COLS;
-              const padsNeeded = shouldPad ? MIN_COLS - groupSize : 0;
+              const shouldPad = numSamples < MIN_COLS && row.length < MIN_COLS;
+              const padsNeeded = shouldPad ? MIN_COLS - row.length : 0;
 
               return (
-                <div key={sample} className="contents">
-                  <SampleImageCard
-                    imageUrl={sample}
-                    numSamples={numSamples}
-                    sampleImages={sampleImages}
-                    alt="Sample Image"
-                    onClick={() => setSelectedSamplePath(sample)}
-                    observerRoot={containerRef.current}
-                  />
-
-                  {isEndOfGroup &&
-                    padsNeeded > 0 &&
-                    Array.from({ length: padsNeeded }).map((_, i) => (
-                      <div key={`pad-${groupIndex}-${i}`} className="invisible" />
-                    ))}
+                // 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>
+            }}
+          />
         )}
       </div>
       <SampleImageViewer