persist control units state on restarts

CHANGELOG.md

@@ -49,6 +49,7 @@
   - *note*: *image2video* requires separate 5b model variant  
 - **backend=original** is now marked as in maintenance-only mode  
 - **python 3.12** improved compatibility, automatically handle `setuptools`  
+- **control** persist/reapply units current state on server restart  
 - massive log cleanup
 - minor ui optimizations  
 

javascript/control.js

@@ -1,4 +1,6 @@
 function controlInputMode(inputMode, ...args) {
+  const updateEl = gradioApp().getElementById('control_update');
+  if (updateEl) updateEl.click();
   const tab = gradioApp().querySelector('#control-tab-input button.selected');
   if (!tab) return ['Select', ...args];
   inputMode = tab.innerText;

modules/control/processors.py

@@ -155,7 +155,7 @@ class Processor():
             for k, v in from_config.items():
                 self.load_config[k] = v
 
-    def load(self, processor_id: str = None) -> str:
+    def load(self, processor_id: str = None, force: bool = True) -> str:
         try:
             t0 = time.time()
             processor_id = processor_id or self.processor_id
@@ -165,6 +165,10 @@ class Processor():
             if self.processor_id != processor_id:
                 self.reset()
                 self.config(processor_id)
+            else:
+                if not force and self.model is not None:
+                    log.debug(f'Control Processor: id={processor_id} already loaded')
+                    return ''
             if processor_id not in config:
                 log.error(f'Control Processor unknown: id="{processor_id}" available={list(config)}')
                 return f'Processor failed to load: {processor_id}'

modules/control/run.py

@@ -71,6 +71,20 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini
                 video_skip_frames: int = 0, video_type: str = 'None', video_duration: float = 2.0, video_loop: bool = False, video_pad: int = 0, video_interpolate: int = 0,
                 *input_script_args
         ):
+    # handle optional initialization via ui
+    for u in units:
+        if not u.enabled:
+            continue
+        if u.process_name is not None and u.process_name != '' and u.process_name != 'None':
+            u.process.load(u.process_name, force=False)
+        if u.model_name is not None and u.model_name != '' and u.model_name != 'None':
+            if u.type == 't2i adapter':
+                u.adapter.load(u.model_name, force=False)
+            else:
+                u.controlnet.load(u.model_name, force=False)
+        if u.process is not None and u.process.override is None and u.override is not None:
+            u.process.override = u.override
+
     global instance, pipe, original_pipeline # pylint: disable=global-statement
     t_start = time.time()
     debug(f'Control: type={unit_type} input={inputs} init={inits} type={input_type}')

modules/control/unit.py

@@ -18,6 +18,7 @@ unit_types = ['t2i adapter', 'controlnet', 'xs', 'lite', 'reference', 'ip']
 class Unit(): # mashup of gradio controls and mapping to actual implementation classes
     def __init__(self,
                  # values
+                 index: int = None,
                  enabled: bool = None,
                  strength: float = None,
                  unit_type: str = None,
@@ -40,15 +41,20 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
                  result_txt = None,
                  extra_controls: list = [],
         ):
+        self.controls = [gr.Label(value=unit_type, visible=False)] # separator
+        self.index = index
         self.enabled = enabled or False
         self.type = unit_type
         self.strength = strength or 1.0
+        self.model_strength = model_strength
         self.start = start or 0
         self.end = end or 1
         self.start = min(self.start, self.end)
         self.end = max(self.start, self.end)
         self.mode = None
         # processor always exists, adapter and controlnet are optional
+        self.model_name = None
+        self.process_name = None
         self.process: processors.Processor = processors.Processor()
         self.adapter: t2iadapter.Adapter = None
         self.controlnet: Union[controlnet.ControlNet, xs.ControlNetXS] = None
@@ -155,6 +161,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
                 if isinstance(model_id, str):
                     self.adapter.load(model_id)
                 else:
+                    self.controls.append(model_id)
                     model_id.change(fn=self.adapter.load, inputs=[model_id], outputs=[result_txt], show_progress=True)
             if extra_controls is not None and len(extra_controls) > 0:
                 extra_controls[0].change(fn=adapter_extra, inputs=extra_controls)
