Organize code

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,15 +1,15 @@
 import os
+import ctypes
 import json
+import subprocess
 import time
 import traceback
 import threading
-import platform
-import subprocess
 import gradio as gr
 
 from datetime import datetime, timezone
 from pydantic import BaseModel
-from typing import Any, Callable, NoReturn, Union, Optional, List, Dict
+from typing import Any, Callable, Union, Optional, List, Dict
 from fastapi import FastAPI
 from PIL import Image
 
@@ -29,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,
@@ -41,9 +44,6 @@ from .task_helpers import (
     map_named_args_to_ui_task_args_list,
 )
 
-is_windows = platform.system() == "Windows"
-is_macos = platform.system() == "Darwin"
-
 
 class OutOfMemoryError(Exception):
     def __init__(self, message="CUDA out of memory") -> None:
@@ -564,7 +564,6 @@ class TaskRunner:
         elif action == "Sleep":
             log.info("[AgentScheduler] Sleeping...")
             if is_windows:
-                import ctypes
                 if not ctypes.windll.PowrProf.SetSuspendState(False, False, False):
                     print(f"Couldn't sleep: {ctypes.GetLastError()}")
             elif is_macos:
@@ -579,8 +578,8 @@ class TaskRunner:
                 command = ["osascript", "-e", 'tell application "Finder" to sleep']
             else:
                 command = ["systemctl", "hibernate"]
-        elif action == "Quit WebUI":
-            log.info("[AgentScheduler] Quitting WebUI...")
+        elif action == "Stop webui":
+            log.info("[AgentScheduler] Stopping webui...")
             _exit(0)
 
         if command:
@@ -647,12 +646,3 @@ def get_instance(block) -> TaskRunner:
             script_callbacks.on_before_reload(on_before_reload)
 
     return TaskRunner.instance
-
-
-def _exit(status: int) -> NoReturn:
-    try:
-        import atexit
-        atexit._run_exitfuncs()
-    except:
-        pass
-    os._exit(status)

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,11 +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", "Quit WebUI"]
+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",
@@ -749,7 +747,7 @@ def on_ui_settings():
             "Action after queue completion",
             gr.Radio,
             lambda: {
-                "choices": completion_action_choices
+                "choices": completion_action_choices,
             },
             section=section,
         ),