novelai-python 0.7.1__tar.gz → 0.7.3__tar.gz

{novelai_python-0.7.1 → novelai_python-0.7.3}/PKG-INFO

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

{novelai_python-0.7.1 → novelai_python-0.7.3}/pyproject.toml

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

{novelai_python-0.7.1 → novelai_python-0.7.3}/src/novelai_python/_exceptions.py

@@ -74,6 +74,14 @@ class ConcurrentGenerationError(APIError):
     pass
 
 
+class DataSerializationError(APIError):
+    """
+    DataSerializationError is raised when the API data is not serializable.
+    """
+
+    pass
+
+
 class AuthError(APIError):
     """
     AuthError is raised when the API returns an error.

{novelai_python-0.7.1 → novelai_python-0.7.3}/src/novelai_python/sdk/ai/augment_image/__init__.py

@@ -11,7 +11,7 @@ from copy import deepcopy
 from io import BytesIO
 from typing import Optional, Union, IO, Any
 from urllib.parse import urlparse
-from zipfile import ZipFile
+from zipfile import ZipFile, BadZipFile
 
 import curl_cffi
 import httpx
@@ -25,13 +25,16 @@ from novelai_python.sdk.ai._cost import CostCalculator
 from novelai_python.sdk.ai._enum import Model
 from ._enum import ReqType, Moods
 from ...schema import ApiBaseModel
-from ...._exceptions import APIError, AuthError, ConcurrentGenerationError, SessionHttpError
+from ...._exceptions import APIError, AuthError, ConcurrentGenerationError, SessionHttpError, DataSerializationError
 from ...._response.ai.generate_image import ImageGenerateResp, RequestParams
 from ....credential import CredentialBase
 from ....utils import try_jsonfy
 
 
 class AugmentImageInfer(ApiBaseModel):
+    """
+    https://docs.novelai.net/image/directortools.html
+    """
     _endpoint: str = PrivateAttr("https://image.novelai.net")
 
     @property
@@ -210,96 +213,104 @@ class AugmentImageInfer(ApiBaseModel):
         :param session:  session
         :return:
         """
-        # Data Build
+        # Prepare request data
         request_data = self.model_dump(mode="json", exclude_none=True)
         async with session if isinstance(session, AsyncSession) else await session.get_session() as sess:
-            # Header
             sess.headers.update(await self.necessary_headers(request_data))
             if override_headers:
                 sess.headers.clear()
                 sess.headers.update(override_headers)
+
+            # Log the request data (sanitize sensitive content)
             try:
                 _log_data = deepcopy(request_data)
-                _log_data.update({
-                    "image": "base64 data"
-                })
-                logger.debug(f"Request Data: {_log_data}")
-                del _log_data
+                if self.image:
+                    _log_data["image"] = "base64 data hidden"
+                logger.debug(f"Request Data: {json.dumps(_log_data, indent=2)}")
             except Exception as e:
-                logger.warning(f"Error when print log data: {e}")
+                logger.warning(f"Failed to log request data: {e}")
+
+            # Perform request and handle response
             try:
-                assert hasattr(sess, "post"), "session must have post method."
+                self.ensure_session_has_post_method(sess)
                 response = await sess.post(
                     self.base_url,
                     data=json.dumps(request_data).encode("utf-8")
                 )
-                if response.headers.get('Content-Type') not in ['binary/octet-stream', 'application/x-zip-compressed']:
-                    logger.warning(
-                        f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
-                    )
-                    try:
-                        _msg = response.json()
-                    except Exception as e:
-                        logger.warning(e)
-                        if not isinstance(response.content, str) and len(response.content) < 50:
-                            raise APIError(
-                                message=f"Unexpected content type: {response.headers.get('Content-Type')}",
-                                request=request_data,
-                                code=response.status_code,
-                                response=try_jsonfy(response.content)
-                            )
-                        else:
-                            _msg = {"statusCode": response.status_code, "message": response.content}
-                    status_code = _msg.get("statusCode", response.status_code)
-                    message = _msg.get("message", "Unknown error")
+                if (
+                        response.headers.get('Content-Type') not in ['binary/octet-stream',
+                                                                     'application/x-zip-compressed']
+                        or response.status_code >= 400
+                ):
+                    error_message = await self.handle_error_response(response, request_data)
+                    status_code = error_message.get("statusCode", response.status_code)
+                    message = error_message.get("message", "Unknown error")
                     if status_code in [400, 401, 402]:
                         # 400 : validation error
                         # 401 : unauthorized
                         # 402 : payment required
                         # 409 : conflict
