novelai-python 0.4.11__tar.gz → 0.4.12__tar.gz

{novelai_python-0.4.11 → novelai_python-0.4.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.4.11
+Version: 0.4.12
 Summary: NovelAI Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache License 2.0

{novelai_python-0.4.11 → novelai_python-0.4.12}/pyproject.toml

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

{novelai_python-0.4.11 → novelai_python-0.4.12}/src/novelai_python/__init__.py

@@ -11,6 +11,7 @@ from ._exceptions import (
     SessionHttpError
 )
 from .credential import JwtCredential, LoginCredential, ApiCredential
+from .sdk import AugmentImageInfer
 from .sdk import GenerateImageInfer, ImageGenerateResp
 from .sdk import Information, InformationResp
 from .sdk import LLM, LLMResp
@@ -31,6 +32,8 @@ __all__ = [
     "GenerateImageInfer",
     "ImageGenerateResp",
 
+    "AugmentImageInfer",
+
     "VoiceGenerate",
     "VoiceResponse",
 

novelai_python-0.4.12/src/novelai_python/credential/ApiToken.py

@@ -0,0 +1,41 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/8 下午3:05
+# @Author  : sudoskys
+# @File    : ApiToken.py
+# @Software: PyCharm
+
+from curl_cffi.requests import AsyncSession
+from loguru import logger
+from pydantic import SecretStr, Field, field_validator
+
+from ._base import CredentialBase, FAKE_UA
+
+
+class ApiCredential(CredentialBase):
+    """
+    ApiCredential is the base class for all credential.
+    """
+    api_token: SecretStr = Field(None, description="api token")
+
+    async def get_session(self, timeout: int = 180, update_headers: dict = None):
+        headers = {
+            "Accept": "*/*",
+            "Accept-Encoding": "gzip, deflate, br",
+            "User-Agent": FAKE_UA.edge,
+            "Authorization": f"Bearer {self.api_token.get_secret_value()}",
+            "Content-Type": "application/json",
+            "Origin": "https://novelai.net",
+            "Referer": "https://novelai.net/",
+        }
+
+        if update_headers:
+            assert isinstance(update_headers, dict), "update_headers must be a dict"
+            headers.update(update_headers)
+
+        return AsyncSession(timeout=timeout, headers=headers, impersonate="edge101")
+
+    @field_validator('api_token')
+    def check_api_token(cls, v: SecretStr):
+        if not v.get_secret_value().startswith("pst"):
+            logger.warning("api token should start with pst-")
+        return v

novelai_python-0.4.12/src/novelai_python/credential/JwtToken.py

@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/26 上午11:04
+# @Author  : sudoskys
+# @File    : JwtToken.py
+# @Software: PyCharm
+from curl_cffi.requests import AsyncSession
+from loguru import logger
+from pydantic import SecretStr, Field, field_validator
+
+from ._base import CredentialBase, FAKE_UA
+
+
+class JwtCredential(CredentialBase):
+    """
+    JwtCredential is the base class for all credential.
+    """
+    jwt_token: SecretStr = Field(None, description="jwt token")
+
+    async def get_session(self, timeout: int = 180, update_headers: dict = None):
+        headers = {
+            "Accept": "*/*",
+            "User-Agent": FAKE_UA.edge,
+            "Accept-Encoding": "gzip, deflate, br",
+            "Authorization": f"Bearer {self.jwt_token.get_secret_value()}",
+            "Content-Type": "application/json",
+            "Origin": "https://novelai.net",
+            "Referer": "https://novelai.net/",
+        }
+
+        if update_headers:
+            assert isinstance(update_headers, dict), "update_headers must be a dict"
+            headers.update(update_headers)
+
+        return AsyncSession(timeout=timeout, headers=headers, impersonate="edge101")
+
+    @field_validator('jwt_token')
+    def check_jwt_token(cls, v: SecretStr):
+        if not v.get_secret_value().startswith("ey"):
+            logger.warning("jwt_token should start with ey")
+        return v

novelai_python-0.4.12/src/novelai_python/credential/UserAuth.py

@@ -0,0 +1,47 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/7 下午12:14
+# @Author  : sudoskys
+# @File    : UserAuth.py
+import time
+from typing import Optional
+
+from curl_cffi.requests import AsyncSession
+from pydantic import SecretStr, Field
+
+from ._base import CredentialBase, FAKE_UA
+
+
+class LoginCredential(CredentialBase):
+    """
+    JwtCredential is the base class for all credential.
+    """
+    username: str = Field(None, description="username")
+    password: SecretStr = Field(None, description="password")
+    _session_headers: dict = {}
+    _update_at: Optional[int] = None
+
+    async def get_session(self, timeout: int = 180, update_headers: dict = None):
+        headers = {
+            "Accept": "*/*",
+            "User-Agent": FAKE_UA.edge,
+            "Accept-Encoding": "gzip, deflate, br",
+            "Authorization": "Bearer ",
+            "Content-Type": "application/json",
+            "Origin": "https://novelai.net",
+            "Referer": "https://novelai.net/",
+        }
+
+        # 30 天有效期
+        if not self._session_headers 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()
+            headers["Authorization"] = f"Bearer {resp.accessToken}"
+            self._session_headers = headers
+            self._update_at = int(time.time())
+        else:
+            headers.update(self._session_headers)
+
+        if update_headers:
+            headers.update(update_headers)
+
+        return AsyncSession(timeout=timeout, headers=headers, impersonate="edge101")

{novelai_python-0.4.11 → novelai_python-0.4.12}/src/novelai_python/sdk/__init__.py

@@ -4,6 +4,7 @@
 # @File    : __init__.py.py
 # @Software: PyCharm
 
+from .ai.augment_image import AugmentImageInfer  # noqa 401
 from .ai.generate import LLM, LLMResp  # noqa 401
 from .ai.generate_image import GenerateImageInfer, ImageGenerateResp  # noqa 401
 from .ai.generate_image.suggest_tags import SuggestTags, SuggestTagsResp  # noqa 401

novelai_python-0.4.12/src/novelai_python/sdk/ai/augment_image/__init__.py

@@ -0,0 +1,248 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/13 下午8:08
+# @Author  : sudoskys
+# @File    : __init__.py.py
+# @Software: PyCharm
+import base64
+import io
+import json
+import pathlib
+from copy import deepcopy
+from io import BytesIO
+from typing import Optional, Union, IO
+from urllib.parse import urlparse
+from zipfile import ZipFile
+
+import curl_cffi
+import httpx
+from PIL import Image
+from curl_cffi.requests import AsyncSession
+from loguru import logger
+from pydantic import ConfigDict, PrivateAttr, model_validator, Field
+from tenacity import retry, stop_after_attempt, wait_random, retry_if_exception
+
+from ._enum import ReqType, Moods
+from ...schema import ApiBaseModel
+from ...._exceptions import APIError, AuthError, ConcurrentGenerationError, SessionHttpError
+from ...._response.ai.generate_image import ImageGenerateResp, RequestParams
+from ....credential import CredentialBase
+from ....utils import try_jsonfy
+
+
+class AugmentImageInfer(ApiBaseModel):
+    _endpoint: str = PrivateAttr("https://image.novelai.net")
+
+    @property
+    def endpoint(self):
+        return self._endpoint
+
+    @endpoint.setter
+    def endpoint(self, value):
+        self._endpoint = value
+
+    req_type: ReqType = Field(..., description="Type of augmentation")
+    width: int = Field(..., description="Width of the image")
+    height: int = Field(..., description="Height of the image")
+    image: Union[str, bytes] = Field(..., description="Base64 encoded image")
+    prompt: Optional[str] = None
+    defry: Optional[int] = Field(0, ge=0, le=5, multiple_of=1)
+    model_config = ConfigDict(extra="ignore")
+
+    @model_validator(mode="after")
+    def image_validator(self):
+        if isinstance(self.image, bytes):
+            self.image = base64.b64encode(self.image).decode("utf-8")
+        if isinstance(self.image, str) and self.image.startswith("data:"):
+            raise ValueError("Invalid `image` format, must be base64 encoded directly.")
+        if isinstance(self.image, str) and self.image.startswith("+vv"):
+            raise ValueError("Invalid `image` format, must be encoded correctly.")
+        return self
+
+    @property
+    def base_url(self):
+        return f"{self.endpoint.strip('/')}/ai/augment-image"
+
+    def calculate_cost(self, is_opus: bool = False) -> float:
+        raise NotImplementedError("This method is not implemented yet")
+
+    @staticmethod
+    def _to_bytes_io(image: Union[bytes, IO, pathlib.Path]) -> io.BytesIO:
+        if isinstance(image, bytes):
+            return io.BytesIO(image)
+        elif isinstance(image, IO):
+            return image
+        elif isinstance(image, pathlib.Path):
+            with open(image, 'rb') as f:
+                return io.BytesIO(f.read())
+        else:
+            raise ValueError("Invalid image input type. Expected bytes, IO, or pathlib.Path.")
+
+    @classmethod
+    def build(cls,
+              image: Union[bytes, IO, pathlib.Path],
+              req_type: ReqType = ReqType.SKETCH,
+              *,
+              defry: Optional[int] = None,
+              prompt: Optional[str] = None,
+              mood: Optional[Moods] = None,
+              ):
+        """
+        Build the request data
+        :param req_type: Request type
+        :param image: Image data
+        :param defry: Defry level
+        :param prompt: Prompt
+        :param mood: Mood
+        :raises ValueError: When image is not valid
+        :return: AugmentImageInfer
+        """
+        try:
+            image_io = cls._to_bytes_io(image)
+            with Image.open(image_io) as img:
+                width, height = img.size
+                image_io.seek(0)
+                image_data = image_io.read()
+        except Exception as e:
+            raise ValueError("Error when opening image") from e
+
+        return cls(
+            req_type=req_type,
+            defry=defry,
+            image=image_data,
+            width=width,
+            height=height,
+            prompt=f"{mood.value};;{prompt}" if mood else prompt,
+        )
+
+    async def necessary_headers(self, request_data) -> dict:
+        """
+        :param request_data:
+        :return:
+        """
+        return {
+            "Host": urlparse(self.endpoint).netloc,
+            "Accept": "*/*",
+            "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",
+            'priority': "u=1, i"
+        }
+
+    @retry(
+        wait=wait_random(min=1, max=3),
+        stop=stop_after_attempt(3),
+        retry=retry_if_exception(lambda e: hasattr(e, "code") and str(e.code) == "500"),
+        reraise=True
+    )
+    async def request(self,
+                      session: Union[AsyncSession, "CredentialBase"],
+                      *,
+                      override_headers: Optional[dict] = None,
+                      ) -> ImageGenerateResp:
+        """
+        生成图片
+        :param override_headers: the headers to override
+        :param session:  session
+        :return:
+        """
+        # Data Build
+        request_data = self.model_dump(mode="json", exclude_none=True)
+        async with session if isinstance(session, AsyncSession) else await session.get_session() as sess:
+            # Header
+            sess.headers.update(await self.necessary_headers(request_data))
+            if override_headers:
+                sess.headers.clear()
+                sess.headers.update(override_headers)
+            try:
+                _log_data = deepcopy(request_data)
+                _log_data.update({
+                    "image": "base64 data"
+                })
+                logger.debug(f"Request Data: {_log_data}")
+                del _log_data
+            except Exception as e:
+                logger.warning(f"Error when print log data: {e}")
+            try:
+                assert hasattr(sess, "post"), "session must have post method."
+                response = await sess.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.warning(
+                        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)
+                        unzip_content.append((filename, data))
+                return ImageGenerateResp(
+                    meta=RequestParams(
+                        endpoint=self.base_url,
+                        raw_request=request_data,
+                    ),
+                    files=unzip_content
+                )
+            except curl_cffi.requests.errors.RequestsError as exc:
+                logger.exception(exc)
+                raise SessionHttpError(
+                    "An AsyncSession RequestsError occurred, maybe SSL error. Try again later!") from exc
+            except httpx.HTTPError as exc:
+                logger.exception(exc)
+                raise SessionHttpError("An HTTPError occurred, maybe SSL error. Try again later!") from exc
+            except APIError as e:
+                raise e
+            except Exception as e:
+                logger.opt(exception=e).exception("An Unexpected error occurred")
+                raise e

novelai_python-0.4.12/src/novelai_python/sdk/ai/augment_image/_enum.py

@@ -0,0 +1,43 @@
+from enum import Enum
+
+
+class ReqType(Enum):
+    """
+    typing.Literal["bg-removal", "colorize", "lineart", "sketch", "emotion", "declutter"]
+    """
+    BG_REMOVAL = "bg-removal"
+    COLORIZE = "colorize"
+    LINEART = "lineart"
+    SKETCH = "sketch"
+    EMOTION = "emotion"
+    DECLUTTER = "declutter"
+
+
+class Moods(Enum):
+    """
+    The mood of the character in the image
+    """
+    Neutral = "neutral"
+    Happy = "happy"
+    Saf = "sad"
+    Angry = "angry"
+    Scared = "scared"
+    Surprised = "surprised"
+    Tired = "tired"
+    Excited = "excited"
+    Nervous = "nervous"
+    Thinking = "thinking"
+    Confused = "confused"
+    Shy = "shy"
+    Disgusted = "disgusted"
+    Smug = "smug"
+    Bored = "bored"
+    Laughing = "laughing"
+    Irritated = "irritated"
+    Aroused = "aroused"
+    Embarrassed = "embarrassed"
+    Worried = "worried"
+    Love = "love"
+    Determined = "determined"
+    Hurt = "hurt"
+    Playful = "playful"

{novelai_python-0.4.11 → novelai_python-0.4.12}/src/novelai_python/sdk/ai/generate/__init__.py

@@ -156,80 +156,77 @@ class LLM(ApiBaseModel):
         request_data = self.model_dump(mode="json", exclude_none=True)
         request_data["order"] = generate_order(request_data) or []
         assert request_data.get("parameters"), "Parameters is required"
-        # 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)
-        # Log
-        logger.debug(f"LLM request data: {request_data}")
-        # Request
-        try:
-            assert hasattr(session, "post"), "session must have post method."
-            response = await session.post(
-                self.base_url,
-                json=request_data,
-                stream=False
-            )
-            header_type = response.headers.get('Content-Type')
-            if "application/json" not in header_type:
-                try:
-                    _msg = response.json()
-                except Exception as e:
-                    _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 == 400:
-                    raise APIError(
-                        f"A validation error occured. {message}",
-                        request=request_data, code=status_code, response=_msg
-                    )
+        async with session if isinstance(session, AsyncSession) else await session.get_session() as sess:
+            # Header
+            sess.headers.update(await self.necessary_headers(request_data))
+            if override_headers:
+                sess.headers.clear()
+                sess.headers.update(override_headers)
+            # Log
+            logger.debug(f"LLM request data: {request_data}")
+            # Request
+            try:
+                assert hasattr(sess, "post"), "session must have post method."
+                response = await sess.post(
+                    self.base_url,
+                    json=request_data,
+                    stream=False
+                )
+                header_type = response.headers.get('Content-Type')
+                if "application/json" not in header_type:
+                    try:
+                        _msg = response.json()
+                    except Exception as e:
+                        _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 == 400:
+                        raise APIError(
+                            f"A validation error occured. {message}",
+                            request=request_data, code=status_code, response=_msg
+                        )
 
-                elif status_code == 401:
-                    raise APIError(
-                        f"Access Token is incorrect. {message}",
-                        request=request_data, code=status_code, response=_msg
-                    )
-                elif status_code == 402:
-                    raise APIError(
-                        f"An active subscription is required to access this endpoint. {message}",
-                        request=request_data, code=status_code, response=_msg
-                    )
-                elif status_code == 409:
-                    raise APIError(
-                        f"A conflict error occured. {message}",
-                        request=request_data, code=status_code, response=_msg
-                    )
+                    elif status_code == 401:
+                        raise APIError(
+                            f"Access Token is incorrect. {message}",
+                            request=request_data, code=status_code, response=_msg
+                        )
+                    elif status_code == 402:
+                        raise APIError(
+                            f"An active subscription is required to access this endpoint. {message}",
+                            request=request_data, code=status_code, response=_msg
+                        )
+                    elif status_code == 409:
+                        raise APIError(
+                            f"A conflict error occured. {message}",
+                            request=request_data, code=status_code, response=_msg
+                        )
+                    else:
+                        raise APIError(
+                            f"An unknown error occured. {response.status_code} {message}",
+                            request=request_data, code=status_code, response=_msg
+                        )
                 else:
-                    raise APIError(
-                        f"An unknown error occured. {response.status_code} {message}",
-                        request=request_data, code=status_code, response=_msg
+                    output = response.json().get("output", None)
+                    assert output, APIError("No Content Returned",
+                                            request=request_data, code=response.status_code, response=response.content
+                                            )
+                    return LLMResp(
+                        output=output,
+                        text=LLMResp.decode_token(model=self.model, token_str=output)
                     )
-            else:
-                output = response.json().get("output", None)
-                assert output, APIError("No Content Returned",
-                                        request=request_data, code=response.status_code, response=response.content
-                                        )
-                return LLMResp(
-                    output=output,
-                    text=LLMResp.decode_token(model=self.model, token_str=output)
+            except curl_cffi.requests.errors.RequestsError as exc:
+                logger.exception(exc)
+                raise SessionHttpError(
+                    "An AsyncSession RequestsError occurred, perhaps it's an SSL error. Please try again later!"
+                )
+            except httpx.HTTPError as exc:
+                logger.exception(exc)
+                raise SessionHttpError(
+                    "An HTTPError occurred, perhaps it's an SSL error. Please try again later!"
                 )
-        except curl_cffi.requests.errors.RequestsError as exc:
-            logger.exception(exc)
-            raise SessionHttpError(
-                "An AsyncSession RequestsError occurred, perhaps it's an SSL error. Please try again later!"
-            )
-        except httpx.HTTPError as exc:
-            logger.exception(exc)
-            raise SessionHttpError(
-                "An HTTPError occurred, perhaps it's an SSL error. Please try again later!"
-            )
-        except APIError as e:
-            raise e
-        except Exception as e:
-            logger.opt(exception=e).exception("An unexpected error occurred")
-            raise e
+            except APIError as e:
+                raise e
+            except Exception as e:
+                logger.opt(exception=e).exception("An unexpected error occurred")
+                raise e