app.py

import torch
from PIL import Image
from RealESRGAN import RealESRGAN
import gradio as gr
import os
import spaces

# 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):
    if image is None:
        raise gr.Error("Image not uploaded")

    try:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        if size == '2x':
            result = model2.predict(image.convert('RGB'))
        elif size == '4x':
            result = model4.predict(image.convert('RGB'))
        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'))

        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>"

# 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_examples=True
)

# Chạy ứng dụng
if __name__ == "__main__":
    iface.launch(debug=True, show_error=True)