Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions api/routers/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
9 changes: 9 additions & 0 deletions electron/main/ipc-handlers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
1 change: 1 addition & 0 deletions electron/preload/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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))
},
Expand Down
21 changes: 21 additions & 0 deletions src/areas/generate/GeneratePage.tsx
Original file line number Diff line number Diff line change
@@ -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'
Expand All @@ -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">
Expand All @@ -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>
</>
)
Expand Down
1 change: 1 addition & 0 deletions src/shared/types/electron.d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down