anlas generation

app.py

@@ -8,13 +8,20 @@ from huggingface_hub import HfApi, snapshot_download
 
 from pnginfo import read_info_from_image, send_paras
 from images_history import img_history_ui
-from utils import set_token, generate_novelai_image, image_from_bytes
+from utils import set_token, generate_novelai_image, image_from_bytes, get_remain_anlas, calculate_cost
 
 client_config = toml.load("config.toml")['client']
 today_count = 0
+anlas = -1
 today = datetime.date.today().strftime('%Y-%m-%d')
 api = HfApi()
 
+def refresh_anlas():
+    global anlas
+    set_token(client_config['token'])
+    anlas = get_remain_anlas()
+    return anlas
+
 def get_count():
     global today_count, today
     now = datetime.date.today().strftime('%Y-%m-%d')
@@ -44,7 +51,7 @@ def control_ui():
             interactive=True
         )
         scale = gr.Slider(label="CFG Scale", value=5.0, minimum=1, maximum=10, step=0.1)
-        steps = gr.Slider(label="步数", value=28, minimum=1, maximum=28, step=1)
+        steps = gr.Slider(label="步数", value=28, minimum=1, maximum=50, step=1)
     with gr.Row():
         seed = gr.Number(label="种子", value=-1, step=1, maximum=2**32-1, minimum=-1, scale=3)
         rand_seed = gr.Button('🎲️', scale=1)
@@ -111,8 +118,6 @@ def control_ui():
     gen_btn = gr.Button(value="生成", variant="primary")
     stop_btn = gr.Button(value="取消", variant="stop", visible=False)
     rand_seed.click(fn=lambda: -1, inputs=None, outputs=seed)
-    width.change(lambda w, h: h if w*h<=1024*1024 else (1024*1024//w//64)*64, [width, height], height)
-    height.change(lambda w, h: w if w*h<=1024*1024 else (1024*1024//h//64)*64, [width, height], width)
     i2i.select(lambda: 'i2i', inputs=None, outputs=selection)
     inp.select(lambda: 'inp', inputs=None, outputs=selection)
     return gen_btn, stop_btn, [prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale, ref_images, info_extracts, ref_strs, i2i_image, i2i_str, i2i_noise, overlay, inp_img, selection], [save, rand_seed, reuse_seed, reuse_img_i2i, reuse_img_inp]
@@ -156,6 +161,15 @@ def rename_save_img(path, payload, save):
         api.upload_folder(folder_path=today_path, path_in_repo=today_path, repo_id="P01yH3dr0n/naimages", repo_type="dataset", token=os.environ.get("hf_token"))
     return renamed_path
 
+def update_btn_cost(w, h, s, sm, dyn, i2i_img, i2i_str, inp_img, selection):
+    if selection == 'i2i' and i2i_img is not None:
+        cost = calculate_cost(w, h, s, False, False, i2i_str)
+    elif selection == 'inp' and inp_img is not None:
+        cost = calculate_cost(w, h, s, False, False)
+    else:
+        cost = calculate_cost(w, h, s, sm, dyn)
+    return gr.Button(value=f"生成(消耗{cost}点数)")
+
 def main_ui():
     with gr.Blocks() as page:
         with gr.Row(variant="panel"):
@@ -163,6 +177,9 @@ def main_ui():
                 gen_btn, stop_btn, paras, others = control_ui()
             with gr.Column():
                 image, info = preview_ui()
+    cost_list = paras[5:8] + paras[10:12] + paras[17:19] + paras[21:23]
+    for component in cost_list:
+        component.change(update_btn_cost, inputs=cost_list, outputs=gen_btn)
     gen_btn.click(lambda: (gr.Button(visible=False), gr.Button(visible=True)), inputs=None, outputs=[gen_btn, stop_btn])
     gen = gen_btn.click(generate, paras, [image, info], concurrency_limit=1)
     others[2].click(lambda o, s: o if len(s) == 0 else s['parameters']['seed'], inputs=[paras[3], info], outputs=paras[3])
@@ -201,7 +218,12 @@ def load_javascript():
     
 def ui():
     load_javascript()
+    refresh_anlas()
     with gr.Blocks(title="NAI Client", theme=gr.themes.Base(), js="() => {document.body.classList.toggle('dark');}") as website:
+        with gr.Row():
+            remain_anlas = gr.Text(label="剩余点数", value=anlas)
+            refresh = gr.Button("🔄")
+            refresh.click(refresh_anlas, inputs=None, outputs=remain_anlas)
         with gr.Tabs():
             with gr.TabItem("图片生成", elem_id="client_ui_main"):
                 _, paras = main_ui()

pnginfo.py

@@ -23,7 +23,7 @@ def to_digit(v):
     except:
         return v
     
-def read_info_from_image(image) -> tuple[str | None, dict]:
+def read_info_from_image(image):
     if image is None:
         return '', {}
     elif type(image) == str:

utils.py

@@ -4,6 +4,7 @@ import zipfile
 import requests
 import json
 import base64
+import math
 
 from PIL import Image
 
@@ -25,6 +26,20 @@ def set_token(token):
             "Referer": "https://novelai.net/"
         }
 
+def get_remain_anlas():
+    data = requests.get("https://api.novelai.net/user/data", headers=headers).content
+    return data['fixedTrainingStepsLeft'] + data['purchasedTrainingSteps']
+
+def calculate_cost(width, height, steps, sm, dyn, strength=1, rmbg=False):
+    pixels = width * height
+    if pixels <= 1048576 and steps <= 28:
+        return 0
+    dyn = sm and dyn
+    L = math.ceil(2951823174884865e-21 * pixels + 5.753298233447344e-7 * pixels * steps)
+    L *= 1.4 if dyn else (1.2 if sm else 1)
+    L = max(math.ceil(L * strength), 2)
+    return L * 3 + 5 if rmbg else L
+
 def generate_novelai_image(
     input_text="", 
     negative_prompt="",