Update app.py

app.py

@@ -10,6 +10,11 @@ from string import ascii_letters, digits
 import requests
 import sys
 import warnings
+import time
+from concurrent.futures import ProcessPoolExecutor
+import threading
+import asyncio
+from queue import Queue as BlockingQueue
 
 # external
 
@@ -25,8 +30,6 @@ from diffusers.models.modeling_utils import ModelMixin
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file, save_file
 from diffusers import DiffusionPipeline, AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler, DDIMScheduler, StableDiffusionXLPipeline, UNet2DConditionModel, AutoencoderKL, UNet3DConditionModel
-from numba import jit, void as vid, int64 as int, float64 as flt, boolean as bol, uint8 as rgb
-from numba.types import unicode_type as str
 from functools import partial
 
 # logging
@@ -45,11 +48,6 @@ formatter = logging.Formatter('\n >>> [%(levelname)s] %(asctime)s %(name)s: %(me
 handler2.setFormatter(formatter)
 root.addHandler(handler2)
 
-# jit decorator shortcuts
-
-math = partial(jit, cache=True, parallel=True)
-loop = partial(jit, cache=True, forceobj=True)
-
 # storage data
 
 last_motion=""
@@ -135,8 +133,128 @@ pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapt
 pipe.enable_vae_slicing()
 pipe.enable_free_init(method="butterworth", use_fast_sampling=fast)
 
+# Threading
+
+class TwoSidedQueue:
+    def __init__(self, queue_in, queue_out):
+        self._queue_in = queue_in
+        self._queue_out = queue_out
+        self._sides = {
+            'empty': queue_out,
+            'full': queue_out,
+            'get': queue_in,
+            'get_nowait': queue_in,
+            'join': queue_out,
+            'put': queue_out,
+            'put_nowait': queue_out,
+            'qsize': queue_out,
+            'task_done': queue_in,
+        }
+    def __getattr__(self, name):
+        return getattr(self._sides.get(name, self._queue_in), name)
+
+
+class LaunchAsync:
+    def __init__(self, coro, *args, **kwargs):
+        self._coro = coro
+        self._args = args
+        self._kwargs = kwargs
+        self._thread = None
+        self._loop = None
+        self._task = None
+        self._queue_in = None
+        self._queue_out = None
+        self._size = 0
+
+    def size(self, size):
+        self._size = size or 0
+        return self
+
+    def put(self, data, *, timeout=None):
+        """
+        `put` data in for the `coro` to `get` out. Will block if the maximum `size` was reached.
+
+        Does nothing if the `coro` is dead.
+        """
+        try:
+            return asyncio.run_coroutine_threadsafe(self._queue_out.put(data), self._loop).result(timeout)
+        except RuntimeError:
+            if self._loop.is_running():
+                raise
+            else:
+                return None
+
+    def get(self, *, timeout=None):
+        """
+        `get` data out of the `coro` it `put` in. Will block if the queue is empty.
+
+        Returns `None` if the `coro` is dead.
+        """
+        try:
+            return asyncio.run_coroutine_threadsafe(self._queue_in.get(), self._loop).result(timeout)
+        except RuntimeError:
+            if self._loop.is_running():
+                raise
+            else:
+                return None
+
+    def dead(self):
+        """
+        Return `true` if the other side is dead (the `coro` has exited, with or without error).
+        """
+        return not self._loop.is_running()
+
+    def __enter__(self):
+        # asyncio.run is used as it's a battle-tested way to safely set up a new loop and tear
+        # it down. However it does mean it's necessary to wait for the task to run before it's
+        # possible to get said loop and task back. For this, the usual blocking queue is used.
+        oneshot = BlockingQueue(1)
+        self._thread = threading.Thread(target=asyncio.run, args=(
+            self._run(self._coro, self._size, oneshot, self._args, self._kwargs),))
+        self._thread.start()
+        self._loop, self._task, self._queue_in, self._queue_out = oneshot.get()
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        try:
+            self._loop.call_soon_threadsafe(self._task.cancel)
+        except RuntimeError:
+            if self._loop.is_running():
+                raise
+        finally:
+            self._thread.join()
+
+    @staticmethod
+    async def _run(coro, size, oneshot, args, kwargs):
+        # asyncio.Queue's are created here so that they pick up the right loop.
+        queue_in, queue_out = asyncio.Queue(size), asyncio.Queue(size)
+        oneshot.put((asyncio.get_event_loop(), asyncio.current_task(), queue_in, queue_out))
+        try:
+            # `queue_in` and `queue_out` are intentionally swapped here.
+            await coro(TwoSidedQueue(queue_out, queue_in), *args, **kwargs)
+        except asyncio.CancelledError:
+            pass
+
+class Command:
+    def __init__(self, func, data=None):
+        self.func = func
+        self.data = data
+
+def parallel(*pairs):
+    async def async_main(queue):
+        while True:
+            command = await queue.get()
+            await queue.put(command.func(*command.data))
+
+    with LaunchAsync(async_main) as queue:
+        for pair in pairs:
+            f = pair.pop()
+            queue.put(Command(f, pair))
+        response = queue.get()
+        return response
+    
 # functionality
-
+        
 def run(cmd):
     return str(subprocess.run(cmd, shell=True, capture_output=True, env=None).stdout)
 
@@ -222,7 +340,6 @@ def calc(img,p1,p2,motion):
         num_frames=(fps*time)
     )
 
-@loop
 def handle(*inp):
     
     inp[1] = translate(inp[1],"english")