add para-attention first-block-cache

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

CHANGELOG.md

@@ -25,6 +25,12 @@
   - **hunyuan video** support for [FastHunyuan](https://huggingface.co/FastVideo/FastHunyuan)  
     simply select model variant and set appropriate parameters  
     recommended: sampler-shift=17, steps=6, resolution=720x1280, frames=125, guidance>6.0  
+- [ParaAttention](https://github.com/chengzeyi/ParaAttention)
+  - first-block caching that can significantly speed up generation by dynamically reusing partial outputs between steps  
+  - available for: flux, hunyuan-video, ltx-video, mochi  
+  - enable in *settings -> pipeline modifiers -> para-attention*  
+  - adjust residual diff threshold to balance the speedup and the accuracy:  
+    higher values leads to more cache hits and speedups, but might also lead to a higher accuracy drop  
 - **Other**:  
   - **upscale**: new [asymmetric vae](Heasterian/AsymmetricAutoencoderKLUpscaler) upscaling method
   - **ipex**: update supported torch versions  
@@ -43,6 +49,8 @@
   - update ui element ids  
   - modernui use local font  
   - unique font family registration  
+  - mochi video number of frames  
+  - mark large models that should offload  
 
 ## Update for 2025-01-15
 

modules/modeldata.py

@@ -29,6 +29,8 @@ def get_model_type(pipe):
         model_type = 'auraflow'
     elif "Flux" in name:
         model_type = 'f1'
+    elif "Mochi" in name:
+        model_type = 'mochi'
     elif "Lumina" in name:
         model_type = 'lumina'
     elif "OmniGen" in name:

modules/para_attention.py

@@ -0,0 +1,20 @@
+from modules import shared
+
+
+supported_models = ['Flux', 'HunyuanVideo', 'CogVideoX', 'Mochi']
+
+
+def apply_first_block_cache(p):
+    if not shared.opts.para_cache_enabled or not shared.native:
+        return
+    if not any(p.sd_model.__class__.__name__.startswith(x) for x in supported_models):
+        return
+    from installer import install
+    install('para_attn')
+    try:
+        from para_attn.first_block_cache import diffusers_adapters
+        diffusers_adapters.apply_cache_on_pipe(p.sd_model, residual_diff_threshold=shared.opts.para_diff_threshold)
+        shared.log.info(f'Applying para-attn first-block-cache: diff-threshold={shared.opts.para_diff_threshold} cls={p.sd_model.__class__.__name__}')
+    except Exception as e:
+        shared.log.error(f'Applying para-attn first-block-cache: {e}')
+        return

modules/processing.py

@@ -4,7 +4,7 @@ import time
 from contextlib import nullcontext
 import numpy as np
 from PIL import Image, ImageOps
-from modules import shared, devices, errors, images, scripts, memstats, lowvram, script_callbacks, extra_networks, detailer, sd_hijack_freeu, sd_models, sd_checkpoint, sd_vae, processing_helpers, timer, face_restoration, token_merge
+from modules import shared, devices, errors, images, scripts, memstats, lowvram, script_callbacks, extra_networks, detailer, sd_models, sd_checkpoint, sd_vae, processing_helpers, timer, face_restoration, token_merge
 from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile_unet
 from modules.processing_class import StableDiffusionProcessing, StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, StableDiffusionProcessingControl # pylint: disable=unused-import
 from modules.processing_info import create_infotext
@@ -168,7 +168,9 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
         shared.prompt_styles.extract_comments(p)
         if shared.opts.cuda_compile_backend == 'none':
             token_merge.apply_token_merging(p.sd_model)
+            from modules import sd_hijack_freeu, para_attention
             sd_hijack_freeu.apply_freeu(p, not shared.native)
+            para_attention.apply_first_block_cache(p)
 
         if p.width is not None:
             p.width = 8 * int(p.width / 8)

modules/sd_offload.py

@@ -10,7 +10,7 @@ from modules.timer import process as process_timer
 
 
 debug_move = shared.log.trace if os.environ.get('SD_MOVE_DEBUG', None) is not None else lambda *args, **kwargs: None
-should_offload = ['sc', 'sd3', 'f1', 'hunyuandit', 'auraflow', 'omnigen']
+should_offload = ['sc', 'sd3', 'f1', 'hunyuandit', 'auraflow', 'omnigen', 'hunyuanvideo', 'cogvideox', 'mochi']
 offload_hook_instance = None
 
 

modules/shared.py

@@ -597,6 +597,10 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), {
     "pag_sep": OptionInfo("<h2>Perturbed-Attention Guidance</h2>", "", gr.HTML),
     "pag_apply_layers": OptionInfo("m0", "PAG layer names"),
 
+    "para_sep": OptionInfo("<h2>Para-Attention</h2>", "", gr.HTML),
+    "para_cache_enabled": OptionInfo(False, "First-block cache enabled"),
+    "para_diff_threshold": OptionInfo(0.1, "Residual diff threshold", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
+
     "hypertile_sep": OptionInfo("<h2>HyperTile</h2>", "", gr.HTML),
     "hypertile_unet_enabled": OptionInfo(False, "UNet Enabled"),
     "hypertile_hires_only": OptionInfo(False, "HiRes pass only"),