Update handler.py

handler.py

@@ -75,13 +75,29 @@ class EndpointHandler:
             path: Path to model weights
         """
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+          
+        # Initialize transformer with Enhance-A-Video injection first
+        transformer = HunyuanVideoTransformer3DModel.from_pretrained(
+            path,
+            subfolder="transformer",
+            torch_dtype=torch.bfloat16
+        )
+        inject_enhance_for_hunyuanvideo(transformer)
         
-        # Initialize HunyuanVideo pipeline
+        # Initialize HunyuanVideo pipeline with the enhanced transformer
         self.pipeline = HunyuanVideoPipeline.from_pretrained(
             path,
+            transformer=transformer,
             torch_dtype=torch.float16,
         ).to(self.device)
         
+        # Initialize HunyuanVideo pipeline
+        self.pipeline = HunyuanVideoPipeline.from_pretrained(
+            path,
+            torch_dtype=torch.float16,
+        ).to(self.device)
+
         # Initialize text encoders in float16
         self.pipeline.text_encoder = self.pipeline.text_encoder.half()
         self.pipeline.text_encoder_2 = self.pipeline.text_encoder_2.half()
@@ -92,11 +108,6 @@ class EndpointHandler:
         # Initialize VAE in float16
         self.pipeline.vae = self.pipeline.vae.half()
 
-        # Initialize Enhance-A-Video
-        inject_enhance_for_hunyuanvideo(self.pipeline.transformer)
-        set_enhance_weight(4.0)  # default weight
-        enable_enhance()
-        
         # Initialize Varnish for post-processing
         self.varnish = Varnish(
             device=self.device,