sdnq add xyz grid support, improve offloading compatibility

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

CHANGELOG.md

@@ -3,14 +3,18 @@
 ## Update for 2025-06-25
 
 - **Changes**  
-  - Use Diffusers version of *OmniGen*  
+  - Add [JoyCaption Beta](https://huggingface.co/fancyfeast/llama-joycaption-beta-one-hf-llava) support  
   - Support Remote VAE with *Omnigen, Lumina 2 and PixArt*  
+  - Use Diffusers version of *OmniGen*  
 
 - **SDNQ Quantization**  
   - Add modules_to_not_convert support for post mode  
   - Fix Qwen 2.5 with int8 matmul  
   - Fix Dora loading  
   - Remove per layer GC  
+  - Improve offload compatibility  
+  - Add support for XYZ grid to test quantization modes  
+    *note*: you need to enable quantization and choose what it applies on, then xyz grid can change quantization mode  
 
 - **API**
   - Add `/sdapi/v1/lora?lora=<lora_name>` endpoint that returns full lora info and metadata  
@@ -27,6 +31,7 @@
   - Case-insensitive sampler name matching  
   - Fix delete file with gallery views  
   - Add `SD_SAVE_DEBUG` env variable to report all params and metadata save operations as they happen  
+  - Fix TAESD model type detection  
 
 ## Update for 2025-06-16
 

modules/api/loras.py

@@ -8,7 +8,6 @@ def get_lora(lora: str) -> dict:
     obj = lora_load.available_networks[lora]
     obj.info = obj.get_info()
     obj.desc = obj.get_desc()
-    print('HERE', obj)
     return obj.__dict__
 
 def get_loras():

modules/hidiffusion/__init__.py

@@ -41,5 +41,5 @@ def apply(p, model_type):
 
 def unapply():
     pipe = shared.sd_model.pipe if hasattr(shared.sd_model, 'pipe') else shared.sd_model
-    if hasattr(pipe, 'unet'):
+    if hasattr(pipe, 'unet') and pipe.unet is not None:
         hidiffusion.remove_hidiffusion(pipe)

modules/interrogate/joycaption.py

@@ -58,9 +58,12 @@ opts = JoyOptions()
 
 
 @torch.no_grad()
-def predict(question: str, image):
+def predict(question: str, image, vqa_model: str = None) -> str:
     global llava_model, processor # pylint: disable=global-statement
     opts.max_new_tokens = shared.opts.interrogate_vlm_max_length
+    if vqa_model is not None and opts.repo != vqa_model:
+        opts.repo = vqa_model
+        llava_model = None
     if llava_model is None:
         shared.log.info(f'Interrogate: type=vlm model="JoyCaption" {str(opts)}')
         processor = AutoProcessor.from_pretrained(opts.repo)

modules/interrogate/vqa.py

@@ -38,7 +38,8 @@ vlm_models = {
     "ToriiGate 0.4 2B": "Minthy/ToriiGate-v0.4-2B",
     "ToriiGate 0.4 7B": "Minthy/ToriiGate-v0.4-7B",
     "ViLT Base": "dandelin/vilt-b32-finetuned-vqa", # 0.5GB
-    "JoyCaption": "fancyfeast/llama-joycaption-alpha-two-hf-llava", # 17.4GB
+    "JoyCaption Alpha": "fancyfeast/llama-joycaption-alpha-two-hf-llava", # 17.4GB
+    "JoyCaption Beta": "fancyfeast/llama-joycaption-beta-one-hf-llava", # 17.4GB
     "JoyTag": "fancyfeast/joytag", # 0.7GB
     "AIDC Ovis2 1B": "AIDC-AI/Ovis2-1B",
     "AIDC Ovis2 2B": "AIDC-AI/Ovis2-2B",
@@ -583,7 +584,7 @@ def interrogate(question:str='', system_prompt:str=None, prompt:str=None, image:
             answer = joytag.predict(image)
         elif 'joycaption' in vqa_model.lower():
             from modules.interrogate import joycaption
-            answer = joycaption.predict(question, image)
+            answer = joycaption.predict(question, image, vqa_model)
         elif 'deepseek' in vqa_model.lower():
             from modules.interrogate import deepseek
             answer = deepseek.predict(question, image, vqa_model)

modules/ipadapter.py

@@ -146,7 +146,7 @@ def unapply(pipe, unload: bool = False): # pylint: disable=arguments-differ
             if unload:
                 shared.log.debug('IP adapter unload')
                 pipe.unload_ip_adapter()
-        if hasattr(pipe, 'unet'):
+        if hasattr(pipe, 'unet') and pipe.unet is not None:
             module = pipe.unet
         elif hasattr(pipe, 'transformer'):
             module = pipe.transformer

modules/model_quant.py

@@ -113,11 +113,15 @@ def create_sdnq_config(kwargs = None, allow_sdnq: bool = True, module: str = 'Mo
             diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
             transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
 
-            if weights_dtype is None:
-                if shared.opts.sdnq_quantize_weights_mode_te != "default" and module in {"TE", "LLM"}:
-                    weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
-                else:
+            if weights_dtype is None and module in {"TE", "LLM"}:
+                if shared.opts.sdnq_quantize_weights_mode_te == "none":
+                    return None
+                elif shared.opts.sdnq_quantize_weights_mode_te == "same as model" or shared.opts.sdnq_quantize_weights_mode_te == "default":
                     weights_dtype = shared.opts.sdnq_quantize_weights_mode
+                else:
+                    weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
+            if weights_dtype is None:
+                return None
 
             if shared.opts.device_map == "gpu":
                 quantization_device = devices.device
@@ -337,6 +341,11 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True):
     else:
         weights_dtype = shared.opts.sdnq_quantize_weights_mode
 
+    if weights_dtype is None or weights_dtype == 'none':
+        return model
+    if debug:
+        log.trace(f'Quantization: type=SDNQ op={op} cls={model.__class__} dtype={weights_dtype} mode{shared.opts.diffusers_offload_mode}')
+
     if shared.opts.diffusers_offload_mode in {"none", "model"}:
         quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu
         return_device = devices.device
@@ -397,7 +406,7 @@ def sdnq_quantize_weights(sd_model):
     try:
         t0 = time.time()
         from modules import shared, devices, sd_models
-        log.info(f"Quantization: type=SDNQ dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} matmul={shared.opts.sdnq_use_quantized_matmul}  group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} modules={shared.opts.sdnq_quantize_weights}")
+        log.debug(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_weights} dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} matmul={shared.opts.sdnq_use_quantized_matmul}  group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} dequantize_fp32={shared.opts.sdnq_dequantize_fp32}")
         global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
 
         sd_model = sd_models.apply_function_to_model(sd_model, sdnq_quantize_model, shared.opts.sdnq_quantize_weights, op="sdnq")

modules/model_te.py

@@ -4,7 +4,6 @@ import torch
 import transformers
 from safetensors.torch import load_file
 from modules import shared, devices, files_cache, errors, model_quant
-from installer import install
 
 
 te_dict = {}
@@ -72,27 +71,32 @@ def load_t5(name=None, cache_dir=None):
     elif 'int8' in name.lower():
         from modules.model_quant import create_sdnq_config
         quantization_config = create_sdnq_config(kwargs=None, allow_sdnq=True, module='any', weights_dtype='int8')
-        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+        if quantization_config is not None:
+            t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
 
     elif 'uint4' in name.lower():
         from modules.model_quant import create_sdnq_config
         quantization_config = create_sdnq_config(kwargs=None, allow_sdnq=True, module='any', weights_dtype='uint4')
-        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+        if quantization_config is not None:
+            t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
 
     elif 'qint4' in name.lower():
         model_quant.load_quanto('Load model: type=T5')
         quantization_config = transformers.QuantoConfig(weights='int4')
-        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+        if quantization_config is not None:
+            t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
 
     elif 'qint8' in name.lower():
         model_quant.load_quanto('Load model: type=T5')
         quantization_config = transformers.QuantoConfig(weights='int8')
-        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+        if quantization_config is not None:
+            t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
 
     elif '/' in name:
         shared.log.debug(f'Load model: type=T5 repo={name}')
         quant_config = model_quant.create_config(module='TE')
-        t5 = transformers.T5EncoderModel.from_pretrained(name, cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_config)
+        if quantization_config is not None:
+            t5 = transformers.T5EncoderModel.from_pretrained(name, cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_config)
 
     else:
         t5 = None

modules/sd_models.py

@@ -660,7 +660,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
         from modules import modelstats
         modelstats.analyze()
 
-    shared.log.info(f"Load {op}: time={timer.summary()} native={get_native(sd_model)} memory={memory_stats()}")
+    shared.log.info(f"Load {op}: family={shared.sd_model_type} time={timer.dct()} native={get_native(sd_model)} memory={memory_stats()}")
 
 
 class DiffusersTaskType(Enum):
@@ -1086,7 +1086,8 @@ def clear_caches():
         lora_common.loaded_networks.clear()
         lora_common.previously_loaded_networks.clear()
         lora_load.lora_cache.clear()
-    from modules import prompt_parser_diffusers, memstats
+    from modules import prompt_parser_diffusers, memstats, sd_offload
+    sd_offload.offload_hook_instance = None
     prompt_parser_diffusers.cache.clear()
     memstats.reset_stats()
 

modules/sd_offload.py

@@ -4,6 +4,7 @@ import time
 import inspect
 import torch
 import accelerate.hooks
+import accelerate.utils.modeling
 from installer import log
 from modules import shared, devices, errors, model_quant
 from modules.timer import process as process_timer
@@ -15,6 +16,16 @@ offload_warn = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'c
 offload_post = ['h1']
 offload_hook_instance = None
 balanced_offload_exclude = ['CogView4Pipeline']
+accelerate_dtype_byte_size = None
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    try:
+        if dtype in [torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz]:
+            dtype = accelerate.utils.modeling.CustomDtype.FP8
+    except Exception: # catch since older torch many not have defined dtypes
+        pass
+    return accelerate_dtype_byte_size(dtype)
 
 
 def get_signature(cls):
@@ -58,6 +69,7 @@ def set_accelerate(sd_model):
 
 
 def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False):
