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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.5.1
+Version: 0.5.2
 Summary: NovelAI Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache License 2.0
@@ -29,6 +29,7 @@ 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
 Provides-Extra: testing
 Requires-Dist: pytest>=7.4.4; extra == "testing"
 Requires-Dist: pytest-asyncio>=0.23.4; extra == "testing"
@@ -75,7 +76,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 +123,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 +143,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 +195,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.2}/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.2}/pyproject.toml

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

{novelai_python-0.5.1 → novelai_python-0.5.2}/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.2}/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.2/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.2/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")

novelai_python-0.5.2/src/novelai_python/tool/image_metadata/bch_utils.py

@@ -0,0 +1,160 @@
+# 纠错码
+# MIT: https://github.com/NovelAI/novelai-image-metadata/blob/main/nai_bch.py
+import functools
+
+import bchlib
+import numpy as np
+
+correctable_bits = 16
+block_length = 2019
+code_block_len = 1920
+
+
+def bit_shuffle(data_bytes, w, h, use_bytes=False):
+    bits = np.frombuffer(data_bytes, dtype=np.uint8)
+    bit_fac = 8
+    if use_bytes:
+        bit_fac = 1
+    else:
+        bits = np.unpackbits(bits)
+    bits = bits.reshape((h, w, 3 * bit_fac))
+    code_block_len = 1920
+    flat_tile_len = (w * h * 3) // code_block_len
+    tile_w = 32
+    if flat_tile_len // tile_w > 100:
+        tile_w = 64
+    tile_h = flat_tile_len // tile_w
+    h_cutoff = (h // tile_h) * tile_h
+    tile_hr = h - h_cutoff
+    easy_tiles = bits[:h_cutoff].reshape(h_cutoff // tile_h, tile_h, w // tile_w, tile_w, 3 * bit_fac)
+    easy_tiles = easy_tiles.swapaxes(1, 2)
+    easy_tiles = easy_tiles.reshape(-1, tile_h * tile_w)
+    easy_tiles = easy_tiles.T
+    rest_tiles = bits[h_cutoff:]
+    rest_tiles = rest_tiles.reshape(tile_hr, 1, w // tile_w, tile_w, 3 * bit_fac)
+    rest_tiles = rest_tiles.swapaxes(1, 2)
+    rest_tiles = rest_tiles.reshape(-1, tile_hr * tile_w)
+    rest_tiles = rest_tiles.T
+    rest_dim = rest_tiles.shape[-1]
+    rest_tiles = np.pad(rest_tiles, ((0, 0), (0, easy_tiles.shape[-1] - rest_tiles.shape[-1])), mode='constant',
+                        constant_values=0)
+    bits = np.concatenate((easy_tiles, rest_tiles), axis=0)
+    dim = bits.shape[-1]
+    bits = bits.reshape((-1,))
+    if not use_bytes:
+        bits = np.packbits(bits)
+    return bytearray(bits.tobytes()), dim, rest_tiles.shape[0], rest_dim
+
+
+def bit_unshuffle(data_bytes, w, h, dim, rest_size, rest_dim, use_bytes=False):
+    bits = np.frombuffer(data_bytes, dtype=np.uint8)
+    bit_fac = 8
+    if use_bytes:
+        bit_fac = 1
+    else:
+        bits = np.unpackbits(bits)
+    code_block_len = 1920
+    flat_tile_len = (w * h * 3) // code_block_len
+    tile_w = 32
+    if flat_tile_len // tile_w > 100:
+        tile_w = 64
+    tile_h = flat_tile_len // tile_w
+    h_cutoff = (h // tile_h) * tile_h
+    tile_hr = h - h_cutoff
+    bits = bits.reshape((-1, dim))
+    rev_cutoff = bits.shape[0] - rest_size
+    rest_tiles = bits[rev_cutoff:]
+    rest_tiles = rest_tiles.reshape((-1, dim))
+    rest_tiles = rest_tiles[:, :rest_dim]
+    rest_tiles = rest_tiles.T
+    rest_tiles = rest_tiles.reshape((tile_hr, w // tile_w, 1, tile_w, 3 * bit_fac))
+    rest_tiles = rest_tiles.swapaxes(1, 2)
+    rest_tiles = rest_tiles.reshape((-1,))
+    easy_tiles = bits[:rev_cutoff]
+    easy_tiles = easy_tiles.T
+    easy_tiles = easy_tiles.reshape((h_cutoff // tile_h, w // tile_w, tile_h, tile_w, 3 * bit_fac))
+    easy_tiles = easy_tiles.swapaxes(1, 2)
+    easy_tiles = easy_tiles.reshape((-1,))
+    data_bytes = np.concatenate((easy_tiles, rest_tiles), axis=0)
+    if not use_bytes:
+        data_bytes = np.packbits(data_bytes)
+    data_bytes = data_bytes.tobytes()
+    return bytearray(data_bytes)
+
+
+@functools.lru_cache(maxsize=512)
+def get_indices(len_db):
+    indices = np.arange(0, len_db)
+    rng = np.random.Generator(np.random.MT19937(31337))
+    indices = rng.permutation(indices)
+    unshuffled = indices.copy()
+    unshuffled[indices] = np.arange(0, len_db)
+    return indices, unshuffled
+
+
+def rand_byte_shuffle(data_bytes):
+    indices, _ = get_indices(len(data_bytes))
+    data_bytes = np.frombuffer(data_bytes, dtype=np.uint8)
+    data_bytes = bytearray(data_bytes[indices].tobytes())
+    return data_bytes
+
+
+def rand_byte_unshuffle(data_bytes):
+    _, indices = get_indices(len(data_bytes))
+    data_bytes = np.frombuffer(data_bytes, dtype=np.uint8)
+    data_bytes = bytearray(data_bytes[indices].tobytes())
+    return data_bytes
+
+
+use_bytes = True  # Bit shuffling provides better resilience, but byte shuffling is much faster and still sufficient
+shuffle_fn = lambda data_bytes, w, h: bit_shuffle(data_bytes, w, h, use_bytes=use_bytes)
+unshuffle_fn = lambda data_bytes, w, h, dim, rest_size, rest_dim: bit_unshuffle(data_bytes, w, h, dim, rest_size,
+                                                                                rest_dim, use_bytes=use_bytes)
+
+
+def split_byte_ranges(data_bytes, n, w, h, shuffle=False):
+    if shuffle:
+        data_bytes, dim, rest_size, rest_dim = shuffle_fn(data_bytes.copy(), w, h)
+    chunks = []
+    for i in range(0, len(data_bytes), n):
+        chunks.append(data_bytes[i:i + n])
+    if shuffle:
+        return chunks, dim, rest_size, rest_dim
+    return chunks
+
+
+def pad(data_bytes):
+    return bytearray(data_bytes + b'\x00' * (2019 - len(data_bytes)))
+
+
+# Returns codes for the data in data_bytes
+def fec_encode(data_bytes, w, h):
+    encoder = bchlib.BCH(16, prim_poly=17475)
+    # import galois
+    # encoder = galois.BCH(16383, 16383-224, d=17, c=224)
+    chunks = [bytearray(encoder.encode(pad(x))) for x in split_byte_ranges(data_bytes, 2019, w, h, shuffle=True)[0]]
+    return b''.join(chunks)
+
+
+# Returns the error corrected data and number of corrected errors or None, None if errors are not correctable
+def fec_decode(data_bytes, codes, w, h):
+    encoder = bchlib.BCH(16, prim_poly=17475)
+    chunks, dim, rest_size, rest_dim = split_byte_ranges(data_bytes, 2019, w, h, shuffle=True)
+    codes = split_byte_ranges(codes, 28, w, h)
+    corrected = []
+    total_errors = 0
+    for i, chunk in enumerate(chunks):
+        c_len = len(chunk)
+        chunk = pad(chunk)
+        code = bytearray(codes[i])
+        n_err = encoder.decode(chunk, code)
+        if n_err > 0:
+            total_errors += n_err
+            encoder.correct(chunk, code)
+        # Ignoring the following case, since it might be caused by corrupted codes which might not even be needed
+        # if n_err < 0:
+        #    raise ValueError("Too many errors to correct")
+        corrected.append(chunk[:c_len])
+    corrected = b''.join(corrected)
+    corrected = unshuffle_fn(corrected, w, h, dim, rest_size, rest_dim)
+    return corrected, total_errors

{novelai_python-0.5.1 → novelai_python-0.5.2}/src/novelai_python/tool/image_metadata/lsb_extractor.py

@@ -1,12 +1,12 @@
 # MIT: https://github.com/NovelAI/novelai-image-metadata/blob/main/nai_meta.py
+import gzip
+import json
 from io import BytesIO
 from pathlib import Path
 from typing import Union
 
-from PIL import Image
 import numpy as np
-import gzip
-import json
+from PIL import Image
 
 
 class LSBExtractor(object):
@@ -59,42 +59,45 @@ class ImageLsbDataExtractor(object):
     def __init__(self):
         self.magic = "stealth_pngcomp"
 
-    def extract_data(self, image: Union[str, bytes, Path, BytesIO, np.ndarray]) -> dict:
-        try:
-            if isinstance(image, Image.Image):
-                image = np.array(image)
-            elif isinstance(image, (str, bytes, Path, BytesIO)):
-                img = Image.open(image)
-                image = np.array(img)
+    def extract_data(self, image: Union[BytesIO, bytes, Path], get_fec: bool = False) -> tuple:
+        if isinstance(image, (bytes, BytesIO, Path)):
+            img = Image.open(image)
+            img.convert("RGBA")
+            image = np.array(img)
+        else:
+            raise TypeError(f"Invalid image type: {type(image)}, only {type(BytesIO)} is supported")
 
+        try:
             if not (image.shape[-1] == 4 and len(image.shape) == 3):
-                raise AssertionError('image format error')
-
+                raise AssertionError('image format error, maybe image already be modified')
             reader = LSBExtractor(image)
-
             read_magic = reader.get_next_n_bytes(len(self.magic)).decode("utf-8")
             if not (self.magic == read_magic):
                 raise AssertionError('magic number mismatch')
-
             read_len = reader.read_32bit_integer() // 8
             json_data = reader.get_next_n_bytes(read_len)
-
             json_data = json.loads(gzip.decompress(json_data).decode("utf-8"))
             if "Comment" in json_data:
                 json_data["Comment"] = json.loads(json_data["Comment"])
-            return json_data
+
+            if not get_fec:
+                return json_data, None
+
+            fec_len = reader.read_32bit_integer()
+            fec_data = None
+            if fec_len != 0xffffffff:
+                fec_data = reader.get_next_n_bytes(fec_len // 8)
+
+            return json_data, fec_data
         except FileNotFoundError:
             # 无法找到文件
             raise Exception(f"The file {image} does not exist.")
-
         except json.JSONDecodeError as e:
             # 无法解析JSON数据
             raise Exception(f"Failed to decode JSON data from image: {image}. Error: {str(e)}")
-
         except AssertionError as err:
             # 魔数不匹配
             raise Exception(f"Failed to extract data from image: {image}. Error: {str(err)}")
-
         except Exception as e:
             # 从图像中提取数据时发生意外错误
             raise Exception(f"Unexpected error happened when extracting data from image: {image}. Error: {str(e)}")

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

@@ -1,225 +0,0 @@
-# -*- coding: utf-8 -*-
-import base64
-import json
-from io import BytesIO
-from pathlib import Path
-from typing import Union, Optional, List
-
-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 .lsb_extractor import ImageLsbDataExtractor
-from .lsb_injector import inject_data
-from novelai_python.sdk.ai._enum import PROMOTION, Model
-
-
-class ImageMetadata(BaseModel):
-    """
-    ImageMetadata is a class that represents the metadata of a NovelAI generated image.
-    """
-
-    class CommentModel(BaseModel):
-        prompt: str
-        steps: int = None
-        height: int = None
-        width: int = None
-        scale: float = None
-        uncond_scale: float = None
-        cfg_rescale: float = None
-        seed: int = None
-        n_samples: int = None
-        hide_debug_overlay: bool = None
-        noise_schedule: str = None
-        legacy_v3_extend: 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: str = None
-        controlnet_strength: float = None
-        controlnet_model: Union[None, str] = None
-        dynamic_thresholding: bool = None
-        dynamic_thresholding_percentile: float = None
-        dynamic_thresholding_mimic_scale: float = None
-        sm: bool = None
-        sm_dyn: bool = None
-        skip_cfg_below_sigma: float = None
-        lora_unet_weights: Union[None, str] = None
-        lora_clip_weights: Union[None, str] = None
-        uc: str = None
-        request_type: str = None
-        signed_hash: str = None
-        model_config = ConfigDict(extra="allow")
-
-        @property
-        def negative_prompt(self):
-            return self.uc
-
-        @property
-        def vibe_transfer_strength(self) -> List[float]:
-            """
-            Get the vibe transfer strength totally
-            :return: List[float]
-            """
-            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
-            :return: List[float]
-            """
-            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
-    """
-    AI generated image
-    silvery white, best quality, amazing quality, very aesthetic, absurdres
-    {'prompt': 'silvery white, best quality, amazing quality, very aesthetic, absurdres', 'steps': 28, 'height': 128, 'width': 128, 'scale': 10.0, 'uncond_scale': 1.0, 'cfg_rescale': 0.38, 'seed': 1148692016, 'n_samples': 1, 'hide_debug_overlay': False, 'noise_schedule': 'native', 'legacy_v3_extend': False, 'reference_information_extracted': 0.87, 'reference_strength': 1.0, 'sampler': 'k_euler', 'controlnet_strength': 1.0, 'controlnet_model': None, 'dynamic_thresholding': False, 'dynamic_thresholding_percentile': 0.999, 'dynamic_thresholding_mimic_scale': 10.0, 'sm': True, 'sm_dyn': True, 'skip_cfg_below_sigma': 0.0, 'lora_unet_weights': None, 'lora_clip_weights': None, 'uc': '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], small', 'request_type': 'PromptGenerateRequest', 'signed_hash': 'GmLfTgQTOI4WRcTloiSmenArH/Y8szJdVJ1DW18yeqYPkqw3o1zellB4xCQFYle+0tuF9KVP8RkHiDNbQXr0DA=='}
-    """
-    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(input_img: BytesIO) -> BytesIO:
-        """
-        Remove LSB from the image, set the alpha channel to 254
-        :param input_img:
-        :return:
-        """
-        if isinstance(input_img, BytesIO):
-            input_img.seek(0)
-        image = Image.open(input_img).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,
-                       origin_image: Image.Image,
-                       *,
-                       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 origin_image:  BytesIO
-        :param inject_lsb:  Inject metadata using LSB
-        :return:  BytesIO
-        """
-        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)
-        # Save original image with metadata (and LSB if inject_lsb is True)
-        new_img = BytesIO()
-        origin_image.save(new_img, format="PNG", pnginfo=metadata, optimize=False, compress_level=0)
-        return new_img
-
-    @classmethod
-    def load_image(cls,
-                   image_io: Union[str, bytes, Path, BytesIO]
-                   ):
-        """
-        Load image and extract metadata using LSB/Metadata
-        :param image_io: str, bytes, Path, BytesIO
-        :return: ImageMetadata
-        :raises ValueError:  Data extraction failed
-        """
-        if isinstance(image_io, BytesIO):
-            image_io.seek(0)
-        try:
-            image_data = ImageLsbDataExtractor().extract_data(image_io)
-            model = cls.model_validate(image_data)
-        except Exception as e:
-            logger.trace(f"Error trying extracting data in LSB: {e}")
-        else:
-            return model
-        with Image.open(image_io) as img:
-            title = img.info.get("Title", None)
-            prompt = img.info.get("Description", None)
-            comment = img.info.get("Comment", None)
-            try:
-                assert isinstance(comment, str), ValueError("Comment Empty")
-                comment = json.loads(comment)
-            except Exception as e:
-                logger.trace(f"Error loading comment: {e}")
-                comment = {}
-        if comment.get("prompt", None) is None:
-            comment["prompt"] = prompt
-        try:
-            comment_model = cls.CommentModel.model_validate(comment)
-            return cls(Title=title, Description=prompt, Comment=comment_model)
-        except Exception as e:
-            logger.debug(f"Error loading comment: {e}")
-            raise ValueError("Data extraction failed")
-
-    @staticmethod
-    def verify_image_is_novelai(
-            image: Union[Image.Image, np.ndarray],
-            verify_key_hex: str = "Y2JcQAOhLwzwSDUJPNgL04nS0Tbqm7cSRc4xk0vRMic="
-    ) -> bool:
-        """
-        Verify if the image is a NovelAI generated image
-        :param image:  Image.Image or np.ndarray
-        :param verify_key_hex:
-        :return:  bool
-        :raises RuntimeError:  No metadata found in image
-        :raises RuntimeError:  Comment not in metadata
-        :raises RuntimeError:  signed_hash not in comment
-        """
-        # MIT:https://github.com/NovelAI/novelai-image-metadata/blob/main/nai_sig.py
-        if isinstance(image, Image.Image):
-            image = np.array(image)
-        metadata = ImageLsbDataExtractor().extract_data(image)
-
-        if metadata is None:
-            raise RuntimeError("No metadata found in image")
-        if "Comment" not in metadata:
-            raise RuntimeError("Comment not in metadata")
-        comment = metadata["Comment"]
-        if "signed_hash" not in comment:
-            raise RuntimeError("signed_hash not in comment")
-        signed_hash = comment["signed_hash"].encode("utf-8")
-        signed_hash = base64.b64decode(signed_hash)
-        del comment["signed_hash"]
-        verify_key_hex = verify_key_hex.encode("utf-8")
-        verify_key = VerifyKey(verify_key_hex, encoder=Base64Encoder)
-        image_and_comment = image[:, :, :3].tobytes() + json.dumps(comment).encode("utf-8")
-        try:
-            verify_key.verify(image_and_comment, signed_hash)
-        except Exception as e:
-            logger.trace(e)
-            return False
-        return True