|
|
import torch |
|
|
from PIL import Image |
|
|
from RealESRGAN import RealESRGAN |
|
|
import gradio as gr |
|
|
from gradio_imageslider import ImageSlider |
|
|
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) |
|
|
|
|
|
|
|
|
@spaces.GPU |
|
|
def inference(image, size): |
|
|
global model2 |
|
|
global model4 |
|
|
global model8 |
|
|
if image is None: |
|
|
raise gr.Error("Image not uploaded") |
|
|
|
|
|
|
|
|
original_image = image.copy() |
|
|
|
|
|
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) |
|
|
result = model2.predict(image.convert('RGB')) |
|
|
elif size == '4x': |
|
|
try: |
|
|
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: |
|
|
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 (original_image, result) |
|
|
|
|
|
|
|
|
title = """<h1 align="center">ProFaker</h1>""" |
|
|
|
|
|
with gr.Blocks() as demo: |
|
|
gr.HTML(title) |
|
|
|
|
|
with gr.Row(): |
|
|
with gr.Column(): |
|
|
input_image = gr.Image(type="pil", label="Input Image") |
|
|
size_select = gr.Radio( |
|
|
["2x", "4x", "8x"], |
|
|
type="value", |
|
|
value="2x", |
|
|
label="Resolution model" |
|
|
) |
|
|
process_btn = gr.Button("Upscale Image") |
|
|
|
|
|
with gr.Column(): |
|
|
result_slider = ImageSlider( |
|
|
interactive=False, |
|
|
label="Before and After Comparison" |
|
|
) |
|
|
|
|
|
process_btn.click( |
|
|
fn=inference, |
|
|
inputs=[input_image, size_select], |
|
|
outputs=result_slider |
|
|
) |
|
|
|
|
|
demo.queue(api_open=True).launch(show_error=True, show_api=True) |