-                        raise AuthError(message, request=request_data, code=status_code, response=_msg)
-                    if status_code in [409]:
-                        # conflict error
-                        raise APIError(message, request=request_data, code=status_code, response=_msg)
-                    if status_code in [429]:
-                        # concurrent error
+                        raise AuthError(message, request=request_data, code=status_code, response=error_message)
+                    elif status_code == 409:
+                        raise APIError(message, request=request_data, code=status_code, response=error_message)
+                    elif status_code == 429:
                         raise ConcurrentGenerationError(
                             message=message,
                             request=request_data,
                             code=status_code,
-                            response=_msg
-                        )
-                    raise APIError(message, request=request_data, code=status_code, response=_msg)
-                zip_file = ZipFile(BytesIO(response.content))
-                unzip_content = []
-                with zip_file as zf:
-                    file_list = zf.namelist()
-                    if not file_list:
-                        raise APIError(
-                            message="No file in zip",
-                            request=request_data,
-                            code=response.status_code,
-                            response=try_jsonfy(response.content)
+                            response=error_message,
                         )
-                    for filename in file_list:
-                        data = zip_file.read(filename)
-                        unzip_content.append((filename, data))
-                return ImageGenerateResp(
-                    meta=RequestParams(
-                        endpoint=self.base_url,
-                        raw_request=request_data,
-                    ),
-                    files=unzip_content
-                )
+                    else:
+                        raise APIError(message, request=request_data, code=status_code, response=error_message)
+
+                # Unpack the ZIP response
+                try:
+                    zip_file = ZipFile(BytesIO(response.content))
+                    unzip_content = []
+                    with zip_file as zf:
+                        file_list = zf.namelist()
+                        if not file_list:
+                            raise DataSerializationError(
+                                message="The ZIP response contains no files.",
+                                request=request_data,
+                                response=try_jsonfy(response.content),
+                                code=response.status_code,
+                            )
+                        for filename in file_list:
+                            data = zip_file.read(filename)
+                            unzip_content.append((filename, data))
+                    return ImageGenerateResp(
+                        meta=RequestParams(
+                            endpoint=self.base_url,
+                            raw_request=request_data,
+                        ),
+                        files=unzip_content,
+                    )
+                except BadZipFile as e:
+                    # Invalid ZIP file - indicate serialization error
+                    logger.exception("The response content is not a valid ZIP file.")
+                    raise DataSerializationError(
+                        message="Invalid ZIP file received from the API.",
+                        request=request_data,
+                        response={},
+                        code=response.status_code,
+                    ) from e
+                except Exception as e:
+                    logger.exception("Unexpected error while unpacking ZIP response.")
+                    raise DataSerializationError(
+                        message="An unexpected error occurred while processing ZIP data.",
+                        request=request_data,
+                        response={},
+                        code=response.status_code,
+                    ) from e
             except curl_cffi.requests.errors.RequestsError as exc:
                 logger.exception(exc)
-                raise SessionHttpError(
-                    "An AsyncSession RequestsError occurred, maybe SSL error. Try again later!") from exc
+                raise SessionHttpError("A RequestsError occurred (e.g., SSL error). Try again later.")
             except httpx.HTTPError as exc:
                 logger.exception(exc)
-                raise SessionHttpError("An HTTPError occurred, maybe SSL error. Try again later!") from exc
+                raise SessionHttpError("An HTTP error occurred. Try again later.")
             except APIError as e:
                 raise e
             except Exception as e:
