Update handler.py

handler.py

@@ -280,6 +280,32 @@ class EndpointHandler:
             enable_mmaudio=True,
         )
 
+        # Store TeaCache config for each model
+        self.text_to_video_teacache = None
+        self.image_to_video_teacache = None
+
+    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,
@@ -444,6 +470,7 @@ class EndpointHandler:
                 
                 # Check if image-to-video generation is requested
                 if input_image:
+                    self._configure_teacache(self.image_to_video, config)
                     processed_image = process_input_image(
                         input_image,
                         config.width,
@@ -453,6 +480,7 @@ class EndpointHandler:
                     generation_kwargs["image"] = processed_image
                     frames = self.image_to_video(**generation_kwargs).frames
                 else:
+                    self._configure_teacache(self.text_to_video, config)
                     frames = self.text_to_video(**generation_kwargs).frames
 
                 try: