add gfpgan

src/facerender/animate.py

@@ -20,6 +20,7 @@ from src.facerender.modules.generator import OcclusionAwareGenerator, OcclusionA
 from src.facerender.modules.make_animation import make_animation 
 
 from pydub import AudioSegment 
+from src.utils.face_enhancer import enhancer_generator_with_len, enhancer_list
 from src.utils.paste_pic import paste_pic
 from src.utils.videoio import save_video_with_watermark
 
@@ -232,22 +233,22 @@ class AnimateFromCoeff():
             full_video_path = av_path 
 
         #### paste back then enhancers
-        # if enhancer:
-        #     video_name_enhancer = x['video_name']  + '_enhanced.mp4'
-        #     enhanced_path = os.path.join(video_save_dir, 'temp_'+video_name_enhancer)
-        #     av_path_enhancer = os.path.join(video_save_dir, video_name_enhancer) 
-        #     return_path = av_path_enhancer
-
-        #     try:
-        #         enhanced_images_gen_with_len = enhancer_generator_with_len(full_video_path, method=enhancer, bg_upsampler=background_enhancer)
-        #         imageio.mimsave(enhanced_path, enhanced_images_gen_with_len, fps=float(25))
-        #     except:
-        #         enhanced_images_gen_with_len = enhancer_list(full_video_path, method=enhancer, bg_upsampler=background_enhancer)
-        #         imageio.mimsave(enhanced_path, enhanced_images_gen_with_len, fps=float(25))
+        if enhancer:
+            video_name_enhancer = x['video_name']  + '_enhanced.mp4'
+            enhanced_path = os.path.join(video_save_dir, 'temp_'+video_name_enhancer)
+            av_path_enhancer = os.path.join(video_save_dir, video_name_enhancer) 
+            return_path = av_path_enhancer
+
+            try:
+                enhanced_images_gen_with_len = enhancer_generator_with_len(full_video_path, method=enhancer, bg_upsampler=background_enhancer)
+                imageio.mimsave(enhanced_path, enhanced_images_gen_with_len, fps=float(25))
+            except:
+                enhanced_images_gen_with_len = enhancer_list(full_video_path, method=enhancer, bg_upsampler=background_enhancer)
+                imageio.mimsave(enhanced_path, enhanced_images_gen_with_len, fps=float(25))
             
-        #     save_video_with_watermark(enhanced_path, new_audio_path, av_path_enhancer, watermark= False)
-        #     print(f'The generated video is named {video_save_dir}/{video_name_enhancer}')
-        #     os.remove(enhanced_path)
+            save_video_with_watermark(enhanced_path, new_audio_path, av_path_enhancer, watermark= False)
+            print(f'The generated video is named {video_save_dir}/{video_name_enhancer}')
+            os.remove(enhanced_path)
 
         os.remove(path)
         os.remove(new_audio_path)

src/facerender/pirender_animate.py

@@ -17,6 +17,7 @@ from src.facerender.pirender.config import Config
 from src.facerender.pirender.face_model import FaceGenerator
 
 from pydub import AudioSegment 
+from src.utils.face_enhancer import enhancer_generator_with_len, enhancer_list
 from src.utils.paste_pic import paste_pic
 from src.utils.videoio import save_video_with_watermark
 
@@ -105,22 +106,22 @@ class AnimateFromCoeff_PIRender():
             full_video_path = av_path 
 
         #### paste back then enhancers
-        # if enhancer:
-        #     video_name_enhancer = x['video_name']  + '_enhanced.mp4'
-        #     enhanced_path = os.path.join(video_save_dir, 'temp_'+video_name_enhancer)
-        #     av_path_enhancer = os.path.join(video_save_dir, video_name_enhancer) 
-        #     return_path = av_path_enhancer
-
-        #     try:
-        #         enhanced_images_gen_with_len = enhancer_generator_with_len(full_video_path, method=enhancer, bg_upsampler=background_enhancer)
-        #         imageio.mimsave(enhanced_path, enhanced_images_gen_with_len, fps=float(25))
-        #     except:
-        #         enhanced_images_gen_with_len = enhancer_list(full_video_path, method=enhancer, bg_upsampler=background_enhancer)
-        #         imageio.mimsave(enhanced_path, enhanced_images_gen_with_len, fps=float(25))
+        if enhancer:
+            video_name_enhancer = x['video_name']  + '_enhanced.mp4'
+            enhanced_path = os.path.join(video_save_dir, 'temp_'+video_name_enhancer)
+            av_path_enhancer = os.path.join(video_save_dir, video_name_enhancer) 
+            return_path = av_path_enhancer
+
+            try:
+                enhanced_images_gen_with_len = enhancer_generator_with_len(full_video_path, method=enhancer, bg_upsampler=background_enhancer)
+                imageio.mimsave(enhanced_path, enhanced_images_gen_with_len, fps=float(25))
+            except:
+                enhanced_images_gen_with_len = enhancer_list(full_video_path, method=enhancer, bg_upsampler=background_enhancer)
+                imageio.mimsave(enhanced_path, enhanced_images_gen_with_len, fps=float(25))
             
-        #     save_video_with_watermark(enhanced_path, new_audio_path, av_path_enhancer, watermark= False)
-        #     print(f'The generated video is named {video_save_dir}/{video_name_enhancer}')
-        #     os.remove(enhanced_path)
+            save_video_with_watermark(enhanced_path, new_audio_path, av_path_enhancer, watermark= False)
+            print(f'The generated video is named {video_save_dir}/{video_name_enhancer}')
+            os.remove(enhanced_path)
 
         os.remove(path)
         os.remove(new_audio_path)

src/utils/face_enhancer.py

@@ -1,7 +1,7 @@
 import os
 import torch 
 
-from gfpgan import GFPGANer
+from gfpgan.utils import GFPGANer
 
 from tqdm import tqdm