-                logger.opt(exception=e).exception("An Unexpected error occurred")
-                raise e
+                logger.opt(exception=e).exception("Unexpected error occurred during the request.")
+                raise Exception("An unexpected error occurred.") from e

{novelai_python-0.7.1 → novelai_python-0.7.3}/src/novelai_python/sdk/ai/generate/__init__.py

@@ -369,7 +369,7 @@ class LLM(ApiBaseModel):
             logger.debug(f"LLM request data: {json.dumps(request_data)}")
             # Request
             try:
-                assert hasattr(sess, "post"), "session must have post method."
+                self.ensure_session_has_post_method(sess)
                 response = await sess.post(
                     self.base_url,
                     json=request_data,

{novelai_python-0.7.1 → novelai_python-0.7.3}/src/novelai_python/sdk/ai/generate_image/__init__.py

@@ -12,7 +12,7 @@ from enum import Enum
 from io import BytesIO
 from typing import Optional, Union, Tuple, List
 from urllib.parse import urlparse
-from zipfile import ZipFile
+from zipfile import ZipFile, BadZipFile
 
 import curl_cffi
 import cv2
@@ -32,7 +32,7 @@ from novelai_python.sdk.ai._enum import Model, Sampler, NoiseSchedule, ControlNe
     get_model_group, ModelGroups, get_supported_params, get_modifiers, ImageBytesTypeAlias
 from .schema import Character, V4Prompt, V4NegativePrompt, PositionMap
 from ...schema import ApiBaseModel
-from ...._exceptions import APIError, AuthError, ConcurrentGenerationError, SessionHttpError
+from ...._exceptions import APIError, AuthError, ConcurrentGenerationError, SessionHttpError, DataSerializationError
 from ...._response.ai.generate_image import ImageGenerateResp, RequestParams
 from ....credential import CredentialBase
 from ....utils import try_jsonfy
@@ -344,6 +344,35 @@ class Params(BaseModel):
 
 class GenerateImageInfer(ApiBaseModel):
     _endpoint: str = PrivateAttr("https://image.novelai.net")
+    _mutual_exclusion: bool = PrivateAttr(False)
+    """Positive words and negative words are mutually exclusive, and conflicting negative words are deleted first."""
+    _quality_modifier: bool = PrivateAttr(True)
+    """Add Quality Modifier To Input"""
+
+    def set_mutual_exclusion(self, value: bool):
+        """
+        **Enable This will modify the negative prompt.**
+        Default is False.
+
+        Positive words and negative words are mutually exclusive, and conflicting negative words are deleted first.
+        :param value: bool
+        :return: self
+        """
+        self._mutual_exclusion = bool(value)
+        return self
+
+    def set_quality_modifier(self, value: bool):
+        """
+        **Enable This will modify the input prompt.**
+        Default is True.
+
+        Add Quality Modifier To Input.
+        Whether to add the quality vocabulary used by the web application.
+        :param value:
+        :return:
+        """
+        self._quality_modifier = bool(value)
+        return self
 
     @property
     def endpoint(self):
@@ -374,10 +403,19 @@ class GenerateImageInfer(ApiBaseModel):
         # Add negative prompt based on ucPreset
         if self.parameters.ucPreset is not None:
             uc_preset = self.parameters.ucPreset
+
             # If ucPreset is Enum, get the value
             if isinstance(self.parameters.ucPreset, Enum):
                 uc_preset = self.parameters.ucPreset.value
             default_negative_prompt = get_default_uc_preset(self.model, uc_preset)
+
+            # Lowres means we don't found any negative prompt.
+            # If the default negative prompt is lowres, and the user has set a negative prompt,
+            # then the default negative prompt is not added.
+            if self.parameters.negative_prompt and default_negative_prompt == "lowres":
+                default_negative_prompt = ""
+
+            # Combine the negative prompt preset and the user's negative prompt
             self.parameters.negative_prompt = ", ".join(
                 filter(None, [default_negative_prompt, self.parameters.negative_prompt])
             )
@@ -429,7 +467,19 @@ class GenerateImageInfer(ApiBaseModel):
                 )
                 return _prompt
 
