File size: 2,349 Bytes
7bed60d |
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 |
import math
import torch
from modules import script_callbacks, images
from modules.shared import opts
import gradio as gr
from PIL import Image
from scripts.aura_sr import AuraSR
def upscale(imageSource, *args):
torch.set_grad_enabled(False)
if imageSource != None:
aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")
upscaledImage = aura_sr.upscale_4x_overlapped(imageSource)
del aura_sr
else:
upscaledImage = None
# re-enable the go button, return result
return gr.Button.update(value='Upscale', variant='primary', interactive=True), upscaledImage
def on_ui_tabs():
def toggleGo ():
# disable the go button while processing
return gr.Button.update(value='...', variant='secondary', interactive=False)
def saveImage (image, suffix):
# use the built-in webui save function
if image is not None:
images.save_image(
image,
opts.outdir_samples or opts.outdir_extras_samples,
'',
extension='png',
suffix=suffix,
)
return
with gr.Blocks() as auraSR_block:
with gr.Row():
with gr.Column():
#show image dimensions?
imageSource = gr.Image(label='image source', sources=['upload'], height=640, type='pil', interactive=True, show_download_button=False, )
go_button = gr.Button(value="Upscale", variant='primary', visible=True)
with gr.Column():
outputImage = gr.Image(label='Output', height=640, type='pil', interactive=False, show_label=False,)
with gr.Row():
filename = gr.Textbox(value='', placeholder='filename suffix for saving ... (regular pattern first)', lines=1, max_lines=1, scale=3, show_label=False)
save_button = gr.Button(value='Save', variant='secondary')
go_button.click(toggleGo, inputs=[], outputs=[go_button])
go_button.click(upscale, inputs=imageSource, outputs=[go_button, outputImage])
save_button.click(fn=saveImage, inputs=[outputImage, filename], outputs=[])
#### UI block name, tab display name, internal name
return [(auraSR_block, "auraSR", "aura_sr")]
script_callbacks.on_ui_tabs(on_ui_tabs)
|