Create app.py

app.py

@@ -0,0 +1,148 @@
+import os
+import datetime
+
+import toml
+import gradio as gr
+from pnginfo import read_info_from_image, send_paras
+
+from utils import set_token, generate_novelai_image, image_from_bytes
+from PIL import PngImagePlugin
+
+today_count = 0
+today = datetime.date.today().strftime('%Y-%m-%d')
+
+def get_count():
+    global today_count, today
+    now = datetime.date.today().strftime('%Y-%m-%d')
+    if now != today:
+        today = now
+        today_count = 0
+    return "Total number of today's image generation: " + str(today_count)
+
+def control_ui():
+    prompt = gr.TextArea(label="Prompt", lines=3)
+    quality_tags = gr.TextArea(
+        label="Quality Tags", lines=1,
+        value=client_config['default_quality'],
+    )
+    neg_prompt = gr.TextArea(
+        label="Negative Prompt", lines=1,
+        value=client_config['default_neg'],
+    )
+    with gr.Row():
+        sampler = gr.Dropdown(
+            choices=[
+                "k_euler", "k_euler_ancestral", "k_dpmpp_2s_ancestral", 
+                "k_dpmpp_2m", "k_dpmpp_sde", "ddim_v3"
+            ],
+            value="k_euler",
+            label="Sampler",
+            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)
+    with gr.Row():
+        seed = gr.Number(label="Seed", 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)
+    with gr.Row():
+        with gr.Column():
+            with gr.Accordion('Advanced Gen Setting', open=False):
+                scheduler = gr.Dropdown(
+                    choices=[
+                        "native", "karras", "exponential", "polyexponential"
+                    ],
+                    value="native",
+                    label="Scheduler",
+                    interactive=True
+                )
+                with gr.Row():
+                    smea = gr.Checkbox(False, label="SMEA")
+                    dyn = gr.Checkbox(False, label="SMEA DYN")
+                with gr.Row():
+                    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=True, label='Always save all generated images')
+    
+    gen_btn = gr.Button(value="Generate", 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], [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):
+    global today_count
+    img_data, payload = await generate_novelai_image(
+        f"{prompt}, {quality_tags}", neg_prompt, seed, scale, 
+        width, height, steps, sampler, scheduler, 
+        smea, dyn, dyn_threshold, cfg_rescale
+    )
+    if not isinstance(img_data, bytes):
+        return None
+    today_count += 1
+    img = image_from_bytes(img_data)
+    
+    return img, payload
+
+
+def preview_ui():
+    with gr.Blocks(css='#preview_image { height: 100%;}') as page:
+        h_slider = gr.Slider(label="Height", value=500, minimum=100, maximum=1200, step=10)
+        image = gr.Image(elem_id='preview_image', interactive=False, height=500)
+        info = gr.JSON(value={}, label="Submitted Payload")
+    h_slider.change(lambda h: gr.Image(height=h), h_slider, image)
+    return image, info
+
+
+def main_ui():
+    with gr.Blocks() as page:
+        with gr.Row(variant="panel"):
+            with gr.Column():
+                gen_btn, paras, others = control_ui()
+            with gr.Column():
+                image, info = preview_ui()
+    gen_btn.click(generate, paras + [others[0]], [image, info])
+    others[2].click(lambda o, s: o if len(s) == 0 else s['parameters']['seed'], inputs=[paras[3], info], outputs=paras[3])
+    return page, paras
+
+
+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")
+            with gr.Column(variant='panel'):
+                info = gr.HTML()
+                items = gr.JSON(value=None, visible=False)
+                png2main = gr.Button('Send to txt2img')
+    return page, png2main, items, info, image
+
+def ui():
+    with gr.Blocks(title="NAI Client") as website:
+        with gr.Tabs():
+            with gr.TabItem("Main", elem_id="client_ui_main"):
+                _, paras = main_ui()
+            with gr.TabItem("PNG Info"):
+                _, png2main, png_items, info, image = util_ui()
+            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)
+            image.change(read_info_from_image, inputs=image, outputs=[info, png_items])
+    return website
+
+
+if __name__ == '__main__':
+    website = ui()
+    website.queue()
+    website.launch(auth=(os.environ.get('account'), os.environ.get('password')))