Update teacache.py

teacache.py

@@ -14,7 +14,7 @@
 
 import re
 from dataclasses import dataclass, field
-from typing import List, Optional
+from typing import List, Optional, Tuple, Dict, Type, Any, Union
 
 import numpy as np
 import torch
@@ -72,12 +72,10 @@ _DEFAULT_SKIP_LAYER_IDENTIFIERS = [
 @dataclass
 class TeaCacheConfig:
     l1_threshold: Optional[float] = None
-
-    # Julian: this is a fix for @dataclass
+    # Fixed with default_factory
     skip_layer_identifiers: List[str] = field(default_factory=lambda: _DEFAULT_SKIP_LAYER_IDENTIFIERS.copy())
-
     _polynomial_coefficients: Optional[List[float]] = None
-
+    # Added missing field
     timestep_modulated_layer_identifier: Optional[str] = None
 
 
@@ -112,13 +110,15 @@ def apply_teacache(
     if denoiser is None:
         denoiser = pipeline.transformer if hasattr(pipeline, "transformer") else pipeline.unet
 
-    # Get model type name for lookup
-    model_type_name = denoiser.__class__.__name__
+    # Fix for isinstance() issue - convert to explicit type checking
+    model_type = denoiser.__class__
+    model_type_name = model_type.__name__
     logger.info(f"Detected model type: {model_type_name}")
 
+    # Modified to avoid using isinstance with dict_keys
     supported_model_types = tuple(_MODEL_TO_POLY_COEFFICIENTS.keys())
     is_supported_model = False
-
+    
     # Check if model_type is one of the supported types
     for supported_type in supported_model_types:
         if model_type == supported_type:
@@ -128,43 +128,57 @@ def apply_teacache(
     if is_supported_model:
         if config.l1_threshold is None:
             logger.info(
-                f"No L1 threshold was provided for {type(denoiser)}. Using default threshold as provided in the TeaCache paper for 1.5x speedup. "
+                f"No L1 threshold was provided for {model_type}. Using default threshold as provided in the TeaCache paper for 1.5x speedup. "
                 f"For higher speedup, increase the threshold."
             )
-            config.l1_threshold = _MODEL_TO_1_POINT_5X_SPEEDUP_THRESHOLD[type(denoiser)]
+            config.l1_threshold = _MODEL_TO_1_POINT_5X_SPEEDUP_THRESHOLD[model_type]
         if config.timestep_modulated_layer_identifier is None:
             logger.info(
-                f"No timestep modulated layer identifier was provided for {type(denoiser)}. Using default identifier as provided in the TeaCache paper."
+                f"No timestep modulated layer identifier was provided for {model_type}. Using default identifier as provided in the TeaCache paper."
             )
-            config.timestep_modulated_layer_identifier = _MODEL_TO_TIMESTEP_MODULATED_LAYER_IDENTIFIER[type(denoiser)]
+            if model_type in _MODEL_TO_TIMESTEP_MODULATED_LAYER_IDENTIFIER:
+                config.timestep_modulated_layer_identifier = _MODEL_TO_TIMESTEP_MODULATED_LAYER_IDENTIFIER[model_type]
+            else:
+                # Use a default for model types without a specific identifier
+                config.timestep_modulated_layer_identifier = "transformer_blocks.0.norm1"
         if config._polynomial_coefficients is None:
             logger.info(
-                f"No polynomial coefficients were provided for {type(denoiser)}. Using default coefficients as provided in the TeaCache paper."
+                f"No polynomial coefficients were provided for {model_type}. Using default coefficients as provided in the TeaCache paper."
             )
-            config._polynomial_coefficients = _MODEL_TO_POLY_COEFFICIENTS[type(denoiser)]
+            config._polynomial_coefficients = _MODEL_TO_POLY_COEFFICIENTS[model_type]
     else:
-        if config.l1_threshold is None:
-            raise ValueError(
-                f"No L1 threshold was provided for {type(denoiser)}. Using TeaCache with this model is not supported "
-                f"in Diffusers. Please provide the L1 threshold in the config by setting the `l1_threshold` attribute."
-            )
-        if config.timestep_modulated_layer_identifier is None:
-            raise ValueError(
-                f"No timestep modulated layer identifier was provided for {type(denoiser)}. Using TeaCache with this model is not supported "
-                f"in Diffusers. Please provide the layer identifier in the config by setting the `timestep_modulated_layer_identifier` attribute."
-            )
-        if config._polynomial_coefficients is None:
-            raise ValueError(
-                f"No polynomial coefficients were provided for {type(denoiser)}. Using TeaCache with this model is not "
-                f"supported in Diffusers. Please provide the polynomial coefficients in the config by setting the "
-                f"`_polynomial_coefficients` attribute. Automatic calibration will be implemented in the future."
-            )
+        # Handle unsupported model types with more helpful defaults for LTX models
+        if model_type_name == "LTXVideoTransformer3DModel" or "LTX" in model_type_name:
+            logger.info(f"Model {model_type_name} appears to be an LTX model variant. Using LTX defaults.")
+            if config.l1_threshold is None:
+                config.l1_threshold = 0.05
+            if config.timestep_modulated_layer_identifier is None:
+                config.timestep_modulated_layer_identifier = "transformer_blocks.0.norm1"
+            if config._polynomial_coefficients is None:
+                config._polynomial_coefficients = [2.14700694e01, -1.28016453e01, 2.31279151e00, 7.92487521e-01, 9.69274326e-03]
+        else:
+            if config.l1_threshold is None:
+                raise ValueError(
+                    f"No L1 threshold was provided for {model_type}. Using TeaCache with this model is not supported "
+                    f"in Diffusers. Please provide the L1 threshold in the config by setting the `l1_threshold` attribute."
+                )
+            if config.timestep_modulated_layer_identifier is None:
+                raise ValueError(
+                    f"No timestep modulated layer identifier was provided for {model_type}. Using TeaCache with this model is not supported "
+                    f"in Diffusers. Please provide the layer identifier in the config by setting the `timestep_modulated_layer_identifier` attribute."
+                )
+            if config._polynomial_coefficients is None:
+                raise ValueError(
+                    f"No polynomial coefficients were provided for {model_type}. Using TeaCache with this model is not "
+                    f"supported in Diffusers. Please provide the polynomial coefficients in the config by setting the "
+                    f"`_polynomial_coefficients` attribute. Automatic calibration will be implemented in the future."
+                )
 
     timestep_modulated_layer_matches = list(
         {
             module
             for name, module in denoiser.named_modules()
-            if re.match(config.timestep_modulated_layer_identifier, name)
+            if re.search(config.timestep_modulated_layer_identifier, name)
         }
     )
 
@@ -190,7 +204,7 @@ def apply_teacache(
         {
             module
             for name, module in denoiser.named_modules()
-            if any(re.match(identifier, name) for identifier in skip_layer_identifiers)
+            if any(re.search(identifier, name) for identifier in skip_layer_identifiers)
         }
     )