+    global accelerate_dtype_byte_size # pylint: disable=global-statement
     t0 = time.time()
     if not shared.native:
         shared.log.warning('Attempting to use offload with backend=original')
@@ -67,6 +79,9 @@ def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False):
         return
     if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate):
         sd_model.has_accelerate = False
+    if accelerate_dtype_byte_size is None:
+        accelerate_dtype_byte_size = accelerate.utils.modeling.dtype_byte_size
+        accelerate.utils.modeling.dtype_byte_size = dtype_byte_size
     if shared.opts.diffusers_offload_mode == "none":
         if shared.sd_model_type in offload_warn or 'video' in shared.sd_model_type:
             shared.log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} type={shared.sd_model.__class__.__name__} large model')
@@ -163,14 +178,18 @@ class OffloadHook(accelerate.hooks.ModelHook):
             max_memory = { device_index: self.gpu, "cpu": self.cpu }
             device_map = getattr(module, "balanced_offload_device_map", None)
             if device_map is None or max_memory != getattr(module, "balanced_offload_max_memory", None):
+                # try:
                 device_map = accelerate.infer_auto_device_map(module, max_memory=max_memory)
+                # except Exception as e:
+                #     shared.log.error(f'Offload: type=balanced module={module.__class__.__name__} {e}')
             offload_dir = getattr(module, "offload_dir", os.path.join(shared.opts.accelerate_offload_path, module.__class__.__name__))
             if devices.backend == "directml":
                 keys = device_map.keys()
                 for v in keys:
                     if isinstance(device_map[v], int):
                         device_map[v] = f"{devices.device.type}:{device_map[v]}" # int implies CUDA or XPU device, but it will break DirectML backend so we add type
