From b4fc7ddd586557e12ac608bddd9115a80e518acd Mon Sep 17 00:00:00 2001
From: Lightning Pixel <lightning.pixel.dev@gmail.com>
Date: Sat, 21 Mar 2026 22:29:16 +0100
Subject: [PATCH] fix(generate): free memory button and model unload-all
 endpoint

---
 api/routers/model.py                |  7 +++++++
 electron/main/ipc-handlers.ts       |  9 +++++++++
 electron/preload/index.ts           |  1 +
 src/areas/generate/GeneratePage.tsx | 21 +++++++++++++++++++++
 src/shared/types/electron.d.ts      |  1 +
 5 files changed, 39 insertions(+)

diff --git a/api/routers/model.py b/api/routers/model.py
index 2748405..9326916 100644
--- a/api/routers/model.py
+++ b/api/routers/model.py
@@ -39,6 +39,13 @@ async def switch_model(model_id: str):
         raise HTTPException(400, str(e))
 
 
+@router.post("/unload-all")
+async def unload_all_models():
+    """Unloads all models from memory to free VRAM/RAM."""
+    generator_registry.unload_all()
+    return {"unloaded": True}
+
+
 @router.post("/unload/{model_id}")
 async def unload_model(model_id: str):
     """Unloads a model from memory so its files can be safely deleted."""
diff --git a/electron/main/ipc-handlers.ts b/electron/main/ipc-handlers.ts
index 6632f84..e17705d 100644
--- a/electron/main/ipc-handlers.ts
+++ b/electron/main/ipc-handlers.ts
@@ -134,6 +134,15 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe
     return result.canceled ? null : result.filePath
   })
 
+  ipcMain.handle('model:unloadAll', async (): Promise<{ success: boolean; error?: string }> => {
+    try {
+      await axios.post(`${API_BASE_URL}/model/unload-all`, {}, { timeout: 10_000 })
+      return { success: true }
+    } catch (err) {
+      return { success: false, error: String(err) }
+    }
+  })
+
   ipcMain.handle('model:delete', async (_, modelId: string): Promise<{ success: boolean; error?: string }> => {
     const modelDir = join(getSettings(app.getPath('userData')).modelsDir, modelId)
     try {
diff --git a/electron/preload/index.ts b/electron/preload/index.ts
index d9d8aa7..bc63204 100644
--- a/electron/preload/index.ts
+++ b/electron/preload/index.ts
@@ -70,6 +70,7 @@ contextBridge.exposeInMainWorld('electron', {
     isDownloaded:   (modelId: string) => ipcRenderer.invoke('model:isDownloaded', modelId),
     download:       (repoId: string, modelId: string) => ipcRenderer.invoke('model:download', { repoId, modelId }),
     delete:         (modelId: string) => ipcRenderer.invoke('model:delete', modelId),
+    unloadAll:      () => ipcRenderer.invoke('model:unloadAll'),
     onProgress:     (cb: (data: { modelId: string; percent: number; file?: string; fileIndex?: number; totalFiles?: number; status?: string }) => void) => {
       ipcRenderer.on('model:downloadProgress', (_event, data) => cb(data))
     },
diff --git a/src/areas/generate/GeneratePage.tsx b/src/areas/generate/GeneratePage.tsx
index 015c150..a99db6f 100644
--- a/src/areas/generate/GeneratePage.tsx
+++ b/src/areas/generate/GeneratePage.tsx
@@ -1,3 +1,4 @@
+import { useState } from 'react'
 import { useAppStore } from '@shared/stores/appStore'
 import { useGeneration } from '@shared/hooks/useGeneration'
 import ImageUpload from './components/ImageUpload'
@@ -12,9 +13,17 @@ export default function GeneratePage(): JSX.Element {
   const { currentJob, startGeneration } = useGeneration()
   const isGenerating = currentJob?.status === 'uploading' || currentJob?.status === 'generating'
 
+  const [unloadStatus, setUnloadStatus] = useState<'idle' | 'done'>('idle')
+
   const canGenerate = !!selectedImagePath && !!modelId && !isGenerating
   const disabledReason = !selectedImagePath ? 'Select an image first' : !modelId ? 'No model selected — install one in the Models tab' : undefined
 
+  async function handleUnloadAll() {
+    await window.electron.model.unloadAll()
+    setUnloadStatus('done')
+    setTimeout(() => setUnloadStatus('idle'), 2000)
+  }
+
   return (
     <>
       <div className="flex flex-col w-80 border-r border-zinc-800 bg-surface-400">
@@ -41,6 +50,18 @@ export default function GeneratePage(): JSX.Element {
         <Viewer3D />
         <GenerationHUD />
         <WorkspacePanel />
+
+        {/* Free memory button — top-left overlay */}
+        <button
+          onClick={handleUnloadAll}
+          title="Free model from memory"
+          className="absolute top-3 left-3 z-20 flex items-center gap-1.5 px-2.5 py-1.5 rounded-lg text-[11px] font-medium bg-zinc-900/70 border border-zinc-700/50 backdrop-blur-sm text-zinc-400 hover:text-zinc-200 hover:border-zinc-600 transition-colors"
+        >
+          <svg width="13" height="13" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="1.75">
+            <path d="M3 6h18M8 6V4h8v2M19 6l-1 14H6L5 6" />
+          </svg>
+          {unloadStatus === 'done' ? 'Freed' : 'Free memory'}
+        </button>
       </div>
     </>
   )
diff --git a/src/shared/types/electron.d.ts b/src/shared/types/electron.d.ts
index 3d9a4c0..95e62d0 100644
--- a/src/shared/types/electron.d.ts
+++ b/src/shared/types/electron.d.ts
@@ -42,6 +42,7 @@ declare global {
         isDownloaded:   (modelId: string) => Promise<boolean>
         download:       (repoId: string, modelId: string) => Promise<{ success: boolean; error?: string }>
         delete:         (modelId: string) => Promise<{ success: boolean; error?: string }>
+        unloadAll:      () => Promise<{ success: boolean; error?: string }>
         onProgress:     (cb: (data: { modelId: string; percent: number; file?: string; fileIndex?: number; totalFiles?: number; status?: string }) => void) => void
         offProgress:    () => void
       }