@@ -163,6 +170,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
                 if isinstance(model_id, str):
                     self.controlnet.load(model_id)
                 else:
+                    self.controls.append(model_id)
                     model_id.change(fn=self.controlnet.load, inputs=[model_id], outputs=[result_txt], show_progress=True)
                     model_id.change(fn=control_mode_show, inputs=[model_id], outputs=[control_mode], show_progress=False)
             if extra_controls is not None and len(extra_controls) > 0:
@@ -172,6 +180,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
                 if isinstance(model_id, str):
                     self.controlnet.load(model_id)
                 else:
+                    self.controls.append(model_id)
                     model_id.change(fn=self.controlnet.load, inputs=[model_id, extra_controls[0]], outputs=[result_txt], show_progress=True)
             if extra_controls is not None and len(extra_controls) > 0:
                 extra_controls[0].change(fn=controlnetxs_extra, inputs=extra_controls)
@@ -180,6 +189,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
                 if isinstance(model_id, str):
                     self.controlnet.load(model_id)
                 else:
+                    self.controls.append(model_id)
                     model_id.change(fn=self.controlnet.load, inputs=[model_id], outputs=[result_txt], show_progress=True)
             if extra_controls is not None and len(extra_controls) > 0:
                 extra_controls[0].change(fn=controlnetxs_extra, inputs=extra_controls)
@@ -189,14 +199,18 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
                 extra_controls[1].change(fn=reference_extra, inputs=extra_controls)
                 extra_controls[2].change(fn=reference_extra, inputs=extra_controls)
                 extra_controls[3].change(fn=reference_extra, inputs=extra_controls)
+
         if enabled_cb is not None:
+            self.controls.append(enabled_cb)
             enabled_cb.change(fn=enabled_change, inputs=[enabled_cb])
         if model_strength is not None:
+            self.controls.append(model_strength)
             model_strength.change(fn=strength_change, inputs=[model_strength])
         if process_id is not None:
             if isinstance(process_id, str):
                 self.process.load(process_id)
             else:
+                self.controls.append(process_id)
                 process_id.change(fn=self.process.load, inputs=[process_id], outputs=[result_txt], show_progress=True)
         if reset_btn is not None:
             reset_btn.click(fn=reset, inputs=[], outputs=[enabled_cb, model_id, process_id, model_strength])
@@ -207,9 +221,13 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
         if image_reuse is not None:
             image_reuse.click(fn=reuse_image, inputs=[preview_process], outputs=[image_preview]) # return list of images for gallery
         if image_preview is not None:
+            self.controls.append(image_preview)
             image_preview.change(fn=set_image, inputs=[image_preview], outputs=[image_preview])
         if control_start is not None and control_end is not None:
+            self.controls.append(control_start)
+            self.controls.append(control_end)
             control_start.change(fn=control_change, inputs=[control_start, control_end])
             control_end.change(fn=control_change, inputs=[control_start, control_end])
         if control_mode is not None:
+            self.controls.append(control_mode)
             control_mode.change(fn=control_mode_change, inputs=[control_mode])

modules/control/units/controlnet.py

@@ -181,7 +181,7 @@ class ControlNet():
         cls = self.get_class()
         self.model = cls.from_single_file(model_path, **self.load_config)
 
-    def load(self, model_id: str = None) -> str:
+    def load(self, model_id: str = None, force: bool = True) -> str:
         try:
             t0 = time.time()
             model_id = model_id or self.model_id
@@ -197,6 +197,9 @@ class ControlNet():
             if model_path is None:
                 log.error(f'Control {what} model load failed: id="{model_id}" error=unknown model id')
                 return
+            if model_id == self.model_id and not force:
+                log.debug(f'Control {what} model: id="{model_id}" path="{model_path}" already loaded')
+                return
             log.debug(f'Control {what} model loading: id="{model_id}" path="{model_path}"')
             if model_path.endswith('.safetensors'):
                 self.load_safetensors(model_path)
@@ -205,6 +208,9 @@ class ControlNet():
                     model_path = model_path.replace('/bin', '')
                     self.load_config['use_safetensors'] = False
                 cls = self.get_class()
