novelai-python 0.6.1__tar.gz → 0.7.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (86) hide show
  1. {novelai_python-0.6.1 → novelai_python-0.7.0}/PKG-INFO +18 -7
  2. {novelai_python-0.6.1 → novelai_python-0.7.0}/README.md +17 -6
  3. {novelai_python-0.6.1 → novelai_python-0.7.0}/pyproject.toml +2 -1
  4. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_enum.py +1 -0
  5. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/generate/__init__.py +2 -1
  6. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/generate/_schema.py +1 -1
  7. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/generate_image/__init__.py +2 -1
  8. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/tokenizer/__init__.py +15 -2
  9. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/tool/image_metadata/__init__.py +160 -111
  10. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/tool/image_metadata/lsb_extractor.py +41 -13
  11. novelai_python-0.7.0/src/novelai_python/tool/random_prompt/__init__.py +95 -0
  12. novelai_python-0.7.0/src/novelai_python/tool/random_prompt/generate_scene_composition.py +2960 -0
  13. novelai_python-0.7.0/src/novelai_python/tool/random_prompt/generate_scene_tags.py +2276 -0
  14. novelai_python-0.7.0/src/novelai_python/tool/random_prompt/generate_tags.py +2517 -0
  15. novelai_python-0.7.0/tests/test_random_prompt.py +61 -0
  16. novelai_python-0.7.0/tests/test_tokenizer.py +32 -0
  17. novelai_python-0.6.1/src/novelai_python/tool/image_metadata/_type.py +0 -21
  18. novelai_python-0.6.1/src/novelai_python/tool/random_prompt/__init__.py +0 -367
  19. novelai_python-0.6.1/src/novelai_python/tool/random_prompt/tag.py +0 -7542
  20. novelai_python-0.6.1/src/novelai_python/tool/random_prompt/tag_character.py +0 -270
  21. novelai_python-0.6.1/src/novelai_python/tool/random_prompt/tag_nsfw.py +0 -984
  22. novelai_python-0.6.1/tests/test_random_prompt.py +0 -70
  23. {novelai_python-0.6.1 → novelai_python-0.7.0}/LICENSE +0 -0
  24. {novelai_python-0.6.1 → novelai_python-0.7.0}/NOTICE.MD +0 -0
  25. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/__init__.py +0 -0
  26. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_exceptions.py +0 -0
  27. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_response/__init__.py +0 -0
  28. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_response/ai/__init__.py +0 -0
  29. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_response/ai/generate.py +0 -0
  30. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_response/ai/generate_image.py +0 -0
  31. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_response/ai/generate_stream.py +0 -0
  32. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_response/ai/generate_voice.py +0 -0
  33. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_response/ai/upscale.py +0 -0
  34. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_response/schema.py +0 -0
  35. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_response/user/__init__.py +0 -0
  36. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_response/user/information.py +0 -0
  37. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_response/user/login.py +0 -0
  38. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/_response/user/subscription.py +0 -0
  39. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/credential/ApiToken.py +0 -0
  40. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/credential/JwtToken.py +0 -0
  41. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/credential/UserAuth.py +0 -0
  42. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/credential/__init__.py +0 -0
  43. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/credential/_base.py +0 -0
  44. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/__init__.py +0 -0
  45. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/LICENSE +0 -0
  46. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/__init__.py +0 -0
  47. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/_const.py +0 -0
  48. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/_cost.py +0 -0
  49. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/_enum.py +0 -0
  50. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/augment_image/__init__.py +0 -0
  51. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/augment_image/_enum.py +0 -0
  52. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/generate/_const.py +0 -0
  53. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/generate/_enum.py +0 -0
  54. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/generate/themes.json +0 -0
  55. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/generate_image/schema.py +0 -0
  56. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/generate_image/suggest_tags.py +0 -0
  57. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/generate_stream.py +0 -0
  58. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/generate_voice/__init__.py +0 -0
  59. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/generate_voice/_enum.py +0 -0
  60. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/ai/upscale.py +0 -0
  61. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/schema.py +0 -0
  62. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/user/__init__.py +0 -0
  63. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/user/information.py +0 -0
  64. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/user/login.py +0 -0
  65. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/sdk/user/subscription.py +0 -0
  66. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/server.py +0 -0
  67. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/tokenizer/clip_simple_tokenizer.py +0 -0
  68. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/tokenizer/nerdstash_v1.model +0 -0
  69. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/tokenizer/nerdstash_v2.model +0 -0
  70. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/tokenizer/novelai.model +0 -0
  71. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/tokenizer/novelai_v2.model +0 -0
  72. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/tool/__init__.py +0 -0
  73. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/tool/image_metadata/bch_utils.py +0 -0
  74. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/tool/image_metadata/lsb_injector.py +0 -0
  75. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/tool/paint_mask/__init__.py +0 -0
  76. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/utils/__init__.py +0 -0
  77. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/utils/encode.py +0 -0
  78. {novelai_python-0.6.1 → novelai_python-0.7.0}/src/novelai_python/utils/useful.py +0 -0
  79. {novelai_python-0.6.1 → novelai_python-0.7.0}/tests/__init__.py +0 -0
  80. {novelai_python-0.6.1 → novelai_python-0.7.0}/tests/test_generate_voice.py +0 -0
  81. {novelai_python-0.6.1 → novelai_python-0.7.0}/tests/test_server.py +0 -0
  82. {novelai_python-0.6.1 → novelai_python-0.7.0}/tests/test_server_run.py +0 -0
  83. {novelai_python-0.6.1 → novelai_python-0.7.0}/tests/test_upscale.py +0 -0
  84. {novelai_python-0.6.1 → novelai_python-0.7.0}/tests/test_user_information.py +0 -0
  85. {novelai_python-0.6.1 → novelai_python-0.7.0}/tests/test_user_login.py +0 -0
  86. {novelai_python-0.6.1 → novelai_python-0.7.0}/tests/test_user_subscription.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: novelai-python
