import os import datetime import gradio as gr from pnginfo import read_info_from_image, send_paras from utils import set_token, generate_novelai_image, image_from_bytes 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, value="{breast expansion},{gigantic breasts},[artist:ningen_mame],artist:ciloranko,[artist:sho_(sho_lwlw)],[artist:foifoi (marfoyfoyfoy)],1girl,skinny,narrow waist,") quality_tags = gr.TextArea( label="Quality Tags", lines=1, value="best quality, amazing quality, very aesthetic, absurdres", ) neg_prompt = gr.TextArea( 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", ) 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.Accordion('风格迁移', open=False): ref_image = gr.Image(label="上传图片", value=None, sources=["upload"], 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): with gr.Row(): use_inp = gr.Checkbox(label='启用', value=False) overlay = gr.Checkbox(label='覆盖原图', value=True) inp_img = gr.ImageEditor(label="上传图片", value=None, sources=["upload"], interactive=True, type="pil", eraser=False, transforms=None, brush="imagemask") inp_str = gr.Slider(label='重绘强度', value=0.7, minimum=0, maximum=0.99, step=0.01) reuse_img_inp = gr.Button(value='使用上一次生成的图片') 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") 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, inp_img, overlay, inp_str], [save, rand_seed, reuse_seed, reuse_img_vibe, use_inp, reuse_img_inp] 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, inp_img, overlay, inp_str, save, use_inp): global today_count set_token(os.environ.get('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, ref_image, info_extract, ref_str, inp_img, overlay, use_inp, inp_str ) if not isinstance(img_data, bytes): return None, payload 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: image = gr.Image(elem_id='preview_image', interactive=False, type='pil') info = gr.JSON(value={}, label="Submitted Payload") 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], others[4]], [image, info]) others[2].click(lambda o, s: o if len(s) == 0 else s['parameters']['seed'], inputs=[paras[3], info], outputs=paras[3]) others[3].click(lambda i: i, inputs=image, outputs=paras[14]) others[5].click(lambda i: gr.ImageEditor(value=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") 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')))