Merge pull request #22 from LlmKira/dev

inpaint well !

playground/enhance.py

@@ -0,0 +1,9 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/13 下午1:58
+# @Author  : sudoskys
+# @File    : enhance.py
+# @Software: PyCharm
+
+# NOTE About Enhance Mode
+# Enhance Mode = origin model name + img2img action + width *1.5(or 1) +height *1.5(or 1) + change seed
+# :)

playground/mask/__init__.py

@@ -0,0 +1,17 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/13 下午12:42
+# @Author  : sudoskys
+# @File    : __init__.py.py
+# @Software: PyCharm
+from novelai_python.utils.useful import create_mask_from_sketch
+
+with open('sk.jpg', 'rb') as f:
+    sk_bytes = f.read()
+
+with open('ori.png', 'rb') as f:
+    ori_bytes = f.read()
+
+return_bytes = create_mask_from_sketch(original_img_bytes=ori_bytes, sketch_img_bytes=sk_bytes, jagged_edges=True)
+
+with open('mask_export.png', 'wb') as f:
+    f.write(return_bytes)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.2.9"
+version = "0.3.0"
 description = "Novelai Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "[email protected]" },
@@ -18,6 +18,7 @@ dependencies = [
     "uvicorn[standard]>=0.27.0.post1",
     "numpy>=1.24.4",
     "argon2-cffi>=23.1.0",
+    "opencv-python>=4.8.1.78",
 ]
 requires-python = ">=3.8"
 readme = "README.md"

src/novelai_python/sdk/ai/generate_image.py

@@ -207,6 +207,14 @@ def height_validator(cls, v: int):
     parameters: Params = Params()
     model_config = ConfigDict(extra="ignore")
 
+    @property
+    def endpoint(self):
+        return self._endpoint
+
+    @endpoint.setter
+    def endpoint(self, value):
+        self._endpoint = value
+
     @override
     def model_post_init(self, *args) -> None:
         """
@@ -241,20 +249,14 @@ def model_post_init(self, *args) -> None:
     def validate_model(self):
         if self.action == Action.INFILL and not self.parameters.mask:
             logger.warning("Mask maybe required for infill mode.")
+        if self.action == Action.INFILL:
+            self.parameters.extra_noise_seed = self.parameters.seed
         return self
 
     @property
     def base_url(self):
         return f"{self.endpoint.strip('/')}/ai/generate-image"
 
-    @property
-    def endpoint(self):
-        return self._endpoint
-
-    @endpoint.setter
-    def endpoint(self, value):
-        self._endpoint = value
-
     def calculate_cost(self, is_opus: bool = False):
         """
         Calculate the Anlas cost of current parameters.

src/novelai_python/utils/useful.py

@@ -7,6 +7,9 @@
 import random
 from typing import List, Union
 
+import cv2 as cv
+import numpy as np
+
 
 def enum_to_list(enum_):
     return list(map(lambda x: x.value, enum_._member_map_.values()))
@@ -42,3 +45,60 @@ def get(self, user_id: Union[int, str]) -> str:
         user_used.append(selected)
 
         return selected
+
+
+def create_mask_from_sketch(original_img_bytes: bytes,
+                            sketch_img_bytes: bytes,
+                            output_format: str = '.png',
+                            jagged_edges: bool = True
+                            ) -> bytes:
+    """
+    Function to create a mask from original and sketch images input as bytes. Returns bytes.
+
+    :param jagged_edges:  Whether to create jagged edges in the mask. Defaults to False.
+    :param original_img_bytes: Bytes corresponding to the original image.
+    :param sketch_img_bytes: Bytes corresponding to the sketch image.
+    :param output_format: Format of the output image. Defaults to '.png'. It could also be '.jpg'
+    :returns bytes: Bytes corresponding to the resultant mask
+    """
+    # Load images
+    ori_img = cv.imdecode(np.frombuffer(original_img_bytes, np.uint8), cv.IMREAD_COLOR)
+    sketch_img = cv.imdecode(np.frombuffer(sketch_img_bytes, np.uint8), cv.IMREAD_COLOR)
+
+    # Check if images have the same size
+    if ori_img.shape != sketch_img.shape:
+        raise ValueError("Images must have the same size.")
+
+    # Calculate difference between the original and sketch images
+    diff_img = cv.absdiff(ori_img, sketch_img)
+
+    # Convert the difference image to grayscale
+    diff_gray = cv.cvtColor(diff_img, cv.COLOR_BGR2GRAY)
+
+    # Threshold to create the mask
+    _, thresh = cv.threshold(diff_gray, 10, 255, cv.THRESH_BINARY)
+
+    if jagged_edges:
+        # Use a bigger kernel for dilation to create larger 'step' effect at the edges
+        kernel = np.ones((7, 7), np.uint8)
+    else:
+        kernel = np.ones((3, 3), np.uint8)
+
+    # Perform morphological opening to remove small noise
+    opening = cv.morphologyEx(thresh, cv.MORPH_OPEN, kernel, iterations=2)
+
+    # Further remove noise with a Gaussian filter
+    smooth = cv.GaussianBlur(opening, (5, 5), 0)
+
+    if jagged_edges:
+        # Apply additional thresholding to create sharper, jagged edges
+        _, smooth = cv.threshold(smooth, 128, 255, cv.THRESH_BINARY)
+
+    # Convert image to BytesIO object
+    is_success, buffer = cv.imencode(output_format, smooth)
+    if is_success:
+        output_io = buffer.tobytes()
+    else:
+        raise ValueError("Error during conversion of image to BytesIO object")
+
+    return output_io