+                if cls is None:
+                    log.error(f'Control {what} model load failed: id="{model_id}" unknown base model')
+                    return
                 self.model = cls.from_pretrained(model_path, **self.load_config)
             if self.dtype is not None:
                 self.model.to(self.dtype)

modules/control/units/lite.py

@@ -78,7 +78,7 @@ class ControlLLLite():
         self.model = None
         self.model_id = None
 
-    def load(self, model_id: str = None) -> str:
+    def load(self, model_id: str = None, force: bool = True) -> str:
         try:
             t0 = time.time()
             model_id = model_id or self.model_id
@@ -94,6 +94,9 @@ class ControlLLLite():
             if model_path is None:
                 log.error(f'Control {what} model load failed: id="{model_id}" error=unknown model id')
                 return
+            if model_id == self.model_id and not force:
+                log.debug(f'Control {what} model: id="{model_id}" path="{model_path}" already loaded')
+                return
             log.debug(f'Control {what} model loading: id="{model_id}" path="{model_path}" {self.load_config}')
             if model_path.endswith('.safetensors'):
                 self.model = ControlNetLLLite(model_path)

modules/control/units/t2iadapter.py

@@ -86,7 +86,7 @@ class Adapter():
         self.model = None
         self.model_id = None
 
-    def load(self, model_id: str = None) -> str:
+    def load(self, model_id: str = None, force: bool = True) -> str:
         try:
             t0 = time.time()
             model_id = model_id or self.model_id
@@ -100,6 +100,9 @@ class Adapter():
             if model_path is None:
                 log.error(f'Control {what} model load failed: id="{model_id}" error=unknown model id')
                 return
+            if model_id == self.model_id and not force:
+                log.debug(f'Control {what} model: id="{model_id}" path="{model_path}" already loaded')
+                return
             log.debug(f'Control {what} model loading: id="{model_id}" path="{model_path}"')
             if model_path.endswith('.pth') or model_path.endswith('.pt') or model_path.endswith('.safetensors') or model_path.endswith('.bin'):
                 from huggingface_hub import hf_hub_download

modules/control/units/xs.py

@@ -74,7 +74,7 @@ class ControlNetXS():
         self.model = None
         self.model_id = None
 
-    def load(self, model_id: str = None, time_embedding_mix: float = 0.0) -> str:
+    def load(self, model_id: str = None, time_embedding_mix: float = 0.0, force: bool = True) -> str:
         try:
             t0 = time.time()
             model_id = model_id or self.model_id
@@ -90,6 +90,9 @@ class ControlNetXS():
             if model_path is None:
                 log.error(f'Control {what} model load failed: id="{model_id}" error=unknown model id')
                 return
+            if model_id == self.model_id and not force:
+                log.debug(f'Control {what} model: id="{model_id}" path="{model_path}" already loaded')
+                return
             self.load_config['time_embedding_mix'] = time_embedding_mix
             log.debug(f'Control {what} model loading: id="{model_id}" path="{model_path}" {self.load_config}')
             if model_path.endswith('.safetensors'):

modules/ui_control.py

@@ -16,6 +16,7 @@ from modules import ui_control_helpers as helpers
 gr_height = None
 max_units = shared.opts.control_max_units
 units: list[unit.Unit] = [] # main state variable
+controls: list[gr.component] = [] # list of gr controls
 debug = shared.log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
 debug('Trace: CONTROL')
 
@@ -41,6 +42,22 @@ def return_controls(res):
         return [None, None, None, None, f'Control: Unexpected response: {type(res)}']
 
 
+def get_units(*values):
+    update = []
+    for c, v in zip(controls, values):
+        if isinstance(c, gr.Label): # unit type indicator
+            what = c.value['label']
+        c.value = v
+        if c.elem_id is not None and c.elem_id.startswith('control_unit'):
+            _prefix, i, name = c.elem_id.split('-')
+            update.append({ 'type': what, 'index': int(i), 'name': name, 'value': v })
+    for u in update:
+        for i in range(len(units)):
+            if units[i].type == u['type'] and units[i].index == u['index']:
+                setattr(units[i], u['name'], u['value'])
+                break
+
+
 def generate_click(job_id: str, active_tab: str, *args):
     while helpers.busy:
         time.sleep(0.01)
