Updates for macOS and coreML / Metal

pull/286/head
Jason Kneen 2024-08-13 13:08:06 +01:00
parent fb695b0208
commit 359848a54e
5 changed files with 139 additions and 111 deletions

View File

@ -5,6 +5,8 @@ if any(arg.startswith('--execution-provider') for arg in sys.argv):
os.environ['OMP_NUM_THREADS'] = '1' os.environ['OMP_NUM_THREADS'] = '1'
# reduce tensorflow log level # reduce tensorflow log level
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
# Force TensorFlow to use Metal
os.environ['TENSORFLOW_METAL'] = '1'
import warnings import warnings
from typing import List from typing import List
import platform import platform
@ -35,23 +37,17 @@ def parse_args() -> None:
program.add_argument('-t', '--target', help='select an target image or video', dest='target_path') program.add_argument('-t', '--target', help='select an target image or video', dest='target_path')
program.add_argument('-o', '--output', help='select output file or directory', dest='output_path') program.add_argument('-o', '--output', help='select output file or directory', dest='output_path')
program.add_argument('--frame-processor', help='pipeline of frame processors', dest='frame_processor', default=['face_swapper'], choices=['face_swapper', 'face_enhancer'], nargs='+') program.add_argument('--frame-processor', help='pipeline of frame processors', dest='frame_processor', default=['face_swapper'], choices=['face_swapper', 'face_enhancer'], nargs='+')
program.add_argument('--keep-fps', help='keep original fps', dest='keep_fps', action='store_true', default=False) program.add_argument('--keep-fps', help='keep original fps', dest='keep_fps', action='store_true', default=True)
program.add_argument('--keep-audio', help='keep original audio', dest='keep_audio', action='store_true', default=True) program.add_argument('--keep-audio', help='keep original audio', dest='keep_audio', action='store_true', default=True)
program.add_argument('--keep-frames', help='keep temporary frames', dest='keep_frames', action='store_true', default=False) program.add_argument('--keep-frames', help='keep temporary frames', dest='keep_frames', action='store_true', default=True)
program.add_argument('--many-faces', help='process every face', dest='many_faces', action='store_true', default=False) program.add_argument('--many-faces', help='process every face', dest='many_faces', action='store_true', default=False)
program.add_argument('--video-encoder', help='adjust output video encoder', dest='video_encoder', default='libx264', choices=['libx264', 'libx265', 'libvpx-vp9']) program.add_argument('--video-encoder', help='adjust output video encoder', dest='video_encoder', default='libx264', choices=['libx264', 'libx265', 'libvpx-vp9'])
program.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=18, choices=range(52), metavar='[0-51]') program.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=50, choices=range(52), metavar='[0-51]')
program.add_argument('--max-memory', help='maximum amount of RAM in GB', dest='max_memory', type=int, default=suggest_max_memory()) program.add_argument('--max-memory', help='maximum amount of RAM in GB', dest='max_memory', type=int, default=suggest_max_memory())
program.add_argument('--execution-provider', help='execution provider', dest='execution_provider', default=['cpu'], choices=suggest_execution_providers(), nargs='+') program.add_argument('--execution-provider', help='execution provider', dest='execution_provider', default=['coreml'], 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-threads', help='number of execution threads', dest='execution_threads', type=int, default=suggest_execution_threads())
program.add_argument('-v', '--version', action='version', version=f'{modules.metadata.name} {modules.metadata.version}') program.add_argument('-v', '--version', action='version', version=f'{modules.metadata.name} {modules.metadata.version}')
# register deprecated args
program.add_argument('-f', '--face', help=argparse.SUPPRESS, dest='source_path_deprecated')
program.add_argument('--cpu-cores', help=argparse.SUPPRESS, dest='cpu_cores_deprecated', type=int)
program.add_argument('--gpu-vendor', help=argparse.SUPPRESS, dest='gpu_vendor_deprecated')
program.add_argument('--gpu-threads', help=argparse.SUPPRESS, dest='gpu_threads_deprecated', type=int)
args = program.parse_args() args = program.parse_args()
modules.globals.source_path = args.source_path modules.globals.source_path = args.source_path
@ -66,10 +62,9 @@ def parse_args() -> None:
modules.globals.video_encoder = args.video_encoder modules.globals.video_encoder = args.video_encoder
modules.globals.video_quality = args.video_quality modules.globals.video_quality = args.video_quality
modules.globals.max_memory = args.max_memory modules.globals.max_memory = args.max_memory
modules.globals.execution_providers = decode_execution_providers(args.execution_provider) modules.globals.execution_providers = ['CoreMLExecutionProvider'] # Force CoreML
modules.globals.execution_threads = args.execution_threads modules.globals.execution_threads = args.execution_threads
#for ENHANCER tumbler:
if 'face_enhancer' in args.frame_processor: if 'face_enhancer' in args.frame_processor:
modules.globals.fp_ui['face_enhancer'] = True modules.globals.fp_ui['face_enhancer'] = True
else: else:
@ -77,36 +72,6 @@ def parse_args() -> None:
modules.globals.nsfw = False modules.globals.nsfw = False
# translate deprecated args
if args.source_path_deprecated:
print('\033[33mArgument -f and --face are deprecated. Use -s and --source instead.\033[0m')
modules.globals.source_path = args.source_path_deprecated
modules.globals.output_path = normalize_output_path(args.source_path_deprecated, modules.globals.target_path, args.output_path)
if args.cpu_cores_deprecated:
print('\033[33mArgument --cpu-cores is deprecated. Use --execution-threads instead.\033[0m')
modules.globals.execution_threads = args.cpu_cores_deprecated
if args.gpu_vendor_deprecated == 'apple':
print('\033[33mArgument --gpu-vendor apple is deprecated. Use --execution-provider coreml instead.\033[0m')
modules.globals.execution_providers = decode_execution_providers(['coreml'])
if args.gpu_vendor_deprecated == 'nvidia':
print('\033[33mArgument --gpu-vendor nvidia is deprecated. Use --execution-provider cuda instead.\033[0m')
modules.globals.execution_providers = decode_execution_providers(['cuda'])
if args.gpu_vendor_deprecated == 'amd':
print('\033[33mArgument --gpu-vendor amd is deprecated. Use --execution-provider cuda instead.\033[0m')
modules.globals.execution_providers = decode_execution_providers(['rocm'])
if args.gpu_threads_deprecated:
print('\033[33mArgument --gpu-threads is deprecated. Use --execution-threads instead.\033[0m')
modules.globals.execution_threads = args.gpu_threads_deprecated
def encode_execution_providers(execution_providers: List[str]) -> List[str]:
return [execution_provider.replace('ExecutionProvider', '').lower() for execution_provider in execution_providers]
def decode_execution_providers(execution_providers: List[str]) -> List[str]:
return [provider for provider, encoded_execution_provider in zip(onnxruntime.get_available_providers(), encode_execution_providers(onnxruntime.get_available_providers()))
if any(execution_provider in encoded_execution_provider for execution_provider in execution_providers)]
def suggest_max_memory() -> int: def suggest_max_memory() -> int:
if platform.system().lower() == 'darwin': if platform.system().lower() == 'darwin':
@ -115,39 +80,22 @@ def suggest_max_memory() -> int:
def suggest_execution_providers() -> List[str]: def suggest_execution_providers() -> List[str]:
return encode_execution_providers(onnxruntime.get_available_providers()) return ['coreml'] # Only suggest CoreML
def suggest_execution_threads() -> int: def suggest_execution_threads() -> int:
if 'DmlExecutionProvider' in modules.globals.execution_providers:
return 1
if 'ROCMExecutionProvider' in modules.globals.execution_providers:
return 1
return 8 return 8
def limit_resources() -> None: def limit_resources() -> None:
# prevent tensorflow memory leak
gpus = tensorflow.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
tensorflow.config.experimental.set_memory_growth(gpu, True)
# limit memory usage
if modules.globals.max_memory: if modules.globals.max_memory:
memory = modules.globals.max_memory * 1024 ** 3
if platform.system().lower() == 'darwin':
memory = modules.globals.max_memory * 1024 ** 6 memory = modules.globals.max_memory * 1024 ** 6
if platform.system().lower() == 'windows':
import ctypes
kernel32 = ctypes.windll.kernel32
kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(memory), ctypes.c_size_t(memory))
else:
import resource import resource
resource.setrlimit(resource.RLIMIT_DATA, (memory, memory)) resource.setrlimit(resource.RLIMIT_DATA, (memory, memory))
def release_resources() -> None: def release_resources() -> None:
if 'CUDAExecutionProvider' in modules.globals.execution_providers: pass # No need to release CUDA resources
torch.cuda.empty_cache()
def pre_check() -> bool: def pre_check() -> bool:
@ -170,8 +118,13 @@ def start() -> None:
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
if not frame_processor.pre_start(): if not frame_processor.pre_start():
return return
# process image to image
if has_image_extension(modules.globals.target_path): if has_image_extension(modules.globals.target_path):
process_image()
else:
process_video()
def process_image():
if modules.globals.nsfw == False: if modules.globals.nsfw == False:
from modules.predicter import predict_image from modules.predicter import predict_image
if predict_image(modules.globals.target_path): if predict_image(modules.globals.target_path):
@ -180,13 +133,13 @@ def start() -> None:
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
update_status('Progressing...', frame_processor.NAME) update_status('Progressing...', frame_processor.NAME)
frame_processor.process_image(modules.globals.source_path, modules.globals.output_path, modules.globals.output_path) frame_processor.process_image(modules.globals.source_path, modules.globals.output_path, modules.globals.output_path)
release_resources()
if is_image(modules.globals.target_path): if is_image(modules.globals.target_path):
update_status('Processing to image succeed!') update_status('Processing to image succeed!')
else: else:
update_status('Processing to image failed!') update_status('Processing to image failed!')
return
# process image to videos
def process_video():
if modules.globals.nsfw == False: if modules.globals.nsfw == False:
from modules.predicter import predict_video from modules.predicter import predict_video
if predict_video(modules.globals.target_path): if predict_video(modules.globals.target_path):
@ -199,8 +152,6 @@ def start() -> None:
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors): for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
update_status('Progressing...', frame_processor.NAME) update_status('Progressing...', frame_processor.NAME)
frame_processor.process_video(modules.globals.source_path, temp_frame_paths) frame_processor.process_video(modules.globals.source_path, temp_frame_paths)
release_resources()
# handles fps
if modules.globals.keep_fps: if modules.globals.keep_fps:
update_status('Detecting fps...') update_status('Detecting fps...')
fps = detect_fps(modules.globals.target_path) fps = detect_fps(modules.globals.target_path)
@ -209,7 +160,6 @@ def start() -> None:
else: else:
update_status('Creating video with 30.0 fps...') update_status('Creating video with 30.0 fps...')
create_video(modules.globals.target_path) create_video(modules.globals.target_path)
# handle audio
if modules.globals.keep_audio: if modules.globals.keep_audio:
if modules.globals.keep_fps: if modules.globals.keep_fps:
update_status('Restoring audio...') update_status('Restoring audio...')
@ -218,7 +168,6 @@ def start() -> None:
restore_audio(modules.globals.target_path, modules.globals.output_path) restore_audio(modules.globals.target_path, modules.globals.output_path)
else: else:
move_temp(modules.globals.target_path, modules.globals.output_path) move_temp(modules.globals.target_path, modules.globals.output_path)
# clean and validate
clean_temp(modules.globals.target_path) clean_temp(modules.globals.target_path)
if is_video(modules.globals.target_path): if is_video(modules.globals.target_path):
update_status('Processing to video succeed!') update_status('Processing to video succeed!')
@ -240,6 +189,69 @@ def run() -> None:
if not frame_processor.pre_check(): if not frame_processor.pre_check():
return return
limit_resources() limit_resources()
print(f"ONNX Runtime version: {onnxruntime.__version__}")
print(f"Available execution providers: {onnxruntime.get_available_providers()}")
print(f"Selected execution provider: CoreMLExecutionProvider")
# Configure ONNX Runtime to use only CoreML
onnxruntime.set_default_logger_severity(3) # Set to WARNING level
options = onnxruntime.SessionOptions()
options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
# Test CoreML with a dummy model
try:
import numpy as np
from onnx import helper, TensorProto
# Create a simple ONNX model
X = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 3, 224, 224])
Y = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 3, 224, 224])
node = helper.make_node('Identity', ['input'], ['output'])
graph = helper.make_graph([node], 'test_model', [X], [Y])
model = helper.make_model(graph)
# Save the model
model_path = 'test_model.onnx'
with open(model_path, 'wb') as f:
f.write(model.SerializeToString())
# Create a CoreML session
session = onnxruntime.InferenceSession(model_path, options, providers=['CoreMLExecutionProvider'])
# Run inference
input_data = np.random.rand(1, 3, 224, 224).astype(np.float32)
output = session.run(None, {'input': input_data})
print("CoreML init successful and being used")
print(f"Input shape: {input_data.shape}, Output shape: {output[0].shape}")
# Clean up
os.remove(model_path)
except Exception as e:
print(f"Error testing CoreML: {str(e)}")
print("The application may not be able to use GPU acceleration")
# Configure TensorFlow to use Metal
try:
tf_devices = tensorflow.config.list_physical_devices()
print("TensorFlow devices:", tf_devices)
if any('GPU' in device.name for device in tf_devices):
print("TensorFlow is using GPU (Metal)")
else:
print("TensorFlow is not using GPU")
except Exception as e:
print(f"Error configuring TensorFlow: {str(e)}")
# Configure PyTorch to use MPS (Metal Performance Shaders)
try:
if torch.backends.mps.is_available():
print("PyTorch is using MPS (Metal Performance Shaders)")
torch.set_default_device('mps')
else:
print("PyTorch MPS is not available")
except Exception as e:
print(f"Error configuring PyTorch: {str(e)}")
if modules.globals.headless: if modules.globals.headless:
start() start()
else: else:

