diff --git a/NOTICE.MD b/NOTICE.MD new file mode 100644 index 0000000..5773a72 --- /dev/null +++ b/NOTICE.MD @@ -0,0 +1,7 @@ +# NOTICE OF THIS PROJECT + +## MIT License + +The MIT license applies to the files in: + + file: "src/novelai_python/sdk/ai/generate_image.py" from https://github.com/HanaokaYuzu/NovelAI-API \ No newline at end of file diff --git a/README.md b/README.md index dea5219..28ca11d 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a - [x] /ai/generate-image - [x] /user/subscription +- [x] /user/login - [ ] /ai/generate-image/suggest-tags - [ ] /ai/annotate-image - [ ] /ai/classify @@ -19,6 +20,8 @@ The goal of this repository is to use Pydantic to build legitimate requests to a ### Usage 🖥️ +More examples can be found in the [playground](/playground) directory. + ```python import asyncio import os @@ -26,7 +29,7 @@ import os from dotenv import load_dotenv from pydantic import SecretStr -from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential +from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential, LoginCredential load_dotenv() @@ -34,7 +37,13 @@ enhance = "year 2023,dynamic angle, best quality, amazing quality, very aesthet async def main(): - globe_s = JwtCredential(jwt_token=SecretStr(os.getenv("NOVELAI_JWT"))) + globe_s = JwtCredential( + jwt_token=SecretStr(os.getenv("NOVELAI_JWT")) + ) + globe_s2 = LoginCredential( + username=os.getenv("NOVELAI_USERNAME"), + password=SecretStr(os.getenv("NOVELAI_PASSWORD")) + ) _res = await GenerateImageInfer.build( prompt=f"1girl,{enhance}").generate( session=globe_s) @@ -69,5 +78,9 @@ python3 -m novelai_python.server -h '0.0.0.0' -p 7888 ## Acknowledgements 🙏 [BackEnd](https://api.novelai.net/docs) + [novelai-api](https://github.com/Aedial/novelai-api) +[NovelAI-API](https://github.com/HanaokaYuzu/NovelAI-API) + + diff --git a/pdm.lock b/pdm.lock index aba63e3..332b172 100644 --- a/pdm.lock +++ b/pdm.lock @@ -5,7 +5,7 @@ groups = ["default", "dev", "testing"] strategy = ["cross_platform", "inherit_metadata"] lock_version = "4.4.1" -content_hash = "sha256:4380c0e0d145f0e2709a296490964bb1abbd62c88a7f64ebb774dd998db100d0" +content_hash = "sha256:9ce409e6d4b57704c11bf0e10f32df5c03636bff59704ef07a8bd82265668c34" [[package]] name = "annotated-types" @@ -38,6 +38,53 @@ files = [ {file = "anyio-4.2.0.tar.gz", hash = "sha256:e1875bb4b4e2de1669f4bc7869b6d3f54231cdced71605e6e64c9be77e3be50f"}, ] +[[package]] +name = "argon2-cffi" +version = "23.1.0" +requires_python = ">=3.7" +summary = "Argon2 for Python" +groups = ["default"] +dependencies = [ + "argon2-cffi-bindings", +] +files = [ + {file = "argon2_cffi-23.1.0-py3-none-any.whl", hash = "sha256:c670642b78ba29641818ab2e68bd4e6a78ba53b7eff7b4c3815ae16abf91c7ea"}, + {file = "argon2_cffi-23.1.0.tar.gz", hash = "sha256:879c3e79a2729ce768ebb7d36d4609e3a78a4ca2ec3a9f12286ca057e3d0db08"}, +] + +[[package]] +name = "argon2-cffi-bindings" +version = "21.2.0" +requires_python = ">=3.6" +summary = "Low-level CFFI bindings for Argon2" +groups = ["default"] +dependencies = [ + "cffi>=1.0.1", +] +files = [ + {file = "argon2-cffi-bindings-21.2.0.tar.gz", hash = "sha256:bb89ceffa6c791807d1305ceb77dbfacc5aa499891d2c55661c6459651fc39e3"}, + {file = "argon2_cffi_bindings-21.2.0-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ccb949252cb2ab3a08c02024acb77cfb179492d5701c7cbdbfd776124d4d2367"}, + {file = "argon2_cffi_bindings-21.2.0-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9524464572e12979364b7d600abf96181d3541da11e23ddf565a32e70bd4dc0d"}, + {file = "argon2_cffi_bindings-21.2.0-cp36-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b746dba803a79238e925d9046a63aa26bf86ab2a2fe74ce6b009a1c3f5c8f2ae"}, + {file = "argon2_cffi_bindings-21.2.0-cp36-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:58ed19212051f49a523abb1dbe954337dc82d947fb6e5a0da60f7c8471a8476c"}, + {file = "argon2_cffi_bindings-21.2.0-cp36-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:bd46088725ef7f58b5a1ef7ca06647ebaf0eb4baff7d1d0d177c6cc8744abd86"}, + {file = "argon2_cffi_bindings-21.2.0-cp36-abi3-musllinux_1_1_i686.whl", hash = "sha256:8cd69c07dd875537a824deec19f978e0f2078fdda07fd5c42ac29668dda5f40f"}, + {file = "argon2_cffi_bindings-21.2.0-cp36-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:f1152ac548bd5b8bcecfb0b0371f082037e47128653df2e8ba6e914d384f3c3e"}, + {file = "argon2_cffi_bindings-21.2.0-cp36-abi3-win32.whl", hash = "sha256:603ca0aba86b1349b147cab91ae970c63118a0f30444d4bc80355937c950c082"}, + {file = "argon2_cffi_bindings-21.2.0-cp36-abi3-win_amd64.whl", hash = "sha256:b2ef1c30440dbbcba7a5dc3e319408b59676e2e039e2ae11a8775ecf482b192f"}, + {file = "argon2_cffi_bindings-21.2.0-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:e415e3f62c8d124ee16018e491a009937f8cf7ebf5eb430ffc5de21b900dad93"}, + {file = "argon2_cffi_bindings-21.2.0-pp37-pypy37_pp73-macosx_10_9_x86_64.whl", hash = "sha256:3e385d1c39c520c08b53d63300c3ecc28622f076f4c2b0e6d7e796e9f6502194"}, + {file = "argon2_cffi_bindings-21.2.0-pp37-pypy37_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2c3e3cc67fdb7d82c4718f19b4e7a87123caf8a93fde7e23cf66ac0337d3cb3f"}, + {file = "argon2_cffi_bindings-21.2.0-pp37-pypy37_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6a22ad9800121b71099d0fb0a65323810a15f2e292f2ba450810a7316e128ee5"}, + {file = "argon2_cffi_bindings-21.2.0-pp37-pypy37_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f9f8b450ed0547e3d473fdc8612083fd08dd2120d6ac8f73828df9b7d45bb351"}, + {file = "argon2_cffi_bindings-21.2.0-pp37-pypy37_pp73-win_amd64.whl", hash = "sha256:93f9bf70084f97245ba10ee36575f0c3f1e7d7724d67d8e5b08e61787c320ed7"}, + {file = "argon2_cffi_bindings-21.2.0-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:3b9ef65804859d335dc6b31582cad2c5166f0c3e7975f324d9ffaa34ee7e6583"}, + {file = "argon2_cffi_bindings-21.2.0-pp38-pypy38_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d4966ef5848d820776f5f562a7d45fdd70c2f330c961d0d745b784034bd9f48d"}, + {file = "argon2_cffi_bindings-21.2.0-pp38-pypy38_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:20ef543a89dee4db46a1a6e206cd015360e5a75822f76df533845c3cbaf72670"}, + {file = "argon2_cffi_bindings-21.2.0-pp38-pypy38_pp73-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ed2937d286e2ad0cc79a7087d3c272832865f779430e0cc2b4f3718d3159b0cb"}, + {file = "argon2_cffi_bindings-21.2.0-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:5e00316dabdaea0b2dd82d141cc66889ced0cdcbfa599e8b471cf22c620c329a"}, +] + [[package]] name = "certifi" version = "2023.11.17" diff --git a/playground/generate_image.py b/playground/generate_image.py index 67748dd..d6cf6c6 100644 --- a/playground/generate_image.py +++ b/playground/generate_image.py @@ -9,7 +9,7 @@ from dotenv import load_dotenv from pydantic import SecretStr -from novelai_python import APIError +from novelai_python import APIError, Login from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential load_dotenv() @@ -21,10 +21,15 @@ async def main(): globe_s = JwtCredential(jwt_token=SecretStr(jwt)) + _res = await Login.build(user_name=os.getenv("NOVELAI_USER"), password=os.getenv("NOVELAI_PASS") + ).request() try: - _res = await GenerateImageInfer.build( - prompt=f"1girl, winter, jacket, sfw, angel, flower,{enhance}").generate( - session=globe_s) + gen = GenerateImageInfer.build( + prompt=f"1girl, winter, jacket, sfw, angel, flower,{enhance}") + print(f"charge: {gen.calculate_cost(is_opus=True)}") + _res = await gen.generate( + session=globe_s, remove_sign=True + ) except APIError as e: print(e.response) return diff --git a/playground/login.py b/playground/login.py new file mode 100644 index 0000000..36281aa --- /dev/null +++ b/playground/login.py @@ -0,0 +1,33 @@ +# -*- coding: utf-8 -*- +# @Time : 2024/2/7 下午12:07 +# @Author : sudoskys +# @File : login.py +# @Software: PyCharm +import asyncio +import os + +from dotenv import load_dotenv +from loguru import logger + +from novelai_python import APIError +from novelai_python import Login, LoginResp + +load_dotenv() + + +async def main(): + try: + _res = await Login.build(user_name=os.getenv("NOVELAI_USER"), password=os.getenv("NOVELAI_PASS") + ).request() + except APIError as e: + logger.exception(e) + print(e.__dict__) + return + + _res: LoginResp + print(_res) + print(f"Access Token: {_res.accessToken}") + + +loop = asyncio.get_event_loop() +loop.run_until_complete(main()) diff --git a/playground/subscription.py b/playground/subscription.py index 6ed7c74..d2b6fe6 100644 --- a/playground/subscription.py +++ b/playground/subscription.py @@ -10,7 +10,7 @@ from loguru import logger from pydantic import SecretStr -from novelai_python import APIError, SubscriptionResp +from novelai_python import APIError, SubscriptionResp, LoginCredential from novelai_python import Subscription, JwtCredential load_dotenv() @@ -26,15 +26,28 @@ async def main(): _res = await Subscription().request( session=globe_s ) + _res: SubscriptionResp + print(f"JwtCredential/Subscription: {_res}") + print(_res.is_active) + print(_res.anlas_left) except APIError as e: logger.exception(e) - print(e.response) + print(e.__dict__) return - _res: SubscriptionResp - print(_res) - print(_res.is_active) - print(_res.anlas_left) + try: + cre = LoginCredential( + username=os.getenv("NOVELAI_USER"), + password=SecretStr(os.getenv("NOVELAI_PASS")) + ) + _res = await Subscription().request( + session=cre + ) + print(f"LoginCredential/User subscription: {_res}") + print(_res.is_active) + print(_res.anlas_left) + except Exception as e: + logger.exception(e) loop = asyncio.get_event_loop() diff --git a/pyproject.toml b/pyproject.toml index ef9be76..8be76a6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "novelai-python" -version = "0.1.8" +version = "0.1.9" description = "Novelai Python Binding With Pydantic" authors = [ { name = "sudoskys", email = "coldlando@hotmail.com" }, @@ -17,6 +17,7 @@ dependencies = [ "fastapi>=0.109.0", "uvicorn[standard]>=0.27.0.post1", "numpy>=1.24.4", + "argon2-cffi>=23.1.0", ] requires-python = ">=3.8" readme = "README.md" diff --git a/src/novelai_python/__init__.py b/src/novelai_python/__init__.py index 8d2521d..d7b6015 100644 --- a/src/novelai_python/__init__.py +++ b/src/novelai_python/__init__.py @@ -5,13 +5,13 @@ # @Software: PyCharm from ._exceptions import ( - APIError, AuthError, - NovelAiError, -) # noqa: F401, F403 -from .credential import JwtCredential # noqa: F401, F403 -from .sdk.ai import GenerateImageInfer, ImageGenerateResp # noqa: F401, F403 -from .sdk.user import Subscription, SubscriptionResp # noqa: F401, F403 + NovelAiError, APIError, +) +from .credential import JwtCredential, LoginCredential +from .sdk import GenerateImageInfer, ImageGenerateResp +from .sdk import Login, LoginResp +from .sdk import Subscription, SubscriptionResp __all__ = [ "GenerateImageInfer", @@ -20,7 +20,11 @@ "Subscription", "SubscriptionResp", + "Login", + "LoginResp", + "JwtCredential", + "LoginCredential", "APIError", "AuthError", diff --git a/src/novelai_python/_exceptions.py b/src/novelai_python/_exceptions.py index cfd73de..f12010b 100644 --- a/src/novelai_python/_exceptions.py +++ b/src/novelai_python/_exceptions.py @@ -32,6 +32,15 @@ def __init__(self, message: str, request: dict, response: Union[dict, str], stat self.code = status_code self.response = response + @property + def __dict__(self): + return { + "message": self.message, + "request": self.request, + "response": self.response, + "code": self.code + } + class AuthError(APIError): """ diff --git a/src/novelai_python/_response/__init__.py b/src/novelai_python/_response/__init__.py index 7f694ba..034d591 100644 --- a/src/novelai_python/_response/__init__.py +++ b/src/novelai_python/_response/__init__.py @@ -4,8 +4,11 @@ # @File : __init__.py.py # @Software: PyCharm from .ai.generate_image import ImageGenerateResp +from .user.login import LoginResp from .user.subscription import SubscriptionResp + __all__ = [ "ImageGenerateResp", "SubscriptionResp", + "LoginResp" ] diff --git a/src/novelai_python/_response/user/login.py b/src/novelai_python/_response/user/login.py new file mode 100644 index 0000000..7aa2622 --- /dev/null +++ b/src/novelai_python/_response/user/login.py @@ -0,0 +1,10 @@ +# -*- coding: utf-8 -*- +# @Time : 2024/2/7 上午11:57 +# @Author : sudoskys +# @File : login.py +# @Software: PyCharm +from pydantic import BaseModel + + +class LoginResp(BaseModel): + accessToken: str diff --git a/src/novelai_python/credential/JwtToken.py b/src/novelai_python/credential/JwtToken.py index 397a779..7ad7615 100644 --- a/src/novelai_python/credential/JwtToken.py +++ b/src/novelai_python/credential/JwtToken.py @@ -4,18 +4,19 @@ # @File : JwtToken.py # @Software: PyCharm from curl_cffi.requests import AsyncSession -from pydantic import BaseModel, SecretStr, Field +from pydantic import SecretStr, Field +from ._base import CredentialBase -class JwtCredential(BaseModel): + +class JwtCredential(CredentialBase): """ 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): + async def get_session(self, timeout: int = 180): if not self._session: self._session = AsyncSession(timeout=timeout, headers={ "Authorization": f"Bearer {self.jwt_token.get_secret_value()}", diff --git a/src/novelai_python/credential/UserAuth.py b/src/novelai_python/credential/UserAuth.py new file mode 100644 index 0000000..187961f --- /dev/null +++ b/src/novelai_python/credential/UserAuth.py @@ -0,0 +1,36 @@ +# -*- coding: utf-8 -*- +# @Time : 2024/2/7 下午12:14 +# @Author : sudoskys +# @File : UserAuth.py +# @Software: PyCharm +import time +from typing import Optional + +from curl_cffi.requests import AsyncSession +from pydantic import SecretStr, Field + +from ._base import CredentialBase + + +class LoginCredential(CredentialBase): + """ + JwtCredential is the base class for all credential. + """ + username: str = Field(None, description="username") + password: SecretStr = Field(None, description="password") + _session: Optional[AsyncSession] = None + _update_at: Optional[int] = None + + async def get_session(self, timeout: int = 180): + # 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={ + "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()) + return self._session diff --git a/src/novelai_python/credential/__init__.py b/src/novelai_python/credential/__init__.py index ffe5bed..2603c82 100644 --- a/src/novelai_python/credential/__init__.py +++ b/src/novelai_python/credential/__init__.py @@ -3,9 +3,15 @@ # @Author : sudoskys # @File : __init__.py.py # @Software: PyCharm -from .JwtToken import JwtCredential from pydantic import SecretStr + +from .JwtToken import JwtCredential +from .UserAuth import LoginCredential +from ._base import CredentialBase + __all__ = [ "JwtCredential", + "LoginCredential", + "CredentialBase", "SecretStr" ] diff --git a/src/novelai_python/credential/_base.py b/src/novelai_python/credential/_base.py new file mode 100644 index 0000000..6e8840f --- /dev/null +++ b/src/novelai_python/credential/_base.py @@ -0,0 +1,14 @@ +# -*- coding: utf-8 -*- +# @Time : 2024/2/7 下午12:14 +# @Author : sudoskys +# @File : _shema.py +# @Software: PyCharm +from curl_cffi.requests import AsyncSession +from pydantic import BaseModel + + +class CredentialBase(BaseModel): + _session: AsyncSession = None + + async def get_session(self, timeout: int = 180): + raise NotImplementedError diff --git a/src/novelai_python/sdk/__init__.py b/src/novelai_python/sdk/__init__.py index bdff43d..3868355 100644 --- a/src/novelai_python/sdk/__init__.py +++ b/src/novelai_python/sdk/__init__.py @@ -5,4 +5,5 @@ # @Software: PyCharm from .ai.generate_image import GenerateImageInfer, ImageGenerateResp # noqa 401 +from .user.login import Login, LoginResp # noqa 401 from .user.subscription import Subscription, SubscriptionResp # noqa 401 diff --git a/src/novelai_python/sdk/ai/LICENSE b/src/novelai_python/sdk/ai/LICENSE new file mode 100644 index 0000000..01eab1e --- /dev/null +++ b/src/novelai_python/sdk/ai/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 GM Development Department + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/src/novelai_python/sdk/ai/generate_image.py b/src/novelai_python/sdk/ai/generate_image.py index 74cf696..e888635 100644 --- a/src/novelai_python/sdk/ai/generate_image.py +++ b/src/novelai_python/sdk/ai/generate_image.py @@ -4,21 +4,22 @@ # @File : generate_image.py # @Software: PyCharm import json +import math import random from io import BytesIO from typing import Optional, Union, Literal from zipfile import ZipFile import httpx -from curl_cffi.requests import AsyncSession +from curl_cffi.requests import AsyncSession, RequestsError from loguru import logger -from novelai_python.utils import NovelAiMetadata -from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, model_validator, Field +from typing_extensions import override from ..._exceptions import APIError, AuthError from ..._response import ImageGenerateResp -from ...credential import JwtCredential -from ...utils import try_jsonfy +from ...utils import try_jsonfy, NovelAiMetadata +from ...credential import CredentialBase class GenerateImageInfer(BaseModel): @@ -26,45 +27,54 @@ class GenerateImageInfer(BaseModel): _charge: bool = PrivateAttr(False) class Params(BaseModel): + # Inpaint add_original_image: Optional[bool] = False - cfg_rescale: Optional[int] = 0 - controlnet_strength: Optional[int] = 1 + mask: Optional[str] = None + + cfg_rescale: Optional[float] = Field(0, ge=0, le=1, multiple_of=0.02) + controlnet_strength: Optional[float] = Field(1.0, ge=0.1, le=2, multiple_of=0.1) dynamic_thresholding: Optional[bool] = False - height: Optional[int] = 1216 + height: Optional[int] = Field(1216, ge=64, le=49152) + # Img2Img image: Optional[str] = None # img2img,base64 - legacy: Optional[bool] = False - legacy_v3_extend: Optional[bool] = False - n_samples: Optional[int] = 1 - negative_prompt: Optional[str] = ( - "nsfw, lowres, {bad}, error, fewer, extra, missing, worst quality, jpeg artifacts, bad quality, " - "watermark, unfinished, displeasing, chromatic aberration, signature, extra digits, artistic error, " - "username, scan, [abstract], bad anatomy, bad hands, @_@, mismatched pupils, heart-shaped pupils, " - "glowing eyes, nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, " - "extra digit, fewer digits, cropped, worst quality, low quality, normal quality, " - "jpeg artifacts, signature, watermark, username, blurry" - ) + strength: Optional[float] = Field(default=0.3, 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 + + n_samples: Optional[int] = Field(1, ge=1, le=8) + negative_prompt: Optional[str] = '' noise_schedule: Optional[Union[str, Literal['native', 'polyexponential', 'exponential']]] = "native" + + # Misc params_version: Optional[int] = 1 + legacy: Optional[bool] = False + legacy_v3_extend: Optional[bool] = False + qualityToggle: Optional[bool] = True sampler: Optional[str] = "k_euler" - scale: Optional[int] = 5 - seed: Optional[int] = -1 + scale: Optional[float] = Field(6.0, ge=0, le=10, multiple_of=0.1) + # Seed + seed: Optional[int] = Field( + default_factory=lambda: random.randint(0, 4294967295 - 7), + gt=0, + le=4294967295 - 7, + ) + extra_noise_seed: Optional[int] = Field( + default_factory=lambda: random.randint(0, 4294967295 - 7), + gt=0, + le=4294967295 - 7, + ) + sm: Optional[bool] = False sm_dyn: Optional[bool] = False - steps: Optional[int] = 28 - ucPreset: Optional[int] = 0 - uncond_scale: Optional[int] = 1 - width: Optional[int] = 832 - - @field_validator('seed') - def seed_validator(cls, v: int): - if v == -1: - v = random.randint(0, 2 ** 32 - 1) - return v + steps: Optional[int] = Field(28, ge=1, le=50) + ucPreset: Optional[Literal[0, 1, 2, 3]] = 0 + uncond_scale: Optional[float] = Field(1.0, ge=0, le=1.5, multiple_of=0.05) + width: Optional[int] = Field(832, ge=64, le=49152) @model_validator(mode="after") def validate_img2img(self): - image = True if self.image else False add_origin = True if self.add_original_image else False if image != add_origin: @@ -112,35 +122,46 @@ def height_validator(cls, v: int): raise ValueError("Invalid height, must be multiple of 64.") return v - @field_validator('n_samples') - def n_samples_validator(cls, v: int): - """ - 小于 8 - :param v: - :return: - """ - if v > 8: - raise ValueError("Invalid n_samples, must be less than 8.") - return v - - action: Union[str, Literal["generate", "img2img"]] = "generate" + action: Union[str, Literal["generate", "img2img", "infill"]] = "generate" input: str = "1girl, best quality, amazing quality, very aesthetic, absurdres" model: Optional[str] = "nai-diffusion-3" parameters: Params = Params() model_config = ConfigDict(extra="ignore") + @override + def model_post_init(self, *args) -> None: + """ + Post-initialization hook. + :return: + """ + if self.parameters.ucPreset == 0: + # 0: 重型 + self.parameters.negative_prompt += str( + ", lowres, {bad}, error, fewer, extra, missing, worst quality, jpeg artifacts, bad quality, " + "watermark, unfinished, displeasing, chromatic aberration, signature, extra digits, artistic error, " + "username, scan, [abstract], bad anatomy, bad hands, @_@, mismatched pupils, heart-shaped pupils, " + "glowing eyes, nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, " + "extra digit, fewer digits, cropped, worst quality, low quality, normal quality, " + "jpeg artifacts, signature, watermark, username, blurry" + ) + elif self.parameters.ucPreset == 1: + # 1: 轻型 + self.parameters.negative_prompt += (", lowres, jpeg artifacts, worst quality" + ", watermark, blurry, very displeasing") + elif self.parameters.ucPreset == 2: + # 2: 人物 + self.parameters.negative_prompt += (", lowres, {bad}, error, fewer, extra, missing, worst quality, " + "jpeg artifacts, bad quality, watermark, unfinished, displeasing, " + "chromatic aberration, signature, extra digits, artistic error, " + "username, scan, [abstract], bad anatomy, bad hands, @_@, mismatched " + "pupils, heart-shaped pupils, glowing eyes") + if self.parameters.qualityToggle: + self.input += ", best quality, amazing quality, very aesthetic, absurdres" + @property def base_url(self): return f"{self.endpoint.strip('/')}/ai/generate-image" - @property - def charge(self): - return self._charge - - @charge.setter - def charge(self, value): - self._charge = value - @property def endpoint(self): return self._endpoint @@ -161,14 +182,46 @@ def valid_wh(): (1024, 1024), ] - def validate_charge(self): - if self.parameters.steps > 28 and not self.charge: - 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 ValueError("Invalid size, must be one of 832x1216, 1216x832, 1024x1024 for free users.") - if self.parameters.n_samples != 1 and not self.charge: - raise ValueError("n_samples must be 1 for free users.") - return self + def calculate_cost(self, is_opus: bool = False): + """ + Calculate the Anlas cost of current parameters. + + Parameters + ---------- + is_opus: `bool`, optional + If the subscription tier is Opus. Opus accounts have access to free generations. + """ + + steps: int = self.parameters.steps + n_samples: int = self.parameters.n_samples + uncond_scale: float = self.parameters.uncond_scale + strength: float = self.action == "img2img" and self.parameters.strength or 1.0 + smea_factor = self.parameters.sm_dyn and 1.4 or self.parameters.sm and 1.2 or 1.0 + resolution = max(self.parameters.width * self.parameters.height, 65536) + + # For normal resolutions, squre is adjusted to the same price as portrait/landscape + if math.prod( + (832, 1216) + ) < resolution <= math.prod((1024, 1024)): + resolution = math.prod((832, 1216)) + per_sample = ( + math.ceil( + 2951823174884865e-21 * resolution + + 5.753298233447344e-7 * resolution * steps + ) + * smea_factor + ) + per_sample = max(math.ceil(per_sample * strength), 2) + + if uncond_scale != 1.0: + per_sample = math.ceil(per_sample * 1.3) + + opus_discount = ( + is_opus + and steps <= 28 + and (resolution <= math.prod((1024, 1024))) + ) + return per_sample * (n_samples - int(opus_discount)) @classmethod def build(cls, @@ -176,42 +229,34 @@ def build(cls, *, model: str = "nai-diffusion-3", action: Literal['generate', 'img2img'] = 'generate', - negative_prompt: Optional[str] = None, - override_negative_prompt: bool = False, - seed: int = -1, + negative_prompt: str = "", + seed: int = None, steps: int = 28, cfg_rescale: int = 0, sampler: str = "k_euler", width: int = 832, height: int = 1216, + qualityToggle: bool = True, + ucPreset: int = 0, ): """ 正负面, step, cfg, 采样方式, seed - :param override_negative_prompt: - :param prompt: - :param model: - :param action: Mode for img generate - :param negative_prompt: - :param seed: - :param steps: - :param cfg_rescale: - :param sampler: - :param width: - :param height: + :param ucPreset: 0: 重型, 1: 轻型, 2: 人物 + :param qualityToggle: 是否开启质量 + :param prompt: 输入 + :param model: 模型 + :param action: Mode for img generate [generate, img2img, infill] + :param negative_prompt: 负面 + :param seed: 随机种子 + :param steps: 步数 + :param cfg_rescale: 0-1 + :param sampler: 采样方式 + :param width: 宽 + :param height: 高 :return: self """ assert isinstance(prompt, str) - _negative_prompt = ("nsfw, lowres, {bad}, error, fewer, extra, missing, worst quality, jpeg artifacts, " - "bad quality, watermark, unfinished, displeasing, chromatic aberration, signature, " - "extra digits, artistic error, username, scan, [abstract], bad anatomy, bad hands, " - "@_@, mismatched pupils, heart-shaped pupils, glowing eyes, nsfw, lowres, bad anatomy, " - "bad hands, text, error,missing fingers, extra digit, fewer digits, cropped," - "worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, " - "username, blurry") - if override_negative_prompt: - _negative_prompt = negative_prompt - else: - _negative_prompt = f"{_negative_prompt}, {negative_prompt}" + _negative_prompt = negative_prompt param = { "negative_prompt": _negative_prompt, "seed": seed, @@ -220,6 +265,8 @@ def build(cls, "sampler": sampler, "width": width, "height": height, + "qualityToggle": qualityToggle, + "ucPreset": ucPreset, } # 清理空值 param = {k: v for k, v in param.items() if v is not None} @@ -230,7 +277,7 @@ def build(cls, parameters=cls.Params(**param) ) - async def generate(self, session: Union[AsyncSession, JwtCredential], + async def generate(self, session: Union[AsyncSession, "CredentialBase"], *, remove_sign: bool = False) -> ImageGenerateResp: """ @@ -239,8 +286,8 @@ async def generate(self, session: Union[AsyncSession, JwtCredential], :param remove_sign: 移除追踪信息 :return: """ - if isinstance(session, JwtCredential): - session = session.session + if isinstance(session, CredentialBase): + session = await session.get_session() request_data = self.model_dump(exclude_none=True) logger.debug(f"Request Data: {request_data}") try: @@ -296,6 +343,9 @@ async def generate(self, session: Union[AsyncSession, JwtCredential], ), files=unzip_content ) + except RequestsError as exc: + logger.exception(exc) + raise RuntimeError(f"An AsyncSession error occurred: {exc}") except httpx.HTTPError as exc: raise RuntimeError(f"An HTTP error occurred: {exc}") except APIError as e: diff --git a/src/novelai_python/sdk/user/login.py b/src/novelai_python/sdk/user/login.py new file mode 100644 index 0000000..c1c128b --- /dev/null +++ b/src/novelai_python/sdk/user/login.py @@ -0,0 +1,98 @@ +# -*- coding: utf-8 -*- +# @Time : 2024/2/7 上午11:46 +# @Author : sudoskys +# @File : login.py +# @Software: PyCharm +import json +from typing import Optional + +import httpx +from curl_cffi.requests import AsyncSession, RequestsError +from loguru import logger +from pydantic import BaseModel, PrivateAttr, Field + +from ..._exceptions import APIError +from ..._response.user.login import LoginResp +from ...utils import try_jsonfy, encode_access_key + + +class Login(BaseModel): + _endpoint: Optional[str] = PrivateAttr("https://api.novelai.net") + key: str = Field(..., description="User's key") + + @property + def endpoint(self): + return self._endpoint + + @endpoint.setter + def endpoint(self, value): + self._endpoint = value + + @property + def base_url(self): + return f"{self.endpoint.strip('/')}/user/login" + + @property + def session(self): + return AsyncSession(timeout=180, headers={ + "Content-Type": "application/json", + "Origin": "https://novelai.net", + "Referer": "https://novelai.net/", + }, impersonate="chrome110") + + @classmethod + def build(cls, *, user_name: str, password: str): + """ + From username and password to build a Login instance + :param user_name: + :param password: + :return: + """ + return cls(key=encode_access_key(user_name, password)) + + async def request(self, + ) -> LoginResp: + """ + Request to get user access token + :return: + """ + request_data = self.model_dump(exclude_none=True) + logger.debug("Login") + try: + assert hasattr(self.session, "post"), "session must have get method." + response = await self.session.post( + self.base_url, + data=json.dumps(request_data).encode("utf-8") + ) + if "application/json" not in response.headers.get('Content-Type') or response.status_code != 201: + logger.error(f"Unexpected content type: {response.headers.get('Content-Type')}") + try: + _msg = response.json() + except Exception: + raise APIError( + message=f"Unexpected content type: {response.headers.get('Content-Type')}", + request=request_data, + status_code=response.status_code, + response=try_jsonfy(response.content) + ) + status_code = _msg.get("statusCode", response.status_code) + message = _msg.get("message", "Unknown error") + if status_code in [400, 401]: + # 400 : A validation error occured. + # 401 : Access Key is incorrect. + raise APIError(message, request=request_data, status_code=status_code, response=_msg) + if status_code in [500]: + # An unknown error occured. + raise APIError(message, request=request_data, status_code=status_code, response=_msg) + raise APIError(message, request=request_data, status_code=status_code, response=_msg) + return LoginResp.model_validate(response.json()) + except RequestsError as exc: + logger.exception(exc) + raise RuntimeError(f"An AsyncSession error occurred: {exc}") + except httpx.HTTPError as exc: + raise RuntimeError(f"An HTTP error occurred: {exc}") + except APIError as e: + raise e + except Exception as e: + logger.opt(exception=e).exception("An Unexpected error occurred") + raise e diff --git a/src/novelai_python/sdk/user/subscription.py b/src/novelai_python/sdk/user/subscription.py index 4dc311c..5545be4 100644 --- a/src/novelai_python/sdk/user/subscription.py +++ b/src/novelai_python/sdk/user/subscription.py @@ -3,16 +3,16 @@ # @Author : sudoskys # @File : subscription.py.py # @Software: PyCharm -from typing import Optional, Union +from typing import Optional, Union, Type import httpx -from curl_cffi.requests import AsyncSession +from curl_cffi.requests import AsyncSession, RequestsError from loguru import logger from pydantic import BaseModel, PrivateAttr -from ... import APIError, AuthError +from ..._exceptions import APIError, AuthError from ..._response.user.subscription import SubscriptionResp -from ...credential import JwtCredential +from ...credential import CredentialBase from ...utils import try_jsonfy @@ -32,15 +32,15 @@ def endpoint(self, value): self._endpoint = value async def request(self, - session: Union[AsyncSession, JwtCredential], + session: Union[AsyncSession, CredentialBase], ) -> SubscriptionResp: """ Request to get user subscription information :param session: :return: """ - if isinstance(session, JwtCredential): - session = session.session + if isinstance(session, CredentialBase): + session = await session.get_session() request_data = {} logger.debug("Subscription") try: @@ -48,7 +48,7 @@ async def request(self, response = await session.get( self.base_url, ) - if "application/json" not in response.headers.get('Content-Type'): + if "application/json" not in response.headers.get('Content-Type') or response.status_code != 200: logger.error(f"Unexpected content type: {response.headers.get('Content-Type')}") try: _msg = response.json() @@ -72,6 +72,9 @@ async def request(self, 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: + logger.exception(exc) + raise RuntimeError(f"An AsyncSession error occurred: {exc}") except httpx.HTTPError as exc: raise RuntimeError(f"An HTTP error occurred: {exc}") except APIError as e: diff --git a/src/novelai_python/server.py b/src/novelai_python/server.py index 3566191..a3134fe 100644 --- a/src/novelai_python/server.py +++ b/src/novelai_python/server.py @@ -15,6 +15,7 @@ from .credential import JwtCredential, SecretStr from .sdk.ai.generate_image import GenerateImageInfer +from .sdk.user.login import Login from .sdk.user.subscription import Subscription app = FastAPI() @@ -37,14 +38,30 @@ async def health(): return {"status": "ok"} +@app.post("/user/login") +async def login( + req: Login +): + """ + 用户登录 + :param req: Login + :return: + """ + try: + _result = await req.request() + return _result.model_dump() + except Exception as e: + logger.exception(e) + return JSONResponse(status_code=500, content=e.__dict__) + + @app.get("/user/subscription") async def subscription( current_token: str = Depends(get_current_token) ): """ - 订阅 + 订阅信息 :param current_token: Authorization - :param req: Subscription :return: """ try: @@ -67,7 +84,7 @@ async def generate_image( :return: """ try: - _result = await req.generate(session=get_session(current_token)) + _result = await req.generate(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: diff --git a/src/novelai_python/utils/__init__.py b/src/novelai_python/utils/__init__.py index e58bffc..b55b5b8 100644 --- a/src/novelai_python/utils/__init__.py +++ b/src/novelai_python/utils/__init__.py @@ -9,6 +9,7 @@ from loguru import logger +from .encode import encode_access_key # noqa 401 from .hash import NovelAiMetadata diff --git a/src/novelai_python/utils/encode.py b/src/novelai_python/utils/encode.py new file mode 100644 index 0000000..ca8ad04 --- /dev/null +++ b/src/novelai_python/utils/encode.py @@ -0,0 +1,37 @@ +# -*- coding: utf-8 -*- +# @Time : 2024/2/7 上午11:41 +# @Author : sudoskys +# @File : encode.py +# @Software: PyCharm +from base64 import urlsafe_b64encode +from hashlib import blake2b + +import argon2 + + +# https://github.com/HanaokaYuzu/NovelAI-API/blob/master/src/novelai/utils.py#L12 +def encode_access_key(username: str, password: str) -> str: + """ + Generate hashed access key from the user's username and password using the blake2 and argon2 algorithms. + :param username: str (plaintext) + :param password: str (plaintext) + :return: str + """ + pre_salt = f"{password[:6]}{username}novelai_data_access_key" + + blake = blake2b(digest_size=16) + blake.update(pre_salt.encode()) + salt = blake.digest() + + raw = argon2.low_level.hash_secret_raw( + secret=password.encode(encoding="utf-8"), + salt=salt, + time_cost=2, + memory_cost=int(2000000 / 1024), + parallelism=1, + hash_len=64, + type=argon2.low_level.Type.ID, + ) + hashed = urlsafe_b64encode(raw).decode() + + return hashed[:64]