novelai-python 0.5.1__tar.gz → 0.5.3__tar.gz

{novelai_python-0.5.1 → novelai_python-0.5.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.5.1
+Version: 0.5.3
 Summary: NovelAI Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache License 2.0
@@ -29,6 +29,8 @@ Requires-Dist: regex>=2023.12.25
 Requires-Dist: tokenizers>=0.15.2
 Requires-Dist: json-repair>=0.29.4
 Requires-Dist: robust-downloader>=0.0.2
+Requires-Dist: bchlib>=2.1.3
+Requires-Dist: arrow>=1.3.0
 Provides-Extra: testing
 Requires-Dist: pytest>=7.4.4; extra == "testing"
 Requires-Dist: pytest-asyncio>=0.23.4; extra == "testing"
@@ -75,7 +77,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 ### Usage 🖥️
 
 ```shell
-pip install novelai-python
+pip -U install novelai-python
 ```
 
 More examples can be found in the [playground](https://github.com/LlmKira/novelai-python/tree/main/playground)
@@ -122,7 +124,8 @@ from dotenv import load_dotenv
 from pydantic import SecretStr
 
 from novelai_python import APIError, LoginCredential
-from novelai_python.sdk.ai.generate import TextLLMModel, LLM, get_default_preset
+from novelai_python.sdk.ai.generate import TextLLMModel, LLM, get_default_preset, AdvanceLLMSetting
+from novelai_python.sdk.ai.generate._enum import get_model_preset
 
 load_dotenv()
 username = os.getenv("NOVELAI_USER", None)
@@ -141,8 +144,17 @@ async def chat(prompt: str):
         agent = LLM.build(
             prompt=prompt,
             model=model,
-            parameters=None  # Auto Select or get from preset
+            # parameters=None,  # Auto Select or get from preset
+            parameters=get_model_preset(TextLLMModel.ERATO).get_all_presets()[0].parameters,  # Select from enum preset
+            advanced_setting=AdvanceLLMSetting(
+                min_length=1,
+                max_length=None,  # Auto
+            )
         )
+        # NOTE:parameter > advanced_setting, which logic in generate/__init__.py
+        # If you not pass the parameter, it will use the default preset.
+        # So if you want to set the generation params, you should pass your own params.
+        # Only if you want to use some params not affect the generation, you can use advanced_setting.
         result = await agent.request(session=login_credential)
     except APIError as e:
         raise Exception(f"Error: {e.message}")
@@ -184,6 +196,14 @@ print(tokenizer_package.decode(tokenizer_package.encode(t_text)))
 
 ```
 
+## About Nsfw 🚫
+
+You might need some solutions for identifying NSFW content and adding a mosaic to prevent operational mishaps.
+
+https://dghs-imgutils.deepghs.org/main/api_doc/detect/nudenet.html
+
+https://dghs-imgutils.deepghs.org/main/api_doc/operate/censor.html
+
 ## Acknowledgements 🙏
 
 [BackEnd](https://api.novelai.net/docs)

{novelai_python-0.5.1 → novelai_python-0.5.3}/README.md

@@ -39,7 +39,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 ### Usage 🖥️
 
 ```shell
-pip install novelai-python
+pip -U install novelai-python
 ```
 
 More examples can be found in the [playground](https://github.com/LlmKira/novelai-python/tree/main/playground)
@@ -86,7 +86,8 @@ from dotenv import load_dotenv
 from pydantic import SecretStr
 
 from novelai_python import APIError, LoginCredential
-from novelai_python.sdk.ai.generate import TextLLMModel, LLM, get_default_preset
+from novelai_python.sdk.ai.generate import TextLLMModel, LLM, get_default_preset, AdvanceLLMSetting
+from novelai_python.sdk.ai.generate._enum import get_model_preset
 
 load_dotenv()
 username = os.getenv("NOVELAI_USER", None)
@@ -105,8 +106,17 @@ async def chat(prompt: str):
         agent = LLM.build(
             prompt=prompt,
             model=model,
-            parameters=None  # Auto Select or get from preset
+            # parameters=None,  # Auto Select or get from preset
+            parameters=get_model_preset(TextLLMModel.ERATO).get_all_presets()[0].parameters,  # Select from enum preset
+            advanced_setting=AdvanceLLMSetting(
+                min_length=1,
+                max_length=None,  # Auto
+            )
         )
+        # NOTE:parameter > advanced_setting, which logic in generate/__init__.py
+        # If you not pass the parameter, it will use the default preset.
+        # So if you want to set the generation params, you should pass your own params.
+        # Only if you want to use some params not affect the generation, you can use advanced_setting.
         result = await agent.request(session=login_credential)
     except APIError as e:
         raise Exception(f"Error: {e.message}")
@@ -148,6 +158,14 @@ print(tokenizer_package.decode(tokenizer_package.encode(t_text)))
 
 ```
 
+## About Nsfw 🚫
+
+You might need some solutions for identifying NSFW content and adding a mosaic to prevent operational mishaps.
+
+https://dghs-imgutils.deepghs.org/main/api_doc/detect/nudenet.html
+
+https://dghs-imgutils.deepghs.org/main/api_doc/operate/censor.html
+
 ## Acknowledgements 🙏
 
 [BackEnd](https://api.novelai.net/docs)

{novelai_python-0.5.1 → novelai_python-0.5.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.5.1"
+version = "0.5.3"
 description = "NovelAI Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "coldlando@hotmail.com" },
@@ -28,6 +28,8 @@ dependencies = [
     "tokenizers>=0.15.2",
     "json-repair>=0.29.4",
     "robust-downloader>=0.0.2",
+    "bchlib>=2.1.3",
+    "arrow>=1.3.0",
 ]
 requires-python = ">=3.9"
 readme = "README.md"

{novelai_python-0.5.1 → novelai_python-0.5.3}/src/novelai_python/credential/ApiToken.py

@@ -4,8 +4,7 @@
 # @File    : ApiToken.py
 # @Software: PyCharm
 
-import datetime
-
+import arrow
 import shortuuid
 from curl_cffi.requests import AsyncSession
 from loguru import logger
@@ -31,7 +30,7 @@ class ApiCredential(CredentialBase):
             "Origin": "https://novelai.net",
             "Referer": "https://novelai.net/",
             "x-correlation-id": self._x_correlation_id,
-            "x-initiated-at": f"{datetime.datetime.now(datetime.UTC).isoformat()}Z",
+            "x-initiated-at": f"{arrow.utcnow().isoformat()}Z",
         }
 
         if update_headers:

{novelai_python-0.5.1 → novelai_python-0.5.3}/src/novelai_python/credential/JwtToken.py

@@ -3,8 +3,8 @@
 # @Author  : sudoskys
 # @File    : JwtToken.py
 # @Software: PyCharm
-import datetime
 
+import arrow
 import shortuuid
 from curl_cffi.requests import AsyncSession
 from loguru import logger
@@ -30,7 +30,7 @@ class JwtCredential(CredentialBase):
             "Origin": "https://novelai.net",
             "Referer": "https://novelai.net/",
             "x-correlation-id": self._x_correlation_id,
-            "x-initiated-at": f"{datetime.datetime.now(datetime.UTC).isoformat()}Z",
+            "x-initiated-at": f"{arrow.utcnow().isoformat()}Z",
         }
 
         if update_headers:

{novelai_python-0.5.1 → novelai_python-0.5.3}/src/novelai_python/credential/UserAuth.py

@@ -2,10 +2,10 @@
 # @Time    : 2024/2/7 下午12:14
 # @Author  : sudoskys
 # @File    : UserAuth.py
-import datetime
 import time
 from typing import Optional
 
+import arrow
 import shortuuid
 from curl_cffi.requests import AsyncSession
 from pydantic import SecretStr, Field
@@ -33,7 +33,7 @@ class LoginCredential(CredentialBase):
             "Origin": "https://novelai.net",
             "Referer": "https://novelai.net/",
             "x-correlation-id": self._x_correlation_id,
-            "x-initiated-at": f"{datetime.datetime.now(datetime.UTC).isoformat()}Z",
+            "x-initiated-at": f"{arrow.utcnow().isoformat()}Z",
         }
 
         # 30 天有效期

{novelai_python-0.5.1 → novelai_python-0.5.3}/src/novelai_python/sdk/ai/generate/__init__.py

@@ -305,7 +305,15 @@ class LLM(ApiBaseModel):
               **kwargs
               ) -> "LLM":
         """
-        Generate Prompt
+        Generate
+        :param prompt: prompt, exp "Hello, World!"
+        :param model: TextLLMModelTypeAlias
+        :param advanced_setting: AdvanceLLMSetting, default None means auto-select
+        :param parameters: LLMGenerationParams, default None means auto-select, use get_model_preset
+        :param default_bias: bool, default True
+        :param logprobs_count: int, default None
+        :param bias_dinkus_asterism: bool, default False
+        :param prefix: str, default "vanilla"
         """
         assert isinstance(prompt, str), "Prompt must be a string"
         created = cls(

{novelai_python-0.5.1 → novelai_python-0.5.3}/src/novelai_python/sdk/ai/generate/_enum.py

@@ -148,30 +148,12 @@ def get_repetition_penalty_whitelist(group: TextTokenizerGroup) -> List[str]:
     return []
 
 
-def get_default_preset(model: TextLLMModel):
-    if model == TextLLMModel.J_6B_V4:
-        return J6bv4.get_all_presets()[0]
-    elif model == TextLLMModel.NEO_2B:
-        return Neo2b.get_all_presets()[0]
-    elif model == TextLLMModel.GENJI_PYTHON_6B:
-        return Genjipython6b.get_all_presets()[0]
-    elif model == TextLLMModel.GENJI_JP_6B_V2:
-        return Genjijp6bv2.get_all_presets()[0]
-    elif model == TextLLMModel.EUTERPE_V2:
-        return Euterpev2.get_all_presets()[0]
-    elif model == TextLLMModel.KRAKE_V1:
-        return Krakev1.get_all_presets()[0]
-    elif model == TextLLMModel.KRAKE_V2:
-        return Krakev2.get_all_presets()[0]
-    elif model in COLORS_LLM:
-        return Colorful.get_all_presets()[0]
-    elif model == TextLLMModel.CLIO:
-        return Clio.get_all_presets()[0]
-    elif model == TextLLMModel.KAYRA:
-        return Kayra.get_all_presets()[0]
-    elif model == TextLLMModel.ERATO:
-        return Erato.get_all_presets()[0]
-    return Default.get_all_presets()[0]
+def get_model_preset(model: TextLLMModel) -> "ModelPresetManager":
+    return PresetMapping.get(model, Default)
+
+
+def get_default_preset(model: TextLLMModel) -> Preset:
+    return get_model_preset(model).get_all_presets()[0]
 
 
 class ModelPresetManager:
@@ -2279,3 +2261,24 @@ Erato.register_preset(
         description="Designed for writing in Japanese. Works fine for English too."
     )
 )
