From b86e357b8eed073d78e714e1e704d8341d001919 Mon Sep 17 00:00:00 2001 From: vladmandic Date: Sat, 28 Mar 2026 09:49:05 +0100 Subject: [PATCH] update nunchaku code Signed-off-by: vladmandic --- CHANGELOG.md | 1 + extensions-builtin/sdnext-modernui | 2 +- installer.py | 4 ++-- modules/mit_nunchaku.py | 16 +++++++++++++--- pipelines/model_z_image.py | 6 ++++-- 5 files changed, 21 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 051361715..726aa1a1b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -129,6 +129,7 @@ Just how big? Some stats: *~530 commits over 880 files* - refactor move `rebmg` to core instead of extensions - remove face restoration - unified command line parsing + - reinstall `nunchaku` with `--reinstall` flag - use explicit icon image references in `gallery`, thanks @awsr - launch use threads to async execute non-critical tasks - switch from deprecated `pkg_resources` to `importlib` diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui index 488ab401c..c7af727f3 160000 --- a/extensions-builtin/sdnext-modernui +++ b/extensions-builtin/sdnext-modernui @@ -1 +1 @@ -Subproject commit 488ab401cfaae83da94821c3f92ba718177dc106 +Subproject commit c7af727f31758c9fc96cf0429bcf3608858a15e8 diff --git a/installer.py b/installer.py index 0b7c491ad..77193d0de 100644 --- a/installer.py +++ b/installer.py @@ -244,13 +244,13 @@ def cleanup_broken_packages(): pass -def pip(arg: str, ignore: bool = False, quiet: bool = True, uv = True): +def pip(arg: str, ignore: bool = False, quiet: bool = True, uv = True) -> tuple[subprocess.CompletedProcess, str]: t_start = time.time() originalArg = arg arg = arg.replace('>=', '==') if opts.get('offline_mode', False): log.warning('Offline mode enabled') - return 'offline' + return None, 'offline' package = arg.replace("install", "").replace("--upgrade", "").replace("--no-deps", "").replace("--force-reinstall", "").replace(" ", " ").strip() uv = uv and args.uv and not package.startswith('git+') pipCmd = "uv pip" if uv else "pip" diff --git a/modules/mit_nunchaku.py b/modules/mit_nunchaku.py index f9501e36c..ea616fe5b 100644 --- a/modules/mit_nunchaku.py +++ b/modules/mit_nunchaku.py @@ -7,8 +7,10 @@ from modules.logger import log ok = False -def check(): +def check(force=False): global ok # pylint: disable=global-statement + if force: + return False if ok: return True try: @@ -26,10 +28,12 @@ def check(): def install_nunchaku(force=False): if not force: - from modules import devices + from modules import devices, shared if devices.backend is None: return False # too early - if not check(): + if shared.cmd_opts.reinstall: + force = True + if not check(force): import os import sys import platform @@ -68,6 +72,9 @@ def install_nunchaku(force=False): url = f'https://github.com/nunchaku-ai/nunchaku/releases/download/v{v}/' fn = f'nunchaku-{v}+{cuda_ver}torch{torch_ver}-cp{python_ver}-cp{python_ver}-{suffix}.whl' result, _output = pip(f'install --upgrade {url+fn}', uv=False, ignore=True, quiet=True) + if (result is None) or (_output == 'offline'): + log.error(f'Nunchaku: install url="{url+fn}" offline mode') + return False if force: log.debug(f'Nunchaku: url="{fn}" code={result.returncode} stdout={result.stdout} stderr={result.stderr} output={_output}') if result.returncode == 0: @@ -75,6 +82,9 @@ def install_nunchaku(force=False): return True fn = f'nunchaku-{v}+torch{torch_ver}-cp{python_ver}-cp{python_ver}-{suffix}.whl' result, _output = pip(f'install --upgrade {url+fn}', uv=False, ignore=True, quiet=True) + if (result is None) or (_output == 'offline'): + log.error(f'Nunchaku: install url="{url+fn}" offline mode') + return False if force: log.debug(f'Nunchaku: url="{fn}" code={result.returncode} stdout={result.stdout} stderr={result.stderr} output={_output}') if result.returncode == 0: diff --git a/pipelines/model_z_image.py b/pipelines/model_z_image.py index e145356e6..bb12a5608 100644 --- a/pipelines/model_z_image.py +++ b/pipelines/model_z_image.py @@ -7,11 +7,13 @@ from pipelines import generic def load_nunchaku(): import nunchaku + if not hasattr(nunchaku, 'NunchakuZImageTransformer2DModel'): # not present in older versions of nunchaku + return None nunchaku_precision = nunchaku.utils.get_precision() nunchaku_rank = 128 nunchaku_repo = f"nunchaku-ai/nunchaku-z-image-turbo/svdq-{nunchaku_precision}_r{nunchaku_rank}-z-image-turbo.safetensors" log.debug(f'Load module: quant=Nunchaku module=transformer repo="{nunchaku_repo}" attention={shared.opts.nunchaku_attention}') - transformer = nunchaku.NunchakuZImageTransformer2DModel.from_pretrained( + transformer = nunchaku.NunchakuZImageTransformer2DModel.from_pretrained( # pylint: disable=no-member nunchaku_repo, torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir, @@ -30,7 +32,7 @@ def load_z_image(checkpoint_info, diffusers_load_config=None): if model_quant.check_nunchaku('Model'): # only available model transformer = load_nunchaku() - else: + if transformer is None: transformer = generic.load_transformer(repo_id, cls_name=diffusers.ZImageTransformer2DModel, load_config=diffusers_load_config) text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3ForCausalLM, load_config=diffusers_load_config)