Merge pull request #191 from MMP0/queue-completion

Add option for action after queue completion

agent_scheduler/helpers.py

@@ -1,14 +1,20 @@
+import os
 import sys
 import abc
+import atexit
 import time
 import logging
+import platform
 import requests
 import traceback
-from typing import Callable, List
+from typing import Callable, List, NoReturn
 
 import gradio as gr
 from gradio.blocks import Block, BlockContext
 
+is_windows = platform.system() == "Windows"
+is_macos = platform.system() == "Darwin"
+
 if logging.getLogger().hasHandlers():
     log = logging.getLogger("sd")
 else:
@@ -183,4 +189,14 @@ def request_with_retry(
         log.error("[ArtVenture] Error while uploading result")
         log.error(e)
         log.debug(traceback.format_exc())
-        return False
+        return False
+
+
+def _exit(status: int) -> NoReturn:
+    try:
+        atexit._run_exitfuncs()
+    except:
+        pass
+    sys.stdout.flush()
+    sys.stderr.flush()
+    os._exit(status)

agent_scheduler/task_runner.py

@@ -1,5 +1,7 @@
 import os
+import ctypes
 import json
+import subprocess
 import time
 import traceback
 import threading
@@ -27,6 +29,9 @@ from .helpers import (
     detect_control_net,
     get_component_by_elem_id,
     get_dict_attribute,
+    is_windows,
+    is_macos,
+    _exit,
 )
 from .task_helpers import (
     encode_image_to_base64,
@@ -404,6 +409,9 @@ class TaskRunner:
 
             task = get_next_task()
             if not task:
+                if not self.paused:
+                    time.sleep(1)
+                    self.__on_completed()
                 break
 
     def execute_pending_tasks_threading(self):
@@ -529,6 +537,56 @@ class TaskRunner:
             self.__saved_images_path.insert(0, data.filename)
         else:
             self.__saved_images_path.append(data.filename)
+    
+    def __on_completed(self):
+        action = getattr(shared.opts, "queue_completion_action", "Do nothing")
+
+        if action == "Do nothing":
+            return
+
+        command = None
+        if action == "Shut down":
+            log.info("[AgentScheduler] Shutting down...")
+            if is_windows:
+                command = ["shutdown", "/s", "/hybrid", "/t", "0"]
+            elif is_macos:
+                command = ["osascript", "-e", 'tell application "Finder" to shut down']
+            else:
+                command = ["systemctl", "poweroff"]
+        elif action == "Restart":
+            log.info("[AgentScheduler] Restarting...")
+            if is_windows:
+                command = ["shutdown", "/r", "/t", "0"]
+            elif is_macos:
+                command = ["osascript", "-e", 'tell application "Finder" to restart']
+            else:
+                command = ["systemctl", "reboot"]
+        elif action == "Sleep":
+            log.info("[AgentScheduler] Sleeping...")
+            if is_windows:
+                if not ctypes.windll.PowrProf.SetSuspendState(False, False, False):
+                    print(f"Couldn't sleep: {ctypes.GetLastError()}")
+            elif is_macos:
+                command = ["osascript", "-e", 'tell application "Finder" to sleep']
+            else:
+                command = ["sh", "-c", 'systemctl hybrid-sleep || (echo "Couldn\'t hybrid sleep, will try to suspend instead: $?"; systemctl suspend)']
+        elif action == "Hibernate":
+            log.info("[AgentScheduler] Hibernating...")
+            if is_windows:
+                command = ["shutdown", "/h"]
+            elif is_macos:
+                command = ["osascript", "-e", 'tell application "Finder" to sleep']
+            else:
+                command = ["systemctl", "hibernate"]
+        elif action == "Stop webui":
+            log.info("[AgentScheduler] Stopping webui...")
+            _exit(0)
+
+        if command:
+            subprocess.Popen(command)
+
+        if action in {"Shut down", "Restart"}:
+            _exit(0)
 
     def on_task_registered(self, callback: Callable):
         """Callback when a task is registered

scripts/task_scheduler.py

@@ -1,6 +1,5 @@
 import os
 import json
-import platform
 import gradio as gr
 from PIL import Image
 from uuid import uuid4
@@ -23,7 +22,7 @@ from modules.generation_parameters_copypaste import (
 )
 
 from agent_scheduler.task_runner import TaskRunner, get_instance
-from agent_scheduler.helpers import log, compare_components_with_ids, get_components_by_ids
+from agent_scheduler.helpers import log, compare_components_with_ids, get_components_by_ids, is_macos
 from agent_scheduler.db import init as init_db, task_manager, TaskStatus
 from agent_scheduler.api import regsiter_apis
 
@@ -42,9 +41,10 @@ ui_placement_append_to_main = "Append to main UI"
 placement_under_generate = "Under Generate button"
 placement_between_prompt_and_generate = "Between Prompt and Generate button"
 
+completion_action_choices = ["Do nothing", "Shut down", "Restart", "Sleep", "Hibernate", "Stop webui"]
+
 task_filter_choices = ["All", "Bookmarked", "Done", "Failed", "Interrupted"]
 
-is_macos = platform.system() == "Darwin"
 enqueue_key_modifiers = [
     "Command" if is_macos else "Ctrl",
     "Control" if is_macos else "Alt",
@@ -740,6 +740,18 @@ def on_ui_settings():
             section=section,
         ),
     )
+    shared.opts.add_option(
+        "queue_completion_action",
+        shared.OptionInfo(
+            "Do nothing",
+            "Action after queue completion",
+            gr.Radio,
+            lambda: {
+                "choices": completion_action_choices,
+            },
+            section=section,
+        ),
+    )
 
 
 def on_app_started(block: gr.Blocks, app):