refactor: birefnet gpu (#130)

* fix app * docs: gpu
2026-09-21 12:43:43 +08:00 · 2024-09-16 02:34:36 +08:00
parent 341f8719dd
commit 02e812c11d
6 changed files with 55 additions and 22 deletions
@@ -155,14 +155,19 @@ python scripts/download_model.py --models all

 ## 5. GPU推理加速（可选）

-如需使用英伟达GPU加速推理，在确保你已经安装CUDA与cuDNN后，根据[文档](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#cuda-12x)找到对应的`onnxruntime-gpu`版本安装，如：
+在当前版本，可被英伟达GPU加速的模型为`birefnet-v1-lite`，并请确保你有16GB左右的显存。
+
+如需使用英伟达GPU加速推理，在确保你已经安装CUDA与cuDNN后，根据[onnxruntime-gpu文档](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#cuda-12x)找到对应的`onnxruntime-gpu`版本安装，以及根据[pytorch官网](https://pytorch.org/get-started/locally/)找到对应的`pytorch`版本安装。

 ```bash
-# CUDA 12.x, cuDNN 8
+# 假如你的电脑安装的是CUDA 12.x, cuDNN 8
+# 安装torch是可选的，如果你始终配置不好cuDNN，那么试试安装torch
 pip install onnxruntime-gpu==1.18.0
+pip install torch --index-url https://download.pytorch.org/whl/cu121
 ```

-完成后，调用如`birefnet-v1-lite`模型将会利用GPU加速推理。
+完成安装后，调用`birefnet-v1-lite`模型即可利用GPU加速推理。
+

 <br>

@@ -151,14 +151,18 @@ Store in the project's `hivision/creator/weights` directory:

 ## 5. GPU Inference Acceleration (Optional)

-If you need to use NVIDIA GPU for accelerated inference, ensure that you have installed CUDA and cuDNN, then find the corresponding `onnxruntime-gpu` version to install according to the [documentation](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#cuda-12x), for example:
+In the current version, the model that can be accelerated by NVIDIA GPUs is `birefnet-v1-lite`, and please ensure you have around 16GB of VRAM.
+
+If you want to use NVIDIA GPU acceleration for inference, after ensuring you have installed CUDA and cuDNN, find the corresponding `onnxruntime-gpu` version to install according to the [onnxruntime-gpu documentation](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#cuda-12x), and find the corresponding `pytorch` version to install according to the [pytorch official website](https://pytorch.org/get-started/locally/).

 ```bash
-# CUDA 12.x, cuDNN 8
+# If your computer is installed with CUDA 12.x and cuDNN 8
+# Installing torch is optional. If you can't configure cuDNN, try installing torch
 pip install onnxruntime-gpu==1.18.0
+pip install torch --index-url https://download.pytorch.org/whl/cu121
 ```

-After completing this, calling models like `birefnet-v1-lite` will utilize GPU acceleration for inference.
+After completing the installation, call the `birefnet-v1-lite` model to utilize GPU acceleration for inference.

 <br>

@@ -147,14 +147,18 @@ python scripts/download_model.py --models all

 ## 5. GPU推論の加速（オプション）

-NVIDIA GPUによる推論加速を使用する場合は、CUDAとcuDNNがインストールされていることを確認し、[文書](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#cuda-12x)に従って対応する`onnxruntime-gpu`バージョンをインストールします。例：
+現在のバージョンでは、NVIDIA GPUで加速可能なモデルは`birefnet-v1-lite`です。約16GBのVRAMが必要であることにご注意ください。
+
+NVIDIA GPUを使用して推論を加速したい場合は、CUDAとcuDNNがインストールされていることを確認した上で、[onnxruntime-gpuのドキュメント](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#cuda-12x)に従って適切な`onnxruntime-gpu`バージョンをインストールし、[PyTorchの公式サイト](https://pytorch.org/get-started/locally/)から適切な`pytorch`バージョンをインストールしてください。

 ```bash
-# CUDA 12.x, cuDNN 8
+# もしコンピュータにCUDA 12.xとcuDNN 8がインストールされている場合
+# torchのインストールは任意です。cuDNNが設定できない場合は、torchを試してみてください
 pip install onnxruntime-gpu==1.18.0
+pip install torch --index-url https://download.pytorch.org/whl/cu121
 ```

-完了後、`birefnet-v1-lite`モデルを呼び出すと、GPUによる推論加速が利用されます。
+インストールが完了したら、`birefnet-v1-lite`モデルを呼び出してGPU加速推論を利用します。

 <br>

@@ -147,14 +147,18 @@ python scripts/download_model.py --models all

 ## 5. GPU 추론 가속 (선택 사항)

-NVIDIA GPU를 통한 추론 가속을 사용하려면 CUDA와 cuDNN이 설치되어 있는지 확인하고, [문서](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#cuda-12x)에 따라 해당 `onnxruntime-gpu` 버전을 설치합니다. 예:
+현재 버전에서 NVIDIA GPU로 가속화할 수 있는 모델은 `birefnet-v1-lite`입니다. 약 16GB의 VRAM이 필요합니다.
+
+NVIDIA GPU를 사용하여 추론을 가속화하려면, CUDA와 cuDNN이 설치되어 있는지 확인한 후, [onnxruntime-gpu 문서](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#cuda-12x)에서 해당하는 `onnxruntime-gpu` 버전을 찾아 설치하고, [PyTorch 공식 웹사이트](https://pytorch.org/get-started/locally/)에서 해당하는 `pytorch` 버전을 찾아 설치하세요.

 ```bash
-# CUDA 12.x, cuDNN 8
+# 컴퓨터에 CUDA 12.x와 cuDNN 8이 설치되어 있는 경우
+# 설치 중 torch를 설치하는 것은 선택 사항입니다. cuDNN을 설정할 수 없는 경우 torch를 설치해 보세요.
 pip install onnxruntime-gpu==1.18.0
+pip install torch --index-url https://download.pytorch.org/whl/cu121
 ```

-완료 후, `birefnet-v1-lite` 모델을 호출하면 GPU에 의한 추론 가속이 이용됩니다.
+설치 완료 후, `birefnet-v1-lite` 모델을 호출하면 GPU에 의한 추론 가속이 이용됩니다.

 <br>

@@ -15,10 +15,17 @@ HUMAN_MATTING_MODELS_EXIST = [
    if file.endswith(".onnx") or file.endswith(".mnn")
 ]
 # 在HUMAN_MATTING_MODELS中的模型才会被加载到Gradio中显示
-HUMAN_MATTING_MODELS = [
+HUMAN_MATTING_MODELS_CHOICE = [
    model for model in HUMAN_MATTING_MODELS if model in HUMAN_MATTING_MODELS_EXIST
 ]

+if len(HUMAN_MATTING_MODELS_CHOICE) == 0:
+    raise ValueError(
+        "未找到任何存在的人像分割模型，请检查 hivision/creator/weights 目录下的文件"
+        + "\n"
+        + "No existing portrait segmentation model was found, please check the files in the hivision/creator/weights directory."
+    )
+
 FACE_DETECT_MODELS = ["face++ (联网Online API)", "mtcnn"]
 FACE_DETECT_MODELS_EXPAND = (
    ["retinaface-resnet50"]
@@ -29,7 +36,7 @@ FACE_DETECT_MODELS_EXPAND = (
    )
    else []
 )
-FACE_DETECT_MODELS += FACE_DETECT_MODELS_EXPAND
+FACE_DETECT_MODELS_CHOICE = FACE_DETECT_MODELS + FACE_DETECT_MODELS_EXPAND

 LANGUAGE = ["zh", "en", "ko", "ja"]

@@ -54,8 +61,8 @@ if __name__ == "__main__":
    demo = create_ui(
        processor,
        root_dir,
-        HUMAN_MATTING_MODELS_EXIST,
-        FACE_DETECT_MODELS,
+        HUMAN_MATTING_MODELS_CHOICE,
+        FACE_DETECT_MODELS_CHOICE,
        LANGUAGE,
    )
    demo.launch(
@@ -37,10 +37,9 @@ WEIGHTS = {
    ),
 }

-ONNX_DEVICE = (
-    "CUDAExecutionProvider"
-    if onnxruntime.get_device() == "GPU"
-    else "CPUExecutionProvider"
+ONNX_DEVICE = onnxruntime.get_device()
+ONNX_PROVIDER = (
+    "CUDAExecutionProvider" if ONNX_DEVICE == "GPU" else "CPUExecutionProvider"
 )

 HIVISION_MODNET_SESS = None
@@ -52,7 +51,7 @@ BIREFNET_V1_LITE_SESS = None
 def load_onnx_model(checkpoint_path, set_cpu=False):
    providers = (
        ["CUDAExecutionProvider", "CPUExecutionProvider"]
-        if ONNX_DEVICE == "CUDAExecutionProvider"
+        if ONNX_PROVIDER == "CUDAExecutionProvider"
        else ["CPUExecutionProvider"]
    )

@@ -365,7 +364,17 @@ def get_birefnet_portrait_matting(input_image, checkpoint_path, ref_size=512):

    if BIREFNET_V1_LITE_SESS is None:
        print("首次加载birefnet-v1-lite模型...")
-        BIREFNET_V1_LITE_SESS = load_onnx_model(checkpoint_path)
+        if ONNX_DEVICE == "GPU":
+            print("onnxruntime-gpu已安装，尝试使用CUDA加载模型")
+            try:
+                import torch
+            except ImportError:
+                print(
+                    "torch未安装，尝试直接使用onnxruntime-gpu加载模型，这需要配置好CUDA和cuDNN"
+                )
+            BIREFNET_V1_LITE_SESS = load_onnx_model(checkpoint_path)
+        else:
+            BIREFNET_V1_LITE_SESS = load_onnx_model(checkpoint_path, set_cpu=True)

    # 记录加载onnx模型的结束时间
    load_end_time = time()