From b0062535a784a48d0feeb84302aaa433040bbc1b Mon Sep 17 00:00:00 2001
From: Shine1i <wasimysdev@gmail.com>
Date: Fri, 13 Feb 2026 13:17:20 +0100
Subject: [PATCH] feat: add image previews in dataset dialog, enable popularity
 sorting in model search, refine training config serialization

---
 .../sections/dataset-preview-dialog.tsx       | 76 +++++++++++++++++++
 .../training/stores/training-config-store.ts  |  4 +
 .../frontend/src/hooks/use-hf-model-search.ts | 24 +++++-
 3 files changed, 103 insertions(+), 1 deletion(-)

diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
index e9d70846be..80fbcb65ee 100644
--- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
@@ -27,6 +27,14 @@ type CheckFormatResponse = {
   total_rows?: number | null;
 };
 
+type PreviewImagePayload = {
+  type: "image";
+  mime?: string;
+  width?: number;
+  height?: number;
+  data?: string;
+};
+
 type DatasetPreviewDialogProps = {
   open: boolean;
   onOpenChange: (open: boolean) => void;
@@ -123,6 +131,36 @@ export function DatasetPreviewDialog({
       ),
       cell: ({ getValue }: { getValue: () => unknown }) => {
         const value = getValue();
+        const images = collectPreviewImages(value);
+        if (images.length > 0) {
+          return (
+            <div className="flex flex-wrap gap-2">
+              {images.slice(0, 4).map((image, index) => {
+                const mime = image.mime || "image/jpeg";
+                const src = image.data ? `data:${mime};base64,${image.data}` : "";
+                const width = image.width ?? 128;
+                const height = image.height ?? 128;
+                return (
+                  <img
+                    key={`${colName}-img-${index}`}
+                    src={src}
+                    alt={`preview-${index}`}
+                    className="h-16 w-auto max-w-40 rounded-md border object-contain bg-muted"
+                    width={width}
+                    height={height}
+                    loading="lazy"
+                  />
+                );
+              })}
+              {images.length > 4 && (
+                <span className="text-xs text-muted-foreground self-end">
+                  +{images.length - 4} more
+                </span>
+              )}
+            </div>
+          );
+        }
+
         const text = formatCell(value);
         if (!text) {
           return (
@@ -285,3 +323,41 @@ function formatCell(value: unknown): string {
     return JSON.stringify(value).slice(0, 500);
   return String(value);
 }
+
+function isPreviewImagePayload(value: unknown): value is PreviewImagePayload {
+  if (!value || typeof value !== "object") return false;
+  const record = value as Record<string, unknown>;
+  return (
+    record.type === "image" &&
+    typeof record.data === "string" &&
+    record.data.length > 0
+  );
+}
+
+function collectPreviewImages(value: unknown): PreviewImagePayload[] {
+  const images: PreviewImagePayload[] = [];
+  const stack: unknown[] = [value];
+  let steps = 0;
+
+  while (stack.length > 0 && steps < 200) {
+    steps += 1;
+    const current = stack.pop();
+    if (isPreviewImagePayload(current)) {
+      images.push(current);
+      continue;
+    }
+
+    if (Array.isArray(current)) {
+      for (const item of current) stack.push(item);
+      continue;
+    }
+
+    if (current && typeof current === "object") {
+      for (const nested of Object.values(current as Record<string, unknown>)) {
+        stack.push(nested);
+      }
+    }
+  }
+
+  return images;
+}
diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts
index 95ded4c211..a9d6d37f42 100644
--- a/studio/frontend/src/features/training/stores/training-config-store.ts
+++ b/studio/frontend/src/features/training/stores/training-config-store.ts
@@ -96,6 +96,10 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
     }),
     {
       name: "unsloth_training_config_v1",
+      partialize: (state) => {
+        const { modelType, ...rest } = state;
+        return rest;
+      },
     },
   ),
 );
diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts
index e9d3b054d0..4745b4dc7b 100644
--- a/studio/frontend/src/hooks/use-hf-model-search.ts
+++ b/studio/frontend/src/hooks/use-hf-model-search.ts
@@ -23,6 +23,28 @@ const EXCLUDED_TAGS = new Set([
   "ctranslate2",
 ]);
 
+function withPopularitySort(
+  input: Parameters<typeof fetch>[0],
+  init?: Parameters<typeof fetch>[1],
+): ReturnType<typeof fetch> {
+  const rawUrl =
+    typeof input === "string"
+      ? input
+      : input instanceof URL
+        ? input.toString()
+        : input.url;
+  const url = new URL(rawUrl);
+
+  if (!url.searchParams.has("sort")) {
+    url.searchParams.set("sort", "downloads");
+  }
+  if (!url.searchParams.has("direction")) {
+    url.searchParams.set("direction", "-1");
+  }
+
+  return fetch(url, init);
+}
+
 function mapModel(raw: unknown): HfModelResult | null {
   const m = raw as {
     name: string;
@@ -53,10 +75,10 @@ export function useHfModelSearch(
       listModels({
         search: {
           ...(query.trim() ? { query } : { owner: "unsloth" }),
-          tags: ["transformers"],
           ...(task ? { task } : {}),
         },
         additionalFields: ["safetensors", "tags"],
+        fetch: withPopularitySort,
         ...(accessToken ? { credentials: { accessToken } } : {}),
       }) as AsyncGenerator<unknown>,
     [query, task, accessToken],