novelai-python 0.2.3__tar.gz → 0.2.5__tar.gz

{novelai_python-0.2.3 → novelai_python-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.2.3
+Version: 0.2.5
 Summary: Novelai Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache License 2.0
@@ -80,7 +80,7 @@ async def main():
     cost = gen.calculate_cost(is_opus=True)
     print(f"charge: {cost} if you are vip3")
 
-    resp = gen.generate(session=globe_s)
+    resp = gen.request(session=globe_s)
     resp: ImageGenerateResp
     print(resp.meta)
 

{novelai_python-0.2.3 → novelai_python-0.2.5}/README.md

@@ -54,7 +54,7 @@ async def main():
     cost = gen.calculate_cost(is_opus=True)
     print(f"charge: {cost} if you are vip3")
 
-    resp = gen.generate(session=globe_s)
+    resp = gen.request(session=globe_s)
     resp: ImageGenerateResp
     print(resp.meta)
 

{novelai_python-0.2.3 → novelai_python-0.2.5}/pyproject.toml

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

{novelai_python-0.2.3 → novelai_python-0.2.5}/src/novelai_python/_exceptions.py

@@ -13,6 +13,13 @@ class NovelAiError(Exception):
     pass
 
 
+class InvalidRequestHeader(NovelAiError):
+    """
+    InvalidRequestHeader is raised when the request header is invalid.
+    """
+    pass
+
+
 class APIError(NovelAiError):
     """
     APIError is raised when the API returns an error.

{novelai_python-0.2.3 → novelai_python-0.2.5}/src/novelai_python/_response/user/subscription.py

@@ -45,8 +45,34 @@ class SubscriptionResp(BaseModel):
 
     @property
     def anlas_left(self):
-        return self.trainingStepsLeft.fixedTrainingStepsLeft
+        return self.trainingStepsLeft.fixedTrainingStepsLeft + self.trainingStepsLeft.purchasedTrainingSteps
 
     @property
     def is_unlimited_image_generation(self):
-        return self.perks.unlimitedImageGeneration
+        return self.perks.unlimitedImageGeneration and self.perks.imageGeneration
+
+    @property
+    def get_tier_name(self):
+        if self.tier == 0:
+            return "Paper"
+        elif self.tier == 1:
+            return "Tablet"
+        elif self.tier == 2:
+            return "Scroll"
+        elif self.tier == 3:
+            return "Opus"
+        else:
+            return "Unknown"
+
+    @property
+    def limit_perks(self):
+        perks = []
+        if not self.perks.imageGeneration:
+            perks.append("imageGeneration")
+        if not self.perks.voiceGeneration:
+            perks.append("voiceGeneration")
+        if not self.perks.unlimitedImageGeneration:
+            perks.append("unlimitedImageGeneration")
+        if not self.perks.unlimitedMaxPriority:
+            perks.append("unlimitedMaxPriority")
+        return perks

{novelai_python-0.2.3 → novelai_python-0.2.5}/src/novelai_python/credential/ApiToken.py

@@ -17,14 +17,20 @@ class ApiCredential(CredentialBase):
     api_token: SecretStr = Field(None, description="api token")
     _session: AsyncSession = None
 
-    async def get_session(self, timeout: int = 180):
+    async def get_session(self, timeout: int = 180, update_headers: dict = None):
+        if update_headers is None:
+            update_headers = {}
         if not self._session:
             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",
+                "Accept-Encoding": "gzip, deflate, br",
                 "Authorization": f"Bearer {self.api_token.get_secret_value()}",
                 "Content-Type": "application/json",
                 "Origin": "https://novelai.net",
                 "Referer": "https://novelai.net/",
             }, impersonate="chrome110")
+        self._session.headers.update(update_headers)
         return self._session
 
     @field_validator('api_token')

{novelai_python-0.2.3 → novelai_python-0.2.5}/src/novelai_python/credential/JwtToken.py

@@ -17,14 +17,21 @@ class JwtCredential(CredentialBase):
     jwt_token: SecretStr = Field(None, description="jwt token")
     _session: AsyncSession = None
 
-    async def get_session(self, timeout: int = 180):
+    async def get_session(self, timeout: int = 180, update_headers: dict = None):
+        if update_headers is None:
+            update_headers = {}
         if not self._session:
             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",
+                "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")
+        assert isinstance(update_headers, dict), "update_headers must be a dict"
+        self._session.headers.update(update_headers)
         return self._session
 
     @field_validator('jwt_token')

{novelai_python-0.2.3 → novelai_python-0.2.5}/src/novelai_python/credential/UserAuth.py

@@ -21,16 +21,22 @@ class LoginCredential(CredentialBase):
     _session: Optional[AsyncSession] = None
     _update_at: Optional[int] = None
 
-    async def get_session(self, timeout: int = 180):
+    async def get_session(self, timeout: int = 180, update_headers: dict = None):
+        if update_headers is None:
+            update_headers = {}
         # 30 天有效期
         if not self._session or int(time.time()) - self._update_at > 29 * 24 * 60 * 60:
             from ..sdk import Login
             resp = await Login.build(user_name=self.username, password=self.password.get_secret_value()).request()
             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",
+                "Accept-Encoding": "gzip, deflate, br",
                 "Authorization": f"Bearer {resp.accessToken}",
                 "Content-Type": "application/json",
                 "Origin": "https://novelai.net",
                 "Referer": "https://novelai.net/",
             }, impersonate="chrome110")
             self._update_at = int(time.time())
+        self._session.headers.update(update_headers)
         return self._session

{novelai_python-0.2.3 → novelai_python-0.2.5}/src/novelai_python/credential/_base.py

@@ -11,5 +11,5 @@ class CredentialBase(BaseModel):
     _session: AsyncSession = None
     """会话"""
 
-    async def get_session(self, timeout: int = 180):
+    async def get_session(self, timeout: int = 180, update_headers: dict = None):
         raise NotImplementedError

{novelai_python-0.2.3 → novelai_python-0.2.5}/src/novelai_python/sdk/ai/generate_image.py

@@ -3,20 +3,24 @@
 # @Author  : sudoskys
 # @File    : generate_image.py
 # @Software: PyCharm
+import base64
 import json
 import math
 import random
 from enum import Enum, IntEnum
 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, RequestsError
+from curl_cffi.requests import AsyncSession
 from loguru import logger
 from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, model_validator, Field
 from typing_extensions import override
 
+from ..schema import ApiBaseModel
 from ..._exceptions import APIError, AuthError
 from ..._response import ImageGenerateResp
 from ...credential import CredentialBase
@@ -48,14 +52,39 @@ class UCPreset(IntEnum):
 
 class Action(Enum):
     GENERATE = "generate"
+    """Generate Image"""
     IMG2IMG = "img2img"
+    """Image to Image"""
     INFILL = "infill"
+    """Inpainting"""
 
 
 class Model(Enum):
     NAI_DIFFUSION_3 = "nai-diffusion-3"
     NAI_DIFFUSION_3_INPAINTING = "nai-diffusion-3-inpainting"
 
+    NAI_DIFFUSION = "nai-diffusion"
+    NAI_DIFFUSION_INPAINTING = "nai-diffusion-inpainting"
+
+    SAFE_DIFFUSION = "safe-diffusion"
+    SAFE_DIFFUSION_INPAINTING = "safe-diffusion-inpainting"
+
+    NAI_DIFFUSION_FURRY = "nai-diffusion-furry"
+    FURRY_DIFFUSION_INPAINTING = "furry-diffusion-inpainting"
+
+
+class ControlNetModel(Enum):
+    HED = "hed"
+    """边缘检测"""
+    MIDAS = "midas"
+    """景深"""
+    FAKE_SCRIBBLE = "fake_scribble"
+    """伪涂鸦"""
+    M_LSD = "mlsd"
+    """(建筑)线条检测"""
+    LANDSCAPER = "uniformer"
+    """风景生成"""
+
 
 class Resolution(Enum):
     RES_512_768 = (512, 768)
@@ -71,25 +100,28 @@ class Resolution(Enum):
     RES_1920_1088 = (1920, 1088)
 
 
-class GenerateImageInfer(BaseModel):
+class GenerateImageInfer(ApiBaseModel):
     _endpoint: Optional[str] = PrivateAttr("https://api.novelai.net")
     _charge: bool = PrivateAttr(False)
 
     class Params(BaseModel):
         # Inpaint
         add_original_image: Optional[bool] = False
-        mask: Optional[str] = None
+        mask: Optional[Union[str, bytes]] = None  # img2img,base64
 
         cfg_rescale: Optional[float] = Field(0, ge=0, le=1, multiple_of=0.02)
+        """Prompt Guidance Rescale"""
         controlnet_strength: Optional[float] = Field(1.0, ge=0.1, le=2, multiple_of=0.1)
+        """ControlNet Strength"""
         dynamic_thresholding: Optional[bool] = False
+        """Reduce artifacts caused by high prompt guidance values."""
         height: Optional[int] = Field(1216, ge=64, le=49152)
         # Img2Img
-        image: Optional[str] = None  # img2img,base64
-        strength: Optional[float] = Field(default=0.3, ge=0.01, le=0.99, multiple_of=0.01)
+        image: Optional[Union[str, bytes]] = None  # img2img,base64
+        strength: Optional[float] = Field(default=0.5, ge=0.01, le=0.99, multiple_of=0.01)
         noise: Optional[float] = Field(default=0, ge=0, le=0.99, multiple_of=0.01)
         controlnet_condition: Optional[str] = None
-        controlnet_model: Optional[str] = None
+        controlnet_model: Optional[ControlNetModel] = None
 
         n_samples: Optional[int] = Field(1, ge=1, le=8)
         negative_prompt: Optional[str] = ''
@@ -103,6 +135,7 @@ class GenerateImageInfer(BaseModel):
         qualityToggle: Optional[bool] = True
         sampler: Optional[Sampler] = Sampler.K_EULER
         scale: Optional[float] = Field(6.0, ge=0, le=10, multiple_of=0.1)
+        """Prompt Guidance"""
         # Seed
         seed: Optional[int] = Field(
             default_factory=lambda: random.randint(0, 4294967295 - 7),
@@ -120,6 +153,7 @@ class GenerateImageInfer(BaseModel):
         steps: Optional[int] = Field(28, ge=1, le=50)
         ucPreset: Optional[UCPreset] = 0
         uncond_scale: Optional[float] = Field(1.0, ge=0, le=1.5, multiple_of=0.05)
+        """Undesired Content Strength"""
         width: Optional[int] = Field(832, ge=64, le=49152)
 
         @model_validator(mode="after")
@@ -130,6 +164,22 @@ class GenerateImageInfer(BaseModel):
                 raise ValueError('Invalid Model Params For img2img2 mode... image should match add_original_image!')
             return self
 
+        @field_validator('mask')
+        def mask_validator(cls, v: Union[str, bytes]):
+            if isinstance(v, str) and v.startswith("data:image/"):
+                raise ValueError("Invalid image format, must be base64 encoded.")
+            if isinstance(v, bytes):
+                return base64.b64encode(v).decode("utf-8")
+            return v
+
+        @field_validator('image')
+        def image_validator(cls, v: Union[str, bytes]):
+            if isinstance(v, str) and v.startswith("data:image/"):
+                raise ValueError("Invalid image format, must be base64 encoded.")
+            if isinstance(v, bytes):
+                return base64.b64encode(v).decode("utf-8")
+            return v
+
         @field_validator('width')
         def width_validator(cls, v: int):
             """
@@ -188,6 +238,12 @@ class GenerateImageInfer(BaseModel):
         if self.parameters.qualityToggle:
             self.input += ", best quality, amazing quality, very aesthetic, absurdres"
 
+    @model_validator(mode="after")
+    def validate_model(self):
+        if self.action == Action.INFILL and not self.parameters.mask:
+            logger.warning("Mask maybe required for infill mode.")
+        return self
+
     @property
     def base_url(self):
         return f"{self.endpoint.strip('/')}/ai/generate-image"
@@ -250,28 +306,42 @@ class GenerateImageInfer(BaseModel):
               negative_prompt: str = "",
               seed: int = None,
               steps: int = 28,
+              scale: float = 6.0,
               cfg_rescale: int = 0,
               sampler: Union[Sampler, str] = Sampler.K_EULER,
               width: int = 832,
               height: int = 1216,
               qualityToggle: bool = True,
               ucPreset: Union[UCPreset, int] = UCPreset.TYPE0,
+              image: Union[str, bytes] = None,
+              add_original_image: bool = None,
+              strength: float = None,
+              mask: Union[str, bytes] = None,
+              controlnet_model: Union[ControlNetModel, str] = None,
+              controlnet_condition: str = None,
               **kwargs
               ):
         """
-        正负面, step, cfg, 采样方式, seed
-        :param ucPreset:  0: 重型, 1: 轻型, 2: 人物
-        :param qualityToggle:  是否开启质量
-        :param prompt: 输入
-        :param model: 模型
+        The build Function, more convenient to create a GenerateImageInfer instance.
+        :param ucPreset:  0: Heavy, 1: Light, 2: Character
+        :param qualityToggle:  Is quality toggle
+        :param prompt: Input prompt for generate image
+        :param model: The model for generate image
         :param action: Mode for img generate [generate, img2img, infill]
-        :param negative_prompt: 负面
-        :param seed: 随机种子
-        :param steps: 步数
-        :param cfg_rescale: 0-1
-        :param sampler: 采样方式
+        :param negative_prompt: The content of negative prompt
+        :param seed: The seed for generate image
+        :param steps: The steps for generate image
+        :param scale: Prompt Guidance
+        :param cfg_rescale: Prompt Guidance Rescale 0-1 lower is more creative
+        :param sampler: The sampler for generate image
         :param width: 宽
         :param height: 高
+        :param image: 图片
+        :param add_original_image: 是否添加原始图片
+        :param strength: IMG2IMG 强度
+        :param mask: Inpainting mask
+        :param controlnet_model: 控制网络模型
+        :param controlnet_condition: 控制网络条件
         :return: self
         """
         assert isinstance(prompt, str)
