File size: 4,253 Bytes
708a0d1 f80801d 708a0d1 f80801d 92a908f 708a0d1 |
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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
from sldl.video import VideoSR
from sldl.image import ImageSR
import gradio as gr
import tempfile
import shutil
import torch
import ffmpeg
import time
from PIL import Image
cc = 2
if torch.backends.mps.is_available():
device = 'mps'
cc = 5
elif torch.cuda.is_available():
device = 'cuda'
cc = 10
else:
device = 'cpu'
vbsrgan = VideoSR('BSRGAN').to(device)
vresrgan = VideoSR('RealESRGAN').to(device)
ibsrgan = ImageSR('BSRGAN').to(device)
iresrgan = ImageSR('RealESRGAN').to(device)
def upscale_video(input_video, output_video, progress, mname):
modelname = mname.lower()
model = vbsrgan
if modelname == 'bsrgan (default)':
# do nothing
pass
elif modelname == 'real esrgan':
model = vresrgan
model(input_video, output_video, progress.tqdm)
def upscale_image(input_image, output_image, mname):
modelname = mname.lower()
model = ibsrgan
if modelname == 'bsrgan (default)':
# do nothing
pass
elif modelname == 'real esrgan':
model = iresrgan
shutil.copy(input_image, output_image)
model(output_image)
# Gradio interface
def video_upscaling_interface(input_text, model_name, progress=gr.Progress()):
if input_text:
temp_dir = tempfile.mkdtemp()
input_video_path = f"{temp_dir}/input_video"
output_video_path = f"{temp_dir}/output_video.mp4"
ffmpeg.input(input_text).output(input_video_path + '.mp4').run()
# Upscale the video
upscale_video(input_video_path + '.mp4', output_video_path, progress, model_name)
return [output_video_path, output_video_path]
else:
return ["no_vid.mp4", "no_vid.mp4"]
def image_upscaling_interface(input_text, model_name):
if input_text:
temp_dir = tempfile.mkdtemp()
input_image_path = f"{temp_dir}/input_image.jpg"
output_image_path = f"{temp_dir}/output_image.jpg"
input_text.save(input_image_path)
upscale_image(input_image_path, output_image_path, model_name)
return [output_image_path, output_image_path]
else:
return ["no_image.jpg", "no_image.jpg"]
css = "footer {display: none !important;} .gradio-container {min-height: 0px !important;}"
with gr.Blocks(css=css) as demo:
gr.Markdown('''
# Upscale
## A CVSYS Project
### NOTICE: This is running on a free Hugging Face Space. That means that it will be VERY slow. Expect it to take HOURS to upscale 5 minutes. Sorry, but that's what 4 vCPUs and 16 GB RAM brings :( - PLEASE be mindful and DO NOT upscale videos longer than 15 seconds!!
Please note that after you upload an image, it may take several minutes before the progress bar appears. This is because we first convert your video to ensure the correct format.
''')
# with gr.Tab("Image"):
# with gr.Row():
# with gr.Column():
# iinp = gr.Image(label="Upload Image", interactive=True, type="pil")
# imod = gr.Dropdown(
# ["BSRGAN (Default)", "Real ESRGAN"],
# value="BSRGAN (Default)",
# interactive=True,
# label="Model"
# )
# with gr.Column():
# iout = gr.Image(label="View Image", interactive=False, type="filepath")
# ifile = gr.File(label="Download Image", interactive=False)
# ibtn = gr.Button(value="Upscale Image")
with gr.Tab("Video"):
with gr.Row():
with gr.Column():
vinp = gr.Video(label="Upload Video", interactive=True)
vmod = gr.Dropdown(
["BSRGAN (Default)", "Real ESRGAN"],
value="BSRGAN (Default)",
interactive=True,
label="Model"
)
with gr.Column():
vout = gr.Video(label="Watch Video", interactive=False)
vfile = gr.File(label="Download Video", interactive=False)
vbtn = gr.Button(value="Upscale Video")
ibtn.click(image_upscaling_interface, [iinp, imod], outputs=[iout, ifile])
vbtn.click(video_upscaling_interface, [vinp, vmod], outputs=[vout, vfile])
demo.queue(concurrency_count=cc)
demo.launch()
|