-            module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir)
+            if device_map is not None:
+                module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir)
             module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access
             module.balanced_offload_device_map = device_map
             module.balanced_offload_max_memory = max_memory

modules/sd_vae_taesd.py

@@ -52,34 +52,36 @@ def warn_once(msg, variant=None):
 def get_model(model_type = 'decoder', variant = None):
     global prev_cls, prev_type, prev_model # pylint: disable=global-statement
     from modules import shared
-    cls = shared.sd_model_type
-    if cls in {'ldm', 'pixartalpha'}:
-        cls = 'sd'
-    elif cls in {'h1', 'lumina2'}:
-        cls = 'f1'
-    elif cls in {'pixartsigma', 'omnigen'}:
-        cls = 'sdxl'
-    elif cls not in supported:
-        warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant)
+    model_cls = shared.sd_model_type
+    if model_cls is None or model_cls == 'none':
+        return None
+    elif model_cls in {'ldm', 'pixartalpha'}:
+        model_cls = 'sd'
+    elif model_cls in {'h1', 'lumina2'}:
+        model_cls = 'f1'
+    elif model_cls in {'pixartsigma', 'omnigen'}:
+        model_cls = 'sdxl'
+    elif model_cls not in supported:
+        warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant)
     variant = variant or shared.opts.taesd_variant
     folder = os.path.join(paths.models_path, "TAESD")
     os.makedirs(folder, exist_ok=True)
     if variant.startswith('TAE'):
         cfg = TAESD_MODELS[variant]
