diff --git a/CHANGELOG.md b/CHANGELOG.md index e2c4fe086..3b44c0ab3 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/TODO.md b/TODO.md index fb1aec1cf..bf27e236f 100644 --- a/TODO.md +++ b/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 diff --git a/modules/mit_nunchaku.py b/modules/mit_nunchaku.py new file mode 100644 index 000000000..395f3099c --- /dev/null +++ b/modules/mit_nunchaku.py @@ -0,0 +1,63 @@ +# 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 diff --git a/modules/model_flux.py b/modules/model_flux.py index 5c0395cf9..25de64c80 100644 --- a/modules/model_flux.py +++ b/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 diff --git a/modules/model_quant.py b/modules/model_quant.py index a32a8e353..99d5d6c9b 100644 --- a/modules/model_quant.py +++ b/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 = {} diff --git a/modules/para_attention.py b/modules/para_attention.py index f5c6e8635..8c00c303d 100644 --- a/modules/para_attention.py +++ b/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 diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py index 5267147c0..14c0ed219 100644 --- a/modules/prompt_parser_diffusers.py +++ b/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: diff --git a/modules/shared.py b/modules/shared.py index 8c51ac40d..601e5ebfa 100644 --- a/modules/shared.py +++ b/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("

Nunchaku Engine

", "", 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"), { diff --git a/wiki b/wiki index 40ac3ec88..a985acf8c 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 40ac3ec884aba0146507eb2e3217ae1fef399170 +Subproject commit a985acf8ca4f8e20c7438f749b4074d37c9df949