2020import yaml
2121import numpy as np
2222import warnings
23+ import shlex
2324from skimage import img_as_ubyte
2425import safetensors
2526import safetensors .torch
@@ -276,7 +277,7 @@ def generate(
276277 ]
277278
278279 video_name = x ["video_name" ] + ".mp4"
279- path = os .path .join (video_save_dir , "temp_" + video_name )
280+ path = os .path .join (video_save_dir , "temp_" + shlex . quote ( video_name ) )
280281
281282 imageio .mimsave (path , result , fps = float (25 ))
282283
@@ -285,7 +286,7 @@ def generate(
285286
286287 audio_path = x ["audio_path" ]
287288 audio_name = os .path .splitext (os .path .split (audio_path )[- 1 ])[0 ]
288- new_audio_path = os .path .join (video_save_dir , audio_name + ".wav" )
289+ new_audio_path = os .path .join (video_save_dir , shlex . quote ( audio_name ) + ".wav" )
289290 start_time = 0
290291 # cog will not keep the .mp3 filename
291292 sound = AudioSegment .from_file (audio_path )
@@ -323,7 +324,7 @@ def generate(
323324 start_time = end_time
324325 if enhancer :
325326 video_name_enhancer = x ["video_name" ] + "_enhanced.mp4"
326- enhanced_path = os .path .join (video_save_dir , "temp_" + video_name_enhancer )
327+ enhanced_path = os .path .join (video_save_dir , "temp_" + shlex . quote ( video_name_enhancer ) )
327328 av_path_enhancer = os .path .join (video_save_dir , video_name_enhancer )
328329 return_path = av_path_enhancer
329330 enhanced_images = face_enhancer (
@@ -333,12 +334,12 @@ def generate(
333334 imageio .mimsave (enhanced_path , enhanced_images , fps = float (25 ))
334335 save_video_with_watermark (enhanced_path , new_audio_path , av_path_enhancer , watermark = False )
335336 print (f"The generated video is named { video_save_dir } /{ video_name_enhancer } " )
336- os .remove (enhanced_path )
337+ os .remove (shlex . quote ( enhanced_path ) )
337338 end_time = time .time ()
338339 print (f"[***6/6***] enhancing takes: { end_time - start_time } sec" )
339340 else :
340341 print (f"[***6/6***] no enhancing" )
341- os .remove (path )
342- os .remove (new_audio_path )
342+ os .remove (shlex . quote ( path ) )
343+ os .remove (shlex . quote ( new_audio_path ) )
343344
344345 return return_path
0 commit comments