-        if (cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
+        if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
             return cfg['model']
-        fn = os.path.join(folder, cfg['fn'] + cls + '_' + model_type + '.pth')
+        fn = os.path.join(folder, cfg['fn'] + model_type + '_' + model_cls + '.pth')
         if not os.path.exists(fn):
             uri = cfg['uri']
             if not uri.endswith('.pth'):
-                uri += '/tae' + cls + '_' + model_type + '.pth'
+                uri += '/tae' + model_cls + '_' + model_type + '.pth'
             try:
                 shared.log.info(f'Decode: type="taesd" variant="{variant}": uri="{uri}" fn="{fn}" download')
                 torch.hub.download_url_to_file(uri, fn)
             except Exception as e:
                 warn_once(f'download uri={uri} {e}', variant=variant)
         if os.path.exists(fn):
-            prev_cls = cls
+            prev_cls = model_cls
             prev_type = model_type
             prev_model = variant
             shared.log.debug(f'Decode: type="taesd" variant="{variant}" fn="{fn}" load')
@@ -97,14 +99,14 @@ def get_model(model_type = 'decoder', variant = None):
                 TAESD_MODELS[variant]['model'] = TAESD(decoder_path=fn if model_type=='decoder' else None, encoder_path=fn if model_type=='encoder' else None)
             return TAESD_MODELS[variant]['model']
     elif variant.startswith('Hybrid'):
-        cfg = CQYAN_MODELS[variant].get(cls, None)
-        if (cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
+        cfg = CQYAN_MODELS[variant].get(model_cls, None)
+        if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
             return cfg['model']
         if cfg is None:
-            warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant)
+            warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant)
             return None
         repo = cfg['repo']
-        prev_cls = cls
+        prev_cls = model_cls
         prev_type = model_type
         prev_model = variant
         shared.log.debug(f'Decode: type="taesd" variant="{variant}" id="{repo}" load')
@@ -116,10 +118,12 @@ def get_model(model_type = 'decoder', variant = None):
             from modules.taesd.hybrid_small import AutoencoderSmall
             vae = AutoencoderSmall.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir, torch_dtype=dtype)
         vae = vae.to(devices.device, dtype=dtype)