@@ -280,12 +350,19 @@ class GenerateImageInfer(BaseModel):
             "negative_prompt": _negative_prompt,
             "seed": seed,
             "steps": steps,
+            "scale": scale,
             "cfg_rescale": cfg_rescale,
             "sampler": sampler,
             "width": width,
             "height": height,
             "qualityToggle": qualityToggle,
             "ucPreset": ucPreset,
+            "image": image,
+            "add_original_image": add_original_image,
+            "strength": strength,
+            "mask": mask,
+            "controlnet_model": controlnet_model,
+            "controlnet_condition": controlnet_condition
         })
         # 清理空值
         param = {k: v for k, v in kwargs.items() if v is not None}
@@ -296,18 +373,52 @@ class GenerateImageInfer(BaseModel):
             parameters=cls.Params(**param)
         )
 
-    async def generate(self, session: Union[AsyncSession, "CredentialBase"],
-                       *,
-                       remove_sign: bool = False) -> ImageGenerateResp:
+    async def necessary_headers(self, request_data) -> dict:
+        """
+        :param request_data:
+        :return:
+        """
+        return {
+            "Host": urlparse(self.endpoint).netloc,
+            "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",
+            "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
+                      ) -> ImageGenerateResp:
         """
         生成图片
+        :param override_headers:
         :param session:  session
         :param remove_sign:  移除追踪信息
         :return:
         """
