File size: 1,685 Bytes
09fb9f0
5941390
d127695
5941390
 
 
8346ea0
 
 
 
 
 
383c5b7
5941390
 
383c5b7
5941390
 
383c5b7
d670e90
5941390
 
 
 
383c5b7
 
 
 
 
 
 
 
09fb9f0
d127695
5941390
d127695
 
 
383c5b7
 
 
 
 
 
 
66eaf5f
d127695
 
66eaf5f
 
d127695
66eaf5f
5941390
383c5b7
 
 
 
5941390
383c5b7
5941390
 
 
 
 
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 gradio as gr
import spaces
from gradio_imageslider import ImageSlider
from image_gen_aux import UpscaleWithModel
from image_gen_aux.utils import load_image

MODELS = {
    "UltraSharp": "OzzyGT/UltraSharp",
    "DAT X4": "OzzyGT/DAT_X4",
    "DAT X3": "OzzyGT/DAT_X3",
    "DAT X2": "OzzyGT/DAT_X2",
}


@spaces.GPU
def upscale_image(image, model_selection):
    original = load_image(image)

    upscaler = UpscaleWithModel.from_pretrained(MODELS[model_selection]).to("cuda")
    image = upscaler(original, tiling=True, tile_width=1024, tile_height=1024)

    return original, image


def clear_result():
    return gr.update(value=None)


title = """<h1 align="center">Image Upscaler</h1>
<div align="center">This space is a showcase of the different super resolution models you can use to upscale with the 
<a href="https://github.com/asomoza/image_gen_aux">Image Generation Auxiliary Tools</a> library.</div>
"""

with gr.Blocks() as demo:
    gr.HTML(title)
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")

            model_selection = gr.Dropdown(
                choices=list(MODELS.keys()),
                value="UltraSharp",
                label="Model",
            )

            run_button = gr.Button("Upscale")
        with gr.Column():
            result = ImageSlider(
                interactive=False,
                label="Generated Image",
            )

    run_button.click(
        fn=clear_result,
        inputs=None,
        outputs=result,
    ).then(
        fn=upscale_image,
        inputs=[input_image, model_selection],
        outputs=result,
    )


demo.launch(share=False)