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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.3.1
+Version: 0.3.3
 Summary: Novelai Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache License 2.0
@@ -39,7 +39,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 - [x] /user/login
 - [x] /user/information
 - [x] /ai/upscale
-- [ ] /ai/generate-image/suggest-tags
+- [x] /ai/generate-image/suggest-tags
 - [ ] /ai/annotate-image
 - [ ] /ai/classify
 - [ ] /ai/generate-prompt

{novelai_python-0.3.1 → novelai_python-0.3.3}/README.md

@@ -12,7 +12,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 - [x] /user/login
 - [x] /user/information
 - [x] /ai/upscale
-- [ ] /ai/generate-image/suggest-tags
+- [x] /ai/generate-image/suggest-tags
 - [ ] /ai/annotate-image
 - [ ] /ai/classify
 - [ ] /ai/generate-prompt

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

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.3.1"
+version = "0.3.3"
 description = "Novelai Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "coldlando@hotmail.com" },

{novelai_python-0.3.1 → novelai_python-0.3.3}/src/novelai_python/__init__.py

@@ -13,6 +13,7 @@ from ._exceptions import (
 )
 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 @@ __all__ = [
     "Login",
     "LoginResp",
 
+    "SuggestTags",
+    "SuggestTagsResp",
+
     "Information",
     "InformationResp",
 

{novelai_python-0.3.1 → novelai_python-0.3.3}/src/novelai_python/_response/ai/generate_image.py

@@ -22,3 +22,12 @@ class ImageGenerateResp(RespBase):
         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

{novelai_python-0.3.1 → novelai_python-0.3.3}/src/novelai_python/credential/ApiToken.py

@@ -25,6 +25,7 @@ class ApiCredential(CredentialBase):
                 "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",

{novelai_python-0.3.1 → novelai_python-0.3.3}/src/novelai_python/credential/JwtToken.py

@@ -24,7 +24,7 @@ class JwtCredential(CredentialBase):
             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",

{novelai_python-0.3.1 → novelai_python-0.3.3}/src/novelai_python/credential/UserAuth.py

@@ -30,7 +30,7 @@ class LoginCredential(CredentialBase):
             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}",

{novelai_python-0.3.1 → novelai_python-0.3.3}/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

novelai_python-0.3.1/src/novelai_python/sdk/ai/generate_image.py → novelai_python-0.3.3/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 loguru import logger
 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):
@@ -309,7 +236,7 @@ class GenerateImageInfer(ApiBaseModel):
               seed: int = None,
               steps: int = 28,
               scale: float = 6.0,
-              cfg_rescale: int = 0,
+              cfg_rescale: float = 0,
               sampler: Union[Sampler, str] = Sampler.K_EULER,
               width: int = 832,
               height: int = 1216,
@@ -395,7 +322,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"
+            "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,

novelai_python-0.3.3/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"

novelai_python-0.3.3/src/novelai_python/sdk/ai/generate_image/suggest_tags.py

