support cpu infer for vsr (PaddlePaddle#725)

wangna11BD · Nov 9, 2022 · 9e22482 · 9e22482
1 parent c2a79fe
commit 9e22482
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 0 deletions.
diff --git a/applications/tools/video-enhance.py b/applications/tools/video-enhance.py
@@ -109,9 +109,15 @@
                     default='none',
                     nargs='+',
                     help='Process order')
+parser.add_argument("--cpu",
+                    dest="cpu",
+                    action="store_true",
+                    help="cpu mode.")
 
 if __name__ == "__main__":
     args = parser.parse_args()
+    if args.cpu:
+        paddle.set_device('cpu')
 
     orders = args.process_order
     temp_video_path = None

diff --git a/docs/en_US/tutorials/video_restore.md b/docs/en_US/tutorials/video_restore.md
@@ -12,6 +12,7 @@ python tools/video-enhance.py --input you_video_path.mp4 --process_order DAIN De
 - `--input (str)`: path of the input video.
 - `--output (str)`: path of the output video.
 - `--process_order`: name and order of called model. For example, if the input is `DAIN DeOldify EDVR`, then `DAINPredictor` `DeOldifyPredictor` `EDVRPredictor` will be called in sequence.
+- `--cpu`: use cpu inference, GPU inference is used by default.
 
 #### Results
 ![](../../imgs/color_sr_peking.gif)

diff --git a/docs/zh_CN/industrial_solution/video_restore_cn.md b/docs/zh_CN/industrial_solution/video_restore_cn.md
@@ -16,6 +16,7 @@ python tools/video-enhance.py --input you_video_path.mp4 --process_order DAIN De
 - `--input (str)`: 输入的视频路径。
 - `--output (str)`: 输出的视频路径。
 - `--process_order`: 调用的模型名字和顺序，比如输入为 `DAIN DeOldify PPMSVSR`，则会顺序调用 `DAINPredictor` `DeOldifyPredictor` `PPMSVSRPredictor` 。
+- `--cpu`: 开启cpu推理，默认使用GPU推理。
 
 <div align='center'>
   <img src='https://user-images.githubusercontent.com/48054808/117925494-e9a70400-b329-11eb-9f38-a48ef946a3a4.gif' width='600'>