From 212cfe998ab8387ff6c2ce1392065b7e01a7a582 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Fri, 29 May 2026 08:56:22 -0600
Subject: [PATCH] Add ability to download or delete the optimizer state from
 the ui

---
 ui/src/app/api/jobs/[jobID]/files/route.ts | 10 +++++
 ui/src/components/FilesWidget.tsx          | 52 +++++++++++++++++++---
 2 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/ui/src/app/api/jobs/[jobID]/files/route.ts b/ui/src/app/api/jobs/[jobID]/files/route.ts
index 575df5e..bc4cc69 100644
--- a/ui/src/app/api/jobs/[jobID]/files/route.ts
+++ b/ui/src/app/api/jobs/[jobID]/files/route.ts
@@ -44,5 +44,15 @@ export async function GET(request: NextRequest, { params }: { params: { jobID: s
     };
   });
 
+  // include the optimizer state if it exists
+  const optimizerPath = path.join(jobFolder, 'optimizer.pt');
+  if (fs.existsSync(optimizerPath)) {
+    const stats = fs.statSync(optimizerPath);
+    fileObjects.push({
+      path: optimizerPath,
+      size: stats.size,
+    });
+  }
+
   return NextResponse.json({ files: fileObjects });
 }
diff --git a/ui/src/components/FilesWidget.tsx b/ui/src/components/FilesWidget.tsx
index f6b9c22..f3322c2 100644
--- a/ui/src/components/FilesWidget.tsx
+++ b/ui/src/components/FilesWidget.tsx
@@ -1,6 +1,6 @@
 import React from 'react';
 import useFilesList from '@/hooks/useFilesList';
-import { Loader2, AlertCircle, Download, Box, Brain, Trash2 } from 'lucide-react';
+import { Loader2, AlertCircle, Download, Box, Brain, Trash2, SlidersHorizontal } from 'lucide-react';
 import { openMergeLoRAsModal } from './MergeLoRAsModal';
 import { getFilename, getFoldername } from '@/utils/basic';
 import { openConfirm } from './ConfirmModal';
@@ -9,6 +9,10 @@ import { apiClient } from '@/utils/api';
 export default function FilesWidget({ jobID, jobName }: { jobID: string; jobName: string }) {
   const { files, status, refreshFiles } = useFilesList(jobID, 5000);
 
+  const isOptimizerFile = (filePath: string) => getFilename(filePath) === 'optimizer.pt';
+  const checkpointFiles = files.filter(file => !isOptimizerFile(file.path));
+  const optimizerFile = files.find(file => isOptimizerFile(file.path));
+
   const cleanSize = (size: number) => {
     if (size < 1024) {
       return `${size} B`;
@@ -47,17 +51,17 @@ export default function FilesWidget({ jobID, jobName }: { jobID: string; jobName
         <div className="flex items-center space-x-2">
           <Brain className="w-5 h-5 text-purple-600 dark:text-purple-400" />
           <h2 className="font-semibold text-gray-100">Checkpoints</h2>
-          <span className="px-2 py-0.5 bg-gray-700 rounded-full text-xs text-gray-300">{files.length}</span>
+          <span className="px-2 py-0.5 bg-gray-700 rounded-full text-xs text-gray-300">{checkpointFiles.length}</span>
         </div>
-        {files.length > 0 && (
+        {checkpointFiles.length > 0 && (
           <span
             className="px-3 py-1 rounded-full text-sm bg-purple-500/10 text-purple-500 uppercase cursor-pointer hover:bg-purple-500/20"
             onClick={() => {
               const outputName = `${jobName}_merged`;
               openMergeLoRAsModal(
-                getFoldername(files[0].path),
+                getFoldername(checkpointFiles[0].path),
                 outputName,
-                files.map(f => ({ path: f.path })),
+                checkpointFiles.map(f => ({ path: f.path })),
                 () => {
                   refreshFiles();
                 },
@@ -85,7 +89,7 @@ export default function FilesWidget({ jobID, jobName }: { jobID: string; jobName
 
         {['success', 'refreshing'].includes(status) && (
           <div className="space-y-1">
-            {files.map((file, index) => {
+            {checkpointFiles.map((file, index) => {
               const fileName = getFilename(file.path);
               const nameWithoutExt = fileName.replace('.safetensors', '');
               return (
@@ -129,6 +133,42 @@ export default function FilesWidget({ jobID, jobName }: { jobID: string; jobName
                 </div>
               );
             })}
+
+            {optimizerFile && (
+              <div className="group flex items-center justify-between px-2 py-1.5 rounded-lg border-t border-gray-800 mt-1 pt-2 hover:bg-gray-800 transition-all duration-200">
+                <a
+                  target="_blank"
+                  href={`/api/files/${encodeURIComponent(optimizerFile.path)}`}
+                  className="flex items-center space-x-2 min-w-0 flex-1"
+                >
+                  <SlidersHorizontal className="w-4 h-4 text-amber-500 flex-shrink-0" />
+                  <div className="flex flex-col min-w-0">
+                    <div className="flex text-sm text-amber-200">
+                      <span className="overflow-hidden text-ellipsis direction-rtl whitespace-nowrap">optimizer</span>
+                    </div>
+                    <span className="text-xs text-amber-600/70">.pt · optimizer state</span>
+                  </div>
+                </a>
+                <div className="flex items-center space-x-3 flex-shrink-0">
+                  <span className="text-xs text-gray-400">{cleanSize(optimizerFile.size)}</span>
+                  <a
+                    target="_blank"
+                    href={`/api/files/${encodeURIComponent(optimizerFile.path)}`}
+                    className="bg-amber-500 bg-opacity-0 group-hover:bg-opacity-10 rounded-full p-1 transition-all"
+                  >
+                    <Download className="w-3 h-3 text-amber-500" />
+                  </a>
+                  <button
+                    type="button"
+                    onClick={() => handleDeleteFile(optimizerFile.path)}
+                    className="bg-red-500 bg-opacity-0 group-hover:bg-opacity-10 hover:!bg-opacity-30 rounded-full p-1 transition-all"
+                    title="Delete optimizer state"
+                  >
+                    <Trash2 className="w-3 h-3 text-red-500" />
+                  </button>
+                </div>
+              </div>
+            )}
           </div>
         )}