Merge pull request #3 from LlmKira/dev

feat(optional):LSB

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.1.6"
+version = "0.1.7"
 description = "Novelai Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "[email protected]" },
@@ -16,6 +16,7 @@ dependencies = [
     "curl-cffi>=0.5.10",
     "fastapi>=0.109.0",
     "uvicorn[standard]>=0.27.0.post1",
+    "numpy>=1.24.4",
 ]
 requires-python = ">=3.8"
 readme = "README.md"

src/novelai_python/sdk/ai/generate_image.py

@@ -12,6 +12,7 @@
 import httpx
 from curl_cffi.requests import AsyncSession
 from loguru import logger
+from novelai_python.utils import NovelAiMetadata
 from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, model_validator
 
 from ..._exceptions import APIError, AuthError
@@ -229,7 +230,15 @@ def build(cls,
             parameters=cls.Params(**param)
         )
 
-    async def generate(self, session: Union[AsyncSession, JwtCredential]) -> ImageGenerateResp:
+    async def generate(self, session: Union[AsyncSession, JwtCredential],
+                       *,
+                       remove_sign: bool = False) -> ImageGenerateResp:
+        """
+        生成图片
+        :param session:  session
+        :param remove_sign:  移除追踪信息
+        :return:
+        """
         if isinstance(session, JwtCredential):
             session = session.session
         request_data = self.model_dump(exclude_none=True)
@@ -274,7 +283,12 @@ async def generate(self, session: Union[AsyncSession, JwtCredential]) -> ImageGe
                         response=try_jsonfy(response.content)
                     )
                 for filename in file_list:
-                    unzip_content.append((filename, zip_file.read(filename)))
+                    data = zip_file.read(filename)
+                    if remove_sign:
+                        data = NovelAiMetadata.rehash(BytesIO(data), remove_stealth=True)
+                        if not isinstance(data, bytes):
+                            data = data.getvalue()
+                    unzip_content.append((filename, data))
             return ImageGenerateResp(
                 meta=ImageGenerateResp.RequestParams(
                     endpoint=self.base_url,

src/novelai_python/utils/__init__.py

@@ -12,13 +12,16 @@
 from .hash import NovelAiMetadata
 
 
-def png_info_reset(img: Union[str, io.BytesIO]):
+def png_info_reset(img: Union[str, io.BytesIO],
+                   *,
+                   remove_stealth: bool = False):
     """
     reset png info hash value
+    :param remove_stealth: LCB
     :param img: BytesIO 对象
     :return: 
     """
-    _fixed = NovelAiMetadata.rehash(img_io=img)
+    _fixed = NovelAiMetadata.rehash(img_io=img, remove_stealth=remove_stealth)
     return _fixed
 
 

src/novelai_python/utils/hash.py

@@ -10,6 +10,7 @@
 from io import BytesIO
 from typing import Union
 
+import numpy as np
 from PIL import Image
 from PIL.PngImagePlugin import PngInfo
 from loguru import logger
@@ -35,16 +36,33 @@ def metadata(self):
         return {"Title": self.title, "Description": self.description, "Comment": self.comment}
 
     @staticmethod
-    def rehash(img_io):
-        cls = NaiPic.read_from_img(img_io)
+    def rehash(img_io, remove_stealth: bool = False):
+        cls = NovelAiMetadata.build_from_img(img_io)
         cls.comment["signed_hash"] = sign_message(json.dumps(cls.description), "novalai-client")
-        cls.write_out(img_io=img_io)
-        return img_io
+        _new_img_io = cls.write_out(img_io=img_io, remove_stealth=remove_stealth)
+        return _new_img_io
 
-    def write_out(self, img_io: BytesIO):
+    @staticmethod
+    def remove_stealth_info(img_io):
+        image = Image.open(img_io).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 write_out(self, img_io: BytesIO, *, remove_stealth: bool = False):
         with Image.open(img_io) as img:
+            if remove_stealth:
+                img = img.convert('RGBA')
+                data = np.array(img)
+                data[..., 3] = 254
+                img = Image.fromarray(data, 'RGBA')
             metadata = PngInfo()
             for k, v in self.metadata.items():
+                if isinstance(v, dict):
+                    v = json.dumps(v)
                 metadata.add_text(k, v)
             _new_img = BytesIO()
             img.save(_new_img, format="PNG", pnginfo=metadata, quality=95, optimize=False)
@@ -88,3 +106,23 @@ def decode_base64(encoded_data):
     # 使用base64解码
     decoded_data = base64.b64decode(byte_data)
     return decoded_data.decode("UTF-8")  # 再次使用UTF-8将字节对象转为字符串
+
+
+def decode_image(img: Union[str, BytesIO], output_path: str):
+    encoded_image = Image.open(img)
+    red_channel = encoded_image.split()[3]
+    x_size = encoded_image.size[0]
+    y_size = encoded_image.size[1]
+    decoded_image = Image.new("RGB", encoded_image.size)
+    pixels = decoded_image.load()
+    for i in range(x_size):
+        for j in range(y_size):
+            if i < 7:  # the left 7 columns
+                r = red_channel.getpixel((i, j))
+                if r > 254:  # above the threshold
+                    pixels[i, j] = (0, 0, 0)  # black
+                else:
+                    pixels[i, j] = (255, 255, 255)  # white
+            else:
+                pixels[i, j] = (255, 255, 255)  # default to white for right part of image
+    decoded_image.save(output_path)