minor refactoring

.gitignore

@@ -3,7 +3,7 @@ output
 model_data
 tokens
 __pycache__
-arty/.first_init
+.first_init
 
 # not done yet
 Dockerfile

arty/bot.py

@@ -4,8 +4,8 @@ import random
 import discord
 import torch
 
-from messages import SORRY_SLOW, NO_GPU
-from thread_queue import ThreadQueue
+from arty.messages import SORRY_SLOW, NO_GPU
+from arty.thread_queue import ThreadQueue
 
 
 class Bot:

arty/make_image.py

@@ -7,7 +7,7 @@ from PIL import Image
 from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler
 from torch import autocast
 
-from utils import get_token
+from arty.utils import get_token
 
 
 class Unsafe(StableDiffusionPipeline):

arty/utils.py

@@ -1,4 +1,6 @@
+import asyncio
 import os
+from threading import Thread
 
 
 def get_token(name: str):
@@ -12,3 +14,21 @@ def get_token(name: str):
         result = token.read()
 
     return result.strip().strip('\n')
+
+
+def run_in_thread(worker):
+    def start_background_loop(loop: asyncio.AbstractEventLoop) -> None:
+        asyncio.set_event_loop(loop)
+        loop.run_forever()
+
+    async def task():
+        await worker.run()
+
+    target_loop = asyncio.new_event_loop()
+
+    thread = Thread(target=start_background_loop, args=(target_loop,), daemon=True)
+    thread.start()
+
+    task = asyncio.run_coroutine_threadsafe(task(), target_loop)
+    task.result()
+

arty/worker.py

@@ -2,7 +2,8 @@ import asyncio
 from io import BytesIO
 
 import discord
-from make_image import make_image
+
+from arty.make_image import make_image
 
 
 class Worker:

arty/main.py → main.py

@@ -1,13 +1,12 @@
 import asyncio
 import os.path
 from concurrent.futures import ThreadPoolExecutor
-from threading import Thread
 
 from arty.thread_queue import ThreadQueue
-from worker import Worker
-from make_image import make_image
-from utils import get_token
-from bot import Bot
+from arty.worker import Worker
+from arty.make_image import make_image
+from arty.utils import get_token, run_in_thread
+from arty.bot import Bot
 
 
 def _init():
@@ -24,28 +23,10 @@ def _init():
     print('first time setup done.')
 
 
-def run_in_thread(worker, target_loop):
-    def start_background_loop(loop: asyncio.AbstractEventLoop) -> None:
-        asyncio.set_event_loop(loop)
-        loop.run_forever()
-
-    async def task():
-        await worker.run()
-
-    thread = Thread(target=start_background_loop, args=(target_loop,), daemon=True)
-    thread.start()
-
-    task = asyncio.run_coroutine_threadsafe(task(), target_loop)
-    task.result()
-
-
 async def _main():
     worker_queue = ThreadQueue()
     bot_queue = ThreadQueue()
 
-    loop1 = asyncio.new_event_loop()
-    loop2 = asyncio.new_event_loop()
-
     worker = Worker(bot_queue, worker_queue)
     bot = Bot(get_token('discord'), bot_queue, worker_queue)
 
@@ -54,8 +35,8 @@ async def _main():
     loop = asyncio.get_event_loop()
     await asyncio.wait(
         fs={
-            loop.run_in_executor(executor, run_in_thread, worker, loop1),
-            loop.run_in_executor(executor, run_in_thread, bot, loop2),
+            loop.run_in_executor(executor, run_in_thread, worker),
+            loop.run_in_executor(executor, run_in_thread, bot),
         },
         return_when=asyncio.ALL_COMPLETED
     )