@@ -197,22 +214,23 @@ def create_ui(_blocks: gr.Blocks=None):
                             enabled = True if i==0 else False
                             with gr.Accordion(f'ControlNet unit {i+1}', visible= i < num_controlnet_units.value, elem_classes='control-unit') as unit_ui:
                                 with gr.Row():
-                                    enabled_cb = gr.Checkbox(enabled, label='', container=False, show_label=False)
-                                    process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None')
-                                    model_id = gr.Dropdown(label="ControlNet", choices=controlnet.list_models(), value='None')
+                                    enabled_cb = gr.Checkbox(enabled, label='', container=False, show_label=False, elem_id=f'control_unit-{i}-enabled')
+                                    process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None', elem_id=f'control_unit-{i}-process_name')
+                                    model_id = gr.Dropdown(label="ControlNet", choices=controlnet.list_models(), value='None', elem_id=f'control_unit-{i}-model_name')
                                     ui_common.create_refresh_button(model_id, controlnet.list_models, lambda: {"choices": controlnet.list_models(refresh=True)}, f'refresh_controlnet_models_{i}')
-                                    model_strength = gr.Slider(label="CN Strength", minimum=0.01, maximum=2.0, step=0.01, value=1.0)
-                                    control_start = gr.Slider(label="Start", minimum=0.0, maximum=1.0, step=0.05, value=0)
-                                    control_end = gr.Slider(label="End", minimum=0.0, maximum=1.0, step=0.05, value=1.0)
-                                    control_mode = gr.Dropdown(label="CN Mode", choices=['', 'Canny', 'Tile', 'Depth', 'Blur', 'Pose', 'Gray', 'LQ'], value=0, type='index', visible=False)
+                                    model_strength = gr.Slider(label="CN Strength", minimum=0.01, maximum=2.0, step=0.01, value=1.0, elem_id=f'control_unit-{i}-strength')
+                                    control_start = gr.Slider(label="Start", minimum=0.0, maximum=1.0, step=0.05, value=0, elem_id=f'control_unit-{i}-start')
+                                    control_end = gr.Slider(label="End", minimum=0.0, maximum=1.0, step=0.05, value=1.0, elem_id=f'control_unit-{i}-end')
+                                    control_mode = gr.Dropdown(label="CN Mode", choices=['', 'Canny', 'Tile', 'Depth', 'Blur', 'Pose', 'Gray', 'LQ'], value=0, type='index', visible=False, elem_id=f'control_unit-{i}-mode')
                                     reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
                                     image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
                                     image_reuse= ui_components.ToolButton(value=ui_symbols.reuse)
                                     process_btn= ui_components.ToolButton(value=ui_symbols.preview)
