Browse Source

Fix art, add imagine

Kirk Trombley 4 years ago
parent
commit
a0b90d3de7
1 changed files with 27 additions and 1 deletions
  1. 27 1
      src/commands/querying.py

+ 27 - 1
src/commands/querying.py

@@ -188,7 +188,7 @@ art_number = 0
 
 
 @as_command
-def art(req: Request, logger: Logger):
+async def art(req: Request, logger: Logger):
     """
     The !art command uses the 9gans gallery at https://9gans.com/ which generates 9 images every hour.
     This command will cycle through those 9 images, so if you fire it 10 times in quick succession, the tenth
@@ -206,3 +206,29 @@ def art(req: Request, logger: Logger):
     except:
         logger.exception("Failed to call art generator")
         RollbotFailure.SERVICE_DOWN.raise_exc(detail="Failed to call art generator")
+
+
+@as_command
+async def imagine(args: Args, api_key: Config("text2img.api_key"), req: Request, logger: Logger):
+    """
+    The !imagine command uses the text2img API at https://deepai.org/machine-learning-model/text2img
+    """
+
+    args = args.strip()
+
+    if len(args) == 0:
+        RollbotFailure.INVALID_ARGUMENTS.raise_exc(
+            detail="The !imagine command needs text to imagine!"
+        )
+
+    try:
+        async with req.post(
+            "https://api.deepai.org/api/text2img", data={"text": args}, headers={"api-key": api_key}
+        ) as res:
+            res.raise_for_status()
+            print(await res.text())
+            js = await res.json()
+            return js.get("output_url", None)
+    except:
+        logger.exception("Failed to call text2img service")
+        RollbotFailure.SERVICE_DOWN.raise_exc(detail="Failed to call text2img service")