Prevent Computer from Sleeping During Image Generation

This pull request introduces functionality to prevent the computer from going to sleep while images are being generated, ensuring consistent performance across Windows, macOS, and Linux. The changes involve integrating system-level calls to manage the sleep state within the `render_internal` function.

**Changes:**
- Added `PreventSleep` class to handle sleep prevention across different operating systems:
  - **Windows**: Utilizes `ctypes` to call `SetThreadExecutionState` to prevent and allow sleep.
  - **macOS**: Uses the `caffeinate` command to prevent sleep.
  - **Linux**: Uses `systemd-inhibit` to prevent sleep.
- Integrated these functions into the `render_internal` function to prevent the system from sleeping during the image generation process and to allow it to sleep once the process is complete.

This enhancement ensures uninterrupted image generation, improving the reliability of the application during long-running tasks on all supported platforms.

**Note**: I only tested this functionality on Windows.
This commit is contained in:
patriceac 2024-07-21 01:33:58 -07:00 committed by GitHub
parent 364902f8a1
commit b815ad01ff
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -7,6 +7,9 @@ import mimetypes
import os
import traceback
from typing import List, Union
import platform
import subprocess
import ctypes
from easydiffusion import app, model_manager, task_manager, package_manager
from easydiffusion.tasks import RenderTask, FilterTask
@ -28,6 +31,13 @@ from pydantic import BaseModel, Extra
from starlette.responses import FileResponse, JSONResponse, StreamingResponse
from pycloudflared import try_cloudflare
import ctypes
# Constants for preventing sleep
ES_CONTINUOUS = 0x80000000
ES_SYSTEM_REQUIRED = 0x00000001
ES_DISPLAY_REQUIRED = 0x00000002
log.info(f"started in {app.SD_DIR}")
log.info(f"started at {datetime.datetime.now():%x %X}")
@ -40,6 +50,34 @@ NOCACHE_HEADERS = {
}
PROTECTED_CONFIG_KEYS = ("block_nsfw",) # can't change these via the HTTP API
# Constants for preventing sleep on Windows
ES_CONTINUOUS = 0x80000000
ES_SYSTEM_REQUIRED = 0x00000001
ES_DISPLAY_REQUIRED = 0x00000002
class PreventSleep:
def __init__(self):
self.os_name = platform.system()
self.inhibit_process = None
def prevent_sleep(self):
if self.os_name == "Windows":
ctypes.windll.kernel32.SetThreadExecutionState(
ES_CONTINUOUS | ES_SYSTEM_REQUIRED | ES_DISPLAY_REQUIRED)
elif self.os_name == "Darwin":
self.inhibit_process = subprocess.Popen(['caffeinate'])
elif self.os_name == "Linux":
self.inhibit_process = subprocess.Popen(['systemd-inhibit', '--what=handle-lid-switch', '--why="Prevent sleep during image generation"', 'bash', '-c', 'sleep infinity'])
else:
raise NotImplementedError("Unsupported OS")
def allow_sleep(self):
if self.os_name == "Windows":
ctypes.windll.kernel32.SetThreadExecutionState(ES_CONTINUOUS)
elif self.inhibit_process:
self.inhibit_process.terminate()
self.inhibit_process = None
class NoCacheStaticFiles(StaticFiles):
def __init__(self, directory: str):
@ -267,7 +305,9 @@ def ping_internal(session_id: str = None):
def render_internal(req: dict):
prevent_sleep_obj = PreventSleep()
try:
prevent_sleep_obj.prevent_sleep()
req = convert_legacy_render_req_to_new(req)
# separate out the request data into rendering and task-specific data
@ -299,6 +339,8 @@ def render_internal(req: dict):
except Exception as e:
log.error(traceback.format_exc())
raise HTTPException(status_code=500, detail=str(e))
finally:
prevent_sleep_obj.allow_sleep()
def filter_internal(req: dict):