Update hawk/conversation/conversation_video.py

hawk/conversation/conversation_video.py

@@ -21,6 +21,8 @@ from hawk.processors.video_processor import ToTHWC,ToUint8,load_video,load_video
 from hawk.processors import Blip2ImageEvalProcessor
             
 from hawk.models.ImageBind.data import load_and_transform_audio_data
+from torch.cuda.amp import autocast
+
 
 class SeparatorStyle(Enum):
     """Different separator style."""
@@ -311,15 +313,16 @@ class Chat:
         else:
             raise NotImplementedError
         
-        
-        # conv.system = "You can understand the video that the user provides.  Follow the instructions carefully and explain your answers in detail."
-        image_emb, _, _ = self.model.encode_videoQformer_visual(video) # 1,32,4096
-        image_motion_emb, _, _ = self.model.encode_videoQformer_visual(video_motion, motion=True) # 1,32,4096
-        image_emb = image_emb.clone().detach()
-        image_motion_emb = image_motion_emb.clone().detach()
-        img_list.append(torch.cat((image_emb, image_motion_emb), dim=1))
-        # img_list.append(image_motion_emb) 
-        conv.append_message(conv.roles[0], "<Video><ImageHere></Video> ")
+        with autocast():
+            # conv.system = "You can understand the video that the user provides.  Follow the instructions carefully and explain your answers in detail."
+            image_emb, _, _ = self.model.encode_videoQformer_visual(video) # 1,32,4096
+            image_motion_emb, _, _ = self.model.encode_videoQformer_visual(video_motion, motion=True) # 1,32,4096
+            image_emb = image_emb.clone().detach()
+            image_motion_emb = image_motion_emb.clone().detach()
+            img_list.append(torch.cat((image_emb, image_motion_emb), dim=1))
+            # img_list.append(image_motion_emb) 
+            conv.append_message(conv.roles[0], "<Video><ImageHere></Video> ")
+            
         return "Received."
 
     def upload_img(self, image, conv, img_list):