Merge pull request #9 from LlmKira/dev

✨ feat(credential): Add ApiCredential class
LlmKira · Feb 8, 2024 · 3ec1d71 · 3ec1d71
2 parents 3f1ee7e + 2156d08
commit 3ec1d71
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 9 deletions.
diff --git a/README.md b/README.md
@@ -10,6 +10,7 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 - [x] /ai/generate-image
 - [x] /user/subscription
 - [x] /user/login
+- [x] /user/information
 - [ ] /ai/generate-image/suggest-tags
 - [ ] /ai/annotate-image
 - [ ] /ai/classify
@@ -29,7 +30,7 @@ import os
 from dotenv import load_dotenv
 from pydantic import SecretStr
 
-from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential, LoginCredential
+from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential, LoginCredential, ApiCredential
 
 load_dotenv()
 
@@ -38,18 +39,26 @@ enhance = "year 2023,dynamic angle,  best quality, amazing quality, very aesthet
 
 async def main():
     globe_s = JwtCredential(
-        jwt_token=SecretStr(os.getenv("NOVELAI_JWT"))
+        jwt_token=SecretStr(os.getenv("NOVELAI_JWT"))  # ey****
+    )
+    globe_s1 = ApiCredential(
+        api_token=SecretStr(os.getenv("NOVELAI_JWT"))  # pst-***
     )
     globe_s2 = LoginCredential(
         username=os.getenv("NOVELAI_USERNAME"),
         password=SecretStr(os.getenv("NOVELAI_PASSWORD"))
     )
-    _res = await GenerateImageInfer.build(
-        prompt=f"1girl,{enhance}").generate(
-        session=globe_s)
-    _res: ImageGenerateResp
-    print(_res.meta)
-    file = _res.files[0]
+
+    gen = await GenerateImageInfer.build(
+        prompt=f"1girl,{enhance}")
+    cost = gen.calculate_cost(is_opus=True)
+    print(f"charge: {cost} if you are vip3")
+
+    resp = gen.generate(session=globe_s)
+    resp: ImageGenerateResp
+    print(resp.meta)
+
+    file = resp.files[0]
     with open(file[0], "wb") as f:
         f.write(file[1])
 

diff --git a/playground/generate_image.py b/playground/generate_image.py
@@ -26,6 +26,8 @@ async def main():
     try:
         gen = GenerateImageInfer.build(
             prompt=f"1girl, winter, jacket, sfw, angel, flower,{enhance}")
+        cost = gen.calculate_cost(is_opus=True)
+        print(f"charge: {cost} if you are vip3")
         print(f"charge: {gen.calculate_cost(is_opus=True)}")
         _res = await gen.generate(
             session=globe_s, remove_sign=True

diff --git a/src/novelai_python/__init__.py b/src/novelai_python/__init__.py
@@ -8,7 +8,7 @@
     AuthError,
     NovelAiError, APIError,
 )
-from .credential import JwtCredential, LoginCredential
+from .credential import JwtCredential, LoginCredential, ApiCredential
 from .sdk import GenerateImageInfer, ImageGenerateResp
 from .sdk import Information, InformationResp
 from .sdk import Login, LoginResp
@@ -29,6 +29,7 @@
 
     "JwtCredential",
     "LoginCredential",
+    "ApiCredential",
 
     "APIError",
     "AuthError",

diff --git a/src/novelai_python/credential/__init__.py b/src/novelai_python/credential/__init__.py
@@ -5,6 +5,7 @@
 # @Software: PyCharm
 from pydantic import SecretStr
 
+from .ApiToken import ApiCredential
 from .JwtToken import JwtCredential
 from .UserAuth import LoginCredential
 from ._base import CredentialBase
@@ -13,5 +14,6 @@
     "JwtCredential",
     "LoginCredential",
     "CredentialBase",
+    "ApiCredential",
     "SecretStr"
 ]