novelai-python 0.3.3__tar.gz → 0.3.5__tar.gz

{novelai_python-0.3.3 → novelai_python-0.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.3.3
+Version: 0.3.5
 Summary: Novelai Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache License 2.0
@@ -20,6 +20,8 @@ Requires-Dist: uvicorn[standard]>=0.27.0.post1
 Requires-Dist: numpy>=1.24.4
 Requires-Dist: argon2-cffi>=23.1.0
 Requires-Dist: opencv-python>=4.8.1.78
+Requires-Dist: fake-useragent>=1.4.0
+Requires-Dist: tenacity>=8.2.3
 Requires-Dist: pytest>=7.4.4; extra == "testing"
 Requires-Dist: pytest-asyncio>=0.23.4; extra == "testing"
 Provides-Extra: testing
@@ -34,6 +36,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

{novelai_python-0.3.3 → novelai_python-0.3.5}/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

{novelai_python-0.3.3 → novelai_python-0.3.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.3.3"
+version = "0.3.5"
 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"

{novelai_python-0.3.3 → novelai_python-0.3.5}/src/novelai_python/credential/ApiToken.py

@@ -7,7 +7,7 @@ from curl_cffi.requests import AsyncSession
 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 @@ class ApiCredential(CredentialBase):
                 "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
 

{novelai_python-0.3.3 → novelai_python-0.3.5}/src/novelai_python/credential/JwtToken.py

@@ -7,7 +7,7 @@ from curl_cffi.requests import AsyncSession
 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 @@ class JwtCredential(CredentialBase):
             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

{novelai_python-0.3.3 → novelai_python-0.3.5}/src/novelai_python/credential/UserAuth.py

@@ -9,7 +9,7 @@ from typing import Optional
 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 @@ class LoginCredential(CredentialBase):
             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

{novelai_python-0.3.3 → novelai_python-0.3.5}/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

{novelai_python-0.3.3 → novelai_python-0.3.5}/src/novelai_python/sdk/ai/generate_image/__init__.py

@@ -18,6 +18,7 @@ import httpx
 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
@@ -33,7 +34,7 @@ class GenerateImageInfer(ApiBaseModel):
 
     class Params(BaseModel):
         # Inpaint
-        add_original_image: Optional[bool] = False
+        add_original_image: Optional[bool] = False  # FIXME: 未知作用
         mask: Optional[Union[str, bytes]] = None  # img2img,base64
 
         cfg_rescale: Optional[float] = Field(0, ge=0, le=1, multiple_of=0.02)
@@ -56,6 +57,20 @@ class GenerateImageInfer(ApiBaseModel):
 
         # Misc
         params_version: Optional[int] = 1
+        reference_image: Optional[Union[str, bytes]] = None
+        reference_information_extracted: Optional[float] = Field(1,
+                                                                 ge=0,
+                                                                 le=1,
+                                                                 multiple_of=0.01,
+                                                                 description="extracting concepts or features"
+                                                                 )
+        reference_strength: Optional[float] = Field(0.6,
+                                                    ge=0,
+                                                    le=1,
+                                                    multiple_of=0.01,
+                                                    description="the stronger the AI will try to emulate visual cues."
+                                                    )
+
         legacy: Optional[bool] = False
         legacy_v3_extend: Optional[bool] = False
 
@@ -83,14 +98,6 @@ class GenerateImageInfer(ApiBaseModel):
         """Undesired Content Strength"""
         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:
-                raise ValueError('Invalid Model Params For img2img2 mode... image should match add_original_image!')
-            return self
-
         @field_validator('mask')
         def mask_validator(cls, v: Union[str, bytes]):
             if isinstance(v, str) and v.startswith("data:image/"):
@@ -107,6 +114,14 @@ class GenerateImageInfer(ApiBaseModel):
                 return base64.b64encode(v).decode("utf-8")
             return v
 
+        @field_validator('reference_image')
+        def image_validator(cls, v: Union[str, bytes]):
+            if isinstance(v, str) and v.startswith("data:image/"):
+                raise ValueError("Invalid image format, must be base64 encoded.")
+            if isinstance(v, bytes):
+                return base64.b64encode(v).decode("utf-8")
+            return v
+
         @field_validator('width')
         def width_validator(cls, v: int):
             """
@@ -175,10 +190,20 @@ class GenerateImageInfer(ApiBaseModel):
 
     @model_validator(mode="after")
     def validate_model(self):
-        if self.action == Action.INFILL and not self.parameters.mask:
-            logger.warning("Mask maybe required for infill mode.")
+        if self.action == Action.INFILL:
+            if not self.parameters.mask:
+                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
+            if not self.parameters.image:
+                raise ValueError("image is must required for img2img mode.")
+            if not self.parameters.add_original_image:
+                raise ValueError("add_original_image is must required for img2img mode.")
+        if self.parameters.image and self.parameters.reference_image:
+            logger.warning("image and reference_image should not be used together.")
         return self
 
     @property
@@ -243,11 +268,13 @@ class GenerateImageInfer(ApiBaseModel):
               qualityToggle: bool = True,
               ucPreset: Union[UCPreset, int] = UCPreset.TYPE0,
               image: Union[str, bytes] = None,
+              reference_mode: bool = False,
               add_original_image: bool = None,
               strength: float = None,
               mask: Union[str, bytes] = None,
               controlnet_model: Union[ControlNetModel, str] = None,
               controlnet_condition: str = None,
+              sm_dyn: bool = False,
               **kwargs
               ):
         """
@@ -260,17 +287,19 @@ class GenerateImageInfer(ApiBaseModel):
         :param negative_prompt: The content of negative prompt
         :param seed: The seed for generate image
         :param steps: The steps for generate image
-        :param scale: Prompt Guidance
+        :param scale: The scale for generate image
         :param cfg_rescale: Prompt Guidance Rescale 0-1 lower is more creative
         :param sampler: The sampler for generate image
         :param width: 宽
         :param height: 高
         :param image: 图片
+        :param reference_mode: 是否是参考模式
         :param add_original_image: 是否添加原始图片
         :param strength: IMG2IMG 强度
         :param mask: Inpainting mask
         :param controlnet_model: 控制网络模型
         :param controlnet_condition: 控制网络条件
+        :param sm_dyn: 是否使用动态sm
         :return: self
         """
         assert isinstance(prompt, str)
@@ -286,13 +315,22 @@ class GenerateImageInfer(ApiBaseModel):
             "height": height,
             "qualityToggle": qualityToggle,
             "ucPreset": ucPreset,
-            "image": image,
             "add_original_image": add_original_image,
-            "strength": strength,
             "mask": mask,
             "controlnet_model": controlnet_model,
-            "controlnet_condition": controlnet_condition
+            "controlnet_condition": controlnet_condition,
+            "sm_dyn": sm_dyn,
         })
+        if reference_mode:
+            kwargs.update({
+                "reference_image": image,
+                "reference_strength": strength,
+            })
+        else:
+            kwargs.update({
+                "image": image,
+                "strength": strength,
+            })
         # 清理空值
         param = {k: v for k, v in kwargs.items() if v is not None}
         return cls(
@@ -322,9 +360,14 @@ class GenerateImageInfer(ApiBaseModel):
             "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"],
                       *,
@@ -352,17 +395,15 @@ class GenerateImageInfer(ApiBaseModel):
         try:
             _log_data = deepcopy(request_data)
             if self.action == Action.GENERATE:
+                _log_data.get("parameters", {}).update({
+                    "reference_image": "base64 data" if self.parameters.reference_image else None,
+                })
                 logger.debug(f"Request Data: {_log_data}")
             else:
                 _log_data.get("parameters", {}).update({
                     "image": "base64 data" if self.parameters.image else None,
-                }
-                )
-                _log_data.get("parameters", {}).update(
-                    {
-                        "mask": "base64 data" if self.parameters.mask else None,
-                    }
-                )
+                    "mask": "base64 data" if self.parameters.mask else None,
+                })
                 logger.debug(f"Request Data: {_log_data}")
             del _log_data
         except Exception as e:
@@ -374,7 +415,7 @@ class GenerateImageInfer(ApiBaseModel):
                 data=json.dumps(request_data).encode("utf-8")
             )
             if response.headers.get('Content-Type') not in ['binary/octet-stream', 'application/x-zip-compressed']:
-                logger.error(
+                logger.warning(
                     f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
                 )
                 try:

{novelai_python-0.3.3 → novelai_python-0.3.5}/src/novelai_python/sdk/ai/generate_image/suggest_tags.py

@@ -41,7 +41,6 @@ class SuggestTags(ApiBaseModel):
         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": "*",
@@ -85,7 +84,7 @@ class SuggestTags(ApiBaseModel):
                 url=self.base_url + "?" + "&".join([f"{k}={v}" for k, v in request_data.items()])
             )
             if "application/json" not in response.headers.get('Content-Type') or response.status_code != 200:
-                logger.error(
+                logger.warning(
                     f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
                 )
                 try:

{novelai_python-0.3.3 → novelai_python-0.3.5}/src/novelai_python/sdk/ai/upscale.py

@@ -15,6 +15,7 @@ import httpx
 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 @@ class Upscale(ApiBaseModel):
         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 @@ class Upscale(ApiBaseModel):
             "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"],
                       *,
@@ -125,7 +131,7 @@ class Upscale(ApiBaseModel):
                 data=json.dumps(request_data).encode("utf-8")
             )
             if response.headers.get('Content-Type') not in ['binary/octet-stream', 'application/x-zip-compressed']:
-                logger.error(
+                logger.warning(
                     f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
                 )
                 try:

{novelai_python-0.3.3 → novelai_python-0.3.5}/src/novelai_python/sdk/user/information.py

@@ -73,7 +73,7 @@ class Information(ApiBaseModel):
                 self.base_url
             )
             if "application/json" not in response.headers.get('Content-Type') or response.status_code != 200:
-                logger.error(
+                logger.warning(
                     f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
                 )
                 try:

{novelai_python-0.3.3 → novelai_python-0.3.5}/src/novelai_python/sdk/user/login.py

@@ -92,7 +92,7 @@ class Login(ApiBaseModel):
                 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(
+                logger.warning(
                     f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
                 )
                 try:

{novelai_python-0.3.3 → novelai_python-0.3.5}/src/novelai_python/sdk/user/subscription.py

@@ -80,7 +80,7 @@ class Subscription(ApiBaseModel):
                 url=self.base_url,
             )
             if "application/json" not in response.headers.get('Content-Type') or response.status_code != 200:
-                logger.error(
+                logger.warning(
                     f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
                 )
                 try:

{novelai_python-0.3.3 → novelai_python-0.3.5}/src/novelai_python/utils/hash.py

@@ -74,17 +74,18 @@ class NovelAiMetadata(BaseModel):
     @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):