import os import json import time import traceback import threading from pydantic import BaseModel from datetime import datetime, timedelta from typing import Any, Callable, Union, Optional from fastapi import FastAPI from PIL import Image from modules import progress, shared, script_callbacks from modules.call_queue import queue_lock, wrap_gradio_call from modules.txt2img import txt2img from modules.img2img import img2img from modules.api.api import Api from modules.api.models import ( StableDiffusionTxt2ImgProcessingAPI, StableDiffusionImg2ImgProcessingAPI, ) from .db import TaskStatus, Task, task_manager from .helpers import ( log, detect_control_net, get_component_by_elem_id, get_dict_attribute, ) from .task_helpers import ( serialize_image, deserialize_image, encode_image_to_base64, serialize_img2img_image_args, deserialize_img2img_image_args, serialize_controlnet_args, deserialize_controlnet_args, serialize_api_task_args, map_ui_task_args_list_to_named_args, map_named_args_to_ui_task_args_list, ) class OutOfMemoryError(Exception): def __init__(self, message="CUDA out of memory") -> None: self.message = message super().__init__(message) task_history_retenion_map = { "7 days": 7, "14 days": 14, "30 days": 30, "90 days": 90, "Keep forever": 0, } class ParsedTaskArgs(BaseModel): is_ui: bool ui_args: list[Any] named_args: dict[str, Any] script_args: list[Any] checkpoint: Optional[str] = None class TaskRunner: instance = None def __init__(self, UiControlNetUnit=None): self.UiControlNetUnit = UiControlNetUnit self.__total_pending_tasks: int = 0 self.__current_thread: threading.Thread = None self.__api = Api(FastAPI(), queue_lock) self.__saved_images_path: list[tuple[str, str]] = [] script_callbacks.on_image_saved(self.__on_image_saved) self.script_callbacks = { "task_registered": [], "task_started": [], "task_finished": [], "task_cleared": [], } # Mark this to True when reload UI self.dispose = False self.interrupted = None if TaskRunner.instance is not None: raise Exception("TaskRunner instance already exists") TaskRunner.instance = self @property def current_task_id(self) -> Union[str, None]: return progress.current_task @property def is_executing_task(self) -> bool: return self.__current_thread and self.__current_thread.is_alive() @property def paused(self) -> bool: return getattr(shared.opts, "queue_paused", False) def __serialize_ui_task_args(self, is_img2img: bool, *args, checkpoint: str = None): named_args, script_args = map_ui_task_args_list_to_named_args( list(args), is_img2img, checkpoint=checkpoint ) # loop through named_args and serialize images if is_img2img: serialize_img2img_image_args(named_args) # loop through script_args and serialize images for i, a in enumerate(script_args): if isinstance(a, Image.Image): script_args[i] = serialize_image(a) elif self.UiControlNetUnit and isinstance(a, self.UiControlNetUnit): script_args[i] = serialize_controlnet_args(a) return json.dumps( { "args": named_args, "script_args": script_args, "checkpoint": checkpoint, "is_ui": True, "is_img2img": is_img2img, } ) def __serialize_api_task_args( self, is_img2img: bool, script_args: list = [], checkpoint: str = None, **api_args, ): named_args = serialize_api_task_args( api_args, is_img2img, checkpoint=checkpoint ) checkpoint = get_dict_attribute( named_args, "override_settings.sd_model_checkpoint", None ) return json.dumps( { "args": named_args, "script_args": script_args, "checkpoint": checkpoint, "is_ui": False, "is_img2img": is_img2img, } ) def __deserialize_ui_task_args( self, is_img2img: bool, named_args: dict, script_args: list ): # loop through image_args and deserialize images if is_img2img: deserialize_img2img_image_args(named_args) # loop through script_args and deserialize images for i, arg in enumerate(script_args): if isinstance(arg, dict) and arg.get("is_cnet", False): script_args[i] = deserialize_controlnet_args(arg) elif isinstance(arg, dict) and arg.get("cls", "") in {"Image", "ndarray"}: script_args[i] = deserialize_image(arg) def __deserialize_api_task_args(self, is_img2img: bool, named_args: dict): # load images from disk if is_img2img: init_images = named_args.get("init_images") for i, img in enumerate(init_images): if isinstance(img, str) and os.path.isfile(img): image = Image.open(img) init_images[i] = encode_image_to_base64(image) named_args.update({"save_images": True, "send_images": False}) def parse_task_args(self, task: Task, deserialization: bool = True): parsed: dict[str, Any] = json.loads(task.params) is_ui = parsed.get("is_ui", True) is_img2img = parsed.get("is_img2img", None) checkpoint = parsed.get("checkpoint", None) named_args: dict[str, Any] = parsed["args"] script_args: list[Any] = parsed.get("script_args", []) if is_ui and deserialization: self.__deserialize_ui_task_args(is_img2img, named_args, script_args) elif deserialization: self.__deserialize_api_task_args(is_img2img, named_args) ui_args = ( map_named_args_to_ui_task_args_list(named_args, script_args, is_img2img) if is_ui else [] ) return ParsedTaskArgs( is_ui=is_ui, ui_args=ui_args, named_args=named_args, script_args=script_args, checkpoint=checkpoint, ) def register_ui_task( self, task_id: str, is_img2img: bool, *args, checkpoint: str = None ): progress.add_task_to_queue(task_id) params = self.__serialize_ui_task_args(is_img2img, *args, checkpoint=checkpoint) task_type = "img2img" if is_img2img else "txt2img" task_manager.add_task(Task(id=task_id, type=task_type, params=params)) self.__run_callbacks( "task_registered", task_id, is_img2img=is_img2img, is_ui=True, args=params ) self.__total_pending_tasks += 1 def register_api_task( self, task_id: str, api_task_id: str, is_img2img: bool, args: dict, checkpoint: str = None, ): progress.add_task_to_queue(task_id) params = self.__serialize_api_task_args( is_img2img, checkpoint=checkpoint, **args ) task_type = "img2img" if is_img2img else "txt2img" task_manager.add_task( Task(id=task_id, api_task_id=api_task_id, type=task_type, params=params) ) self.__run_callbacks( "task_registered", task_id, is_img2img=is_img2img, is_ui=False, args=params ) self.__total_pending_tasks += 1 def execute_task(self, task: Task, get_next_task: Callable[[], Task]): while True: if self.dispose: break if progress.current_task is None: task_id = task.id is_img2img = task.type == "img2img" log.info(f"[AgentScheduler] Executing task {task_id}") task_args = self.parse_task_args(task) task_meta = { "is_img2img": is_img2img, "is_ui": task_args.is_ui, "api_task_id": task.api_task_id, } self.interrupted = None self.__saved_images_path = [] self.__run_callbacks("task_started", task_id, **task_meta) # enable image saving samples_save = shared.opts.samples_save shared.opts.samples_save = True res = self.__execute_task(task_id, is_img2img, task_args) # disable image saving shared.opts.samples_save = samples_save if not res or isinstance(res, Exception): if isinstance(res, OutOfMemoryError): log.error( f"[AgentScheduler] Task {task_id} failed: CUDA OOM. Queue will be paused." ) shared.opts.queue_paused = True else: log.error(f"[AgentScheduler] Task {task_id} failed: {res}") log.debug(traceback.format_exc()) task_manager.update_task( id=task_id, status=TaskStatus.FAILED, result=str(res) if res else None, ) self.__run_callbacks( "task_finished", task_id, status=TaskStatus.FAILED, **task_meta ) else: is_interrupted = self.interrupted == task_id if is_interrupted: log.info(f"\n[AgentScheduler] Task {task.id} interrupted") task_manager.update_task( id=task_id, status=TaskStatus.INTERRUPTED, ) self.__run_callbacks( "task_finished", task_id, status=TaskStatus.INTERRUPTED, **task_meta, ) else: result = { "images": [], "infotexts": [], } for filename, pnginfo in self.__saved_images_path: result["images"].append(filename) result["infotexts"].append(pnginfo) task_manager.update_task( id=task_id, status=TaskStatus.DONE, result=json.dumps(result), ) self.__run_callbacks( "task_finished", task_id, status=TaskStatus.DONE, result=result, **task_meta, ) self.__saved_images_path = [] else: time.sleep(2) continue task = get_next_task() if not task: break def execute_pending_tasks_threading(self): if self.paused: log.info("[AgentScheduler] Runner is paused") return if self.is_executing_task: log.info("[AgentScheduler] Runner already started") return pending_task = self.__get_pending_task() if pending_task: # Start the infinite loop in a separate thread self.__current_thread = threading.Thread( target=self.execute_task, args=( pending_task, self.__get_pending_task, ), ) self.__current_thread.daemon = True self.__current_thread.start() def __execute_task(self, task_id: str, is_img2img: bool, task_args: ParsedTaskArgs): if task_args.is_ui: return self.__execute_ui_task(task_id, is_img2img, *task_args.ui_args) else: return self.__execute_api_task( task_id, is_img2img, **task_args.named_args, ) def __execute_ui_task(self, task_id: str, is_img2img: bool, *args): func = wrap_gradio_call(img2img if is_img2img else txt2img, add_stats=True) with queue_lock: shared.state.begin() progress.start_task(task_id) res = None try: result = func(*args) if ( result[0] is None and hasattr(shared.state, "oom") and shared.state.oom ): res = OutOfMemoryError() elif "CUDA out of memory" in result[2]: res = OutOfMemoryError() else: res = result[1] except Exception as e: res = e finally: progress.finish_task(task_id) shared.state.end() return res def __execute_api_task(self, task_id: str, is_img2img: bool, **kwargs): progress.start_task(task_id) res = None try: result = ( self.__api.img2imgapi(StableDiffusionImg2ImgProcessingAPI(**kwargs)) if is_img2img else self.__api.text2imgapi( StableDiffusionTxt2ImgProcessingAPI(**kwargs) ) ) res = result.info except Exception as e: if "CUDA out of memory" in str(e): res = OutOfMemoryError() else: res = e finally: progress.finish_task(task_id) return res def __get_pending_task(self): if self.dispose: return None if self.paused: log.info("[AgentScheduler] Runner is paused") return None # delete task that are too old retention_days = 30 if ( getattr(shared.opts, "queue_history_retention_days", None) and shared.opts.queue_history_retention_days in task_history_retenion_map ): retention_days = task_history_retenion_map[ shared.opts.queue_history_retention_days ] if retention_days > 0: deleted_rows = task_manager.delete_tasks_before( datetime.now() - timedelta(days=retention_days) ) if deleted_rows > 0: log.debug( f"[AgentScheduler] Deleted {deleted_rows} tasks older than {retention_days} days" ) self.__total_pending_tasks = task_manager.count_tasks(status="pending") # get more task if needed if self.__total_pending_tasks > 0: log.info( f"[AgentScheduler] Total pending tasks: {self.__total_pending_tasks}" ) pending_tasks = task_manager.get_tasks(status="pending", limit=1) if len(pending_tasks) > 0: return pending_tasks[0] else: log.info("[AgentScheduler] Task queue is empty") self.__run_callbacks("task_cleared") def __on_image_saved(self, data: script_callbacks.ImageSaveParams): self.__saved_images_path.append( (data.filename, data.pnginfo.get("parameters", "")) ) def on_task_registered(self, callback: Callable): """Callback when a task is registered Callback signature: callback(task_id: str, is_img2img: bool, is_ui: bool, args: dict) """ self.script_callbacks["task_registered"].append(callback) def on_task_started(self, callback: Callable): """Callback when a task is started Callback signature: callback(task_id: str, is_img2img: bool, is_ui: bool) """ self.script_callbacks["task_started"].append(callback) def on_task_finished(self, callback: Callable): """Callback when a task is finished Callback signature: callback(task_id: str, is_img2img: bool, is_ui: bool, status: TaskStatus, result: dict) """ self.script_callbacks["task_finished"].append(callback) def on_task_cleared(self, callback: Callable): self.script_callbacks["task_cleared"].append(callback) def __run_callbacks(self, name: str, *args, **kwargs): for callback in self.script_callbacks[name]: callback(*args, **kwargs) def get_instance(block) -> TaskRunner: if TaskRunner.instance is None: if block is not None: txt2img_submit_button = get_component_by_elem_id(block, "txt2img_generate") UiControlNetUnit = detect_control_net(block, txt2img_submit_button) TaskRunner(UiControlNetUnit) else: TaskRunner() def on_before_reload(): # Tell old instance to stop TaskRunner.instance.dispose = True # force recreate the instance TaskRunner.instance = None script_callbacks.on_before_reload(on_before_reload) return TaskRunner.instance