add teacache for flux

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

CHANGELOG.md

@@ -4,18 +4,22 @@
 
 - Video: add FasterCache and PAB support to WanDB and LTX models  
 - ZLUDA: add more GPUs to recognized list  
-- LoRA: obey configured device when performing calculations  
 - Progress: add additional fields to progress API  
 - Progress: use batch-count for progress  
 - Grid: add of max-rows and max-columns in settings to control grid format  
+- LoRA: add option to force-reload LoRA on every generate  
 - Gallery: add max-columns in settings for gradio gallery components  
 - Styles: resize and bring quick-ui to forward on hover  
-- Logging: fix debug logging  
 - Logging: logging cleanup  
 - Params: Reset default guidance-rescale from 0.7 to 0.0  
-- Diag: add get-server-status to ui generate context menu  
+- Diag: add get-server-status to UI generate context menu  
 - Pipe: [SoftFill](https://github.com/zacheryvaughn/softfill-pipelines)  
   select in scripts, available for sdxl in inpaint model  
+- Flux: TeaCache for Flux.1  
+- Fix: LoRA obey configured device when performing calculations  
+- Fix: ZLUDA startup
+- Fix: balanced offload remove non-blocking move op  
+- Fix: debug logging  
 
 ## Update for 2025-04-03
 

TODO.md

@@ -17,6 +17,7 @@ N/A
 - Video: API support  
 - Video: STG: <https://github.com/huggingface/diffusers/blob/main/examples/community/README.md#spatiotemporal-skip-guidance>  
 - Video: SmoothCache: https://github.com/huggingface/diffusers/issues/11135  
+- TeaCache: https://github.com/ali-vilab/TeaCache
 
 ## Code TODO
 

modules/lora/extra_networks_lora.py

@@ -139,6 +139,8 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
         return [f'{name}:{te}:{unet}' for name, te, unet in zip(names, te_multipliers, unet_multipliers)]
 
     def changed(self, requested: List[str], include: List[str], exclude: List[str]):
+        if shared.opts.lora_force_reload:
+            return True
         sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
         if not hasattr(sd_model, 'loaded_loras'):
             sd_model.loaded_loras = {}

modules/model_flux.py

@@ -223,6 +223,10 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
     shared.sd_model = None
     devices.torch_gc(force=True)
 
+    if shared.opts.teacache_enabled:
+        from modules import teacache
+        diffusers.FluxTransformer2DModel.forward = teacache.teacache_forward
+
     # load overrides if any
     if shared.opts.sd_unet != 'Default':
         try:

modules/para_attention.py

@@ -4,17 +4,17 @@ from modules import shared
 supported_models = ['Flux', 'HunyuanVideo', 'CogVideoX', 'Mochi']
 
 
-def apply_first_block_cache(p):
+def apply_first_block_cache():
     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):
+    if not any(shared.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__}')
+        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__}')
     except Exception as e:
-        shared.log.error(f'Applying para-attn first-block-cache: {e}')
+        shared.log.error(f'Transformers cache: type=paraattn {e}')
         return

modules/processing.py

@@ -169,9 +169,10 @@ 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
+            from modules import sd_hijack_freeu, para_attention, teacache
             sd_hijack_freeu.apply_freeu(p, not shared.native)
-            para_attention.apply_first_block_cache(p)
+            para_attention.apply_first_block_cache()
+            teacache.apply_teacache(p)
 
         if p.width is not None:
             p.width = 8 * int(p.width / 8)

modules/shared.py

@@ -560,13 +560,13 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), {
     "pag_apply_layers": OptionInfo("m0", "PAG layer names"),
 
     "pab_sep": OptionInfo("<h2>PAB: Pyramid attention broadcast </h2>", "", gr.HTML),
-    "pab_enabled": OptionInfo(False, "Attention cache enabled"),
-    "pab_spacial_skip_range": OptionInfo(2, "FC spacial skip range", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),
-    "pab_spacial_skip_start": OptionInfo(100, "FC spacial skip start", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}),
-    "pab_spacial_skip_end": OptionInfo(800, "FC spacial skip end", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}),
+    "pab_enabled": OptionInfo(False, "PAB cache enabled"),
+    "pab_spacial_skip_range": OptionInfo(2, "PAB spacial skip range", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),
+    "pab_spacial_skip_start": OptionInfo(100, "PAB spacial skip start", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}),
+    "pab_spacial_skip_end": OptionInfo(800, "PAB spacial skip end", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}),
 
     "faster_cache__sep": OptionInfo("<h2>Faster Cache</h2>", "", gr.HTML),
-    "faster_cache_enabled": OptionInfo(False, "Faster cache enabled"),
+    "faster_cache_enabled": OptionInfo(False, "FC cache enabled"),
     "fc_spacial_skip_range": OptionInfo(2, "FC spacial skip range", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),
     "fc_spacial_skip_start": OptionInfo(0, "FC spacial skip start", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}),
     "fc_spacial_skip_end": OptionInfo(681, "FC spacial skip end", gr.Slider, {"minimum": 0, "maximum": 1.0, "step": 0.01}),
