towsty/aigen
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Add a Sleep GPU button that stops Comfy to free VRAM.

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Towsty
2026-08-30 19:19:41 -05:00
co-authored by Cursor
parent 3a90892859
commit 64179dd516
4 changed files with 184 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
components/SettingsModal.vue
+66
View File
@@ -10,6 +10,36 @@
<button type="button" class="rounded-full border border-white/10 px-3 py-1 text-sm text-zinc-400 hover:text-white" @click="emit('close')">Close</button>
</div>
<section class="mt-6">
<h3 class="text-sm font-semibold uppercase tracking-wider text-zinc-500">GPU</h3>
<div class="mt-2 space-y-3 rounded-2xl border border-white/10 bg-zinc-950/50 px-4 py-3">
<p class="text-sm text-zinc-400">
Stops Comfy on the desktop GPU so it drops out of VRAM. The PC stays on. Generate or Wake GPU will start it again.
</p>
<div class="flex flex-wrap items-center gap-2">
<button
type="button"
class="rounded-xl border px-3 py-2 text-sm disabled:cursor-not-allowed disabled:opacity-60"
:class="gpuSleeping ? 'border-amber-300/40 text-amber-100' : 'border-white/15 text-zinc-200 hover:border-amber-300/50 hover:text-white'"
:disabled="gpuSleeping || gpuWaking"
@click="sleepGpu"
>
{{ gpuSleeping ? 'Stopping GPU…' : 'Sleep GPU' }}
</button>
<button
type="button"
class="rounded-xl border px-3 py-2 text-sm disabled:cursor-not-allowed disabled:opacity-60"
:class="gpuWaking ? 'border-amber-300/40 text-amber-100' : 'border-amber-400/40 text-amber-200 hover:border-amber-300 hover:text-amber-50'"
:disabled="gpuWaking || gpuSleeping"
@click="wakeGpu"
>
{{ gpuWaking ? 'Starting GPU…' : 'Wake GPU' }}
</button>
</div>
<p v-if="gpuMessage" class="text-xs" :class="gpuError ? 'text-red-300' : 'text-zinc-500'">{{ gpuMessage }}</p>
</div>
</section>
<section class="mt-6">
<h3 class="text-sm font-semibold uppercase tracking-wider text-zinc-500">Storage</h3>
<div class="mt-2 rounded-2xl border border-white/10 bg-zinc-950/50 px-4 py-3">
@@ -256,6 +286,10 @@ const loginNext = ref('')
const drafts = reactive<Record<string, { name: string; protect: boolean; password: string; currentPassword: string }>>({})
const storageLoading = ref(true)
const storageError = ref('')
const gpuSleeping = ref(false)
const gpuWaking = ref(false)
const gpuError = ref(false)
const gpuMessage = ref('')
const storage = ref<{
libraryBytes: number
breakdown: { clips: number; stills: number; tmp: number; other: number }
@@ -392,6 +426,38 @@ async function loadStorage() {
}
}
async function sleepGpu() {
if (gpuSleeping.value || gpuWaking.value) return
gpuSleeping.value = true
gpuError.value = false
gpuMessage.value = 'Stopping Comfy on the GPU…'
try {
await $fetch('/api/comfy/sleep', { method: 'POST' })
gpuMessage.value = 'Comfy is stopped. VRAM is free until you generate or wake the GPU.'
} catch (error: any) {
gpuError.value = true
gpuMessage.value = error?.data?.statusMessage || error?.statusMessage || 'Could not stop Comfy.'
} finally {
gpuSleeping.value = false
}
}
async function wakeGpu() {
if (gpuWaking.value || gpuSleeping.value) return
gpuWaking.value = true
gpuError.value = false
gpuMessage.value = 'Starting Comfy on the GPU…'
try {
await $fetch('/api/comfy/wake', { method: 'POST' })
gpuMessage.value = 'Wake sent. The first job can still take a couple of minutes if Comfy is cold.'
} catch (error: any) {
gpuError.value = true
gpuMessage.value = error?.data?.statusMessage || error?.statusMessage || 'Could not wake the GPU.'
} finally {
gpuWaking.value = false
}
}
async function loadLoraKeys() {
loraRowsLoading.value = true
try {
pages/index.vue
+42 -2
View File
@@ -30,6 +30,17 @@
>
{{ gpuWaking ? 'Starting GPU…' : 'Wake GPU' }}
</button>
<button
v-if="showSleepGpu"
type="button"
class="rounded-full border px-3 py-1 hover:border-amber-300/50 hover:text-white disabled:cursor-not-allowed disabled:opacity-80"
:class="gpuSleeping ? 'border-amber-300/40 text-amber-100' : 'border-white/15 text-zinc-300'"
:disabled="gpuSleeping || rigBusy"
:title="gpuWakeError || (rigBusy ? 'Cancel the running job first' : (gpuSleeping ? 'Stopping Comfy on the GPU' : 'Stop Comfy on the GPU to free VRAM. Does not shut the PC down.'))"
@click="sleepGpu"
>
{{ gpuSleeping ? 'Stopping GPU…' : 'Sleep GPU' }}
</button>
</div>
<span class="hidden md:block text-zinc-400">{{ userName }}</span>
<button type="button" class="rounded-full border border-white/10 px-3 py-1 text-zinc-300 hover:border-amber-300/50 hover:text-white" @click="settingsOpen = true">
@@ -2453,6 +2464,7 @@ const comfyOk = ref(false)
const comfyBusy = ref(false)
const comfyAsleep = ref(false)
const gpuWaking = ref(false)
const gpuSleeping = ref(false)
const gpuWakeError = ref('')
const userName = ref('')
const authEnabled = ref(false)
@@ -2726,8 +2738,17 @@ const liveMaxStep = computed(() => (
const rigBusy = computed(() => (
comfyBusy.value || imageComfyBusy.value || videoBusy.value || editBusy.value
))
const showWakeGpu = computed(() => gpuWaking.value || (!comfyOk.value && !imageComfyOk.value))
const showWakeGpu = computed(() => gpuWaking.value || (!gpuSleeping.value && !comfyOk.value && !imageComfyOk.value))
const showSleepGpu = computed(() => gpuSleeping.value || ((comfyOk.value || imageComfyOk.value) && !gpuWaking.value))
const rigBadge = computed(() => {
if (gpuSleeping.value) {
return {
label: 'Stopping…',
className: 'border-amber-400/40 text-amber-200',
dot: 'bg-amber-400 animate-pulse',
title: 'Stopping Comfy on the GPU'
}
}
if (gpuWaking.value) {
return {
label: 'Starting…',
@@ -2857,6 +2878,7 @@ const recommendHint = computed(() => {
return 'Send this still and draft to Beast Qwen VL.'
})
const gpuSleepHint = computed(() => {
if (gpuSleeping.value) return 'Stopping Comfy on the GPU to free VRAM. The PC stays on.'
if (comfyOk.value || imageComfyOk.value) return ''
if (gpuWaking.value) return 'Starting Comfy on the GPU. Generate will wait in the job queue until it answers.'
if (comfyAsleep.value) return 'GPU is asleep. Generate will wake it — the first job can take a few minutes. Or click Wake GPU to pre-warm.'
@@ -3528,7 +3550,7 @@ async function pollComfyHealth() {
}
async function pokeGpu() {
if (gpuWaking.value) return
if (gpuWaking.value || gpuSleeping.value) return
gpuWaking.value = true
gpuWakeError.value = ''
try {
@@ -3550,6 +3572,24 @@ async function pokeGpu() {
}
}
async function sleepGpu() {
if (gpuSleeping.value || gpuWaking.value || rigBusy.value) return
gpuSleeping.value = true
gpuWakeError.value = ''
try {
await $fetch('/api/comfy/sleep', { method: 'POST' })
comfyOk.value = false
imageComfyOk.value = false
imageBeastOk.value = false
comfyAsleep.value = true
await pollComfyHealth()
} catch (error: any) {
gpuWakeError.value = error?.data?.statusMessage || error?.statusMessage || error?.message || 'Could not stop Comfy.'
} finally {
gpuSleeping.value = false
}
}
onMounted(async () => {
try {
withSound.value = localStorage.getItem('aigen-generate-sound') !== 'false'
server/api/comfy/sleep.post.ts
+22
View File
@@ -0,0 +1,22 @@
import { listJobs } from '~/server/utils/jobs'
import { listOwnersWithStudioQueues, listStudioJobs } from '~/server/utils/studioQueue'
import { requestComfySleep } from '~/server/utils/comfyLifecycle'
export default defineEventHandler(async () => {
const live = listJobs().filter(job => job.status === 'queued' || job.status === 'uploading' || job.status === 'running')
if (live.length) {
throw createError({
statusCode: 409,
statusMessage: 'A job is still running. Cancel it first, then sleep the GPU.'
})
}
for (const owner of listOwnersWithStudioQueues()) {
if (listStudioJobs(owner).some(job => job.status === 'running')) {
throw createError({
statusCode: 409,
statusMessage: 'A job is still running. Cancel it first, then sleep the GPU.'
})
}
}
return await requestComfySleep()
})
server/utils/comfyLifecycle.ts
+54 -4
View File
@@ -81,7 +81,7 @@ export async function fetchQueue() {
}
}
async function controlRequest(path: string, method = 'GET') {
async function controlRequest(path: string, method = 'GET', timeoutMs = 5000) {
const { controlUrl, controlToken } = settings()
if (!controlUrl) return null
try {
@@ -91,10 +91,9 @@ async function controlRequest(path: string, method = 'GET') {
Accept: 'application/json',
...(controlToken ? { Authorization: `Bearer ${controlToken}` } : {})
},
signal: AbortSignal.timeout(5000)
signal: AbortSignal.timeout(timeoutMs)
})
if (!res.ok) return { ok: false, status: res.status }
return await res.json() as {
const body = await res.json().catch(() => null) as {
ok?: boolean
http?: boolean
process?: boolean
@@ -108,9 +107,23 @@ async function controlRequest(path: string, method = 'GET') {
port?: number | null
proxyPort?: number | null
stopped?: boolean
killed?: unknown
error?: string
running?: number
pending?: number
status?: number
} | null
if (!res.ok) {
return {
ok: false as const,
status: res.status,
error: body?.error,
running: body?.running,
pending: body?.pending,
killed: body?.killed
}
}
return body || { ok: true }
} catch {
return null
}
@@ -151,6 +164,43 @@ export async function requestComfyWake() {
}
}
export async function requestComfySleep() {
const { controlUrl } = settings()
if (!controlUrl) {
throw createError({
statusCode: 503,
statusMessage: 'COMFY_CONTROL_URL is not set. The studio cannot poke the Beast host agent.'
})
}
const remote = await controlRequest('/stop', 'POST', 25_000)
if (!remote) {
throw createError({
statusCode: 503,
statusMessage: 'Beast host agent is unreachable. Confirm it is running on the desktop.'
})
}
if (remote.status === 409 || remote.error === 'queue-busy') {
const running = Number(remote.running || 0)
const pending = Number(remote.pending || 0)
throw createError({
statusCode: 409,
statusMessage: `Comfy is still working (${running} running, ${pending} queued). Cancel the job first, then sleep the GPU.`
})
}
if (remote.ok === false) {
throw createError({
statusCode: Number(remote.status) || 502,
statusMessage: 'Beast host agent refused to stop Comfy.'
})
}
return {
ok: true,
stopped: true,
asleep: true,
killed: remote.killed ?? null
}
}
async function findLocalComfyProcess() {
const platform = process.platform
try {