diff --git a/example_modules/demo_modules/rembg.module.py b/example_modules/demo_modules/rembg.module.py index dab055c..0d7f0aa 100644 --- a/example_modules/demo_modules/rembg.module.py +++ b/example_modules/demo_modules/rembg.module.py @@ -19,6 +19,43 @@ def reg_or_not(): from jarvis.functional_modules.caller_context import CallerContext from jarvis.functional_modules.functional_module import functional_module from jarvis.logger import logger + import threading + from queue import Queue + import asyncio + + class RembgThread: + _task_queue: Queue + _work_thread: threading.Thread + + def __init__(self) -> None: + self._task_queue = Queue() + self._work_thread = threading.Thread(target=self._work_thread_routine, daemon=True) + self._work_thread.start() + + def _work_thread_routine(self): + while True: + img_base64, fut = self._task_queue.get() + try: + bytes_io = io.BytesIO(base64.b64decode(img_base64)) + img = Image.open(bytes_io) + output = remove(img, session=session) + output_bytesio = io.BytesIO() + output.save(output_bytesio, 'PNG') + output_bytes = output_bytesio.getvalue() + result = base64.b64encode(output_bytes).decode() + fut.set_result(result) + except BaseException as ex: + if isinstance(ex, InterruptedError): + break + fut.set_exception(ex) + logger.debug("Rembg thread exit") + + def add_task(self, img_base64: str): + fut = asyncio.Future(loop=asyncio.get_event_loop()) + self._task_queue.put((img_base64, fut)) + return fut + rembg_thread = RembgThread() + @functional_module( name="remove_bg", description="Remove the background of last image", @@ -28,14 +65,10 @@ def reg_or_not(): if img_base64 is None: await context.reply_text("You need to give me an image first") return "Failed" - bytes_io = io.BytesIO(base64.b64decode(img_base64)) - img = Image.open(bytes_io) - output = remove(img, session=session) - output_bytesio = io.BytesIO() - output.save(output_bytesio, 'PNG') - output_bytes = output_bytesio.getvalue() - result = base64.b64encode(output_bytes).decode() - + + await context.reply_text("Removing the background, please wait") + fut = rembg_thread.add_task(img_base64) + result = await fut context.set_last_image(result) await context.reply_image_base64(result) return "Success"