Introduce dynamic execution queue count

README.md

@@ -47,9 +47,11 @@ python run.py [options]
 --output-video-encoder {libx264,libx265,libvpx-vp9,h264_nvenc,hevc_nvenc} encoder used for the output video
 --output-video-quality [0-100] quality used for the output video
 --max-memory MAX_MEMORY maximum amount of RAM in GB
---execution-providers {tensorrt,cuda,cpu} [{tensorrt,cuda,cpu} ...] list of available execution providers (choices: cpu, ...)
---execution-threads EXECUTION_THREADS number of execution threads
+--execution-providers {cpu} [{cpu} ...] list of available execution providers (choices: cpu, ...)
+--execution-thread-count EXECUTION_THREAD_COUNT number of execution threads
+--execution-queue-count EXECUTION_QUEUE_COUNT number of execution queries
 -v, --version show program's version number and exit
+
 ```
 
 

roop/core.py

@@ -1,12 +1,11 @@
 #!/usr/bin/env python3
 
 import os
-import sys
-# single thread doubles cuda performance - needs to be set before torch import
-if any(arg.startswith('--execution-providers') for arg in sys.argv):
- os.environ['OMP_NUM_THREADS'] = '1'
+# single thread doubles cuda performance
+os.environ['OMP_NUM_THREADS'] = '1'
 # reduce tensorflow log level
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+import sys
 import warnings
 from typing import List
 import platform
@@ -46,7 +45,8 @@ def parse_args() -> None:
  program.add_argument('--output-video-quality', help='quality used for the output video', dest='output_video_quality', type=int, default=35, choices=range(101), metavar='[0-100]')
  program.add_argument('--max-memory', help='maximum amount of RAM in GB', dest='max_memory', type=int)
  program.add_argument('--execution-providers', help='list of available execution providers (choices: cpu, ...)', dest='execution_providers', default=['cpu'], choices=suggest_execution_providers(), nargs='+')
- program.add_argument('--execution-threads', help='number of execution threads', dest='execution_threads', type=int, default=suggest_execution_threads())
+ program.add_argument('--execution-thread-count', help='number of execution threads', dest='execution_thread_count', type=int, default=suggest_execution_thread_count())
+ program.add_argument('--execution-queue-count', help='number of execution queries', dest='execution_queue_count', type=int, default=1)
  program.add_argument('-v', '--version', action='version', version=f'{roop.metadata.name} {roop.metadata.version}')
 
  args = program.parse_args()
@@ -70,7 +70,8 @@ def parse_args() -> None:
  roop.globals.output_video_quality = args.output_video_quality
  roop.globals.max_memory = args.max_memory
  roop.globals.execution_providers = decode_execution_providers(args.execution_providers)
- roop.globals.execution_threads = args.execution_threads
+ roop.globals.execution_thread_count = args.execution_thread_count
+ roop.globals.execution_queue_count = args.execution_queue_count
 
 
 def encode_execution_providers(execution_providers: List[str]) -> List[str]:
@@ -90,7 +91,7 @@ def suggest_ui_layouts() -> List[str]:
  return list_module_names('roop/uis/__layouts__')
 
 
-def suggest_execution_threads() -> int:
+def suggest_execution_thread_count() -> int:
  if 'CUDAExecutionProvider' in onnxruntime.get_available_providers():
  return 8
  return 1

roop/globals.py

@@ -19,5 +19,6 @@
 output_video_quality: Optional[int] = None
 max_memory: Optional[int] = None
 execution_providers: List[str] = []
-execution_threads: Optional[int] = None
+execution_thread_count: Optional[int] = None
+execution_queue_count: Optional[int] = None
 log_level: str = 'error'

roop/processors/frame/core.py

@@ -57,10 +57,10 @@ def list_frame_processors_names() -> Optional[List[str]]:
 
 
 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:
+ with ThreadPoolExecutor(max_workers=roop.globals.execution_thread_count) as executor:
  futures = []
  queue = create_queue(temp_frame_paths)
- queue_per_future = max(len(temp_frame_paths) // roop.globals.execution_threads, 1)
+ queue_per_future = max(len(temp_frame_paths) // roop.globals.execution_thread_count, 1) * roop.globals.execution_queue_count
  while not queue.empty():
  future = executor.submit(process_frames, source_path, pick_queue(queue, queue_per_future), update)
  futures.append(future)
@@ -96,7 +96,8 @@ def update_progress(progress: Any = None) -> None:
  progress.set_postfix({
  'memory_usage': '{:.2f}'.format(memory_usage).zfill(5) + 'GB',
  'execution_providers': roop.globals.execution_providers,
- 'execution_threads': roop.globals.execution_threads
+ 'execution_thread_count': roop.globals.execution_thread_count,
+ 'execution_queue_count': roop.globals.execution_queue_count
  })
  progress.refresh()
  progress.update(1)

roop/uis/__components__/settings.py

@@ -8,7 +8,8 @@
 
 FRAME_PROCESSORS_CHECKBOX_GROUP: Optional[gradio.CheckboxGroup] = None
 EXECUTION_PROVIDERS_CHECKBOX_GROUP: Optional[gradio.CheckboxGroup] = None
-EXECUTION_THREADS_SLIDER: Optional[gradio.Slider] = None
+EXECUTION_THREAD_COUNT_SLIDER: Optional[gradio.Slider] = None
+EXECUTION_QUEUE_COUNT_SLIDER: Optional[gradio.Slider] = None
 KEEP_FPS_CHECKBOX: Optional[gradio.Checkbox] = None
 KEEP_TEMP_CHECKBOX: Optional[gradio.Checkbox] = None
 SKIP_AUDIO_CHECKBOX: Optional[gradio.Checkbox] = None
@@ -18,7 +19,8 @@
 def render() -> None:
  global FRAME_PROCESSORS_CHECKBOX_GROUP
  global EXECUTION_PROVIDERS_CHECKBOX_GROUP
- global EXECUTION_THREADS_SLIDER
+ global EXECUTION_THREAD_COUNT_SLIDER
+ global EXECUTION_QUEUE_COUNT_SLIDER
  global KEEP_FPS_CHECKBOX
  global KEEP_TEMP_CHECKBOX
  global SKIP_AUDIO_CHECKBOX
@@ -38,12 +40,19 @@ def render() -> None:
  choices=onnxruntime.get_available_providers(),
  value=roop.globals.execution_providers
  )
- EXECUTION_THREADS_SLIDER = gradio.Slider(
- label='execution_threads',
- value=roop.globals.execution_threads,
+ EXECUTION_THREAD_COUNT_SLIDER = gradio.Slider(
+ label='execution_thread_count',
+ value=roop.globals.execution_thread_count,
  step=1,
  minimum=1,
- maximum=64
+ maximum=128
+ )
+ EXECUTION_QUEUE_COUNT_SLIDER = gradio.Slider(
+ label='execution_queue_count',
+ value=roop.globals.execution_queue_count,
+ step=1,
+ minimum=1,
+ maximum=16
  )
  with gradio.Box():
  KEEP_FPS_CHECKBOX = gradio.Checkbox(
@@ -68,7 +77,8 @@ def render() -> None:
 def listen() -> None:
  FRAME_PROCESSORS_CHECKBOX_GROUP.change(update_frame_processors, inputs=FRAME_PROCESSORS_CHECKBOX_GROUP, outputs=FRAME_PROCESSORS_CHECKBOX_GROUP)
  EXECUTION_PROVIDERS_CHECKBOX_GROUP.change(update_execution_providers, inputs=EXECUTION_PROVIDERS_CHECKBOX_GROUP, outputs=EXECUTION_PROVIDERS_CHECKBOX_GROUP)
- EXECUTION_THREADS_SLIDER.change(update_execution_threads, inputs=EXECUTION_THREADS_SLIDER, outputs=EXECUTION_THREADS_SLIDER)
+ EXECUTION_THREAD_COUNT_SLIDER.change(update_execution_thread_count, inputs=EXECUTION_THREAD_COUNT_SLIDER, outputs=EXECUTION_THREAD_COUNT_SLIDER)
+ EXECUTION_QUEUE_COUNT_SLIDER.change(update_execution_queue_count, inputs=EXECUTION_QUEUE_COUNT_SLIDER, outputs=EXECUTION_QUEUE_COUNT_SLIDER)
  KEEP_FPS_CHECKBOX.change(lambda value: update_checkbox('keep_fps', value), inputs=KEEP_FPS_CHECKBOX, outputs=KEEP_FPS_CHECKBOX)
  KEEP_TEMP_CHECKBOX.change(lambda value: update_checkbox('keep_temp', value), inputs=KEEP_TEMP_CHECKBOX, outputs=KEEP_TEMP_CHECKBOX)
  SKIP_AUDIO_CHECKBOX.change(lambda value: update_checkbox('skip_audio', value), inputs=SKIP_AUDIO_CHECKBOX, outputs=SKIP_AUDIO_CHECKBOX)
@@ -94,9 +104,13 @@ def update_execution_providers(execution_providers: List[str]) -> Dict[Any, Any]
  return gradio.update(value=execution_providers)
 
 
-def update_execution_threads(execution_threads: int = 1) -> Dict[Any, Any]:
- roop.globals.execution_threads = execution_threads
- return gradio.update(value=execution_threads)
+def update_execution_thread_count(execution_thread_count: int = 1) -> Dict[Any, Any]:
+ roop.globals.execution_thread_count = execution_thread_count
+ return gradio.update(value=execution_thread_count)
+
+def update_execution_queue_count(execution_queue_count: int = 1) -> Dict[Any, Any]:
+ roop.globals.execution_queue_count = execution_queue_count
+ return gradio.update(value=execution_queue_count)
 
 
 def update_checkbox(name: str, value: bool) -> Dict[Any, Any]:

roop/uis/__layouts__/default.py

@@ -6,12 +6,12 @@
 def render() -> gradio.Blocks:
  with gradio.Blocks() as layout:
  with gradio.Row():
- with gradio.Column(scale=1):
+ with gradio.Column(scale=2):
  settings.render()
  with gradio.Column(scale=1):
  source.render()
  target.render()
- with gradio.Column(scale=2):
+ with gradio.Column(scale=3):
  preview.render()
  reference.render()
  with gradio.Row():