cache-dit support

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

CHANGELOG.md

@@ -14,6 +14,11 @@
     note that nunchaku optimized and prequantized unet is replacement for base unet, so its only applicable to base models, not any of finetunes  
     *how to use*: enable nunchaku in settings -> quantization and then load either sdxl-base or sdxl-base-turbo reference models  
     *note*: sdxl support for nunchaku is not in released version of `nunchaku==1.0.0`, so you need to build [nunchaku](https://nunchaku.tech/docs/nunchaku/installation/installation.html) from source
+- **Features**
+  - [Cache-DiT](https://github.com/vipshop/cache-dit)  
+    cache-dit is a unified, flexible and training-free cache acceleration framework  
+    compatible with many dit-based models such as FLUX.1, Qwen, HunyuanImage, Wan2.2, Chroma, etc.  
+    enable in *settings -> pipeline modifers -> cache-dit*  
   - [Nunchaku Flux.1 PulID](https://nunchaku.tech/docs/nunchaku/python_api/nunchaku.pipeline.pipeline_flux_pulid.html)  
     automatically enabled if loaded model is FLUX.1 with Nunchaku engine enabled and when PulID script is enabled  
 - **Offloading**
@@ -37,8 +42,6 @@
     accepts optional `page` parameter to search specific networks page  
   - **reference models** additional example images, thanks @liutyi  
   - **video** support for configurable multi-stage models such as WAN-2.2-14B  
-  - [cache-dit](https://github.com/vipshop/cache-dit)  
-    *experimental support*, enable in *settings -> pipeline modifers -> cache-dit*  
 - **Fixes**
   - framepack: add explicit hf-login before framepack load  
   - benchmark: remove forced sampler from system info benchmark  

modules/cachedit.py

@@ -1,3 +1,4 @@
+import os
 from installer import install
 from modules import shared
 
@@ -6,6 +7,7 @@ def apply_cache_dit(pipe):
     if not shared.opts.cache_dit_enabled:
         return
     install('git+https://github.com/vipshop/cache-dit', 'cache_dit')
+    os.environ.setdefault("CACHE_DIT_LOG_LEVEL", "error")
     try:
         import cache_dit
     except Exception as e:
@@ -21,7 +23,14 @@ def apply_cache_dit(pipe):
         unapply_cache_dir(pipe)
 
     config_args = {}
-
+    if shared.opts.cache_dit_fcompute >= 0:
+        config_args['Fn_compute_blocks'] = int(shared.opts.cache_dit_fcompute)
+    if shared.opts.cache_dit_bcompute >= 0:
+        config_args['Bn_compute_blocks'] = int(shared.opts.cache_dit_bcompute)
+    if shared.opts.cache_dit_threshold >= 0:
+        config_args['residual_diff_threshold'] = float(shared.opts.cache_dit_threshold)
+    if shared.opts.cache_dit_warmup >= 0:
+        config_args['max_warmup_steps'] = int(shared.opts.cache_dit_warmup)
     cache_config = cache_dit.BasicCacheConfig(**config_args)
     if shared.opts.cache_dit_calibrator == "TaylorSeer":
         calibrator_config = cache_dit.TaylorSeerCalibratorConfig(taylorseer_order=1)
@@ -29,17 +38,6 @@ def apply_cache_dit(pipe):
         calibrator_config = cache_dit.FoCaCalibratorConfig()
     else:
         calibrator_config = None
-    """
-        Fn_compute_blocks=shared.opts.cache_dit_Fn_compute_blocks, # 8
-        Bn_compute_blocks=shared.opts.cache_dit_Bn_compute_blocks, # 0
-        residual_diff_threshold=shared.opts.cache_dit_residual_diff_threshold, # 0.08
-        max_warmup_steps=shared.opts.cache_dit_max_warmup_steps, # 8
-        max_cached_steps=shared.opts.cache_dit_max_cached_steps, # -1
-        max_continuous_cached_steps=shared.opts.cache_dit_max_continuous_cached_steps, # -1
-        enable_separate_cfg=shared.opts.cache_dit_enable_separate_cfg, # False
-        cfg_compute_first=shared.opts.cache_dit_cfg_compute_first, # False
-    )
-    """
     shared.log.info(f'Apply Cache-DiT: config="{cache_config.strify()}" calibrator="{calibrator_config.strify() if calibrator_config else "None"}"')
     cache_dit.enable_cache(
         pipe,
@@ -54,8 +52,8 @@ def unapply_cache_dir(pipe):
         return
     try:
         import cache_dit
-        stats = cache_dit.summary(pipe)
-        shared.log.critical(f'Unapply Cache-DiT: {stats}')
+        # stats = cache_dit.summary(pipe)
+        # shared.log.critical(f'Unapply Cache-DiT: {stats}')
         cache_dit.disable_cache(pipe)
         pipe.has_cache_dit = False
     except Exception:

modules/shared.py

@@ -359,6 +359,10 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), {
     "cache_dit_sep": OptionInfo("<h2>Cache-DiT</h2>", "", gr.HTML),
     "cache_dit_enabled": OptionInfo(False, "Cache-DiT enabled"),
     "cache_dit_calibrator": OptionInfo("None", "Cache-DiT calibrator", gr.Radio, {"choices": ["None", "TaylorSeer", "FoCa"]}),
+    "cache_dit_fcompute": OptionInfo(-1, "Cache-DiT F-compute blocks", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),
+    "cache_dit_bcompute": OptionInfo(-1, "Cache-DiT B-compute blocks", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),
+    "cache_dit_threshold": OptionInfo(-1, "Cache-DiT residual diff threshold", gr.Slider, {"minimum": -1.0, "maximum": 1.0, "step": 0.01}),
+    "cache_dit_warmup": OptionInfo(-1, "Cache-DiT warmup steps", gr.Slider, {"minimum": -1, "maximum": 50, "step": 1}),
 
     "faster_cache__sep": OptionInfo("<h2>Faster Cache</h2>", "", gr.HTML),
     "faster_cache_enabled": OptionInfo(False, "FasterCache cache enabled"),