Merge pull request #3875 from vladmandic/dev

merge dev

CHANGELOG.md

@@ -1,22 +1,34 @@
 # Change Log for SD.Next
 
-## Update for 2025-04-14
+## Update for 2025-04-16
 
-- [CFG-Zero](https://github.com/WeichenFan/CFG-Zero-star) new guidance method optimized for flow-matching models  
-  implemented for **FLUX.1, HiDream-I1, SD3.x, CogView4, HunyuanVideo, WanAI**  
-  enable and configure in *settings -> pipeline modifiers -> cfg zero*  
-  experiment with CFGZero support in XYZ-grid  
-- add **UniPC FlowMatch** scheduler  
-- **HiDream** add HF gated access auth check  
-- clenup **CogView3** and **CogView4** model loader  
-- add explicit offload after encode prompt  
-  configure in *settings -> text encoder -> offload*  
-- networks: set which networks to skip when scanning civitai  
-  in *settings -> networks -> network scan*  
-  comma-separate list of regex patterns to skip  
-- ui display reference models with subdued color  
-- xyz grid support bool  
-- fix debug logging  
+- **Features**
+  - [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/wiki/Nunchaku) for details  
+  - [CFG-Zero](https://github.com/WeichenFan/CFG-Zero-star) new guidance method optimized for flow-matching models  
+    implemented for **FLUX.1, HiDream-I1, SD3.x, CogView4, HunyuanVideo, WanAI**  
+    enable and configure in *settings -> pipeline modifiers -> cfg zero*  
+    experiment with CFGZero support in XYZ-grid  
+- **Optimizations**
+  - **HiDream** optimized offloading and prompt-encode caching  
+    it now works in 12GB VRAM / 26GB RAM!  
+  - **CogView3** and **CogView4** model loader optimizations  
+  - add explicit offload after encode prompt  
+    configure in *settings -> text encoder -> offload*  
+- **Other**
+  - **HiDream** add HF gated access auth check  
+  - add **UniPC FlowMatch** scheduler  
+  - add **LCM FlowMatch** scheduler  
+  - networks: set which networks to skip when scanning civitai  
+    in *settings -> networks -> network scan*  
+    comma-separate list of regex patterns to skip  
+  - ui display reference models with subdued color  
+  - xyz grid support bool  
+- **Fixes**
+  - NNCF with TE-only quant  
+  - **HunyuanVideo-I2V** with latest transformers  
+  - trace logging  
 
 ## Update for 2025-04-12
 

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

installer.py

@@ -538,7 +538,7 @@ def check_diffusers():
     t_start = time.time()
     if args.skip_all or args.skip_git or args.experimental:
         return
-    sha = 'a8f5134c113da402a93580ef7a021557e816c98d' # diffusers commit hash
+    sha = 'ce1063acfa0cbc2168a7e9dddd4282ab8013b810' # diffusers commit hash
     pkg = pkg_resources.working_set.by_key.get('diffusers', None)
     minor = int(pkg.version.split('.')[1] if pkg is not None else 0)
     cur = opts.get('diffusers_version', '') if minor > 0 else ''

modules/ggml/__init__.py

@@ -47,7 +47,8 @@ def load_gguf_state_dict(path: str, compute_dtype: torch.dtype) -> dict:
 
 def load_gguf(path, cls, compute_dtype: torch.dtype):
     _gguf = install_gguf()
-    module = cls.from_single_file(
+    loader = cls.from_single_file if hasattr(cls, 'from_single_file') else cls.from_pretrained
+    module = loader(
         path,
         quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype=compute_dtype),
         torch_dtype=compute_dtype,

modules/lora/lora_overrides.py

@@ -28,6 +28,7 @@ force_diffusers = [ # forced always
 force_models = [ # forced always
     # 'sd3',
     'sc',
+    'h1',
     'kandinsky',
     'hunyuandit',
     'auraflow',

modules/mit_nunchaku.py

@@ -0,0 +1,64 @@
+# 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 False # 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 False
+        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 False
+        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
+    return True

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)
@@ -136,7 +150,6 @@ def load_transformer(file_path): # triggered by opts.sd_unet change
     if quant is not None and quant != 'none':
         shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} prequant={quant} dtype={devices.dtype}')
     if 'gguf' in file_path.lower():
-        # _transformer, _text_encoder_2 = load_flux_gguf(file_path)
         from modules import ggml
         _transformer = ggml.load_gguf(file_path, cls=diffusers.FluxTransformer2DModel, compute_dtype=devices.dtype)
         if _transformer is not None:
@@ -198,7 +211,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 +323,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_hidream.py

@@ -1,31 +1,46 @@
+import os
 import transformers
 import diffusers
 from modules import shared, devices, sd_models, model_quant, modelloader, sd_hijack_te
 
 
-def load_hidream(checkpoint_info, diffusers_load_config={}):
-    login = modelloader.hf_login()
-    repo_id = sd_models.path_to_repo(checkpoint_info.name)
-
-    from huggingface_hub import auth_check
-    try:
-        auth_check(shared.opts.model_h1_llama_repo)
-    except Exception as e:
-        shared.log.error(f'Load model: type=HiDream te4="{shared.opts.model_h1_llama_repo}" login={login} {e}')
-        return False
-
+def load_transformer(repo_id, diffusers_load_config={}):
     load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer', device_map=True)
-    shared.log.debug(f'Load model: type=HiDream transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
-    transformer = diffusers.HiDreamImageTransformer2DModel.from_pretrained(
-        repo_id,
-        subfolder="transformer",
-        cache_dir=shared.opts.hfcache_dir,
-        **load_args,
-        **quant_args,
-    )
-    if shared.opts.diffusers_offload_mode != 'none':
-        transformer = transformer.to(devices.cpu)
+    fn = None
 
+    if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default':
+        from modules import sd_unet
+        if shared.opts.sd_unet not in list(sd_unet.unet_dict):
+            shared.log.error(f'Load module: type=Transformer not found: {shared.opts.sd_unet}')
+            return None
+        fn = sd_unet.unet_dict[shared.opts.sd_unet] if os.path.exists(sd_unet.unet_dict[shared.opts.sd_unet]) else None
+
+    if fn is not None and 'gguf' in fn.lower():
+        shared.log.error('Load model: type=HiDream format="gguf" unsupported')
+        transformer = None
+        # from modules import ggml
+        # transformer = ggml.load_gguf(fn, cls=diffusers.HiDreamImageTransformer2DModel, compute_dtype=devices.dtype)
+    elif fn is not None and 'safetensors' in fn.lower():
+        shared.log.debug(f'Load model: type=HiDream transformer="{repo_id}" quant="{model_quant.get_quant(repo_id)}" args={load_args}')
+        transformer = diffusers.HiDreamImageTransformer2DModel.from_single_file(fn, cache_dir=shared.opts.hfcache_dir, **load_args)
+    # elif model_quant.check_nunchaku('Transformer'):
+    #     shared.log.error(f'Load model: type=HiDream transformer="{repo_id}" quant="Nunchaku" unsupported')
+    #     transformer = None
+    else:
+        shared.log.debug(f'Load model: type=HiDream transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
+        transformer = diffusers.HiDreamImageTransformer2DModel.from_pretrained(
+            repo_id,
+            subfolder="transformer",
+            cache_dir=shared.opts.hfcache_dir,
+            **load_args,
+            **quant_args,
+        )
+    if shared.opts.diffusers_offload_mode != 'none' and transformer is not None:
+        sd_models.move_model(transformer, devices.cpu)
+    return transformer
+
+
+def load_text_encoders(repo_id, diffusers_load_config={}):
     load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
     shared.log.debug(f'Load model: type=HiDream te3="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
     text_encoder_3 = transformers.T5EncoderModel.from_pretrained(
@@ -35,8 +50,8 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
         **load_args,
         **quant_args,
     )
-    if shared.opts.diffusers_offload_mode != 'none':
-        text_encoder_3 = text_encoder_3.to(devices.cpu)
+    if shared.opts.diffusers_offload_mode != 'none' and text_encoder_3 is not None:
+        sd_models.move_model(text_encoder_3, devices.cpu)
 
     load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='LLM', device_map=True)
     shared.log.debug(f'Load model: type=HiDream te4="{shared.opts.model_h1_llama_repo}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
@@ -54,21 +69,42 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
         cache_dir=shared.opts.hfcache_dir,
         **load_args,
     )
-    if shared.opts.diffusers_offload_mode != 'none':
-        text_encoder_4 = text_encoder_4.to(devices.cpu)
+    if shared.opts.diffusers_offload_mode != 'none' and text_encoder_4 is not None:
+        sd_models.move_model(text_encoder_4, devices.cpu)
+    return text_encoder_3, text_encoder_4, tokenizer_4
+
+
+def load_hidream(checkpoint_info, diffusers_load_config={}):
+    login = modelloader.hf_login()
+    repo_id = sd_models.path_to_repo(checkpoint_info.name)
+
+    from huggingface_hub import auth_check
+    try:
+        auth_check(shared.opts.model_h1_llama_repo)
+    except Exception as e:
+        shared.log.error(f'Load model: type=HiDream te4="{shared.opts.model_h1_llama_repo}" login={login} {e}')
+        return False
+
+    transformer = load_transformer(repo_id, diffusers_load_config)
+    text_encoder_3, text_encoder_4, tokenizer_4 = load_text_encoders(repo_id, diffusers_load_config)
 
     load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
     shared.log.debug(f'Load model: type=HiDream model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
+
     pipe = diffusers.HiDreamImagePipeline.from_pretrained(
         repo_id,
+        transformer=transformer,
         text_encoder_3=text_encoder_3,
         text_encoder_4=text_encoder_4,
         tokenizer_4=tokenizer_4,
-        transformer=transformer,
         cache_dir=shared.opts.diffusers_dir,
         **load_args,
     )
     sd_hijack_te.init_hijack(pipe)
+    del text_encoder_3
+    del text_encoder_4
+    del tokenizer_4
+    del transformer
 
     devices.torch_gc()
     return pipe

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/processing_args.py

@@ -143,7 +143,8 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
     if (prompt_attention != 'fixed') and ('Onnx' not in model.__class__.__name__) and ('prompt' not in p.task_args) and (
         'StableDiffusion' in model.__class__.__name__ or
         'StableCascade' in model.__class__.__name__ or
-        'Flux' in model.__class__.__name__
+        'Flux' in model.__class__.__name__ or
+        'HiDreamImage' in model.__class__.__name__
     ):
         try:
             prompt_parser_diffusers.embedder = prompt_parser_diffusers.PromptEmbedder(prompts, negative_prompts, steps, clip_skip, p)
@@ -162,25 +163,31 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
             prompts = [p.replace('|image|', '<|image_1|>') for p in prompts]
         if hasattr(model, 'text_encoder') and hasattr(model, 'tokenizer') and 'prompt_embeds' in possible and prompt_parser_diffusers.embedder is not None:
             args['prompt_embeds'] = prompt_parser_diffusers.embedder('prompt_embeds')
-            if 'StableCascade' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['prompt_embeds_pooled'] = prompt_parser_diffusers.embedder('positive_pooleds').unsqueeze(0)
-            elif 'XL' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
-            elif 'StableDiffusion3' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
-            elif 'Flux' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
+            if prompt_parser_diffusers.embedder is not None:
+                if 'StableCascade' in model.__class__.__name__:
+                    args['prompt_embeds_pooled'] = prompt_parser_diffusers.embedder('positive_pooleds').unsqueeze(0)
+                elif 'XL' in model.__class__.__name__:
+                    args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
+                elif 'StableDiffusion3' in model.__class__.__name__:
+                    args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
+                elif 'Flux' in model.__class__.__name__:
+                    args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
+                elif 'HiDreamImage' in model.__class__.__name__:
+                    args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
         else:
             args['prompt'] = prompts
     if 'negative_prompt' in possible:
         if hasattr(model, 'text_encoder') and hasattr(model, 'tokenizer') and 'negative_prompt_embeds' in possible and prompt_parser_diffusers.embedder is not None:
             args['negative_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_prompt_embeds')
-            if 'StableCascade' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['negative_prompt_embeds_pooled'] = prompt_parser_diffusers.embedder('negative_pooleds').unsqueeze(0)
-            if 'XL' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
-            if 'StableDiffusion3' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None:
-                args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
+            if prompt_parser_diffusers.embedder is not None:
+                if 'StableCascade' in model.__class__.__name__:
+                    args['negative_prompt_embeds_pooled'] = prompt_parser_diffusers.embedder('negative_pooleds').unsqueeze(0)
+                elif 'XL' in model.__class__.__name__:
+                    args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
+                elif 'StableDiffusion3' in model.__class__.__name__:
+                    args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
+                elif 'HiDreamImage' in model.__class__.__name__:
+                    args['negative_pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('negative_pooleds')
         else:
             if 'PixArtSigmaPipeline' in model.__class__.__name__: # pixart-sigma pipeline throws list-of-list for negative prompt
                 args['negative_prompt'] = negative_prompts[0]

modules/prompt_parser_diffusers.py

@@ -26,7 +26,8 @@ def prompt_compatible(pipe = None):
         'StableDiffusion' not in pipe.__class__.__name__ and
         'DemoFusion' not in pipe.__class__.__name__ and
         'StableCascade' not in pipe.__class__.__name__ and
-        'Flux' not in pipe.__class__.__name__
+        'Flux' not in pipe.__class__.__name__ and
+        'HiDreamImage' not in pipe.__class__.__name__
     ):
         shared.log.warning(f"Prompt parser not supported: {pipe.__class__.__name__}")
         return False
@@ -190,14 +191,36 @@ class PromptEmbedder:
     def __call__(self, key, step=0):
         batch = getattr(self, key)
         res = []
-        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:
@@ -457,21 +480,43 @@ def split_prompts(pipe, prompt, SD3 = False):
     else:
         prompt3 = prompt
 
+    if prompt.find("TE4:") != -1:
+        prompt, prompt4 = prompt.split("TE4:")
+    elif prompt2.find("TE4:") != -1:
+        prompt2, prompt4 = prompt2.split("TE4:")
+    elif prompt3.find("TE4:") != -1:
+        prompt3, prompt4 = prompt3.split("TE4:")
+    else:
+        prompt4 = prompt
+
     prompt = prompt.strip()
     prompt2 = " " if prompt2.strip() == "" else prompt2.strip()
     prompt3 = " " if prompt3.strip() == "" else prompt3.strip()
+    prompt4 = " " if prompt4.strip() == "" else prompt4.strip()
 
     if SD3 and prompt3 != " ":
         ps, _ws = get_prompts_with_weights(pipe, prompt3)
         prompt3 = " ".join(ps)
-    return prompt, prompt2, prompt3
+    return prompt, prompt2, prompt3, prompt4
 
 
 def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", clip_skip: int = None):
     device = devices.device
-    SD3 = hasattr(pipe, 'text_encoder_3')
-    prompt, prompt_2, prompt_3 = split_prompts(pipe, prompt, SD3)
-    neg_prompt, neg_prompt_2, neg_prompt_3 = split_prompts(pipe, neg_prompt, SD3)
+    SD3 = bool(hasattr(pipe, 'text_encoder_3') and not hasattr(pipe, 'text_encoder_4'))
+    prompt, prompt_2, prompt_3, prompt_4 = split_prompts(pipe, prompt, SD3)
+    neg_prompt, neg_prompt_2, neg_prompt_3, neg_prompt_4 = split_prompts(pipe, neg_prompt, SD3)
+
+    if "Flux" in pipe.__class__.__name__: # clip is only used for the pooled embeds
+        prompt_embeds, pooled_prompt_embeds, _ = pipe.encode_prompt(prompt=prompt, prompt_2=prompt_2, device=device, num_images_per_prompt=1)
+        return prompt_embeds, pooled_prompt_embeds, None, None # no negative support
+
+    if "HiDreamImage" in pipe.__class__.__name__: # clip is only used for the pooled embeds
+        prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds = pipe.encode_prompt(
+            prompt=prompt, prompt_2=prompt_2, prompt_3=prompt_3, prompt_4=prompt_4,
+            negative_prompt=neg_prompt, negative_prompt_2=neg_prompt_2, negative_prompt_3=neg_prompt_3, negative_prompt_4=neg_prompt_4,
+            device=device, num_images_per_prompt=1,
+        )
+        return prompt_embeds, pooled_prompt_embeds, negative_prompt_embeds, negative_pooled_prompt_embeds
 
     if prompt != prompt_2:
         ps = [get_prompts_with_weights(pipe, p) for p in [prompt, prompt_2]]
@@ -488,10 +533,6 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c
         negatives.pop(0)
         negative_weights.pop(0)
 
-    if "Flux" in pipe.__class__.__name__: # clip is only used for the pooled embeds
-        prompt_embeds, pooled_prompt_embeds, _ = pipe.encode_prompt(prompt=prompt, prompt_2=prompt_2, device=device, num_images_per_prompt=1)
-        return prompt_embeds, pooled_prompt_embeds, None, None # no negative support
-
     embedding_providers = prepare_embedding_providers(pipe, clip_skip)
     empty_embedding_providers = None
     if 'StableCascade' in pipe.__class__.__name__:
@@ -593,8 +634,8 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c
 
 def get_xhinker_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", clip_skip: int = None):
     is_sd3 = hasattr(pipe, 'text_encoder_3')
-    prompt, prompt_2, _prompt_3 = split_prompts(pipe, prompt, is_sd3)
-    neg_prompt, neg_prompt_2, _neg_prompt_3 = split_prompts(pipe, neg_prompt, is_sd3)
+    prompt, prompt_2, _prompt_3, _ = split_prompts(pipe, prompt, is_sd3)
+    neg_prompt, neg_prompt_2, _neg_prompt_3, _ = split_prompts(pipe, neg_prompt, is_sd3)
     try:
         prompt = pipe.maybe_convert_prompt(prompt, pipe.tokenizer)
         neg_prompt = pipe.maybe_convert_prompt(neg_prompt, pipe.tokenizer)

modules/sd_hijack_te.py

@@ -15,6 +15,8 @@ def hijack_encode_prompt(*args, **kwargs):
         res = None
     t1 = time.time()
     timer.process.add('te', t1-t0)
+    if hasattr(shared.sd_model, "maybe_free_model_hooks"):
+        shared.sd_model.maybe_free_model_hooks()
     shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
     return res
 

modules/sd_models.py

@@ -225,6 +225,8 @@ def move_model(model, device=None, force=False):
                 pass # ignore model move if sequential offload is enabled
             elif 'Params4bit' in str(e0) or 'Params8bit' in str(e0):
                 pass # ignore model move if quantization is enabled
+            elif 'already been set to the correct devices' in str(e0):
+                pass # ignore errors on pre-quant models
             else:
                 raise e0
         t1 = time.time()
@@ -600,7 +602,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
         prompt_parser_diffusers.cache.clear()
 
         set_diffuser_options(sd_model, vae, op, offload=False)
-        if 'Model' in shared.opts.nncf_compress_weights and not ('Model' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"):
+        if shared.opts.nncf_compress_weights and not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"):
             sd_model = model_quant.nncf_compress_weights(sd_model) # run this before move model so it can be compressed in CPU
         if shared.opts.optimum_quanto_weights:
             sd_model = model_quant.optimum_quanto_weights(sd_model) # run this before move model so it can be compressed in CPU

modules/sd_models_utils.py

@@ -197,6 +197,15 @@ def apply_function_to_model(sd_model, function, options, op=None):
                         dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16
                     )
             sd_model.text_encoder_3 = function(sd_model.text_encoder_3, op="text_encoder_3", sd_model=sd_model)
+        if hasattr(sd_model, 'text_encoder_4') and hasattr(sd_model.text_encoder_4, 'config'):
+            if op == "nncf" and sd_model.text_encoder_4.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}:
+                from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32
+                for i in range(len(sd_model.text_encoder_4.encoder.block)):
+                    sd_model.text_encoder_4.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense(
+                        sd_model.text_encoder_4.encoder.block[i].layer[1].DenseReluDense,
+                        dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16
+                    )
+            sd_model.text_encoder_4 = function(sd_model.text_encoder_4, op="text_encoder_4", sd_model=sd_model)
         if hasattr(sd_model, 'prior_pipe') and hasattr(sd_model.prior_pipe, 'text_encoder') and hasattr(sd_model.prior_pipe.text_encoder, 'config'):
             sd_model.prior_pipe.text_encoder = function(sd_model.prior_pipe.text_encoder, op="prior_pipe.text_encoder", sd_model=sd_model)
     if "VAE" in options:

modules/sd_offload.py

@@ -9,8 +9,10 @@ from modules import shared, devices, errors, model_quant
 from modules.timer import process as process_timer
 
 
-debug_move = log.trace if os.environ.get('SD_MOVE_DEBUG', None) is not None else lambda *args, **kwargs: None
-should_offload = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'cogview4']
+debug = os.environ.get('SD_MOVE_DEBUG', None) is not None
+debug_move = log.trace if debug else lambda *args, **kwargs: None
+offload_warn = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'cogview4']
+offload_post = ['h1']
 offload_hook_instance = None
 balanced_offload_exclude = ['OmniGenPipeline', 'CogView4Pipeline']
 
@@ -66,7 +68,7 @@ def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False):
     if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate):
         sd_model.has_accelerate = False
     if shared.opts.diffusers_offload_mode == "none":
-        if shared.sd_model_type in should_offload or 'video' in shared.sd_model_type:
+        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')
         else:
             shared.log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}')
