novelai-python 0.2.7__tar.gz → 0.2.9__tar.gz

{novelai_python-0.2.7 → novelai_python-0.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.2.7
+Version: 0.2.9
 Summary: Novelai Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache License 2.0
@@ -26,7 +26,7 @@ Description-Content-Type: text/markdown
 
 # novelai-python
 
-[![PyPI version](https://badge.fury.io/py/novelai_python.svg)](https://badge.fury.io/py/novelai_python)
+[![PyPI version](https://badge.fury.io/py/novelai-python.svg)](https://badge.fury.io/py/novelai-python)
 [![Downloads](https://pepy.tech/badge/novelai_python)](https://pepy.tech/project/novelai_python)
 
 The goal of this repository is to use Pydantic to build legitimate requests to access the Novelai API service.
@@ -37,16 +37,20 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 - [x] /user/subscription
 - [x] /user/login
 - [x] /user/information
+- [x] /ai/upscale
 - [ ] /ai/generate-image/suggest-tags
 - [ ] /ai/annotate-image
 - [ ] /ai/classify
-- [ ] /ai/upscale
 - [ ] /ai/generate-prompt
 - [ ] /ai/generate
 - [ ] /ai/generate-voice
 
 ### Usage 🖥️
 
+```shell
+pip install novelai-python
+```
+
 More examples can be found in the [playground](https://github.com/LlmKira/novelai-python/tree/main/playground) directory.
 
 ```python

{novelai_python-0.2.7 → novelai_python-0.2.9}/README.md

@@ -1,6 +1,6 @@
 # novelai-python
 
-[![PyPI version](https://badge.fury.io/py/novelai_python.svg)](https://badge.fury.io/py/novelai_python)
+[![PyPI version](https://badge.fury.io/py/novelai-python.svg)](https://badge.fury.io/py/novelai-python)
 [![Downloads](https://pepy.tech/badge/novelai_python)](https://pepy.tech/project/novelai_python)
 
 The goal of this repository is to use Pydantic to build legitimate requests to access the Novelai API service.
@@ -11,16 +11,20 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 - [x] /user/subscription
 - [x] /user/login
 - [x] /user/information
+- [x] /ai/upscale
 - [ ] /ai/generate-image/suggest-tags
 - [ ] /ai/annotate-image
 - [ ] /ai/classify
-- [ ] /ai/upscale
 - [ ] /ai/generate-prompt
 - [ ] /ai/generate
 - [ ] /ai/generate-voice
 
 ### Usage 🖥️
 
+```shell
+pip install novelai-python
+```
+
 More examples can be found in the [playground](https://github.com/LlmKira/novelai-python/tree/main/playground) directory.
 
 ```python

{novelai_python-0.2.7 → novelai_python-0.2.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.2.7"
+version = "0.2.9"
 description = "Novelai Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "coldlando@hotmail.com" },

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/__init__.py

@@ -8,18 +8,23 @@ from ._exceptions import (
     NovelAiError,
     APIError,
     AuthError,
-    ConcurrentGenerationError
+    ConcurrentGenerationError,
+    SessionHttpError
 )
 from .credential import JwtCredential, LoginCredential, ApiCredential
 from .sdk import GenerateImageInfer, ImageGenerateResp
 from .sdk import Information, InformationResp
 from .sdk import Login, LoginResp
 from .sdk import Subscription, SubscriptionResp
+from .sdk import Upscale, UpscaleResp
 
 __all__ = [
     "GenerateImageInfer",
     "ImageGenerateResp",
 
+    "Upscale",
+    "UpscaleResp",
+
     "Subscription",
     "SubscriptionResp",
 
@@ -34,6 +39,7 @@ __all__ = [
     "ApiCredential",
 
     "APIError",
+    "SessionHttpError",
     "AuthError",
     "NovelAiError",
     "ConcurrentGenerationError"

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/_exceptions.py

@@ -10,7 +10,17 @@ class NovelAiError(Exception):
     """
     NovelAiError is the base exception for all novelai_python errors.
     """
-    pass
+    message: str
+    exception: Optional[Exception] = None
+
+    def __init__(self, message: str) -> None:
+        self.message = message
+
+    @property
+    def __dict__(self):
+        return {
+            "message": self.message,
+        }
 
 
 class InvalidRequestHeader(NovelAiError):
@@ -20,33 +30,36 @@ class InvalidRequestHeader(NovelAiError):
     pass
 
 
+class SessionHttpError(NovelAiError):
+    """
+    HTTPError is raised when a request to the API fails.
+    """
+    pass
+
+
 class APIError(NovelAiError):
     """
     APIError is raised when the API returns an error.
     """
-    message: str
     request: Any
     code: Optional[str] = None
     response: Union[Dict[str, Any], str] = None
 
-    def __init__(self, message: str, request: dict, response: Union[dict, str], status_code: str) -> None:
-        if not isinstance(response, dict):
-            response = {"error": f"data type error, should be dict, but got {type(response)}"}
-        if not isinstance(request, dict):
-            request = {"error": f"data type error, should be dict, but got {type(request)}"}
+    def __init__(self, message: str, request: Any, response: Union[Dict[str, Any], str], code: Optional[str]) -> None:
+        super().__init__(message)
         self.request = request
-        self.message = message
-        self.code = status_code
         self.response = response
+        self.code = code
 
     @property
     def __dict__(self):
-        return {
-            "message": self.message,
+        parent_dict = super().__dict__
+        parent_dict.update({
             "request": self.request,
             "response": self.response,
             "code": self.code
-        }
+        })
+        return parent_dict
 
 
 class ConcurrentGenerationError(APIError):
@@ -54,8 +67,7 @@ class ConcurrentGenerationError(APIError):
     ConcurrentGenerationError is raised when the API returns an error.
     """
 
-    def __init__(self, message: str, request: dict, response: Union[dict, str], status_code: str) -> None:
-        super().__init__(message, request, response, status_code)
+    pass
 
 
 class AuthError(APIError):
@@ -63,5 +75,4 @@ class AuthError(APIError):
     AuthError is raised when the API returns an error.
     """
 
-    def __init__(self, message: str, request: dict, response: Union[dict, str], status_code: str) -> None:
-        super().__init__(message, request, response, status_code)
+    pass

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/_response/ai/generate_image.py

@@ -7,8 +7,10 @@ from typing import Tuple, List
 
 from pydantic import BaseModel
 
+from ..schema import RespBase
 
-class ImageGenerateResp(BaseModel):
+
+class ImageGenerateResp(RespBase):
     class RequestParams(BaseModel):
         endpoint: str
         raw_request: dict = None

novelai_python-0.2.9/src/novelai_python/_response/ai/upscale.py

@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/13 上午11:29
+# @Author  : sudoskys
+# @File    : upscale.py
+# @Software: PyCharm
+from typing import Tuple
+
+from pydantic import BaseModel
+
+from ..schema import RespBase
+
+
+class UpscaleResp(RespBase):
+    class RequestParams(BaseModel):
+        endpoint: str
+        raw_request: dict = None
+
+    meta: RequestParams
+    files: Tuple[str, bytes] = None

novelai_python-0.2.9/src/novelai_python/_response/schema.py

@@ -0,0 +1,11 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/13 上午11:30
+# @Author  : sudoskys
+# @File    : schema.py
+# @Software: PyCharm
+
+from pydantic import BaseModel
+
+
+class RespBase(BaseModel):
+    pass

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/_response/user/information.py

@@ -4,10 +4,12 @@
 # @File    : information.py
 # @Software: PyCharm
 
-from pydantic import BaseModel, Field
+from pydantic import Field
 
+from ..schema import RespBase
 
-class InformationResp(BaseModel):
+
+class InformationResp(RespBase):
     emailVerified: bool = Field(..., description="Email verification status")
     emailVerificationLetterSent: bool = Field(..., description="Email verification letter sent status")
     trialActivated: bool = Field(..., description="Trial activation status")

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/_response/user/login.py

@@ -3,8 +3,9 @@
 # @Author  : sudoskys
 # @File    : login.py
 # @Software: PyCharm
-from pydantic import BaseModel
 
+from ..schema import RespBase
 
-class LoginResp(BaseModel):
+
+class LoginResp(RespBase):
     accessToken: str

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/_response/user/subscription.py

@@ -7,6 +7,8 @@ from typing import Optional, Dict, Any, List
 
 from pydantic import BaseModel, Field
 
+from ..schema import RespBase
+
 
 class TrainingSteps(BaseModel):
     fixedTrainingStepsLeft: int
@@ -30,7 +32,7 @@ class Perks(BaseModel):
     contextTokens: int
 
 
-class SubscriptionResp(BaseModel):
+class SubscriptionResp(RespBase):
     tier: int = Field(..., description="Subscription tier")
     active: bool = Field(..., description="Subscription status")
     expiresAt: int = Field(..., description="Subscription expiration time")

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/credential/JwtToken.py

@@ -24,6 +24,7 @@ class JwtCredential(CredentialBase):
             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 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0",
                 "Accept-Encoding": "gzip, deflate, br",
                 "Authorization": f"Bearer {self.jwt_token.get_secret_value()}",
                 "Content-Type": "application/json",

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/credential/UserAuth.py

@@ -30,6 +30,7 @@ class LoginCredential(CredentialBase):
             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 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0",
                 "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}",

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/sdk/__init__.py

@@ -8,3 +8,4 @@ from .ai.generate_image import GenerateImageInfer, ImageGenerateResp  # noqa 401
 from .user.information import Information, InformationResp  # noqa 401
 from .user.login import Login, LoginResp  # noqa 401
 from .user.subscription import Subscription, SubscriptionResp  # noqa 401
+from .ai.upscale import Upscale, UpscaleResp  # noqa 401

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/sdk/ai/generate_image.py

@@ -21,8 +21,8 @@ from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, model_
 from typing_extensions import override
 
 from ..schema import ApiBaseModel
-from ..._exceptions import APIError, AuthError, ConcurrentGenerationError
-from ..._response import ImageGenerateResp
+from ..._exceptions import APIError, AuthError, ConcurrentGenerationError, SessionHttpError
+from ..._response.ai.generate_image import ImageGenerateResp
 from ...credential import CredentialBase
 from ...utils import try_jsonfy, NovelAiMetadata
 
@@ -102,7 +102,6 @@ class Resolution(Enum):
 
 class GenerateImageInfer(ApiBaseModel):
     _endpoint: Optional[str] = PrivateAttr("https://api.novelai.net")
-    _charge: bool = PrivateAttr(False)
 
     class Params(BaseModel):
         # Inpaint
@@ -412,14 +411,30 @@ class GenerateImageInfer(ApiBaseModel):
         request_data = self.model_dump(mode="json", exclude_none=True)
         # Header
         if isinstance(session, AsyncSession):
-            session.headers.update(self.necessary_headers(request_data))
+            session.headers.update(await self.necessary_headers(request_data))
         elif isinstance(session, CredentialBase):
             update_header = await self.necessary_headers(request_data)
             session = await session.get_session(update_headers=update_header)
         if override_headers:
             session.headers.clear()
             session.headers.update(override_headers)
-        logger.debug(f"Request Data: {request_data}")
+        try:
+            _log_data = request_data.copy()
+            if self.action == Action.GENERATE:
+                logger.debug(f"Request Data: {_log_data}")
+            else:
+                _log_data.get("parameters", {}).update({
+                    "image": "base64 data" if self.parameters.image else "None",
+                }
+                )
+                _log_data.get("parameters", {}).update(
+                    {
+                        "mask": "base64 data" if self.parameters.mask else "None",
+                    }
+                )
+                logger.debug(f"Request Data: {request_data}")
+        except Exception as e:
+            logger.warning(f"Error when print log data: {e}")
         try:
             assert hasattr(session, "post"), "session must have post method."
             response = await session.post(
@@ -432,13 +447,17 @@ class GenerateImageInfer(ApiBaseModel):
                 )
                 try:
                     _msg = response.json()
-                except Exception:
-                    raise APIError(
-                        message=f"Unexpected content type: {response.headers.get('Content-Type')}",
-                        request=request_data,
-                        status_code=response.status_code,
-                        response=try_jsonfy(response.content)
-                    )
+                except Exception as e:
+                    logger.warning(e)
+                    if not isinstance(response.content, str) and len(response.content) < 50:
+                        raise APIError(
+                            message=f"Unexpected content type: {response.headers.get('Content-Type')}",
+                            request=request_data,
+                            code=response.status_code,
+                            response=try_jsonfy(response.content)
+                        )
+                    else:
+                        _msg = {"statusCode": response.status_code, "message": response.content}
                 status_code = _msg.get("statusCode", response.status_code)
                 message = _msg.get("message", "Unknown error")
                 if status_code in [400, 401, 402]:
@@ -446,19 +465,19 @@ class GenerateImageInfer(ApiBaseModel):
                     # 401 : unauthorized
                     # 402 : payment required
                     # 409 : conflict
-                    raise AuthError(message, request=request_data, status_code=status_code, response=_msg)
+                    raise AuthError(message, request=request_data, code=status_code, response=_msg)
                 if status_code in [409]:
                     # conflict error
-                    raise APIError(message, request=request_data, status_code=status_code, response=_msg)
+                    raise APIError(message, request=request_data, code=status_code, response=_msg)
                 if status_code in [429]:
                     # concurrent error
                     raise ConcurrentGenerationError(
                         message=message,
                         request=request_data,
-                        status_code=status_code,
+                        code=status_code,
                         response=_msg
                     )
-                raise APIError(message, request=request_data, status_code=status_code, response=_msg)
+                raise APIError(message, request=request_data, code=status_code, response=_msg)
             zip_file = ZipFile(BytesIO(response.content))
             unzip_content = []
             with zip_file as zf:
@@ -467,7 +486,7 @@ class GenerateImageInfer(ApiBaseModel):
                     raise APIError(
                         message="No file in zip",
                         request=request_data,
-                        status_code=response.status_code,
+                        code=response.status_code,
                         response=try_jsonfy(response.content)
                     )
                 for filename in file_list:
@@ -486,9 +505,10 @@ class GenerateImageInfer(ApiBaseModel):
             )
         except curl_cffi.requests.errors.RequestsError as exc:
             logger.exception(exc)
-            raise RuntimeError(f"An AsyncSession error occurred: {exc}")
+            raise SessionHttpError("An AsyncSession RequestsError occurred, maybe SSL error. Try again later!")
         except httpx.HTTPError as exc:
-            raise RuntimeError(f"An HTTP error occurred: {exc}")
+            logger.exception(exc)
+            raise SessionHttpError("An HTTPError occurred, maybe SSL error. Try again later!")
         except APIError as e:
             raise e
         except Exception as e:

novelai_python-0.2.9/src/novelai_python/sdk/ai/upscale.py

@@ -0,0 +1,204 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/13 上午10:48
+# @Author  : sudoskys
+# @File    : upscale.py
+# @Software: PyCharm
+import base64
+import json
+from io import BytesIO
+from typing import Optional, Union
+from urllib.parse import urlparse
+from zipfile import ZipFile
+
+import curl_cffi
+import httpx
+from curl_cffi.requests import AsyncSession
+from loguru import logger
+from pydantic import ConfigDict, PrivateAttr, model_validator
+
+from ..schema import ApiBaseModel
+from ..._exceptions import APIError, AuthError, SessionHttpError
+from ..._response.ai.upscale import UpscaleResp
+from ...credential import CredentialBase
+from ...utils import try_jsonfy, NovelAiMetadata
+
+
+class Upscale(ApiBaseModel):
+    _endpoint: Optional[str] = PrivateAttr("https://api.novelai.net")
+    image: Union[str, bytes]  # base64
+    width: Optional[int] = None
+    height: Optional[int] = None
+    scale: float = 4
+    model_config = ConfigDict(extra="ignore")
+
+    @model_validator(mode="after")
+    def validate_model(self):
+        if isinstance(self.image, str) and self.image.startswith("data:image/"):
+            raise ValueError("Invalid image format, must be base64 encoded.")
+        if isinstance(self.image, bytes):
+            self.image = base64.b64encode(self.image).decode("utf-8")
+        # Auto detect image size
+        try:
+            from PIL import Image
+            with Image.open(BytesIO(base64.b64decode(self.image))) as img:
+                width, height = img.size
+        except Exception as e:
+            logger.warning(f"Error when validate image size: {e}")
+            if self.width is None or self.height is None:
+                raise ValueError("Invalid image size and cant auto detect, must be set width and height.")
+        else:
+            if self.width is None:
+                self.width = width
+            if self.height is None:
+                self.height = height
+        return self
+
+    @property
+    def base_url(self):
+        return f"{self.endpoint.strip('/')}/ai/upscale"
+
+    @property
+    def endpoint(self):
+        return self._endpoint
+
+    @endpoint.setter
+    def endpoint(self, value):
+        self._endpoint = value
+
+    async def necessary_headers(self, request_data) -> dict:
+        """
+        :param request_data:
+        :return:
+        """
+        return {
+            "Host": urlparse(self.endpoint).netloc,
+            "Accept": "*/*",
+            "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0",
+            "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/",
+            "Content-Type": "application/json",
+            "Origin": "https://novelai.net",
+            "Content-Length": str(len(json.dumps(request_data).encode("utf-8"))),
+            "Connection": "keep-alive",
+            "Sec-Fetch-Dest": "empty",
+            "Sec-Fetch-Mode": "cors",
+            "Sec-Fetch-Site": "same-site",
+            "Pragma": "no-cache",
+            "Cache-Control": "no-cache",
+        }
+
+    async def request(self,
+                      session: Union[AsyncSession, "CredentialBase"],
+                      *,
+                      override_headers: Optional[dict] = None,
+                      remove_sign: bool = False
+                      ) -> UpscaleResp:
+        """
+        生成图片
+        :param override_headers:
+        :param session:  session
+        :param remove_sign:  移除追踪信息
+        :return:
+        """
+        # Data Build
+        request_data = self.model_dump(mode="json", exclude_none=True)
+        # Header
+        if isinstance(session, AsyncSession):
+            session.headers.update(await self.necessary_headers(request_data))
+        elif isinstance(session, CredentialBase):
+            update_header = await self.necessary_headers(request_data)
+            session = await session.get_session(update_headers=update_header)
+        if override_headers:
+            session.headers.clear()
+            session.headers.update(override_headers)
+        try:
+            _log_data = request_data.copy()
+            _log_data.update({"image": "base64 data"}) if isinstance(_log_data.get("image"), str) else None
+            logger.info(f"Upscale request data: {_log_data}")
+        except Exception as e:
+            logger.warning(f"Error when print log data: {e}")
+        try:
+            assert hasattr(session, "post"), "session must have post method."
+            response = await session.post(
+                self.base_url,
+                data=json.dumps(request_data).encode("utf-8")
+            )
+            if response.headers.get('Content-Type') not in ['binary/octet-stream', 'application/x-zip-compressed']:
+                logger.error(
+                    f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
+                )
+                try:
+                    _msg = response.json()
+                except Exception as e:
+                    logger.warning(e)
+                    if not isinstance(response.content, str) and len(response.content) < 50:
+                        raise APIError(
+                            message=f"Unexpected content type: {response.headers.get('Content-Type')}",
+                            request=request_data,
+                            code=response.status_code,
+                            response=try_jsonfy(response.content)
+                        )
+                    else:
+                        _msg = {"statusCode": response.status_code, "message": response.content}
+                status_code = _msg.get("statusCode", response.status_code)
+                message = _msg.get("message", "Unknown error")
+                if status_code in [400, 401, 402]:
+                    # 400 : validation error
+                    # 401 : unauthorized
+                    # 402 : payment required
+                    # 409 : conflict
+                    raise AuthError(message, request=request_data, code=status_code, response=_msg)
+                if status_code in [409]:
+                    # conflict error
+                    raise APIError(message, request=request_data, code=status_code, response=_msg)
+                """
+                if status_code in [429]:
+                    # concurrent error
+                    raise ConcurrentGenerationError(
+                        message=message,
+                        request=request_data,
+                        code=status_code,
+                        response=_msg
+                    )
+                """
+                raise APIError(message, request=request_data, code=status_code, response=_msg)
+            zip_file = ZipFile(BytesIO(response.content))
+            unzip_content = []
+            with zip_file as zf:
+                file_list = zf.namelist()
+                if not file_list:
+                    raise APIError(
+                        message="No file in zip",
+                        request=request_data,
+                        code=response.status_code,
+                        response=try_jsonfy(response.content)
+                    )
+                for filename in file_list:
+                    data = zip_file.read(filename)
+                    if remove_sign:
+                        try:
+                            data = NovelAiMetadata.rehash(BytesIO(data), remove_stealth=True)
+                            if not isinstance(data, bytes):
+                                data = data.getvalue()
+                        except Exception as e:
+                            logger.exception(f"SdkWarn:Remove sign error: {e}")
+                    unzip_content.append((filename, data))
+            return UpscaleResp(
+                meta=UpscaleResp.RequestParams(
+                    endpoint=self.base_url,
+                    raw_request=request_data,
+                ),
+                files=unzip_content[0]
+            )
+        except curl_cffi.requests.errors.RequestsError as exc:
+            logger.exception(exc)
+            raise SessionHttpError("An AsyncSession RequestsError occurred, maybe SSL error. Try again later!")
+        except httpx.HTTPError as exc:
+            logger.exception(exc)
+            raise SessionHttpError("An HTTPError occurred, maybe SSL error. Try again later!")
+        except APIError as e:
+            raise e
+        except Exception as e:
+            logger.opt(exception=e).exception("An Unexpected error occurred")
+            raise e

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/sdk/user/information.py

@@ -12,7 +12,7 @@ from loguru import logger
 from pydantic import PrivateAttr
 
 from ..schema import ApiBaseModel
-from ..._exceptions import APIError, AuthError
+from ..._exceptions import APIError, AuthError, SessionHttpError
 from ..._response.user.information import InformationResp
 from ...credential import CredentialBase
 from ...utils import try_jsonfy
@@ -77,13 +77,17 @@ class Information(ApiBaseModel):
                 )
                 try:
                     _msg = response.json()
-                except Exception:
-                    raise APIError(
-                        message=f"Unexpected content type: {response.headers.get('Content-Type')}",
-                        request=request_data,
-                        status_code=response.status_code,
-                        response=try_jsonfy(response.content)
-                    )
+                except Exception as e:
+                    logger.warning(e)
+                    if not isinstance(response.content, str) and len(response.content) < 50:
+                        raise APIError(
+                            message=f"Unexpected content type: {response.headers.get('Content-Type')}",
+                            request=request_data,
+                            code=response.status_code,
+                            response=try_jsonfy(response.content)
+                        )
+                    else:
+                        _msg = {"statusCode": response.status_code, "message": response.content}
                 status_code = _msg.get("statusCode", response.status_code)
                 message = _msg.get("message", "Unknown error")
                 if status_code in [400, 401, 402]:
@@ -91,17 +95,18 @@ class Information(ApiBaseModel):
                     # 401 : unauthorized
                     # 402 : payment required
                     # 409 : conflict
-                    raise AuthError(message, request=request_data, status_code=status_code, response=_msg)
+                    raise AuthError(message, request=request_data, code=status_code, response=_msg)
                 if status_code in [500]:
                     # An unknown error occured.
-                    raise APIError(message, request=request_data, status_code=status_code, response=_msg)
-                raise APIError(message, request=request_data, status_code=status_code, response=_msg)
+                    raise APIError(message, request=request_data, code=status_code, response=_msg)
+                raise APIError(message, request=request_data, code=status_code, response=_msg)
             return InformationResp.model_validate(response.json())
         except curl_cffi.requests.errors.RequestsError as exc:
             logger.exception(exc)
-            raise RuntimeError(f"An AsyncSession error occurred: {exc}")
+            raise SessionHttpError("An AsyncSession RequestsError occurred, maybe SSL error. Try again later!")
         except httpx.HTTPError as exc:
-            raise RuntimeError(f"An HTTP error occurred: {exc}")
+            logger.exception(exc)
+            raise SessionHttpError("An HTTPError occurred, maybe SSL error. Try again later!")
         except APIError as e:
             raise e
         except Exception as e:

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/sdk/user/login.py

@@ -13,7 +13,7 @@ from loguru import logger
 from pydantic import PrivateAttr, Field
 
 from ..schema import ApiBaseModel
-from ..._exceptions import APIError
+from ..._exceptions import APIError, SessionHttpError
 from ..._response.user.login import LoginResp
 from ...credential import CredentialBase
 from ...utils import try_jsonfy, encode_access_key
@@ -77,7 +77,7 @@ class Login(ApiBaseModel):
         # Data Build
         request_data = self.model_dump(mode="json", exclude_none=True)
         if isinstance(session, AsyncSession):
-            session.headers.update(self.necessary_headers(request_data))
+            session.headers.update(await self.necessary_headers(request_data))
         elif isinstance(session, CredentialBase):
             session = await session.get_session(update_headers=await self.necessary_headers(request_data))
         # Header
@@ -97,29 +97,34 @@ class Login(ApiBaseModel):
                 )
                 try:
                     _msg = response.json()
-                except Exception:
-                    raise APIError(
-                        message=f"Unexpected content type: {response.headers.get('Content-Type')}",
-                        request=request_data,
-                        status_code=response.status_code,
-                        response=try_jsonfy(response.content)
-                    )
+                except Exception as e:
+                    logger.warning(e)
+                    if not isinstance(response.content, str) and len(response.content) < 50:
+                        raise APIError(
+                            message=f"Unexpected content type: {response.headers.get('Content-Type')}",
+                            request=request_data,
+                            code=response.status_code,
+                            response=try_jsonfy(response.content)
+                        )
+                    else:
+                        _msg = {"statusCode": response.status_code, "message": response.content}
                 status_code = _msg.get("statusCode", response.status_code)
                 message = _msg.get("message", "Unknown error")
                 if status_code in [400, 401]:
                     # 400 : A validation error occured.
                     # 401 : Access Key is incorrect.
-                    raise APIError(message, request=request_data, status_code=status_code, response=_msg)
+                    raise APIError(message, request=request_data, code=status_code, response=_msg)
                 if status_code in [500]:
                     # An unknown error occured.
-                    raise APIError(message, request=request_data, status_code=status_code, response=_msg)
-                raise APIError(message, request=request_data, status_code=status_code, response=_msg)
+                    raise APIError(message, request=request_data, code=status_code, response=_msg)
+                raise APIError(message, request=request_data, code=status_code, response=_msg)
             return LoginResp.model_validate(response.json())
         except curl_cffi.requests.errors.RequestsError as exc:
             logger.exception(exc)
-            raise RuntimeError(f"An AsyncSession error occurred: {exc}")
+            raise SessionHttpError("An AsyncSession RequestsError occurred, maybe SSL error. Try again later!")
         except httpx.HTTPError as exc:
-            raise RuntimeError(f"An HTTP error occurred: {exc}")
+            logger.exception(exc)
+            raise SessionHttpError("An HTTPError occurred, maybe SSL error. Try again later!")
         except APIError as e:
             raise e
         except Exception as e:

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/sdk/user/subscription.py

@@ -12,7 +12,7 @@ from loguru import logger
 from pydantic import PrivateAttr
 
 from ..schema import ApiBaseModel
-from ..._exceptions import APIError, AuthError
+from ..._exceptions import APIError, AuthError, SessionHttpError
 from ..._response.user.subscription import SubscriptionResp
 from ...credential import CredentialBase
 from ...utils import try_jsonfy
@@ -84,13 +84,17 @@ class Subscription(ApiBaseModel):
                 )
                 try:
                     _msg = response.json()
-                except Exception:
-                    raise APIError(
-                        message=f"Unexpected content type: {response.headers.get('Content-Type')}",
-                        request=request_data,
-                        status_code=response.status_code,
-                        response=try_jsonfy(response.content)
-                    )
+                except Exception as e:
+                    logger.warning(e)
+                    if not isinstance(response.content, str) and len(response.content) < 50:
+                        raise APIError(
+                            message=f"Unexpected content type: {response.headers.get('Content-Type')}",
+                            request=request_data,
+                            code=response.status_code,
+                            response=try_jsonfy(response.content)
+                        )
+                    else:
+                        _msg = {"statusCode": response.status_code, "message": response.content}
                 status_code = _msg.get("statusCode", response.status_code)
                 message = _msg.get("message", "Unknown error")
                 if status_code in [400, 401, 402]:
@@ -98,17 +102,18 @@ class Subscription(ApiBaseModel):
                     # 401 : unauthorized
                     # 402 : payment required
                     # 409 : conflict
-                    raise AuthError(message, request=request_data, status_code=status_code, response=_msg)
+                    raise AuthError(message, request=request_data, code=status_code, response=_msg)
                 if status_code in [500]:
                     # An unknown error occured.
-                    raise APIError(message, request=request_data, status_code=status_code, response=_msg)
-                raise APIError(message, request=request_data, status_code=status_code, response=_msg)
+                    raise APIError(message, request=request_data, code=status_code, response=_msg)
+                raise APIError(message, request=request_data, code=status_code, response=_msg)
             return SubscriptionResp.model_validate(response.json())
         except curl_cffi.requests.errors.RequestsError as exc:
             logger.exception(exc)
-            raise RuntimeError(f"An AsyncSession error occurred: {exc}")
+            raise SessionHttpError("An AsyncSession RequestsError occurred, maybe SSL error. Try again later!")
         except httpx.HTTPError as exc:
-            raise RuntimeError(f"An HTTP error occurred: {exc}")
+            logger.exception(exc)
+            raise SessionHttpError("An HTTPError occurred, maybe SSL error. Try again later!")
         except APIError as e:
             raise e
         except Exception as e:

{novelai_python-0.2.7 → novelai_python-0.2.9}/src/novelai_python/utils/hash.py

@@ -28,6 +28,9 @@ class NovelAiMetadata(BaseModel):
     title: str = "AI generated image"
     description: str
     comment: Union[dict, str] = {}
+    """
+    {'prompt': 'very very aesthetic, absurdres', 'steps': 28, 'height': 1216, 'width': 832, 'scale': 6.0, 'uncond_scale': 1.0, 'cfg_rescale': 0.0, 'seed': 121910159, 'n_samples': 1, 'hide_debug_overlay': False, 'noise_schedule': 'native', 'legacy_v3_extend': False, 'sampler': 'k_euler', 'controlnet_strength': 1.0, 'controlnet_model': None, 'dynamic_thresholding': False, 'dynamic_thresholding_percentile': 0.999, 'dynamic_thresholding_mimic_scale': 10.0, 'sm': False, 'sm_dyn': False, 'skip_cfg_below_sigma': 0.0, 'lora_unet_weights': None, 'lora_clip_weights': None, 'uc': ', lowres, {bad}, error, fewer, extra, missing, worst quality, jpeg artifacts, bad quality, watermark, unfinished, displeasing, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract], bad anatomy, bad hands, @_@, mismatched pupils, heart-shaped pupils, glowing eyes, nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry', 'request_type': 'PromptGenerateRequest', 'signed_hash': '11111='}
+    """
 
     @property
     def metadata(self):

novelai_python-0.2.9/src/novelai_python/utils/useful.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/9 下午10:41
+# @Author  : sudoskys
+# @File    : useful.py
+# @Software: PyCharm
+import collections
+import random
+from typing import List, Union
+
+
+def enum_to_list(enum_):
+    return list(map(lambda x: x.value, enum_._member_map_.values()))
+
+
+class QueSelect(object):
+    def __init__(self, data: List[str]):
+        """
+        A queue selector
+        :param data:
+        """
+        self.data = collections.deque(data)
+        self.used = collections.deque()
+        self.users = {}
+
+    def get(self, user_id: Union[int, str]) -> str:
+        user_id = str(user_id)
+        if user_id not in self.users:
+            self.users[user_id] = {'data': self.data.copy(), 'used': collections.deque()}
+
+        user_data = self.users[user_id]['data']
+        user_used = self.users[user_id]['used']
+
+        if len(user_data) == 0:
+            user_data, user_used = user_used, user_data
+            # 随机掉 User data
+            if len(user_data) > 2:
+                random.shuffle(user_data)
+            self.users[user_id]['data'] = user_data
+            self.users[user_id]['used'] = user_used
+
+        selected = user_data.popleft()
+        user_used.append(selected)
+
+        return selected

novelai_python-0.2.9/tests/test_upscale.py

@@ -0,0 +1,78 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/13 下午12:03
+# @Author  : sudoskys
+# @File    : test_upscale.py
+# @Software: PyCharm
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/14 下午4:20
+# @Author  : sudoskys
+# @File    : test_upscale.py
+# @Software: PyCharm
+from unittest import mock
+
+import pytest
+from curl_cffi.requests import AsyncSession
+
+from novelai_python import APIError, Upscale, AuthError
+
+
+@pytest.mark.asyncio
+async def test_validation_error_during_upscale():
+    validation_error_response = mock.Mock()
+    validation_error_response.headers = {"Content-Type": "application/json"}
+    validation_error_response.status_code = 400
+    validation_error_response.json.return_value = {
+        "statusCode": 400,
+        "message": "A validation error occurred."
+    }
+    session = mock.MagicMock(spec=AsyncSession)
+    session.post = mock.AsyncMock(return_value=validation_error_response)
+    session.headers = {}
+
+    upscale = Upscale(image="base64_encoded_image", height=100, width=100)
+    with pytest.raises(AuthError) as e:
+        await upscale.request(session=session)
+    assert e.type is AuthError
+    expected_message = 'A validation error occurred.'
+    assert expected_message == str(e.value)
+
+
+@pytest.mark.asyncio
+async def test_unauthorized_error_during_upscale():
+    unauthorized_error_response = mock.Mock()
+    unauthorized_error_response.headers = {"Content-Type": "application/json"}
+    unauthorized_error_response.status_code = 401
+    unauthorized_error_response.json.return_value = {
+        "statusCode": 401,
+        "message": "Unauthorized."
+    }
+    session = mock.MagicMock(spec=AsyncSession)
+    session.post = mock.AsyncMock(return_value=unauthorized_error_response)
+    session.headers = {}
+
+    upscale = Upscale(image="base64_encoded_image", height=100, width=100)
+    with pytest.raises(APIError) as e:
+        await upscale.request(session=session)
+    assert e.type is AuthError
+    expected_message = 'Unauthorized.'
+    assert expected_message == str(e.value)
+
+
+@pytest.mark.asyncio
+async def test_unknown_error_during_upscale():
+    unknown_error_response = mock.Mock()
+    unknown_error_response.headers = {"Content-Type": "application/json"}
+    unknown_error_response.status_code = 500
+    unknown_error_response.json.return_value = {
+        "statusCode": 500,
+        "message": "Unknown error occurred."
+    }
+    session = mock.MagicMock(spec=AsyncSession)
+    session.post = mock.AsyncMock(return_value=unknown_error_response)
+    session.headers = {}
+
+    upscale = Upscale(image="base64_encoded_image", height=100, width=100)
+    with pytest.raises(APIError) as e:
+        await upscale.request(session=session)
+    expected_message = 'Unknown error occurred.'
+    assert expected_message == str(e.value)

novelai_python-0.2.7/src/novelai_python/utils/useful.py

@@ -1,7 +0,0 @@
-# -*- coding: utf-8 -*-
-# @Time    : 2024/2/9 下午10:41
-# @Author  : sudoskys
-# @File    : useful.py
-# @Software: PyCharm
-def enum_to_list(enum_):
-    return list(map(lambda x: x.value, enum_._member_map_.values()))