From b815ad01ff674d963b822df90c134eb5d5c823f8 Mon Sep 17 00:00:00 2001
From: patriceac <48073125+patriceac@users.noreply.github.com>
Date: Sun, 21 Jul 2024 01:33:58 -0700
Subject: [PATCH 1/2] 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.
---
 ui/easydiffusion/server.py | 42 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/ui/easydiffusion/server.py b/ui/easydiffusion/server.py
index a251ede6..ef74be67 100644
--- a/ui/easydiffusion/server.py
+++ b/ui/easydiffusion/server.py
@@ -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):

From ab3e20fbd67af8cb752ebd83d1fd96e9f0df8082 Mon Sep 17 00:00:00 2001
From: patriceac <48073125+patriceac@users.noreply.github.com>
Date: Sun, 21 Jul 2024 01:37:25 -0700
Subject: [PATCH 2/2] Removing duplicate import

Removing duplicate import.
---
 ui/easydiffusion/server.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/ui/easydiffusion/server.py b/ui/easydiffusion/server.py
index ef74be67..82184adb 100644
--- a/ui/easydiffusion/server.py
+++ b/ui/easydiffusion/server.py
@@ -31,8 +31,6 @@ 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