Merge pull request #26 from LlmKira/dev

feat: auto retry when return StatusCode ==500

README.md

@@ -7,6 +7,8 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 
 ### Roadmap 🚧
 
+- [x] utils.NovelAiMetadata
+- [x] utils.random_prompt
 - [x] /ai/generate-image
 - [x] /user/subscription
 - [x] /user/login

playground/generate_image.py

@@ -38,6 +38,7 @@ async def main():
             session=globe_s, remove_sign=True
         )
     except APIError as e:
+        print(str(e))
         print(e.response)
         return
 

playground/read_nai_tag.py

@@ -9,8 +9,11 @@
 
 if not os.path.exists("generate_image.png"):
     raise FileNotFoundError("generate_image.png not found,pls run generate_image.py first")
+try:
+    meta = NovelAiMetadata.build_from_img(image_io="generate_image.png")  # OR BytesIO(data)
+except ValueError:
+    raise LookupError("Cant find a MetaData")
 
-ba = NovelAiMetadata.build_from_img(image_io="generate_image.png")
-print(ba.title)
-print(ba.description)
-print(ba.comment)
+print(meta.title)
+print(meta.description)
+print(meta.comment)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.3.3"
+version = "0.3.4"
 description = "Novelai Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "[email protected]" },
@@ -19,6 +19,8 @@ dependencies = [
     "numpy>=1.24.4",
     "argon2-cffi>=23.1.0",
     "opencv-python>=4.8.1.78",
+    "fake-useragent>=1.4.0",
+    "tenacity>=8.2.3",
 ]
 requires-python = ">=3.8"
 readme = "README.md"

src/novelai_python/credential/ApiToken.py

@@ -7,7 +7,7 @@
 from loguru import logger
 from pydantic import SecretStr, Field, field_validator
 
-from ._base import CredentialBase
+from ._base import CredentialBase, FAKE_UA
 
 
 class ApiCredential(CredentialBase):
@@ -25,12 +25,12 @@ async def get_session(self, timeout: int = 180, update_headers: dict = None):
                 "Accept": "*/*",
                 "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
                 "Accept-Encoding": "gzip, deflate, br",
-                "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
+                "User-Agent": FAKE_UA.edge,
                 "Authorization": f"Bearer {self.api_token.get_secret_value()}",
                 "Content-Type": "application/json",
                 "Origin": "https://novelai.net",
                 "Referer": "https://novelai.net/",
-            }, impersonate="chrome110")
+            }, impersonate="edge101")
         self._session.headers.update(update_headers)
         return self._session
 

src/novelai_python/credential/JwtToken.py

@@ -7,7 +7,7 @@
 from loguru import logger
 from pydantic import SecretStr, Field, field_validator
 
-from ._base import CredentialBase
+from ._base import CredentialBase, FAKE_UA
 
 
 class JwtCredential(CredentialBase):
@@ -24,13 +24,13 @@ async def get_session(self, timeout: int = 180, update_headers: dict = None):
             self._session = AsyncSession(timeout=timeout, headers={
                 "Accept": "*/*",
                 "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
-                "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
+                "User-Agent": FAKE_UA.edge,
                 "Accept-Encoding": "gzip, deflate, br",
                 "Authorization": f"Bearer {self.jwt_token.get_secret_value()}",
                 "Content-Type": "application/json",
                 "Origin": "https://novelai.net",
                 "Referer": "https://novelai.net/",
-            }, impersonate="chrome110")
+            }, impersonate="edge101")
         assert isinstance(update_headers, dict), "update_headers must be a dict"
         self._session.headers.update(update_headers)
         return self._session

src/novelai_python/credential/UserAuth.py

@@ -9,7 +9,7 @@
 from curl_cffi.requests import AsyncSession
 from pydantic import SecretStr, Field
 
-from ._base import CredentialBase
+from ._base import CredentialBase, FAKE_UA
 
 
 class LoginCredential(CredentialBase):
@@ -30,14 +30,14 @@ async def get_session(self, timeout: int = 180, update_headers: dict = None):
             resp = await Login.build(user_name=self.username, password=self.password.get_secret_value()).request()
             self._session = AsyncSession(timeout=timeout, headers={
                 "Accept": "*/*",
-                "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
+                "User-Agent": FAKE_UA.edge,
                 "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
                 "Accept-Encoding": "gzip, deflate, br",
                 "Authorization": f"Bearer {resp.accessToken}",
                 "Content-Type": "application/json",
                 "Origin": "https://novelai.net",
                 "Referer": "https://novelai.net/",
-            }, impersonate="chrome110")
+            }, impersonate="edge101")
             self._update_at = int(time.time())
         self._session.headers.update(update_headers)
         return self._session

