novelai-python 0.1.6__tar.gz → 0.1.8__tar.gz

{novelai_python-0.1.6 → novelai_python-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.1.6
+Version: 0.1.8
 Summary: Novelai Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache 2.0
@@ -17,6 +17,7 @@ Requires-Dist: Pillow>=10.2.0
 Requires-Dist: curl-cffi>=0.5.10
 Requires-Dist: fastapi>=0.109.0
 Requires-Dist: uvicorn[standard]>=0.27.0.post1
+Requires-Dist: numpy>=1.24.4
 Requires-Dist: pytest>=7.4.4; extra == "testing"
 Provides-Extra: testing
 Description-Content-Type: text/markdown
@@ -31,6 +32,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 ### Roadmap 🚧
 
 - [x] /ai/generate-image
+- [x] /user/subscription
 - [ ] /ai/generate-image/suggest-tags
 - [ ] /ai/annotate-image
 - [ ] /ai/classify

{novelai_python-0.1.6 → novelai_python-0.1.8}/README.md

@@ -8,6 +8,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 ### Roadmap 🚧
 
 - [x] /ai/generate-image
+- [x] /user/subscription
 - [ ] /ai/generate-image/suggest-tags
 - [ ] /ai/annotate-image
 - [ ] /ai/classify

{novelai_python-0.1.6 → novelai_python-0.1.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.1.6"
+version = "0.1.8"
 description = "Novelai Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "coldlando@hotmail.com" },
@@ -16,6 +16,7 @@ dependencies = [
     "curl-cffi>=0.5.10",
     "fastapi>=0.109.0",
     "uvicorn[standard]>=0.27.0.post1",
+    "numpy>=1.24.4",
 ]
 requires-python = ">=3.8"
 readme = "README.md"

{novelai_python-0.1.6 → novelai_python-0.1.8}/src/novelai_python/__init__.py

@@ -9,14 +9,19 @@ from ._exceptions import (
     AuthError,
     NovelAiError,
 )  # noqa: F401, F403
-from ._response import ImageGenerateResp  # noqa: F401, F403
 from .credential import JwtCredential  # noqa: F401, F403
-from .sdk.ai import GenerateImageInfer  # noqa: F401, F403
+from .sdk.ai import GenerateImageInfer, ImageGenerateResp  # noqa: F401, F403
+from .sdk.user import Subscription, SubscriptionResp  # noqa: F401, F403
 
 __all__ = [
     "GenerateImageInfer",
     "ImageGenerateResp",
+
+    "Subscription",
+    "SubscriptionResp",
+
     "JwtCredential",
+
     "APIError",
     "AuthError",
     "NovelAiError",

{novelai_python-0.1.6 → novelai_python-0.1.8}/src/novelai_python/_response/__init__.py

@@ -4,7 +4,8 @@
 # @File    : __init__.py.py
 # @Software: PyCharm
 from .ai.generate_image import ImageGenerateResp
-
+from .user.subscription import SubscriptionResp
 __all__ = [
-    "ImageGenerateResp"
+    "ImageGenerateResp",
+    "SubscriptionResp",
 ]

novelai_python-0.1.8/src/novelai_python/_response/user/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/7 上午9:57
+# @Author  : sudoskys
+# @File    : __init__.py.py
+# @Software: PyCharm

novelai_python-0.1.8/src/novelai_python/_response/user/subscription.py

@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/7 上午9:57
+# @Author  : sudoskys
+# @File    : subscription.py
+# @Software: PyCharm
+from typing import Optional, Dict, Any, List
+
+from pydantic import BaseModel, Field
+
+
+class TrainingSteps(BaseModel):
+    fixedTrainingStepsLeft: int
+    purchasedTrainingSteps: int
+
+
+class ImageGenerationLimit(BaseModel):
+    resolution: int
+    maxPrompts: int
+
+
+class Perks(BaseModel):
+    maxPriorityActions: int
+    startPriority: int
+    moduleTrainingSteps: int
+    unlimitedMaxPriority: bool
+    voiceGeneration: bool
+    imageGeneration: bool
+    unlimitedImageGeneration: bool
+    unlimitedImageGenerationLimits: List[ImageGenerationLimit]
+    contextTokens: int
+
+
+class SubscriptionResp(BaseModel):
+    tier: int = Field(..., description="Subscription tier")
+    active: bool = Field(..., description="Subscription status")
+    expiresAt: int = Field(..., description="Subscription expiration time")
+    perks: Perks = Field(..., description="Subscription perks")
+    paymentProcessorData: Optional[Dict[Any, Any]]
+    trainingStepsLeft: TrainingSteps = Field(..., description="Training steps left")
+    accountType: int = Field(..., description="Account type")
+
+    @property
+    def is_active(self):
+        return self.active
+
+    @property
+    def anlas_left(self):
+        return self.trainingStepsLeft.fixedTrainingStepsLeft
+
+    @property
+    def is_unlimited_image_generation(self):
+        return self.perks.unlimitedImageGeneration

{novelai_python-0.1.6 → novelai_python-0.1.8}/src/novelai_python/sdk/__init__.py

@@ -5,4 +5,4 @@
 # @Software: PyCharm
 
 from .ai.generate_image import GenerateImageInfer, ImageGenerateResp  # noqa 401
-
+from .user.subscription import Subscription, SubscriptionResp  # noqa 401

{novelai_python-0.1.6 → novelai_python-0.1.8}/src/novelai_python/sdk/ai/generate_image.py

@@ -12,6 +12,7 @@ from zipfile import ZipFile
 import httpx
 from curl_cffi.requests import AsyncSession
 from loguru import logger
+from novelai_python.utils import NovelAiMetadata
 from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, model_validator
 
 from ..._exceptions import APIError, AuthError
@@ -229,7 +230,15 @@ class GenerateImageInfer(BaseModel):
             parameters=cls.Params(**param)
         )
 
-    async def generate(self, session: Union[AsyncSession, JwtCredential]) -> ImageGenerateResp:
+    async def generate(self, session: Union[AsyncSession, JwtCredential],
+                       *,
+                       remove_sign: bool = False) -> ImageGenerateResp:
+        """
+        生成图片
+        :param session:  session
+        :param remove_sign:  移除追踪信息
+        :return:
+        """
         if isinstance(session, JwtCredential):
             session = session.session
         request_data = self.model_dump(exclude_none=True)
@@ -274,7 +283,12 @@ class GenerateImageInfer(BaseModel):
                         response=try_jsonfy(response.content)
                     )
                 for filename in file_list:
-                    unzip_content.append((filename, zip_file.read(filename)))
+                    data = zip_file.read(filename)
+                    if remove_sign:
+                        data = NovelAiMetadata.rehash(BytesIO(data), remove_stealth=True)
+                        if not isinstance(data, bytes):
+                            data = data.getvalue()
+                    unzip_content.append((filename, data))
             return ImageGenerateResp(
                 meta=ImageGenerateResp.RequestParams(
                     endpoint=self.base_url,

novelai_python-0.1.8/src/novelai_python/sdk/user/__init__.py

@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/7 上午10:03
+# @Author  : sudoskys
+# @File    : __init__.py.py
+# @Software: PyCharm
+
+from .subscription import *  # noqa: F403

novelai_python-0.1.8/src/novelai_python/sdk/user/subscription.py

@@ -0,0 +1,81 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/7 上午10:04
+# @Author  : sudoskys
+# @File    : subscription.py.py
+# @Software: PyCharm
+from typing import Optional, Union
+
+import httpx
+from curl_cffi.requests import AsyncSession
+from loguru import logger
+from pydantic import BaseModel, PrivateAttr
+
+from ... import APIError, AuthError
+from ..._response.user.subscription import SubscriptionResp
+from ...credential import JwtCredential
+from ...utils import try_jsonfy
+
+
+class Subscription(BaseModel):
+    _endpoint: Optional[str] = PrivateAttr("https://api.novelai.net")
+
+    @property
+    def base_url(self):
+        return f"{self.endpoint.strip('/')}/user/subscription"
+
+    @property
+    def endpoint(self):
+        return self._endpoint
+
+    @endpoint.setter
+    def endpoint(self, value):
+        self._endpoint = value
+
+    async def request(self,
+                      session: Union[AsyncSession, JwtCredential],
+                      ) -> SubscriptionResp:
+        """
+        Request to get user subscription information
+        :param session: 
+        :return: 
+        """
+        if isinstance(session, JwtCredential):
+            session = session.session
+        request_data = {}
+        logger.debug("Subscription")
+        try:
+            assert hasattr(session, "get"), "session must have get method."
+            response = await session.get(
+                self.base_url,
+            )
+            if "application/json" not in response.headers.get('Content-Type'):
+                logger.error(f"Unexpected content type: {response.headers.get('Content-Type')}")
+                try:
+                    _msg = response.json()
+                except Exception:
+                    raise APIError(
+                        message=f"Unexpected content type: {response.headers.get('Content-Type')}",
+                        request=request_data,
+                        status_code=response.status_code,
+                        response=try_jsonfy(response.content)
+                    )
+                status_code = _msg.get("statusCode", response.status_code)
+                message = _msg.get("message", "Unknown error")
+                if status_code in [400, 401, 402]:
+                    # 400 : validation error
+                    # 401 : unauthorized
+                    # 402 : payment required
+                    # 409 : conflict
+                    raise AuthError(message, request=request_data, status_code=status_code, response=_msg)
+                if status_code in [500]:
+                    # An unknown error occured.
+                    raise APIError(message, request=request_data, status_code=status_code, response=_msg)
+                raise APIError(message, request=request_data, status_code=status_code, response=_msg)
+            return SubscriptionResp.model_validate(response.json())
+        except httpx.HTTPError as exc:
+            raise RuntimeError(f"An HTTP error occurred: {exc}")
+        except APIError as e:
+            raise e
+        except Exception as e:
+            logger.opt(exception=e).exception("An Unexpected error occurred")
+            raise e

{novelai_python-0.1.6 → novelai_python-0.1.8}/src/novelai_python/server.py

@@ -10,10 +10,12 @@ import zipfile
 import uvicorn
 from fastapi import FastAPI, Depends, Security
 from fastapi.security import APIKeyHeader
+from loguru import logger
 from starlette.responses import JSONResponse, StreamingResponse
 
 from .credential import JwtCredential, SecretStr
 from .sdk.ai.generate_image import GenerateImageInfer
+from .sdk.user.subscription import Subscription
 
 app = FastAPI()
 token_key = APIKeyHeader(name="Authorization")
@@ -35,6 +37,24 @@ async def health():
     return {"status": "ok"}
 
 
+@app.get("/user/subscription")
+async def subscription(
+        current_token: str = Depends(get_current_token)
+):
+    """
+    订阅
+    :param current_token: Authorization
+    :param req: Subscription
+    :return:
+    """
+    try:
+        _result = await Subscription().request(session=get_session(current_token))
+        return _result.model_dump()
+    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,
@@ -58,6 +78,7 @@ async def generate_image(
             'Content-Disposition': 'attachment;filename=multiple_files.zip'
         })
     except Exception as e:
+        logger.exception(e)
         return JSONResponse(status_code=500, content=e.__dict__)
 
 

{novelai_python-0.1.6 → novelai_python-0.1.8}/src/novelai_python/utils/__init__.py

@@ -12,13 +12,16 @@ from loguru import logger
 from .hash import NovelAiMetadata
 
 
-def png_info_reset(img: Union[str, io.BytesIO]):
+def png_info_reset(img: Union[str, io.BytesIO],
+                   *,
+                   remove_stealth: bool = False):
     """
     reset png info hash value
+    :param remove_stealth: LCB
     :param img: BytesIO 对象
     :return: 
     """
-    _fixed = NovelAiMetadata.rehash(img_io=img)
+    _fixed = NovelAiMetadata.rehash(img_io=img, remove_stealth=remove_stealth)
     return _fixed
 
 

{novelai_python-0.1.6 → novelai_python-0.1.8}/src/novelai_python/utils/hash.py

@@ -10,6 +10,7 @@ import json
 from io import BytesIO
 from typing import Union
 
+import numpy as np
 from PIL import Image
 from PIL.PngImagePlugin import PngInfo
 from loguru import logger
@@ -35,16 +36,33 @@ class NovelAiMetadata(BaseModel):
         return {"Title": self.title, "Description": self.description, "Comment": self.comment}
 
     @staticmethod
-    def rehash(img_io):
-        cls = NaiPic.read_from_img(img_io)
+    def rehash(img_io, remove_stealth: bool = False):
+        cls = NovelAiMetadata.build_from_img(img_io)
         cls.comment["signed_hash"] = sign_message(json.dumps(cls.description), "novalai-client")
-        cls.write_out(img_io=img_io)
-        return img_io
+        _new_img_io = cls.write_out(img_io=img_io, remove_stealth=remove_stealth)
+        return _new_img_io
 
-    def write_out(self, img_io: BytesIO):
+    @staticmethod
+    def remove_stealth_info(img_io):
+        image = Image.open(img_io).convert('RGBA')
+        data = np.array(image)
+        data[..., 3] = 254
+        new_image = Image.fromarray(data, 'RGBA')
+        _new_img_io = BytesIO()
+        new_image.save(_new_img_io, format="PNG")
+        return _new_img_io
+
+    def write_out(self, img_io: BytesIO, *, remove_stealth: bool = False):
         with Image.open(img_io) as img:
+            if remove_stealth:
+                img = img.convert('RGBA')
+                data = np.array(img)
+                data[..., 3] = 254
+                img = Image.fromarray(data, 'RGBA')
             metadata = PngInfo()
             for k, v in self.metadata.items():
+                if isinstance(v, dict):
+                    v = json.dumps(v)
                 metadata.add_text(k, v)
             _new_img = BytesIO()
             img.save(_new_img, format="PNG", pnginfo=metadata, quality=95, optimize=False)
@@ -88,3 +106,23 @@ def decode_base64(encoded_data):
     # 使用base64解码
     decoded_data = base64.b64decode(byte_data)
     return decoded_data.decode("UTF-8")  # 再次使用UTF-8将字节对象转为字符串
+
+
+def decode_image(img: Union[str, BytesIO], output_path: str):
+    encoded_image = Image.open(img)
+    red_channel = encoded_image.split()[3]
+    x_size = encoded_image.size[0]
+    y_size = encoded_image.size[1]
+    decoded_image = Image.new("RGB", encoded_image.size)
+    pixels = decoded_image.load()
+    for i in range(x_size):
+        for j in range(y_size):
+            if i < 7:  # the left 7 columns
+                r = red_channel.getpixel((i, j))
+                if r > 254:  # above the threshold
+                    pixels[i, j] = (0, 0, 0)  # black
+                else:
+                    pixels[i, j] = (255, 255, 255)  # white
+            else:
+                pixels[i, j] = (255, 255, 255)  # default to white for right part of image
+    decoded_image.save(output_path)

novelai_python-0.1.8/tests/test_server_run.py

@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/1/30 下午11:52
+# @Author  : sudoskys
+# @File    : test_server_run.py
+# @Software: PyCharm
+
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+
+from src.novelai_python.server import app, get_session
+
+client = TestClient(app)
+
+
+def test_health_check():
+    response = client.get("/health")
+    assert response.status_code == 200
+    assert response.json() == {"status": "ok"}
+
+
+@patch('src.novelai_python.server.Subscription')
+def test_subscription_without_api_key(mock_subscription):
+    mock_subscription.return_value.request.return_value.model_dump.return_value = {"status": "subscribed"}
+    response = client.get("/user/subscription")
+    assert response.status_code == 403
+
+
+@patch('src.novelai_python.server.GenerateImageInfer')
+def test_generate_image_without_api_key(mock_generate_image):
+    mock_generate_image.return_value.generate.return_value = {"status": "image generated"}
+    response = client.post("/ai/generate_image")
+    assert response.status_code == 403
+
+
+def test_get_session_new_token():
+    token = "new_token"
+    session = get_session(token)
+    assert session.jwt_token.get_secret_value() == token
+
+
+def test_get_session_existing_token():
+    token = "existing_token"
+    get_session(token)
+    session = get_session(token)
+    assert session.jwt_token.get_secret_value() == token

novelai_python-0.1.6/tests/test_server_run.py

@@ -1,28 +0,0 @@
-# -*- coding: utf-8 -*-
-# @Time    : 2024/1/30 下午11:52
-# @Author  : sudoskys
-# @File    : test_server_run.py
-# @Software: PyCharm
-from fastapi.testclient import TestClient
-
-from src.novelai_python.sdk.ai.generate_image import GenerateImageInfer
-from src.novelai_python.server import app, get_session
-
-client = TestClient(app)
-
-
-def test_health_check():
-    response = client.get("/health")
-    assert response.status_code == 200
-    assert response.json() == {"status": "ok"}
-
-
-def test_generate_image_with_valid_token():
-    valid_token = "valid_token"
-    get_session(valid_token)  # to simulate a valid session
-    response = client.post(
-        "/ai/generate_image",
-        headers={"Authorization": valid_token},
-        json=GenerateImageInfer(input="1girl").model_dump()
-    )
-    assert response.status_code == 500