Update engine.py

engine.py

@@ -18,6 +18,8 @@ import tqdm
 from tqdm import tqdm as loader
 
 import cv2
+import torch
+from basicsr.archs.srvgg_arch import SRVGGNetCompact
 
 from liveportrait.config.argument_config import ArgumentConfig
 from liveportrait.utils.camera import get_rotation_matrix
@@ -31,6 +33,21 @@ logger = logging.getLogger(__name__)
 # Global constants
 DATA_ROOT = os.environ.get('DATA_ROOT', '/tmp/data')
 MODELS_DIR = os.path.join(DATA_ROOT, "models")
+os.system("pip freeze")
+
+if not os.path.exists('RestoreFormer.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth -P .")
+
+if not os.path.exists('realesr-general-x4v3.pth'):
+    os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
+    
+model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+model_path = 'realesr-general-x4v3.pth'
+half = True if torch.cuda.is_available() else False
+upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
+
+enhancer = GFPGANer(
+                model_path='RestoreFormer.pth', upscale=2, arch='RestoreFormer', channel_multiplier=2, bg_upsampler=upsampler)
 
 def base64_data_uri_to_PIL_Image(base64_string: str) -> Image.Image:
     """
@@ -113,7 +130,9 @@ class Engine:
             
             _, frame = await self.transform_frame(processed_data,  params)
             bgr_frame = cv2.cvtColor(np.array(frame), cv2.COLOR_BGR2RGB)
-            video_writer.write(cv2.cvtColor(bgr_frame, cv2.COLOR_BGR2RGB))
+            new_frame = cv2.cvtColor(bgr_frame, cv2.COLOR_BGR2RGB)
+            _, _, output = enhancer.enhance(new_frame, has_aligned=False, only_center_face=False, paste_back=True)
+            video_writer.write(output)
             
         video_writer.release()
         cap.release()