diff --git a/pdm.lock b/pdm.lock index 637b037..d699a32 100644 --- a/pdm.lock +++ b/pdm.lock @@ -5,7 +5,7 @@ groups = ["default", "dev", "testing"] strategy = ["cross_platform", "inherit_metadata"] lock_version = "4.4.1" -content_hash = "sha256:601014f72651e21a1c2c5592a2c66fa00ce682f0d7fc1b9ddcfd698ad7eeba30" +content_hash = "sha256:9aa6958e83642a27ca1aae9b488f8f01f014246d5f42b49f2cf7a04fc6328cb1" [[package]] name = "annotated-types" @@ -576,6 +576,32 @@ files = [ {file = "numpy-1.24.4.tar.gz", hash = "sha256:80f5e3a4e498641401868df4208b74581206afbee7cf7b8329daae82676d9463"}, ] +[[package]] +name = "opencv-python" +version = "4.8.1.78" +requires_python = ">=3.6" +summary = "Wrapper package for OpenCV python bindings." +groups = ["default"] +dependencies = [ + "numpy>=1.17.0; python_version >= \"3.7\"", + "numpy>=1.17.3; python_version >= \"3.8\"", + "numpy>=1.19.3; python_version >= \"3.6\" and platform_system == \"Linux\" and platform_machine == \"aarch64\"", + "numpy>=1.19.3; python_version >= \"3.9\"", + "numpy>=1.21.0; python_version <= \"3.9\" and platform_system == \"Darwin\" and platform_machine == \"arm64\"", + "numpy>=1.21.2; python_version >= \"3.10\"", + "numpy>=1.21.4; python_version >= \"3.10\" and platform_system == \"Darwin\"", + "numpy>=1.23.5; python_version >= \"3.11\"", +] +files = [ + {file = "opencv-python-4.8.1.78.tar.gz", hash = "sha256:cc7adbbcd1112877a39274106cb2752e04984bc01a031162952e97450d6117f6"}, + {file = "opencv_python-4.8.1.78-cp37-abi3-macosx_10_16_x86_64.whl", hash = "sha256:91d5f6f5209dc2635d496f6b8ca6573ecdad051a09e6b5de4c399b8e673c60da"}, + {file = "opencv_python-4.8.1.78-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:bc31f47e05447da8b3089faa0a07ffe80e114c91ce0b171e6424f9badbd1c5cd"}, + {file = "opencv_python-4.8.1.78-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9814beca408d3a0eca1bae7e3e5be68b07c17ecceb392b94170881216e09b319"}, + {file = "opencv_python-4.8.1.78-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c4c406bdb41eb21ea51b4e90dfbc989c002786c3f601c236a99c59a54670a394"}, + {file = "opencv_python-4.8.1.78-cp37-abi3-win32.whl", hash = "sha256:a7aac3900fbacf55b551e7b53626c3dad4c71ce85643645c43e91fcb19045e47"}, + {file = "opencv_python-4.8.1.78-cp37-abi3-win_amd64.whl", hash = "sha256:b983197f97cfa6fcb74e1da1802c7497a6f94ed561aba6980f1f33123f904956"}, +] + [[package]] name = "packaging" version = "23.2" diff --git a/playground/enhance.py b/playground/enhance.py new file mode 100644 index 0000000..9c3b2c4 --- /dev/null +++ b/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 +# :) diff --git a/playground/mask/__init__.py b/playground/mask/__init__.py new file mode 100644 index 0000000..afb3fe2 --- /dev/null +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 856ef44..c3af90f 100644 --- a/pyproject.toml +++ b/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 = "coldlando@hotmail.com" }, @@ -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" diff --git a/src/novelai_python/sdk/ai/generate_image.py b/src/novelai_python/sdk/ai/generate_image.py index 231910d..2fff470 100644 --- a/src/novelai_python/sdk/ai/generate_image.py +++ b/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. diff --git a/src/novelai_python/utils/useful.py b/src/novelai_python/utils/useful.py index 5a4a7e2..db52059 100644 --- a/src/novelai_python/utils/useful.py +++ b/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