Update app.py

app.py

@@ -1,11 +1,15 @@
 import os
 import datetime
 
+import toml
 import gradio as gr
-from pnginfo import read_info_from_image, send_paras
+from PIL import PngImagePlugin
 
+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
 
+client_config = toml.load("config.toml")['client']
 today_count = 0
 today = datetime.date.today().strftime('%Y-%m-%d')
 
@@ -15,17 +19,17 @@ def get_count():
     if now != today:
         today = now
         today_count = 0
-    return "Total number of today's image generation: " + str(today_count)
+    return "今日已生成图片" + str(today_count) + "张"
 
 def control_ui():
-    prompt = gr.TextArea(elem_id='txt2img_prompt', label="Prompt", lines=3, value="{breast expansion},{gigantic breasts},[artist:ningen_mame],artist:ciloranko,[artist:sho_(sho_lwlw)],[artist:foifoi (marfoyfoyfoy)],1girl,skinny,narrow waist,")
+    prompt = gr.TextArea(elem_id='txt2img_prompt', label="提示词", lines=3)
     quality_tags = gr.TextArea(
-        elem_id='txt2img_qua_prompt', label="Quality Tags", lines=1,
-        value="best quality, amazing quality, very aesthetic, absurdres",
+        elem_id='txt2img_qua_prompt', label="质量词", lines=1,
+        value=client_config['default_quality'],
     )
     neg_prompt = gr.TextArea(
-        elem_id='txt2img_neg_prompt', label="Negative Prompt", lines=1,
-        value="wide hips,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",
+        elem_id='txt2img_neg_prompt', label="负面词", lines=1,
+        value=client_config['default_neg'],
     )
     with gr.Row():
         sampler = gr.Dropdown(
@@ -34,27 +38,27 @@ def control_ui():
                 "k_dpmpp_2m", "k_dpmpp_sde", "ddim_v3"
             ],
             value="k_euler",
-            label="Sampler",
+            label="采样器",
             interactive=True
         )
-        scale = gr.Slider(label="Scale", value=5.0, minimum=1, maximum=10, step=0.1)
-        steps = gr.Slider(label="Steps", value=28, minimum=1, maximum=28, step=1)
+        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)
     with gr.Row():
-        seed = gr.Number(label="Seed", value=-1, step=1, maximum=2**32-1, minimum=-1, scale=3)
+        seed = gr.Number(label="种子", value=-1, step=1, maximum=2**32-1, minimum=-1, scale=3)
         rand_seed = gr.Button('🎲️', scale=1)
         reuse_seed = gr.Button('♻️', scale=1)
     with gr.Row():
-        width = gr.Slider(label="Width", value=1024, minimum=64, maximum=2048, step=64)
-        height = gr.Slider(label="Height", value=1024, minimum=64, maximum=2048, step=64)
+        width = gr.Slider(label="宽度", value=1024, minimum=64, maximum=2048, step=64)
+        height = gr.Slider(label="高度", value=1024, minimum=64, maximum=2048, step=64)
     with gr.Row():
         with gr.Column():
             with gr.Accordion('风格迁移', open=False):
-                ref_image = gr.Image(label="上传图片", value=None, interactive=True, type="pil")
+                ref_image = gr.Image(label="上传图片", value=None, sources=["upload", "clipboard", "webcam"], interactive=True, type="pil")
                 info_extract = gr.Slider(label='参考信息提取', value=1, minimum=0, maximum=1, step=0.1)
                 ref_str = gr.Slider(label='参考强度', value=0.6, minimum=0, maximum=1, step=0.1)
                 reuse_img_vibe = gr.Button(value='使用上一次生成的图片')
             with gr.Accordion('图生图', open=False):
-                i2i_image = gr.Image(label="上传图片", value=None, interactive=True, type="pil")
+                i2i_image = gr.Image(label="上传图片", value=None, sources=["upload", "clipboard", "webcam"], interactive=True, type="pil")
                 i2i_str = gr.Slider(label='去噪强度', value=0.7, minimum=0, maximum=0.99, step=0.01)
                 i2i_noise = gr.Slider(label='噪声', value=0, minimum=0, maximum=1, step=0.1)
                 reuse_img_i2i = gr.Button(value='使用上一次生成的图片')
