From 146ba3d2681544cb036c4879bf6de0acb877aef1 Mon Sep 17 00:00:00 2001
From: Aedial <aedial.dev@gmail.com>
Date: Sat, 18 Nov 2023 00:30:53 +0100
Subject: [PATCH] [API] Fix Quality Tags for v2 and v3

---
 novelai_api/_high_level.py | 14 +++++++++++++-
 novelai_api/utils.py       | 10 +++++++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/novelai_api/_high_level.py b/novelai_api/_high_level.py
index f041198..a7c91e8 100644
--- a/novelai_api/_high_level.py
+++ b/novelai_api/_high_level.py
@@ -477,7 +477,19 @@ async def generate_image(
 
         quality_toggle = settings["qualityToggle"]
         if quality_toggle:
-            prompt = f"masterpiece, best quality, {prompt}"
+            if model in (
+                ImageModel.Anime_Full,
+                ImageModel.Anime_Curated,
+                ImageModel.Furry,
+                ImageModel.Inpainting_Anime_Full,
+                ImageModel.Inainting_Anime_Curated,
+                ImageModel.Inpainting_Furry,
+            ):
+                prompt = f"masterpiece, best quality, {prompt}"
+            elif model is ImageModel.Anime_v2:
+                prompt = f"very aesthetic, best quality, absurdres, {prompt}"
+            elif model in (ImageModel.Anime_v3, ImageModel.Inpainting_Anime_v3):
+                prompt = f"{prompt}, best quality, amazing quality, very aesthetic, absurdres"
 
         async for e in self._parent.low_level.generate_image(prompt, model, action, settings):
             yield e
diff --git a/novelai_api/utils.py b/novelai_api/utils.py
index 26d7d8d..17fead3 100644
--- a/novelai_api/utils.py
+++ b/novelai_api/utils.py
@@ -1,7 +1,7 @@
 import json
 from base64 import b64decode, b64encode, urlsafe_b64encode
 from hashlib import blake2b
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, AsyncGenerator, AsyncIterable, Dict, Iterable, List, Optional, Tuple, Union
 from zlib import MAX_WBITS, Z_BEST_COMPRESSION
 from zlib import compressobj as deflate_obj
 from zlib import decompress as inflate
@@ -367,4 +367,12 @@ def tokenize_if_not(model: Model, o: Union[str, List[int]]) -> List[int]:
     return Tokenizer.encode(model, o)
 
 
+async def gather_asyncgenerator(agen: Union[AsyncGenerator[Any, None], AsyncIterable[Any]]) -> List[Any]:
+    """
+    Gather all the items of an async generator into a list
+    """
+
+    return [item async for item in agen]
+
+
 # TODO: story tree builder