add BASICSR_EXT env

README.md

@@ -58,17 +58,22 @@ These pipelines/commands cannot cover all the cases and more details are in the
 - [PyTorch >= 1.7](https://pytorch.org/)
 - NVIDIA GPU + [CUDA](https://developer.nvidia.com/cuda-downloads)
 
-### Pip install
+### Option 1: Pip install
 
 ```bash
 pip install basicsr
 ```
 
-- pip installation does not compile cuda extensions.
-- If you want to use cuda extensions, set environment variable `BASICSR_JIT=True`. Note that every time you run the model, it will compile the extensions just time.
+- If you want to compile cuda extensions when installing, please set up the environment variable `BASICSR_EXT=True`.
+
+  ```bash
+  BASICSR_EXT=True pip install basicsr
+  ```
+
+- If you want to use cuda extensions during running, set environment variable `BASICSR_JIT=True`. Note that every time you run the model, it will compile the extensions just time.
   - Example: StyleGAN2 inference colab.
 
-### Git clone and compile
+### Option 2: Git clone and compile
 
 1. Clone repo
 
@@ -90,11 +95,11 @@ pip install basicsr
     If you do need the cuda extensions: <br>
     &emsp;[*dcn* for EDVR](basicsr/ops)<br>
     &emsp;[*upfirdn2d* and *fused_act* for StyleGAN2](basicsr/ops)<br>
-    please add `--cuda_ext` when installing.<br>
+    please set up the environment variable `BASICSR_EXT=True` when installing.<br>
     If you use the EDVR and StyleGAN2 model, the above cuda extensions are necessary.
 
     ```bash
-    python setup.py develop --cuda_ext
+    BASICSR_EXT=True python setup.py develop
     ```
 
     Otherwise, install without compiling cuda extensions

README_CN.md

@@ -57,17 +57,22 @@ BasicSR (**Basic** **S**uper **R**estoration) 是一个基于 PyTorch 的开源
 - [PyTorch >= 1.7](https://pytorch.org/)
 - NVIDIA GPU + [CUDA](https://developer.nvidia.com/cuda-downloads)
 
-### 通过 pip 安装
+### 选项1: 通过 pip 安装
 
 ```bash
 pip install basicsr
 ```
 
-- 通过 pip 安装不会编译 cuda 算子.
-- 如果你想使用 cuda 算子, 则需要设置环境变量 `BASICSR_JIT=True`. 注意: 通过这种方式, 在每一次运行模型时候, 它都需要编译, (比较耗时).
+- 如果要在安装的时候编译 cuda 算子, 请设置环境变量 `BASICSR_EXT=True`。
+
+  ```bash
+  BASICSR_EXT=True pip install basicsr
+  ```
+
+- 如果你想在运行的时候使用 cuda 算子（安装的时候没有编译）, 则需要设置环境变量 `BASICSR_JIT=True`. 注意: 通过这种方式, 在每一次运行模型时候, 它都需要编译, (比较耗时).
   - 例子: StyleGAN2 inference colab.
 
-### 通过 git clone 然后 编译
+### 选项2: 通过 git clone 然后 编译
 
 1. Clone repo
 
@@ -89,11 +94,11 @@ pip install basicsr
     如果你需要以下 cuda 扩展算子: <br>
     &emsp;[*dcn* for EDVR](basicsr/ops)<br>
     &emsp;[*upfirdn2d* and *fused_act* for StyleGAN2](basicsr/ops)<br>
-    在安装命令后添加 `--cuda_ext`.<br>
+    在安装命令时设置环境变量 后添加 `BASICSR_EXT=True`.<br>
     如果使用 EDVR 和 StyleGAN2 模型, 需要使用上面的 cuda 扩展算子.
 
     ```bash
-    python setup.py develop --cuda_ext
+    BASICSR_EXT=True python setup.py develop
     ```
 
     否则, 安装的时候不会编译 cuda 算子.

setup.py

@@ -4,7 +4,6 @@
 
 import os
 import subprocess
-import sys
 import time
 import torch
 from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension
@@ -115,7 +114,8 @@ def get_requirements(filename='requirements.txt'):
 
 
 if __name__ == '__main__':
-    if '--cuda_ext' in sys.argv:
+    cuda_ext = os.getenv('BASICSR_EXT')  # whether compile cuda ext
+    if cuda_ext == 'True':
         ext_modules = [
             make_cuda_ext(
                 name='deform_conv_ext',
@@ -133,7 +133,6 @@ def get_requirements(filename='requirements.txt'):
                 sources=['src/upfirdn2d.cpp'],
                 sources_cuda=['src/upfirdn2d_kernel.cu']),
         ]
-        sys.argv.remove('--cuda_ext')
     else:
         ext_modules = []