+
+PresetMapping = {
+    TextLLMModel.J_6B_V4: J6bv4,
+    TextLLMModel.NEO_2B: Neo2b,
+    TextLLMModel.GENJI_PYTHON_6B: Genjipython6b,
+    TextLLMModel.GENJI_JP_6B_V2: Genjijp6bv2,
+    TextLLMModel.EUTERPE_V2: Euterpev2,
+    TextLLMModel.KRAKE_V1: Krakev1,
+    TextLLMModel.KRAKE_V2: Krakev2,
+    TextLLMModel.BLUE: Colorful,
+    TextLLMModel.RED: Colorful,
+    TextLLMModel.GREEN: Colorful,
+    TextLLMModel.PURPLE: Colorful,
+    TextLLMModel.PINK: Colorful,
+    TextLLMModel.YELLOW: Colorful,
+    TextLLMModel.WHITE: Colorful,
+    TextLLMModel.BLACK: Colorful,
+    TextLLMModel.CLIO: Clio,
+    TextLLMModel.KAYRA: Kayra,
+    TextLLMModel.ERATO: Erato,
+}

novelai_python-0.5.3/src/novelai_python/tool/image_metadata/__init__.py

@@ -0,0 +1,297 @@
+# -*- coding: utf-8 -*-
+import base64
+import json
+from io import BytesIO
+from typing import Union, Optional, List, Tuple
+
+import numpy as np
+from PIL import Image
+from PIL.PngImagePlugin import PngInfo
+from loguru import logger
+from nacl.encoding import Base64Encoder
+from nacl.signing import VerifyKey
+from pydantic import BaseModel, ConfigDict
+
+from novelai_python.sdk.ai._enum import PROMOTION, Model
+from ._type import IMAGE_INPUT_TYPE, get_image_bytes
+from .bch_utils import fec_decode
+from .lsb_extractor import ImageLsbDataExtractor
+from .lsb_injector import inject_data
+
+
+class ImageMetadata(BaseModel):
+    """
+    ImageMetadata is a class that represents the metadata of a NovelAI generated image.
+    """
+
+    class CommentModel(BaseModel):
+        prompt: Optional[str] = None
+        steps: Optional[int] = None
+        height: Optional[int] = None
+        width: Optional[int] = None
+        scale: Optional[float] = None
+        uncond_scale: Optional[float] = None
+        cfg_rescale: Optional[float] = None
+        seed: Optional[int] = None
+        n_samples: Optional[int] = None
+        hide_debug_overlay: Optional[bool] = None
+        noise_schedule: Optional[str] = None
+        legacy_v3_extend: Optional[bool] = None
+        reference_information_extracted: Optional[float] = None
+        reference_strength: Optional[float] = None
+        reference_strength_multiple: Optional[List[float]] = None
+        reference_information_extracted_multiple: Optional[List[float]] = None
+        sampler: Optional[str] = None
+        controlnet_strength: Optional[float] = None
+        controlnet_model: Optional[Union[None, str]] = None
+        dynamic_thresholding: Optional[bool] = None
+        dynamic_thresholding_percentile: Optional[float] = None
+        dynamic_thresholding_mimic_scale: Optional[float] = None
+        sm: Optional[bool] = None
+        sm_dyn: Optional[bool] = None
+        skip_cfg_below_sigma: Optional[float] = None
+        lora_unet_weights: Optional[Union[None, str]] = None
+        lora_clip_weights: Optional[Union[None, str]] = None
+        uc: Optional[str] = None
+        request_type: Optional[str] = None
+        signed_hash: Optional[str] = None
+        skip_cfg_above_sigma: Optional[float] = None
+        deliberate_euler_ancestral_bug: Optional[bool] = None
+        prefer_brownian: Optional[bool] = None
+        cfg_sched_eligibility: Optional[str] = None
+        explike_fine_detail: Optional[bool] = None
+        minimize_sigma_inf: Optional[bool] = None
+        uncond_per_vibe: Optional[bool] = None
+        wonky_vibe_correlation: Optional[bool] = None
+        version: Optional[int] = None
+
+        model_config = ConfigDict(extra="allow")
+
+        @property
+        def generate_method(self):
+            return self.request_type or "Unknown"
+
+        @property
+        def negative_prompt(self):
+            return self.uc
+
+        @property
+        def vibe_transfer_strength(self) -> List[float]:
+            """
+            Get the vibe transfer strength totally
+            """
+            if self.reference_strength_multiple:
+                return self.reference_strength_multiple
+            reference_strength = [] if self.reference_strength is None else [self.reference_strength]
+            return reference_strength
+
+        @property
+        def vibe_transfer_information(self) -> List[float]:
+            """
+            Get the vibe transfer information totally
+            """
+            if self.reference_information_extracted_multiple:
+                return self.reference_information_extracted_multiple
+            reference_information = [] if self.reference_information_extracted is None else [
+                self.reference_information_extracted]
+            return reference_information
+
+    Title: str = "AI generated image"
+    Software: str = "NovelAI"
+    Source: str = None
+    Description: str
+    Comment: CommentModel
+
+    model_config = ConfigDict(extra="allow")
+
+    @property
+    def used_model(self) -> Union[Model, None]:
+        """
+        Get the model used to generate the image from the Model code
+        :return:  Model or None
+        """
+        return PROMOTION.get(self.Source, None)
+
+    @staticmethod
+    def reset_alpha(image: IMAGE_INPUT_TYPE) -> BytesIO:
+        """
+        Remove LSB from the image, set the alpha channel to 254
+        :param image: Type: Union[str, bytes, Path, BytesIO]
+        :return: BytesIO
+        """
+        image_data = get_image_bytes(image)
+        image = Image.open(image_data).convert('RGBA')
+        data = np.array(image)
+        data[..., 3] = 254
+        new_image = Image.fromarray(data, 'RGBA')
+        _new_img_io = BytesIO()
+        new_image.save(_new_img_io, format="PNG")
+        return _new_img_io
+
+    def apply_to_image(self, image: IMAGE_INPUT_TYPE, *, inject_lsb: bool = True) -> BytesIO:
+        """
+        Write metadata to origin_image
+        If you set inject_lsb to True, the image will be injected with metadata using LSB.
+
+        **But if you set inject_lsb to False, the image will be reset to the 254 alpha channel**
+        :param image:  Union[str, bytes, Path, BytesIO]
+        :param inject_lsb:  Inject metadata using LSB
+        :return:  BytesIO
+        """
+        image_data = get_image_bytes(image)
+        origin_image = Image.open(image_data)
+        metadata = PngInfo()
+        for k, v in self.model_dump(mode="json").items():
+            if isinstance(v, dict):
+                v = json.dumps(v)
+            metadata.add_text(k, v)
+        if inject_lsb:
+            # Inject metadata using LSB
+            origin_image = inject_data(origin_image, metadata)
+
+        # Prepare image to be saved
+        be_copy_image = BytesIO()
+        origin_image.save(be_copy_image, format="PNG", pnginfo=metadata, optimize=False, compress_level=0)
+        return be_copy_image
+
+    @classmethod
+    def _extract_metadata_from_lsb(cls, image: IMAGE_INPUT_TYPE):
+        """
+        Extract metadata using LSB extraction method.
+        :param image: IMAGE_INPUT_TYPE
+        """
+        try:
+            image_meta_data, fec_data = ImageLsbDataExtractor().extract_data(image)
+            return cls.model_validate(image_meta_data)
+        except Exception as e:
+            logger.trace(f"Error extracting data in LSB: {e}")
+            return None
+
+    @classmethod
+    def _extract_metadata_from_comments(cls, image: IMAGE_INPUT_TYPE):
+        """
+        Extract metadata from image comments and other info.
+        :param image: Union[str, bytes, Path, BytesIO]
+        """
+        with Image.open(image) as img:
+            title = img.info.get("Title", None)
+            description = img.info.get("Description", None)
+            comment = img.info.get("Comment", "{}")
+            try:
+                comment_dict = json.loads(comment)
+            except Exception as e:
+                logger.trace(f"Error loading comment: {e}")
+                comment_dict = {}
+        try:
+            comment_dict.setdefault("prompt", description)
+            comment_model = cls.CommentModel.model_validate(comment_dict)
+            return cls(Title=title, Description=description, Comment=comment_model)
+        except Exception as e:
+            logger.debug(f"Error loading comment: {e}")
+            return None
+
+    @classmethod
+    def load_image(cls, image: IMAGE_INPUT_TYPE):
+        """
+        Load image and extract metadata using LSB/Metadata
+        :param image: str, bytes, Path, BytesIO
+        :return: ImageMetadata
+        :raises ValueError: Data extraction failed
+        """
+        image_data = get_image_bytes(image)
+        metadata = cls._extract_metadata_from_lsb(image_data)
+        if metadata:
+            return metadata
+        metadata = cls._extract_metadata_from_comments(image)
+        if metadata:
+            return metadata
+        raise ValueError("No metadata found")
+
+
+class ImageVerifier:
+
+    @staticmethod
+    def verify_latents(image: Image.Image, signed_hash: bytes, verify_key: VerifyKey):
+        image.load()
+        sig = None
+        latents = None
+        try:
+            for cid, data, after_idat in image.private_chunks:
+                if after_idat:
+                    if cid == b'ltns':
+                        sig = data
+                    elif cid == b'ltnt':
+                        latents = data
+        except Exception as e:
+            logger.trace(f"Error extracting latents: {e}")
+            return True, False, None
+        if sig is None or latents is None:
+            return True, False, None
+        if not sig.startswith(b'NovelAI_ltntsig'):
+            return False, False, None
+        sig = sig[len(b'NovelAI_ltntsig'):]
+        if not latents.startswith(b'NovelAI_latents'):
+            return False, False, None
+        latents = latents[len(b'NovelAI_latents'):]
+        if len(sig) != 64:
+            return False, False, None
+        w, h = image.size
+        base_size = (w // 8) * (h // 8) * 4
+        if len(latents) != base_size * 4 and len(latents) != base_size * 2:
+            return False, False, None
+        float_dim = 4 if len(latents) == base_size * 4 else 2
+        try:
+            verify_key.verify(signed_hash + latents, sig)
+            return True, True, (float_dim, latents)
+        except Exception as e:
+            logger.trace(f"Error verifying latents: {e}")
+            return False, False, None
+
+    def verify(self,
+               image,
+               verify_key_hex: str = "Y2JcQAOhLwzwSDUJPNgL04nS0Tbqm7cSRc4xk0vRMic="
+               ) -> Tuple[bool, bool]:
+        """
+        Verify if the image is a NovelAI generated image
+
+        :param image: Union[str, bytes, Path, BytesIO] - The input image to verify.
+        :param verify_key_hex: str - The verification key in base64 format.
+        :return: Tuple[bool, bool] - The first bool indicates if the image is verified, the second bool indicates if the image has latents.
+        :raises ValueError: If the required metadata or signed hash is missing.
+        """
+        image_obj = Image.open(get_image_bytes(image))
+        w, h = image_obj.size
+        image_array = np.array(image_obj)
+        try:
+            metadata, fec_data = ImageLsbDataExtractor().extract_data(image, get_fec=True)
+        except Exception as e:
+            logger.trace(f"Error extracting data in LSB: {e}")
+            metadata = None
+            fec_data = None
+        if not metadata or not metadata.get("Comment") or not metadata["Comment"].get("signed_hash"):
+            raise ValueError(
+                "Bad image lsb or metadata. Comment or Comment.signed_hash is missing, the image is be tampered or not generated by NovelAI."
+            )
+        parameter = metadata["Comment"]
+        signed_hash = base64.b64decode(parameter.pop("signed_hash").encode("utf-8"))
+        # Build verify key
+        verify_key = VerifyKey(verify_key_hex.encode("utf-8"), encoder=Base64Encoder)
+        # Verify latents
+        good_latents, have_latents, latents = self.verify_latents(image_obj, signed_hash, verify_key)
+        if not good_latents:
+            return False, False
+        rgb = image_array[:, :, :3].tobytes()
+        parameter = json.dumps(parameter).encode("utf-8")
+        image_and_comment = rgb + parameter
+        try:
+            verify_key.verify(image_and_comment, signed_hash)
+        except Exception as e:
+            logger.trace(f"Error verifying image [1]: {e}")
+            try:
+                rgb, errs = fec_decode(bytearray(rgb), bytearray(fec_data), w, h)
+                image_and_comment = rgb + parameter
+                verify_key.verify(image_and_comment, signed_hash)
+            except Exception as e:
+                logger.trace(f"Error verifying image [2]: {e}")
+                return False, False
+        return True, have_latents

novelai_python-0.5.3/src/novelai_python/tool/image_metadata/_type.py

@@ -0,0 +1,21 @@
+from io import BytesIO
+from pathlib import Path
+from typing import Union
+
+IMAGE_INPUT_TYPE = Union[str, bytes, Path, BytesIO]
+
+
+def get_image_bytes(image: IMAGE_INPUT_TYPE) -> BytesIO:
+    if isinstance(image, (str, Path)):
+        try:
+            with open(image, "rb") as f:
+                return BytesIO(f.read())
+        except FileNotFoundError:
+            raise FileNotFoundError(f"Image not found: {image}")
+    elif isinstance(image, BytesIO):
+        image.seek(0)
+        return image
+    elif isinstance(image, bytes):
+        return BytesIO(image)
+    else:
+        raise TypeError(f"Invalid image type: {type(image)}, only {IMAGE_INPUT_TYPE} is supported")