novelai-python 0.2.9__tar.gz → 0.3.1__tar.gz

{novelai_python-0.2.9 → novelai_python-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.2.9
+Version: 0.3.1
 Summary: Novelai Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache License 2.0
@@ -19,6 +19,7 @@ Requires-Dist: fastapi>=0.109.0
 Requires-Dist: uvicorn[standard]>=0.27.0.post1
 Requires-Dist: numpy>=1.24.4
 Requires-Dist: argon2-cffi>=23.1.0
+Requires-Dist: opencv-python>=4.8.1.78
 Requires-Dist: pytest>=7.4.4; extra == "testing"
 Requires-Dist: pytest-asyncio>=0.23.4; extra == "testing"
 Provides-Extra: testing

{novelai_python-0.2.9 → novelai_python-0.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.2.9"
+version = "0.3.1"
 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"

{novelai_python-0.2.9 → novelai_python-0.3.1}/src/novelai_python/sdk/ai/generate_image.py

@@ -7,6 +7,7 @@ import base64
 import json
 import math
 import random
+from copy import deepcopy
 from enum import Enum, IntEnum
 from io import BytesIO
 from typing import Optional, Union
@@ -207,6 +208,14 @@ class GenerateImageInfer(ApiBaseModel):
     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 +250,14 @@ class GenerateImageInfer(ApiBaseModel):
     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.GENERATE:
+            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.
@@ -392,6 +395,7 @@ class GenerateImageInfer(ApiBaseModel):
             "Sec-Fetch-Site": "same-site",
             "Pragma": "no-cache",
             "Cache-Control": "no-cache",
+            "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0"
         }
 
     async def request(self,
@@ -419,20 +423,21 @@ class GenerateImageInfer(ApiBaseModel):
             session.headers.clear()
             session.headers.update(override_headers)
         try:
-            _log_data = request_data.copy()
+            _log_data = deepcopy(request_data)
             if self.action == Action.GENERATE:
                 logger.debug(f"Request Data: {_log_data}")
             else:
                 _log_data.get("parameters", {}).update({
-                    "image": "base64 data" if self.parameters.image else "None",
+                    "image": "base64 data" if self.parameters.image else None,
                 }
                 )
                 _log_data.get("parameters", {}).update(
                     {
-                        "mask": "base64 data" if self.parameters.mask else "None",
+                        "mask": "base64 data" if self.parameters.mask else None,
                     }
                 )
-                logger.debug(f"Request Data: {request_data}")
+                logger.debug(f"Request Data: {_log_data}")
+            del _log_data
         except Exception as e:
             logger.warning(f"Error when print log data: {e}")
         try:

{novelai_python-0.2.9 → novelai_python-0.3.1}/src/novelai_python/server.py

@@ -15,6 +15,7 @@ from starlette.responses import JSONResponse, StreamingResponse
 
 from .credential import JwtCredential, SecretStr
 from .sdk.ai.generate_image import GenerateImageInfer
+from .sdk.ai.upscale import Upscale
 from .sdk.user.information import Information
 from .sdk.user.login import Login
 from .sdk.user.subscription import Subscription
@@ -90,6 +91,33 @@ async def subscription(
         return JSONResponse(status_code=500, content=e.__dict__)
 
 
+@app.post("/ai/upscale")
+async def upscale(
+        req: Upscale,
+        current_token: str = Depends(get_current_token)
+):
+    """
+    生成图片
+    :param current_token: Authorization
+    :param req: Upscale
+    :return:
+    """
+    try:
+        _result = await req.request(session=get_session(current_token), remove_sign=True)
+        zip_file_bytes = io.BytesIO()
+        with zipfile.ZipFile(zip_file_bytes, mode="w", compression=zipfile.ZIP_DEFLATED) as zip_file:
+            file = _result.files  # ONLY TUPLE
+            zip_file.writestr(zinfo_or_arcname=file[0], data=file[1])
+        # return the zip file
+        zip_file_bytes.seek(0)
+        return StreamingResponse(zip_file_bytes, media_type='application/zip', headers={
+            'Content-Disposition': 'attachment;filename=image.zip'
+        })
+    except Exception as e:
+        logger.exception(e)
+        return JSONResponse(status_code=500, content=e.__dict__)
+
+
 @app.post("/ai/generate_image")
 async def generate_image(
         req: GenerateImageInfer,
@@ -110,7 +138,7 @@ async def generate_image(
         # return the zip file
         zip_file_bytes.seek(0)
         return StreamingResponse(zip_file_bytes, media_type='application/zip', headers={
-            'Content-Disposition': 'attachment;filename=multiple_files.zip'
+            'Content-Disposition': 'attachment;filename=image.zip'
         })
     except Exception as e:
         logger.exception(e)

novelai_python-0.3.1/src/novelai_python/utils/useful.py

@@ -0,0 +1,104 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/9 下午10:41
+# @Author  : sudoskys
+# @File    : useful.py
+# @Software: PyCharm
+import collections
+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()))
+
+
+class QueSelect(object):
+    def __init__(self, data: List[str]):
+        """
+        A queue selector
+        :param data:
+        """
+        self.data = collections.deque(data)
+        self.used = collections.deque()
+        self.users = {}
+
+    def get(self, user_id: Union[int, str]) -> str:
+        user_id = str(user_id)
+        if user_id not in self.users:
+            self.users[user_id] = {'data': self.data.copy(), 'used': collections.deque()}
+
+        user_data = self.users[user_id]['data']
+        user_used = self.users[user_id]['used']
+
+        if len(user_data) == 0:
+            user_data, user_used = user_used, user_data
+            # 随机掉 User data
+            if len(user_data) > 2:
+                random.shuffle(user_data)
+            self.users[user_id]['data'] = user_data
+            self.users[user_id]['used'] = user_used
+
+        selected = user_data.popleft()
+        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

novelai_python-0.2.9/src/novelai_python/utils/useful.py

@@ -1,44 +0,0 @@
-# -*- coding: utf-8 -*-
-# @Time    : 2024/2/9 下午10:41
-# @Author  : sudoskys
-# @File    : useful.py
-# @Software: PyCharm
-import collections
-import random
-from typing import List, Union
-
-
-def enum_to_list(enum_):
-    return list(map(lambda x: x.value, enum_._member_map_.values()))
-
-
-class QueSelect(object):
-    def __init__(self, data: List[str]):
-        """
-        A queue selector
-        :param data:
-        """
-        self.data = collections.deque(data)
-        self.used = collections.deque()
-        self.users = {}
-
-    def get(self, user_id: Union[int, str]) -> str:
-        user_id = str(user_id)
-        if user_id not in self.users:
-            self.users[user_id] = {'data': self.data.copy(), 'used': collections.deque()}
-
-        user_data = self.users[user_id]['data']
-        user_used = self.users[user_id]['used']
-
-        if len(user_data) == 0:
-            user_data, user_used = user_used, user_data
-            # 随机掉 User data
-            if len(user_data) > 2:
-                random.shuffle(user_data)
-            self.users[user_id]['data'] = user_data
-            self.users[user_id]['used'] = user_used
-
-        selected = user_data.popleft()
-        user_used.append(selected)
-
-        return selected