novelai-python 0.1.0__tar.gz → 0.1.1__tar.gz

novelai_python-0.1.1/PKG-INFO

@@ -0,0 +1,75 @@
+Metadata-Version: 2.1
+Name: novelai-python
+Version: 0.1.1
+Summary: Novelai Python Binding With Pydantic
+Author-Email: sudoskys <coldlando@hotmail.com>
+License: Apache 2.0
+Project-URL: Repository, https://github.com/LlmKira/novelai-python/
+Project-URL: Issues, https://github.com/LlmKira/novelai-python/issues
+Requires-Python: >=3.8
+Requires-Dist: elara>=0.5.5
+Requires-Dist: loguru>=0.7.2
+Requires-Dist: pydantic-settings>=2.1.0
+Requires-Dist: pydantic>=2.5.3
+Requires-Dist: httpx>=0.26.0
+Requires-Dist: shortuuid>=1.0.11
+Requires-Dist: Pillow>=10.2.0
+Requires-Dist: curl-cffi>=0.5.10
+Description-Content-Type: text/markdown
+
+# novelai-python
+
+[![PyPI version](https://badge.fury.io/py/novelai_python.svg)](https://badge.fury.io/py/novelai_gen)
+[![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.
+
+### Roadmap 🚧
+
+- [x] /ai/generate-image
+- [ ] /ai/generate-image/suggest-tags
+- [ ] /ai/annotate-image
+- [ ] /ai/classify
+- [ ] /ai/upscale
+- [ ] /ai/generate-prompt
+- [ ] /ai/generate
+- [ ] /ai/generate-voice
+
+### Usage 🖥️
+
+```python
+import asyncio
+import os
+
+from dotenv import load_dotenv
+from pydantic import SecretStr
+
+from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential
+
+load_dotenv()
+
+enhance = "year 2023,dynamic angle,  best quality, amazing quality, very aesthetic, absurdres"
+
+
+async def main():
+    globe_s = JwtCredential(jwt_token=SecretStr(os.getenv("NOVELAI_JWT")))
+    _res = await GenerateImageInfer.build(
+        prompt=f"1girl,{enhance}").generate(
+        session=globe_s)
+    _res: ImageGenerateResp
+    print(_res.meta)
+    file = _res.files[0]
+    with open(file[0], "wb") as f:
+        f.write(file[1])
+
+
+loop = asyncio.get_event_loop()
+loop.run_until_complete(main())
+
+```
+
+## Acknowledgements 🙏
+
+[BackEnd](https://api.novelai.net/docs)
+[novelai-api](https://github.com/Aedial/novelai-api)
+

novelai_python-0.1.1/README.md

@@ -0,0 +1,56 @@
+# novelai-python
+
+[![PyPI version](https://badge.fury.io/py/novelai_python.svg)](https://badge.fury.io/py/novelai_gen)
+[![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.
+
+### Roadmap 🚧
+
+- [x] /ai/generate-image
+- [ ] /ai/generate-image/suggest-tags
+- [ ] /ai/annotate-image
+- [ ] /ai/classify
+- [ ] /ai/upscale
+- [ ] /ai/generate-prompt
+- [ ] /ai/generate
+- [ ] /ai/generate-voice
+
+### Usage 🖥️
+
+```python
+import asyncio
+import os
+
+from dotenv import load_dotenv
+from pydantic import SecretStr
+
+from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential
+
+load_dotenv()
+
+enhance = "year 2023,dynamic angle,  best quality, amazing quality, very aesthetic, absurdres"
+
+
+async def main():
+    globe_s = JwtCredential(jwt_token=SecretStr(os.getenv("NOVELAI_JWT")))
+    _res = await GenerateImageInfer.build(
+        prompt=f"1girl,{enhance}").generate(
+        session=globe_s)
+    _res: ImageGenerateResp
+    print(_res.meta)
+    file = _res.files[0]
+    with open(file[0], "wb") as f:
+        f.write(file[1])
+
+
+loop = asyncio.get_event_loop()
+loop.run_until_complete(main())
+
+```
+
+## Acknowledgements 🙏
+
+[BackEnd](https://api.novelai.net/docs)
+[novelai-api](https://github.com/Aedial/novelai-api)
+

novelai_python-0.1.1/pyproject.toml

@@ -0,0 +1,35 @@
+[project]
+name = "novelai-python"
+version = "0.1.1"
+description = "Novelai Python Binding With Pydantic"
+authors = [
+    { name = "sudoskys", email = "coldlando@hotmail.com" },
+]
+dependencies = [
+    "elara>=0.5.5",
+    "loguru>=0.7.2",
+    "pydantic-settings>=2.1.0",
+    "pydantic>=2.5.3",
+    "httpx>=0.26.0",
+    "shortuuid>=1.0.11",
+    "Pillow>=10.2.0",
+    "curl-cffi>=0.5.10",
+]
+requires-python = ">=3.8"
+readme = "README.md"
+
+[project.license]
+text = "Apache 2.0"
+
+[project.urls]
+Repository = "https://github.com/LlmKira/novelai-python/"
+Issues = "https://github.com/LlmKira/novelai-python/issues"
+
+[build-system]
+requires = [
+    "pdm-backend",
+]
+build-backend = "pdm.backend"
+
+[tool.pdm]
+distribution = true

novelai_python-0.1.1/src/novelai_python/__init__.py

@@ -0,0 +1,8 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2023/11/18 上午12:18
+# @Author  : sudoskys
+# @File    : __init__.py.py
+# @Software: PyCharm
+
+from .credential import *  # noqa: F401, F403
+from .sdk import *  # noqa: F401, F403

novelai_python-0.1.1/src/novelai_python/_exceptions.py

@@ -0,0 +1,49 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/26 上午10:55
+# @Author  : sudoskys
+# @File    : _exceptions.py
+# @Software: PyCharm
+from typing import Any, Optional, Union
+from typing import TypeGuard
+
+
+def is_dict(obj: object) -> TypeGuard[dict[object, object]]:
+    """
+    Check if the object is a dict.
+    :param obj:  The object to check.
+    :return:   True if the object is a dict, False otherwise.
+    """
+    return isinstance(obj, dict)
+
+
+class NovelAiError(Exception):
+    """
+    NovelAiError is the base exception for all novelai_python errors.
+    """
+    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: Any, response: Any, status_code: str) -> None:
+        super().__init__(message)
+        self.request = request
+        self.message = message
+        self.code = status_code
+        self.response = response
+
+
+class AuthError(APIError):
+    """
+    AuthError is raised when the API returns an error.
+    """
+
+    def __init__(self, message: str, request: Any, response: Any, status_code: str) -> None:
+        super().__init__(message, request, response, status_code)

novelai_python-0.1.1/src/novelai_python/_response/__init__.py

@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/26 上午10:51
+# @Author  : sudoskys
+# @File    : __init__.py.py
+# @Software: PyCharm
+from .ai.generate_image import ImageGenerateResp
+
+__all__ = [
+    "ImageGenerateResp"
+]

novelai_python-0.1.1/src/novelai_python/_response/ai/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/26 上午11:20
+# @Author  : sudoskys
+# @File    : __init__.py
+# @Software: PyCharm

novelai_python-0.1.1/src/novelai_python/_response/ai/generate_image.py

@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/26 上午11:16
+# @Author  : sudoskys
+# @File    : text2image.py
+# @Software: PyCharm
+from typing import Tuple, List
+
+from pydantic import BaseModel
+
+
+class ImageGenerateResp(BaseModel):
+    class RequestParams(BaseModel):
+        endpoint: str
+        raw_request: dict = None
+
+    meta: RequestParams
+    files: List[Tuple[str, bytes]] = None
+
+    def query_params(self, key: str, default=None):
+        if not isinstance(self.meta.raw_request.get("parameters"), dict):
+            raise Exception("Resp parameters is not dict")
+        return self.meta.raw_request.get("parameters").get(key, default)

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

@@ -0,0 +1,26 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/26 上午11:04
+# @Author  : sudoskys
+# @File    : JwtToken.py
+# @Software: PyCharm
+from curl_cffi.requests import AsyncSession
+from pydantic import BaseModel, SecretStr, Field
+
+
+class JwtCredential(BaseModel):
+    """
+    JwtCredential is the base class for all credential.
+    """
+    jwt_token: SecretStr = Field(None, description="jwt token")
+    _session: AsyncSession = None
+
+    @property
+    def session(self, timeout: int = 180):
+        if not self._session:
+            self._session = AsyncSession(timeout=timeout, headers={
+                "Authorization": f"Bearer {self.jwt_token.get_secret_value()}",
+                "Content-Type": "application/json",
+                "Origin": "https://novelai.net",
+                "Referer": "https://novelai.net/",
+            }, impersonate="chrome110")
+        return self._session

novelai_python-0.1.1/src/novelai_python/credential/__init__.py

@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/26 上午10:51
+# @Author  : sudoskys
+# @File    : __init__.py.py
+# @Software: PyCharm
+from .JwtToken import JwtCredential
+
+__all__ = [
+    "JwtCredential"
+]

novelai_python-0.1.1/src/novelai_python/sdk/__init__.py

@@ -0,0 +1,8 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/26 上午10:51
+# @Author  : sudoskys
+# @File    : __init__.py.py
+# @Software: PyCharm
+
+from .ai.generate_image import GenerateImageInfer, ImageGenerateResp  # noqa 401
+

novelai_python-0.1.1/src/novelai_python/sdk/ai/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/26 上午11:20
+# @Author  : sudoskys
+# @File    : __init__.py.py
+# @Software: PyCharm

novelai_python-0.1.0/src/novelai_python/__init__.py → novelai_python-0.1.1/src/novelai_python/sdk/ai/generate_image.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
-# @Time    : 2023/11/18 上午12:18
+# @Time    : 2024/1/26 上午11:21
 # @Author  : sudoskys
-# @File    : __init__.py.py
+# @File    : generate_image.py
 # @Software: PyCharm
 import json
 import random
@@ -10,40 +10,17 @@ from typing import Optional, Union
 from zipfile import ZipFile
 
 import httpx
-import shortuuid
 from curl_cffi.requests import AsyncSession
-from dotenv import load_dotenv
 from loguru import logger
-from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, SecretStr
+from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator
 
-from .error import ServerError, AuthError
-from .hash import NaiPic
-from .schema import NaiResult
+from ..._exceptions import APIError, AuthError
+from ..._response import ImageGenerateResp
+from ...credential import JwtCredential
+from ...utils import try_jsonfy
 
-load_dotenv()
 
-
-class CheckError(Exception):
-    pass
-
-
-class CurlSession(BaseModel):
-    jwt_token: SecretStr
-    _session: AsyncSession = None
-
-    @property
-    def session(self):
-        if not self._session:
-            self._session = AsyncSession(timeout=180, headers={
-                "Authorization": f"Bearer {self.jwt_token.get_secret_value()}",
-                "Content-Type": "application/json",
-                "Origin": "https://novelai.net",
-                "Referer": "https://novelai.net/",
-            }, impersonate="chrome110")
-        return self._session
-
-
-class NovelAiInference(BaseModel):
+class GenerateImageInfer(BaseModel):
     _endpoint: Optional[str] = PrivateAttr("https://api.novelai.net")
     _charge: bool = PrivateAttr(False)
 
@@ -174,11 +151,11 @@ class NovelAiInference(BaseModel):
 
     def validate_charge(self):
         if self.parameters.steps > 28 and not self.charge:
-            raise CheckError("steps must be less than 28 for free users.")
+            raise ValueError("steps must be less than 28 for free users.")
         if (self.parameters.width, self.parameters.height) not in self.valid_wh() and not self.charge:
-            raise CheckError("Invalid size, must be one of 832x1216, 1216x832, 1024x1024 for free users.")
+            raise ValueError("Invalid size, must be one of 832x1216, 1216x832, 1024x1024 for free users.")
         if self.parameters.n_samples != 1 and not self.charge:
-            raise CheckError("n_samples must be 1 for free users.")
+            raise ValueError("n_samples must be 1 for free users.")
         return self
 
     @classmethod
@@ -235,60 +212,63 @@ class NovelAiInference(BaseModel):
             parameters=cls.Params(**param)
         )
 
-    async def generate(self, session: Union[AsyncSession, CurlSession]) -> NaiResult:
-        if isinstance(session, CurlSession):
+    async def generate(self, session: Union[AsyncSession, JwtCredential]) -> ImageGenerateResp:
+        if isinstance(session, JwtCredential):
             session = session.session
-        self.validate_charge()
         request_data = self.model_dump()
-        logger.info(f"Request Data: {request_data}")
+        logger.debug(f"Request Data: {request_data}")
         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') in ['binary/octet-stream', 'application/x-zip-compressed']:
-                zip_file = ZipFile(BytesIO(response.content))
-                with zip_file as zf:
-                    file_list = zf.namelist()
-                    if not file_list:
-                        raise ServerError(msg="Nai Returned zip file is empty.")
-                    with zf.open(file_list[0]) as file:
-                        _img = BytesIO(file.read())
-            else:
-                logger.error(f"response: {response.text}")
-                # 解析错误内容
+            if response.headers.get('Content-Type') not in ['binary/octet-stream', 'application/x-zip-compressed']:
                 try:
                     _msg = response.json()
-                    message = _msg["message"]
                 except Exception:
-                    raise ServerError(msg=f"Unexpected content type: {response.headers.get('Content-Type')}")
-                code = _msg.get("statusCode", 200)
-                if message == "Error generating image":
-                    raise ServerError(msg=f"[Nai Server] Generate image error: {code}")
-                raise AuthError(msg=f"[Nai Server]{message}")
-            try:
-                _img_exif = NaiPic.read_from_img(_img)
-            except Exception as e:
-                logger.debug(e)
-                _img_exif = NaiPic.read_from_param(prompt=self.input,
-                                                   neg_prompt=self.parameters.negative_prompt,
-                                                   seed=self.parameters.seed,
-                                                   steps=self.parameters.steps,
-                                                   cfg_rescale=self.parameters.cfg_rescale,
-                                                   sampler=self.parameters.sampler,
-                                                   )
-            _img = _img_exif.generate_to_img_io(img_io=_img)
-            _img.seek(0)
-            return NaiResult(
-                meta=NaiResult.RequestParams(
+                    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)
+                    )
+                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, status_code=status_code, response=message)
+                if status_code in [409]:
+                    # conflict error
+                    raise APIError(message, request=request_data, status_code=status_code, response=message)
+                raise APIError(message, request=request_data, status_code=status_code, response=message)
+            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,
+                        status_code=response.status_code,
+                        response=try_jsonfy(response.content)
+                    )
+                for filename in file_list:
+                    unzip_content.append((filename, zip_file.read(filename)))
+            return ImageGenerateResp(
+                meta=ImageGenerateResp.RequestParams(
                     endpoint=self.base_url,
                     raw_request=request_data,
                 ),
-                files=(f"{str(self.input[:7])}-s{shortuuid.uuid()[:5]}.png", _img.getvalue())
+                files=unzip_content
             )
         except httpx.HTTPError as exc:
             raise RuntimeError(f"An HTTP error occurred: {exc}")
-        except ServerError as e:
+        except APIError as e:
             raise e
         except Exception as e:
-            raise RuntimeError(f"An error occurred: {e}")
+            logger.opt(exception=e).exception("An Unexpected error occurred")
+            raise e

novelai_python-0.1.1/src/novelai_python/utils/__init__.py

@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/26 上午10:51
+# @Author  : sudoskys
+# @File    : __init__.py.py
+# @Software: PyCharm
+import io
+import json
+from typing import Union
+
+from .hash import NovelAiMetadata
+
+
+def png_info_reset(img: Union[str, io.BytesIO]):
+    """
+    reset png info hash value
+    :param img: BytesIO 对象
+    :return: 
+    """
+    _fixed = NovelAiMetadata.rehash(img_io=img)
+    return _fixed
+
+
+def try_jsonfy(obj: Union[str, dict, list, tuple]):
+    """
+    try to jsonfy a object
+    :param obj:
+    :return:
+    """
+    try:
+        return json.loads(obj)
+    except json.JSONDecoder:
+        return f"Decode Error {type(obj)}"

{novelai_python-0.1.0/src/novelai_python → novelai_python-0.1.1/src/novelai_python/utils}/hash.py

@@ -23,28 +23,35 @@ def sign_message(message, key):
     return signed_hash
 
 
-class NaiPic(BaseModel):
+class NovelAiMetadata(BaseModel):
     title: str = "AI generated image"
     description: str
     comment: Union[dict, str] = {}
 
-    def generate_exif(self):
-        self.comment["signed_hash"] = sign_message(json.dumps(self.description), "novalai-bot")
-        self.comment = json.dumps(self.comment)
+    @property
+    def metadata(self):
+        if isinstance(self.comment, dict):
+            _comment = json.dumps(self.comment)
         return {"Title": self.title, "Description": self.description, "Comment": self.comment}
 
-    def generate_to_img_io(self, img_io):
+    @staticmethod
+    def rehash(img_io):
+        cls = NaiPic.read_from_img(img_io)
+        cls.comment["signed_hash"] = sign_message(json.dumps(cls.description), "novalai-client")
+        cls.write_out(img_io=img_io)
+        return img_io
+
+    def write_out(self, img_io: BytesIO):
         with Image.open(img_io) as img:
             metadata = PngInfo()
-            _info = self.generate_exif()
-            for k, v in _info.items():
+            for k, v in self.metadata.items():
                 metadata.add_text(k, v)
             _new_img = BytesIO()
             img.save(_new_img, format="PNG", pnginfo=metadata, quality=95, optimize=False)
         return _new_img
 
     @classmethod
-    def read_from_img(cls, image_io):
+    def build_from_img(cls, image_io):
         with Image.open(image_io) as img:
             title = img.info.get("Title")
             if not title == 'AI generated image':
@@ -59,7 +66,7 @@ class NaiPic(BaseModel):
             return cls(title=title, description=prompt, comment=comment)
 
     @classmethod
-    def read_from_param(cls, prompt, neg_prompt, **kwargs):
+    def build_from_param(cls, prompt, neg_prompt, **kwargs):
         _comment = {"uc": neg_prompt}
         _comment.update(kwargs)
         return cls(title="AI generated image", description=prompt, comment=_comment)

novelai_python-0.1.0/PKG-INFO

@@ -1,10 +0,0 @@
-Metadata-Version: 2.1
-Name: novelai-python
-Version: 0.1.0
-Summary: Novelai Python Binding With Pydantic
-Author-Email: sudoskys <coldlando@hotmail.com>
-License: Apache 2.0
-Requires-Python: >=3.8
-Description-Content-Type: text/markdown
-
-# novelai-python

novelai_python-0.1.0/README.md

@@ -1 +0,0 @@
-# novelai-python

novelai_python-0.1.0/pyproject.toml

@@ -1,22 +0,0 @@
-[project]
-name = "novelai-python"
-version = "0.1.0"
-description = "Novelai Python Binding With Pydantic"
-authors = [
-    { name = "sudoskys", email = "coldlando@hotmail.com" },
-]
-dependencies = []
-requires-python = ">=3.8"
-readme = "README.md"
-
-[project.license]
-text = "Apache 2.0"
-
-[build-system]
-requires = [
-    "pdm-backend",
-]
-build-backend = "pdm.backend"
-
-[tool.pdm]
-distribution = true

novelai_python-0.1.0/src/novelai_python/error.py

@@ -1,29 +0,0 @@
-# -*- coding: utf-8 -*-
-# @Time    : 2023/11/18 上午12:58
-# @Author  : sudoskys
-# @File    : error.py
-# @Software: PyCharm
-
-
-class ServerError(Exception):
-    def __init__(self, msg: str = None):
-        self.msg = msg
-
-    def __str__(self):
-        return f"ServerError: {self.msg}"
-
-
-class ValidationError(Exception):
-    def __init__(self, msg: str = None):
-        self.msg = msg
-
-    def __str__(self):
-        return f"ValidationError: {self.msg}"
-
-
-class AuthError(Exception):
-    def __init__(self, msg: str = None):
-        self.msg = msg
-
-    def __str__(self):
-        return f"AuthError: {self.msg}"

novelai_python-0.1.0/src/novelai_python/schema.py

@@ -1,24 +0,0 @@
-# -*- coding: utf-8 -*-
-# @Time    : 2023/11/18 上午1:10
-# @Author  : sudoskys
-# @File    : schema.py
-# @Software: PyCharm
-from typing import Tuple
-
-from pydantic import BaseModel
-
-
-class NaiResult(BaseModel):
-    class RequestParams(BaseModel):
-        endpoint: str
-        raw_request: dict = None
-
-    meta: RequestParams
-    files: Tuple[str, bytes] = None
-
-    def query_params(self, key: str, default=None):
-        if isinstance(self.meta.raw_request, str):
-            return default
-        if isinstance(self.meta.raw_request.get("parameters"), dict):
-            return self.meta.raw_request.get("parameters").get(key, default)
-        return default