View File

@ -17,7 +17,7 @@ NAME = 'DLC.FACE-SWAPPER'
def pre_check() -> bool: def pre_check() -> bool:
download_directory_path = resolve_relative_path('../models') download_directory_path = resolve_relative_path('../models')
conditional_download(download_directory_path, ['https://huggingface.co/hacksider/deep-live-cam/blob/main/inswapper_128_fp16.onnx']) conditional_download(download_directory_path, ['https://huggingface.co/hacksider/deep-live-cam/blob/main/inswapper_128.onnx'])
return True return True
@ -39,7 +39,7 @@ def get_face_swapper() -> Any:
with THREAD_LOCK: with THREAD_LOCK:
if FACE_SWAPPER is None: if FACE_SWAPPER is None:
model_path = resolve_relative_path('../models/inswapper_128_fp16.onnx') model_path = resolve_relative_path('../models/inswapper_128.onnx')
FACE_SWAPPER = insightface.model_zoo.get_model(model_path, providers=modules.globals.execution_providers) FACE_SWAPPER = insightface.model_zoo.get_model(model_path, providers=modules.globals.execution_providers)
return FACE_SWAPPER return FACE_SWAPPER

View File

@ -257,11 +257,11 @@ def webcam_preview():
global preview_label, PREVIEW global preview_label, PREVIEW
cap = cv2.VideoCapture(0) # Use index for the webcam (adjust the index accordingly if necessary) cap = cv2.VideoCapture(0) # Use index for the webcam (adjust the index accordingly if necessary)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 960) # Set the width of the resolution cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1024) # Set the width of the resolution
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 540) # Set the height of the resolution cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 768) # Set the height of the resolution
cap.set(cv2.CAP_PROP_FPS, 60) # Set the frame rate of the webcam cap.set(cv2.CAP_PROP_FPS, 60) # Set the frame rate of the webcam
PREVIEW_MAX_WIDTH = 960 PREVIEW_MAX_WIDTH = 1024
PREVIEW_MAX_HEIGHT = 540 PREVIEW_MAX_HEIGHT = 768
preview_label.configure(image=None) # Reset the preview image before startup preview_label.configure(image=None) # Reset the preview image before startup

