From d2e12366eb2e3b036208cb5309a08a79bf63cf7e Mon Sep 17 00:00:00 2001
From: Lightning Pixel <lightning.pixel.dev@gmail.com>
Date: Sun, 22 Mar 2026 16:42:19 +0100
Subject: [PATCH] fix(generate): improve memory release on unload and disable
 free-memory during generation

---
 api/routers/model.py                 |  10 +++
 api/services/generators/base.py      |  19 +++++
 api/services/generators/hunyuan3d.py |   9 ---
 electron/main/python-bridge.ts       |  11 ++-
 src/areas/generate/GeneratePage.tsx  |   3 +-
 src/areas/models/ModelsPage.tsx      | 112 +++++++++++++++++++++++----
 6 files changed, 140 insertions(+), 24 deletions(-)

diff --git a/api/routers/model.py b/api/routers/model.py
index 2e3bee0..6ababd5 100644
--- a/api/routers/model.py
+++ b/api/routers/model.py
@@ -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}
 
 
diff --git a/api/services/generators/base.py b/api/services/generators/base.py
index 0180f93..a23b01d 100644
--- a/api/services/generators/base.py
+++ b/api/services/generators/base.py
@@ -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
diff --git a/api/services/generators/hunyuan3d.py b/api/services/generators/hunyuan3d.py
index 2a2b247..5f63570 100644
--- a/api/services/generators/hunyuan3d.py
+++ b/api/services/generators/hunyuan3d.py
@@ -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
     # ------------------------------------------------------------------ #
diff --git a/electron/main/python-bridge.ts b/electron/main/python-bridge.ts
index dde2d66..d04d028 100644
--- a/electron/main/python-bridge.ts
+++ b/electron/main/python-bridge.ts
@@ -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
@@ -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 })
       }
     })
@@ -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
diff --git a/src/areas/generate/GeneratePage.tsx b/src/areas/generate/GeneratePage.tsx
index 9bb1747..4aee54c 100644
--- a/src/areas/generate/GeneratePage.tsx
+++ b/src/areas/generate/GeneratePage.tsx
@@ -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" />
diff --git a/src/areas/models/ModelsPage.tsx b/src/areas/models/ModelsPage.tsx
index 701f641..12198ae 100644
--- a/src/areas/models/ModelsPage.tsx
+++ b/src/areas/models/ModelsPage.tsx
@@ -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'
@@ -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)
@@ -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 ────────────────────────────────────────────────────────────────
@@ -287,7 +301,7 @@ export default function ModelsPage(): JSX.Element {
                       }
                     })
                   }}
-                  onUninstall={(extId) => setUninstallTarget(extId)}
+                  onUninstall={(extId) => openUninstallModal(extId)}
                 />
               ))}
             </div>
@@ -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>
   )
 }