@@ -0,0 +1,127 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/13 下午8:09
+# @Author  : sudoskys
+# @File    : suggest-tags.py
+# @Software: PyCharm
+from typing import Optional, Union
+from urllib.parse import urlparse
+
+import curl_cffi
+import httpx
+from curl_cffi.requests import AsyncSession
+from loguru import logger
+from pydantic import PrivateAttr
+
+from ._enum import Model
+from ...schema import ApiBaseModel
+from ...._exceptions import APIError, AuthError, SessionHttpError
+from ...._response.ai.generate_image import SuggestTagsResp
+from ....credential import CredentialBase
+from ....utils import try_jsonfy
+
+
+class SuggestTags(ApiBaseModel):
+    _endpoint: Optional[str] = PrivateAttr("https://api.novelai.net")
+    model: Model = Model.NAI_DIFFUSION_3
+    prompt: str = "landscape"
+
+    @property
+    def endpoint(self):
+        return self._endpoint
+
+    @endpoint.setter
+    def endpoint(self, value):
+        self._endpoint = value
+
+    @property
+    def base_url(self):
+        return f"{self.endpoint.strip('/')}/ai/generate-image/suggest-tags"
+
+    async def necessary_headers(self, request_data) -> dict:
+        return {
+            "Host": urlparse(self.endpoint).netloc,
+            "Accept": "*/*",
+            "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",
+            "Access-Control-Allow-Origin": "*",
+            "Referer": "https://novelai.net/",
+            "Content-Type": "application/json",
+            "Origin": "https://novelai.net",
+            "Connection": "keep-alive",
+            "Sec-Fetch-Dest": "empty",
+            "Sec-Fetch-Mode": "cors",
+            "Sec-Fetch-Site": "same-site",
+            "Pragma": "no-cache",
+            "Cache-Control": "no-cache",
+            "TE": "trailers",
+        }
+
+    async def request(self,
+                      session: Union[AsyncSession, CredentialBase],
+                      *,
+                      override_headers: Optional[dict] = None
+                      ) -> SuggestTagsResp:
+        """
+        Request to get user subscription information
+        :param override_headers:
+        :param session:
+        :return:
+        """
+        # Data Build
+        request_data = self.model_dump(mode="json", exclude_none=True)
+        if isinstance(session, AsyncSession):
+            session.headers.update(await self.necessary_headers(request_data))
+        elif isinstance(session, CredentialBase):
+            session = await session.get_session(update_headers=await self.necessary_headers(request_data))
+        # Header
+        if override_headers:
+            session.headers.clear()
+            session.headers.update(override_headers)
+        logger.debug("SuggestTags")
+        try:
+            assert hasattr(session, "get"), "session must have get method."
+            response = await session.get(
+                url=self.base_url + "?" + "&".join([f"{k}={v}" for k, v in request_data.items()])
+            )
+            if "application/json" not in response.headers.get('Content-Type') or response.status_code != 200:
+                logger.error(
+                    f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
+                )
+                try:
+                    _msg = response.json()
+                except Exception as e:
+                    logger.warning(e)
+                    if not isinstance(response.content, str) and len(response.content) < 50:
+                        raise APIError(
+                            message=f"Unexpected content type: {response.headers.get('Content-Type')}",
+                            request=request_data,
+                            code=response.status_code,
+                            response=try_jsonfy(response.content)
+                        )
+                    else:
+                        _msg = {"statusCode": response.status_code, "message": 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, code=status_code, response=_msg)
+                if status_code in [500]:
+                    # An unknown error occured.
+                    raise APIError(message, request=request_data, code=status_code, response=_msg)
+                raise APIError(message, request=request_data, code=status_code, response=_msg)
+            return SuggestTagsResp.model_validate(response.json())
+        except curl_cffi.requests.errors.RequestsError as exc:
+            logger.exception(exc)
+            raise SessionHttpError("An AsyncSession RequestsError occurred, maybe SSL error. Try again later!")
+        except httpx.HTTPError as exc:
+            logger.exception(exc)
+            raise SessionHttpError("An HTTPError occurred, maybe SSL error. Try again later!")
+        except APIError as e:
+            raise e
+        except Exception as e:
+            logger.opt(exception=e).exception("An Unexpected error occurred")
+            raise e

{novelai_python-0.3.1 → novelai_python-0.3.3}/src/novelai_python/sdk/ai/upscale.py