View File

@ -9,6 +9,7 @@ import urllib
from pathlib import Path from pathlib import Path
from typing import List, Any from typing import List, Any
from tqdm import tqdm from tqdm import tqdm
import cv2
import modules.globals import modules.globals
@ -44,7 +45,19 @@ def detect_fps(target_path: str) -> float:
def extract_frames(target_path: str) -> None: def extract_frames(target_path: str) -> None:
temp_directory_path = get_temp_directory_path(target_path) temp_directory_path = get_temp_directory_path(target_path)
run_ffmpeg(['-i', target_path, '-pix_fmt', 'rgb24', os.path.join(temp_directory_path, '%04d.png')]) cap = cv2.VideoCapture(target_path)
frame_count = 0
while True:
ret, frame = cap.read()
if not ret:
break
# Save the frame
cv2.imwrite(os.path.join(temp_directory_path, f'{frame_count:04d}.png'), frame)
frame_count += 1
cap.release()
def create_video(target_path: str, fps: float = 30.0) -> None: def create_video(target_path: str, fps: float = 30.0) -> None:

View File

@ -1,23 +1,26 @@
--extra-index-url https://download.pytorch.org/whl/cu118 # Deep Live Cam requirements
numpy==1.23.5 # Core dependencies
numpy==1.26.4
onnxruntime-silicon==1.16.3
opencv-python==4.8.1.78 opencv-python==4.8.1.78
onnx==1.16.0
insightface==0.7.3
psutil==5.9.8
tk==0.1.0
customtkinter==5.2.2
pillow==9.5.0 pillow==9.5.0
torch==2.0.1+cu118; sys_platform != 'darwin' insightface==0.7.3
torch==2.0.1; sys_platform == 'darwin' torch==2.1.0 # Add the specific version you're using
torchvision==0.15.2+cu118; sys_platform != 'darwin' tensorflow==2.16.1 # Add the specific version you're using
torchvision==0.15.2; sys_platform == 'darwin'
onnxruntime==1.18.0; sys_platform == 'darwin' and platform_machine != 'arm64' # Image processing
onnxruntime-silicon==1.16.3; sys_platform == 'darwin' and platform_machine == 'arm64' scikit-image==0.24.0
onnxruntime-gpu==1.18.0; sys_platform != 'darwin' matplotlib==3.9.1.post1
tensorflow==2.13.0rc1; sys_platform == 'darwin'
tensorflow==2.12.1; sys_platform != 'darwin' # Machine learning
opennsfw2==0.10.2 scikit-learn==1.5.1
protobuf==4.23.2
# Utilities
tqdm==4.66.4 tqdm==4.66.4
gfpgan==1.3.8 requests==2.32.3
prettytable==3.11.0
# Optional dependencies (comment out if not needed)
# albumentations==1.4.13
# coloredlogs==15.0.1