-        if isinstance(session, CredentialBase):
-            session = await session.get_session()
+        # Data Build
         request_data = self.model_dump(mode="json", exclude_none=True)
+        # Header
+        if isinstance(session, AsyncSession):
+            session.headers.update(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:
             assert hasattr(session, "post"), "session must have post method."
@@ -365,7 +476,7 @@ class GenerateImageInfer(BaseModel):
                 ),
                 files=unzip_content
             )
-        except RequestsError as exc:
+        except curl_cffi.requests.errors.RequestsError as exc:
             logger.exception(exc)
             raise RuntimeError(f"An AsyncSession error occurred: {exc}")
         except httpx.HTTPError as exc:

novelai_python-0.2.5/src/novelai_python/sdk/schema.py

@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/10 上午11:15
+# @Author  : sudoskys
+# @File    : schema.py
+# @Software: PyCharm
+from abc import ABC, abstractmethod
+from typing import Optional, Union
+
+from curl_cffi.requests import AsyncSession
+from pydantic import BaseModel, PrivateAttr
+
+from ..credential import CredentialBase
+
+
+class ApiBaseModel(BaseModel, ABC):
+    _endpoint: Optional[str] = PrivateAttr()
+
+    @property
+    def endpoint(self):
+        return self._endpoint
+
+    @endpoint.setter
+    def endpoint(self, value):
+        self._endpoint = value
+
+    @abstractmethod
+    async def necessary_headers(self, request_data) -> dict:
+        raise NotImplementedError()
+
+    @abstractmethod
+    async def request(self,
+                      session: Union[AsyncSession, CredentialBase],
+                      *,
+                      override_headers: Optional[dict] = None
+                      ):
+        raise NotImplementedError()