@@ -73,7 +73,7 @@ class Upscale(ApiBaseModel):
         return {
             "Host": urlparse(self.endpoint).netloc,
             "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",
             "Referer": "https://novelai.net/",

{novelai_python-0.3.1 → novelai_python-0.3.3}/src/novelai_python/sdk/user/information.py

@@ -4,6 +4,7 @@
 # @File    : information.py
 # @Software: PyCharm
 from typing import Optional, Union
+from urllib.parse import urlparse
 
 import curl_cffi
 import httpx
@@ -35,7 +36,7 @@ class Information(ApiBaseModel):
 
     async def necessary_headers(self, request_data) -> dict:
         return {
-            "Host": "api.novelai.net",
+            "Host": urlparse(self.endpoint).netloc,
             "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",

{novelai_python-0.3.1 → novelai_python-0.3.3}/src/novelai_python/sdk/user/subscription.py

@@ -4,6 +4,7 @@
 # @File    : subscription.py.py
 # @Software: PyCharm
 from typing import Optional, Union
+from urllib.parse import urlparse
 
 import curl_cffi
 import httpx
@@ -35,7 +36,7 @@ class Subscription(ApiBaseModel):
 
     async def necessary_headers(self, request_data) -> dict:
         return {
-            "Host": "api.novelai.net",
+            "Host": urlparse(self.endpoint).netloc,
             "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",

{novelai_python-0.3.1 → novelai_python-0.3.3}/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.generate_image.suggest_tags import SuggestTags
 from .sdk.ai.upscale import Upscale
 from .sdk.user.information import Information
 from .sdk.user.login import Login
@@ -118,6 +119,25 @@ async def upscale(
         return JSONResponse(status_code=500, content=e.__dict__)
 
 
+@app.get("/ai/generate_image/suggest_tags")
+async def suggest_tags(
+        req: SuggestTags,
+        current_token: str = Depends(get_current_token)
+):
+    """
+    生成建议
+    :param current_token: Authorization
+    :param req: SuggestTags
+    :return:
+    """
+    try:
+        _result = await req.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,

{novelai_python-0.3.1 → novelai_python-0.3.3}/src/novelai_python/utils/random_prompt/__init__.py

@@ -13,7 +13,6 @@ from .tag import (
     socks,
     action, expression, footwears, bottoms, color as colors
 )
-from .tag_artist import rankArtist
 from .tag_character import rankMoods, rankCharacter, rankIdentity
 from .tag_nsfw import nsfw
 
@@ -30,9 +29,11 @@ class RandomPromptGenerator(object):
         :param existing_tags:  a list of existing tags
         :return:  a tag
         """
-        valid_tags = [tag
-                      for tag in tags
-                      if len(tag) < 3 or not tag[2] or any(sub_tag in existing_tags for sub_tag in tag[2])]
+        valid_tags = [
+            tag
+            for tag in tags
+            if len(tag) < 3 or not tag[2] or any(sub_tag in existing_tags for sub_tag in tag[2])
+        ]
         total_weight = sum(tagr[1] for tagr in valid_tags if len(tagr) > 1)
         if total_weight == 0:
             if isinstance(tags, list):
@@ -70,7 +71,7 @@ class RandomPromptGenerator(object):
         has_unique_feature = any(feature in features for feature in unique_features)
         if random.random() < 0.1 and enable_skin_color:
             features.append(self.get_weighted_choice(sinkColor, features))
-        if random.random() < 0.8:
+        if random.random() < 0.7:
             features.append(self.get_weighted_choice(eyeColors, features))
         if random.random() < 0.1:
             features.append(self.get_weighted_choice(eyeCharacteristics, features))
@@ -230,12 +231,6 @@ class RandomPromptGenerator(object):
             return ', '.join(tags)
         if random.random() < 0.3:
             tags.append(self.get_weighted_choice(artStyle, tags))
-        if random.random() < 0.6:
-            tags.append("{" + self.get_weighted_choice(rankArtist, tags) + "}")
-            if random.random() < 0.5:
-                tags.append("[" + self.get_weighted_choice(rankArtist, tags) + "]")
-                if random.random() < 0.4:
-                    tags.append("{{" + self.get_weighted_choice(rankArtist, tags) + "}}")
         c_count = 0
         d_count = 0
         u_count = 0
@@ -286,23 +281,23 @@ class RandomPromptGenerator(object):
                 features = []
                 if random.random() < 0.2:
                     features = self.character_features(nsfw['fu'], None, True, irs)
-                if random.random() < 0.6:
+                if random.random() < 0.3:
                     tags.append(self.get_weighted_choice(nsfw["sex"], tags))
-                if random.random() < 0.6:
+                if random.random() < 0.5:
                     tags.append(self.get_weighted_choice(nsfw["pussy"], tags))
-                if random.random() < 0.6:
+                if random.random() < 0.3:
                     tags.append(self.get_weighted_choice(nsfw["sexMod"], tags))
-                if random.random() < 0.6:
+                if random.random() < 0.5:
                     tags.append(self.get_weighted_choice(nsfw["sexActionMode"], tags))
-                if random.random() < 0.2:
+                if random.random() < 0.1:
                     tags.append(self.get_weighted_choice(nsfw["bdsm"], tags))
-                if random.random() < 0.3:
+                if random.random() < 0.2:
                     tags.append(self.get_weighted_choice(nsfw["sexAccessories"], tags))
                 tags.extend(features)
-        if random.random() < 0.65 and enable_moods:
+        if random.random() < 0.4 and enable_moods:
             # 心情
             tags.append("[" + self.get_weighted_choice(rankMoods, tags) + "]")
-        if random.random() < 0.5 and enable_identity:
+        if random.random() < 0.4 and enable_identity:
             # 身份
             tags.append("[" + self.get_weighted_choice(rankIdentity, tags) + "]")
         if random.random() < 0.2 and enable_character:
@@ -315,9 +310,9 @@ class RandomPromptGenerator(object):
                 num_objects = random.randint(1, 3)
                 for _ in range(num_objects):
                     tags.append(self.get_weighted_choice(backgroundObjects, tags))
-        if random.random() < 0.3:
+        if random.random() < 0.1:
             tags.append(self.get_weighted_choice(cameraPerspective, tags))
-        if random.random() < 0.7:
+        if random.random() < 0.5:
             tags.append(self.get_weighted_choice(cameraAngle, tags))
         for _ in range(c_count):
             if random.random() < 0.2:

{novelai_python-0.3.1 → novelai_python-0.3.3}/src/novelai_python/utils/useful.py

@@ -47,20 +47,27 @@ class QueSelect(object):
         return selected
 
 
+import cv2 as cv
+import numpy as np
+
+
 def create_mask_from_sketch(original_img_bytes: bytes,
                             sketch_img_bytes: bytes,
-                            output_format: str = '.png',
-                            jagged_edges: bool = True
+                            min_block_size: int = 15,
+                            jagged_edges: bool = True,
+                            output_format: str = '.png'
                             ) -> 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
+    Function to create a mask from original and sketch images input as bytes. Returns BytesIO object.
+
+    :param original_img_bytes:  Bytes corresponding to the original image.
+    :param sketch_img_bytes:  Bytes corresponding to the sketch image.
+    :param min_block_size:  The minimum size of the pixel blocks. Default is 1, i.e., no pixelation.
+    :param jagged_edges:  If set to True, the edges of the resulting mask will be more jagged.
+    :param output_format:  Format of the output image. Defaults to '.png'. It could also be '.jpg'
+    :return:  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)
@@ -78,6 +85,7 @@ def create_mask_from_sketch(original_img_bytes: bytes,
     # Threshold to create the mask
     _, thresh = cv.threshold(diff_gray, 10, 255, cv.THRESH_BINARY)
 
+    # Create bigger kernel for morphological operations
     if jagged_edges:
         # Use a bigger kernel for dilation to create larger 'step' effect at the edges
         kernel = np.ones((7, 7), np.uint8)
@@ -87,8 +95,20 @@ def create_mask_from_sketch(original_img_bytes: bytes,
     # Perform morphological opening to remove small noise
     opening = cv.morphologyEx(thresh, cv.MORPH_OPEN, kernel, iterations=2)
 
+    # Perform morphological dilation
+    dilation = cv.dilate(opening, kernel, iterations=2)
+
+    # Perform morphological closing to connect separated areas
+    closing = cv.morphologyEx(dilation, cv.MORPH_CLOSE, kernel, iterations=3)
+
     # Further remove noise with a Gaussian filter
-    smooth = cv.GaussianBlur(opening, (5, 5), 0)
+    smooth = cv.GaussianBlur(closing, (5, 5), 0)
+
+    if min_block_size > 1:
+        # Resize to smaller image, then resize back to original size to create a pixelated effect
+        small = cv.resize(smooth, (smooth.shape[1] // min_block_size, smooth.shape[0] // min_block_size),
+                          interpolation=cv.INTER_LINEAR)
+        smooth = cv.resize(small, (smooth.shape[1], smooth.shape[0]), interpolation=cv.INTER_NEAREST)
 
     if jagged_edges:
         # Apply additional thresholding to create sharper, jagged edges

novelai_python-0.3.1/src/novelai_python/utils/random_prompt/tag_artist.py

@@ -1,398 +0,0 @@
-# -*- coding: utf-8 -*-
-# @Time    : 2024/1/26 下午5:08
-# @Author  : sudoskys
-# @File    : tag_artist.py
-# @Software: PyCharm
-
-rankArtist = [
-    [
-        "redi_(rasec_asdjh)",
-        50
-    ],
-    [
-        "suoiresnu",
-        41
-    ],
-    [
-        "kulo_jawa",
-        30
-    ],
-    [
-        "sakanaya_(sakanaya952)",
-        25
-    ],
-    [
-        "momisan",
-        20
-    ],
-    [
-        "tansuan_(ensj3875)",
-        19
-    ],
-    [
-        "hitmanatee",
-        18
-    ],
-    [
-        "yoru0409",
-        18
-    ],
-    [
-        "zovokia",
-        14
-    ],
-    [
-        "memeh",
-        13
-    ],
-    [
-        "honashi",
-        13
-    ],
-    [
-        "pinkmill",
-        13
-    ],
-    [
-        "kanjy00u",
-        12
-    ],
-    [
-        "eckert&eich",
-        12
-    ],
-    [
-        "ningen mame",
-        100
-    ],
-    [
-        "ciloranko",
-        100
-    ],
-    [
-        "sho (sho lwlw)",
-        80
-    ],
-    [
-        "tianliang duohe fangdongye",
-        100
-    ],
-    [
-        "alphonse (white datura)",
-        80
-    ],
-    [
-        "rei (sanbonzakura)",
-        100
-    ],
-    [
-        "mamyouda",
-        80
-    ],
-    [
-        "ogipote",
-        100
-    ],
-    [
-        "fuzichoco",
-        100
-    ],
-    [
-        "norza",
-        80
-    ],
-    [
-        "AkiZero1510",
-        100
-    ],
-    [
-        "ATDAN",
-        100
-    ],
-    [
-        "mignon",
-        80
-    ],
-    [
-        "shexyo",
-        100
-    ],
-    [
-        "musashi (dishwasher1910)",
-        100
-    ],
-    [
-        "modare",
-        80
-    ],
-    [
-        "asteroidill",
-        80
-    ],
-    [
-        "potg(piotegu)",
-        80
-    ],
-    [
-        "kedama milk",
-        100
-    ],
-    [
-        "jun (aousa0328)",
-        80
-    ],
-    [
-        "sky freedom",
-        80
-    ],
-    [
-        "ask (askzy)",
-        100
-    ],
-    [
-        "toridamono",
-        100
-    ],
-    [
-        "sencha (senchat)",
-        80
-    ],
-    [
-        "void_0",
-        80
-    ],
-    [
-        "mandrill",
-        100
-    ],
-    [
-        "tokkyu",
-        100
-    ],
-    [
-        "houkisei",
-        80
-    ],
-    [
-        "konya karasue",
-        100
-    ],
-    [
-        "akakura",
-        100
-    ],
-    [
-        "kawacy",
-        100
-    ],
-    [
-        "reoen",
-        100
-    ],
-    [
-        "nixeu",
-        100
-    ],
-    [
-        "cierra_(ra-bit)",
-        100
-    ],
-    [
-        "artistask_(askzy)",
-        100
-    ],
-    [
-        "gomzi",
-        100
-    ],
-    [
-        "you_shimizu",
-        100
-    ],
-    [
-        "demizu_posuka",
-        100
-    ],
-    [
-        "asakuraf",
-        100
-    ],
-    [
-        "ningen_mame",
-        100
-    ],
-    [
-        "sho_(sho_lwlw)",
-        100
-    ],
-    [
-        "alphonse_(white_datura)",
-        100
-    ],
-    [
-        "rei_(sanbonzakura)",
-        100
-    ],
-    [
-        "wlop",
-        100
-    ],
-    [
-        "as109",
-        100
-    ],
-    [
-        "Milkychu",
-        100
-    ],
-    [
-        "Cogecha",
-        100
-    ],
-    [
-        "liduke",
-        100
-    ],
-    [
-        "infukun",
-        100
-    ],
-    [
-        "amazuyu_tatsuki",
-        100
-    ],
-    [
-        "Happ\u014dbi Jin",
-        100
-    ],
-    [
-        "Rella",
-        100
-    ],
-    [
-        "omone hokoma agm",
-        100
-    ],
-    [
-        "toosaka asagi",
-        100
-    ],
-    [
-        "Mochizuki Kei",
-        100
-    ],
-    [
-        "Namie",
-        100
-    ],
-    [
-        "Hoji",
-        100
-    ],
-    [
-        "sy4",
-        100
-    ],
-    [
-        "mika_pikazo",
-        100
-    ],
-    [
-        "momoko(momopoco)",
-        100
-    ],
-    [
-        "rurudo",
-        100
-    ],
-    [
-        "missile228",
-        100
-    ],
-    [
-        "miyase_mahiro",
-        100
-    ],
-    [
-        "touzai(poppin phl95)",
-        100
-    ],
-    [
-        "kukka",
-        100
-    ],
-    [
-        "kantoku",
-        100
-    ],
-    [
-        "poco(asahi age)",
-        100
-    ],
-    [
-        "refeia",
-        100
-    ],
-    [
-        "ke-ta",
-        100
-    ],
-    [
-        "anmi",
-        100
-    ],
-    [
-        "tiv",
-        100
-    ],
-    [
-        "chen_bin",
-        100
-    ],
-    [
-        "lm7 (op-center)",
-        100
-    ],
-    [
-        "shion (mirudakemann)",
-        100
-    ],
-    [
-        "asteroid ill",
-        100
-    ],
-    [
-        "potg (piotegu)",
-        100
-    ],
-    [
-        "KinHasu",
-        100
-    ],
-    [
-        "ryou_(ryoutarou)",
-        100
-    ],
-    [
-        "mamimi_(mamamimi)",
-        100
-    ],
-]
-
-if __name__ == "__main__":
-    var = vars().copy()
-
-    # 获取tags列表
-    tags = []
-    for key, value in var.items():
-        if not key.startswith('_') and key != 'var':
-            tags.extend(value)
-    print(tags)
-
-    # 展开二级列表
-    tags = [tag[0] for tag in tags]
-    print(tags)
-    import json
-
-    with open('../../../../playground/cos/tag_artist.json', 'w') as f:
-        json.dump(tags, f)