Merge pull request #24 from LlmKira/dev

Dev

playground/mask/__init__.py

@@ -11,7 +11,11 @@
 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)
+return_bytes = create_mask_from_sketch(original_img_bytes=ori_bytes,
+                                       sketch_img_bytes=sk_bytes,
+                                       jagged_edges=True,
+                                       min_block_size=15
+                                       )
 
 with open('mask_export.png', 'wb') as f:
     f.write(return_bytes)

playground/suggest_tag.py

@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/13 下午8:25
+# @Author  : sudoskys
+# @File    : suggest_tag.py
+# @Software: PyCharm
+import asyncio
+import os
+
+from dotenv import load_dotenv
+from loguru import logger
+from pydantic import SecretStr
+
+from novelai_python import APIError
+from novelai_python import SuggestTags, SuggestTagsResp, JwtCredential
+
+load_dotenv()
+
+token = None
+jwt = os.getenv("NOVELAI_JWT") or token
+
+
+async def main():
+    globe_s = JwtCredential(jwt_token=SecretStr(jwt))
+    try:
+        _res = await SuggestTags().request(
+            session=globe_s
+        )
+        _res: SuggestTagsResp
+        print(f"Information: {_res}")
+        print(_res.model_dump())
+    except APIError as e:
+        logger.exception(e)
+        print(e.__dict__)
+        return
+
+
+loop = asyncio.get_event_loop()
+loop.run_until_complete(main())

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.3.1"
+version = "0.3.2"
 description = "Novelai Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "[email protected]" },

src/novelai_python/__init__.py

@@ -13,6 +13,7 @@
 )
 from .credential import JwtCredential, LoginCredential, ApiCredential
 from .sdk import GenerateImageInfer, ImageGenerateResp
+from .sdk import SuggestTags, SuggestTagsResp
 from .sdk import Information, InformationResp
 from .sdk import Login, LoginResp
 from .sdk import Subscription, SubscriptionResp
@@ -31,6 +32,9 @@
     "Login",
     "LoginResp",
 
+    "SuggestTags",
+    "SuggestTagsResp",
+
     "Information",
     "InformationResp",
 

src/novelai_python/_response/ai/generate_image.py

@@ -22,3 +22,12 @@ def query_params(self, key: str, default=None):
         if not isinstance(self.meta.raw_request.get("parameters"), dict):
             raise Exception("Resp parameters is not dict")
         return self.meta.raw_request.get("parameters").get(key, default)
+
+
+class SuggestTagsResp(RespBase):
+    class Tag(BaseModel):
+        tag: str
+        count: int
+        confidence: float
+
+    tags: List[Tag] = None

src/novelai_python/credential/ApiToken.py

@@ -25,6 +25,7 @@ async def get_session(self, timeout: int = 180, update_headers: dict = None):
                 "Accept": "*/*",
                 "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
                 "Accept-Encoding": "gzip, deflate, br",
+                "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
                 "Authorization": f"Bearer {self.api_token.get_secret_value()}",
                 "Content-Type": "application/json",
                 "Origin": "https://novelai.net",

src/novelai_python/credential/JwtToken.py