{novelai_python-0.2.3 → novelai_python-0.2.5}/src/novelai_python/sdk/user/__init__.py

@@ -3,3 +3,4 @@
 # @Author  : sudoskys
 # @File    : __init__.py.py
 # @Software: PyCharm
+

{novelai_python-0.2.3 → novelai_python-0.2.5}/src/novelai_python/sdk/user/information.py

@@ -5,18 +5,20 @@
 # @Software: PyCharm
 from typing import Optional, Union
 
+import curl_cffi
 import httpx
-from curl_cffi.requests import AsyncSession, RequestsError
+from curl_cffi.requests import AsyncSession
 from loguru import logger
-from pydantic import BaseModel, PrivateAttr
+from pydantic import PrivateAttr
 
+from ..schema import ApiBaseModel
 from ..._exceptions import APIError, AuthError
 from ..._response.user.information import InformationResp
 from ...credential import CredentialBase
 from ...utils import try_jsonfy
 
 
-class Information(BaseModel):
+class Information(ApiBaseModel):
     _endpoint: Optional[str] = PrivateAttr("https://api.novelai.net")
 
     @property
@@ -31,22 +33,43 @@ class Information(BaseModel):
     def endpoint(self, value):
         self._endpoint = value
 
+    async def necessary_headers(self, request_data) -> dict:
+        return {
+            "Host": "api.novelai.net",
+            "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",
+            "Referer": "https://novelai.net/",
+            "Content-Type": "application/json",
+            "Origin": "https://novelai.net"
+        }
+
     async def request(self,
                       session: Union[AsyncSession, CredentialBase],
+                      *,
+                      override_headers: Optional[dict] = None
                       ) -> InformationResp:
         """
         Request to get user subscription information
-        :param session: 
+        :param override_headers:
+        :param session:
         :return: 
         """