-        CQYAN_MODELS[variant][cls]['model'] = vae
+        CQYAN_MODELS[variant][model_cls]['model'] = vae
         return vae
+    elif variant is None:
+        warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} variant is none', variant=variant)
     else:
-        warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant)
+        warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant)
     return None
 
 

modules/shared.py

@@ -66,6 +66,7 @@ dir_timestamps = {}
 dir_cache = {}
 max_workers = 8
 default_hfcache_dir = os.environ.get("SD_HFCACHEDIR", None) or os.path.join(os.path.expanduser('~'), '.cache', 'huggingface', 'hub')
+sdnq_quant_modes = ["int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"]
 
 
 class Backend(Enum):
@@ -518,8 +519,8 @@ options_templates.update(options_section(("quantization", "Quantization Settings
     "sdnq_quantize_sep": OptionInfo("<h2>SDNQ: SD.Next Quantization</h2>", "", gr.HTML),
     "sdnq_quantize_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
     "sdnq_quantize_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ["pre", "post"], "visible": native}),
-    "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": ["int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"], "visible": native}),
-    "sdnq_quantize_weights_mode_te": OptionInfo("default", "Quantization type for Text Encoders", gr.Dropdown, {"choices": ["default", "int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"], "visible": native}),
+    "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": sdnq_quant_modes, "visible": native}),
+    "sdnq_quantize_weights_mode_te": OptionInfo("default", "Quantization type for Text Encoders", gr.Dropdown, {"choices": ['default'] + sdnq_quant_modes, "visible": native}),
     "sdnq_quantize_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1, "visible": native}),
     "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize the convolutional layers", gr.Checkbox, {"visible": native}),
     "sdnq_dequantize_compile": OptionInfo(devices.has_triton(), "Dequantize using torch.compile", gr.Checkbox, {"visible": native}),

modules/timer.py

@@ -44,8 +44,7 @@ class Timer:
     def summary(self, min_time=default_min_time, total=True):
         if self.profile:
             min_time = -1
-        if self.total <= 0:
-            self.total = sum(self.records.values())
+        self.total = sum(self.records.values())
         res = f"total={self.total:.2f} " if total else ''
         additions = [x for x in self.records.items() if x[1] >= min_time]
         additions = sorted(additions, key=lambda x: x[1], reverse=True)
@@ -60,6 +59,8 @@ class Timer:
     def dct(self, min_time=default_min_time):
         if self.profile:
             res = {k: round(v, 4) for k, v in self.records.items()}
+        self.total = sum(self.records.values())
+        self.records['total'] = self.total
         res = {k: round(v, 2) for k, v in self.records.items() if v >= min_time}
         res = {k: v for k, v in sorted(res.items(), key=lambda x: x[1], reverse=True)} # noqa: C416 # pylint: disable=unnecessary-comprehension
         return res

scripts/xyz_grid_classes.py

@@ -1,4 +1,4 @@
-from scripts.xyz_grid_shared import apply_field, apply_task_arg, apply_task_args, apply_setting, apply_prompt_primary, apply_prompt_refine, apply_prompt_detailer, apply_prompt_all, apply_order, apply_sampler, apply_hr_sampler_name, confirm_samplers, apply_checkpoint, apply_refiner, apply_unet, apply_dict, apply_clip_skip, apply_vae, list_lora, apply_lora, apply_lora_strength, apply_te, apply_styles, apply_upscaler, apply_context, apply_detailer, apply_override, apply_processing, apply_options, apply_seed, format_value_add_label, format_bool, format_value, format_value_join_list, do_nothing, format_nothing, str_permutations # pylint: disable=no-name-in-module, unused-import
+from scripts.xyz_grid_shared import apply_field, apply_task_arg, apply_task_args, apply_setting, apply_prompt_primary, apply_prompt_refine, apply_prompt_detailer, apply_prompt_all, apply_order, apply_sampler, apply_hr_sampler_name, confirm_samplers, apply_checkpoint, apply_refiner, apply_unet, apply_dict, apply_clip_skip, apply_vae, list_lora, apply_lora, apply_lora_strength, apply_te, apply_styles, apply_upscaler, apply_context, apply_detailer, apply_override, apply_processing, apply_options, apply_seed, apply_sdnq_quant,  apply_sdnq_quant_te, format_value_add_label, format_bool, format_value, format_value_join_list, do_nothing, format_nothing, str_permutations # pylint: disable=no-name-in-module, unused-import
 from modules import shared, shared_items, sd_samplers, ipadapter, sd_models, sd_vae, sd_unet
 
 