@@ -175,6 +177,29 @@ class OffloadHook(accelerate.hooks.ModelHook):
         return args, kwargs
 
     def post_forward(self, module, output):
+        if getattr(module, "offload_post", False) and module.device != devices.cpu:
+            used_gpu, used_ram = devices.torch_gc(fast=True)
+            perc_gpu = used_gpu / shared.gpu_memory
+            try:
+                module_size = self.model_size()
+                prev_gpu = used_gpu
+                offload_now = perc_gpu > shared.opts.diffusers_offload_min_gpu_memory
+                if offload_now:
+                    module = module.to(devices.cpu)
+                    used_gpu -= module_size
+                if debug:
+                    cls = module.__class__.__name__
+                    quant = getattr(module, "quantization_method", None)
+                    debug_move(f'Offload: type=balanced op={"post" if offload_now else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={cls} size={module_size:.3f}')
+            except Exception as e:
+                if 'out of memory' in str(e):
+                    devices.torch_gc(fast=True, force=True, reason='oom')
+                elif 'bitsandbytes' in str(e):
+                    pass
+                else:
+                    shared.log.error(f'Offload: type=balanced op=apply module={module.__name__} {e}')
+                if os.environ.get('SD_MOVE_DEBUG', None):
+                    errors.display(e, f'Offload: type=balanced op=apply module={module.__name__}')
         return output
 
     def detach_hook(self, module):
