Update app.py

app.py

@@ -2,6 +2,7 @@ import torch
 from PIL import Image
 from RealESRGAN import RealESRGAN
 import gradio as gr
+import gc
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model2 = RealESRGAN(device, scale=2)
@@ -11,6 +12,9 @@ model4.load_weights('weights/RealESRGAN_x4.pth', download=True)
 model8 = RealESRGAN(device, scale=8)
 model8.load_weights('weights/RealESRGAN_x8.pth', download=True)
 
+if torch.cuda.is_available():
+    torch.cuda.empty_cache()
+    gc.collect()
 
 def inference(image, size):
     if image is None:
@@ -20,9 +24,6 @@ def inference(image, size):
     if width >= 5000 or height >= 5000:
         raise gr.Error("The image is too large.")
         
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        
     if size == '2x':
         result = model2.predict(image.convert('RGB'))
     elif size == '4x':