3
- Version: 0.6.1
3
+ Version: 0.7.0
4
4
  Summary: NovelAI Python Binding With Pydantic
5
5
  Author-Email: sudoskys <coldlando@hotmail.com>
6
6
  License: Apache License 2.0
@@ -45,7 +45,8 @@ Description-Content-Type: text/markdown
45
45
 
46
46
  ✨ NovelAI api python sdk with Pydantic, modern and user-friendly.
47
47
 
48
- The goal of this repository is to use Pydantic to build legitimate requests to access the NovelAI API service.
48
+ The goal of this repository is to use Pydantic to build legitimate requests to access
49
+ the [NovelAI API service](https://api.novelai.net/docs).
49
50
 
50
51
  > Python >= 3.9 is required.
51
52
 
@@ -60,6 +61,8 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
60
61
  pip -U install novelai-python
61
62
  ```
62
63
 
64
+ All API users must adhere to the NovelAI Terms of Service: https://novelai.net/terms.
65
+
63
66
  **More examples can be found in the [playground](https://github.com/LlmKira/novelai-python/tree/main/playground)
64
67
  directory, read code as documentation.**
65
68
 
@@ -76,6 +79,7 @@ from novelai_python.sdk.ai.generate_image.schema import PositionMap
76
79
 
77
80
  load_dotenv()
78
81
  session = ApiCredential(api_token=SecretStr(os.getenv("NOVELAI_JWT"))) # pst-***
82
+ # For security reasons, storing user credentials in plaintext is strongly discouraged.
79
83
 
80
84
  prompt = "1girl, year 2023,dynamic angle, best quality, amazing quality, very aesthetic, absurdres"
81
85
 
@@ -175,8 +179,12 @@ loop.run_until_complete(chat("Hello"))
175
179
  ```python
176
180
  from novelai_python.tool.random_prompt import RandomPromptGenerator
177
181
 
178
- prompt = RandomPromptGenerator(nsfw_enabled=False).random_prompt()
179
- print(prompt)
182
+ generator = RandomPromptGenerator()
183
+ for i in range(10):
184
+ print(generator.generate_common_tags(nsfw=False))
185
+ print(generator.generate_scene_tags())
186
+ print(generator.generate_scene_composition())
187
+ print(generator.get_holiday_themed_tags())
180
188
  ```
181
189
 
182
190
  #### 📦 Run A Server
@@ -189,10 +197,15 @@ python3 -m novelai_python.server -h '127.0.0.1' -p 7888
189
197
  #### 📦 Tokenizer
190
198
 
191
199
  ```python
192
- from novelai_python._enum import get_tokenizer_model, TextLLMModel
200
+ from novelai_python._enum import get_tokenizer_model, TextLLMModel, TextTokenizerGroup
193
201
  from novelai_python.tokenizer import NaiTokenizer
194
202
 
203
+ # Through llm model name to get the tokenizer
195
204
  tokenizer_package = NaiTokenizer(get_tokenizer_model(TextLLMModel.ERATO))
205
+ # Directly use the tokenizer
206
+ clip_tokenizer = NaiTokenizer(TextTokenizerGroup.CLIP)
207
+
208
+ # Tokenize a text
196
209
  t_text = "a fox jumped over the lazy dog"
197
210
  encode_tokens = tokenizer_package.encode(t_text)
198
211
  print(tokenizer_package.tokenize_text(t_text))
@@ -236,8 +249,6 @@ https://dghs-imgutils.deepghs.org/main/api_doc/operate/censor.html
236
249
 
237
250
  ## 🙏 Acknowledgements
238
251
 
239
- [BackEnd](https://api.novelai.net/docs)
240
-
241
252
  [novelai-api](https://github.com/Aedial/novelai-api)
242
253
 
243
254
  [NovelAI-API](https://github.com/HanaokaYuzu/NovelAI-API)
@@ -7,7 +7,8 @@
7
7
 
8
8
  ✨ NovelAI api python sdk with Pydantic, modern and user-friendly.
9
9
 
10
- The goal of this repository is to use Pydantic to build legitimate requests to access the NovelAI API service.
10
+ The goal of this repository is to use Pydantic to build legitimate requests to access
11
+ the [NovelAI API service](https://api.novelai.net/docs).
11
12
 
12
13
  > Python >= 3.9 is required.
13
14
 
@@ -22,6 +23,8 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
22
23
  pip -U install novelai-python
23
24
  ```
24
25
 
26
+ All API users must adhere to the NovelAI Terms of Service: https://novelai.net/terms.
27
+
25
28
  **More examples can be found in the [playground](https://github.com/LlmKira/novelai-python/tree/main/playground)
26
29
  directory, read code as documentation.**
27
30
 
@@ -38,6 +41,7 @@ from novelai_python.sdk.ai.generate_image.schema import PositionMap
38
41
 
39
42
  load_dotenv()
40
43
  session = ApiCredential(api_token=SecretStr(os.getenv("NOVELAI_JWT"))) # pst-***
44
+ # For security reasons, storing user credentials in plaintext is strongly discouraged.
41
45
 
42
46
  prompt = "1girl, year 2023,dynamic angle, best quality, amazing quality, very aesthetic, absurdres"
43
47
 
@@ -137,8 +141,12 @@ loop.run_until_complete(chat("Hello"))
137
141
  ```python
138
142
  from novelai_python.tool.random_prompt import RandomPromptGenerator
139
143
 
140
- prompt = RandomPromptGenerator(nsfw_enabled=False).random_prompt()
141
- print(prompt)
144
+ generator = RandomPromptGenerator()
145
+ for i in range(10):
146
+ print(generator.generate_common_tags(nsfw=False))
147
+ print(generator.generate_scene_tags())
148
+ print(generator.generate_scene_composition())
149
+ print(generator.get_holiday_themed_tags())
142
150
  ```
143
151
 
144
152
  #### 📦 Run A Server
@@ -151,10 +159,15 @@ python3 -m novelai_python.server -h '127.0.0.1' -p 7888
151
159
  #### 📦 Tokenizer
152
160
 
153
161
  ```python
154
- from novelai_python._enum import get_tokenizer_model, TextLLMModel
162
+ from novelai_python._enum import get_tokenizer_model, TextLLMModel, TextTokenizerGroup
155
163
  from novelai_python.tokenizer import NaiTokenizer
156
164
 
165
+ # Through llm model name to get the tokenizer
157
166
  tokenizer_package = NaiTokenizer(get_tokenizer_model(TextLLMModel.ERATO))
167
+ # Directly use the tokenizer
168
+ clip_tokenizer = NaiTokenizer(TextTokenizerGroup.CLIP)
169
+
170
+ # Tokenize a text
158
171
  t_text = "a fox jumped over the lazy dog"
159
172
  encode_tokens = tokenizer_package.encode(t_text)
160
173
  print(tokenizer_package.tokenize_text(t_text))
@@ -198,8 +211,6 @@ https://dghs-imgutils.deepghs.org/main/api_doc/operate/censor.html
198
211
 
199
212
  ## 🙏 Acknowledgements
200
213
 
201
- [BackEnd](https://api.novelai.net/docs)
202
-
203
214
  [novelai-api](https://github.com/Aedial/novelai-api)
204
215
 
205
216
  [NovelAI-API](https://github.com/HanaokaYuzu/NovelAI-API)
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "novelai-python"
3
- version = "0.6.1"
3
+ version = "0.7.0"
4
4
  description = "NovelAI Python Binding With Pydantic"
5
5
  authors = [
6
6
  { name = "sudoskys", email = "coldlando@hotmail.com" },
@@ -59,4 +59,5 @@ distribution = true
59
59
  [tool.pdm.dev-dependencies]
60
60
  dev = [
61
61
  "pre-commit>=3.5.0",
62
+ "wasmtime>=28.0.0",
62
63
  ]
@@ -40,6 +40,7 @@ class TextTokenizerGroup(object):
40
40
  LLAMA3 = "llama3_tokenizer.def"
41
41
  GPT2 = "gpt2_tokenizer.def"
42
42
  CLIP = "clip_tokenizer.def"
43
+ T5 = "t5_tokenizer.def"
43
44
 
44
45
 
45
46
  TextLLMModelTypeAlias = Union[TextLLMModel, str]
@@ -341,7 +341,8 @@ class LLM(ApiBaseModel):
341
341
  override_headers: Optional[dict] = None
342
342
  ) -> LLMResp:
343
343
  """
344
- 生成图片
344
+ Generate text using NovelAI's large language models. According to our Terms of Service, all generation requests must be initiated by a human action. Automating text or image generation to create excessive load on our systems is not allowed.
345
+
345
346
  :param override_headers:
346
347
  :param session: session
347
348
  :return: LLMStreamResp
@@ -115,7 +115,7 @@ class LLMGenerationParams(BaseModel):
115
115
  repetition_penalty_slope: Optional[Union[float, int]] = Field(default=None, ge=0, allow_inf_nan=False)
116
116
 
117
117
  eos_token_id: int = None
118
- bad_words_ids: List[List[int]] = None
118
+ bad_words_ids: Optional[List[List[int]]] = None
119
119
  logit_bias_groups: Optional[List[LogitBiasGroup]] = Field(default_factory=list)
120
120
 
121
121
  repetition_penalty_frequency: Optional[Union[float, int]] = Field(default=None, allow_inf_nan=False)
@@ -891,7 +891,8 @@ class GenerateImageInfer(ApiBaseModel):
891
891
  override_headers: Optional[dict] = None,
892
892
  ) -> ImageGenerateResp:
893
893
  """
894
- 生成图片
894
+ Generate images using NovelAI's diffusion models. According to our Terms of Service, all generation requests must be initiated by a human action. Automating text or image generation to create excessive load on our systems is not allowed.
895
+
895
896
  :param override_headers: the headers to override
896
897
  :param session: session
897
898
  :return:
@@ -1,5 +1,6 @@
1
1
  import gzip
2
2
  import io
3
+ import json
3
4
  import pathlib
4
5
  import zlib
5
6
  from typing import Dict, List, Optional
@@ -56,6 +57,14 @@ class LLMTokenizerSetting(TokenizerSetting):
56
57
  return self
57
58
 
58
59
 
60
+ def is_clip_model(model_name: str) -> bool:
61
+ return model_name.lower().startswith("clip_tokenizer")
62
+
63
+
64
+ def is_t5_model(model_name: str) -> bool:
65
+ return model_name.lower().startswith("t5_tokenizer")
66
+
67
+
59
68
  class NaiTokenizer:
60
69
  def __init__(self,
61
70
  model: str,
@@ -129,8 +138,10 @@ class NaiTokenizer:
129
138
  decoded_str = self._read_compressed_def(model_bytes)
130
139
  repaired_data = repair_json(decoded_str, return_objects=True)
131
140
  try:
132
- if "clip" in self.model_full_name.lower():
141
+ if is_clip_model(self.model_full_name):
133
142
  self.tokenizer_setting = ClipTokenizerSetting.model_validate(repaired_data)
143
+ elif is_t5_model(self.model_full_name):
144
+ self.tokenizer_setting = json.dumps(repaired_data)
134
145
  else:
135
146
  self.tokenizer_setting = LLMTokenizerSetting.model_validate(repaired_data)
136
147
  except Exception as e:
@@ -140,8 +151,10 @@ class NaiTokenizer:
140
151
  )
141
152
 
142
153
  def _create_tokenizer(self):
143
- if "clip" in self.model_full_name.lower():
154
+ if is_clip_model(self.model_full_name):
144
155
  self._create_clip_tokenizer()
156
+ elif is_t5_model(self.model_full_name):
157
+ self.tokenizer = Tokenizer.from_str(self.tokenizer_setting)
145
158
  else:
146
159
  self._create_bpe_tokenizer()
147
160
 
@@ -13,94 +13,109 @@ from nacl.signing import VerifyKey
13
13
  from pydantic import BaseModel, ConfigDict
14
14
 
15
15
  from novelai_python.sdk.ai._enum import PROMOTION, Model
16
- from ._type import IMAGE_INPUT_TYPE, get_image_bytes
17
16
  from .bch_utils import fec_decode
18
17
  from .lsb_extractor import ImageLsbDataExtractor
19
18
  from .lsb_injector import inject_data
20
19
 
21
20
 
21
+ class BrokenMetaDataError(Exception):
22
+ pass
23
+
24
+
25
+ class BrokenCommentError(Exception):
26
+ pass
27
+
28
+
29
+ class CommentModel(BaseModel):
30
+ prompt: Optional[str] = None
31
+ steps: Optional[int] = None
32
+ height: Optional[int] = None
33
+ width: Optional[int] = None
34
+ scale: Optional[float] = None
35
+ uncond_scale: Optional[float] = None
36
+ cfg_rescale: Optional[float] = None
37
+ seed: Optional[int] = None
38
+ n_samples: Optional[int] = None
39
+ hide_debug_overlay: Optional[bool] = None
40
+ noise_schedule: Optional[str] = None
41
+ legacy_v3_extend: Optional[bool] = None
42
+ reference_information_extracted: Optional[float] = None
43
+ reference_strength: Optional[float] = None
44
+ reference_strength_multiple: Optional[List[float]] = None
45
+ reference_information_extracted_multiple: Optional[List[float]] = None
46
+ sampler: Optional[str] = None
47
+ controlnet_strength: Optional[float] = None
48
+ controlnet_model: Optional[Union[None, str]] = None
49
+ dynamic_thresholding: Optional[bool] = None
50
+ dynamic_thresholding_percentile: Optional[float] = None
51
+ dynamic_thresholding_mimic_scale: Optional[float] = None
52
+ sm: Optional[bool] = None
53
+ sm_dyn: Optional[bool] = None
54
+ skip_cfg_below_sigma: Optional[float] = None
55
+ lora_unet_weights: Optional[Union[None, str]] = None
56
+ lora_clip_weights: Optional[Union[None, str]] = None
57
+ uc: Optional[str] = None
58
+ request_type: Optional[str] = None
59
+ signed_hash: Optional[str] = None
60
+ skip_cfg_above_sigma: Optional[float] = None
61
+ deliberate_euler_ancestral_bug: Optional[bool] = None
62
+ prefer_brownian: Optional[bool] = None
63
+ cfg_sched_eligibility: Optional[str] = None
64
+ explike_fine_detail: Optional[bool] = None
65
+ minimize_sigma_inf: Optional[bool] = None
66
+ uncond_per_vibe: Optional[bool] = None
67
+ wonky_vibe_correlation: Optional[bool] = None
68
+ version: Optional[int] = None
69
+
70
+ extra_passthrough_testing: Optional[dict] = None
71
+ """More test infomation"""
72
+ v4_prompt: Optional[dict] = None
73
+ v4_negative_prompt: Optional[dict] = None
74
+
75
+ model_config = ConfigDict(extra="allow")
76
+
77
+ @property
78
+ def generate_method(self):
79
+ return self.request_type or "Unknown"
80
+
81
+ @property
82
+ def negative_prompt(self):
83
+ return self.uc
84
+
85
+ @property
86
+ def vibe_transfer_strength(self) -> List[float]:
87
+ """
88
+ Get the vibe transfer strength totally
89
+ """
90
+ if self.reference_strength_multiple:
91
+ return self.reference_strength_multiple
92
+ reference_strength = [] if self.reference_strength is None else [self.reference_strength]
93
+ return reference_strength
94
+
95
+ @property
96
+ def vibe_transfer_information(self) -> List[float]:
97
+ """
98
+ Get the vibe transfer information totally
99
+ """
100
+ if self.reference_information_extracted_multiple:
101
+ return self.reference_information_extracted_multiple
102
+ reference_information = [] if self.reference_information_extracted is None else [
103
+ self.reference_information_extracted]
104
+ return reference_information
105
+
106
+
22
107
  class ImageMetadata(BaseModel):
23
108
  """
24
109
  ImageMetadata is a class that represents the metadata of a NovelAI generated image.
25
110
  """
26
111
 
27
- class CommentModel(BaseModel):
28
- prompt: Optional[str] = None
29
- steps: Optional[int] = None
30
- height: Optional[int] = None
31
- width: Optional[int] = None
32
- scale: Optional[float] = None
33
- uncond_scale: Optional[float] = None
34
- cfg_rescale: Optional[float] = None
35
- seed: Optional[int] = None
36
- n_samples: Optional[int] = None
37
- hide_debug_overlay: Optional[bool] = None
38
- noise_schedule: Optional[str] = None
39
- legacy_v3_extend: Optional[bool] = None
40
- reference_information_extracted: Optional[float] = None
41
- reference_strength: Optional[float] = None
42
- reference_strength_multiple: Optional[List[float]] = None
43
- reference_information_extracted_multiple: Optional[List[float]] = None
44
- sampler: Optional[str] = None
45
- controlnet_strength: Optional[float] = None
46
- controlnet_model: Optional[Union[None, str]] = None
47
- dynamic_thresholding: Optional[bool] = None
48
- dynamic_thresholding_percentile: Optional[float] = None
49
- dynamic_thresholding_mimic_scale: Optional[float] = None
50
- sm: Optional[bool] = None
51
- sm_dyn: Optional[bool] = None
52
- skip_cfg_below_sigma: Optional[float] = None
53
- lora_unet_weights: Optional[Union[None, str]] = None
54
- lora_clip_weights: Optional[Union[None, str]] = None
55
- uc: Optional[str] = None
56
- request_type: Optional[str] = None
57
- signed_hash: Optional[str] = None
58
- skip_cfg_above_sigma: Optional[float] = None
59
- deliberate_euler_ancestral_bug: Optional[bool] = None
60
- prefer_brownian: Optional[bool] = None
61
- cfg_sched_eligibility: Optional[str] = None
62
- explike_fine_detail: Optional[bool] = None
63
- minimize_sigma_inf: Optional[bool] = None
64
- uncond_per_vibe: Optional[bool] = None
65
- wonky_vibe_correlation: Optional[bool] = None
66
- version: Optional[int] = None
67
-
68
- model_config = ConfigDict(extra="allow")
69
-
70
- @property
71
- def generate_method(self):
72
- return self.request_type or "Unknown"
73
-
74
- @property
75
- def negative_prompt(self):
76
- return self.uc
77
-
78
- @property
79
- def vibe_transfer_strength(self) -> List[float]:
80
- """
81
- Get the vibe transfer strength totally
82
- """
83
- if self.reference_strength_multiple:
84
- return self.reference_strength_multiple
85
- reference_strength = [] if self.reference_strength is None else [self.reference_strength]
86
- return reference_strength
87
-
88
- @property
89
- def vibe_transfer_information(self) -> List[float]:
90
- """
91
- Get the vibe transfer information totally
92
- """
93
- if self.reference_information_extracted_multiple:
94
- return self.reference_information_extracted_multiple
95
- reference_information = [] if self.reference_information_extracted is None else [
96
- self.reference_information_extracted]
97
- return reference_information
98
-
99
- Title: str = "AI generated image"
100
- Software: str = "NovelAI"
101
- Source: str = None
102
- Description: str
103
- Comment: CommentModel
112
+ Generation_time: Optional[float] = None
113
+ Comment: Optional[CommentModel] = None
114
+ Title: Optional[str] = None
115
+ Description: Optional[str] = None
116
+ Software: Optional[str] = None
117
+ Source: Optional[str] = None
118
+ dpi: Optional[Tuple[float, float]] = None
104
119
 
105
120
  model_config = ConfigDict(extra="allow")
106
121
 
@@ -113,14 +128,13 @@ class ImageMetadata(BaseModel):
113
128
  return PROMOTION.get(self.Source, None)
114
129
 
115
130
  @staticmethod
116
- def reset_alpha(image: IMAGE_INPUT_TYPE) -> BytesIO:
131
+ def reset_alpha(image: Image.Image) -> BytesIO:
117
132
  """
118
133
  Remove LSB from the image, set the alpha channel to 254
119
134
  :param image: Type: Union[str, bytes, Path, BytesIO]
120
135
  :return: BytesIO
121
136
  """
122
- image_data = get_image_bytes(image)
123
- image = Image.open(image_data).convert('RGBA')
137
+ image = image.convert('RGBA')
124
138
  data = np.array(image)
125
139
  data[..., 3] = 254
126
140
  new_image = Image.fromarray(data, 'RGBA')
@@ -128,7 +142,7 @@ class ImageMetadata(BaseModel):
128
142
  new_image.save(_new_img_io, format="PNG")
129
143
  return _new_img_io
130
144
 
131
- def apply_to_image(self, image: IMAGE_INPUT_TYPE, *, inject_lsb: bool = True) -> BytesIO:
145
+ def apply_to_image(self, image: Image.Image, *, inject_lsb: bool = True) -> BytesIO:
132
146
  """
133
147
  Write metadata to origin_image
134
148
  If you set inject_lsb to True, the image will be injected with metadata using LSB.
@@ -138,10 +152,9 @@ class ImageMetadata(BaseModel):
138
152
  :param inject_lsb: Inject metadata using LSB
139
153
  :return: BytesIO
140
154
  """
141
- image_data = get_image_bytes(image)
142
- origin_image = Image.open(image_data)
155
+ origin_image = image
143
156
  metadata = PngInfo()
144
- for k, v in self.model_dump(mode="json").items():
157
+ for k, v in self.model_dump(mode="json", exclude_none=True).items():
145
158
  if isinstance(v, dict):
146
159
  v = json.dumps(v)
147
160
  metadata.add_text(k, v)
@@ -155,7 +168,7 @@ class ImageMetadata(BaseModel):
155
168
  return be_copy_image
156
169
 
157
170
  @classmethod
158
- def _extract_metadata_from_lsb(cls, image: IMAGE_INPUT_TYPE):
171
+ def _extract_metadata_from_lsb(cls, image: Image.Image):
159
172
  """
160
173
  Extract metadata using LSB extraction method.
161
174
  :param image: IMAGE_INPUT_TYPE
@@ -168,38 +181,50 @@ class ImageMetadata(BaseModel):
168
181
  return None
169
182
 
170
183
  @classmethod
171
- def _extract_metadata_from_comments(cls, image: IMAGE_INPUT_TYPE):
184
+ def _extract_metadata_from_comments(cls, image: Image.Image):
172
185
  """
173
186
  Extract metadata from image comments and other info.
174
187
  :param image: Union[str, bytes, Path, BytesIO]
175
188
  """
176
- with Image.open(image) as img:
177
- title = img.info.get("Title", None)
178
- description = img.info.get("Description", None)
179
- comment = img.info.get("Comment", "{}")
180
- try:
181
- comment_dict = json.loads(comment)
182
- except Exception as e:
183
- logger.trace(f"Error loading comment: {e}")
184
- comment_dict = {}
189
+ generate_map = image.info
190
+ comment = generate_map.get("Comment", "{}")
185
191
  try:
186
- comment_dict.setdefault("prompt", description)
187
- comment_model = cls.CommentModel.model_validate(comment_dict)
188
- return cls(Title=title, Description=description, Comment=comment_model)
192
+ comment_dict = json.loads(comment)
189
193
  except Exception as e:
190
- logger.debug(f"Error loading comment: {e}")
191
- return None
194
+ raise BrokenMetaDataError(
195
+ f"ImageMetadata: non-standard comment json cant be loaded, given image maybe be modified!"
196
+ ) from e
197
+ comment_dict["prompt"] = comment_dict.get("prompt", "")
198
+ try:
199
+ comment_model = CommentModel.model_validate(comment_dict)
200
+ except Exception as e:
201
+ raise BrokenCommentError(
202
+ f"ImageMetadata: comment model validate failed, given image maybe be modified."
203
+ "If you are sure that the image is from novelai.net, please issue this problem to "
204
+ "https://github.com/LlmKira/novelai-python/issues."
205
+ ) from e
206
+ try:
207
+ return cls.model_validate(
208
+ {
209
+ **generate_map,
210
+ "Comment": comment_model.model_dump(mode="python"),
211
+ }
212
+ )
213
+ except Exception as e:
214
+ raise BrokenMetaDataError(
215
+ f"ImageMetadata: some error occurred when validate the model, are the image from novelai.net?"
216
+ ) from e
192
217
 
193
218
  @classmethod
194
- def load_image(cls, image: IMAGE_INPUT_TYPE):
219
+ def load_image(cls, image: Image.Image):
195
220
  """
196
221
  Load image and extract metadata using LSB/Metadata
197
222
  :param image: str, bytes, Path, BytesIO
198
223
  :return: ImageMetadata
199
224
  :raises ValueError: Data extraction failed
200
225
  """
201
- image_data = get_image_bytes(image)
202
- metadata = cls._extract_metadata_from_lsb(image_data)
226
+
227
+ metadata = cls._extract_metadata_from_lsb(image)
203
228
  if metadata:
204
229
  return metadata
205
230
  metadata = cls._extract_metadata_from_comments(image)
@@ -207,12 +232,37 @@ class ImageMetadata(BaseModel):
207
232
  return metadata
208
233
  raise ValueError("No metadata found")
209
234
 
235
+ @classmethod
236
+ def load_from_watermark(cls, image: Image.Image):
237
+ """
238
+ Load image and extract metadata using LSB/Metadata
239
+ :param image: str, bytes, Path, BytesIO
240
+ :return: ImageMetadata
241
+ :raises ValueError: Data extraction failed
242
+ """
243
+ metadata = cls._extract_metadata_from_lsb(image)
244
+ if metadata:
245
+ return metadata
246
+ raise ValueError("No metadata found")
247
+
248
+ @classmethod
249
+ def load_from_pnginfo(cls, image: Image.Image):
250
+ """
251
+ Load image and extract metadata using LSB/Metadata
252
+ :param image: str, bytes, Path, BytesIO
253
+ :return: ImageMetadata
254
+ :raises ValueError: Data extraction failed
255
+ """
256
+ metadata = cls._extract_metadata_from_comments(image)
257
+ if metadata:
258
+ return metadata
259
+ raise ValueError("No metadata found")
260
+
210
261
 
211
262
  class ImageVerifier:
212
263
 
213
264
  @staticmethod
214
265
  def verify_latents(image: Image.Image, signed_hash: bytes, verify_key: VerifyKey):
215
- image.load()
216
266
  sig = None
217
267
  latents = None
218
268
  try:
@@ -248,20 +298,19 @@ class ImageVerifier:
248
298
  return False, False, None
249
299
 
250
300
  def verify(self,
251
- image,
301
+ image: Image.Image,
252
302
  verify_key_hex: str = "Y2JcQAOhLwzwSDUJPNgL04nS0Tbqm7cSRc4xk0vRMic="
253
303
  ) -> Tuple[bool, bool]:
254
304
  """
255
305
  Verify if the image is a NovelAI generated image
256
306
 
257
- :param image: Union[str, bytes, Path, BytesIO] - The input image to verify.
307
+ :param image: PIL.Image.Image - The image to verify.
258
308
  :param verify_key_hex: str - The verification key in base64 format.
259
309
  :return: Tuple[bool, bool] - The first bool indicates if the image is verified, the second bool indicates if the image has latents.
260
310
  :raises ValueError: If the required metadata or signed hash is missing.
261
311
  """
262
- image_obj = Image.open(get_image_bytes(image))
263
- w, h = image_obj.size
264
- image_array = np.array(image_obj)
312
+ w, h = image.size
313
+ image_array = np.array(image)
265
314
  try:
266
315
  metadata, fec_data = ImageLsbDataExtractor().extract_data(image, get_fec=True)
267
316
  except Exception as e:
@@ -277,7 +326,7 @@ class ImageVerifier:
277
326
  # Build verify key
278
327
  verify_key = VerifyKey(verify_key_hex.encode("utf-8"), encoder=Base64Encoder)
279
328
  # Verify latents
280
- good_latents, have_latents, latents = self.verify_latents(image_obj, signed_hash, verify_key)
329
+ good_latents, have_latents, latents = self.verify_latents(image, signed_hash, verify_key)
281
330
  if not good_latents:
282
331
  return False, False
283
332
  rgb = image_array[:, :, :3].tobytes()