Add IPAdapter advanced weighting [Instant Style] (#2767)

* rename variable

* nit

* add transformer index

* nit

* API support

* Fix transformer_index alignment issue

* Add weight control panel

* fix sdxl

* free edit

* clear advanced weighting

* nit

* Add assertion

Author: huchenlei

internal_controlnet/external_code.py

@@ -195,6 +195,9 @@ class ControlNetUnit:
     # For T2IAdapter
     # - SD1.5: 5 weights (4 encoder block + 1 middle block)
     # - SDXL: 4 weights (3 encoder block + 1 middle block)
+    # For IPAdapter
+    # - SD15: 16 (6 input blocks + 9 output blocks + 1 middle block)
+    # - SDXL: 11 weights (4 input blocks + 6 output blocks + 1 middle block)
     # Note1: Setting advanced weighting will disable `soft_injection`, i.e.
     # It is recommended to set ControlMode = BALANCED when using `advanced_weighting`.
     # Note2: The field `weight` is still used in some places, e.g. reference_only,

requirements.txt

@@ -4,4 +4,5 @@ opencv-python>=4.8.0
 svglib
 addict
 yapf
-albumentations==1.4.3
+albumentations==1.4.3
+matplotlib

scripts/controlnet.py

@@ -927,7 +927,7 @@ def controlnet_main_entry(self, p):
         logger.info(f"unit_separate = {batch_option_uint_separate}, style_align = {batch_option_style_align}")
 
         detected_maps = []
-        forward_params = []
+        forward_params: List[ControlParams] = []
         post_processors = []
 
         # cache stuff
@@ -1177,10 +1177,21 @@ def recolor_intensity_post_processing(x, i):
 
         for i, param in enumerate(forward_params):
             if param.control_model_type == ControlModelType.IPAdapter:
+                if param.advanced_weighting is not None:
+                    logger.info(f"IP-Adapter using advanced weighting {param.advanced_weighting}")
+                    assert len(param.advanced_weighting) == global_state.get_sd_version().transformer_block_num
+                    # Convert advanced weighting list to dict
+                    weight = {
+                        i: w
+                        for i, w in enumerate(param.advanced_weighting)
+                        if w > 0
+                    }
+                else:
+                    weight = param.weight
                 param.control_model.hook(
                     model=unet,
                     preprocessor_outputs=param.hint_cond,
-                    weight=param.weight,
+                    weight=weight,
                     dtype=torch.float32,
                     start=param.start_guidance_percent,
                     end=param.stop_guidance_percent

scripts/controlnet_ui/advanced_weight_control.py

@@ -0,0 +1,230 @@
+import gradio as gr
+import matplotlib.pyplot as plt
+from matplotlib.patches import Patch
+import io
+import json
+from PIL import Image
+from typing import List
+
+from scripts.enums import StableDiffusionVersion
+from scripts.global_state import get_sd_version
+from scripts.ipadapter.weight import calc_weights
+
+
+INPUT_BLOCK_COLOR = "#61bdee"
+MIDDLE_BLOCK_COLOR = "#e2e2e2"
+OUTPUT_BLOCK_COLOR = "#dc6e55"
+
+
+def get_bar_colors(
+    sd_version: StableDiffusionVersion, input_color, middle_color, output_color
+):
+    middle_block_idx = 4 if sd_version == StableDiffusionVersion.SDXL else 6
+
+    def get_color(idx):
+        if idx < middle_block_idx:
+            return input_color
+        elif idx == middle_block_idx:
+            return middle_color
+        else:
+            return output_color
+
+    return [get_color(i) for i in range(sd_version.transformer_block_num)]
+
+
+def plot_weights(
+    numbers: List[float],
+    colors: List[str],
+):
+    # Create a bar chart
+    plt.figure(figsize=(8, 4))
+    plt.bar(range(len(numbers)), numbers, color=colors)
+    plt.xlabel("Transformer Index")
+    plt.ylabel("Weight")
+    plt.legend(
+        handles=[
+            Patch(color=color, label=label)
+            for color, label in (
+                (INPUT_BLOCK_COLOR, "Input Block"),
+                (MIDDLE_BLOCK_COLOR, "Middle Block"),
+                (OUTPUT_BLOCK_COLOR, "Output Block"),
+            )
+        ],
+        loc="best",
+    )
+
+    # Save the plot to a BytesIO buffer
+    buffer = io.BytesIO()
+    plt.savefig(buffer, format="png")
+    plt.close()
+    buffer.seek(0)
+
+    # Convert the buffer to a PIL image and return it
+    image = Image.open(buffer)
+    return image
+
+
+class AdvancedWeightControl:
+    def __init__(self):
+        self.group = None
+        self.weight_type = None
+        self.weight_plot = None
+        self.weight_editor = None
+        self.weight_composition = None
+
+    def render(self):
+        with gr.Group(visible=False) as self.group:
+            with gr.Row():
+                self.weight_type = gr.Dropdown(
+                    choices=[
+                        "normal",
+                        "ease in",
+                        "ease out",
+                        "ease in-out",
+                        "reverse in-out",
+                        "weak input",
+                        "weak output",
+                        "weak middle",
+                        "strong middle",
+                        "style transfer",
+                        "composition",
+                        "strong style transfer",
+                        "style and composition",
+                        "strong style and composition",
+                    ],
+                    label="Weight Type",
+                    value="normal",
+                )
+                self.weight_composition = gr.Slider(
+                    label="Composition Weight",
+                    minimum=0,
+                    maximum=2.0,
+                    value=1.0,
+                    step=0.01,
+                    visible=False,
+                )
+                self.weight_editor = gr.Textbox(label="Weights", visible=False)
+
+            self.weight_plot = gr.Image(
+                value=None,
+                label="Weight Plot",
+                interactive=False,
+                visible=False,
+            )
+
+    def register_callbacks(
+        self,
+        weight_input: gr.Slider,
+        advanced_weighting: gr.State,
+        control_type: gr.Radio,
+        update_unit_counter: gr.Number,
+    ):
+        def advanced_weighting_supported(control_type: str) -> bool:
+            return control_type in ("IP-Adapter", "Instant-ID")
+
+        self.weight_type.change(
+            fn=lambda weight_type: gr.update(
+                visible=weight_type
+                in ("style and composition", "strong style and composition")
+            ),
+            inputs=[self.weight_type],
+            outputs=[self.weight_composition],
+        )
+
+        def update_weight_textbox(
+            control_type: str,
+            weight_type: str,
+            weight: float,
+            weight_composition: float,
+        ):
+            if not advanced_weighting_supported(control_type):
+                return gr.update()
+
+            sd_version = get_sd_version()
+            weights = calc_weights(weight_type, weight, sd_version, weight_composition)
+            return gr.update(value=str([round(w, 2) for w in weights]), visible=True)
+
+        trigger_inputs = [self.weight_type, weight_input, self.weight_composition]
+        for trigger_input in trigger_inputs:
+            trigger_input.change(
+                fn=update_weight_textbox,
+                inputs=[
+                    control_type,
+                    self.weight_type,
+                    weight_input,
+                    self.weight_composition,
+                ],
+                outputs=[self.weight_editor],
+            )
+
+        def update_plot(weights_string: str):
+            try:
+                weights = json.loads(weights_string)
+                assert isinstance(weights, list)
+            except Exception:
+                return gr.update(visible=False)
+
+            sd_version = get_sd_version()
+            weight_plot = plot_weights(
+                weights,
+                get_bar_colors(
+                    sd_version,
+                    input_color=INPUT_BLOCK_COLOR,
+                    middle_color=MIDDLE_BLOCK_COLOR,
+                    output_color=OUTPUT_BLOCK_COLOR,
+                ),
+            )
+            return gr.update(value=weight_plot, visible=True)
+
+        def update_advanced_weighting(weights_string: str):
+            try:
+                weights = json.loads(weights_string)
+                assert isinstance(weights, list)
+            except Exception:
+                return None
+            return weights
+
+        self.weight_editor.change(
+            fn=update_plot,
+            inputs=[self.weight_editor],
+            outputs=[self.weight_plot],
+        )
+
+        self.weight_editor.change(
+            fn=update_advanced_weighting,
+            inputs=[self.weight_editor],
+            outputs=[advanced_weighting],
+        ).then(
+            fn=lambda x: gr.update(value=x + 1),
+            inputs=[update_unit_counter],
+            outputs=[update_unit_counter],
+        )  # Necessary to flush gr.State change to unit state.
+
+        # TODO: Expose advanced weighting control for other control types.
+        def control_type_change(control_type: str, old_weights):
+            supported = advanced_weighting_supported(control_type)
+            if supported:
+                return (
+                    gr.update(visible=supported),
+                    old_weights,
+                    gr.update(),
+                    gr.update(),
+                )
+            else:
+                return (
+                    gr.update(visible=supported),
+                    None,
+                    gr.update(visible=False),
+                    gr.update(visible=False),
+                )
+
+        control_type.change(
+            fn=control_type_change,
+            inputs=[control_type, advanced_weighting],
+            outputs=[
+                self.group,
+                advanced_weighting,
+                self.weight_editor,
+                self.weight_plot,
+            ],
+        )

scripts/controlnet_ui/controlnet_ui_group.py

@@ -18,6 +18,7 @@
 from scripts.controlnet_ui.preset import ControlNetPresetUI
 from scripts.controlnet_ui.tool_button import ToolButton
 from scripts.controlnet_ui.photopea import Photopea
+from scripts.controlnet_ui.advanced_weight_control import AdvancedWeightControl
 from scripts.enums import InputMode
 from modules import shared
 from modules.ui_components import FormRow
@@ -289,6 +290,7 @@ def __init__(
         self.input_mode = gr.State(InputMode.SIMPLE)
         self.inpaint_crop_input_image = None
         self.hr_option = None
+        self.advanced_weight_control = AdvancedWeightControl()
         self.batch_image_dir_state = None
         self.output_dir_state = None
 
@@ -640,6 +642,8 @@ def render(self, tabname: str, elem_id_tabname: str) -> None:
             visible=False,
         )
 
+        self.advanced_weight_control.render()
+
         self.preset_panel = ControlNetPresetUI(
             id_prefix=f"{elem_id_tabname}_{tabname}_"
         )
@@ -675,6 +679,8 @@ def render(self, tabname: str, elem_id_tabname: str) -> None:
             self.control_mode,
             self.inpaint_crop_input_image,
             self.hr_option,
+            self.save_detected_map,
+            self.advanced_weighting,
         )
 
         unit = gr.State(self.default_unit)
@@ -1256,6 +1262,12 @@ def register_core_callbacks(self):
                 for key in vars(external_code.ControlNetUnit()).keys()
             ],
         )
+        self.advanced_weight_control.register_callbacks(
+            self.weight,
+            self.advanced_weighting,
+            self.type_filter,
+            self.update_unit_counter,
+        )
         if self.is_img2img:
             self.register_img2img_same_input()