novelai-python 0.2.8__tar.gz → 0.3.0__tar.gz

{novelai_python-0.2.8 → novelai_python-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.2.8
+Version: 0.3.0
 Summary: Novelai Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache License 2.0
@@ -19,6 +19,7 @@ Requires-Dist: fastapi>=0.109.0
 Requires-Dist: uvicorn[standard]>=0.27.0.post1
 Requires-Dist: numpy>=1.24.4
 Requires-Dist: argon2-cffi>=23.1.0
+Requires-Dist: opencv-python>=4.8.1.78
 Requires-Dist: pytest>=7.4.4; extra == "testing"
 Requires-Dist: pytest-asyncio>=0.23.4; extra == "testing"
 Provides-Extra: testing
@@ -26,7 +27,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 +38,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.8 → novelai_python-0.3.0}/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.8 → novelai_python-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.2.8"
+version = "0.3.0"
 description = "Novelai Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "coldlando@hotmail.com" },
@@ -18,6 +18,7 @@ dependencies = [
     "uvicorn[standard]>=0.27.0.post1",
     "numpy>=1.24.4",
     "argon2-cffi>=23.1.0",
+    "opencv-python>=4.8.1.78",
 ]
 requires-python = ">=3.8"
 readme = "README.md"

{novelai_python-0.2.8 → novelai_python-0.3.0}/src/novelai_python/__init__.py

@@ -16,11 +16,15 @@ 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",
 

{novelai_python-0.2.8 → novelai_python-0.3.0}/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.3.0/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.3.0/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.8 → novelai_python-0.3.0}/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.8 → novelai_python-0.3.0}/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.8 → novelai_python-0.3.0}/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.8 → novelai_python-0.3.0}/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.8 → novelai_python-0.3.0}/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.8 → novelai_python-0.3.0}/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.8 → novelai_python-0.3.0}/src/novelai_python/sdk/ai/generate_image.py

@@ -22,7 +22,7 @@ from typing_extensions import override
 
 from ..schema import ApiBaseModel
 from ..._exceptions import APIError, AuthError, ConcurrentGenerationError, SessionHttpError
-from ..._response import ImageGenerateResp
+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
@@ -208,6 +207,14 @@ class GenerateImageInfer(ApiBaseModel):
     parameters: Params = Params()
     model_config = ConfigDict(extra="ignore")
 
+    @property
+    def endpoint(self):
+        return self._endpoint
+
+    @endpoint.setter
+    def endpoint(self, value):
+        self._endpoint = value
+
     @override
     def model_post_init(self, *args) -> None:
         """
@@ -242,20 +249,14 @@ class GenerateImageInfer(ApiBaseModel):
     def validate_model(self):
         if self.action == Action.INFILL and not self.parameters.mask:
             logger.warning("Mask maybe required for infill mode.")
+        if self.action == Action.INFILL:
+            self.parameters.extra_noise_seed = self.parameters.seed
         return self
 
     @property
     def base_url(self):
         return f"{self.endpoint.strip('/')}/ai/generate-image"
 
-    @property
-    def endpoint(self):
-        return self._endpoint
-
-    @endpoint.setter
-    def endpoint(self, value):
-        self._endpoint = value
-
     def calculate_cost(self, is_opus: bool = False):
         """
         Calculate the Anlas cost of current parameters.
@@ -412,14 +413,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(

novelai_python-0.3.0/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.8 → novelai_python-0.3.0}/src/novelai_python/sdk/user/login.py

@@ -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

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

@@ -0,0 +1,104 @@
+# -*- 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
+
+import cv2 as cv
+import numpy as np
+
+
+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
+
+
+def create_mask_from_sketch(original_img_bytes: bytes,
+                            sketch_img_bytes: bytes,
+                            output_format: str = '.png',
+                            jagged_edges: bool = True
+                            ) -> bytes:
+    """
+    Function to create a mask from original and sketch images input as bytes. Returns bytes.
+
+    :param jagged_edges:  Whether to create jagged edges in the mask. Defaults to False.
+    :param original_img_bytes: Bytes corresponding to the original image.
+    :param sketch_img_bytes: Bytes corresponding to the sketch image.
+    :param output_format: Format of the output image. Defaults to '.png'. It could also be '.jpg'
+    :returns bytes: Bytes corresponding to the resultant mask
+    """
+    # Load images
+    ori_img = cv.imdecode(np.frombuffer(original_img_bytes, np.uint8), cv.IMREAD_COLOR)
+    sketch_img = cv.imdecode(np.frombuffer(sketch_img_bytes, np.uint8), cv.IMREAD_COLOR)
+
+    # Check if images have the same size
+    if ori_img.shape != sketch_img.shape:
+        raise ValueError("Images must have the same size.")
+
+    # Calculate difference between the original and sketch images
+    diff_img = cv.absdiff(ori_img, sketch_img)
+
+    # Convert the difference image to grayscale
+    diff_gray = cv.cvtColor(diff_img, cv.COLOR_BGR2GRAY)
+
+    # Threshold to create the mask
+    _, thresh = cv.threshold(diff_gray, 10, 255, cv.THRESH_BINARY)
+
+    if jagged_edges:
+        # Use a bigger kernel for dilation to create larger 'step' effect at the edges
+        kernel = np.ones((7, 7), np.uint8)
+    else:
+        kernel = np.ones((3, 3), np.uint8)
+
+    # Perform morphological opening to remove small noise
+    opening = cv.morphologyEx(thresh, cv.MORPH_OPEN, kernel, iterations=2)
+
+    # Further remove noise with a Gaussian filter
+    smooth = cv.GaussianBlur(opening, (5, 5), 0)
+
+    if jagged_edges:
+        # Apply additional thresholding to create sharper, jagged edges
+        _, smooth = cv.threshold(smooth, 128, 255, cv.THRESH_BINARY)
+
+    # Convert image to BytesIO object
+    is_success, buffer = cv.imencode(output_format, smooth)
+    if is_success:
+        output_io = buffer.tobytes()
+    else:
+        raise ValueError("Error during conversion of image to BytesIO object")
+
+    return output_io

novelai_python-0.3.0/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.8/src/novelai_python/utils/useful.py

@@ -1,46 +0,0 @@
-# -*- coding: utf-8 -*-
-# @Time    : 2024/2/9 下午10:41
-# @Author  : sudoskys
-# @File    : useful.py
-# @Software: PyCharm
-import collections
-import random
-from typing import List
-
-
-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: int) -> str:
-        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
-
-
-