-        if isinstance(session, CredentialBase):
-            session = await session.get_session()
         request_data = {}
+        # Data Build
+        if isinstance(session, AsyncSession):
+            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
+        if override_headers:
+            session.headers.clear()
+            session.headers.update(override_headers)
         logger.debug("Information")
         try:
             assert hasattr(session, "get"), "session must have get method."
             response = await session.get(
-                self.base_url,
+                self.base_url
             )
             if "application/json" not in response.headers.get('Content-Type') or response.status_code != 200:
                 logger.error(
@@ -74,7 +97,7 @@ class Information(BaseModel):
                     raise APIError(message, request=request_data, status_code=status_code, response=_msg)
                 raise APIError(message, request=request_data, status_code=status_code, response=_msg)
             return InformationResp.model_validate(response.json())
-        except RequestsError as exc:
+        except curl_cffi.requests.errors.RequestsError as exc:
             logger.exception(exc)
             raise RuntimeError(f"An AsyncSession error occurred: {exc}")
         except httpx.HTTPError as exc:

{novelai_python-0.2.3 → novelai_python-0.2.5}/src/novelai_python/sdk/user/login.py

@@ -4,19 +4,22 @@
 # @File    : login.py
 # @Software: PyCharm
 import json
-from typing import Optional
+from typing import Optional, Union
 
+import curl_cffi
 import httpx
-from curl_cffi.requests import AsyncSession, RequestsError
+from curl_cffi.requests import AsyncSession
 from loguru import logger
-from pydantic import BaseModel, PrivateAttr, Field
+from pydantic import PrivateAttr, Field
 
+from ..schema import ApiBaseModel
 from ..._exceptions import APIError
 from ..._response.user.login import LoginResp
+from ...credential import CredentialBase
 from ...utils import try_jsonfy, encode_access_key
 
 
-class Login(BaseModel):
+class Login(ApiBaseModel):
     _endpoint: Optional[str] = PrivateAttr("https://api.novelai.net")
     _session: Optional[AsyncSession] = PrivateAttr(None)
     key: str = Field(..., description="User's key")
@@ -59,13 +62,28 @@ class Login(BaseModel):
         """
         return cls(key=encode_access_key(user_name, password))
 
+    async def necessary_headers(self, request_data) -> dict:
+        return {}
+
     async def request(self,
+                      session: Union[AsyncSession, CredentialBase] = None,
+                      *,
+                      override_headers: Optional[dict] = None,
                       ) -> LoginResp:
         """
         Request to get user access token
         :return:
         """
+        # Data Build
         request_data = self.model_dump(mode="json", exclude_none=True)
+        if isinstance(session, AsyncSession):
+            session.headers.update(self.necessary_headers(request_data))
+        elif isinstance(session, CredentialBase):
+            session = await session.get_session(update_headers=await self.necessary_headers(request_data))
+        # Header
+        if override_headers:
+            session.headers.clear()
+            session.headers.update(override_headers)
         logger.debug("Login")
         try:
             assert hasattr(self.session, "post"), "session must have get method."
@@ -97,7 +115,7 @@ class Login(BaseModel):
                     raise APIError(message, request=request_data, status_code=status_code, response=_msg)
                 raise APIError(message, request=request_data, status_code=status_code, response=_msg)
             return LoginResp.model_validate(response.json())
-        except RequestsError as exc:
+        except curl_cffi.requests.errors.RequestsError as exc:
             logger.exception(exc)
             raise RuntimeError(f"An AsyncSession error occurred: {exc}")
         except httpx.HTTPError as exc:

{novelai_python-0.2.3 → novelai_python-0.2.5}/src/novelai_python/sdk/user/subscription.py

@@ -5,18 +5,20 @@
 # @Software: PyCharm
 from typing import Optional, Union
 
+import curl_cffi
 import httpx
-from curl_cffi.requests import AsyncSession, RequestsError
+from curl_cffi.requests import AsyncSession
 from loguru import logger
-from pydantic import BaseModel, PrivateAttr
+from pydantic import PrivateAttr
 
+from ..schema import ApiBaseModel
 from ..._exceptions import APIError, AuthError
 from ..._response.user.subscription import SubscriptionResp
 from ...credential import CredentialBase
 from ...utils import try_jsonfy
 
 
-class Subscription(BaseModel):
+class Subscription(ApiBaseModel):
     _endpoint: Optional[str] = PrivateAttr("https://api.novelai.net")
 
     @property
@@ -31,22 +33,50 @@ class Subscription(BaseModel):
     def endpoint(self, value):
         self._endpoint = value
 
+    async def necessary_headers(self, request_data) -> dict:
+        return {
+            "Host": "api.novelai.net",
+            "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",
+            "Referer": "https://novelai.net/",
+            "Content-Type": "application/json",
+            "Origin": "https://novelai.net",
+            "Connection": "keep-alive",
+            "Sec-Fetch-Dest": "empty",
+            "Sec-Fetch-Mode": "cors",
+            "Sec-Fetch-Site": "same-site",
+            "Pragma": "no-cache",
+            "Cache-Control": "no-cache",
+            "TE": "trailers",
+        }
+
     async def request(self,
                       session: Union[AsyncSession, CredentialBase],
+                      *,
+                      override_headers: Optional[dict] = None
                       ) -> SubscriptionResp:
         """
         Request to get user subscription information
-        :param session: 
+        :param override_headers:
+        :param session:
         :return: 
         """
-        if isinstance(session, CredentialBase):
-            session = await session.get_session()
+        # Data Build
         request_data = {}
+        if isinstance(session, AsyncSession):
+            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
+        if override_headers:
+            session.headers.clear()
+            session.headers.update(override_headers)
         logger.debug("Subscription")
         try:
             assert hasattr(session, "get"), "session must have get method."
             response = await session.get(
-                self.base_url,
+                url=self.base_url,
             )
             if "application/json" not in response.headers.get('Content-Type') or response.status_code != 200:
                 logger.error(
@@ -74,7 +104,7 @@ class Subscription(BaseModel):
                     raise APIError(message, request=request_data, status_code=status_code, response=_msg)
                 raise APIError(message, request=request_data, status_code=status_code, response=_msg)
             return SubscriptionResp.model_validate(response.json())
-        except RequestsError as exc:
+        except curl_cffi.requests.errors.RequestsError as exc:
             logger.exception(exc)
             raise RuntimeError(f"An AsyncSession error occurred: {exc}")
         except httpx.HTTPError as exc:

{novelai_python-0.2.3 → novelai_python-0.2.5}/src/novelai_python/server.py

@@ -102,7 +102,7 @@ async def generate_image(
     :return:
     """
     try:
-        _result = await req.generate(session=get_session(current_token), remove_sign=True)
+        _result = await req.request(session=get_session(current_token), remove_sign=True)
         zip_file_bytes = io.BytesIO()
         with zipfile.ZipFile(zip_file_bytes, mode="w", compression=zipfile.ZIP_DEFLATED) as zip_file:
             for file in _result.files:

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

@@ -0,0 +1,7 @@
+# -*- 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()))

{novelai_python-0.2.3 → novelai_python-0.2.5}/tests/test_server_run.py

@@ -28,7 +28,7 @@ def test_subscription_without_api_key(mock_subscription):
 
 @patch('src.novelai_python.server.GenerateImageInfer')
 def test_generate_image_without_api_key(mock_generate_image):
-    mock_generate_image.return_value.generate.return_value = {"status": "image generated"}
+    mock_generate_image.return_value.request.return_value = {"status": "image generated"}
     response = client.post("/ai/generate_image")
     assert response.status_code == 403
 

{novelai_python-0.2.3 → novelai_python-0.2.5}/tests/test_user_information.py

@@ -34,6 +34,7 @@ async def test_request_method_successful():
     }
     session = mock.MagicMock(spec=AsyncSession)
     session.get = mock.AsyncMock(return_value=successful_response)