@@ -247,15 +272,16 @@ def apply_balanced_offload(sd_model=None, exclude=[]):
             perc_gpu = used_gpu / shared.gpu_memory
             try:
                 prev_gpu = used_gpu
-                do_offload = (perc_gpu > shared.opts.diffusers_offload_min_gpu_memory) and (module.device != devices.cpu)
-                if do_offload:
+                offload_now = (perc_gpu > shared.opts.diffusers_offload_min_gpu_memory) and (module.device != devices.cpu)
+                if offload_now:
                     module = module.to(devices.cpu)
                     used_gpu -= module_size
                 cls = module.__class__.__name__
                 quant = getattr(module, "quantization_method", None)
                 if not cached:
                     shared.log.debug(f'Model module={module_name} type={cls} dtype={module.dtype} quant={quant} params={offload_hook_instance.param_map[module_name]:.3f} size={offload_hook_instance.offload_map[module_name]:.3f}')
-                debug_move(f'Offload: type=balanced op={"move" if do_offload else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={cls} size={module_size:.3f}')
+                if debug:
+                    debug_move(f'Offload: type=balanced op={"move" if offload_now else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={cls} size={module_size:.3f}')
             except Exception as e:
                 if 'out of memory' in str(e):
                     devices.torch_gc(fast=True, force=True, reason='oom')
