fix the --gpu mess

This commit is contained in:
Somdev Sangwan 2023-05-30 07:00:37 +05:30 committed by GitHub
parent df9ad81a10
commit 25444da59d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

3
run.py
View File

@ -36,13 +36,14 @@ parser.add_argument('-f', '--face', help='use this face', dest='source_img')
parser.add_argument('-t', '--target', help='replace this face', dest='target_path')
parser.add_argument('-o', '--output', help='save output to this file', dest='output_file')
parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False)
if torch.cuda.is_available():
parser.add_argument('--gpu', help='use gpu', dest='gpu', action='store_true', default=False)
parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False)
for name, value in vars(parser.parse_args()).items():
args[name] = value
if torch.cuda.is_available() not args['gpu']:
quit("You are using --gpu flag but CUDA isn't available on your system.")
sep = "/"
if os.name == "nt":