src/novelai_python/credential/_base.py

@@ -4,8 +4,11 @@
 # @File    : _shema.py
 # @Software: PyCharm
 from curl_cffi.requests import AsyncSession
+from fake_useragent import UserAgent
 from pydantic import BaseModel
 
+FAKE_UA = UserAgent()
+
 
 class CredentialBase(BaseModel):
     _session: AsyncSession = None

src/novelai_python/sdk/ai/generate_image/__init__.py

@@ -18,6 +18,7 @@
 from curl_cffi.requests import AsyncSession
 from loguru import logger
 from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, model_validator, Field
+from tenacity import retry, stop_after_attempt, wait_random, retry_if_exception
 from typing_extensions import override
 
 from ._enum import Model, Sampler, NoiseSchedule, ControlNetModel, Action, UCPreset
@@ -179,6 +180,9 @@ def validate_model(self):
             logger.warning("Mask maybe required for infill mode.")
         if self.action != Action.GENERATE:
             self.parameters.extra_noise_seed = self.parameters.seed
+        if self.action == Action.IMG2IMG:
+            self.parameters.sm = False
+            self.parameters.sm_dyn = False
         return self
 
     @property
@@ -322,9 +326,14 @@ async def necessary_headers(self, request_data) -> dict:
             "Sec-Fetch-Site": "same-site",
             "Pragma": "no-cache",
             "Cache-Control": "no-cache",
-            "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
         }
 
+    @retry(
+        wait=wait_random(min=1, max=3),
+        stop=stop_after_attempt(3),
+        retry=retry_if_exception(lambda e: hasattr(e, "code") and str(e.code) == "500"),
+        reraise=True
+    )
     async def request(self,
                       session: Union[AsyncSession, "CredentialBase"],
                       *,

src/novelai_python/sdk/ai/generate_image/suggest_tags.py

@@ -41,7 +41,6 @@ async def necessary_headers(self, request_data) -> dict:
         return {
             "Host": urlparse(self.endpoint).netloc,
             "Accept": "*/*",
-            "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
             "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
             "Accept-Encoding": "gzip, deflate, br",
             "Access-Control-Allow-Origin": "*",

src/novelai_python/sdk/ai/upscale.py

@@ -15,6 +15,7 @@
 from curl_cffi.requests import AsyncSession
 from loguru import logger
 from pydantic import ConfigDict, PrivateAttr, model_validator
+from tenacity import wait_random, retry, stop_after_attempt, retry_if_exception
 
 from ..schema import ApiBaseModel
 from ..._exceptions import APIError, AuthError, SessionHttpError
@@ -73,7 +74,6 @@ async def necessary_headers(self, request_data) -> dict:
         return {
             "Host": urlparse(self.endpoint).netloc,
             "Accept": "*/*",
-            "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
             "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
             "Accept-Encoding": "gzip, deflate, br",
             "Referer": "https://novelai.net/",
@@ -88,6 +88,12 @@ async def necessary_headers(self, request_data) -> dict:
             "Cache-Control": "no-cache",
         }
 
+    @retry(
+        wait=wait_random(min=1, max=3),
+        stop=stop_after_attempt(3),
+        retry=retry_if_exception(lambda e: hasattr(e, "code") and str(e.code) == "500"),
+        reraise=True
+    )
     async def request(self,
                       session: Union[AsyncSession, "CredentialBase"],
                       *,

src/novelai_python/utils/hash.py

@@ -74,17 +74,18 @@ def write_out(self, img_io: BytesIO, *, remove_stealth: bool = False):
     @classmethod
     def build_from_img(cls, image_io):
         with Image.open(image_io) as img:
-            title = img.info.get("Title")
-            if not title == 'AI generated image':
+            title = img.info.get("Title", "Empty")
+            if title != 'AI generated image':
                 raise ValueError("Not a NaiPic")
             prompt = img.info.get("Description")
-            comment = img.info.get("Comment")
+            comment = img.info.get("Comment", None)
+            assert isinstance(comment, str), ValueError("Comment Empty")
             try:
                 comment = json.loads(comment)
             except Exception as e:
                 logger.debug(e)
                 comment = {}
-            return cls(title=title, description=prompt, comment=comment)
+        return cls(title=title, description=prompt, comment=comment)
 
     @classmethod
     def build_from_param(cls, prompt, neg_prompt, **kwargs):