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
10 changes: 10 additions & 0 deletions api/routers/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,16 @@ async def switch_model(model_id: str):
async def unload_all_models():
"""Unloads all models from memory to free VRAM/RAM."""
generator_registry.unload_all()
# Force Python to release memory back to the OS
import gc
gc.collect()
try:
import ctypes, sys
if sys.platform == "win32":
k32 = ctypes.windll.kernel32
k32.SetProcessWorkingSetSizeEx(k32.GetCurrentProcess(), -1, -1, 0)
except Exception:
pass
return {"unloaded": True}


Expand Down
19 changes: 19 additions & 0 deletions api/services/generators/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,25 @@ def load(self) -> None:
def unload(self) -> None:
"""Release memory. Can be overridden if needed."""
self._model = None
import gc
gc.collect()
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except ImportError:
pass
# Force the OS to reclaim unused memory from this process
try:
import ctypes
import sys
if sys.platform == "win32":
kernel32 = ctypes.windll.kernel32
kernel32.SetProcessWorkingSetSizeEx(
kernel32.GetCurrentProcess(), -1, -1, 0
)
except Exception:
pass

def is_loaded(self) -> bool:
return self._model is not None
Expand Down
9 changes: 0 additions & 9 deletions api/services/generators/hunyuan3d.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,15 +69,6 @@ def load(self) -> None:
self._model = pipeline
print(f"[Hunyuan3DGenerator] Loaded on {device}.")

def unload(self) -> None:
super().unload()
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except ImportError:
pass

# ------------------------------------------------------------------ #
# Inference
# ------------------------------------------------------------------ #
Expand Down
11 changes: 10 additions & 1 deletion electron/main/python-bridge.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ export class PythonBridge {
private ready = false
private startPromise: Promise<void> | null = null
private getWindow: (() => BrowserWindow | null) | null = null
private intentionalStop = false

setWindowGetter(fn: () => BrowserWindow | null): void {
this.getWindow = fn
Expand Down Expand Up @@ -78,7 +79,7 @@ export class PythonBridge {
console.log('[PythonBridge] Process exited with code', code)
this.ready = false
this.process = null
if (wasReady) {
if (wasReady && !this.intentionalStop) {
this.getWindow()?.webContents.send('python:crashed', { code })
}
})
Expand All @@ -100,6 +101,14 @@ export class PythonBridge {
console.log('[PythonBridge] Stopped')
}

async restart(): Promise<void> {
console.log('[PythonBridge] Restarting to free memory…')
this.intentionalStop = true
await this.stop()
this.intentionalStop = false
await this.start()
}

private emitTqdmLog(raw: string): void {
if (/INFO/.test(raw)) return
if (!raw.trim()) return
Expand Down
3 changes: 2 additions & 1 deletion src/areas/generate/GeneratePage.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,9 @@ export default function GeneratePage(): JSX.Element {
{/* Free memory button — top-left overlay */}
<button
onClick={handleUnloadAll}
disabled={isGenerating}
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"
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 disabled:opacity-30 disabled:pointer-events-none"
>
<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" />
Expand Down
112 changes: 99 additions & 13 deletions src/areas/models/ModelsPage.tsx
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { useEffect, useState } from 'react'
import { createPortal } from 'react-dom'
import { useAppStore } from '@shared/stores/appStore'
import { useNavStore } from '@shared/stores/navStore'
import { useExtensionsStore } from '@shared/stores/extensionsStore'
Expand Down Expand Up @@ -36,6 +37,7 @@ export default function ModelsPage(): JSX.Element {
const [deleteTarget, setDeleteTarget] = useState<LocalModel | null>(null)
const [deleteError, setDeleteError] = useState<string | null>(null)
const [uninstallTarget, setUninstallTarget] = useState<string | null>(null)
const [modelsToDelete, setModelsToDelete] = useState<Set<string>>(new Set())

// GitHub extension install form
const [showGHForm, setShowGHForm] = useState(false)
Expand Down Expand Up @@ -104,13 +106,25 @@ export default function ModelsPage(): JSX.Element {

// ── Uninstall extension ────────────────────────────────────────────────────

function openUninstallModal(extId: string) {
const ext = extensions.find((e) => e.id === extId)
const installedModels = ext?.models.filter((v) => installedVariantIds.includes(v.id)) ?? []
setModelsToDelete(new Set(installedModels.map((v) => v.id)))
setUninstallTarget(extId)
}

async function handleUninstallExtension(extId: string) {
// Delete selected models first
for (const modelId of modelsToDelete) {
await window.electron.model.delete(modelId)
}
const result = await uninstallExt(extId)
setUninstallTarget(null)
setModelsToDelete(new Set())
if (!result.success) {
// Surface the error — for now just log; could show a toast
console.error('[extensions:uninstall]', result.error)
}
refresh()
}

// ── Helpers ────────────────────────────────────────────────────────────────
Expand Down Expand Up @@ -287,7 +301,7 @@ export default function ModelsPage(): JSX.Element {
}
})
}}
onUninstall={(extId) => setUninstallTarget(extId)}
onUninstall={(extId) => openUninstallModal(extId)}
/>
))}
</div>
Expand Down Expand Up @@ -348,17 +362,89 @@ export default function ModelsPage(): JSX.Element {
)}

