Merge pull request #65 from LlmKira/dev

✨ fix: Correct data extraction error handling in ImageMetadata
LlmKira · Jul 6, 2024 · 671f1a9 · 671f1a9
2 parents 4c8cad3 + c4de261
commit 671f1a9
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 8 deletions.
diff --git a/playground/image_metadata/read_nai_tag.py b/playground/image_metadata/read_nai_tag.py
@@ -2,13 +2,16 @@
 # @Time    : 2024/2/8 下午4:57
 # @Author  : sudoskys
 # @File    : read_nai_tag.py
-
-import os
+from io import BytesIO
 from pathlib import Path
 
 from novelai_python.tool.image_metadata import ImageMetadata
 
 image = Path(__file__).parent.joinpath("sample-0316.png")
+image_clear = ImageMetadata.reset_alpha(
+    input_img=BytesIO(image.read_bytes())
+)
+
 try:
     meta = ImageMetadata.load_image(image)
 except ValueError:
@@ -18,6 +21,14 @@
 print(meta.Description)
 print(meta.Comment)
 
+try:
+    meta = ImageMetadata.load_image(image_clear)
+    print(meta.Title)
+    print(meta.Description)
+    print(meta.Comment)
+except ValueError:
+    print("Cant find a MetaData")
+
 image = Path(__file__).parent.joinpath("sample-0317.png")
 try:
     meta = ImageMetadata.load_image(image)

diff --git a/src/novelai_python/sdk/ai/generate_image/__init__.py b/src/novelai_python/sdk/ai/generate_image/__init__.py
@@ -453,7 +453,7 @@ def calculate_cost(self, is_opus: bool = False):
                 )
         per_sample = max(math.ceil(per_sample * strength), 2)
         # uncond_scale is not 1.0
-        if not math.isclose(uncond_scale, 1.0, rel_tol=1e-5):
+        if not math.isclose(uncond_scale, 1.0, rel_tol=1e-10):
             per_sample = math.ceil(per_sample * 1.3)
 
         return per_sample * n_samples

diff --git a/src/novelai_python/tool/image_metadata/__init__.py b/src/novelai_python/tool/image_metadata/__init__.py
@@ -155,15 +155,15 @@ def load_image(cls,
         Load image and extract metadata using LSB/Metadata
         :param image_io: str, bytes, Path, BytesIO
         :return: ImageMetadata
-        :raises ValidationError:  Data extraction failed
+        :raises ValueError:  Data extraction failed
         """
         if isinstance(image_io, BytesIO):
             image_io.seek(0)
         try:
             image_data = ImageLsbDataExtractor().extract_data(image_io)
-            model = cls(**image_data)
+            model = cls.model_validate(image_data)
         except Exception as e:
-            logger.debug(f"Error trying extracting data in LSB: {e}")
+            logger.trace(f"Error trying extracting data in LSB: {e}")
         else:
             return model
         with Image.open(image_io) as img:
@@ -174,9 +174,16 @@ def load_image(cls,
                 assert isinstance(comment, str), ValueError("Comment Empty")
                 comment = json.loads(comment)
             except Exception as e:
-                logger.debug(f"Error loading comment: {e}")
+                logger.trace(f"Error loading comment: {e}")
                 comment = {}
-        return cls(Title=title, Description=prompt, Comment=cls.CommentModel(**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(