Update handler.py

handler.py

@@ -13,11 +13,9 @@ import os
 import numpy as np
 import torch
 from diffusers import LTXPipeline, LTXImageToVideoPipeline
+from diffusers.hooks import apply_enhance_a_video, EnhanceAVideoConfig
 from PIL import Image
 
-from enhance_a_video import enable_enhance, set_enhance_weight
-from enhance import inject_enhance_for_ltx
-
 from teacache import TeaCacheConfig, enable_teacache, disable_teacache
 from varnish import Varnish
 from varnish.utils import is_truthy, process_input_image
@@ -150,6 +148,7 @@ class EndpointHandler:
                 model_path,
                 torch_dtype=torch.bfloat16
             ).to("cuda")
+
         else:
             # Initialize models with bfloat16 precision
             self.text_to_video = LTXPipeline.from_pretrained(
@@ -447,9 +446,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
 
                 try: