TeaCache_Lumina2.py

import torch
import numpy as np
from comfy.ldm.common_dit import pad_to_patch_size  # noqa
from unittest.mock import patch
import re

#DEFAULT_COEFFICIENTS = [393.76566581, -603.50993606, 209.10239044, -23.00726601, 0.86377344]
DEFAULT_COEFFICIENTS = [225.7042019806413, -608.8453716535591, 304.1869942338369, 124.21267720116742, -1.4089066892956552]

# referenced from https://github.com/spawner1145/TeaCache/blob/main/TeaCache4Lumina2/teacache_lumina2.py
# firstly transplanted by @fexli https://github.com/fexli
# retransplanted by @spawner1145 https://github.com/spawner1145
def teacache_forward_working(
    self, x, timesteps, context, num_tokens, attention_mask=None, transformer_options={}, **kwargs
):
    if not hasattr(self, 'teacache_state'):
        self.teacache_state = {
            "cnt": 0,
            "num_steps": transformer_options.get("num_steps"),
            "cache": transformer_options.get("cache", {}),
            "uncond_seq_len": transformer_options.get("uncond_seq_len")
        }
    if not isinstance(self.teacache_state.get("cache"), dict):
        self.teacache_state["cache"] = {}

    if self.teacache_state.get("num_steps") is None and transformer_options.get("num_steps") is not None:
        self.teacache_state["num_steps"] = transformer_options.get("num_steps")

    cap_feats = context
    cap_mask = attention_mask
    ref_latents = kwargs.get("ref_latents", [])
    ref_contexts = kwargs.get("ref_contexts", [])
    siglip_feats = kwargs.get("siglip_feats", [])
    transformer_options = kwargs.get("transformer_options", transformer_options) or {}
    bs, c_channels, h_img, w_img = x.shape
    x = pad_to_patch_size(x, (self.patch_size, self.patch_size))
    t = (1.0 - timesteps).to(dtype=x.dtype)
    t_in = t * getattr(self, "time_scale", 1.0)

    t_emb = self.t_embedder(t_in, dtype=x.dtype)
    adaln_input = t_emb

    if getattr(self, "clip_text_pooled_proj", None) is not None:
        pooled = kwargs.get("clip_text_pooled", None)
        if pooled is not None:
            pooled = self.clip_text_pooled_proj(pooled)
        else:
            clip_text_dim = getattr(self, "clip_text_dim", None)
            if clip_text_dim is None:
                clip_text_dim = t_emb.shape[-1]
            pooled = torch.zeros((x.shape[0], clip_text_dim), device=x.device, dtype=x.dtype)
        adaln_input = self.time_text_embed(torch.cat((t_emb, pooled), dim=-1))

    try:
        patchify_out = self.patchify_and_embed(
            x,
            cap_feats,
            cap_mask,
            adaln_input,
            num_tokens,
            ref_latents=ref_latents,
            ref_contexts=ref_contexts,
            siglip_feats=siglip_feats,
            transformer_options=transformer_options,
        )
    except TypeError:
        patchify_out = self.patchify_and_embed(x, cap_feats, cap_mask, adaln_input, num_tokens)

    if len(patchify_out) == 6:
        x, mask, img_size, cap_size, freqs_cis, timestep_zero_index = patchify_out
    else:
        x, mask, img_size, cap_size, freqs_cis = patchify_out
        timestep_zero_index = None
    freqs_cis = freqs_cis.to(x.device)
    max_seq_len = x.shape[1]
    should_calc = True

    enable_teacache = transformer_options.get('enable_teacache', False)
    current_cache = None
    modulated_inp = None
    
    if enable_teacache:
        cache_key = max_seq_len
        if cache_key not in self.teacache_state['cache']:
            self.teacache_state['cache'][cache_key] = {
                "accumulated_rel_l1_distance": 0.0,
                "previous_modulated_input": None,
                "previous_residual": None,
            }
        current_cache = self.teacache_state['cache'][cache_key]
        
        try:
           if self.layers and hasattr(self.layers[0], 'adaLN_modulation'):
               mod_result = self.layers[0].adaLN_modulation(adaln_input.clone())
               if isinstance(mod_result, (list, tuple)) and len(mod_result) > 0:
                    modulated_inp = mod_result[0]
               elif torch.is_tensor(mod_result):
                   modulated_inp = mod_result
               else:
                    raise ValueError("adaLN_modulation returned unexpected type or empty list/tuple")
           else:
                raise AttributeError("Layer 0 or adaLN_modulation not found")
        except Exception as e:
             print(f"Warning: TeaCache - Failed to get modulated_inp: {e}. Disabling cache for this step.")
             enable_teacache = False
             should_calc = True
             modulated_inp = None
             if current_cache:
                  current_cache["previous_modulated_input"] = None
                  current_cache["accumulated_rel_l1_distance"] = 0.0

    if enable_teacache and modulated_inp is not None and current_cache is not None:
        num_steps_in_state = self.teacache_state.get("num_steps")
        if num_steps_in_state is None or num_steps_in_state == 0:
            should_calc = True
            current_cache["accumulated_rel_l1_distance"] = 0.0
        elif self.teacache_state['cnt'] == 0 or self.teacache_state['cnt'] == num_steps_in_state - 1:
            should_calc = True
            current_cache["accumulated_rel_l1_distance"] = 0.0
        else:
            if current_cache.get("previous_modulated_input") is not None:
                # coefficients = [393.76566581, -603.50993606, 209.10239044, -23.00726601, 0.86377344]
                coefficients = transformer_options.get('coefficients', DEFAULT_COEFFICIENTS)
                if not isinstance(coefficients, (list, tuple)):
                      coefficients = DEFAULT_COEFFICIENTS
                
                try:
                    rescale_func = np.poly1d(coefficients)
                except Exception as e:
                    print(f"Warning: TeaCache np.poly1d failed with coefficients {coefficients}: {e}. Using default for this step.")
                    rescale_func = np.poly1d(DEFAULT_COEFFICIENTS)

                prev_mod_input = current_cache["previous_modulated_input"]
                if prev_mod_input.shape != modulated_inp.shape:
                     print(f"Warning: TeaCache - modulated input shape mismatch: prev={prev_mod_input.shape}, curr={modulated_inp.shape}. Forcing recalculation.")
                     should_calc = True
                     current_cache["accumulated_rel_l1_distance"] = 0.0
                     rel_l1_change = float('inf')
                else:
                    prev_mean = prev_mod_input.abs().mean()
                    if prev_mean.item() > 1e-9:
                        rel_l1_change = ((modulated_inp - prev_mod_input).abs().mean() / prev_mean).cpu().item()
                    else:
                        rel_l1_change = 0.0 if modulated_inp.abs().mean().item() < 1e-9 else float('inf')
                    
                    rescaled_value = rescale_func(rel_l1_change)
                    if np.isnan(rescaled_value) or np.isinf(rescaled_value):
                         current_cache["accumulated_rel_l1_distance"] = float('inf')
                    else:
                         current_cache["accumulated_rel_l1_distance"] += rescaled_value

                    if current_cache["accumulated_rel_l1_distance"] < transformer_options.get('rel_l1_thresh', 0.3):
                        should_calc = False
                    else:
                        should_calc = True
                        current_cache["accumulated_rel_l1_distance"] = 0.0
            else:
                 should_calc = True
                 current_cache["accumulated_rel_l1_distance"] = 0.0

        current_cache["previous_modulated_input"] = modulated_inp.clone()

        if self.teacache_state.get('uncond_seq_len') is None:
            self.teacache_state['uncond_seq_len'] = cache_key

        if num_steps_in_state is not None and cache_key != self.teacache_state.get('uncond_seq_len'):
            self.teacache_state['cnt'] += 1
            if self.teacache_state['cnt'] >= num_steps_in_state:
                self.teacache_state['cnt'] = 0

    can_reuse_residual = (enable_teacache and 
                         not should_calc and 
                         current_cache and 
                         current_cache.get("previous_residual") is not None and
                         current_cache["previous_residual"].shape == x.shape) # 检查形状

    if can_reuse_residual:
         processed_x = x + current_cache["previous_residual"]
    else:
        if enable_teacache and not should_calc and current_cache and current_cache.get("previous_residual") is not None and current_cache["previous_residual"].shape != x.shape:
             print(f"Warning: TeaCache - Residual shape mismatch: cache={current_cache['previous_residual'].shape}, input={x.shape}. Forcing recalculation.")
             if current_cache:
                  current_cache["accumulated_rel_l1_distance"] = 0.0

        original_x = x.clone()
        current_x_for_processing = x
        for layer in self.layers:
            try:
                current_x_for_processing = layer(
                    current_x_for_processing,
                    mask,
                    freqs_cis,
                    adaln_input,
                    timestep_zero_index=timestep_zero_index,
                    transformer_options=transformer_options,
                )
            except TypeError:
                current_x_for_processing = layer(current_x_for_processing, mask, freqs_cis, adaln_input)

        if enable_teacache and current_cache:
            current_cache["previous_residual"] = current_x_for_processing - original_x
            current_cache["accumulated_rel_l1_distance"] = 0.0
        processed_x = current_x_for_processing

    try:
        output = self.final_layer(processed_x, adaln_input, timestep_zero_index=timestep_zero_index)
    except TypeError:
        output = self.final_layer(processed_x, adaln_input)
    output = self.unpatchify(output, img_size, cap_size, return_tensor=True)[:, :, :h_img, :w_img]

    return -output


