Update roop/processors/frame/core.py

roop/processors/frame/core.py

@@ -29,11 +29,12 @@ def load_frame_processor_module(frame_processor: str) -> Any:
         frame_processor_module = importlib.import_module(f'roop.processors.frame.{frame_processor}')
         for method_name in FRAME_PROCESSORS_INTERFACE:
             if not hasattr(frame_processor_module, method_name):
-                raise NotImplementedError
+                raise NotImplementedError(f"Missing method: {method_name}")
         logging.info(f'Successfully loaded frame processor module: {frame_processor}')
     except (ImportError, NotImplementedError) as e:
         logging.error(f'Error loading frame processor {frame_processor}: {e}')
-        quit(f'Frame processor {frame_processor} crashed.')
+        raise SystemExit(f'Frame processor {frame_processor} crashed.')
+
     return frame_processor_module
 
 def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType]:
@@ -45,51 +46,62 @@ def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType
             FRAME_PROCESSORS_MODULES.append(frame_processor_module)
     return FRAME_PROCESSORS_MODULES
 
-def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None], update: Callable[[], None]) -> None:
-    with ThreadPoolExecutor(max_workers=roop.globals.execution_threads) as executor:
-        futures = []
-        queue = create_queue(temp_frame_paths)
-        queue_per_future = len(temp_frame_paths) // roop.globals.execution_threads
-        while not queue.empty():
-            future = executor.submit(process_frames, source_path, pick_queue(queue, queue_per_future), update)
-            futures.append(future)
-            logging.info(f'Submitted future for processing frames.')
-        for future in as_completed(futures):
-            try:
-                future.result()
-                logging.info('Frame processing completed for a future.')
-            except Exception as e:
-                logging.error(f'Error in processing frame: {e}')
+def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Callable[[], None]], None], update: Callable[[], None]) -> None:
+    try:
+        with ThreadPoolExecutor(max_workers=roop.globals.execution_threads) as executor:
+            futures = []
+            queue = create_queue(temp_frame_paths)
+            queue_per_future = len(temp_frame_paths) // roop.globals.execution_threads
+            while not queue.empty():
+                future = executor.submit(process_frames, source_path, pick_queue(queue, queue_per_future), update)
+                futures.append(future)
+                logging.info('Submitted future for processing frames.')
+                
+            for future in as_completed(futures):
+                try:
+                    future.result()
+                    logging.info('Frame processing completed for a future.')
+                except Exception as e:
+                    logging.error(f'Error in processing frame: {e}')
+    except Exception as e:
+        logging.error(f'Error in multi-process frame: {e}')
 
-def create_queue(temp_frame_paths: List[str]) -> Queue[str]:
-    queue: Queue[str] = Queue()
+def create_queue(temp_frame_paths: List[str]) -> Queue:
+    queue = Queue()
     for frame_path in temp_frame_paths:
         queue.put(frame_path)
     logging.info('Queue created with frame paths.')
     return queue
 
-def pick_queue(queue: Queue[str], queue_per_future: int) -> List[str]:
-    queues = []
+def pick_queue(queue: Queue, queue_per_future: int) -> List[str]:
+    picked_items = []
     for _ in range(queue_per_future):
         if not queue.empty():
-            queues.append(queue.get())
-    logging.info(f'Picked {len(queues)} items from queue for processing.')
-    return queues
+            picked_items.append(queue.get())
+    logging.info(f'Picked {len(picked_items)} items from queue for processing.')
+    return picked_items
 
-def process_video(source_path: str, frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None]) -> None:
+def process_video(source_path: str, frame_paths: List[str], process_frames: Callable[[str, List[str], Callable[[], None]], None]) -> None:
     progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
     total = len(frame_paths)
-    with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format, mininterval=0.1) as progress:
-        multi_process_frame(source_path, frame_paths, process_frames, lambda: update_progress(progress))
+    try:
+        with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format, mininterval=0.1) as progress:
+            multi_process_frame(source_path, frame_paths, process_frames, lambda: update_progress(progress))
+    except Exception as e:
+        logging.error(f'Error in processing video: {e}')
 
 def update_progress(progress: Any = None) -> None:
-    process = psutil.Process(os.getpid())
-    memory_usage = process.memory_info().rss / 1024 / 1024 / 1024
-    progress.set_postfix({
-        'memory_usage': '{:.2f}'.format(memory_usage).zfill(5) + 'GB',
-        'execution_providers': roop.globals.execution_providers,
-        'execution_threads': roop.globals.execution_threads
-    })
-    progress.refresh()
-    progress.update(1)
-    logging.info(f'Updated progress: {progress.n}/{progress.total} frames processed. Memory usage: {memory_usage:.2f}GB')
+    try:
+        process = psutil.Process(os.getpid())
+        memory_usage = process.memory_info().rss / 1024 / 1024 / 1024
+        if progress:
+            progress.set_postfix({
+                'memory_usage': '{:.2f}'.format(memory_usage).zfill(5) + 'GB',
+                'execution_providers': roop.globals.execution_providers,
+                'execution_threads': roop.globals.execution_threads
+            })
+            progress.refresh()
+            progress.update(1)
+        logging.info(f'Updated progress: {progress.n}/{progress.total} frames processed. Memory usage: {memory_usage:.2f}GB')
+    except Exception as e:
+        logging.error(f'Error updating progress: {e}')