diff --git a/README-FAQ.md b/README-FAQ.md
index 2661433e..35987248 100644
--- a/README-FAQ.md
+++ b/README-FAQ.md
@@ -33,3 +33,5 @@ A: Use [**VBench**](https://github.com/Vchitect/VBench?tab=readme-ov-file#usage)
Use [**VBench-Long**](https://github.com/Vchitect/VBench/tree/master/vbench2_beta_long) for evaluating videos 5.0 seconds or longer (≥ 5.0s).
Each benchmark is optimized for its respective video length to ensure fair and consistent evaluation.
+**Q: Does VBench support running on Huawei Ascend NPU?**
+A: Yes. [Ascend Extension for PyTorch](https://github.com/Ascend/pytorch) develops `torch_npu` to adapt Ascend NPU to PyTorch so that developers who use the PyTorch can obtain powerful compute capabilities of Ascend AI Processors. Please install `torch` and `torch_npu` before installing vbench if you are using Ascend NPU devices. For more details, please refer to the official [installation guide](https://github.com/Ascend/pytorch#installation).
diff --git a/README.md b/README.md
index 738f1ff9..95e13fb2 100755
--- a/README.md
+++ b/README.md
@@ -246,7 +246,7 @@ vbench evaluate --videos_path "sampled_videos/lavie/human_action" --dimension "h
##### Python
```python
from vbench import VBench
-my_VBench = VBench(device, , )
+my_VBench = VBench(, , )
my_VBench.evaluate(
videos_path = ,
name = ,
diff --git a/setup.py b/setup.py
index 13e28fa1..f6cb0100 100644
--- a/setup.py
+++ b/setup.py
@@ -4,6 +4,8 @@
import os
def check_torch_version():
+ use_cuda = False
+
try:
import torch
if not torch.cuda.is_available():
@@ -13,6 +15,8 @@ def check_torch_version():
print(cuda_version)
if cuda_version not in ["11.6", "11.7", "11.8", "12.1"]:
raise RuntimeError(f"\033[91mUnsupported CUDA version: {cuda_version}. Please install PyTorch with 11.6<=CUDA<=12.1.\033[0m")
+
+ use_cuda = True
except:
print("\033[93mPlease install torch and torchvision compiled with cuda 11.8 before installing vbench\033[0m")
print("\033[93mFor CUDA 11.8, run:\033[0m")
@@ -21,7 +25,23 @@ def check_torch_version():
print("\033[93m pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121\033[0m")
print("\033[93mFor CUDA 11 with PyTorch < 2.0 run:\033[0m")
print("\033[93m pip install torch==1.13.1 torchvision==0.14.1\033[0m")
- raise
+ print("\033[93mTry to use Ascend NPU.\033[0m")
+
+ if use_cuda:
+ return
+
+ try:
+ import torch
+ import torch_npu
+
+ if not torch.npu.is_available():
+ raise RuntimeError("\033[91mAscend NPU is not available.\033[0m")
+
+ print("\033[92mUsing Ascend NPU to accelerate the process.\033[0m")
+ except:
+ print("\033[93mPlease install torch torch_npu and torchvision before installing vbench if you are using Ascend NPU devices.\033[0m")
+ print("\033[93mFor more details, please refer to the official installation guide: https://github.com/Ascend/pytorch#installation\033[0m")
+
def fetch_readme():
with open('README-pypi.md', encoding='utf-8') as f:
diff --git a/vbench/__init__.py b/vbench/__init__.py
index 83e284e0..6675891a 100755
--- a/vbench/__init__.py
+++ b/vbench/__init__.py
@@ -14,6 +14,7 @@ def __init__(self, device, full_info_dir, output_path):
self.full_info_dir = full_info_dir # full json file that VBench originally provides
self.output_path = output_path # output directory to save VBench results
os.makedirs(self.output_path, exist_ok=True)
+ print0(f"Use device: {self.device}")
def build_full_dimension_list(self, ):
return ["subject_consistency", "background_consistency", "aesthetic_quality", "imaging_quality", "object_class", "multiple_objects", "color", "spatial_relationship", "scene", "temporal_style", 'overall_consistency', "human_action", "temporal_flickering", "motion_smoothness", "dynamic_degree", "appearance_style"]