add sd35 controlnets

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

modules/control/units/controlnet.py

@@ -85,6 +85,9 @@ predefined_f1 = {
     "XLabs-AI HED": 'XLabs-AI/flux-controlnet-hed-diffusers'
 }
 predefined_sd3 = {
+    "StabilityAI Canny": 'diffusers-internal-dev/sd35-controlnet-canny-8b',
+    "StabilityAI Depth": 'diffusers-internal-dev/sd35-controlnet-depth-8b',
+    "StabilityAI Blur": 'diffusers-internal-dev/sd35-controlnet-blur-8b',
     "InstantX Canny": 'InstantX/SD3-Controlnet-Canny',
     "InstantX Pose": 'InstantX/SD3-Controlnet-Pose',
     "InstantX Depth": 'InstantX/SD3-Controlnet-Depth',

modules/model_quant.py

@@ -5,6 +5,7 @@ from installer import install, log
 
 bnb = None
 quanto = None
+ao = None
 
 
 def create_bnb_config(kwargs = None, allow_bnb: bool = True):
@@ -12,6 +13,8 @@ def create_bnb_config(kwargs = None, allow_bnb: bool = True):
     if len(shared.opts.bnb_quantization) > 0 and allow_bnb:
         if 'Model' in shared.opts.bnb_quantization:
             load_bnb()
+            if bnb is None:
+                return kwargs
             bnb_config = diffusers.BitsAndBytesConfig(
                 load_in_8bit=shared.opts.bnb_quantization_type in ['fp8'],
                 load_in_4bit=shared.opts.bnb_quantization_type in ['nf4', 'fp4'],
@@ -28,6 +31,44 @@ def create_bnb_config(kwargs = None, allow_bnb: bool = True):
     return kwargs
 
 
+def create_ao_config(kwargs = None, allow_ao: bool = True):
+    from modules import shared
+    if len(shared.opts.torchao_quantization) > 0 and shared.opts.torchao_quantization_mode == 'pre' and allow_ao:
+        if 'Model' in shared.opts.torchao_quantization:
+            load_torchao()
+            if ao is None:
+                return kwargs
+            ao_config = {}
+            # ao_config = diffusers.TorchAoConfig("int8wo") # TODO torchao
+            shared.log.debug(f'Quantization: module=all type=bnb dtype={shared.opts.torchao_quantization_type}')
+            if kwargs is None:
+                return ao_config
+            else:
+                kwargs['quantization_config'] = ao_config
+                return kwargs
+    return kwargs
+
+
+def load_torchao(msg='', silent=False):
+    global ao # pylint: disable=global-statement
+    if ao is not None:
+        return ao
+    install('torchao', quiet=True)
+    try:
+        import torchao
+        ao = torchao
+        fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
+        log.debug(f'Quantization: type=quanto version={ao.__version__} fn={fn}') # pylint: disable=protected-access
+        return ao
+    except Exception as e:
+        if len(msg) > 0:
+            log.error(f"{msg} failed to import optimum.quanto: {e}")
+        ao = None
+        if not silent:
+            raise
+    return None
+
+
 def load_bnb(msg='', silent=False):
     global bnb # pylint: disable=global-statement
     if bnb is not None:

modules/model_sd3.py

@@ -150,6 +150,7 @@ def load_sd3(checkpoint_info, cache_dir=None, config=None):
     shared.log.debug(f'Load model: type=SD3 kwargs={list(kwargs)} repo="{repo_id}"')
 
     kwargs = model_quant.create_bnb_config(kwargs)
+    kwargs = model_quant.create_ao_config(kwargs)
     pipe = loader(
         repo_id,
         torch_dtype=devices.dtype,

modules/sd_models.py

@@ -279,7 +279,7 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model', offload=True):
                 model.eval()
             return model
         sd_model = sd_models_compile.apply_compile_to_model(sd_model, eval_model, ["Model", "VAE", "Text Encoder"], op="eval")
-    if len(shared.opts.torchao_quantization) > 0:
+    if len(shared.opts.torchao_quantization) > 0 and shared.opts.torchao_quantization_mode != 'post':
         sd_model = sd_models_compile.torchao_quantization(sd_model)
 
     if shared.opts.opt_channelslast and hasattr(sd_model, 'unet'):

modules/sd_models_compile.py

@@ -535,7 +535,6 @@ def torchao_quantization(sd_model):
         if hasattr(sd_model, 'transformer') and 'Model' in shared.opts.torchao_quantization:
             modules.append('transformer')
             q.quantize_(sd_model.transformer, fn(), device=devices.device)
-            # sd_model.transformer = q.autoquant(sd_model.transformer, error_on_unseen=False)
         if hasattr(sd_model, 'vae') and 'VAE' in shared.opts.torchao_quantization:
             modules.append('vae')
             q.quantize_(sd_model.vae, fn(), device=devices.device)

modules/shared.py

@@ -477,7 +477,7 @@ options_templates.update(options_section(('sd', "Execution & Models"), {
     "sd_checkpoint_autoload": OptionInfo(True, "Model autoload on start"),
     "sd_checkpoint_autodownload": OptionInfo(True, "Model auto-download on demand"),
     "sd_textencoder_cache": OptionInfo(True, "Cache text encoder results", gr.Checkbox, {"visible": False}),
-    "sd_textencoder_cache_size": OptionInfo(4, "Text encoder results LRU cache size", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
+    "sd_textencoder_cache_size": OptionInfo(4, "Text encoder cache size", gr.Slider, {"minimum": 0, "maximum": 16, "step": 1}),
     "stream_load": OptionInfo(False, "Load models using stream loading method", gr.Checkbox, {"visible": not native }),
     "prompt_mean_norm": OptionInfo(False, "Prompt attention normalization", gr.Checkbox),
     "comma_padding_backtrack": OptionInfo(20, "Prompt padding", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1, "visible": not native }),
@@ -590,6 +590,7 @@ options_templates.update(options_section(('quantization', "Quantization Settings
     "optimum_quanto_weights_type": OptionInfo("qint8", "Optimum.quanto quantization type", gr.Radio, {"choices": ['qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2', 'qint4', 'qint2'], "visible": native}),
     "optimum_quanto_activations_type": OptionInfo("none", "Optimum.quanto quantization activations ", gr.Radio, {"choices": ['none', 'qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2'], "visible": native}),
     "torchao_quantization": OptionInfo([], "TorchAO quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}),
+    "torchao_quantization_mode": OptionInfo("pre", "TorchAO quantization mode", gr.Radio, {"choices": ['pre', 'post'], "visible": native}),
     "torchao_quantization_type": OptionInfo("int8", "TorchAO quantization type", gr.Radio, {"choices": ["int8+act", "int8", "int4", "fp8+act", "fp8", "fpx"], "visible": native}),
     "nncf_compress_weights": OptionInfo([], "NNCF compression enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}),
     "nncf_compress_weights_mode": OptionInfo("INT8", "NNCF compress mode", gr.Radio, {"choices": ['INT8', 'INT8_SYM', 'INT4_ASYM', 'INT4_SYM', 'NF4'] if cmd_opts.use_openvino else ['INT8']}),