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