File size: 3,008 Bytes
0ba1507 db587e8 ff31910 0ba1507 ff31910 0ba1507 a6a25ce 0ba1507 a6a25ce 0ba1507 a6a25ce 0ba1507 a6a25ce 0ba1507 a6a25ce |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
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")
# Store original image for comparison
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 tuple of original and processed images for the slider
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) |