Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 
+os.system('nvidia-smi')
 os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/RetinaFace-R50.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116085&Signature=GlUNW6%2B8FxvxWmE9jKIZYOOciKQ%3D" -O weights/RetinaFace-R50.pth')
 os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-BFR-512.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116208&Signature=hBgvVvKVSNGeXqT8glG%2Bd2t2OKc%3D" -O weights/GPEN-512.pth')
 os.system('wget "https://public-vigen-video.oss-cn-shanghai.aliyuncs.com/robin/models/GPEN-Colorization-1024.pth?OSSAccessKeyId=LTAI4G6bfnyW4TA4wFUXTYBe&Expires=1961116315&Signature=9tPavW2h%2F1LhIKiXj73sTQoWqcc%3D" -O weights/GPEN-1024-Color.pth ')
@@ -18,24 +19,27 @@ from face_colorization import FaceColorization
   
 
 def inference(file, mode):
-    im = cv2.imread(file, cv2.IMREAD_COLOR)
-    im = cv2.resize(im, (0,0), fx=2, fy=2)
-    faceenhancer = FaceEnhancement(size=512, model='GPEN-512', channel_multiplier=2, device='cuda')
-    img, orig_faces, enhanced_faces = faceenhancer.process(im)
-    
+
     if mode == "enhance":
-        # model = {'name':'GPEN-512', 'size':512}
-        # im = cv2.imread(file, cv2.IMREAD_COLOR)
-        # im = cv2.resize(im, (0,0), fx=2, fy=2)
-        # faceenhancer = FaceEnhancement(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
-        # img, orig_faces, enhanced_faces = faceenhancer.process(im)
+        model = {'name':'GPEN-512', 'size':512}
+        im = cv2.imread(file, cv2.IMREAD_COLOR)
+        im = cv2.resize(im, (0,0), fx=2, fy=2)
+        faceenhancer = FaceEnhancement(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
+        img, orig_faces, enhanced_faces = faceenhancer.process(im)
         cv2.imwrite(os.path.join("output.png"), img)
         return os.path.join("output.png")
     else:
+        emodel = {'name':'GPEN-512', 'size':512}
+        im = cv2.imread(file, cv2.IMREAD_COLOR)
+        im = cv2.resize(im, (0,0), fx=2, fy=2)
+        faceenhancer = FaceEnhancement(size=model['size'], model=emodel['name'], channel_multiplier=2, device='cpu')
+        img, orig_faces, enhanced_faces = faceenhancer.process(im)
+        cv2.imwrite(os.path.join("enhanced.png"), img)
+        
         model = {'name':'GPEN-1024-Color', 'size':1024}
         grayf = cv2.imread("enhanced.png", cv2.IMREAD_GRAYSCALE)
         grayf = cv2.cvtColor(grayf, cv2.COLOR_GRAY2BGR) # channel: 1->3
-        facecolorizer = FaceColorization(size=model['size'], model=model['name'], channel_multiplier=2, device='cuda')
+        facecolorizer = FaceColorization(size=model['size'], model=model['name'], channel_multiplier=2, device='cpu')
         colorf = facecolorizer.process(grayf)
 
         colorf = cv2.resize(colorf, (grayf.shape[1], grayf.shape[0]))