ar08 commited on
Commit
7a48a20
·
verified ·
1 Parent(s): 3aff7a7

Update roop/processors/frame/core.py

Browse files
Files changed (1) hide show
  1. roop/processors/frame/core.py +18 -11
roop/processors/frame/core.py CHANGED
@@ -5,10 +5,14 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
5
  from queue import Queue
6
  from types import ModuleType
7
  from typing import Any, List, Callable
 
8
  from tqdm import tqdm
9
 
10
  import roop
11
 
 
 
 
12
  FRAME_PROCESSORS_MODULES: List[ModuleType] = []
13
  FRAME_PROCESSORS_INTERFACE = [
14
  'pre_check',
@@ -20,18 +24,18 @@ FRAME_PROCESSORS_INTERFACE = [
20
  'post_process'
21
  ]
22
 
23
-
24
  def load_frame_processor_module(frame_processor: str) -> Any:
25
  try:
26
  frame_processor_module = importlib.import_module(f'roop.processors.frame.{frame_processor}')
27
  for method_name in FRAME_PROCESSORS_INTERFACE:
28
  if not hasattr(frame_processor_module, method_name):
29
  raise NotImplementedError
30
- except (ImportError, NotImplementedError):
 
 
31
  quit(f'Frame processor {frame_processor} crashed.')
32
  return frame_processor_module
33
 
34
-
35
  def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType]:
36
  global FRAME_PROCESSORS_MODULES
37
 
@@ -41,7 +45,6 @@ def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType
41
  FRAME_PROCESSORS_MODULES.append(frame_processor_module)
42
  return FRAME_PROCESSORS_MODULES
43
 
44
-
45
  def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None], update: Callable[[], None]) -> None:
46
  with ThreadPoolExecutor(max_workers=roop.globals.execution_threads) as executor:
47
  futures = []
@@ -50,32 +53,35 @@ def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_f
50
  while not queue.empty():
51
  future = executor.submit(process_frames, source_path, pick_queue(queue, queue_per_future), update)
52
  futures.append(future)
 
53
  for future in as_completed(futures):
54
- future.result()
55
-
 
 
 
56
 
57
  def create_queue(temp_frame_paths: List[str]) -> Queue[str]:
58
  queue: Queue[str] = Queue()
59
  for frame_path in temp_frame_paths:
60
  queue.put(frame_path)
 
61
  return queue
62
 
63
-
64
  def pick_queue(queue: Queue[str], queue_per_future: int) -> List[str]:
65
  queues = []
66
  for _ in range(queue_per_future):
67
  if not queue.empty():
68
  queues.append(queue.get())
 
69
  return queues
70
 
71
-
72
- def process_video(source_path: str, frame_paths: list[str], process_frames: Callable[[str, List[str], Any], None]) -> None:
73
  progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
74
  total = len(frame_paths)
75
- with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress:
76
  multi_process_frame(source_path, frame_paths, process_frames, lambda: update_progress(progress))
77
 
78
-
79
  def update_progress(progress: Any = None) -> None:
80
  process = psutil.Process(os.getpid())
81
  memory_usage = process.memory_info().rss / 1024 / 1024 / 1024
@@ -86,3 +92,4 @@ def update_progress(progress: Any = None) -> None:
86
  })
87
  progress.refresh()
88
  progress.update(1)
 
 
5
  from queue import Queue
6
  from types import ModuleType
7
  from typing import Any, List, Callable
8
+ import logging
9
  from tqdm import tqdm
10
 
11
  import roop
12
 
13
+ # Initialize logging
14
+ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
15
+
16
  FRAME_PROCESSORS_MODULES: List[ModuleType] = []
17
  FRAME_PROCESSORS_INTERFACE = [
18
  'pre_check',
 
24
  'post_process'
25
  ]
26
 
 
27
  def load_frame_processor_module(frame_processor: str) -> Any:
28
  try:
29
  frame_processor_module = importlib.import_module(f'roop.processors.frame.{frame_processor}')
30
  for method_name in FRAME_PROCESSORS_INTERFACE:
31
  if not hasattr(frame_processor_module, method_name):
32
  raise NotImplementedError
33
+ logging.info(f'Successfully loaded frame processor module: {frame_processor}')
34
+ except (ImportError, NotImplementedError) as e:
35
+ logging.error(f'Error loading frame processor {frame_processor}: {e}')
36
  quit(f'Frame processor {frame_processor} crashed.')
37
  return frame_processor_module
38
 
 
39
  def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType]:
40
  global FRAME_PROCESSORS_MODULES
41
 
 
45
  FRAME_PROCESSORS_MODULES.append(frame_processor_module)
46
  return FRAME_PROCESSORS_MODULES
47
 
 
48
  def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None], update: Callable[[], None]) -> None:
49
  with ThreadPoolExecutor(max_workers=roop.globals.execution_threads) as executor:
50
  futures = []
 
53
  while not queue.empty():
54
  future = executor.submit(process_frames, source_path, pick_queue(queue, queue_per_future), update)
55
  futures.append(future)
56
+ logging.info(f'Submitted future for processing frames.')
57
  for future in as_completed(futures):
58
+ try:
59
+ future.result()
60
+ logging.info('Frame processing completed for a future.')
61
+ except Exception as e:
62
+ logging.error(f'Error in processing frame: {e}')
63
 
64
  def create_queue(temp_frame_paths: List[str]) -> Queue[str]:
65
  queue: Queue[str] = Queue()
66
  for frame_path in temp_frame_paths:
67
  queue.put(frame_path)
68
+ logging.info('Queue created with frame paths.')
69
  return queue
70
 
 
71
  def pick_queue(queue: Queue[str], queue_per_future: int) -> List[str]:
72
  queues = []
73
  for _ in range(queue_per_future):
74
  if not queue.empty():
75
  queues.append(queue.get())
76
+ logging.info(f'Picked {len(queues)} items from queue for processing.')
77
  return queues
78
 
79
+ def process_video(source_path: str, frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None]) -> None:
 
80
  progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
81
  total = len(frame_paths)
82
+ with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format, mininterval=0.1) as progress:
83
  multi_process_frame(source_path, frame_paths, process_frames, lambda: update_progress(progress))
84
 
 
85
  def update_progress(progress: Any = None) -> None:
86
  process = psutil.Process(os.getpid())
87
  memory_usage = process.memory_info().rss / 1024 / 1024 / 1024
 
92
  })
93
  progress.refresh()
94
  progress.update(1)
95
+ logging.info(f'Updated progress: {progress.n}/{progress.total} frames processed. Memory usage: {memory_usage:.2f}GB')