Update handler.py

handler.py

@@ -7,6 +7,7 @@ import random
 import traceback
 import torch
 import os
+import gc
 
 # note: there is no HunyuanImageToVideoPipeline yet in Diffusers
 from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel, FasterCacheConfig
@@ -276,7 +277,7 @@ class EndpointHandler:
             #else:
             #    disable_teacache(self.pipeline.transformer)
 
-            with torch.inference_mode():
+            with torch.amp.autocast_mode.autocast('cuda', torch.bfloat16), torch.no_grad(), torch.inference_mode():
 
                 # Prepare generation parameters
                 generation_kwargs = {
@@ -360,7 +361,11 @@ class EndpointHandler:
                     asyncio.set_event_loop(loop)
                 
                 video_uri, metadata = loop.run_until_complete(self.process_frames(frames, config))
-                      
+                
+                torch.cuda.empty_cache()
+                torch.cuda.reset_peak_memory_stats()
+                gc.collect()
+                
                 return {
                     "video": video_uri,
                     "content-type": "video/mp4",