import os import datetime import toml import gradio as gr import shutil 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 client_config = toml.load("config.toml")['client'] today_count = 0 today = datetime.date.today().strftime('%Y-%m-%d') api = HfApi() def get_count(): global today_count, today now = datetime.date.today().strftime('%Y-%m-%d') if now != today: today = now today_count = 0 return today_count def control_ui(): prompt = gr.TextArea(value="{breast expansion},{gigantic breasts},[artist:ningen_mame],artist:ciloranko,[artist:sho_(sho_lwlw)],[artist:foifoi (marfoyfoyfoy)],1girl,skinny,narrow waist,loli,", elem_id='txt2img_prompt', label="提示词", lines=3) quality_tags = gr.TextArea( elem_id='txt2img_qua_prompt', label="质量词", lines=1, value=client_config['default_quality'], ) neg_prompt = gr.TextArea( elem_id='txt2img_neg_prompt', label="负面词", 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="采样器", 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) 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) reuse_seed = gr.Button('♻️', scale=1) with gr.Row(): 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_images = gr.Gallery(label="上传单(多)张图片", format="png", value=None, interactive=True, type="pil", show_share_button=False) info_extracts = gr.State([]) ref_strs = gr.State([]) @gr.render(inputs=ref_images, triggers=[ref_images.upload, ref_images.change]) def multiple_vibes(images): if images is None: return else: exts = [] strs = [] for i, _ in enumerate(images): with gr.Row(): extract = gr.Slider(label=f"图片{i + 1} 参考信息提取", value=1, minimum=0, maximum=1, step=0.1, interactive=True) strength = gr.Slider(label=f"图片{i + 1} 参考强度", value=0.6, minimum=0, maximum=1, step=0.1, interactive=True) exts.append(extract) strs.append(strength) def dynparas(*args): return list(args) for e, s in zip(exts, strs): e.change(fn=dynparas, inputs=exts, outputs=info_extracts) s.change(fn=dynparas, inputs=strs, outputs=ref_strs) ref_images.change(lambda l: ([], []) if l is None else ([1]*len(l), [0.6]*len(l)), inputs=ref_images, outputs=[info_extracts, ref_strs]) ref_images.upload(lambda l: ([], []) if l is None else ([1]*len(l), [0.6]*len(l)), inputs=ref_images, outputs=[info_extracts, ref_strs]) with gr.Accordion('附加输入', open=False): with gr.Tab('图生图') as i2i: i2i_image = gr.Image(label="上传图片", value=None, sources=["upload", "clipboard", "webcam"], interactive=True, type="pil", show_share_button=False) 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='使用上一次生成的图片') with gr.Tab('局部重绘') as inp: overlay = gr.Checkbox(label='覆盖原图', value=True) inp_img = gr.ImageMask(label="上传图片", value=None, sources=["upload", "clipboard", "webcam"], interactive=True, type="pil", eraser=False, transforms=None, brush=gr.Brush(colors=['#FFFFFF'], color_mode='fixed'), layers=False, show_share_button=False) reuse_img_inp = gr.Button(value='使用上一次生成的图片') selection = gr.Radio(choices=['i2i', 'inp'], value='i2i', visible=False) with gr.Row(): with gr.Column(): with gr.Accordion('高级选项', 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(): save = gr.Checkbox(value=True, label='云端保存图片') 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] def generate(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): 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_images, info_extracts, ref_strs, i2i_image, i2i_str, i2i_noise, overlay, inp_img, selection ) if not isinstance(img_data, bytes): return gr.Image(value=None), payload today_count = get_count() + 1 img = image_from_bytes(img_data) return img, payload def preview_ui(): with gr.Blocks(css='#preview_image { height: 100%;}'): image = gr.Image(format='png', elem_id='preview_image', interactive=False, type='filepath', show_share_button=False) info = gr.JSON(value={}, label="生成信息") return image, info def rename_save_img(path, payload, save): if path is None: return None default = os.path.basename(path) filename = str(today_count).rjust(5, '0') + '-' + str(payload['parameters']['seed']) + '.png' renamed_path = path.replace(default, filename) if os.path.exists(renamed_path): return renamed_path os.replace(path, renamed_path) if save: save_path = client_config['save_path'] today = datetime.date.today().strftime('%Y-%m-%d') today_path = os.path.join(save_path, today) os.makedirs(today_path, exist_ok=True) shutil.copy(renamed_path, os.path.join(today_path, filename)) 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 main_ui(): with gr.Blocks() as page: with gr.Row(variant="panel"): with gr.Column(): gen_btn, stop_btn, paras, others = control_ui() with gr.Column(): image, info = preview_ui() 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]) others[3].click(lambda i: i, inputs=image, outputs=paras[17]) others[4].click(lambda i: gr.ImageEditor(value=i), inputs=image, outputs=paras[21]) ren = image.change(rename_save_img, inputs=[image, info, others[0]], outputs=image, trigger_mode="once") image.change(lambda: (gr.Button(visible=True), gr.Button(visible=False)), inputs=None, outputs=[gen_btn, stop_btn]) stop_btn.click(lambda: (gr.Button(visible=True), gr.Button(visible=False)), inputs=None, outputs=[gen_btn, stop_btn], cancels=[gen, ren]) 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="上传图片", sources=["upload"], interactive=True, type="pil") with gr.Column(variant='panel'): info = gr.HTML() items = gr.JSON(value=None, visible=False) png2main = gr.Button('参数发送到文生图') return page, png2main, items, info, image def load_javascript(): head = '' for f in sorted(os.listdir('./tagcomplete/javascript')): head += f'\n' head += '\n' share = gr.routes.templates.TemplateResponse def template_response(*args, **kwargs): res = share(*args, **kwargs) res.body = res.body.replace(b'', f'{head}'.encode("utf8")) res.init_headers() return res gr.routes.templates.TemplateResponse = template_response def ui(): load_javascript() with gr.Blocks(title="NAI Client", theme=gr.themes.Base(), js="() => {document.body.classList.toggle('dark');}") as website: with gr.Tabs(): with gr.TabItem("图片生成", elem_id="client_ui_main"): _, paras = main_ui() with gr.TabItem("图片信息读取"): _, png2main, png_items, info, image = util_ui() with gr.TabItem("云端图片浏览") as tab: gal2main, gal_items = img_history_ui(tab) with gr.TabItem("设置"): switchLightDark = gr.Button(value="切换浅色/深色模式") switchLightDark.click(fn=None, js="() => {document.body.classList.toggle('dark');}") loadTagComplete = gr.Button(value="重新加载tag补全") loadTagComplete.click(fn=None, js="() => {run();}") clearTagCache = gr.Button(value="清除tag补全缓存") clearTagCache.click(fn=None, js="() => {localStorage.clear();}") 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 if __name__ == '__main__': snapshot_download(repo_id="P01yH3dr0n/naimages", repo_type="dataset", local_dir="./", token=os.environ.get("hf_token"), allow_patterns="*.png") website = ui() website.queue(default_concurrency_limit=5) website.launch(auth=(os.environ.get('account'), os.environ.get('password')), allowed_paths=['tagcomplete'], debug=True)