Merge 890beb0eae into 181144ce33

modules/processors/frame/core.py

@@ -42,18 +42,29 @@ def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType
 
 def set_frame_processors_modules_from_ui(frame_processors: List[str]) -> None:
     global FRAME_PROCESSORS_MODULES
+    current_processor_names = [proc.__name__.split('.')[-1] for proc in FRAME_PROCESSORS_MODULES]
+
     for frame_processor, state in modules.globals.fp_ui.items():
-        if state == True and frame_processor not in frame_processors:
+        if state == True and frame_processor not in current_processor_names:
-            frame_processor_module = load_frame_processor_module(frame_processor)
-            FRAME_PROCESSORS_MODULES.append(frame_processor_module)
-            modules.globals.frame_processors.append(frame_processor)
-        if state == False:
             try:
                 frame_processor_module = load_frame_processor_module(frame_processor)
-                FRAME_PROCESSORS_MODULES.remove(frame_processor_module)
+                FRAME_PROCESSORS_MODULES.append(frame_processor_module)
+                if frame_processor not in modules.globals.frame_processors:
+                     modules.globals.frame_processors.append(frame_processor)
+            except SystemExit:
+                 print(f"Warning: Failed to load frame processor {frame_processor} requested by UI state.")
+            except Exception as e:
+                 print(f"Warning: Error loading frame processor {frame_processor} requested by UI state: {e}")
+
+        elif state == False and frame_processor in current_processor_names:
+            try:
+                module_to_remove = next((mod for mod in FRAME_PROCESSORS_MODULES if mod.__name__.endswith(f'.{frame_processor}')), None)
+                if module_to_remove:
+                    FRAME_PROCESSORS_MODULES.remove(module_to_remove)
+                if frame_processor in modules.globals.frame_processors:
                     modules.globals.frame_processors.remove(frame_processor)
-            except:
+            except Exception as e:
-                pass
+                 print(f"Warning: Error removing frame processor {frame_processor}: {e}")
 
 def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None], progress: Any = None) -> None:
     with ThreadPoolExecutor(max_workers=modules.globals.execution_threads) as executor:

modules/processors/frame/face_enhancer.py

@@ -48,6 +48,17 @@ def pre_start() -> bool:
     return True
 
 
+TENSORRT_AVAILABLE = False
+try:
+    import torch_tensorrt
+    TENSORRT_AVAILABLE = True
+except ImportError as im:
+    print(f"TensorRT is not available: {im}")
+    pass
+except Exception as e:
+    print(f"TensorRT is not available: {e}")
+    pass
+
 def get_face_enhancer() -> Any:
     global FACE_ENHANCER
 
@@ -55,16 +66,26 @@ def get_face_enhancer() -> Any:
         if FACE_ENHANCER is None:
             model_path = os.path.join(models_dir, "GFPGANv1.4.pth")
             
-            match platform.system():
+            selected_device = None
-                case "Darwin":  # Mac OS
+            device_priority = []
-                    if torch.backends.mps.is_available():
-                        mps_device = torch.device("mps")
-                        FACE_ENHANCER = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=mps_device)  # type: ignore[attr-defined]
-                    else:
-                        FACE_ENHANCER = gfpgan.GFPGANer(model_path=model_path, upscale=1)  # type: ignore[attr-defined]
-                case _:  # Other OS
-                    FACE_ENHANCER = gfpgan.GFPGANer(model_path=model_path, upscale=1)  # type: ignore[attr-defined]
 
+            if TENSORRT_AVAILABLE and torch.cuda.is_available():
+                selected_device = torch.device("cuda")
+                device_priority.append("TensorRT+CUDA")
+            elif torch.cuda.is_available():
+                selected_device = torch.device("cuda")
+                device_priority.append("CUDA")
+            elif torch.backends.mps.is_available() and platform.system() == "Darwin":
+                selected_device = torch.device("mps")
+                device_priority.append("MPS")
+            elif not torch.cuda.is_available():
+                selected_device = torch.device("cpu")
+                device_priority.append("CPU")
+            
+            FACE_ENHANCER = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=selected_device)
+
+            # for debug:
+            print(f"Selected device: {selected_device} and device priority: {device_priority}")
     return FACE_ENHANCER