@@ -273,6 +299,7 @@ def apply_balanced_offload(sd_model=None, exclude=[]):
             if device_map and max_memory:
                 module.balanced_offload_device_map = device_map
                 module.balanced_offload_max_memory = max_memory
+            module.offload_post = shared.sd_model_type in [offload_post] and shared.opts.te_hijack and module_name.startswith("text_encoder")
         devices.torch_gc(fast=True, force=True, reason='offload')
 
     apply_balanced_offload_to_module(sd_model)

modules/sd_samplers_diffusers.py

@@ -30,6 +30,7 @@ try:
         HeunDiscreteScheduler,
         FlowMatchHeunDiscreteScheduler,
         LCMScheduler,
+        FlowMatchLCMScheduler,
         PNDMScheduler,
         IPNDMScheduler,
         DDPMScheduler,
@@ -103,12 +104,12 @@ config = {
 
     'Heun': { 'use_beta_sigmas': False, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'timestep_spacing': 'linspace' },
     'Heun FlowMatch': { 'timestep_spacing': "linspace", 'shift': 1 },
+    'LCM FlowMatch': { 'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': "scaled_linear", 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'thresholding': False, 'timestep_spacing': 'linspace' },
 
     'DEIS': { 'solver_order': 2, 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "deis", 'solver_type': "logrho", 'lower_order_final': True, 'timestep_spacing': 'linspace', 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_flow_sigmas': False, 'use_beta_sigmas': False },
     'SA Solver': {'predictor_order': 2, 'corrector_order': 2, 'thresholding': False, 'lower_order_final': True, 'use_karras_sigmas': False, 'use_flow_sigmas': False, 'use_exponential_sigmas': False, 'use_beta_sigmas': False, 'timestep_spacing': 'linspace'},
     'DC Solver': { 'beta_start': 0.0001, 'beta_end': 0.02, 'solver_order': 2, 'prediction_type': "epsilon", 'thresholding': False, 'solver_type': 'bh2', 'lower_order_final': True, 'dc_order': 2, 'disable_corrector': [0] },
     'VDM Solver': { 'clip_sample_range': 2.0, },
-    'LCM': { 'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': "scaled_linear", 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'thresholding': False, 'timestep_spacing': 'linspace' },
     'TCD': { 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'beta_schedule': 'scaled_linear' },
     'TDD': { },
     'PeRFlow': { 'prediction_type': 'ddim_eps' },
@@ -179,6 +180,7 @@ samplers_data_diffusers = [
     SamplerData('CMSI', lambda model: DiffusionSampler('CMSI', CMStochasticIterativeScheduler, model), [], {}),
 
     SamplerData('LCM', lambda model: DiffusionSampler('LCM', LCMScheduler, model), [], {}),
+    SamplerData('LCM FlowMatch', lambda model: DiffusionSampler('LCM FlowMatch', FlowMatchLCMScheduler, model), [], {}),
     SamplerData('TCD', lambda model: DiffusionSampler('TCD', TCDScheduler, model), [], {}),
     SamplerData('TDD', lambda model: DiffusionSampler('TDD', TDDScheduler, model), [], {}),
     SamplerData('PeRFlow', lambda model: DiffusionSampler('PeRFlow', PeRFlowScheduler, model), [], {}),

modules/sd_unet.py

@@ -10,10 +10,10 @@ debug = os.environ.get('SD_LOAD_DEBUG', None) is not None
 
 def load_unet(model):
     global loaded_unet # pylint: disable=global-statement
-    if shared.opts.sd_unet == 'Default':
+    if shared.opts.sd_unet == 'Default' or shared.opts.sd_unet == 'None':
         return
     if shared.opts.sd_unet not in list(unet_dict):
-        shared.log.error(f'UNet model not found: {shared.opts.sd_unet}')
+        shared.log.error(f'Load module: type=UNet not found: {shared.opts.sd_unet}')
         return
     config_file = os.path.splitext(unet_dict[shared.opts.sd_unet])[0] + '.json'
     if os.path.exists(config_file):
@@ -34,7 +34,7 @@ def load_unet(model):
             if prior_text_encoder is not None:
                 model.prior_pipe.text_encoder = None # Prevent OOM
                 model.prior_pipe.text_encoder = prior_text_encoder.to(devices.device, dtype=devices.dtype)
-        elif "Flux" in model.__class__.__name__ or "StableDiffusion3" in model.__class__.__name__:
+        elif "Flux" in model.__class__.__name__ or "StableDiffusion3" in model.__class__.__name__ or "HiDream" in model.__class__.__name__:
             loaded_unet = shared.opts.sd_unet
             sd_models.load_diffuser() # TODO model load: force-reloading entire model as loading transformers only leads to massive memory usage
             """
@@ -51,7 +51,7 @@ def load_unet(model):
             """
         else:
             if not hasattr(model, 'unet') or model.unet is None:
-                shared.log.error('UNet not found in current model')
+                shared.log.error('Load module: type=UNET not found in current model')
                 return
             shared.log.info(f'Load module: type=UNet name="{shared.opts.sd_unet}" file="{unet_dict[shared.opts.sd_unet]}" config="{config_file}"')
             from diffusers import UNet2DConditionModel

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"), {

modules/shared_items.py

@@ -69,7 +69,7 @@ def refresh_vae_list():
 
 def sd_unet_items():
     import modules.sd_unet
-    return ["None"] + list(modules.sd_unet.unet_dict)
+    return ['Default'] + list(modules.sd_unet.unet_dict)
 
 
 def refresh_unet_list():
@@ -79,7 +79,7 @@ def refresh_unet_list():
 
 def sd_te_items():
     import modules.model_te
-    predefined = ['None']
+    predefined = ['Default']
     return predefined + list(modules.model_te.te_dict)
 
 

requirements.txt

@@ -45,14 +45,14 @@ accelerate==1.6.0
 opencv-contrib-python-headless==4.9.0.80
 einops==0.4.1
 gradio==3.43.2
-huggingface_hub==0.30.1
+huggingface_hub==0.30.2
 numexpr==2.8.8
 numpy==1.26.4
 numba==0.59.1
 protobuf==4.25.3
 pytorch_lightning==1.9.4
 tokenizers==0.21.1
-transformers==4.51.1
+transformers==4.51.3
 urllib3==1.26.19
 Pillow==10.4.0
 timm==0.9.16