-        self.input = enhance_message(self.input)
+        if self._quality_modifier:
+            logger.trace("Enhancing input with quality modifier, will modify input prompt.")
+            self.input = enhance_message(self.input)
+
+        if self._mutual_exclusion:
+            logger.trace("Mutual exclusion is enabled, will modify negative prompt.")
+            input_prompt = {x.strip(): x for x in self.input.split(",")}
+            uc_prompt = {x.strip(): x for x in self.parameters.negative_prompt.split(",")}
+            # Remove conflicting negative words
+            for key in input_prompt:
+                if key in uc_prompt:
+                    uc_prompt.pop(key)
+            self.parameters.negative_prompt = ",".join(uc_prompt.values())
 
     @model_validator(mode="after")
     def _build_nai4_prompt(self):
@@ -467,7 +517,9 @@ class GenerateImageInfer(ApiBaseModel):
     @model_validator(mode="after")
     def _backend_logic(self):
         """
-        Backend logic for GenerateImageInfer.
+        Calibration results.
+
+        **NOTE: It is forbidden to write logic that does not belong to novelai.net here**
         :return: self
         """
         if self.action == Action.GENERATE:
@@ -943,111 +995,132 @@ class GenerateImageInfer(ApiBaseModel):
         retry=retry_if_exception(lambda e: hasattr(e, "code") and str(e.code) == "500"),
         reraise=True
     )
-    async def request(self,
-                      session: Union[AsyncSession, "CredentialBase"],
-                      *,
-                      override_headers: Optional[dict] = None,
-                      ) -> ImageGenerateResp:
+    async def request(
+            self,
+            session: Union[AsyncSession, "CredentialBase"],
+            *,
+            override_headers: Optional[dict] = None,
+    ) -> ImageGenerateResp:
         """
-        Generate images using NovelAI's diffusion models. According to our Terms of Service, all generation requests must be initiated by a human action. Automating text or image generation to create excessive load on our systems is not allowed.
-
-        :param override_headers: the headers to override
-        :param session:  session
-        :return:
+        **Generate images using NovelAI's diffusion models.**
+
+        According to our Terms of Service, all generation requests must be initiated by a human action. Automating text
+        or image generation to create excessive load on our systems is not allowed.
+
+        :param override_headers: Headers to override the default headers.
+        :param session: Async session object or credential-based session.
+        :return: ImageGenerateResp containing the response data and metadata.
+        :raises AuthError: If the request is unauthorized.
+        :raises APIError: If the API returns an error.
+        :raises ConcurrentGenerationError: If the request is rate-limited.
+        :raises SessionHttpError: If an HTTP error occurs.
+        :raises DataSerializationError: If an error occurs while processing the response data.
         """
-        # Data Build
+        # Prepare request data
         request_data = self.model_dump(mode="json", exclude_none=True)
         async with session if isinstance(session, AsyncSession) else await session.get_session() as sess:
-            # Header
             sess.headers.update(await self.necessary_headers(request_data))
             if override_headers:
                 sess.headers.clear()
                 sess.headers.update(override_headers)
+
+            # Log the request data (sanitize sensitive content)
             try:
                 _log_data = deepcopy(request_data)
                 if self.parameters.image:
-                    _log_data["parameters"]["image"] = "base64 data"
+                    _log_data["parameters"]["image"] = "base64 data hidden"
                 if self.parameters.mask:
-                    _log_data["parameters"]["mask"] = "base64 data"
+                    _log_data["parameters"]["mask"] = "base64 data hidden"
                 if self.parameters.reference_image_multiple:
-                    _log_data["parameters"]["reference_image_multiple"] = ["base64 data"] * len(
-                        self.parameters.reference_image_multiple)
+                    _log_data["parameters"]["reference_image_multiple"] = ["base64 data hidden"] * len(
+                        self.parameters.reference_image_multiple
+                    )
                 logger.debug(f"Request Data: {json.dumps(_log_data, indent=2)}")
-                del _log_data
             except Exception as e:
