|
import torch |
|
from PIL import Image |
|
from RealESRGAN import RealESRGAN |
|
import gradio as gr |
|
import os |
|
import spaces |
|
|
|
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") |
|
|
|
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) |
|
|
|
@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)}") |
|
|
|
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." |
|
|
|
|
|
|
|
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, |
|
flagging_mode="never", |
|
cache_examples=True |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
iface.launch(debug=True, show_error=True) |