-                                    image_preview = gr.Image(label="Input", type="pil", source="upload", height=128, width=128, visible=False, interactive=True, show_label=False, show_download_button=False, container=False)
+                                    image_preview = gr.Image(label="Input", type="pil", source="upload", height=128, width=128, visible=False, interactive=True, show_label=False, show_download_button=False, container=False, elem_id=f'control_unit-{i}-override')
                             controlnet_ui_units.append(unit_ui)
                             units.append(unit.Unit(
                                 unit_type = 'controlnet',
+                                index = i,
                                 enabled = enabled,
                                 result_txt = result_txt,
                                 enabled_cb = enabled_cb,
@@ -247,19 +265,20 @@ def create_ui(_blocks: gr.Blocks=None):
                             enabled = True if i==0 else False
                             with gr.Accordion(f'T2I-Adapter unit {i+1}', visible= i < num_adapter_units.value, elem_classes='control-unit') as unit_ui:
                                 with gr.Row():
-                                    enabled_cb = gr.Checkbox(enabled, label='', container=False, show_label=False)
-                                    process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None')
-                                    model_id = gr.Dropdown(label="Adapter", choices=t2iadapter.list_models(), value='None')
+                                    enabled_cb = gr.Checkbox(enabled, label='', container=False, show_label=False, elem_id=f'control_unit-{i}-enabled')
+                                    process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None', elem_id=f'control_unit-{i}-process_name')
+                                    model_id = gr.Dropdown(label="Adapter", choices=t2iadapter.list_models(), value='None', elem_id=f'control_unit-{i}-model_name')
                                     ui_common.create_refresh_button(model_id, t2iadapter.list_models, lambda: {"choices": t2iadapter.list_models(refresh=True)}, f'refresh_adapter_models_{i}')
-                                    model_strength = gr.Slider(label="T2I Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0)
+                                    model_strength = gr.Slider(label="T2I Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0, elem_id=f'control_unit-{i}-strength')
                                     reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
                                     image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
                                     image_reuse= ui_components.ToolButton(value=ui_symbols.reuse)
                                     process_btn= ui_components.ToolButton(value=ui_symbols.preview)
-                                    image_preview = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=128, width=128, visible=False)
+                                    image_preview = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=128, width=128, visible=False, elem_id=f'control_unit-{i}-override')
                             adapter_ui_units.append(unit_ui)
                             units.append(unit.Unit(
                                 unit_type = 't2i adapter',
+                                index = i,
                                 enabled = enabled,
                                 result_txt = result_txt,
                                 enabled_cb = enabled_cb,
@@ -291,21 +310,22 @@ def create_ui(_blocks: gr.Blocks=None):
                             enabled = True if i==0 else False
                             with gr.Accordion(f'ControlNet-XS unit {i+1}', visible= i < num_controlnet_units.value, elem_classes='control-unit') as unit_ui:
                                 with gr.Row():
-                                    enabled_cb = gr.Checkbox(enabled, label='', container=False, show_label=False)
-                                    process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None')
-                                    model_id = gr.Dropdown(label="ControlNet-XS", choices=xs.list_models(), value='None')
+                                    enabled_cb = gr.Checkbox(enabled, label='', container=False, show_label=False, elem_id=f'control_unit-{i}-enabled')
+                                    process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None', elem_id=f'control_unit-{i}-process_name')
+                                    model_id = gr.Dropdown(label="ControlNet-XS", choices=xs.list_models(), value='None', elem_id=f'control_unit-{i}-model_name')
                                     ui_common.create_refresh_button(model_id, xs.list_models, lambda: {"choices": xs.list_models(refresh=True)}, f'refresh_xs_models_{i}')
-                                    model_strength = gr.Slider(label="CN Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0)
-                                    control_start = gr.Slider(label="Start", minimum=0.0, maximum=1.0, step=0.05, value=0)
-                                    control_end = gr.Slider(label="End", minimum=0.0, maximum=1.0, step=0.05, value=1.0)
+                                    model_strength = gr.Slider(label="CN Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0, elem_id=f'control_unit-{i}-strength')
+                                    control_start = gr.Slider(label="Start", minimum=0.0, maximum=1.0, step=0.05, value=0, elem_id=f'control_unit-{i}-start')
+                                    control_end = gr.Slider(label="End", minimum=0.0, maximum=1.0, step=0.05, value=1.0, elem_id=f'control_unit-{i}-end')
                                     reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
                                     image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
                                     image_reuse= ui_components.ToolButton(value=ui_symbols.reuse)
                                     process_btn= ui_components.ToolButton(value=ui_symbols.preview)
-                                    image_preview = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=128, width=128, visible=False)
+                                    image_preview = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=128, width=128, visible=False, elem_id=f'control_unit-{i}-override')
                             controlnetxs_ui_units.append(unit_ui)
                             units.append(unit.Unit(
                                 unit_type = 'xs',
+                                index = 1,
                                 enabled = enabled,
                                 result_txt = result_txt,
                                 enabled_cb = enabled_cb,
@@ -338,19 +358,20 @@ def create_ui(_blocks: gr.Blocks=None):
                             enabled = True if i==0 else False
                             with gr.Accordion(f'Control-LLLite unit {i+1}', visible= i < num_lite_units.value, elem_classes='control-unit') as unit_ui:
                                 with gr.Row():
-                                    enabled_cb = gr.Checkbox(enabled, label='', container=False, show_label=False)
-                                    process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None')
-                                    model_id = gr.Dropdown(label="Model", choices=lite.list_models(), value='None')
+                                    enabled_cb = gr.Checkbox(enabled, label='', container=False, show_label=False, elem_id=f'control_unit-{i}-enabled')
+                                    process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None', elem_id=f'control_unit-{i}-process_name')
+                                    model_id = gr.Dropdown(label="Model", choices=lite.list_models(), value='None', elem_id=f'control_unit-{i}-model_name')
                                     ui_common.create_refresh_button(model_id, lite.list_models, lambda: {"choices": lite.list_models(refresh=True)}, f'refresh_lite_models_{i}')
-                                    model_strength = gr.Slider(label="CN Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0)
+                                    model_strength = gr.Slider(label="CN Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0, elem_id=f'control_unit-{i}-strength')
                                     reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
                                     image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
                                     image_reuse= ui_components.ToolButton(value=ui_symbols.reuse)
-                                    image_preview = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=128, width=128, visible=False)
+                                    image_preview = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=128, width=128, visible=False, elem_id=f'control_unit-{i}-override')
                                     process_btn= ui_components.ToolButton(value=ui_symbols.preview)
                             lite_ui_units.append(unit_ui)
                             units.append(unit.Unit(
                                 unit_type = 'lite',
+                                index = i,
                                 enabled = enabled,
                                 result_txt = result_txt,
                                 enabled_cb = enabled_cb,
@@ -383,16 +404,17 @@ def create_ui(_blocks: gr.Blocks=None):
                             enabled = True if i==0 else False
                             with gr.Accordion(f'Reference unit {i+1}', visible=True, elem_classes='control-unit') as unit_ui:
                                 with gr.Row():
-                                    enabled_cb = gr.Checkbox(enabled, label='', container=False, show_label=False)
-                                    model_id = gr.Dropdown(label="Reference", choices=reference.list_models(), value='Reference', visible=False)
-                                    model_strength = gr.Slider(label="CN Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0, visible=False)
+                                    enabled_cb = gr.Checkbox(enabled, label='', container=False, show_label=False, elem_id=f'control_unit-{i}-enabled')
+                                    model_id = gr.Dropdown(label="Reference", choices=reference.list_models(), value='Reference', visible=False, elem_id=f'control_unit-{i}-model_name')
+                                    model_strength = gr.Slider(label="CN Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0, visible=False, elem_id=f'control_unit-{i}-strength')
                                     reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
                                     image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
                                     image_reuse= ui_components.ToolButton(value=ui_symbols.reuse)
-                                    image_preview = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=128, width=128, visible=False)
+                                    image_preview = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=128, width=128, visible=False, elem_id=f'control_unit-{i}-override')
                                     process_btn= ui_components.ToolButton(value=ui_symbols.preview)
                             units.append(unit.Unit(
                                 unit_type = 'reference',
+                                index = i,
                                 enabled = enabled,
                                 result_txt = result_txt,
                                 enabled_cb = enabled_cb,
@@ -468,6 +490,12 @@ def create_ui(_blocks: gr.Blocks=None):
                 btn.click(fn=helpers.copy_input, inputs=[input_mode, btn, input_image, input_resize, input_inpaint], outputs=[input_image, input_resize, input_inpaint], _js='controlInputMode')
                 btn.click(fn=helpers.transfer_input, inputs=[btn], outputs=[input_image, input_resize, input_inpaint] + input_buttons)
 
+            # hidden button to update gradio control values
+            for u in units:
+                controls.extend(u.controls)
+            btn_update = gr.Button('Update', interactive=True, visible=False, elem_id='control_update')
+            btn_update.click(fn=get_units, inputs=controls, outputs=[], show_progress=True, queue=False)
+
             show_preview.change(fn=lambda x: gr.update(visible=x), inputs=[show_preview], outputs=[column_preview])
             input_type.change(fn=lambda x: gr.update(visible=x == 2), inputs=[input_type], outputs=[column_init])
             btn_prompt_counter.click(fn=call_queue.wrap_queued_call(ui_common.update_token_counter), inputs=[prompt, steps], outputs=[prompt_counter])
@@ -485,6 +513,7 @@ def create_ui(_blocks: gr.Blocks=None):
                 show_progress=True,
                 queue=False,
             )
+
             prompt.submit(**select_dict)
             btn_generate.click(**select_dict)
             for ctrl in [input_image, input_resize, input_video, input_batch, input_folder, init_image, init_video, init_batch, init_folder, tab_image, tab_video, tab_batch, tab_folder, tab_image_init, tab_video_init, tab_batch_init, tab_folder_init]: