add nunchaku prototype

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

CHANGELOG.md

@@ -2,7 +2,12 @@
 
 ## Update for 2025-04-15
 
-- **HiDream** optimized offloading, now works in 12GB VRAM / 26GB RAM
+- [Nunchaku](https://github.com/mit-han-lab/nunchaku) inference engine with custom **SVDQuant** 4-bit execution  
+  highly experimental and with limited support, but when it works, its magic: **Flux.1 at 5.90 it/s** *(not sec/it)*!  
+  see [Nunchaku Wiki](https://github.com/vladmandic/sdnext/Nunchaku) for details  
+- **HiDream** optimized offloading and prompt-encode caching  
+  it now works in 12GB VRAM / 26GB RAM!  
+- fix: NNCF for TE-only quant  
 
 ## Update for 2025-04-14
 

TODO.md

@@ -39,3 +39,6 @@ N/A
 - modules/lora/lora_extract.py:185:9: W0511: TODO: lora support pre-quantized flux
 - control: support scripts via api
 - modernui: monkey-patch for missing tabs.select event
+- nunchaku: cache-dir for transformer and t5 loader
+- nunchaku: batch support
+- nunchaku: LoRA support

modules/mit_nunchaku.py

@@ -0,0 +1,63 @@
+# MIT-Han-Lab Nunchaku: <https://github.com/mit-han-lab/nunchaku>
+# TODO nunchaku: cache-dir for transformer and t5 loader
+# TODO nunchaku: batch support
+# TODO nunchaku: LoRA support
+
+from installer import log, pip
+from modules import devices
+
+
+ver = '0.2.0'
+ok = False
+
+
+def check():
+    global ok # pylint: disable=global-statement
+    if ok:
+        return True
+    try:
+        import nunchaku
+        import nunchaku.utils
+        log.info(f'Nunchaku: path={nunchaku.__path__} precision={nunchaku.utils.get_precision()}')
+        ok = True
+        return True
+    except Exception as e:
+        log.error(f'Nunchaku: {e}')
+        ok = False
+        return False
+
+
+def install_nunchaku():
+    if devices.backend is None:
+        return # too early
+    if not check():
+        import sys
+        import platform
+        import importlib
+        import pkg_resources
+        import torch
+        python_ver = f'{sys.version_info.major}{sys.version_info.minor}'
+        if python_ver not in ['311', '312', '313']:
+            log.error(f'Nunchaku: python={sys.version_info} unsupported')
+            return
+        arch = platform.system().lower()
+        if arch not in ['linux', 'windows']:
+            log.error(f'Nunchaku: platform={arch} unsupported')
+            return
+        if devices.backend not in ['cuda']:
+            log.error(f'Nunchaku: backend={devices.backend} unsupported')
+            return
+        torch_ver = torch.__version__[:3]
+        if torch_ver not in ['2.5', '2.6', '2.7', '2.8']:
+            log.error(f'Nunchaku: torch={torch.__version__} unsupported')
+        suffix = 'x86_64' if arch == 'linux' else 'win_amd64'
+        url = f'https://huggingface.co/mit-han-lab/nunchaku/resolve/main/nunchaku-{ver}'
+        url += f'+torch{torch_ver}-cp{python_ver}-cp{python_ver}-{arch}_{suffix}.whl'
+        cmd = f'install --upgrade {url}'
+        # pip install https://huggingface.co/mit-han-lab/nunchaku/resolve/main/nunchaku-0.2.0+torch2.6-cp311-cp311-linux_x86_64.whl
+        log.debug(f'Nunchaku: url={url}')
+        pip(cmd, ignore=False, uv=False)
+        importlib.reload(pkg_resources)
+    if not check():
+        log.error('Nunchaku: install failed')
+        return False

modules/model_flux.py

@@ -109,11 +109,25 @@ def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unu
 
 def load_quants(kwargs, repo_id, cache_dir, allow_quant):
     try:
-        if 'transformer' not in kwargs and (('Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization or 'Model' in shared.opts.quanto_quantization) or ('Transformer' in shared.opts.bnb_quantization or 'Transformer' in shared.opts.torchao_quantization or 'Transformer' in shared.opts.quanto_quantization)):
+        if 'transformer' not in kwargs and model_quant.check_nunchaku('Transformer'):
+            import nunchaku
+            nunchaku_precision = nunchaku.utils.get_precision()
+            nunchaku_repo = f"mit-han-lab/svdq-{nunchaku_precision}-flux.1-dev" if 'dev' in repo_id else f"mit-han-lab/svdq-{nunchaku_precision}-flux.1-schnell"
+            shared.log.debug(f'Load module: quant=Nunchaku module=transformer repo="{nunchaku_repo}" precision={nunchaku_precision} attention={shared.opts.nunchaku_attention}')
+            kwargs['transformer'] = nunchaku.NunchakuFluxTransformer2dModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype)
+            if shared.opts.nunchaku_attention:
+                kwargs['transformer'].set_attention_impl("nunchaku-fp16")
+        elif 'transformer' not in kwargs and model_quant.check_quant('Transformer'):
             quant_args = model_quant.create_config(allow=allow_quant, module='Transformer')
             if quant_args:
                 kwargs['transformer'] = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
-        if 'text_encoder_2' not in kwargs and ('TE' in shared.opts.bnb_quantization or 'TE' in shared.opts.torchao_quantization or 'TE' in shared.opts.quanto_quantization):
+        if 'text_encoder_2' not in kwargs and model_quant.check_nunchaku('TE'):
+            import nunchaku
+            nunchaku_precision = nunchaku.utils.get_precision()
+            nunchaku_repo = 'mit-han-lab/svdq-flux.1-t5'
+            shared.log.debug(f'Load module: quant=Nunchaku module=t5 repo="{nunchaku_repo}" precision={nunchaku_precision}')
+            kwargs['text_encoder_2'] = nunchaku.NunchakuT5EncoderModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype)
+        elif 'text_encoder_2' not in kwargs and model_quant.check_quant('TE'):
             quant_args = model_quant.create_config(allow=allow_quant, module='TE')
             if quant_args:
                 kwargs['text_encoder_2'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
@@ -198,7 +212,7 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
     if shared.opts.teacache_enabled:
         from modules import teacache
         shared.log.debug(f'Transformers cache: type=teacache patch=forward cls={diffusers.FluxTransformer2DModel.__name__}')
-        diffusers.FluxTransformer2DModel.forward = teacache.teacache_flux_forward
+        diffusers.FluxTransformer2DModel.forward = teacache.teacache_flux_forward # patch must be done before transformer is loaded
 
     # load overrides if any
     if shared.opts.sd_unet != 'Default':
@@ -310,6 +324,10 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
     else:
         pipe = cls.from_pretrained(repo_id, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config)
 
+    if shared.opts.teacache_enabled and model_quant.check_nunchaku('Transformer'):
+        from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe
+        apply_cache_on_pipe(pipe, residual_diff_threshold=0.12)
+
     # release memory
     transformer = None
     text_encoder_1 = None

modules/model_quant.py

@@ -100,6 +100,26 @@ def create_quanto_config(kwargs = None, allow_quanto: bool = True, module: str =
     return kwargs
 
 
+def check_quant(module: str = ''):
+    from modules import shared
+    if 'Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization or 'Model' in shared.opts.quanto_quantization:
+        return True
+    if module in shared.opts.bnb_quantization or module in shared.opts.torchao_quantization or module in shared.opts.quanto_quantization:
+        return True
+    return False
+
+
+def check_nunchaku(module: str = ''):
+    from modules import shared
+    if 'Model' not in shared.opts.nunchaku_quantization and module not in shared.opts.nunchaku_quantization:
+        return False
+    from modules import mit_nunchaku
+    mit_nunchaku.install_nunchaku()
+    if not mit_nunchaku.ok:
+        return False
+    return True
+
+
 def create_config(kwargs = None, allow: bool = True, module: str = 'Model'):
     if kwargs is None:
         kwargs = {}

modules/para_attention.py

@@ -12,9 +12,13 @@ def apply_first_block_cache():
     from installer import install
     install('para_attn')
     try:
-        from para_attn.first_block_cache import diffusers_adapters
-        diffusers_adapters.apply_cache_on_pipe(shared.sd_model, residual_diff_threshold=shared.opts.para_diff_threshold)
-        shared.log.info(f'Transformers cache: type=paraattn rdt={shared.opts.para_diff_threshold} cls={shared.sd_model.__class__.__name__}')
+        if 'Nunchaku' in shared.sd_model.transformer.__class__.__name__:
+            from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe
+            shared.log.info(f'Transformers cache: type=nunchaku rdt={shared.opts.para_diff_threshold} cls={shared.sd_model.transformer.__class__.__name__}')
+        else:
+            from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
+            shared.log.info(f'Transformers cache: type=paraattn rdt={shared.opts.para_diff_threshold} cls={shared.sd_model.transformer.__class__.__name__}')
+        apply_cache_on_pipe(shared.sd_model, residual_diff_threshold=shared.opts.para_diff_threshold)
     except Exception as e:
         shared.log.error(f'Transformers cache: type=paraattn {e}')
         return

modules/prompt_parser_diffusers.py

@@ -191,32 +191,36 @@ class PromptEmbedder:
     def __call__(self, key, step=0):
         batch = getattr(self, key)
         res = []
-        if isinstance(batch[0][0], list) and len(batch[0][0]) == 2 and isinstance(batch[0][0][1], torch.Tensor) and batch[0][0][1].shape[0] == 32:
-            # hidream uses a list of t5 + llama prompt embeds: [t5_embeds, llama_embeds]
-            # t5_embeds shape: [batch_size, seq_len, dim]
-            # llama_embeds shape: [number_of_hidden_states, batch_size, seq_len, dim]
-            res2 = []
-            for i in range(self.batchsize):
-                if len(batch[i]) == 0:  # if asking for a null key, ie pooled on SD1.5
-                    return None
-                try:
-                    res.append(batch[i][step][0])
-                    res2.append(batch[i][step][1])
-                except IndexError:
-                    # if not scheduled, return default
-                    res.append(batch[i][0][0])
-                    res2.append(batch[i][0][1])
-            res = [torch.cat(res, dim=0), torch.cat(res2, dim=1)]
-            return res
-        else:
-            for i in range(self.batchsize):
-                if len(batch[i]) == 0:  # if asking for a null key, ie pooled on SD1.5
-                    return None
-                try:
-                    res.append(batch[i][step])
-                except IndexError:
-                    res.append(batch[i][0])  # if not scheduled, return default
-            return torch.cat(res)
+        try:
+            if isinstance(batch[0][0], list) and len(batch[0][0]) == 2 and isinstance(batch[0][0][1], torch.Tensor) and batch[0][0][1].shape[0] == 32:
+                # hidream uses a list of t5 + llama prompt embeds: [t5_embeds, llama_embeds]
+                # t5_embeds shape: [batch_size, seq_len, dim]
+                # llama_embeds shape: [number_of_hidden_states, batch_size, seq_len, dim]
+                res2 = []
+                for i in range(self.batchsize):
+                    if len(batch[i]) == 0:  # if asking for a null key, ie pooled on SD1.5
+                        return None
+                    try:
+                        res.append(batch[i][step][0])
+                        res2.append(batch[i][step][1])
+                    except IndexError:
+                        # if not scheduled, return default
+                        res.append(batch[i][0][0])
+                        res2.append(batch[i][0][1])
+                res = [torch.cat(res, dim=0), torch.cat(res2, dim=1)]
+                return res
+            else:
+                for i in range(self.batchsize):
+                    if len(batch[i]) == 0:  # if asking for a null key, ie pooled on SD1.5
+                        return None
+                    try:
+                        res.append(batch[i][step])
+                    except IndexError:
+                        res.append(batch[i][0])  # if not scheduled, return default
+                return torch.cat(res)
+        except Exception:
+            pass
+        return None
 
 
 def compel_hijack(self, token_ids: torch.Tensor, attention_mask: typing.Optional[torch.Tensor] = None) -> torch.Tensor:

modules/shared.py

@@ -548,6 +548,10 @@ options_templates.update(options_section(('quantization', "Quantization Settings
     "layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "TE"], "visible": native}),
     "layerwise_quantization_storage": OptionInfo("float8_e4m3fn", "Layerwise casting storage", gr.Dropdown, {"choices": ["float8_e4m3fn", "float8_e5m2"], "visible": native}),
     "layerwise_quantization_nonblocking": OptionInfo(False, "Layerwise non-blocking operations", gr.Checkbox, {"visible": native}),
+
+    "nunchaku_sep": OptionInfo("<h2>Nunchaku Engine</h2>", "", gr.HTML),
+    "nunchaku_quantization": OptionInfo([], "SVDQuant enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
+    "nunchaku_attention": OptionInfo(False, "Nunchaku attention", gr.Checkbox, {"visible": native}),
 }))
 
 options_templates.update(options_section(('advanced', "Pipeline Modifiers"), {