repo_description: Powerful & Easy-to-Use Video Face Swapping and Editing Software
repo_topics: ai computer-vision deepfake face-editor faceswap live-portrait video-editor
code: import threading
  import os
  import subprocess as sp
  import gc
  import traceback
  from typing import Dict, TYPE_CHECKING
  
  from packaging import version
  import numpy as np
  import onnxruntime
  import torch
  import onnx
  from torchvision.transforms import v2
  from PySide6 import QtCore
  try:
      import tensorrt as trt
      TENSORRT_AVAILABLE = True
  except ModuleNotFoundError:
      print("No TensorRT Found")
      TENSORRT_AVAILABLE = False
  
  from app.processors.utils.engine_builder import onnx_to_trt as onnx2trt
  from app.processors.utils.tensorrt_predictor import TensorRTPredictor
  from app.processors.face_detectors import FaceDetectors
  from app.processors.face_landmark_detectors import FaceLandmarkDetectors
  from app.processors.face_masks import FaceMasks
  from app.processors.face_restorers import FaceRestorers
  from app.processors.face_swappers import FaceSwappers
  from app.processors.frame_enhancers import FrameEnhancers
  from app.processors.face_editors import FaceEditors
  from app.processors.utils.dfm_model import DFMModel
  from app.processors.models_data import models_list, arcface_mapping_model_dict, models_trt_list
  from app.helpers.miscellaneous import is_file_exists
  from app.helpers.downloader import download_file
  
  if TYPE_CHECKING:
      from app.ui.main_ui import MainWindow
  
  onnxruntime.set_default_logger_severity(4)
  onnxruntime.log_verbosity_level = -1
  lock = threading.Lock()
  
  class ModelsProcessor(QtCore.QObject):
      processing_complete = QtCore.Signal()
      model_loaded = QtCore.Signal()  # Signal emitted with Onnx InferenceSession
  
      def __init__(self, main_window: 'MainWindow', device='cuda'):
          super().__init__()
          self.main_window = main_window
          self.provider_name = 'TensorRT'
          self.device = device
          self.model_lock = threading.RLock()  # Reentrant lock for model access
          self.trt_ep_options = {
              # 'trt_max_workspace_size': 3 << 30,  # Dimensione massima dello spazio di lavoro in bytes
              'trt_engine_cache_enable': True,
              'trt_engine_cache_path': "tensorrt-engines",
              'trt_timing_cache_enable': True,
              'trt_timing_cache_path': "tensorrt-engines",
              'trt_dump_ep_context_model': True,
              'trt_ep_context_file_path': "tensorrt-engines",
              'trt_layer_norm_fp32_fallback': True,
              'trt_builder_optimization_level': 5,
          }
          self.providers = [
              ('CUDAExecutionProvider'),
              ('CPUExecutionProvider')
          ]       
          self.nThreads = 2
          self.syncvec = torch.empty((1, 1), dtype=torch.float32, device=self.device)
  
          # Initialize models and models_path
          self.models: Dict[str, onnxruntime.InferenceSession] = {}
          self.models_path = {}
          self.models_data = {}
          for model_data in models_list:
              model_name, model_path = model_data['model_name'], model_data['local_path']
              self.models[model_name] = None #Model Instance
              self.models_path[model_name] = model_path
              self.models_data[model_name] = {'local_path': model_data['local_path'], 'hash': model_data['hash'], 'url': model_data.get('url')}
  
          self.dfm_models: Dict[str, DFMModel] = {}
  
          if TENSORRT_AVAILABLE:
              # Initialize models_trt and models_trt_path
              self.models_trt = {}
              self.models_trt_path = {}
              for model_data in models_trt_list:
                  model_name, model_path = model_data['model_name'], model_data['local_path']
                  self.models_trt[model_name] = None #Model Instance
                  self.models_trt_path[model_name] = model_path
  
          self.face_detectors = FaceDetectors(self)
          self.face_landmark_detectors = FaceLandmarkDetectors(self)
          self.face_masks = FaceMasks(self)
          self.face_restorers = FaceRestorers(self)
          self.face_swappers = FaceSwappers(self)
          self.frame_enhancers = FrameEnhancers(self)
          self.face_editors = FaceEditors(self)
  
          self.clip_session = []
          self.arcface_dst = np.array( [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366], [41.5493, 92.3655], [70.7299, 92.2041]], dtype=np.float32)
          self.FFHQ_kps = np.array([[ 192.98138, 239.94708 ], [ 318.90277, 240.1936 ], [ 256.63416, 314.01935 ], [ 201.26117, 371.41043 ], [ 313.08905, 371.15118 ] ])
          self.mean_lmk = []
          self.anchors  = []
          self.emap = []
          self.LandmarksSubsetIdxs = [
              0, 1, 4, 5, 6, 7, 8, 10, 13, 14, 17, 21, 33, 37, 39,
              40, 46, 52, 53, 54, 55, 58, 61, 63, 65, 66, 67, 70, 78, 80,
              81, 82, 84, 87, 88, 91, 93, 95, 103, 105, 107, 109, 127, 132, 133,
              136, 144, 145, 146, 148, 149, 150, 152, 153, 154, 155, 157, 158, 159, 160,
              161, 162, 163, 168, 172, 173, 176, 178, 181, 185, 191, 195, 197, 234, 246,
              249, 251, 263, 267, 269, 270, 276, 282, 283, 284, 285, 288, 291, 293, 295,
              296, 297, 300, 308, 310, 311, 312, 314, 317, 318, 321, 323, 324, 332, 334,
              336, 338, 356, 361, 362, 365, 373, 374, 375, 377, 378, 379, 380, 381, 382,
              384, 385, 386, 387, 388, 389, 390, 397, 398, 400, 402, 405, 409, 415, 454,
              466, 468, 469, 470, 471, 472, 473, 474, 475, 476, 477
          ]
  
          self.normalize = v2.Normalize(mean = [ 0., 0., 0. ],
                                        std = [ 1/1.0, 1/1.0, 1/1.0 ])
          
          self.lp_mask_crop = self.face_editors.lp_mask_crop
          self.lp_lip_array = self.face_editors.lp_lip_array
  
      def load_model(self, model_name, session_options=None):
          with self.model_lock:
              self.main_window.model_loading_signal.emit()
              # QApplication.processEvents()
              # if not is_file_exists(self.models_path[model_name]):
              #     download_file(model_name, self.models_path[model_name], self.models_data[model_name]['hash'], self.models_data[model_name]['url'])
              if session_options is None:
                  model_instance = onnxruntime.InferenceSession(self.models_path[model_name], providers=self.providers)
              else:
                  model_instance = onnxruntime.InferenceSession(self.models_path[model_name], sess_options=session_options, providers=self.providers)
  
              # Check if another thread has already loaded an instance for this model, if yes then delete the current one and return that instead
              if self.models[model_name]:
                  del model_instance
                  gc.collect()
                  return self.models[model_name]
              self.main_window.model_loaded_signal.emit()
  
              return model_instance
  
      def load_dfm_model(self, dfm_model):
          with self.model_lock:
              if not self.dfm_models.get(dfm_model):
                  self.main_window.model_loading_signal.emit()
                  max_models_to_keep = self.main_window.control['MaxDFMModelsSlider']
                  total_loaded_models = len(self.dfm_models)
                  if total_loaded_models==max_models_to_keep:
                      print("Clearing DFM Model")
                      model_name, model_instance = list(self.dfm_models.items())[0]
                      del model_instance
                      self.dfm_models.pop(model_name)
                      gc.collect()
                  try:
                      self.dfm_models[dfm_model] = DFMModel(self.main_window.dfm_models_data[dfm_model], self.providers, self.device)
                  except:
                      traceback.print_exc()   
                      self.dfm_models[dfm_model] = None         
                  self.main_window.model_loaded_signal.emit()
              return self.dfm_models[dfm_model]
  
  
      def load_model_trt(self, model_name, custom_plugin_path=None, precision='fp16', debug=False):
          # self.showModelLoadingProgressBar()
          #time.sleep(0.5)
          self.main_window.model_loading_signal.emit()
  
          if not os.path.exists(self.models_trt_path[model_name]):
              onnx2trt(onnx_model_path=self.models_path[model_name],
                       trt_model_path=self.models_trt_path[model_name],
                       precision=precision,
                       custom_plugin_path=custom_plugin_path,
                       verbose=False
                      )
          model_instance = TensorRTPredictor(model_path=self.models_trt_path[model_name], custom_plugin_path=custom_plugin_path, pool_size=self.nThreads, device=self.device, debug=debug)
  
          self.main_window.model_loaded_signal.emit()
          return model_instance
  
      def delete_models(self):
          for model_name, model_instance in self.models.items():
              del model_instance
              self.models[model_name] = None
          self.clip_session = []
          gc.collect()
  
      def delete_models_trt(self):
          if TENSORRT_AVAILABLE:
              for model_data in models_trt_list:
                  model_name = model_data['model_name']
                  if isinstance(self.models_trt[model_name], TensorRTPredictor):
                      # È un'istanza di TensorRTPredictor
                      self.models_trt[model_name].cleanup()
                      del self.models_trt[model_name]
                      self.models_trt[model_name] = None #Model Instance
              gc.collect()
  
      def delete_models_dfm(self):
          keys_to_remove = []
          for model_name, model_instance in self.dfm_models.items():
              del model_instance
              keys_to_remove.append(model_name)
          
          for model_name in keys_to_remove:
              self.dfm_models.pop(model_name)
          
          self.clip_session = []
          gc.collect()
  
      def showModelLoadingProgressBar(self):
          self.main_window.model_load_dialog.show()
  
      def hideModelLoadProgressBar(self):
          if self.main_window.model_load_dialog:
              self.main_window.model_load_dialog.close()
  
      def switch_providers_priority(self, provider_name):
          match provider_name:
              case "TensorRT" | "TensorRT-Engine":
                  providers = [
                                  ('TensorrtExecutionProvider', self.trt_ep_options),
                                  ('CUDAExecutionProvider'),
                                  ('CPUExecutionProvider')
                              ]
                  self.device = 'cuda'
                  if version.parse(trt.__version__) < version.parse("10.2.0") and provider_name == "TensorRT-Engine":
                      print("TensorRT-Engine provider cannot be used when TensorRT version is lower than 10.2.0.")
                      provider_name = "TensorRT"
  
              case "CPU":
                  providers = [
                                  ('CPUExecutionProvider')
                              ]
                  self.device = 'cpu'
              case "CUDA":
                  providers = [
                                  ('CUDAExecutionProvider'),
                                  ('CPUExecutionProvider')
                              ]
                  self.device = 'cuda'
              #case _:
  
          self.providers = providers
          self.provider_name = provider_name
          self.lp_mask_crop = self.lp_mask_crop.to(self.device)
  
          return self.provider_name
  
      def set_number_of_threads(self, value):
          self.nThreads = value
          self.delete_models_trt()
  
      def get_gpu_memory(self):
          command = "nvidia-smi --query-gpu=memory.total --format=csv"
          memory_total_info = sp.check_output(command.split()).decode('ascii').split('\n')[:-1][1:]
          memory_total = [int(x.split()[0]) for i, x in enumerate(memory_total_info)]
  
          command = "nvidia-smi --query-gpu=memory.free --format=csv"
          memory_free_info = sp.check_output(command.split()).decode('ascii').split('\n')[:-1][1:]
          memory_free = [int(x.split()[0]) for i, x in enumerate(memory_free_info)]
  
          memory_used = memory_total[0] - memory_free[0]
  
          return memory_used, memory_total[0]
      
      def clear_gpu_memory(self):
          self.delete_models()
          self.delete_models_dfm()
          self.delete_models_trt()
          torch.cuda.empty_cache()
  
  
      def load_inswapper_iss_emap(self, model_name):
          with self.model_lock:
              if not self.models[model_name]:
                  self.main_window.model_loading_signal.emit()
                  graph = onnx.load(self.models_path[model_name]).graph
                  self.emap = onnx.numpy_helper.to_array(graph.initializer[-1])
                  self.main_window.model_loaded_signal.emit()
  
      def run_detect(self, img, detect_mode='RetinaFace', max_num=1, score=0.5, input_size=(512, 512), use_landmark_detection=False, landmark_detect_mode='203', landmark_score=0.5, from_points=False, rotation_angles=None):
          rotation_angles = rotation_angles or [0]
          return self.face_detectors.run_detect(img, detect_mode, max_num, score, input_size, use_landmark_detection, landmark_detect_mode, landmark_score, from_points, rotation_angles)
      
      def run_detect_landmark(self, img, bbox, det_kpss, detect_mode='203', score=0.5, from_points=False):
          return self.face_landmark_detectors.run_detect_landmark(img, bbox, det_kpss, detect_mode, score, from_points)
  
      def get_arcface_model(self, face_swapper_model): 
          if face_swapper_model in arcface_mapping_model_dict:
              return arcface_mapping_model_dict[face_swapper_model]
          else:
              raise ValueError(f"Face swapper model {face_swapper_model} not found.")
  
      def run_recognize_direct(self, img, kps, similarity_type='Opal', arcface_model='Inswapper128ArcFace'):
          return self.face_swappers.run_recognize_direct(img, kps, similarity_type, arcface_model)
  
      def calc_inswapper_latent(self, source_embedding):
          return self.face_swappers.calc_inswapper_latent(source_embedding)
  
      def run_inswapper(self, image, embedding, output):
          self.face_swappers.run_inswapper(image, embedding, output)
  
      def calc_swapper_latent_iss(self, source_embedding, version="A"):
          return self.face_swappers.calc_swapper_latent_iss(source_embedding, version)
  
      def run_iss_swapper(self, image, embedding, output, version="A"):
          self.face_swappers.run_iss_swapper(image, embedding, output, version)
  
      def calc_swapper_latent_simswap512(self, source_embedding):
          return self.face_swappers.calc_swapper_latent_simswap512(source_embedding)
  
      def run_swapper_simswap512(self, image, embedding, output):
          self.face_swappers.run_swapper_simswap512(image, embedding, output)
  
      def calc_swapper_latent_ghost(self, source_embedding):
          return self.face_swappers.calc_swapper_latent_ghost(source_embedding)
  
      def run_swapper_ghostface(self, image, embedding, output, swapper_model='GhostFace-v2'):
          self.face_swappers.run_swapper_ghostface(image, embedding, output, swapper_model)
  
      def calc_swapper_latent_cscs(self, source_embedding):
          return self.face_swappers.calc_swapper_latent_cscs(source_embedding)
  
      def run_swapper_cscs(self, image, embedding, output):
          self.face_swappers.run_swapper_cscs(image, embedding, output)
  
      def run_enhance_frame_tile_process(self, img, enhancer_type, tile_size=256, scale=1):
          return self.frame_enhancers.run_enhance_frame_tile_process(img, enhancer_type, tile_size, scale)
  
      def run_deoldify_artistic(self, image, output):
          return self.frame_enhancers.run_deoldify_artistic(image, output)
  
      def run_deoldify_stable(self, image, output):
          return self.frame_enhancers.run_deoldify_artistic(image, output)
      
      def run_deoldify_video(self, image, output):
          return self.frame_enhancers.run_deoldify_video(image, output)
      
      def run_ddcolor_artistic(self, image, output):
          return self.frame_enhancers.run_ddcolor_artistic(image, output)
  
      def run_ddcolor(self, tensor_gray_rgb, output_ab):
          return self.frame_enhancers.run_ddcolor(tensor_gray_rgb, output_ab)
  
      def run_occluder(self, image, output):
          self.face_masks.run_occluder(image, output)
  
      def run_dfl_xseg(self, image, output):
          self.face_masks.run_dfl_xseg(image, output)
  
      def run_faceparser(self, image, output):
          self.face_masks.run_faceparser(image, output)
  
      def run_CLIPs(self, img, CLIPText, CLIPAmount):
          return self.face_masks.run_CLIPs(img, CLIPText, CLIPAmount)
      
      def lp_motion_extractor(self, img, face_editor_type='Human-Face', **kwargs) -> dict:
          return self.face_editors.lp_motion_extractor(img, face_editor_type, **kwargs)
  
      def lp_appearance_feature_extractor(self, img, face_editor_type='Human-Face'):
          return self.face_editors.lp_appearance_feature_extractor(img, face_editor_type)
  
      def lp_retarget_eye(self, kp_source: torch.Tensor, eye_close_ratio: torch.Tensor, face_editor_type='Human-Face') -> torch.Tensor:
          return self.face_editors.lp_retarget_eye(kp_source, eye_close_ratio, face_editor_type)
  
      def lp_retarget_lip(self, kp_source: torch.Tensor, lip_close_ratio: torch.Tensor, face_editor_type='Human-Face') -> torch.Tensor:
          return self.face_editors.lp_retarget_lip(kp_source, lip_close_ratio, face_editor_type)
  
      def lp_stitch(self, kp_source: torch.Tensor, kp_driving: torch.Tensor, face_editor_type='Human-Face') -> torch.Tensor:
          return self.face_editors.lp_stitch(kp_source, kp_driving, face_editor_type)
  
      def lp_stitching(self, kp_source: torch.Tensor, kp_driving: torch.Tensor, face_editor_type='Human-Face') -> torch.Tensor:
          return self.face_editors.lp_stitching(kp_source, kp_driving, face_editor_type)
  
      def lp_warp_decode(self, feature_3d: torch.Tensor, kp_source: torch.Tensor, kp_driving: torch.Tensor, face_editor_type='Human-Face') -> torch.Tensor:
          return self.face_editors.lp_warp_decode(feature_3d, kp_source, kp_driving, face_editor_type)
  
      def findCosineDistance(self, vector1, vector2):
          vector1 = vector1.ravel()
          vector2 = vector2.ravel()
          cos_dist = 1 - np.dot(vector1, vector2)/(np.linalg.norm(vector1)*np.linalg.norm(vector2)) # 2..0
          return 100-cos_dist*50
  
      def apply_facerestorer(self, swapped_face_upscaled, restorer_det_type, restorer_type, restorer_blend, fidelity_weight, detect_score):
          return self.face_restorers.apply_facerestorer(swapped_face_upscaled, restorer_det_type, restorer_type, restorer_blend, fidelity_weight, detect_score)
  
      def apply_occlusion(self, img, amount):
          return self.face_masks.apply_occlusion(img, amount)
      
      def apply_dfl_xseg(self, img, amount):
          return self.face_masks.apply_dfl_xseg(img, amount)
      
      def apply_face_parser(self, img, parameters):
          return self.face_masks.apply_face_parser(img, parameters)
      
      def apply_face_makeup(self, img, parameters):
          return self.face_editors.apply_face_makeup(img, parameters)
      
      def restore_mouth(self, img_orig, img_swap, kpss_orig, blend_alpha=0.5, feather_radius=10, size_factor=0.5, radius_factor_x=1.0, radius_factor_y=1.0, x_offset=0, y_offset=0):
          return self.face_masks.restore_mouth(img_orig, img_swap, kpss_orig, blend_alpha, feather_radius, size_factor, radius_factor_x, radius_factor_y, x_offset, y_offset)
  
      def restore_eyes(self, img_orig, img_swap, kpss_orig, blend_alpha=0.5, feather_radius=10, size_factor=3.5, radius_factor_x=1.0, radius_factor_y=1.0, x_offset=0, y_offset=0, eye_spacing_offset=0):
          return self.face_masks.restore_eyes(img_orig, img_swap, kpss_orig, blend_alpha, feather_radius, size_factor, radius_factor_x, radius_factor_y, x_offset, y_offset, eye_spacing_offset)
  
      def apply_fake_diff(self, swapped_face, original_face, DiffAmount):
          return self.face_masks.apply_fake_diff(swapped_face, original_face, DiffAmount)
  