change style dropdowns to multiselect

4 files changed, 43 insertions, 30 deletions

diff --git a/modules/img2img.py b/modules/img2img.py
index f62783c6..79382cc1 100644
--- a/modules/img2img.py
+++ b/modules/img2img.py
@@ -59,7 +59,7 @@ def process_batch(p, input_dir, output_dir, args):
                 processed_image.save(os.path.join(output_dir, filename))

 

 

-def img2img(mode: int, prompt: str, negative_prompt: str, prompt_style: str, prompt_style2: str, init_img, sketch, init_img_with_mask, inpaint_color_sketch, inpaint_color_sketch_orig, init_img_inpaint, init_mask_inpaint, steps: int, sampler_index: int, mask_blur: int, mask_alpha: float, inpainting_fill: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, *args):

+def img2img(mode: int, prompt: str, negative_prompt: str, prompt_styles, init_img, sketch, init_img_with_mask, inpaint_color_sketch, inpaint_color_sketch_orig, init_img_inpaint, init_mask_inpaint, steps: int, sampler_index: int, mask_blur: int, mask_alpha: float, inpainting_fill: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, *args):

     is_batch = mode == 5

 

     if mode == 0:  # img2img

@@ -101,7 +101,7 @@ def img2img(mode: int, prompt: str, negative_prompt: str, prompt_style: str, pro
         outpath_grids=opts.outdir_grids or opts.outdir_img2img_grids,

         prompt=prompt,

         negative_prompt=negative_prompt,

-        styles=[prompt_style, prompt_style2],

+        styles=prompt_styles,

         seed=seed,

         subseed=subseed,

         subseed_strength=subseed_strength,

diff --git a/modules/styles.py b/modules/styles.py
index ce6e71ca..990d5623 100644
--- a/modules/styles.py
+++ b/modules/styles.py
@@ -40,12 +40,18 @@ def apply_styles_to_prompt(prompt, styles):
 class StyleDatabase:

     def __init__(self, path: str):

         self.no_style = PromptStyle("None", "", "")

-        self.styles = {"None": self.no_style}

+        self.styles = {}

+        self.path = path

 

-        if not os.path.exists(path):

+        self.reload()

+

+    def reload(self):

+        self.styles.clear()

+

+        if not os.path.exists(self.path):

             return

 

-        with open(path, "r", encoding="utf-8-sig", newline='') as file:

+        with open(self.path, "r", encoding="utf-8-sig", newline='') as file:

             reader = csv.DictReader(file)

             for row in reader:

                 # Support loading old CSV format with "name, text"-columns

diff --git a/modules/txt2img.py b/modules/txt2img.py
index 38b5f591..5a71793b 100644
--- a/modules/txt2img.py
+++ b/modules/txt2img.py
@@ -8,13 +8,13 @@ import modules.processing as processing
 from modules.ui import plaintext_to_html

 

 

-def txt2img(prompt: str, negative_prompt: str, prompt_style: str, prompt_style2: str, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int, hr_resize_x: int, hr_resize_y: int, *args):

+def txt2img(prompt: str, negative_prompt: str, prompt_styles,  steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int, hr_resize_x: int, hr_resize_y: int, *args):

     p = StableDiffusionProcessingTxt2Img(

         sd_model=shared.sd_model,

         outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples,

         outpath_grids=opts.outdir_grids or opts.outdir_txt2img_grids,

         prompt=prompt,

-        styles=[prompt_style, prompt_style2],

+        styles=prompt_styles,

         negative_prompt=negative_prompt,

         seed=seed,

         subseed=subseed,

diff --git a/modules/ui.py b/modules/ui.py
index 202e84e5..db198a47 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -180,7 +180,7 @@ def add_style(name: str, prompt: str, negative_prompt: str):
     # reserialize all styles every time we save them

     shared.prompt_styles.save_styles(shared.styles_filename)

 

-    return [gr.Dropdown.update(visible=True, choices=list(shared.prompt_styles.styles)) for _ in range(4)]

+    return [gr.Dropdown.update(visible=True, choices=list(shared.prompt_styles.styles)) for _ in range(2)]

 

 

 def calc_resolution_hires(enable, width, height, hr_scale, hr_resize_x, hr_resize_y):

@@ -197,11 +197,11 @@ def calc_resolution_hires(enable, width, height, hr_scale, hr_resize_x, hr_resiz
     return f"resize: from <span class='resolution'>{p.width}x{p.height}</span> to <span class='resolution'>{p.hr_resize_x or p.hr_upscale_to_x}x{p.hr_resize_y or p.hr_upscale_to_y}</span>"

 

 

-def apply_styles(prompt, prompt_neg, style1_name, style2_name):

-    prompt = shared.prompt_styles.apply_styles_to_prompt(prompt, [style1_name, style2_name])

-    prompt_neg = shared.prompt_styles.apply_negative_styles_to_prompt(prompt_neg, [style1_name, style2_name])

+def apply_styles(prompt, prompt_neg, styles):

+    prompt = shared.prompt_styles.apply_styles_to_prompt(prompt, styles)

+    prompt_neg = shared.prompt_styles.apply_negative_styles_to_prompt(prompt_neg, styles)

 

-    return [gr.Textbox.update(value=prompt), gr.Textbox.update(value=prompt_neg), gr.Dropdown.update(value="None"), gr.Dropdown.update(value="None")]

+    return [gr.Textbox.update(value=prompt), gr.Textbox.update(value=prompt_neg), gr.Dropdown.update(value=[])]

 

 

 def interrogate(image):

@@ -374,13 +374,10 @@ def create_toprow(is_img2img):
                 )

 

             with gr.Row():

-                with gr.Column(scale=1, elem_id="style_pos_col"):

-                    prompt_style = gr.Dropdown(label="Style 1", elem_id=f"{id_part}_style_index", choices=[k for k, v in shared.prompt_styles.styles.items()], value=next(iter(shared.prompt_styles.styles.keys())))

+                prompt_styles = gr.Dropdown(label="Styles", elem_id=f"{id_part}_styles", choices=[k for k, v in shared.prompt_styles.styles.items()], value=[], multiselect=True)

+                create_refresh_button(prompt_styles, shared.prompt_styles.reload, lambda: {"choices": [k for k, v in shared.prompt_styles.styles.items()]}, f"refresh_{id_part}_styles")

 

-                with gr.Column(scale=1, elem_id="style_neg_col"):

-                    prompt_style2 = gr.Dropdown(label="Style 2", elem_id=f"{id_part}_style2_index", choices=[k for k, v in shared.prompt_styles.styles.items()], value=next(iter(shared.prompt_styles.styles.keys())))

-

-    return prompt, prompt_style, negative_prompt, prompt_style2, submit, button_interrogate, button_deepbooru, prompt_style_apply, save_style, paste, token_counter, token_button

+    return prompt, prompt_styles, negative_prompt, submit, button_interrogate, button_deepbooru, prompt_style_apply, save_style, paste, token_counter, token_button

 

 

 def setup_progressbar(*args, **kwargs):

@@ -590,7 +587,7 @@ def create_ui():
     modules.scripts.scripts_txt2img.initialize_scripts(is_img2img=False)

 

     with gr.Blocks(analytics_enabled=False) as txt2img_interface:

-        txt2img_prompt, txt2img_prompt_style, txt2img_negative_prompt, txt2img_prompt_style2, submit, _, _,txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, token_counter, token_button = create_toprow(is_img2img=False)

+        txt2img_prompt, txt2img_prompt_styles, txt2img_negative_prompt, submit, _, _,txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, token_counter, token_button = create_toprow(is_img2img=False)

 

         dummy_component = gr.Label(visible=False)

         txt_prompt_img = gr.File(label="", elem_id="txt2img_prompt_image", file_count="single", type="bytes", visible=False)

@@ -684,8 +681,7 @@ def create_ui():
                 inputs=[

                     txt2img_prompt,

                     txt2img_negative_prompt,

-                    txt2img_prompt_style,

-                    txt2img_prompt_style2,

+                    txt2img_prompt_styles,

                     steps,

                     sampler_index,

                     restore_faces,

@@ -780,7 +776,7 @@ def create_ui():
     modules.scripts.scripts_img2img.initialize_scripts(is_img2img=True)

 

     with gr.Blocks(analytics_enabled=False) as img2img_interface:

-        img2img_prompt, img2img_prompt_style, img2img_negative_prompt, img2img_prompt_style2, submit, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste,token_counter, token_button = create_toprow(is_img2img=True)

+        img2img_prompt, img2img_prompt_styles, img2img_negative_prompt, submit, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste,token_counter, token_button = create_toprow(is_img2img=True)

 

         with gr.Row(elem_id='img2img_progress_row'):

             img2img_prompt_img = gr.File(label="", elem_id="img2img_prompt_image", file_count="single", type="bytes", visible=False)

@@ -921,8 +917,7 @@ def create_ui():
                     dummy_component,

                     img2img_prompt,

                     img2img_negative_prompt,

-                    img2img_prompt_style,

-                    img2img_prompt_style2,

+                    img2img_prompt_styles,

                     init_img,

                     sketch,

                     init_img_with_mask,

@@ -977,7 +972,7 @@ def create_ui():
             )

 

             prompts = [(txt2img_prompt, txt2img_negative_prompt), (img2img_prompt, img2img_negative_prompt)]

-            style_dropdowns = [(txt2img_prompt_style, txt2img_prompt_style2), (img2img_prompt_style, img2img_prompt_style2)]

+            style_dropdowns = [txt2img_prompt_styles, img2img_prompt_styles]

             style_js_funcs = ["update_txt2img_tokens", "update_img2img_tokens"]

 

             for button, (prompt, negative_prompt) in zip([txt2img_save_style, img2img_save_style], prompts):

@@ -987,15 +982,15 @@ def create_ui():
                     # Have to pass empty dummy component here, because the JavaScript and Python function have to accept

                     # the same number of parameters, but we only know the style-name after the JavaScript prompt

                     inputs=[dummy_component, prompt, negative_prompt],

-                    outputs=[txt2img_prompt_style, img2img_prompt_style, txt2img_prompt_style2, img2img_prompt_style2],

+                    outputs=[txt2img_prompt_styles, img2img_prompt_styles],

                 )

 

-            for button, (prompt, negative_prompt), (style1, style2), js_func in zip([txt2img_prompt_style_apply, img2img_prompt_style_apply], prompts, style_dropdowns, style_js_funcs):

+            for button, (prompt, negative_prompt), styles, js_func in zip([txt2img_prompt_style_apply, img2img_prompt_style_apply], prompts, style_dropdowns, style_js_funcs):

                 button.click(

                     fn=apply_styles,

                     _js=js_func,

-                    inputs=[prompt, negative_prompt, style1, style2],

-                    outputs=[prompt, negative_prompt, style1, style2],

+                    inputs=[prompt, negative_prompt, styles],

+                    outputs=[prompt, negative_prompt, styles],

                 )

 

             token_button.click(fn=update_token_counter, inputs=[img2img_prompt, steps], outputs=[token_counter])

@@ -1530,6 +1525,7 @@ def create_ui():
 

         previous_section = None

         current_tab = None

+        current_row = None

         with gr.Tabs(elem_id="settings"):

             for i, (k, item) in enumerate(opts.data_labels.items()):

                 section_must_be_skipped = item.section[0] is None

@@ -1538,10 +1534,14 @@ def create_ui():
                     elem_id, text = item.section

 

                     if current_tab is not None:

+                        current_row.__exit__()

                         current_tab.__exit__()

 

+                    gr.Group()

                     current_tab = gr.TabItem(elem_id="settings_{}".format(elem_id), label=text)

                     current_tab.__enter__()

+                    current_row = gr.Column(variant='compact')

+                    current_row.__enter__()

 

                     previous_section = item.section

 

@@ -1556,6 +1556,7 @@ def create_ui():
                     components.append(component)

 

             if current_tab is not None:

+                current_row.__exit__()

                 current_tab.__exit__()

 

             with gr.TabItem("Actions"):

@@ -1774,7 +1775,13 @@ def create_ui():
             apply_field(x, 'value')

 

         if type(x) == gr.Dropdown:

-            apply_field(x, 'value', lambda val: val in x.choices, getattr(x, 'init_field', None))

+            def check_dropdown(val):

+                if x.multiselect:

+                    return all([value in x.choices for value in val])

+                else:

+                    return val in x.choices

+

+            apply_field(x, 'value', check_dropdown, getattr(x, 'init_field', None))

 

     visit(txt2img_interface, loadsave, "txt2img")

     visit(img2img_interface, loadsave, "img2img")