@@ -24,7 +24,7 @@ async def get_session(self, timeout: int = 180, update_headers: dict = None):
             self._session = AsyncSession(timeout=timeout, headers={
                 "Accept": "*/*",
                 "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
-                "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0",
+                "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
                 "Accept-Encoding": "gzip, deflate, br",
                 "Authorization": f"Bearer {self.jwt_token.get_secret_value()}",
                 "Content-Type": "application/json",

src/novelai_python/credential/UserAuth.py

@@ -30,7 +30,7 @@ async def get_session(self, timeout: int = 180, update_headers: dict = None):
             resp = await Login.build(user_name=self.username, password=self.password.get_secret_value()).request()
             self._session = AsyncSession(timeout=timeout, headers={
                 "Accept": "*/*",
-                "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:123.0) Gecko/20100101 Firefox/123.0",
+                "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
                 "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
                 "Accept-Encoding": "gzip, deflate, br",
                 "Authorization": f"Bearer {resp.accessToken}",

src/novelai_python/sdk/__init__.py

@@ -5,7 +5,8 @@
 # @Software: PyCharm
 
 from .ai.generate_image import GenerateImageInfer, ImageGenerateResp  # noqa 401
+from .ai.generate_image.suggest_tags import SuggestTags, SuggestTagsResp  # noqa 401
+from .ai.upscale import Upscale, UpscaleResp  # noqa 401
 from .user.information import Information, InformationResp  # noqa 401
 from .user.login import Login, LoginResp  # noqa 401
 from .user.subscription import Subscription, SubscriptionResp  # noqa 401
-from .ai.upscale import Upscale, UpscaleResp  # noqa 401

src/novelai_python/sdk/ai/generate_image.py → src/novelai_python/sdk/ai/generate_image/__init__.py

@@ -1,14 +1,13 @@
 # -*- coding: utf-8 -*-
-# @Time    : 2024/1/26 上午11:21
+# @Time    : 2024/2/13 下午8:08
 # @Author  : sudoskys
-# @File    : generate_image.py
+# @File    : __init__.py.py
 # @Software: PyCharm
 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
 from urllib.parse import urlparse
@@ -21,84 +20,12 @@
 from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator, model_validator, Field
 from typing_extensions import override
 
-from ..schema import ApiBaseModel
-from ..._exceptions import APIError, AuthError, ConcurrentGenerationError, SessionHttpError
-from ..._response.ai.generate_image import ImageGenerateResp
-from ...credential import CredentialBase
-from ...utils import try_jsonfy, NovelAiMetadata
-
-
-class Sampler(Enum):
-    K_EULER = "k_euler"
-    K_EULER_ANCESTRAL = "k_euler_ancestral"
-    K_DPMPP_2S_ANCESTRAL = "k_dpmpp_2s_ancestral"
-    K_DPMPP_2M = "k_dpmpp_2m"
-    K_DPMPP_SDE = "k_dpmpp_sde"
-    DDIM_V3 = "ddim_v3"
-
-
-class NoiseSchedule(Enum):
-    NATIVE = "native"
-    KARRAS = "karras"
-    EXPONENTIAL = "exponential"
-    POLYEXPONENTIAL = "polyexponential"
-
-
-class UCPreset(IntEnum):
-    TYPE0 = 0
-    TYPE1 = 1
-    TYPE2 = 2
-    TYPE3 = 3
-
-
-class Action(Enum):
-    GENERATE = "generate"
-    """Generate Image"""
-    IMG2IMG = "img2img"
-    """Image to Image"""
-    INFILL = "infill"
-    """Inpainting"""
-
-
-class Model(Enum):
-    NAI_DIFFUSION_3 = "nai-diffusion-3"
-    NAI_DIFFUSION_3_INPAINTING = "nai-diffusion-3-inpainting"
-
-    NAI_DIFFUSION = "nai-diffusion"
-    NAI_DIFFUSION_INPAINTING = "nai-diffusion-inpainting"
-
-    SAFE_DIFFUSION = "safe-diffusion"
-    SAFE_DIFFUSION_INPAINTING = "safe-diffusion-inpainting"
-
-    NAI_DIFFUSION_FURRY = "nai-diffusion-furry"
-    FURRY_DIFFUSION_INPAINTING = "furry-diffusion-inpainting"
-
-
-class ControlNetModel(Enum):
-    HED = "hed"
-    """边缘检测"""
-    MIDAS = "midas"
-    """景深"""
-    FAKE_SCRIBBLE = "fake_scribble"
-    """伪涂鸦"""
-    M_LSD = "mlsd"
-    """(建筑)线条检测"""
-    LANDSCAPER = "uniformer"
-    """风景生成"""
-
-
-class Resolution(Enum):
-    RES_512_768 = (512, 768)
-    RES_768_512 = (768, 512)
-    RES_640_640 = (640, 640)
-    RES_832_1216 = (832, 1216)
-    RES_1216_832 = (1216, 832)
-    RES_1024_1024 = (1024, 1024)
-    RES_1024_1536 = (1024, 1536)
-    RES_1536_1024 = (1536, 1024)
-    RES_1472_1472 = (1472, 1472)
-    RES_1088_1920 = (1088, 1920)
-    RES_1920_1088 = (1920, 1088)
+from ._enum import Model, Sampler, NoiseSchedule, ControlNetModel, Action, UCPreset
+from ...schema import ApiBaseModel
+from ...._exceptions import APIError, AuthError, ConcurrentGenerationError, SessionHttpError
+from ...._response.ai.generate_image import ImageGenerateResp
+from ....credential import CredentialBase
+from ....utils import try_jsonfy, NovelAiMetadata
 
 
 class GenerateImageInfer(ApiBaseModel):
@@ -395,7 +322,7 @@ async def necessary_headers(self, request_data) -> dict:
             "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"
+            "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 13_2_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/110.0.0.0 Safari/537.36",
         }
 
     async def request(self,

src/novelai_python/sdk/ai/generate_image/_enum.py

@@ -0,0 +1,78 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/13 下午8:10
+# @Author  : sudoskys
+# @File    : _enum.py
+# @Software: PyCharm
+from enum import Enum, IntEnum
+
+
+class Sampler(Enum):
+    K_EULER = "k_euler"
+    K_EULER_ANCESTRAL = "k_euler_ancestral"
+    K_DPMPP_2S_ANCESTRAL = "k_dpmpp_2s_ancestral"
+    K_DPMPP_2M = "k_dpmpp_2m"
+    K_DPMPP_SDE = "k_dpmpp_sde"
+    DDIM_V3 = "ddim_v3"
+
+
+class NoiseSchedule(Enum):
+    NATIVE = "native"
+    KARRAS = "karras"
+    EXPONENTIAL = "exponential"
+    POLYEXPONENTIAL = "polyexponential"
+
+
+class UCPreset(IntEnum):
+    TYPE0 = 0
+    TYPE1 = 1
+    TYPE2 = 2
+    TYPE3 = 3
+
+
+class Action(Enum):
+    GENERATE = "generate"
+    """Generate Image"""
+    IMG2IMG = "img2img"
+    """Image to Image"""
+    INFILL = "infill"
+    """Inpainting"""
+
+
+class ControlNetModel(Enum):
+    HED = "hed"
+    """边缘检测"""
+    MIDAS = "midas"
+    """景深"""
+    FAKE_SCRIBBLE = "fake_scribble"
+    """伪涂鸦"""
+    M_LSD = "mlsd"
+    """(建筑)线条检测"""
+    LANDSCAPER = "uniformer"
+    """风景生成"""
+
+
+class Resolution(Enum):
+    RES_512_768 = (512, 768)
+    RES_768_512 = (768, 512)
+    RES_640_640 = (640, 640)
+    RES_832_1216 = (832, 1216)
+    RES_1216_832 = (1216, 832)
+    RES_1024_1024 = (1024, 1024)
+    RES_1024_1536 = (1024, 1536)
+    RES_1536_1024 = (1536, 1024)
+    RES_1472_1472 = (1472, 1472)
+    RES_1088_1920 = (1088, 1920)
+    RES_1920_1088 = (1920, 1088)
+
+class Model(Enum):
+    NAI_DIFFUSION_3 = "nai-diffusion-3"
+    NAI_DIFFUSION_3_INPAINTING = "nai-diffusion-3-inpainting"
+
+    NAI_DIFFUSION = "nai-diffusion"
+    NAI_DIFFUSION_INPAINTING = "nai-diffusion-inpainting"
+
+    SAFE_DIFFUSION = "safe-diffusion"
+    SAFE_DIFFUSION_INPAINTING = "safe-diffusion-inpainting"
+
+    NAI_DIFFUSION_FURRY = "nai-diffusion-furry"
+    FURRY_DIFFUSION_INPAINTING = "furry-diffusion-inpainting"