Merge branch 'ae'

1 files changed, 54 insertions, 1 deletions

diff --git a/modules/ui.py b/modules/ui.py
index f6a92ddc..381ca925 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -25,7 +25,9 @@ import gradio.routes
 

 from modules import sd_hijack, sd_models, localization

 from modules.paths import script_path

-from modules.shared import opts, cmd_opts, restricted_opts

+

+from modules.shared import opts, cmd_opts, restricted_opts, aesthetic_embeddings

+

 if cmd_opts.deepdanbooru:

     from modules.deepbooru import get_deepbooru_tags

 import modules.shared as shared

@@ -41,8 +43,11 @@ from modules import prompt_parser
 from modules.images import save_image

 import modules.textual_inversion.ui

 import modules.hypernetworks.ui

+

+import modules.aesthetic_clip as aesthetic_clip

 import modules.images_history as img_his

 

+

 # this is a fix for Windows users. Without it, javascript files will be served with text/html content-type and the browser will not show any UI

 mimetypes.init()

 mimetypes.add_type('application/javascript', '.js')

@@ -655,6 +660,8 @@ def create_ui(wrap_gradio_gpu_call):
 

                 seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox = create_seed_inputs()

 

+                aesthetic_weight, aesthetic_steps, aesthetic_lr, aesthetic_slerp, aesthetic_imgs, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative = aesthetic_clip.create_ui()

+

                 with gr.Group():

                     custom_inputs = modules.scripts.scripts_txt2img.setup_ui(is_img2img=False)

 

@@ -709,7 +716,16 @@ def create_ui(wrap_gradio_gpu_call):
                     denoising_strength,

                     firstphase_width,

                     firstphase_height,

+                    aesthetic_lr,

+                    aesthetic_weight,

+                    aesthetic_steps,

+                    aesthetic_imgs,

+                    aesthetic_slerp,

+                    aesthetic_imgs_text,

+                    aesthetic_slerp_angle,

+                    aesthetic_text_negative

                 ] + custom_inputs,

+

                 outputs=[

                     txt2img_gallery,

                     generation_info,

@@ -870,6 +886,8 @@ def create_ui(wrap_gradio_gpu_call):
 

                 seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox = create_seed_inputs()

 

+                aesthetic_weight_im, aesthetic_steps_im, aesthetic_lr_im, aesthetic_slerp_im, aesthetic_imgs_im, aesthetic_imgs_text_im, aesthetic_slerp_angle_im, aesthetic_text_negative_im = aesthetic_clip.create_ui()

+

                 with gr.Group():

                     custom_inputs = modules.scripts.scripts_img2img.setup_ui(is_img2img=True)

 

@@ -960,6 +978,14 @@ def create_ui(wrap_gradio_gpu_call):
                     inpainting_mask_invert,

                     img2img_batch_input_dir,

                     img2img_batch_output_dir,

+                    aesthetic_lr_im,

+                    aesthetic_weight_im,

+                    aesthetic_steps_im,

+                    aesthetic_imgs_im,

+                    aesthetic_slerp_im,

+                    aesthetic_imgs_text_im,

+                    aesthetic_slerp_angle_im,

+                    aesthetic_text_negative_im,

                 ] + custom_inputs,

                 outputs=[

                     img2img_gallery,

@@ -1220,6 +1246,18 @@ def create_ui(wrap_gradio_gpu_call):
                         with gr.Column():

                             create_embedding = gr.Button(value="Create embedding", variant='primary')

 

+                with gr.Tab(label="Create aesthetic images embedding"):

+

+                    new_embedding_name_ae = gr.Textbox(label="Name")

+                    process_src_ae = gr.Textbox(label='Source directory')

+                    batch_ae = gr.Slider(minimum=1, maximum=1024, step=1, label="Batch size", value=256)

+                    with gr.Row():

+                        with gr.Column(scale=3):

+                            gr.HTML(value="")

+

+                        with gr.Column():

+                            create_embedding_ae = gr.Button(value="Create images embedding", variant='primary')

+

                 with gr.Tab(label="Create hypernetwork"):

                     new_hypernetwork_name = gr.Textbox(label="Name")

                     new_hypernetwork_sizes = gr.CheckboxGroup(label="Modules", value=["768", "320", "640", "1280"], choices=["768", "320", "640", "1280"])

@@ -1309,6 +1347,21 @@ def create_ui(wrap_gradio_gpu_call):
             ]

         )

 

+        create_embedding_ae.click(

+            fn=aesthetic_clip.generate_imgs_embd,

+            inputs=[

+                new_embedding_name_ae,

+                process_src_ae,

+                batch_ae

+            ],

+            outputs=[

+                aesthetic_imgs,

+                aesthetic_imgs_im,

+                ti_output,

+                ti_outcome,

+            ]

+        )

+

         create_hypernetwork.click(

             fn=modules.hypernetworks.ui.create_hypernetwork,

             inputs=[