Update handler.py

handler.py

@@ -8,6 +8,7 @@ from diffusers import HunyuanVideoPipeline
 from varnish import Varnish
 
 from enhance_a_video import enable_enhance, inject_enhance_for_hunyuanvideo, set_enhance_weight
+from teacache import enable_teacache, disable_teacache
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -42,6 +43,10 @@ class GenerationConfig:
     audio_prompt: str = ""
     audio_negative_prompt: str = "voices, voice, talking, speaking, speech"
 
+    # TeaCache settings
+    enable_teacache: bool = True
+    teacache_threshold: float = 0.15
+
     def validate_and_adjust(self) -> 'GenerationConfig':
         """Validate and adjust parameters"""
         # Ensure num_frames follows 4k + 1 format
@@ -128,6 +133,8 @@ class EndpointHandler:
             enable_audio=params.get("enable_audio", False),
             audio_prompt=params.get("audio_prompt", ""),
             audio_negative_prompt=params.get("audio_negative_prompt", "voices, voice, talking, speaking, speech"),
+            enable_teacache=params.get("enable_teacache", True),
+            teacache_threshold=params.get("teacache_threshold", 0.15)
         ).validate_and_adjust()
 
         try:
@@ -139,6 +146,16 @@ class EndpointHandler:
             else:
                 generator = None
 
+            # Configure TeaCache
+            if config.enable_teacache:
+                enable_teacache(
+                    self.pipeline.transformer,
+                    num_inference_steps=config.num_inference_steps,
+                    rel_l1_thresh=config.teacache_threshold
+                )
+            else:
+                disable_teacache(self.pipeline.transformer)
+
             # Generate video frames
             with torch.inference_mode():
                 output = self.pipeline(
@@ -192,6 +209,8 @@ class EndpointHandler:
                         "fps": result.metadata.fps,
                         "duration": result.metadata.duration,
                         "seed": config.seed,
+                        "teacache_enabled": config.enable_teacache,
+                        "teacache_threshold": config.teacache_threshold if config.enable_teacache else 0,
                     }
                 }