-                logger.warning(f"Error when print log data: {e}")
+                logger.warning(f"Failed to log request data: {e}")
+
+            # Perform request and handle response
             try:
-                assert hasattr(sess, "post"), "session must have post method."
+                self.ensure_session_has_post_method(sess)
                 response = await sess.post(
                     self.base_url,
-                    data=json.dumps(request_data).encode("utf-8")
+                    data=json.dumps(request_data).encode("utf-8"),
                 )
-                if response.headers.get('Content-Type') not in ['binary/octet-stream', 'application/x-zip-compressed']:
-                    logger.warning(
-                        f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
-                    )
-                    try:
-                        _msg = response.json()
-                    except Exception as e:
-                        logger.warning(e)
-                        if not isinstance(response.content, str) and len(response.content) < 50:
-                            raise APIError(
-                                message=f"Unexpected content type: {response.headers.get('Content-Type')}",
-                                request=request_data,
-                                code=response.status_code,
-                                response=try_jsonfy(response.content)
-                            )
-                        else:
-                            _msg = {"statusCode": response.status_code, "message": response.content}
-                    status_code = _msg.get("statusCode", response.status_code)
-                    message = _msg.get("message", "Unknown error")
+                # Validate response content type and status code
+                if (
+                        response.headers.get("Content-Type")
+                        not in ["binary/octet-stream", "application/x-zip-compressed"]
+                        or response.status_code >= 400
+                ):
+                    error_message = await self.handle_error_response(response, request_data)
+                    status_code = error_message.get("statusCode", response.status_code)
+                    message = error_message.get("message", "Unknown error")
                     if status_code in [400, 401, 402]:
                         # 400 : validation error
                         # 401 : unauthorized
                         # 402 : payment required
                         # 409 : conflict
-                        raise AuthError(message, request=request_data, code=status_code, response=_msg)
-                    if status_code in [409]:
-                        # conflict error
-                        raise APIError(message, request=request_data, code=status_code, response=_msg)
-                    if status_code in [429]:
-                        # concurrent error
+                        raise AuthError(message, request=request_data, code=status_code, response=error_message)
+                    elif status_code == 409:
+                        raise APIError(message, request=request_data, code=status_code, response=error_message)
+                    elif status_code == 429:
                         raise ConcurrentGenerationError(
                             message=message,
                             request=request_data,
                             code=status_code,
-                            response=_msg
+                            response=error_message,
                         )
-                    raise APIError(message, request=request_data, code=status_code, response=_msg)
-                zip_file = ZipFile(BytesIO(response.content))
-                unzip_content = []
-                with zip_file as zf:
-                    file_list = zf.namelist()
-                    if not file_list:
-                        raise APIError(
-                            message="No file in zip",
-                            request=request_data,
-                            code=response.status_code,
-                            response=try_jsonfy(response.content)
-                        )
-                    for filename in file_list:
-                        data = zip_file.read(filename)
-                        unzip_content.append((filename, data))
-                return ImageGenerateResp(
-                    meta=RequestParams(
-                        endpoint=self.base_url,
-                        raw_request=request_data,
-                    ),
-                    files=unzip_content
-                )
+                    else:
+                        raise APIError(message, request=request_data, code=status_code, response=error_message)
+
+                # Unpack the ZIP response
+                try:
+                    zip_file = ZipFile(BytesIO(response.content))
+                    unzip_content = []
+                    with zip_file as zf:
+                        file_list = zf.namelist()
+                        if not file_list:
+                            raise DataSerializationError(
+                                message="The ZIP response contains no files.",
+                                request=request_data,
+                                response=try_jsonfy(response.content),
+                                code=response.status_code,
+                            )
+                        for filename in file_list:
+                            data = zip_file.read(filename)
+                            unzip_content.append((filename, data))
+                    return ImageGenerateResp(
+                        meta=RequestParams(
+                            endpoint=self.base_url,
+                            raw_request=request_data,
+                        ),
+                        files=unzip_content,
+                    )
+                except BadZipFile as e:
+                    # Invalid ZIP file - indicate serialization error
+                    logger.exception("The response content is not a valid ZIP file.")
+                    raise DataSerializationError(
+                        message="Invalid ZIP file received from the API.",
+                        request=request_data,
+                        response={},
+                        code=response.status_code,
+                    ) from e
+                except Exception as e:
+                    logger.exception("Unexpected error while unpacking ZIP response.")
+                    raise DataSerializationError(
+                        message="An unexpected error occurred while processing ZIP data.",
+                        request=request_data,
+                        response={},
+                        code=response.status_code,
+                    ) from e
             except curl_cffi.requests.errors.RequestsError as exc:
                 logger.exception(exc)
