novelai-python 0.4.15__tar.gz → 0.5.0__tar.gz

{novelai_python-0.4.15 → novelai_python-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: novelai-python
-Version: 0.4.15
+Version: 0.5.0
 Summary: NovelAI Python Binding With Pydantic
 Author-Email: sudoskys <coldlando@hotmail.com>
 License: Apache License 2.0
@@ -27,6 +27,8 @@ Requires-Dist: pynacl>=1.5.0
 Requires-Dist: ftfy>=6.2.0
 Requires-Dist: regex>=2023.12.25
 Requires-Dist: tokenizers>=0.15.2
+Requires-Dist: json-repair>=0.29.4
+Requires-Dist: robust-downloader>=0.0.2
 Provides-Extra: testing
 Requires-Dist: pytest>=7.4.4; extra == "testing"
 Requires-Dist: pytest-asyncio>=0.23.4; extra == "testing"
@@ -50,6 +52,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 - [x] tool.random_prompt
 - [x] tool.paint_mask
 - [x] tool.image_metadata
+- [x] tokenizer
 - [x] /ai/generate-image
 - [x] /user/subscription
 - [x] /user/login
@@ -119,7 +122,7 @@ from dotenv import load_dotenv
 from pydantic import SecretStr
 
 from novelai_python import APIError, LoginCredential
-from novelai_python.sdk.ai.generate import TextLLMModel, LLM
+from novelai_python.sdk.ai.generate import TextLLMModel, LLM, get_default_preset
 
 load_dotenv()
 username = os.getenv("NOVELAI_USER", None)
@@ -133,7 +136,13 @@ login_credential = LoginCredential(
 
 async def chat(prompt: str):
     try:
-        agent = LLM.build(prompt=prompt, model=TextLLMModel.Kayra)
+        model = TextLLMModel.ERATO  # llama3
+        parameters = get_default_preset(model).parameters
+        agent = LLM.build(
+            prompt=prompt,
+            model=model,
+            parameters=None  # Auto Select or get from preset
+        )
         result = await agent.request(session=login_credential)
     except APIError as e:
         raise Exception(f"Error: {e.message}")
@@ -160,6 +169,21 @@ pip install novelai_python
 python3 -m novelai_python.server -h '127.0.0.1' -p 7888
 ```
 
+#### Tokenizer
+
+```python
+from novelai_python._enum import get_tokenizer_model, TextLLMModel
+from novelai_python.tokenizer import NaiTokenizer
+
+tokenizer_package = NaiTokenizer(get_tokenizer_model(TextLLMModel.ERATO))
+t_text = "a fox jumped over the lazy dog"
+encode_tokens = tokenizer_package.encode(t_text)
+print(tokenizer_package.tokenize_text(t_text))
+print(f"Tokenized text: {encode_tokens}")
+print(tokenizer_package.decode(tokenizer_package.encode(t_text)))
+
+```
+
 ## Acknowledgements 🙏
 
 [BackEnd](https://api.novelai.net/docs)

{novelai_python-0.4.15 → novelai_python-0.5.0}/README.md

@@ -16,6 +16,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 - [x] tool.random_prompt
 - [x] tool.paint_mask
 - [x] tool.image_metadata
+- [x] tokenizer
 - [x] /ai/generate-image
 - [x] /user/subscription
 - [x] /user/login
@@ -85,7 +86,7 @@ from dotenv import load_dotenv
 from pydantic import SecretStr
 
 from novelai_python import APIError, LoginCredential
-from novelai_python.sdk.ai.generate import TextLLMModel, LLM
+from novelai_python.sdk.ai.generate import TextLLMModel, LLM, get_default_preset
 
 load_dotenv()
 username = os.getenv("NOVELAI_USER", None)
@@ -99,7 +100,13 @@ login_credential = LoginCredential(
 
 async def chat(prompt: str):
     try:
-        agent = LLM.build(prompt=prompt, model=TextLLMModel.Kayra)
+        model = TextLLMModel.ERATO  # llama3
+        parameters = get_default_preset(model).parameters
+        agent = LLM.build(
+            prompt=prompt,
+            model=model,
+            parameters=None  # Auto Select or get from preset
+        )
         result = await agent.request(session=login_credential)
     except APIError as e:
         raise Exception(f"Error: {e.message}")
@@ -126,6 +133,21 @@ pip install novelai_python
 python3 -m novelai_python.server -h '127.0.0.1' -p 7888
 ```
 
+#### Tokenizer
+
+```python
+from novelai_python._enum import get_tokenizer_model, TextLLMModel
+from novelai_python.tokenizer import NaiTokenizer
+
+tokenizer_package = NaiTokenizer(get_tokenizer_model(TextLLMModel.ERATO))
+t_text = "a fox jumped over the lazy dog"
+encode_tokens = tokenizer_package.encode(t_text)
+print(tokenizer_package.tokenize_text(t_text))
+print(f"Tokenized text: {encode_tokens}")
+print(tokenizer_package.decode(tokenizer_package.encode(t_text)))
+
+```
+
 ## Acknowledgements 🙏
 
 [BackEnd](https://api.novelai.net/docs)

{novelai_python-0.4.15 → novelai_python-0.5.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.4.15"
+version = "0.5.0"
 description = "NovelAI Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "coldlando@hotmail.com" },
@@ -26,6 +26,8 @@ dependencies = [
     "ftfy>=6.2.0",
     "regex>=2023.12.25",
     "tokenizers>=0.15.2",
+    "json-repair>=0.29.4",
+    "robust-downloader>=0.0.2",
 ]
 requires-python = ">=3.9"
 readme = "README.md"

{novelai_python-0.4.15 → novelai_python-0.5.0}/src/novelai_python/__init__.py

@@ -2,7 +2,6 @@
 # @Time    : 2023/11/18 上午12:18
 # @Author  : sudoskys
 # @File    : __init__.py
-
 from ._exceptions import (
     NovelAiError,
     APIError,

novelai_python-0.5.0/src/novelai_python/_enum.py

@@ -0,0 +1,94 @@
+from enum import Enum
+from typing import Optional, Union
+
+
+class TextLLMModel(Enum):
+    NEO_2B = "2.7B"
+    J_6B = "6B"
+    J_6B_V3 = "6B-v3"
+    J_6B_V4 = "6B-v4"
+    GENJI_PYTHON_6B = "genji-python-6b"
+    GENJI_JP_6B = "genji-jp-6b"
+    GENJI_JP_6B_V2 = "genji-jp-6b-v2"
+    EUTERPE_V0 = "euterpe-v0"
+    EUTERPE_V2 = "euterpe-v2"
+    KRAKE_V1 = "krake-v1"
+    KRAKE_V2 = "krake-v2"
+    BLUE = "blue"
+    RED = "red"
+    GREEN = "green"
+    PURPLE = "purple"
+    PINK = "pink"
+    YELLOW = "yellow"
+    WHITE = "white"
+    BLACK = "black"
+    CASSANDRA = "cassandra"
+    COMMENT_BOT = "hypebot"
+    INFILL = "infillmodel"
+    CLIO = "clio-v1"
+    KAYRA = "kayra-v1"
+    ERATO = "llama-3-erato-v1"
+
+
+class TextTokenizerGroup(object):
+    GENJI = "genji_tokenizer.def"
+    PILE = "pile_tokenizer.def"
+    PILE_NAI = "pile_tokenizer.def"
+    NAI_INLINE = "gpt2_tokenizer.def"
+    NERDSTASH_V2 = "nerdstash_tokenizer_v2.def"
+    NERDSTASH = "nerdstash_tokenizer.def"
+    LLAMA3 = "llama3_tokenizer.def"
+    GPT2 = "gpt2_tokenizer.def"
+    CLIP = "clip_tokenizer.def"
+
+
+TextLLMModelTypeAlias = Union[TextLLMModel, str]
+
+TOKENIZER_MODEL_MAP = {
+    TextLLMModel.GENJI_JP_6B_V2: TextTokenizerGroup.GENJI,
+    TextLLMModel.CASSANDRA: TextTokenizerGroup.PILE,
+    TextLLMModel.KRAKE_V2: TextTokenizerGroup.PILE,
+    TextLLMModel.INFILL: TextTokenizerGroup.NAI_INLINE,
+    TextLLMModel.KAYRA: TextTokenizerGroup.NERDSTASH_V2,
+    TextLLMModel.BLUE: TextTokenizerGroup.NERDSTASH_V2,
+    TextLLMModel.PINK: TextTokenizerGroup.NERDSTASH_V2,
+    TextLLMModel.YELLOW: TextTokenizerGroup.NERDSTASH_V2,
+    TextLLMModel.RED: TextTokenizerGroup.NERDSTASH_V2,
+    TextLLMModel.GREEN: TextTokenizerGroup.NERDSTASH_V2,
+    TextLLMModel.BLACK: TextTokenizerGroup.NERDSTASH_V2,
+    TextLLMModel.CLIO: TextTokenizerGroup.NERDSTASH,
+    TextLLMModel.PURPLE: TextTokenizerGroup.LLAMA3,
+    TextLLMModel.WHITE: TextTokenizerGroup.LLAMA3,
+    TextLLMModel.ERATO: TextTokenizerGroup.LLAMA3,
+}
+
+COLORS_LLM = [
+    TextLLMModel.BLUE,
+    TextLLMModel.RED,
+    TextLLMModel.GREEN,
+    TextLLMModel.PURPLE,
+    TextLLMModel.PINK,
+    TextLLMModel.YELLOW,
+    TextLLMModel.WHITE,
+    TextLLMModel.BLACK,
+]
+
+
+def get_llm_group(model: TextLLMModel) -> Optional[TextTokenizerGroup]:
+    if isinstance(model, str):
+        model = TextLLMModel(model)
+    return TOKENIZER_MODEL_MAP.get(model, None)
+
+
+def get_tokenizer_model(model: TextLLMModel) -> str:
+    if isinstance(model, str):
+        model = TextLLMModel(model)
+    group = TOKENIZER_MODEL_MAP.get(model, TextTokenizerGroup.GPT2)
+    return group
+
+
+def get_tokenizer_model_url(model: TextLLMModel) -> str:
+    model_name = get_tokenizer_model(model)
+    if not model_name.endswith(".def"):
+        model_name = f"{model_name}.def"
+    return f"https://novelai.net/tokenizer/compressed/{model_name}?v=2&static=true"

{novelai_python-0.4.15 → novelai_python-0.5.0}/src/novelai_python/_response/ai/generate.py

@@ -2,8 +2,8 @@ from typing import TYPE_CHECKING
 
 from pydantic import BaseModel, ConfigDict
 
-from novelai_python.sdk.ai.generate._enum import TOKENIZER, TextLLMModel  # noqa
-from novelai_python.tokenizer import LLMTokenizer
+from novelai_python._enum import get_tokenizer_model, TextLLMModel
+from novelai_python.tokenizer import NaiTokenizer
 from novelai_python.utils.encode import b64_to_tokens
 
 if TYPE_CHECKING:
@@ -20,4 +20,7 @@ class LLMResp(BaseModel):
 
     @staticmethod
     def decode_token(token_str, model: TextLLMModel) -> str:
-        return LLMTokenizer().decode(b64_to_tokens(token_str), tokenizer_name=TOKENIZER.get(model))
+        dtype = 'uint32' if model in [TextLLMModel.ERATO] else 'uint16'
+        return NaiTokenizer(model=get_tokenizer_model(model)).decode(
+            b64_to_tokens(token_str, dtype=dtype)
+        )

{novelai_python-0.4.15 → novelai_python-0.5.0}/src/novelai_python/_response/ai/generate_stream.py

@@ -2,8 +2,8 @@ from typing import Optional, Any
 
 from pydantic import BaseModel, ConfigDict
 
-from novelai_python.sdk.ai.generate._enum import TOKENIZER, TextLLMModel  # noqa
-from novelai_python.tokenizer import LLMTokenizer
+from novelai_python._enum import get_tokenizer_model, TextLLMModel
+from novelai_python.tokenizer import NaiTokenizer
 from novelai_python.utils.encode import b64_to_tokens
 
 
@@ -20,4 +20,7 @@ class LLMStreamResp(BaseModel):
 
     @staticmethod
     def decode(token_str, model: TextLLMModel) -> str:
-        return LLMTokenizer().decode(b64_to_tokens(token_str), tokenizer_name=TOKENIZER.get(model))
+        dtype = 'uint32' if model in [TextLLMModel.ERATO] else 'uint16'
+        return NaiTokenizer(model=get_tokenizer_model(model)).decode(
+            b64_to_tokens(token_str, dtype=dtype)
+        )

{novelai_python-0.4.15 → novelai_python-0.5.0}/src/novelai_python/sdk/ai/_cost.py

@@ -1,11 +1,11 @@
 import math
 import random
-from typing import List, Optional
+from typing import List, Optional, Union
 
 from pydantic import BaseModel
 
 from novelai_python.sdk.ai._const import map, initialN, initial_n, step, newN
-from novelai_python.sdk.ai._enum import Sampler, Model, ModelGroups, get_model_group
+from novelai_python.sdk.ai._enum import Sampler, Model, ModelGroups, get_model_group, ModelTypeAlias
 
 
 class Args(BaseModel):
@@ -43,9 +43,10 @@ class CostCalculator:
             is_sm_dynamic: bool,
             is_account_active: bool,
             sampler: Optional[Sampler],
-            model: Optional[Model],
+            model: ModelTypeAlias,
             tool: str = None,
-            is_tool_active: bool = False) -> int:
+            is_tool_active: bool = False
+    ) -> int:
         return CostCalculator.calculate_cost(
             Args(
                 height=height,
@@ -59,7 +60,7 @@ class CostCalculator:
                 sampler=sampler,
                 tool=tool
             ),
-            model_group=get_model_group(model.value) if model else None,
+            model_group=get_model_group(model) if model else None,
             is_account_active=is_account_active,
             account_tier=account_tier,
             is_tool_active=is_tool_active

{novelai_python-0.4.15 → novelai_python-0.5.0}/src/novelai_python/sdk/ai/_enum.py

@@ -4,7 +4,7 @@
 # @File    : _enum.py
 # @Software: PyCharm
 from enum import Enum, IntEnum
-from typing import List
+from typing import List, Optional, Union
 
 from pydantic.dataclasses import dataclass
 
@@ -36,56 +36,6 @@ class NoiseSchedule(Enum):
     POLYEXPONENTIAL = "polyexponential"
 
 
-def get_supported_noise_schedule(sample_type: Sampler) -> List[NoiseSchedule]:
-    """
-    Get supported noise schedule for a given sample type
-    :param sample_type: Sampler
-    :return: List[NoiseSchedule]
-    """
-    if sample_type in [
-        Sampler.K_EULER_ANCESTRAL,
-        Sampler.K_DPMPP_2S_ANCESTRAL,
-        Sampler.K_DPMPP_2M,
-        Sampler.K_DPMPP_2M_SDE,
-        Sampler.K_DPMPP_SDE,
-        Sampler.K_EULER
-    ]:
-        return [
-            NoiseSchedule.NATIVE,
-            NoiseSchedule.KARRAS,
-            NoiseSchedule.EXPONENTIAL,
-            NoiseSchedule.POLYEXPONENTIAL
-        ]
-    elif sample_type in [Sampler.K_DPM_2]:
-        return [
-            NoiseSchedule.EXPONENTIAL,
-            NoiseSchedule.POLYEXPONENTIAL
-        ]
-    else:
-        return []
-
-
-def get_default_noise_schedule(sample_type: Sampler) -> NoiseSchedule:
-    """
-    Get default noise schedule for a given sample type
-    :param sample_type: Sampler
-    :return: NoiseSchedule
-    """
-    if sample_type in [
-        Sampler.K_EULER_ANCESTRAL,
-        Sampler.K_DPMPP_2S_ANCESTRAL,
-        Sampler.K_DPMPP_2M,
-        Sampler.K_DPMPP_2M_SDE,
-        Sampler.K_DPMPP_SDE,
-        Sampler.K_EULER
-    ]:
-        return NoiseSchedule.KARRAS
-    elif sample_type in [Sampler.K_DPM_2]:
-        return NoiseSchedule.EXPONENTIAL
-    else:
-        return NoiseSchedule.NATIVE
-
-
 class UCPreset(IntEnum):
     TYPE0 = 0
     TYPE1 = 1
@@ -171,30 +121,6 @@ class ModelGroups(Enum):
     STABLE_DIFFUSION_XL_FURRY = "stable_diffusion_xl_furry"
 
 
-def get_model_group(model: str) -> ModelGroups:
-    if isinstance(model, Enum):
-        model = model.value
-    mapping = {
-        "stable-diffusion": ModelGroups.STABLE_DIFFUSION,
-        "nai-diffusion": ModelGroups.STABLE_DIFFUSION,
-        "safe-diffusion": ModelGroups.STABLE_DIFFUSION,
-        "waifu-diffusion": ModelGroups.STABLE_DIFFUSION,
-        "nai-diffusion-furry": ModelGroups.STABLE_DIFFUSION,
-        "curated-diffusion-test": ModelGroups.STABLE_DIFFUSION,
-        "nai-diffusion-inpainting": ModelGroups.STABLE_DIFFUSION,
-        "safe-diffusion-inpainting": ModelGroups.STABLE_DIFFUSION,
-        "furry-diffusion-inpainting": ModelGroups.STABLE_DIFFUSION,
-        "nai-diffusion-2": ModelGroups.STABLE_DIFFUSION_GROUP_2,
-        "nai-diffusion-xl": ModelGroups.STABLE_DIFFUSION_XL,
-        "nai-diffusion-3": ModelGroups.STABLE_DIFFUSION_XL,
-        "nai-diffusion-3-inpainting": ModelGroups.STABLE_DIFFUSION_XL,
-        "custom": ModelGroups.STABLE_DIFFUSION_XL,
-        "nai-diffusion-furry-3": ModelGroups.STABLE_DIFFUSION_XL_FURRY,
-        "nai-diffusion-furry-3-inpainting": ModelGroups.STABLE_DIFFUSION_XL_FURRY,
-    }
-    return mapping.get(model, ModelGroups.STABLE_DIFFUSION)
-
-
 INPAINTING_MODEL_LIST = [
     Model.NAI_DIFFUSION_3_INPAINTING,
     Model.NAI_DIFFUSION_INPAINTING,
@@ -224,8 +150,88 @@ PROMOTION = {
     "Stable Diffusion XL 9CC2F394": Model.NAI_DIFFUSION_FURRY_3,
 }
 
+ModelTypeAlias = Optional[Union[Model, str]]
+ImageBytesTypeAlias = Optional[Union[str, bytes]]
+UCPresetTypeAlias = Optional[Union[UCPreset, int]]
+
+
+def get_supported_noise_schedule(sample_type: Sampler) -> List[NoiseSchedule]:
+    """
+    Get supported noise schedule for a given sample type
+    :param sample_type: Sampler
+    :return: List[NoiseSchedule]
+    """
+    if sample_type in [
+        Sampler.K_EULER_ANCESTRAL,
+        Sampler.K_DPMPP_2S_ANCESTRAL,
+        Sampler.K_DPMPP_2M,
+        Sampler.K_DPMPP_2M_SDE,
+        Sampler.K_DPMPP_SDE,
+        Sampler.K_EULER
+    ]:
+        return [
+            NoiseSchedule.NATIVE,
+            NoiseSchedule.KARRAS,
+            NoiseSchedule.EXPONENTIAL,
+            NoiseSchedule.POLYEXPONENTIAL
+        ]
+    elif sample_type in [Sampler.K_DPM_2]:
+        return [
+            NoiseSchedule.EXPONENTIAL,
+            NoiseSchedule.POLYEXPONENTIAL
+        ]
+    else:
+        return []
+
+
+def get_default_noise_schedule(sample_type: Sampler) -> NoiseSchedule:
+    """
+    Get default noise schedule for a given sample type
+    :param sample_type: Sampler
+    :return: NoiseSchedule
+    """
+    if sample_type in [
+        Sampler.K_EULER_ANCESTRAL,
+        Sampler.K_DPMPP_2S_ANCESTRAL,
+        Sampler.K_DPMPP_2M,
+        Sampler.K_DPMPP_2M_SDE,
+        Sampler.K_DPMPP_SDE,
+        Sampler.K_EULER
+    ]:
+        return NoiseSchedule.KARRAS
+    elif sample_type in [Sampler.K_DPM_2]:
+        return NoiseSchedule.EXPONENTIAL
+    else:
+        return NoiseSchedule.NATIVE
+
+
+def get_model_group(model: ModelTypeAlias) -> ModelGroups:
+    if isinstance(model, Enum):
+        model = model.value
+    else:
+        model = str(model)
+    mapping = {
+        "stable-diffusion": ModelGroups.STABLE_DIFFUSION,
+        "nai-diffusion": ModelGroups.STABLE_DIFFUSION,
+        "safe-diffusion": ModelGroups.STABLE_DIFFUSION,
+        "waifu-diffusion": ModelGroups.STABLE_DIFFUSION,
+        "nai-diffusion-furry": ModelGroups.STABLE_DIFFUSION,
+        "curated-diffusion-test": ModelGroups.STABLE_DIFFUSION,
+        "nai-diffusion-inpainting": ModelGroups.STABLE_DIFFUSION,
+        "safe-diffusion-inpainting": ModelGroups.STABLE_DIFFUSION,
+        "furry-diffusion-inpainting": ModelGroups.STABLE_DIFFUSION,
+        "nai-diffusion-2": ModelGroups.STABLE_DIFFUSION_GROUP_2,
+        "nai-diffusion-xl": ModelGroups.STABLE_DIFFUSION_XL,
+        "nai-diffusion-3": ModelGroups.STABLE_DIFFUSION_XL,
+        "nai-diffusion-3-inpainting": ModelGroups.STABLE_DIFFUSION_XL,
+        "custom": ModelGroups.STABLE_DIFFUSION_XL,
+        "nai-diffusion-furry-3": ModelGroups.STABLE_DIFFUSION_XL_FURRY,
+        "nai-diffusion-furry-3-inpainting": ModelGroups.STABLE_DIFFUSION_XL_FURRY,
+    }
+    return mapping.get(model, ModelGroups.STABLE_DIFFUSION)
+
 
-def get_default_uc_preset(model: str, uc_preset: int) -> str:
+def get_default_uc_preset(model: ModelTypeAlias, uc_preset: int) -> str:
     if isinstance(model, Enum):
         model = model.value
     if isinstance(uc_preset, Enum):