Updated code to work

florence_sam/{pipeline.py → florancesam_pipeline.py}

@@ -12,6 +12,8 @@ from utils.sam import load_sam_image_model, load_sam_video_model, run_sam_infere
 class VideoProcessor:
     def __init__(self, device=None):
         self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.autocast = torch.autocast(device_type=self.device.type, dtype=torch.bfloat16)
+        self.autocast.__enter__()
         self._enable_mixed_precision()
 
         # Load models
@@ -26,11 +28,20 @@ class VideoProcessor:
         )
 
     def _enable_mixed_precision(self):
-        torch.autocast(device_type=self.device.type, dtype=torch.bfloat16).__enter__()
         if torch.cuda.is_available() and torch.cuda.get_device_properties(0).major >= 8:
             torch.backends.cuda.matmul.allow_tf32 = True
             torch.backends.cudnn.allow_tf32 = True
 
+    def _reset_mixed_precision(self):
+        # Exit the autocast context
+        self.autocast.__exit__(None, None, None)
+        
+        # Reset CUDA settings
+        if torch.cuda.is_available():
+            torch.backends.cuda.matmul.allow_tf32 = False
+            torch.backends.cudnn.allow_tf32 = False
+            torch.cuda.empty_cache()  # Clear the CUDA cache
+
     def process_video(self, video_path, scale_factor, prompt):
         self.scale_factor = scale_factor
 
@@ -170,8 +181,6 @@ class VideoProcessor:
 
         return output_video_path, session_path
 
-video_processor = VideoProcessor()
-
 #Example usage
 # output_video = video_processor.process_video(
 #     video_path="videos/clip-07-camera-2.mp4", 

florence_sam/main.py

@@ -1,7 +1,11 @@
 import os
-from pipeline import video_processor
+import torch
 from propainter_pipeline import process_video
+from florancesam_pipeline import VideoProcessor
 
+video_processor = VideoProcessor()
+
+print("Processing video with FlorenceSam...")
 session_path, fps = video_processor.process_video(
     video_path="/home/ubuntu/ahmedghani/clip-07-camera-2.mp4", 
     scale_factor=0.5, 
@@ -9,4 +13,9 @@ session_path, fps = video_processor.process_video(
 )
 print(f"Processed video saved at: {session_path}") 
 print(f"FPS: {fps}")
-process_video(video=os.path.join(session_path, "input_frames"), mask=os.path.join(session_path, "output_frames"), resize_ratio=0.5, save_fps=int(fps), fp16=True)
+
+video_processor._reset_mixed_precision()
+del video_processor
+torch.cuda.empty_cache()
+print("Processing video with ProPainter...")
+process_video(video=os.path.join(session_path, "input_frames"), mask=os.path.join(session_path, "output_frames"), save_fps=int(fps), fp16=True)