persist control units state on restarts

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'):