Update app.py

app.py

@@ -2,77 +2,85 @@ import torch
 from PIL import Image
 from RealESRGAN import RealESRGAN
 import gradio as gr
+import os
+import spaces
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model2 = RealESRGAN(device, scale=2)
-model2.load_weights('weights/RealESRGAN_x2.pth', download=True)
-model4 = RealESRGAN(device, scale=4)
-model4.load_weights('weights/RealESRGAN_x4.pth', download=True)
-model8 = RealESRGAN(device, scale=8)
-model8.load_weights('weights/RealESRGAN_x8.pth', download=True)
+# Kiểm tra và cấu hình GPU
+if torch.cuda.is_available():
+    print(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
+    device = torch.device("cuda")
+else:
+    print("CUDA is not available. Using CPU.")
+    device = torch.device("cpu")
 
+# Lazy loading cho các mô hình
+class LazyRealESRGAN:
+    def __init__(self, device, scale):
+        self.device = device
+        self.scale = scale
+        self.model = None
 
+    def load_model(self):
+        if self.model is None:
+            self.model = RealESRGAN(self.device, scale=self.scale)
+            self.model.load_weights(f'weights/RealESRGAN_x{self.scale}.pth', download=True)
+
+    def predict(self, img):
+        self.load_model()
+        return self.model.predict(img)
+
+model2 = LazyRealESRGAN(device, scale=2)
+model4 = LazyRealESRGAN(device, scale=4)
+model8 = LazyRealESRGAN(device, scale=8)
+
+# Hàm inference chính
+@spaces.GPU
 def inference(image, size):
-    global model2
-    global model4
-    global model8
     if image is None:
         raise gr.Error("Image not uploaded")
-        
 
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    
-    if size == '2x':
-        try:
-            result = model2.predict(image.convert('RGB'))
-        except torch.cuda.OutOfMemoryError as e:
-            print(e)
-            model2 = RealESRGAN(device, scale=2)
-            model2.load_weights('weights/RealESRGAN_x2.pth', download=False)
+    try:
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        
+        if size == '2x':
             result = model2.predict(image.convert('RGB'))
-    elif size == '4x':
-        try:
+        elif size == '4x':
             result = model4.predict(image.convert('RGB'))
-        except torch.cuda.OutOfMemoryError as e:
-            print(e)
-            model4 = RealESRGAN(device, scale=4)
-            model4.load_weights('weights/RealESRGAN_x4.pth', download=False)
-            result = model2.predict(image.convert('RGB'))
-    else:
-        try:
+        else:
             width, height = image.size
             if width >= 5000 or height >= 5000:
                 raise gr.Error("The image is too large.")
             result = model8.predict(image.convert('RGB'))
-        except torch.cuda.OutOfMemoryError as e:
-            print(e)
-            model8 = RealESRGAN(device, scale=8)
-            model8.load_weights('weights/RealESRGAN_x8.pth', download=False)
-            result = model2.predict(image.convert('RGB'))
-            
-    print(f"Image size ({device}): {size} ... OK")
-    return result
 
+        print(f"Image size ({device}): {size} ... OK")
+        return result
+    except torch.cuda.OutOfMemoryError:
+        raise gr.Error("GPU out of memory. Try a smaller image or lower upscaling factor.")
+    except Exception as e:
+        raise gr.Error(f"An error occurred: {str(e)}")
 
+# Cấu hình giao diện Gradio
 title = "Face Real ESRGAN UpScale: 2x 4x 8x"
 description = "This is an unofficial demo for Real-ESRGAN. Scales the resolution of a photo. This model shows better results on faces compared to the original version.<br>Telegram BOT: https://t.me/restoration_photo_bot"
 article = "<div style='text-align: center;'>Twitter <a href='https://twitter.com/DoEvent' target='_blank'>Max Skobeev</a> | <a href='https://huggingface.co/sberbank-ai/Real-ESRGAN' target='_blank'>Model card</a><div>"
 
-
-gr.Interface(inference,
-    [gr.Image(type="pil"), 
-    gr.Radio(["2x", "4x", "8x"], 
-    type="value",
-    value="2x",
-    label="Resolution model")], 
+# Khởi tạo và chạy giao diện Gradio
+iface = gr.Interface(
+    inference,
+    [
+        gr.Image(type="pil"),
+        gr.Radio(["2x", "4x", "8x"], type="value", value="2x", label="Resolution model")
+    ],
     gr.Image(type="pil", label="Output"),
     title=title,
     description=description,
     article=article,
     examples=[["groot.jpeg", "2x"]],
     flagging_mode="never",
-    cache_mode="lazy",
-    delete_cache=(44000, 44000),
-    ).queue(api_open=True).launch(show_error=True, show_api=True)
-    
+    cache_examples=True
+)
+
+# Chạy ứng dụng
+if __name__ == "__main__":
+    iface.launch(debug=True, show_error=True)