optimize api

app.py

@@ -2,6 +2,7 @@ import os
 import shutil
 import uuid
 import cv2
+import gc
 import gradio as gr
 import torch
 from basicsr.archs.rrdbnet_arch import RRDBNet
@@ -16,24 +17,23 @@ if not os.path.exists('model_zoo/gan/GFPGANv1.4.pth'):
 if not os.path.exists('model_zoo/swinir/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth'):
     os.system('wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth -P model_zoo/swinir')
 
-model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
-model_path = 'model_zoo/real/RealESRGAN_x4plus.pth'
-netscale = 4
-tile = 400 if torch.cuda.is_available() else 0
-dni_weight = None
-# restorer
-upsampler = RealESRGANer(
-    scale=netscale,
-    model_path=model_path,
-    dni_weight=dni_weight,
-    model=model,
-    tile=tile,
-    tile_pad=10,
-    pre_pad=0,
-    half=False, #Use fp32 precision during inference. Default: fp16 (half precision).
-    gpu_id=None) #gpu device to use (default=None) can be 0,1,2 for multi-gpu
-
 def inference(img, scale):
+    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+    model_path = 'model_zoo/real/RealESRGAN_x4plus.pth'
+    netscale = 4
+    tile = 400 if torch.cuda.is_available() else 0
+    dni_weight = None
+    # restorer
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=dni_weight,
+        model=model,
+        tile=tile,
+        tile_pad=10,
+        pre_pad=0,
+        half=False, #Use fp32 precision during inference. Default: fp16 (half precision).
+        gpu_id=None) #gpu device to use (default=None) can be 0,1,2 for multi-gpu
     # background enhancer with RealESRGAN
     os.makedirs('output', exist_ok=True)
     if scale > 4:
@@ -52,22 +52,16 @@ def inference(img, scale):
         h, w = img.shape[0:2]
         if h < 300:
             img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
-        try:
-            face_enhancer = GFPGANer(
-                model_path='model_zoo/gan/GFPGANv1.4.pth', upscale=scale, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
-            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
 
+        face_enhancer = GFPGANer(
+            model_path='model_zoo/gan/GFPGANv1.4.pth', upscale=scale, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
+        _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
 
-        except RuntimeError as error:
-            print('Error', error)
+        if scale != 2:
+            interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
+            h, w = img.shape[0:2]
+            output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
 
-        try:
-            if scale != 2:
-                interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
-                h, w = img.shape[0:2]
-                output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
-        except Exception as error:
-            print('wrong scale input.', error)
         if img_mode == 'RGBA':  # RGBA images should be saved in png format
             extension = 'png'
         else:
@@ -82,6 +76,23 @@ def inference(img, scale):
     except Exception as error:
         print('global exception', error)
         return None, None
+    finally:
+        #clean_folder('output')
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+def clean_folder(folder):
+    for filename in os.listdir(folder):
+        file_path = os.path.join(folder, filename)
+        try:
+            if os.path.isfile(file_path) or os.path.islink(file_path):
+                os.unlink(file_path)
+            elif os.path.isdir(file_path):
+                shutil.rmtree(file_path)
+        except Exception as e:
+            print('Failed to delete %s. Reason: %s' % (file_path, e))
+
 title = "Real Esrgan Restore Ai Face Restoration by appsgenz.com"
 description = ""
 article = "AppsGenz"
@@ -97,4 +108,4 @@ grApp = gr.Interface(
     description=description,
     article=article)
 grApp.queue(concurrency_count=2)
-grApp.launch()
+grApp.launch(share=True)