File size: 2,014 Bytes
a998f86 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import torch
from PIL import Image
from RealESRGAN import RealESRGAN
import gradio as gr
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_scales = {'2x': 2, '4x': 4, '8x': 8}
# Load RealESRGAN models for different scales
models = {scale: RealESRGAN(device, scale=scale) for scale in model_scales.values()}
def inference(images, scale):
results = []
if images is None or len(images) == 0:
raise gr.Error("No image uploaded. Please upload at least one image.")
for image in images:
width, height = 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()
# Select the appropriate model based on the chosen scale
model = models[model_scales[scale]]
result = model.predict(image.convert('RGB'))
print(f"Image size ({device}): {scale} ... OK")
results.append(result)
return results
title = "Advanced Real ESRGAN UpScale: 2x 4x 8x"
description = (
"This advanced demo for Real-ESRGAN allows you to upscale multiple images "
"with different models and resolutions. Choose the scale and upload images for high-resolution enhancement."
)
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", label="Upload Image", multiple=True),
gr.Radio(
list(model_scales.keys()),
type="value",
value='2x',
label='Resolution model',
),
],
gr.Image(type="pil", label="Output"),
title=title,
description=description,
article=article,
examples=[['groot.jpeg', '2x']],
allow_flagging='never',
cache_examples=False,
).launch()
|