mirror of
https://github.com/YaoFANGUK/video-subtitle-remover.git
synced 2026-04-14 16:47:31 +08:00
f78e985e1c
添加任务列表组件并优化视频加载逻辑 支持可视化显示字幕区域 整理所有模型, 分别为STTN智能擦除, STTN字幕检测, LAMA, ProPainter, OpenCV 提高处理性能 新增CPU运行模式并优化多语言支持 修复Propainter模式部分视频报错 本次提交新增了CPU运行模式，适用于无GPU加速的场景。同时，优化了多语言支持，新增了日语、韩语、越南语等语言配置文件，并更新了README文档以反映新的运行模式和多语言支持。此外，修复了部分代码逻辑，提升了系统的稳定性和兼容性。
120 lines
4.3 KiB
Python
120 lines
4.3 KiB
Python
import traceback
|
||
import importlib.util
|
||
|
||
import torch
|
||
|
||
from backend.config import tr
|
||
|
||
class HardwareAccelerator:
|
||
|
||
# 类变量，用于存储单例实例
|
||
_instance = None
|
||
|
||
@classmethod
|
||
def instance(cls):
|
||
"""获取单例实例"""
|
||
if cls._instance is None:
|
||
cls._instance = HardwareAccelerator()
|
||
cls._instance.initialize()
|
||
return cls._instance
|
||
|
||
def __init__(self):
|
||
self.__cuda = False
|
||
self.__dml = False
|
||
self.__onnx_providers = []
|
||
self.__enabled = True
|
||
self.__device = None
|
||
|
||
def initialize(self):
|
||
self.check_directml_available()
|
||
self.check_cuda_available()
|
||
self.load_onnx_providers()
|
||
|
||
def check_directml_available(self):
|
||
self.__dml = importlib.util.find_spec("torch_directml")
|
||
|
||
def check_cuda_available(self):
|
||
self.__cuda = torch.cuda.is_available()
|
||
|
||
def load_onnx_providers(self):
|
||
try:
|
||
import onnxruntime as ort
|
||
available_providers = ort.get_available_providers()
|
||
for provider in available_providers:
|
||
if provider in [
|
||
"CPUExecutionProvider"
|
||
]:
|
||
continue
|
||
if provider not in [
|
||
"DmlExecutionProvider", # DirectML，适用于 Windows GPU
|
||
"ROCMExecutionProvider", # AMD ROCm
|
||
"MIGraphXExecutionProvider", # AMD MIGraphX
|
||
"VitisAIExecutionProvider", # AMD VitisAI，适用于 RyzenAI & Windows, 实测和DirectML性能似乎差不多
|
||
"OpenVINOExecutionProvider", # Intel GPU
|
||
"MetalExecutionProvider", # Apple macOS
|
||
"CoreMLExecutionProvider", # Apple macOS
|
||
"CUDAExecutionProvider", # Nvidia GPU
|
||
]:
|
||
print(tr['Main']['OnnxExectionProviderNotSupportedSkipped'].format(provider))
|
||
continue
|
||
print(tr['Main']['OnnxExecutionProviderDetected'].format(provider))
|
||
self.__onnx_providers.append(provider)
|
||
except ModuleNotFoundError as e:
|
||
print(tr['Main']['OnnxRuntimeNotInstall'])
|
||
|
||
def has_accelerator(self):
|
||
if not self.__enabled:
|
||
return False
|
||
return self.__cuda or self.__dml or len(self.__onnx_providers) > 0
|
||
|
||
@property
|
||
def accelerator_name(self):
|
||
if not self.__enabled:
|
||
return "CPU"
|
||
if self.__dml:
|
||
return "DirectML"
|
||
if self.__cuda:
|
||
return "GPU"
|
||
elif len(self.__onnx_providers) > 0:
|
||
return ", ".join(self.__onnx_providers)
|
||
else:
|
||
return "CPU"
|
||
|
||
@property
|
||
def onnx_providers(self):
|
||
if not self.__enabled:
|
||
return []
|
||
return self.__onnx_providers
|
||
|
||
def has_cuda(self):
|
||
if not self.__enabled:
|
||
return False
|
||
return self.__cuda
|
||
|
||
def set_enabled(self, enable):
|
||
self.__enabled = enable
|
||
|
||
@property
|
||
def device(self):
|
||
"""
|
||
onnxruntime-directml 1.21.1-1.22.0(往上未测试) 和 torch-directml 不能同时初始化, 会相互影响
|
||
提示site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 266, in run
|
||
return self._sess.run(output_names, input_feed, run_options)
|
||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||
UnicodeDecodeError: 'utf-8' codec can't decode byte 0xb2 in position 344: invalid start bn 344: invalid start byte
|
||
onnxruntime-directml 1.21.1 则正常, 但Win10跑不起来, Win11正常
|
||
为了避免冲突以及避免重写一个QPT智能部署流程, 这里采用延迟初始化的方式+继续使用onnxruntime-directml 1.20.1
|
||
当然SubtitleDetect放到一个独立进程去操作也是可以的
|
||
"""
|
||
if self.__enabled:
|
||
if self.__dml:
|
||
try:
|
||
import torch_directml
|
||
return torch_directml.device(torch_directml.default_device())
|
||
self.__dml = True
|
||
except:
|
||
traceback.print_exc()
|
||
self.__dml = False
|
||
if self.__cuda:
|
||
return torch.device("cuda:0")
|
||
return torch.device("cpu") |