@@ -581,6 +581,10 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), {
     "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}),
 
+    "teacache_sep": OptionInfo("<h2>TeaCache</h2>", "", gr.HTML),
+    "teacache_enabled": OptionInfo(False, "TC cache enabled"),
+    "teacache_thresh": OptionInfo(0.6, "TC L1 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"),
@@ -932,6 +936,7 @@ options_templates.update(options_section(('extra_networks', "Networks"), {
     "lora_fuse_diffusers": OptionInfo(True, "LoRA fuse directly to model"),
     "lora_apply_gpu": OptionInfo(False, "LoRA load directly on GPU"),
     "lora_legacy": OptionInfo(not native, "LoRA load using legacy method"),
+    "lora_force_reload": OptionInfo(False, "LoRA force reload always"),
     "lora_force_diffusers": OptionInfo(False if not cmd_opts.use_openvino else True, "LoRA load using Diffusers method"),
     "lora_maybe_diffusers": OptionInfo(False, "LoRA load using Diffusers method for selected models"),
     "lora_apply_tags": OptionInfo(0, "LoRA auto-apply tags", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),

modules/teacache/__init__.py

@@ -0,0 +1 @@
+from .teacache_flux import apply_teacache, teacache_forward

modules/teacache/teacache_flux.py

@@ -0,0 +1,322 @@
+from typing import Any, Dict, Optional, Union
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
+import torch
+import numpy as np
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def teacache_forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        pooled_projections: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        guidance: torch.Tensor = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_block_samples=None,
+        controlnet_single_block_samples=None,
+        return_dict: bool = True,
+        controlnet_blocks_repeat: bool = False,
+    ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+    """
+    The [`FluxTransformer2DModel`] forward method.
+
+    Args:
+        hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
+            Input `hidden_states`.
+        encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
+            Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+        pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+            from the embeddings of input conditions.
+        timestep ( `torch.LongTensor`):
+            Used to indicate denoising step.
+        block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+            A list of tensors that if specified are added to the residuals of transformer blocks.
+        joint_attention_kwargs (`dict`, *optional*):
+            A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+            `self.processor` in
+            [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+        return_dict (`bool`, *optional*, defaults to `True`):
+            Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+            tuple.
+
+    Returns:
+        If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+        `tuple` where the first element is the sample tensor.
+    """
+    if joint_attention_kwargs is not None:
+        joint_attention_kwargs = joint_attention_kwargs.copy()
+        lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+    else:
+        lora_scale = 1.0
+
+    if USE_PEFT_BACKEND:
+        # weight the lora layers by setting `lora_scale` for each PEFT layer
+        scale_lora_layers(self, lora_scale)
+    else:
+        if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+            logger.warning(
+                "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+            )
+
+    hidden_states = self.x_embedder(hidden_states)
+
+    timestep = timestep.to(hidden_states.dtype) * 1000
+    if guidance is not None:
+        guidance = guidance.to(hidden_states.dtype) * 1000
+    else:
+        guidance = None
+
+    temb = (
+        self.time_text_embed(timestep, pooled_projections)
+        if guidance is None
+        else self.time_text_embed(timestep, guidance, pooled_projections)
+    )
+    encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+    if txt_ids.ndim == 3:
+        logger.warning(
+            "Passing `txt_ids` 3d torch.Tensor is deprecated."
+            "Please remove the batch dimension and pass it as a 2d torch Tensor"
+        )
+        txt_ids = txt_ids[0]
+    if img_ids.ndim == 3:
+        logger.warning(
+            "Passing `img_ids` 3d torch.Tensor is deprecated."
+            "Please remove the batch dimension and pass it as a 2d torch Tensor"
+        )
+        img_ids = img_ids[0]
+
+    ids = torch.cat((txt_ids, img_ids), dim=0)
+    image_rotary_emb = self.pos_embed(ids)
+
+    if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs:
+        ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds")
+        ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds)
+        joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
+
+    if self.enable_teacache:
+        inp = hidden_states.clone()
+        temb_ = temb.clone()
+        modulated_inp, _gate_msa, _shift_mlp, _scale_mlp, _gate_mlp = self.transformer_blocks[0].norm1(inp, emb=temb_)
+        if self.cnt == 0 or self.cnt == self.num_steps-1:
+            should_calc = True
+            self.accumulated_rel_l1_distance = 0
+        else:
+            coefficients = [4.98651651e+02, -2.83781631e+02,  5.58554382e+01, -3.82021401e+00, 2.64230861e-01]
+            rescale_func = np.poly1d(coefficients)
+            self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item())
+            if self.accumulated_rel_l1_distance < self.rel_l1_thresh:
+                should_calc = False
+            else:
+                should_calc = True
+                self.accumulated_rel_l1_distance = 0
+        self.previous_modulated_input = modulated_inp
+        self.cnt += 1
+        if self.cnt == self.num_steps:
+            self.cnt = 0
+
+    if self.enable_teacache:
+        if not should_calc:
+            hidden_states += self.previous_residual
+        else:
+            ori_hidden_states = hidden_states.clone()
+            for index_block, block in enumerate(self.transformer_blocks):
+                if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                    def create_custom_forward4(module, return_dict=None):
+                        def custom_forward(*inputs):
+                            if return_dict is not None:
+                                return module(*inputs, return_dict=return_dict)
+                            else:
+                                return module(*inputs)
+
+                        return custom_forward
+
+                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                    encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward4(block),
+                        hidden_states,
+                        encoder_hidden_states,
+                        temb,
+                        image_rotary_emb,
+                        **ckpt_kwargs,
+                    )
+
+                else:
+                    encoder_hidden_states, hidden_states = block(
+                        hidden_states=hidden_states,
+                        encoder_hidden_states=encoder_hidden_states,
+                        temb=temb,
+                        image_rotary_emb=image_rotary_emb,
+                        joint_attention_kwargs=joint_attention_kwargs,
+                    )
+
+                # controlnet residual
+                if controlnet_block_samples is not None:
+                    interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                    interval_control = int(np.ceil(interval_control))
+                    # For Xlabs ControlNet.
+                    if controlnet_blocks_repeat:
+                        hidden_states = (
+                            hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+                        )
+                    else:
+                        hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+            hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+            for index_block, block in enumerate(self.single_transformer_blocks):
+                if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                    def create_custom_forward2(module, return_dict=None):
+                        def custom_forward(*inputs):
+                            if return_dict is not None:
+                                return module(*inputs, return_dict=return_dict)
+                            else:
+                                return module(*inputs)
+
+                        return custom_forward
+
+                    ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward2(block),
+                        hidden_states,
+                        temb,
+                        image_rotary_emb,
+                        **ckpt_kwargs,
+                    )
+
+                else:
+                    hidden_states = block(
+                        hidden_states=hidden_states,
+                        temb=temb,
+                        image_rotary_emb=image_rotary_emb,
+                        joint_attention_kwargs=joint_attention_kwargs,
+                    )
+
+                # controlnet residual
+                if controlnet_single_block_samples is not None:
+                    interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                    interval_control = int(np.ceil(interval_control))
+                    hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                        hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+                        + controlnet_single_block_samples[index_block // interval_control]
+                    )
+
+            hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+            self.previous_residual = hidden_states - ori_hidden_states
+    else:
+        for index_block, block in enumerate(self.transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                def create_custom_forward1(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward1(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_block_samples is not None:
+                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                # For Xlabs ControlNet.
+                if controlnet_blocks_repeat:
+                    hidden_states = (
+                        hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+                    )
+                else:
+                    hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+                def create_custom_forward3(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward3(block),
+                    hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_single_block_samples is not None:
+                interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                    hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+                    + controlnet_single_block_samples[index_block // interval_control]
+                )
+
+        hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+
+    hidden_states = self.norm_out(hidden_states, temb)
+    output = self.proj_out(hidden_states)
+
+    if USE_PEFT_BACKEND:
+        # remove `lora_scale` from each PEFT layer
+        unscale_lora_layers(self, lora_scale)
+
+    if not return_dict:
+        return (output,)
+
+    return Transformer2DModelOutput(sample=output)
+
+
+def apply_teacache(p):
+    from modules import shared
+    if not shared.native or not shared.opts.teacache_enabled or not shared.sd_model.__class__.__name__.startswith('Flux'):
+        return
+    shared.sd_model.transformer.__class__.enable_teacache = shared.opts.teacache_thresh > 0
+    shared.sd_model.transformer.__class__.cnt = 0
+    shared.sd_model.transformer.__class__.num_steps = p.steps
+    shared.sd_model.transformer.__class__.rel_l1_thresh = shared.opts.teacache_thresh # 0.25 for 1.5x speedup, 0.4 for 1.8x speedup, 0.6 for 2.0x speedup, 0.8 for 2.25x speedup
+    shared.sd_model.transformer.__class__.accumulated_rel_l1_distance = 0
+    shared.sd_model.transformer.__class__.previous_modulated_input = None
+    shared.sd_model.transformer.__class__.previous_residual = None
+    shared.log.info(f'Transformers cache: type=teacache thresh={shared.opts.teacache_thresh} cls={shared.sd_model.__class__.__name__}')

scripts/xyz_grid_classes.py

@@ -212,4 +212,5 @@ axis_options = [
     AxisOption("[IY] Scale", float, apply_task_arg('infusenet_conditioning_scale')),
     AxisOption("[IY] Start", float, apply_task_arg('infusenet_guidance_start')),
     AxisOption("[IY] End", float, apply_task_arg('infusenet_guidance_end')),
+    AxisOption("[TeaCache] Threshold", float, apply_setting('teacache_thresh')),
 ]