{/* ── Confirm uninstall extension ──────────────────────────────────── */}
{uninstallTarget && (
<ConfirmModal
title={`Uninstall extension "${uninstallTarget}"?`}
description="The extension folder will be deleted. Downloaded model weights will not be affected."
confirmLabel="Uninstall"
cancelLabel="Cancel"
variant="danger"
onConfirm={() => handleUninstallExtension(uninstallTarget)}
onCancel={() => setUninstallTarget(null)}
/>
)}
{uninstallTarget && (() => {
const ext = extensions.find((e) => e.id === uninstallTarget)
const installedModels = ext?.models.filter((v) => installedVariantIds.includes(v.id)) ?? []

return createPortal(
<div
className="fixed inset-0 z-[9999] flex items-center justify-center"
onMouseDown={(e) => { if (e.target === e.currentTarget) { setUninstallTarget(null); setModelsToDelete(new Set()) } }}
>
<div className="absolute inset-0 bg-zinc-950/70 backdrop-blur-sm animate-fade-in" />
<div className="relative w-96 rounded-2xl bg-zinc-900 border border-accent/20 shadow-2xl shadow-accent/5 overflow-hidden animate-slide-up-center">
<div className="px-5 py-5 flex flex-col gap-4">
<div className="flex items-start gap-3">
<div className="w-10 h-10 rounded-xl flex items-center justify-center shrink-0 bg-accent/10 border border-accent/20">
<svg width="17" height="17" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round" className="text-accent-light">
<polyline points="3 6 5 6 21 6" />
<path d="M19 6l-1 14a2 2 0 01-2 2H8a2 2 0 01-2-2L5 6" />
<path d="M10 11v6M14 11v6" />
<path d="M9 6V4a1 1 0 011-1h4a1 1 0 011 1v2" />
</svg>
</div>
<div className="flex flex-col gap-1 pt-0.5">
<h2 className="text-base font-semibold text-zinc-100 leading-tight">
Uninstall extension &ldquo;{ext?.name ?? uninstallTarget}&rdquo;?
</h2>
<p className="text-xs text-zinc-500 leading-relaxed">
The extension folder will be deleted.
</p>
</div>
</div>

{installedModels.length > 0 && (
<div className="flex flex-col gap-2 px-1">
<p className="text-[11px] font-medium text-zinc-400">
Also delete downloaded model weights:
</p>
{installedModels.map((v) => {
const checked = modelsToDelete.has(v.id)
return (
<label
key={v.id}
className="flex items-center gap-2.5 px-3 py-2 rounded-lg bg-zinc-800/60 border border-zinc-700/40 cursor-pointer hover:border-zinc-600/60 transition-colors"
>
<input
type="checkbox"
checked={checked}
onChange={() => {
setModelsToDelete((prev) => {
const next = new Set(prev)
if (checked) next.delete(v.id)
else next.add(v.id)
return next
})
}}
className="accent-accent w-3.5 h-3.5 rounded"
/>
<span className="text-xs text-zinc-200">{formatModelName(v.id)}</span>
</label>
)
})}
</div>
)}

<div className="flex gap-2.5">
<button
onClick={() => { setUninstallTarget(null); setModelsToDelete(new Set()) }}
className="flex-1 py-2.5 rounded-xl bg-zinc-800 hover:bg-zinc-700/80 text-zinc-400 hover:text-zinc-200 text-sm font-medium transition-colors border border-zinc-700/50"
>
Cancel
</button>
<button
onClick={() => handleUninstallExtension(uninstallTarget)}
className="flex-1 py-2.5 rounded-xl bg-accent hover:bg-accent-dark text-white text-sm font-semibold transition-colors shadow-lg shadow-accent/20"
>
Uninstall
</button>
</div>
</div>
</div>
</div>,
document.body
)
})()}
</div>
)
}