-                raise SessionHttpError("An AsyncSession RequestsError occurred, maybe SSL error. Try again later!")
+                raise SessionHttpError("A RequestsError occurred (e.g., SSL error). Try again later.")
             except httpx.HTTPError as exc:
                 logger.exception(exc)
-                raise SessionHttpError("An HTTPError occurred, maybe SSL error. Try again later!")
+                raise SessionHttpError("An HTTP error occurred. Try again later.")
             except APIError as e:
                 raise e
             except Exception as e:
-                logger.opt(exception=e).exception("An Unexpected error occurred")
-                raise e
+                logger.opt(exception=e).exception("Unexpected error occurred during the request.")
+                raise Exception("An unexpected error occurred.") from e

{novelai_python-0.7.1 → novelai_python-0.7.3}/src/novelai_python/sdk/ai/generate_image/suggest_tags.py

@@ -77,41 +77,28 @@ class SuggestTags(ApiBaseModel):
         if override_headers:
             session.headers.clear()
             session.headers.update(override_headers)
-        logger.debug("SuggestTags")
         try:
             assert hasattr(session, "get"), "session must have get method."
             response = await session.get(
                 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.warning(
-                    f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
-                )
-                try:
-                    _msg = response.json()
-                except Exception as e:
-                    logger.warning(e)
-                    if not isinstance(response.content, str) and len(response.content) < 50:
-                        raise APIError(
-                            message=f"Unexpected content type: {response.headers.get('Content-Type')}",
-                            request=request_data,
-                            code=response.status_code,
-                            response=try_jsonfy(response.content)
-                        )
-                    else:
-                        _msg = {"statusCode": response.status_code, "message": response.content}
-                status_code = _msg.get("statusCode", response.status_code)
-                message = _msg.get("message", "Unknown error")
+            if (
+                    "application/json" not in response.headers.get('Content-Type')
+                    or response.status_code != 200
+            ):
+                error_message = await self.handle_error_response(response=response, request_data=request_data)
+                status_code = error_message.get("statusCode", response.status_code)
+                message = error_message.get("message", "Unknown error")
                 if status_code in [400, 401, 402]:
                     # 400 : validation error
                     # 401 : unauthorized
                     # 402 : payment required
                     # 409 : conflict
-                    raise AuthError(message, request=request_data, code=status_code, response=_msg)
+                    raise AuthError(message, request=request_data, code=status_code, response=error_message)
                 if status_code in [500]:
                     # An unknown error occured.
-                    raise APIError(message, request=request_data, code=status_code, response=_msg)
-                raise APIError(message, request=request_data, code=status_code, response=_msg)
+                    raise APIError(message, request=request_data, code=status_code, response=error_message)
+                raise APIError(message, request=request_data, code=status_code, response=error_message)
             return SuggestTagsResp.model_validate(response.json())
         except curl_cffi.requests.errors.RequestsError as exc:
             logger.exception(exc)

{novelai_python-0.7.1 → novelai_python-0.7.3}/src/novelai_python/sdk/ai/generate_stream.py

@@ -95,7 +95,7 @@ class LLMStream(LLM):
             logger.debug(f"StreamLLM request data: {request_data}")
             # Request
             try:
-                assert hasattr(sess, "post"), "session must have post method."
+                self.ensure_session_has_post_method(sess)
                 response = await sess.post(
                     self.base_url,
                     json=request_data,