diff --git a/README.md b/README.md index 1590a94..62dc09a 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,8 @@ The goal of this repository is to use Pydantic to build legitimate requests to a ### Roadmap 🚧 +- [x] utils.NovelAiMetadata +- [x] utils.random_prompt - [x] /ai/generate-image - [x] /user/subscription - [x] /user/login diff --git a/pdm.lock b/pdm.lock index d699a32..33edffc 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:9aa6958e83642a27ca1aae9b488f8f01f014246d5f42b49f2cf7a04fc6328cb1" +content_hash = "sha256:90a2e0d19ce5af6c64699269d1f82ca785c5ef752121c570efdee808a295bb3d" [[package]] name = "annotated-types" @@ -296,6 +296,19 @@ files = [ {file = "exceptiongroup-1.2.0.tar.gz", hash = "sha256:91f5c769735f051a4290d52edd0858999b57e5876e9f85937691bd4c9fa3ed68"}, ] +[[package]] +name = "fake-useragent" +version = "1.4.0" +summary = "Up-to-date simple useragent faker with real world database" +groups = ["default"] +dependencies = [ + "importlib-resources>=5.0; python_version < \"3.10\"", +] +files = [ + {file = "fake-useragent-1.4.0.tar.gz", hash = "sha256:5426e4015d8ccc5bb25f64d3dfcfd3915eba30ffebd31b86b60dc7a4c5d65528"}, + {file = "fake_useragent-1.4.0-py3-none-any.whl", hash = "sha256:9acce439ee2c6cf9c3772fa6c200f62dc8d56605063327a4d8c5d0e47f414b85"}, +] + [[package]] name = "fastapi" version = "0.109.0" @@ -434,6 +447,21 @@ files = [ {file = "idna-3.6.tar.gz", hash = "sha256:9ecdbbd083b06798ae1e86adcbfe8ab1479cf864e4ee30fe4e46a003d12491ca"}, ] +[[package]] +name = "importlib-resources" +version = "6.1.1" +requires_python = ">=3.8" +summary = "Read resources from Python packages" +groups = ["default"] +marker = "python_version < \"3.10\"" +dependencies = [ + "zipp>=3.1.0; python_version < \"3.10\"", +] +files = [ + {file = "importlib_resources-6.1.1-py3-none-any.whl", hash = "sha256:e8bf90d8213b486f428c9c39714b920041cb02c184686a3dee24905aaa8105d6"}, + {file = "importlib_resources-6.1.1.tar.gz", hash = "sha256:3893a00122eafde6894c59914446a512f728a0c1a45f9bb9b63721b6bacf0b4a"}, +] + [[package]] name = "iniconfig" version = "2.0.0" @@ -1028,6 +1056,17 @@ files = [ {file = "starlette-0.35.1.tar.gz", hash = "sha256:3e2639dac3520e4f58734ed22553f950d3f3cb1001cd2eaac4d57e8cdc5f66bc"}, ] +[[package]] +name = "tenacity" +version = "8.2.3" +requires_python = ">=3.7" +summary = "Retry code until it succeeds" +groups = ["default"] +files = [ + {file = "tenacity-8.2.3-py3-none-any.whl", hash = "sha256:ce510e327a630c9e1beaf17d42e6ffacc88185044ad85cf74c0a8887c6a0f88c"}, + {file = "tenacity-8.2.3.tar.gz", hash = "sha256:5398ef0d78e63f40007c1fb4c0bff96e1911394d2fa8d194f77619c05ff6cc8a"}, +] + [[package]] name = "tomli" version = "2.0.1" @@ -1325,3 +1364,15 @@ files = [ {file = "win32_setctime-1.1.0-py3-none-any.whl", hash = "sha256:231db239e959c2fe7eb1d7dc129f11172354f98361c4fa2d6d2d7e278baa8aad"}, {file = "win32_setctime-1.1.0.tar.gz", hash = "sha256:15cf5750465118d6929ae4de4eb46e8edae9a5634350c01ba582df868e932cb2"}, ] + +[[package]] +name = "zipp" +version = "3.17.0" +requires_python = ">=3.8" +summary = "Backport of pathlib-compatible object wrapper for zip files" +groups = ["default"] +marker = "python_version < \"3.10\"" +files = [ + {file = "zipp-3.17.0-py3-none-any.whl", hash = "sha256:0e923e726174922dce09c53c59ad483ff7bbb8e572e00c7f7c46b88556409f31"}, + {file = "zipp-3.17.0.tar.gz", hash = "sha256:84e64a1c28cf7e91ed2078bb8cc8c259cb19b76942096c8d7b84947690cabaf0"}, +] diff --git a/playground/generate_image.py b/playground/generate_image.py index 12ea1e6..c16b703 100644 --- a/playground/generate_image.py +++ b/playground/generate_image.py @@ -38,6 +38,7 @@ async def main(): session=globe_s, remove_sign=True ) except APIError as e: + print(str(e)) print(e.response) return diff --git a/playground/read_nai_tag.py b/playground/read_nai_tag.py index a774610..b927630 100644 --- a/playground/read_nai_tag.py +++ b/playground/read_nai_tag.py @@ -9,8 +9,11 @@ if not os.path.exists("generate_image.png"): raise FileNotFoundError("generate_image.png not found,pls run generate_image.py first") +try: + meta = NovelAiMetadata.build_from_img(image_io="generate_image.png") # OR BytesIO(data) +except ValueError: + raise LookupError("Cant find a MetaData") -ba = NovelAiMetadata.build_from_img(image_io="generate_image.png") -print(ba.title) -print(ba.description) -print(ba.comment) +print(meta.title) +print(meta.description) +print(meta.comment) diff --git a/pyproject.toml b/pyproject.toml index 2bc3797..96e5d75 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "novelai-python" -version = "0.3.3" +version = "0.3.4" description = "Novelai Python Binding With Pydantic" authors = [ { name = "sudoskys", email = "coldlando@hotmail.com" }, @@ -19,6 +19,8 @@ dependencies = [ "numpy>=1.24.4", "argon2-cffi>=23.1.0", "opencv-python>=4.8.1.78", + "fake-useragent>=1.4.0", + "tenacity>=8.2.3", ] requires-python = ">=3.8" readme = "README.md" diff --git a/src/novelai_python/credential/ApiToken.py b/src/novelai_python/credential/ApiToken.py index 525df5f..81f6c63 100644 --- a/src/novelai_python/credential/ApiToken.py +++ b/src/novelai_python/credential/ApiToken.py @@ -7,7 +7,7 @@ from loguru import logger from pydantic import SecretStr, Field, field_validator -from ._base import CredentialBase +from ._base import CredentialBase, FAKE_UA class ApiCredential(CredentialBase): @@ -25,12 +25,12 @@ async def get_session(self, timeout: int = 180, update_headers: dict = None): "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", - "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36", + "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/", - }, impersonate="chrome110") + }, impersonate="edge101") self._session.headers.update(update_headers) return self._session diff --git a/src/novelai_python/credential/JwtToken.py b/src/novelai_python/credential/JwtToken.py index b7247ad..5e6e75e 100644 --- a/src/novelai_python/credential/JwtToken.py +++ b/src/novelai_python/credential/JwtToken.py @@ -7,7 +7,7 @@ from loguru import logger from pydantic import SecretStr, Field, field_validator -from ._base import CredentialBase +from ._base import CredentialBase, FAKE_UA class JwtCredential(CredentialBase): @@ -24,13 +24,13 @@ async def get_session(self, timeout: int = 180, update_headers: dict = None): self._session = AsyncSession(timeout=timeout, headers={ "Accept": "*/*", "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2", - "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36", + "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/", - }, impersonate="chrome110") + }, impersonate="edge101") assert isinstance(update_headers, dict), "update_headers must be a dict" self._session.headers.update(update_headers) return self._session diff --git a/src/novelai_python/credential/UserAuth.py b/src/novelai_python/credential/UserAuth.py index ff82499..5fd1e86 100644 --- a/src/novelai_python/credential/UserAuth.py +++ b/src/novelai_python/credential/UserAuth.py @@ -9,7 +9,7 @@ from curl_cffi.requests import AsyncSession from pydantic import SecretStr, Field -from ._base import CredentialBase +from ._base import CredentialBase, FAKE_UA class LoginCredential(CredentialBase): @@ -30,14 +30,14 @@ async def get_session(self, timeout: int = 180, update_headers: dict = None): resp = await Login.build(user_name=self.username, password=self.password.get_secret_value()).request() self._session = AsyncSession(timeout=timeout, headers={ "Accept": "*/*", - "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36", + "User-Agent": FAKE_UA.edge, "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") + }, impersonate="edge101") self._update_at = int(time.time()) self._session.headers.update(update_headers) return self._session diff --git a/src/novelai_python/credential/_base.py b/src/novelai_python/credential/_base.py index e79fc4b..94ec5d6 100644 --- a/src/novelai_python/credential/_base.py +++ b/src/novelai_python/credential/_base.py @@ -4,8 +4,11 @@ # @File : _shema.py # @Software: PyCharm from curl_cffi.requests import AsyncSession +from fake_useragent import UserAgent from pydantic import BaseModel +FAKE_UA = UserAgent() + class CredentialBase(BaseModel): _session: AsyncSession = None diff --git a/src/novelai_python/sdk/ai/generate_image/__init__.py b/src/novelai_python/sdk/ai/generate_image/__init__.py index 4ba6742..cde9550 100644 --- a/src/novelai_python/sdk/ai/generate_image/__init__.py +++ b/src/novelai_python/sdk/ai/generate_image/__init__.py @@ -18,6 +18,7 @@ from curl_cffi.requests import AsyncSession from loguru import logger from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, model_validator, Field +from tenacity import retry, stop_after_attempt, wait_random, retry_if_exception from typing_extensions import override from ._enum import Model, Sampler, NoiseSchedule, ControlNetModel, Action, UCPreset @@ -179,6 +180,9 @@ def validate_model(self): logger.warning("Mask maybe required for infill mode.") if self.action != Action.GENERATE: self.parameters.extra_noise_seed = self.parameters.seed + if self.action == Action.IMG2IMG: + self.parameters.sm = False + self.parameters.sm_dyn = False return self @property @@ -322,9 +326,14 @@ async def necessary_headers(self, request_data) -> dict: "Sec-Fetch-Site": "same-site", "Pragma": "no-cache", "Cache-Control": "no-cache", - "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36", } + @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"], *, diff --git a/src/novelai_python/sdk/ai/generate_image/suggest_tags.py b/src/novelai_python/sdk/ai/generate_image/suggest_tags.py index 0406376..a56cb65 100644 --- a/src/novelai_python/sdk/ai/generate_image/suggest_tags.py +++ b/src/novelai_python/sdk/ai/generate_image/suggest_tags.py @@ -41,7 +41,6 @@ async def necessary_headers(self, request_data) -> dict: return { "Host": urlparse(self.endpoint).netloc, "Accept": "*/*", - "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36", "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", "Access-Control-Allow-Origin": "*", diff --git a/src/novelai_python/sdk/ai/upscale.py b/src/novelai_python/sdk/ai/upscale.py index dc4d7b8..f21fb65 100644 --- a/src/novelai_python/sdk/ai/upscale.py +++ b/src/novelai_python/sdk/ai/upscale.py @@ -15,6 +15,7 @@ from curl_cffi.requests import AsyncSession from loguru import logger from pydantic import ConfigDict, PrivateAttr, model_validator +from tenacity import wait_random, retry, stop_after_attempt, retry_if_exception from ..schema import ApiBaseModel from ..._exceptions import APIError, AuthError, SessionHttpError @@ -73,7 +74,6 @@ async def necessary_headers(self, request_data) -> dict: return { "Host": urlparse(self.endpoint).netloc, "Accept": "*/*", - "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36", "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/", @@ -88,6 +88,12 @@ async def necessary_headers(self, request_data) -> dict: "Cache-Control": "no-cache", } + @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"], *, diff --git a/src/novelai_python/utils/hash.py b/src/novelai_python/utils/hash.py index 7ae83e0..8068eb4 100644 --- a/src/novelai_python/utils/hash.py +++ b/src/novelai_python/utils/hash.py @@ -74,17 +74,18 @@ def write_out(self, img_io: BytesIO, *, remove_stealth: bool = False): @classmethod 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': + title = img.info.get("Title", "Empty") + if title != 'AI generated image': raise ValueError("Not a NaiPic") prompt = img.info.get("Description") - comment = img.info.get("Comment") + comment = img.info.get("Comment", None) + assert isinstance(comment, str), ValueError("Comment Empty") try: comment = json.loads(comment) except Exception as e: logger.debug(e) comment = {} - return cls(title=title, description=prompt, comment=comment) + return cls(title=title, description=prompt, comment=comment) @classmethod def build_from_param(cls, prompt, neg_prompt, **kwargs):