Testing text2img command (will be replaced, only for tests); added exception handling, queue and database table with generated images

bot/modules/api/objects/prompt_request.py

@@ -0,0 +1,14 @@
+import dataclasses
+
+
+@dataclasses.dataclass
+class Prompt:
+    prompt: str
+    negative_prompt: str = None
+    steps: int = 20
+    cfg_scale: int = 7
+    width: int = 768
+    height: int = 768
+    restore_faces: bool = True
+    sampler: str = "Euler a"
+    creator: int = None

bot/modules/api/status.py

@@ -0,0 +1,31 @@
+from bot.db import db, DBTables
+import aiohttp
+import asyncio
+import time
+
+
+async def job_exists(endpoint):
+    async with aiohttp.ClientSession() as session:
+        r = await session.get(
+            endpoint + "/sdapi/v1/progress",
+            json={
+                "skip_current_image": True,
+            }
+        )
+        if r.status != 200:
+            return None
+        return (await r.json()).get('state').get('job_count') > 0
+
+
+async def wait_for_status(ignore_exceptions: bool = False):
+    endpoint = db[DBTables.config].get('endpoint')
+    try:
+        while await job_exists(endpoint):
+            while db[DBTables.cooldown].get('_last_time_status_checked', 0) + 5 > time.time():
+                await asyncio.sleep(5)
+            db[DBTables.cooldown]['_last_time_status_checked'] = time.time()
+        return
+    except Exception as e:
+        if not ignore_exceptions:
+            raise e
+        return

bot/modules/api/txt2img.py

@@ -1,26 +1,25 @@
 import aiohttp
 from bot.db import db, DBTables
+from .objects.prompt_request import Prompt
 import json
 import base64
 
 
-async def txt2img(prompt: str, negative_prompt: str = None, steps: int = 20,
-                  cfg_scale: int = 7, width: int = 768, height: int = 768,
-                  restore_faces: bool = True, sampler: str = "Euler a") -> list[bytes, dict] | None:
+async def txt2img(prompt: Prompt, ignore_exceptions: bool = False) -> list[bytes, dict] | None:
     endpoint = db[DBTables.config].get('endpoint')
     try:
         async with aiohttp.ClientSession() as session:
             r = await session.post(
                 endpoint + "/sdapi/v1/txt2img",
                 json={
-                    "prompt": prompt,
-                    "steps": steps,
-                    "cfg_scale": cfg_scale,
-                    "width": width,
-                    "height": height,
-                    "restore_faces": restore_faces,
-                    "negative_prompt": negative_prompt,
-                    "sampler_index": sampler
+                    "prompt": prompt.prompt,
+                    "steps": prompt.steps,
+                    "cfg_scale": prompt.cfg_scale,
+                    "width": prompt.width,
+                    "height": prompt.height,
+                    "restore_faces": prompt.restore_faces,
+                    "negative_prompt": prompt.negative_prompt,
+                    "sampler_index": prompt.sampler
                 }
             )
             if r.status != 200:
@@ -28,5 +27,6 @@ async def txt2img(prompt: str, negative_prompt: str = None, steps: int = 20,
             return [base64.b64decode((await r.json())["images"][0]),
                     json.loads((await r.json())["info"])]
     except Exception as e:
-        assert e
-        return None
+        if not ignore_exceptions:
+            raise e
+        return