@@ -69,7 +73,7 @@ def control_ui():
             '''
     with gr.Row():
         with gr.Column():
-            with gr.Accordion('Advanced Gen Setting', open=False):
+            with gr.Accordion('高级选项', open=False):
                 scheduler = gr.Dropdown(
                     choices=[
                         "native", "karras", "exponential", "polyexponential"
@@ -85,15 +89,14 @@ def control_ui():
                     dyn_threshold = gr.Checkbox(False, label="Dynamic Thresholding")
                     cfg_rescale = gr.Slider(0, 1, 0, step=0.01, label="CFG rescale")
         with gr.Column():
-            gr.Textbox(value=get_count, label='Usage count', every=10)
-            save = gr.Checkbox(value=False, label='Always save all generated images', visible=False)
-    gen_btn = gr.Button(value="Generate", variant="primary")
+            gr.Textbox(value=get_count, label='使用统计', every=10)
+            save = gr.Checkbox(value=True, label='云端保存图片')
+    gen_btn = gr.Button(value="生成", variant="primary")
     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)
     return gen_btn,[prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale, ref_image, info_extract, ref_str, i2i_image, i2i_str, i2i_noise], [save, rand_seed, reuse_seed, reuse_img_vibe, reuse_img_i2i]
 
-
 def generate(prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale, ref_image, info_extract, ref_str, i2i_image, i2i_str, i2i_noise, save):
     global today_count
     set_token(os.environ.get('token'))
@@ -104,19 +107,28 @@ def generate(prompt, quality_tags, neg_prompt, seed, scale, width, height, steps
         i2i_image, i2i_str, i2i_noise
     )
     if not isinstance(img_data, bytes):
-        return None, payload
+        return gr.Image(value=None), payload
     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
 
-
 def preview_ui():
     with gr.Blocks(css='#preview_image { height: 100%;}'):
         image = gr.Image(elem_id='preview_image', interactive=False, type='pil')
-        info = gr.JSON(value={}, label="Submitted Payload")
+        info = gr.JSON(value={}, label="生成信息")
     return image, info
 
-
 def main_ui():
     with gr.Blocks() as page:
         with gr.Row(variant="panel"):
@@ -130,34 +142,41 @@ def main_ui():
     others[4].click(lambda i: i, inputs=image, outputs=paras[17])
     return page, paras[:14]
 
-
 def util_ui():
     with gr.Blocks(analytics_enabled=False) as page:
         with gr.Row(equal_height=False):
             with gr.Column(variant='panel'):
-                image = gr.Image(label="Source", sources=["upload"], interactive=True, type="pil")
+                image = gr.Image(label="上传图片", sources=["upload"], interactive=True, type="pil")
             with gr.Column(variant='panel'):
                 info = gr.HTML()
                 items = gr.JSON(value=None, visible=False)
-                png2main = gr.Button('Send to txt2img')
+                png2main = gr.Button('参数发送到文生图')
     return page, png2main, items, info, image
 
 def ui():
     head = ''
     for f in sorted(os.listdir('./tagcomplete/javascript')):
         head += f'<script type="text/javascript" src="file=tagcomplete/javascript/{f}"></script>\n'
-    with gr.Blocks(title="NAI Client", head=head, theme=gr.themes.Base()) as website:
+    with gr.Blocks(title="NAI Client", head=head) as website:
         with gr.Tabs():
-            with gr.TabItem("Main", elem_id="client_ui_main"):
+            with gr.TabItem("文生图", elem_id="client_ui_main"):
                 _, paras = main_ui()
-            with gr.TabItem("PNG Info"):
+            with gr.TabItem("图片信息读取"):
                 _, png2main, png_items, info, image = util_ui()
+            with gr.TabItem("云端图片浏览") as tab:
+                gal2main, gal_items = img_history_ui(tab)
             png2main.click(fn=send_paras,
                            inputs=[png_items] + paras,
                            outputs=paras)
             png2main.click(fn=None,
                            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