Merge pull request #39 from LlmKira/dev

✨ feat: add apply_lsb script, update enum with new model, and fix lsb…

.gitignore

@@ -175,3 +175,4 @@ cython_debug/
 /playground/newtag/
 /playground/oldtag/
 /playground/art_assert/
+/playground/unpack/

playground/image_metadata/apply_lsb.py

@@ -0,0 +1,21 @@
+from pathlib import Path
+
+from novelai_python.tool.image_metadata import ImageMetadata, ImageLsbDataExtractor
+from PIL import Image
+
+image = Path(__file__).parent.joinpath("sample-0316.png")
+write_out = Path(__file__).parent.joinpath("sample-0316-out.png")
+try:
+    meta = ImageMetadata.load_image(image)
+    with Image.open(image) as img:
+        new_io = meta.apply_to_image(img, inject_lsb=True)
+        with open(write_out, 'wb') as f:
+            f.write(new_io.getvalue())
+    new_meta = ImageMetadata.load_image(write_out)
+    data = ImageLsbDataExtractor().extract_data(write_out)
+except ValueError:
+    raise LookupError("Cant find a MetaData")
+print(data)
+print(new_meta)
+print(new_meta.used_model)
+print(ImageMetadata.verify_image_is_novelai(Image.open(write_out)))

playground/image_metadata/read_image.py

@@ -10,6 +10,8 @@
     title = img.info.get("Title", None)
     prompt = img.info.get("Description", None)
     comment = img.info.get("Comment", None)
+    print(img.info)
+
     assert isinstance(comment, str), ValueError("Comment Empty")
     try:
         comment = json.loads(comment)

playground/image_metadata/read_lsb.py

@@ -1,11 +1,11 @@
 import os
 from pathlib import Path
 
-from novelai_python.tool.image_metadata.lsb_extractor import ImageDataExtractor
+from novelai_python.tool.image_metadata.lsb_extractor import ImageLsbDataExtractor
 
 image = Path(__file__).parent.joinpath("sample-0316.png")
 try:
-    data = ImageDataExtractor().extract_data(image)
+    data = ImageLsbDataExtractor().extract_data(image)
 except ValueError:
     raise LookupError("Cant find a MetaData")
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.4.0"
+version = "0.4.1"
 description = "NovelAI Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "[email protected]" },

src/novelai_python/sdk/ai/generate_image/_enum.py

@@ -70,6 +70,7 @@ class Model(Enum):
     NAI_DIFFUSION_3_INPAINTING = "nai-diffusion-3-inpainting"
 
     NAI_DIFFUSION = "nai-diffusion"
+    NAI_DIFFUSION_2 = "nai-diffusion-2"
     NAI_DIFFUSION_INPAINTING = "nai-diffusion-inpainting"
 
     SAFE_DIFFUSION = "safe-diffusion"
@@ -84,3 +85,19 @@ class Model(Enum):
                          Model.SAFE_DIFFUSION_INPAINTING,
                          Model.FURRY_DIFFUSION_INPAINTING
                          ]
+
+PROMOTION = {
+    "Stable Diffusion 1D44365E": Model.SAFE_DIFFUSION,
+    "Stable Diffusion F4D50568": Model.SAFE_DIFFUSION,
+    "Stable Diffusion 81274D13": Model.NAI_DIFFUSION,
+    "Stable Diffusion 3B3287AF": Model.NAI_DIFFUSION,
+    "Stable Diffusion 4CC42576": Model.NAI_DIFFUSION_FURRY,
+    "Stable Diffusion 1D09C008": Model.NAI_DIFFUSION_FURRY,
+    "Stable Diffusion 1D09D794": Model.NAI_DIFFUSION_FURRY,
+    "Stable Diffusion F64BA557": Model.NAI_DIFFUSION_FURRY,
+    "Stable Diffusion 49BFAF6A": Model.NAI_DIFFUSION_2,
+    "Stable Diffusion F1022D28": Model.NAI_DIFFUSION_2,
+    "Stable Diffusion XL B0BDF6C1": Model.NAI_DIFFUSION_3,
+    "Stable Diffusion XL C1E1DE52": Model.NAI_DIFFUSION_3,
+    "Stable Diffusion XL 8BA2AF87": Model.NAI_DIFFUSION_3,
+}

src/novelai_python/tool/image_metadata/__init__.py

@@ -14,6 +14,8 @@
 from pydantic import BaseModel, ConfigDict
 
 from .lsb_extractor import ImageLsbDataExtractor
+from .lsb_injector import inject_data
+from ...sdk.ai.generate_image._enum import PROMOTION, Model
 
 
 class ImageMetadata(BaseModel):
@@ -69,11 +71,20 @@ def negative_prompt(self):
     model_config = ConfigDict(extra="allow")
 
     @property
-    def jsonfy(self):
-        return {"Title": self.title, "Description": self.description, "Comment": self.comment.model_dump()}
+    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:
+        """
         image = Image.open(input_img).convert('RGBA')
         data = np.array(image)
         data[..., 3] = 254
@@ -82,40 +93,49 @@ def reset_alpha(input_img: BytesIO) -> BytesIO:
         new_image.save(_new_img_io, format="PNG")
         return _new_img_io
 
-    def write_out(self,
-                  img_io: BytesIO,
-                  *,
-                  remove_lsb: bool = False
-                  ):
-        if remove_lsb:
-            img_io = self.reset_alpha(img_io)
-        with Image.open(img_io) as img:
-            metadata = PngInfo()
-            for k, v in self.jsonfy.items():
-                if isinstance(v, dict):
-                    v = json.dumps(v)
-                metadata.add_text(k, v)
-            _new_img = BytesIO()
-            # Save original image with metadata
-            img.save(_new_img, format="PNG", pnginfo=metadata, optimize=False, compress_level=0)
-        return _new_img
+    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:
-        :return:
+        :param image_io: str, bytes, Path, BytesIO
+        :return: ImageMetadata
+        :raises ValidationError:  Data extraction failed
         """
         try:
             image_data = ImageLsbDataExtractor().extract_data(image_io)
             model = cls(**image_data)
         except Exception as e:
             logger.debug(f"Error trying extracting data in LSB: {e}")
         else:
-            print(model)
             return model
         with Image.open(image_io) as img:
             title = img.info.get("Title", None)
@@ -127,13 +147,22 @@ def load_image(cls,
             except Exception as e:
                 logger.debug(f"Error loading comment: {e}")
                 comment = {}
-            return cls(Title=title, Description=prompt, Comment=cls.CommentModel(**comment))
+        return cls(Title=title, Description=prompt, Comment=cls.CommentModel(**comment))
 
     @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)