Update handler.py

handler.py

@@ -9,10 +9,10 @@ import torch
 
 # note: there is no HunyuanImageToVideoPipeline yet in Diffusers
 from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
+from diffusers.hooks import apply_enhance_a_video, EnhanceAVideoConfig
 from varnish import Varnish
 from varnish.utils import is_truthy, process_input_image
 
-from enhance_a_video import enable_enhance, inject_enhance_for_hunyuanvideo, set_enhance_weight
 from teacache import enable_teacache, disable_teacache
 
 # Configure logging
@@ -258,13 +258,6 @@ class EndpointHandler:
             #    disable_teacache(self.pipeline.transformer)
 
             with torch.inference_mode():
-                # Configure Enhance-A-Video weight if enabled
-                if config.enable_enhance_a_video:
-                    set_enhance_weight(config.enhance_a_video_weight)
-                    enable_enhance()
-                else:
-                    # Reset enhance weight to 0 to effectively disable it
-                    set_enhance_weight(0)
 
                 # Prepare generation parameters
                 generation_kwargs = {
@@ -314,8 +307,6 @@ class EndpointHandler:
                 if config.lora_model_trigger:
                     generation_kwargs["prompt"] = f"{config.lora_model_trigger} {generation_kwargs['prompt']}"
 
- 
-
                 # Check if image-to-video generation is requested
                 if support_image_prompt and input_image:
                     self._configure_teacache(self.image_to_video, config)
@@ -326,9 +317,23 @@ class EndpointHandler:
                         config.input_image_quality,
                     )
                     generation_kwargs["image"] = processed_image
+
+                    apply_enhance_a_video(pipe.image_to_video.transformer, EnhanceAVideoConfig(
+                        weight=config.enhance_a_video_weight if config.enable_enhance_a_video else 0.0,
+                        num_frames_callback=lambda: (config.num_frames - 1),
+                        _attention_type=1
+                    ))
+                    
                     frames = self.image_to_video(**generation_kwargs).frames
                 else:
                     self._configure_teacache(self.text_to_video, config)
+
+                    apply_enhance_a_video(pipe.text_to_video.transformer, EnhanceAVideoConfig(
+                        weight=config.enhance_a_video_weight if config.enable_enhance_a_video else 0.0,
+                        num_frames_callback=lambda: (config.num_frames - 1),
+                        _attention_type=1
+                    ))
+                    
                     frames = self.text_to_video(**generation_kwargs).frames