+    session.headers = {}
 
     info = Information()
     resp = await info.request(session)
@@ -48,6 +49,7 @@ async def test_request_method_unauthorized_error():
     auth_response.json.return_value = {"statusCode": 401, "message": "Access Token is incorrect."}
 
     session = mock.MagicMock(spec=AsyncSession)
+    session.headers = {}
     session.get = mock.AsyncMock(return_value=auth_response)
 
     info = Information()
@@ -66,6 +68,7 @@ async def test_request_method_unknown_error():
     unknown_error_response.json.return_value = {"statusCode": 500, "message": "An unknown error occurred."}
 
     session = mock.MagicMock(spec=AsyncSession)
+    session.headers = {}
     session.get = mock.AsyncMock(return_value=unknown_error_response)
 
     info = Information()

{novelai_python-0.2.3 → novelai_python-0.2.5}/tests/test_user_login.py

@@ -21,6 +21,7 @@ async def test_successful_user_login():
     }
     session = mock.MagicMock(spec=AsyncSession)
     session.post = mock.AsyncMock(return_value=successful_login_response)
+    session.headers = {}
 
     login = Login(key="encoded_key")
     login.session = session
@@ -40,6 +41,7 @@ async def test_validation_error_during_login():
     }
     session = mock.MagicMock(spec=AsyncSession)
     session.post = mock.AsyncMock(return_value=validation_error_response)