class TeaCache_Lumina2:

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "rel_l1_thresh": ("FLOAT", {"default": 6.0, "min": 0.0, "step": 0.001}),
                "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01,
                                            "tooltip": "The start percentage of the steps that will apply TeaCache. / TeaCache开始应用的步数百分比。"}),
                "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01,
                                          "tooltip": "The end percentage of the steps that will apply TeaCache. / TeaCache停止应用的步数百分比。"}),
                 "coefficients_string": ("STRING", {
                    "multiline": True, 
                    "default": str(DEFAULT_COEFFICIENTS),
                    "tooltip": "Coefficients for np.poly1d. Format: 393.7, -603.5, 209.1, -23.0, 0.86 (with or without brackets []) / 用于 np.poly1d 的系数。格式: 393.7, -603.5, 209.1, -23.0, 0.86 (可带或不带方括号[])"
                    }),
            }
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "patch_teacache"
    CATEGORY = "utils"

    def patch_teacache(self, model, rel_l1_thresh, start_percent, end_percent, coefficients_string):
        if rel_l1_thresh == 0:
            try:
               diffusion_model = model.get_model_object("diffusion_model")
               if hasattr(diffusion_model, 'teacache_state'):
                   delattr(diffusion_model, 'teacache_state')
            except:
                pass
            return (model,)

        parsed_coefficients = DEFAULT_COEFFICIENTS
        try:
            s = re.sub(r'[\[\]\s]', '', coefficients_string.strip())
            if not s:
                 parsed_coefficients = DEFAULT_COEFFICIENTS
            else:
                coeff_list = [float(item) for item in s.split(',') if item]
                if not coeff_list:
                     parsed_coefficients = DEFAULT_COEFFICIENTS
                else:
                    parsed_coefficients = coeff_list
        except ValueError:
            print(f"Warning: TeaCache - Could not parse coefficients string: '{coefficients_string}'. Using default coefficients: {DEFAULT_COEFFICIENTS}")
            parsed_coefficients = DEFAULT_COEFFICIENTS
        except Exception as e:
             print(f"Warning: TeaCache - Error parsing coefficients '{coefficients_string}': {e}. Using default coefficients: {DEFAULT_COEFFICIENTS}")
             parsed_coefficients = DEFAULT_COEFFICIENTS


        new_model = model.clone()
        if 'transformer_options' not in new_model.model_options:
            new_model.model_options['transformer_options'] = {}

        new_model.model_options["transformer_options"]["cache"] = {}
        new_model.model_options["transformer_options"]["uncond_seq_len"] = None
        new_model.model_options["transformer_options"]["rel_l1_thresh"] = rel_l1_thresh
        new_model.model_options["transformer_options"]["coefficients"] = parsed_coefficients
        
        diffusion_model = new_model.get_model_object("diffusion_model")

        if hasattr(diffusion_model, 'teacache_state'):
            delattr(diffusion_model, 'teacache_state')

        if hasattr(diffusion_model, "_forward"):
            context_patch_manager = patch.multiple(
                diffusion_model,
                _forward=teacache_forward_working.__get__(diffusion_model, diffusion_model.__class__)
            )
        else:
            context_patch_manager = patch.multiple(
                diffusion_model,
                forward=teacache_forward_working.__get__(diffusion_model, diffusion_model.__class__)
            )

        old_wrapper = new_model.model_options.get("model_function_wrapper")

        def unet_wrapper_function(model_function, kwargs):
            input_val = kwargs["input"]
            timestep = kwargs["timestep"]
            c_condition_dict = kwargs["c"]
            cond_or_uncond = kwargs["cond_or_uncond"]

            if not isinstance(c_condition_dict, dict): c_condition_dict = {}
            if "transformer_options" not in c_condition_dict or not isinstance(c_condition_dict["transformer_options"], dict):
                c_condition_dict["transformer_options"] = {}

            for key, value in new_model.model_options["transformer_options"].items():
                if key not in c_condition_dict["transformer_options"]:
                    c_condition_dict["transformer_options"][key] = value

            current_step_index = 0
            if "sample_sigmas" not in c_condition_dict["transformer_options"] or \
               c_condition_dict["transformer_options"]["sample_sigmas"] is None:
                print("warning: TeaCache - 'sample_sigmas' not found in c.transformer_options.TeaCache might not work correctly.")
                c_condition_dict["transformer_options"]["enable_teacache"] = False
                c_condition_dict["transformer_options"]["num_steps"] = 1
                if hasattr(diffusion_model, 'teacache_state'):
                     delattr(diffusion_model, 'teacache_state')
            else:
                sigmas = c_condition_dict["transformer_options"]["sample_sigmas"]
                total_sampler_steps = len(sigmas)
                if total_sampler_steps > 0:
                   c_condition_dict["transformer_options"]["num_steps"] = total_sampler_steps
                else:
                   c_condition_dict["transformer_options"]["num_steps"] = 1
                   c_condition_dict["transformer_options"]["enable_teacache"] = False


                if hasattr(diffusion_model, 'teacache_state') and diffusion_model.teacache_state is not None:
                    if diffusion_model.teacache_state.get("num_steps") != total_sampler_steps and total_sampler_steps > 0:
                        delattr(diffusion_model, 'teacache_state')
                        c_condition_dict["transformer_options"]["cache"] = {}

                current_timestep = timestep[0].to(device=sigmas.device, dtype=sigmas.dtype)
                # 使用 torch.isclose 和 any() 来处理浮点数比较
                close_mask = torch.isclose(sigmas, current_timestep, atol=1e-6)
                if close_mask.any():
                     matched_step_index = torch.nonzero(close_mask, as_tuple=True)[0]
                     current_step_index = matched_step_index[0].item()
                else:
                    current_step_index = 0 
                    if total_sampler_steps > 1:
                         try:
                            indices = torch.where(sigmas >= current_timestep)[0]
                            if len(indices) > 0:
                                 current_step_index = indices[-1].item()
                            else:
                                 current_step_index = total_sampler_steps -1
                         except:
                            for i in range(total_sampler_steps - 1):
                                if (sigmas[i] - current_timestep) * (sigmas[i + 1] - current_timestep) <= 0:
                                    current_step_index = i
                                    break

                current_percent = 0.0
                if total_sampler_steps > 1:
                    current_percent = current_step_index / max(1, (total_sampler_steps - 1))
                # elif total_sampler_steps == 1:
                #    current_percent = 0.0
                elif total_sampler_steps <= 0 :
                    # current_percent = 0.0
                    c_condition_dict["transformer_options"]["enable_teacache"] = False

                if start_percent <= current_percent <= end_percent and total_sampler_steps > 0:
                    c_condition_dict["transformer_options"]["enable_teacache"] = True
                else:
                    c_condition_dict["transformer_options"]["enable_teacache"] = False

            if current_step_index == 0:
                 # print("Debug: TeaCache - Resetting state at step 0")
                 if hasattr(diffusion_model, 'teacache_state'):
                     delattr(diffusion_model, 'teacache_state')
                 c_condition_dict["transformer_options"]["cache"] = {}


            with context_patch_manager:
                if old_wrapper:
                    return old_wrapper(model_function, {"input": input_val, "timestep": timestep, "c": c_condition_dict, "cond_or_uncond": cond_or_uncond})
                return model_function(input_val, timestep, **c_condition_dict)

        new_model.set_model_unet_function_wrapper(unet_wrapper_function)
        return (new_model,)