reduce memory leak on control unit change

Signed-off-by: Vladimir Mandic <mandic00@live.com>

TODO.md

@@ -74,7 +74,6 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
 - control: support scripts via api
 - fc: autodetect distilled based on model
 - fc: autodetect tensor format based on model
-- flux transformer from-single-file with quant
 - flux: loader for civitai nf4 models
 - hypertile: vae breaks when using non-standard sizes
 - install: enable ROCm for windows when available
@@ -84,9 +83,8 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
 - lora: add t5 key support for sd35/f1
 - lora: maybe force imediate quantization
 - lora: support pre-quantized flux
-- model fix: cogview4: balanced offload does not work for GlmModel
+- model load: cogview4: balanced offload does not work for GlmModel
 - model load: add ChromaFillPipeline, ChromaControlPipeline, ChromaImg2ImgPipeline etc when available
-- model load: chroma transformer from-single-file with quant
 - model load: force-reloading entire model as loading transformers only leads to massive memory usage
 - model load: implement model in-memory caching
 - modernui: monkey-patch for missing tabs.select event

modules/control/processors.py

@@ -139,8 +139,9 @@ class Processor():
     def reset(self, processor_id: str = None):
         if self.model is not None:
             debug(f'Control Processor unloaded: id="{self.processor_id}"')
-        self.model = None
-        self.processor_id = processor_id
+            self.model = None
+            self.processor_id = processor_id
+            devices.torch_gc(force=True, reason='processor')
         # self.override = None
         # devices.torch_gc()
         self.load_config = { 'cache_dir': cache_dir }

modules/control/unit.py

@@ -93,16 +93,6 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
         # control tile
         self.tile = '1x1'
 
-        def reset():
-            if self.process is not None:
-                self.process.reset()
-            if self.adapter is not None:
-                self.adapter.reset()
-            if self.controlnet is not None:
-                self.controlnet.reset()
-            self.override = None
-            return [True, 'None', 'None', 1.0] # reset ui values
-
         def enabled_change(val):
             self.enabled = val
 
@@ -241,7 +231,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
                 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])
+            reset_btn.click(fn=self.reset, inputs=[], outputs=[enabled_cb, model_id, process_id, model_strength])
         if preview_btn is not None:
             preview_btn.click(fn=self.process.preview, inputs=[], outputs=[preview_process]) # return list of images for gallery
         if image_upload is not None:
@@ -262,3 +252,13 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
         if control_tile is not None:
             self.controls.append(control_tile)
             control_tile.change(fn=control_tile_change, inputs=[control_tile])
+
+    def reset(self):
+        if self.process is not None:
+            self.process.reset()
+        if self.adapter is not None:
+            self.adapter.reset()
+        if self.controlnet is not None:
+            self.controlnet.reset()
+        self.override = None
+        return [True, 'None', 'None', 1.0] # reset ui values

modules/control/units/controlnet.py

@@ -195,8 +195,9 @@ class ControlNet():
     def reset(self):
         if self.model is not None:
             debug_log(f'Control {what} model unloaded')
-        self.model = None
-        self.model_id = None
+            self.model = None
+            self.model_id = None
+            devices.torch_gc(force=True, reason='controlnet')
 
     def get_class(self, model_id:str=''):
         from modules import shared
@@ -226,7 +227,7 @@ class ControlNet():
             return None, None
         return cls, config
 
-    def load_safetensors(self, model_id, model_path):
+    def load_safetensors(self, model_id, model_path, cls, config):
         name = os.path.splitext(model_path)[0]
         config_path = None
         if not os.path.exists(model_path):
@@ -251,11 +252,7 @@ class ControlNet():
             config_path = f'{name}.json'
         if config_path is not None:
             self.load_config['original_config_file '] = config_path
-        cls, config = self.get_class(model_id)
-        if cls is None:
-            log.error(f'Control {what} model load: unknown base model')
-        else:
-            self.model = cls.from_single_file(model_path, config=config, **self.load_config)
+        self.model = cls.from_single_file(model_path, config=config, **self.load_config)
 
     def load(self, model_id: str = None, force: bool = True) -> str:
         with load_lock:
@@ -281,9 +278,13 @@ class ControlNet():
                     # 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}"')
-                cls, _config = self.get_class(model_id)
+                cls, config = self.get_class(model_id)
+                if cls is None:
+                    log.error(f'Control {what} model load: id="{model_id}" unknown base model')
+                    return
+                self.reset()
                 if model_path.endswith('.safetensors'):
-                    self.load_safetensors(model_id, model_path)
+                    self.load_safetensors(model_id, model_path, cls, config)
                 else:
                     kwargs = {}
                     if '/bin' in model_path:
@@ -291,9 +292,6 @@ class ControlNet():
                         self.load_config['use_safetensors'] = False
                     else:
                         self.load_config['use_safetensors'] = True
-                    if cls is None:
-                        log.error(f'Control {what} model load: id="{model_id}" unknown base model')
-                        return
                     if variants.get(model_id, None) is not None:
                         kwargs['variant'] = variants[model_id]
                     try:

pipelines/model_cogview.py

@@ -56,12 +56,12 @@ def load_cogview4(checkpoint_info, diffusers_load_config={}):
 
     load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
     shared.log.debug(f'Load model: type=CogView4 te="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
-    text_encoder = transformers.AutoModelForCausalLM.from_pretrained(
+    text_encoder = transformers.AutoModelForCausalLM.from_pretrained( # TODO model load: cogview4 balanced offload does not work for GlmModel
         repo_id,
         subfolder="text_encoder",
         cache_dir=shared.opts.diffusers_dir,
         **load_args,
-        **quant_args,
+        # **quant_args,
     )
 
     load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
@@ -76,6 +76,6 @@ def load_cogview4(checkpoint_info, diffusers_load_config={}):
     if shared.opts.diffusers_eval:
         pipe.text_encoder.eval()
         pipe.transformer.eval()
-    pipe.enable_model_cpu_offload() # TODO model fix: cogview4: balanced offload does not work for GlmModel
+    pipe.enable_model_cpu_offload()
     devices.torch_gc()
     return pipe

scripts/xyz/xyz_grid_shared.py

@@ -321,6 +321,8 @@ def apply_control(field):
             )
             shared.log.debug(f'XYZ grid apply control: {field}="{x}" unit={unit}')
             if len(run.unit.current) > 0:
+                if hasattr(run.unit.current[0], 'reset'):
+                    run.unit.current[0].reset()
                 run.unit.current[0] = unit
             else:
                 run.unit.current = [unit]