From 2156d08ce42c82019cb80d88f3dbce65503b18c6 Mon Sep 17 00:00:00 2001
From: sudoskys <coldlando@hotmail.com>
Date: Thu, 8 Feb 2024 16:46:13 +0800
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat(playground):=20add=20cost=20ca?=
 =?UTF-8?q?lculation=20for=20VIP3=20users?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

✨ feat(playground): add cost calculation for VIP3 users

- Added cost calculation for VIP3 users in the `generate_image.py` file
- Output the charge amount if the user is a VIP3 member
---
 README.md                    | 25 +++++++++++++++++--------
 playground/generate_image.py |  2 ++
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 572c49b..079b445 100644
--- 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
index d6cf6c6..c37430e 100644
--- 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