Update app.py

app.py

@@ -63,7 +63,16 @@ def control_ui():
                 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 = 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=gr.Brush(default_size=5, colors=["#000000"], color_mode="fixed"))
         with gr.Column():
             gr.Textbox(value=get_count, label='Usage count', every=10)
             save = gr.Checkbox(value=True, label='Always save all generated images')
@@ -72,16 +81,17 @@ def control_ui():
     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]
+    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], [save, rand_seed, reuse_seed, reuse_img, use_inp]
 
 
-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, ref_image, info_extract, ref_str, inp_img, overlay, 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
+        smea, dyn, dyn_threshold, cfg_rescale, info_extract, ref_str,
+        inp_img, overlay, use_inp
     )
     if not isinstance(img_data, bytes):
         return None, payload
@@ -105,8 +115,9 @@ def main_ui():
                 gen_btn, paras, others = control_ui()
             with gr.Column():
                 image, info = preview_ui()
-    gen_btn.click(generate, paras + [others[0]], [image, info])
+    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(inputs=image, outputs=paras[14])
     return page, paras