From 2d15e697d0f201fcc70bc6b953c1fd91a0eb463e Mon Sep 17 00:00:00 2001
From: Shuming Liu <smshine@qq.com>
Date: Sat, 31 Jan 2026 16:00:38 +0800
Subject: [PATCH] fixed a bug

---
 video_cleaner.py | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/video_cleaner.py b/video_cleaner.py
index f86d29b..b06f07c 100644
--- a/video_cleaner.py
+++ b/video_cleaner.py
@@ -15,17 +15,29 @@ except ImportError:
     sys.exit(1)
 
 class VideoCleaner:
-    def __init__(self, model_path='yolo26n.pt', brightness_threshold=25, age_days=30):
-        print(f"Initializing YOLO model: {model_path}...")
-        try:
-            self.model = YOLO(model_path)
-        except Exception as e:
-            print(f"Failed to load YOLO model: {e}")
-            sys.exit(1)
-            
+    def __init__(self, model_path='yolo26n.pt', brightness_threshold=25, age_days=30, workers=4):
+        self.model_path = model_path
         self.brightness_threshold = brightness_threshold
         self.age_limit = timedelta(days=age_days)
+        self.workers = workers
         self.supported_extensions = ('.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv')
+        
+        # We'll initialize the model inside each process for thread/process safety
+        self.model = None 
+
+    def _get_model(self):
+        """Lazy initialization of the model for each process."""
+        if self.model is None:
+            # print(f"[{os.getpid()}] Initializing YOLO model: {self.model_path}...")
+            try:
+                self.model = YOLO(self.model_path)
+                # Ensure it's on GPU if available
+                if torch.cuda.is_available():
+                    self.model.to('cuda')
+            except Exception as e:
+                print(f"Failed to load YOLO model in process {os.getpid()}: {e}")
+                sys.exit(1)
+        return self.model
 
     def is_frame_color_and_bright(self, frame):
         """Checks if a single frame is in 'Day/Lights-on' mode."""