@@ -58,6 +58,7 @@ class SharedSettingsStackHelper(object):
     extra_networks_default_multiplier = None
     disable_apply_metadata = None
     disable_apply_params = None
+    sdnq_quant_mode = None
 
     def __enter__(self):
         # Save overridden settings so they can be restored later
@@ -89,6 +90,8 @@ class SharedSettingsStackHelper(object):
         self.teacache_thresh = shared.opts.teacache_thresh
         self.disable_apply_metadata = shared.opts.disable_apply_metadata
         self.disable_apply_params = shared.opts.disable_apply_params
+        self.sdnq_quant_mode = shared.opts.sdnq_quantize_weights_mode
+
         shared.opts.data["disable_apply_metadata"] = []
         shared.opts.data["disable_apply_params"] = ''
 
@@ -135,6 +138,9 @@ class SharedSettingsStackHelper(object):
         if self.sd_unet != shared.opts.sd_unet:
             shared.opts.data["sd_unet"] = self.sd_unet
             sd_unet.load_unet(shared.sd_model)
+        if self.sdnq_quant_mode != shared.opts.sdnq_quantize_weights_mode:
+            shared.opts.data["sdnq_quantize_weights_mode"] = self.sdnq_quant_mode
+            sd_models.reload_model_weights(op='model')
 
 
 axis_options = [
@@ -193,6 +199,8 @@ axis_options = [
     AxisOption("[Postprocess] Context", str, apply_context, choices=lambda: ["Add with forward", "Remove with forward", "Add with backward", "Remove with backward"]),
     AxisOption("[Postprocess] Detailer", str, apply_detailer, fmt=format_value_add_label),
     AxisOption("[Postprocess] Detailer strength", str, apply_field("detailer_strength")),
+    AxisOption("[Quant] SDNQ quant mode", str, apply_sdnq_quant, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared.sdnq_quant_modes)),
+    AxisOption("[Quant] SDNQ quant mode TE", str, apply_sdnq_quant_te, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared.sdnq_quant_modes)),
     AxisOption("[HDR] Mode", int, apply_field("hdr_mode")),
     AxisOption("[HDR] Brightness", float, apply_field("hdr_brightness")),
     AxisOption("[HDR] Color", float, apply_field("hdr_color")),

scripts/xyz_grid_shared.py

@@ -147,6 +147,18 @@ def confirm_samplers(p, xs):
             shared.log.warning(f"XYZ grid: unknown sampler: {x}")
 
 
+def apply_sdnq_quant(p, x, xs):
+    shared.opts.sdnq_quantize_weights_mode = x
+    sd_models.unload_model_weights(op='model') # reload will happen on-demand
+    shared.log.debug(f'XYZ grid apply sdnq quant: mode="{x}"')
+
+
+def apply_sdnq_quant_te(p, x, xs):
+    shared.opts.sdnq_quantize_weights_mode_te = x
+    sd_models.unload_model_weights(op='model') # reload will happen on-demand
+    shared.log.debug(f'XYZ grid apply sdnq quant te: mode="{x}"')
+
+
 def apply_checkpoint(p, x, xs):
     if x == shared.opts.sd_model_checkpoint:
         return