Update app.py

app.py

@@ -1,11 +1,16 @@
 import os
 import datetime
 
+import toml
 import gradio as gr
 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 PIL import PngImagePlugin
 
+
+client_config = toml.load("config.toml")['client']
 today_count = 0
 today = datetime.date.today().strftime('%Y-%m-%d')
 
@@ -21,11 +26,11 @@ def control_ui():
     prompt = gr.TextArea(label="Prompt", lines=3)
     quality_tags = gr.TextArea(
         label="Quality Tags", lines=1,
-        value="best quality, amazing quality, very aesthetic, absurdres",
+        value=client_config['default_quality'],
     )
     neg_prompt = gr.TextArea(
         label="Negative Prompt", lines=1,
-        value="wide hip,thick thighs,plump,lowres, {bad}, error, fewer, extra, missing, worst quality, jpeg artifacts, bad quality, watermark, unfinished, displeasing, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract],pubic hair",
+        value=client_config['default_neg'],
     )
     with gr.Row():
         sampler = gr.Dropdown(
@@ -75,10 +80,10 @@ def control_ui():
     return gen_btn,[prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale], [save, rand_seed, reuse_seed]
 
 
-async def generate(prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale, save):
+def generate(prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale, save):
     global today_count
-    set_token(os.environ.get('token'))
-    img_data, payload = await generate_novelai_image(
+    set_token(client_config['token'])
+    img_data, payload = generate_novelai_image(
         f"{prompt}, {quality_tags}", neg_prompt, seed, scale, 
         width, height, steps, sampler, scheduler, 
         smea, dyn, dyn_threshold, cfg_rescale
@@ -88,6 +93,18 @@ async def generate(prompt, quality_tags, neg_prompt, seed, scale, width, height,
     today_count += 1
     img = image_from_bytes(img_data)
     
+    if save:
+        save_path = client_config['save_path']
+        today = datetime.date.today().strftime('%Y-%m-%d')
+        today_path = os.path.join(save_path, today)
+        if not os.path.exists(today_path):
+            os.makedirs(today_path, mode=777, exist_ok=True)
+        filename = str(today_count).rjust(5, '0') + '-' + str(payload['parameters']['seed']) + '.png'
+        pnginfo_data = PngImagePlugin.PngInfo()
+        for k, v in img.info.items():
+            pnginfo_data.add_text(k, str(v))
+        img.save(os.path.join(today_path, filename), pnginfo=pnginfo_data)
+    
     return img, payload
 
 
@@ -130,6 +147,8 @@ def ui():
                 _, paras = main_ui()
             with gr.TabItem("PNG Info"):
                 _, png2main, png_items, info, image = util_ui()
+            with gr.TabItem("Image Browser") as tab:
+                gal2main, gal_items = img_history_ui(tab)
             png2main.click(fn=send_paras,
                            inputs=[png_items] + paras,
                            outputs=paras)
@@ -137,6 +156,13 @@ def ui():
                            js="(x) => { if (x !== null) document.getElementById('client_ui_main-button').click(); "
                               "return null; }",
                            inputs=image)
+            gal2main.click(fn=send_paras,
+                           inputs=[gal_items] + paras,
+                           outputs=paras)
+            gal2main.click(fn=None,
+                           js="(x) => { if (x !== null) document.getElementById('client_ui_main-button').click(); "
+                              "return null; }",
+                           inputs=gal_items)
             image.change(read_info_from_image, inputs=image, outputs=[info, png_items])
     return website
 
@@ -144,4 +170,4 @@ def ui():
 if __name__ == '__main__':
     website = ui()
     website.queue()
-    website.launch(auth=(os.environ.get('account'), os.environ.get('password')))
+    website.launch()