+    session.headers = {}
 
     login = Login(key="encoded_key")
     login.session = session
@@ -61,6 +63,7 @@ async def test_incorrect_access_key_during_login():
     }
     session = mock.MagicMock(spec=AsyncSession)
     session.post = mock.AsyncMock(return_value=incorrect_key_response)
+    session.headers = {}
 
     login = Login(key="encoded_key")
     login.session = session
@@ -82,6 +85,8 @@ async def test_unknown_error_during_login():
     }
     session = mock.MagicMock(spec=AsyncSession)
     session.post = mock.AsyncMock(return_value=unknown_error_response)
+    session.headers = {}
+
     login = Login(key="encoded_key")
     login.session = session
     with pytest.raises(APIError) as e:

{novelai_python-0.2.3 → novelai_python-0.2.5}/tests/test_user_subscription.py

@@ -49,6 +49,8 @@ async def test_successful_subscription_request():
     }
     session = mock.MagicMock(spec=AsyncSession)
     session.get = mock.AsyncMock(return_value=successful_response)
+    session.headers = {}
+
     subscription = Subscription()
     resp = await subscription.request(session)
     assert isinstance(resp, SubscriptionResp)
@@ -65,6 +67,7 @@ async def test_incorrect_access_token_subscription_request():
     }
     session = mock.MagicMock(spec=AsyncSession)
     session.get = mock.AsyncMock(return_value=incorrect_token_response)
+    session.headers = {}
 
     subscription = Subscription()
     with pytest.raises(AuthError) as e:
@@ -85,6 +88,7 @@ async def test_unknown_error_subscription_request():
     }
     session = mock.MagicMock(spec=AsyncSession)
     session.get = mock.AsyncMock(return_value=unknown_error_response)
+    session.headers = {}
 
     subscription = Subscription()
     with pytest.raises(APIError) as e: