Update handler.py

handler.py

@@ -16,7 +16,6 @@ from diffusers import LTXPipeline, LTXImageToVideoPipeline
 from diffusers.hooks import apply_enhance_a_video, EnhanceAVideoConfig
 from PIL import Image
 
-from teacache import TeaCacheConfig, enable_teacache, disable_teacache
 from varnish import Varnish
 from varnish.utils import is_truthy, process_input_image
 
@@ -149,12 +148,16 @@ class EndpointHandler:
                 torch_dtype=torch.bfloat16
             ).to("cuda")
 
+            apply_teacache(self.image_to_video)
+
         else:
             # Initialize models with bfloat16 precision
             self.text_to_video = LTXPipeline.from_pretrained(
                 model_path,
                 torch_dtype=torch.bfloat16
             ).to("cuda")
+            
+            apply_teacache(self.text_to_video)
 
         # Initialize LoRA tracking
         self._current_lora_model = None
@@ -195,32 +198,11 @@ class EndpointHandler:
             enable_mmaudio=True,
         )
 
-        # Store TeaCache config for each model
-        self.text_to_video_teacache = None
-        self.image_to_video_teacache = None
+        # Determine if TeaCache is already installed or not
+        self.text_to_video_teacache = False
+        self.image_to_video_teacache = False
+
 
-    def _configure_teacache(self, model, config: GenerationConfig):
-        """Configure TeaCache for a model based on generation config
-        
-        Args:
-            model: The model to configure TeaCache for
-            config: Generation configuration
-        """
-        if config.enable_teacache:
-            # Create and enable TeaCache if it should be enabled
-            teacache_config = TeaCacheConfig(
-                enabled=True,
-                rel_l1_thresh=config.teacache_threshold,
-                num_inference_steps=config.num_inference_steps
-            )
-            enable_teacache(model.transformer.__class__, teacache_config)
-            logger.info(f"TeaCache enabled with threshold {config.teacache_threshold}")
-        else:
-            # Disable TeaCache if it was previously enabled
-            if hasattr(model.transformer.__class__, 'teacache_config'):
-                disable_teacache(model.transformer.__class__)
-                logger.info("TeaCache disabled")
-    
     async def process_frames(
         self,
         frames: torch.Tensor,
@@ -451,7 +433,6 @@ class EndpointHandler:
                 
                 # Check if image-to-video generation is requested
                 if support_image_prompt and input_image:
-                    self._configure_teacache(self.image_to_video, config)
                     processed_image = process_input_image(
                         input_image,
                         config.width,
@@ -463,7 +444,6 @@ class EndpointHandler:
                     # apply_enhance_a_video(self.image_to_video.transformer, enhance_a_video_config)
                     frames = self.image_to_video(**generation_kwargs).frames
                 else:
-                    self._configure_teacache(self.text_to_video, config)
                     # disabled (we cannot install the hook multiple times, we would have to uninstall it first or find another way to dynamically enable it, eg. using the weight only)
                     # apply_enhance_a_video(self.text_to_video.transformer, enhance_a_video_config)
                     frames = self.text_to_video(**generation_kwargs).frames