File size: 2,919 Bytes
43404dd
2625963
43404dd
8b0d315
2625963
8b0d315
1be1a61
2625963
 
 
43404dd
 
 
1be1a61
ed068ef
 
 
8b0d315
1be1a61
 
 
 
 
2625963
8b0d315
 
ed068ef
1be1a61
ed068ef
1be7ff1
2625963
1be1a61
1be7ff1
1be1a61
8b0d315
ed068ef
 
 
1be1a61
2625963
 
 
 
 
 
 
 
8b0d315
2625963
 
 
 
 
 
 
ed068ef
2625963
8b0d315
 
ed068ef
 
 
8b0d315
2625963
8b0d315
2625963
 
1be1a61
2625963
8b0d315
2625963
8b0d315
2625963
1be1a61
2625963
1be1a61
2625963
43404dd
2625963
 
1be1a61
ed068ef
 
 
 
43404dd
2625963
 
1be1a61
 
2625963
 
 
1be1a61
 
2625963
1be1a61
2625963
8b0d315
ed068ef
 
 
2625963
 
ed068ef
2625963
1be1a61
2625963
 
 
 
 
8b0d315
2625963
ed068ef
2625963
 
 
1be1a61
8b0d315
2625963
8b0d315
 
ed068ef
1be1a61
ed068ef
2625963
8b0d315
 
2625963
1be1a61
2625963
ed068ef
1be7ff1
ed068ef
1be7ff1
8b0d315
ed068ef
1be1a61
2625963
1be1a61
 
8b0d315
ed068ef
2625963
 
 
8b0d315
 
1be1a61
ed068ef
2625963
ed068ef
 
 
1be1a61
2625963
ed068ef
 
 
2625963
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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
import os
import cv2
import zipfile
import tempfile
import gradio as gr
import torch

from pathlib import Path
from urllib.request import urlretrieve

from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer


# =========================================
# MODEL
# =========================================

MODEL_URL = (
    "https://github.com/xinntao/Real-ESRGAN/releases/download/"
    "v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"
)

MODEL_PATH = "RealESRGAN_x4plus_anime_6B.pth"


# =========================================
# DOWNLOAD MODEL
# =========================================

if not os.path.exists(MODEL_PATH):
    print("Downloading model...")
    urlretrieve(MODEL_URL, MODEL_PATH)


# =========================================
# LOAD MODEL
# =========================================

model = RRDBNet(
    num_in_ch=3,
    num_out_ch=3,
    num_feat=64,
    num_block=6,
    num_grow_ch=32,
    scale=4,
)

upsampler = RealESRGANer(
    scale=4,
    model_path=MODEL_PATH,
    model=model,
    tile=256,
    tile_pad=10,
    pre_pad=0,
    half=torch.cuda.is_available()
)


# =========================================
# UPSCALE FUNCTION
# =========================================

def upscale(files):

    if not files:
        return None

    temp_dir = tempfile.mkdtemp()

    output_paths = []

    for file in files:

        input_path = file.name

        img = cv2.imread(input_path, cv2.IMREAD_COLOR)

        if img is None:
            continue

        output, _ = upsampler.enhance(
            img,
            outscale=2
        )

        out_name = (
            Path(input_path).stem + "_2x.png"
        )

        out_path = os.path.join(
            temp_dir,
            out_name
        )

        cv2.imwrite(out_path, output)

        output_paths.append(out_path)

    if len(output_paths) == 0:
        return None

    zip_path = os.path.join(
        temp_dir,
        "upscaled_images.zip"
    )

    with zipfile.ZipFile(
        zip_path,
        "w",
        zipfile.ZIP_DEFLATED
    ) as zipf:

        for file in output_paths:

            zipf.write(
                file,
                arcname=os.path.basename(file)
            )

    return zip_path


# =========================================
# UI
# =========================================

with gr.Blocks() as demo:

    gr.Markdown("# Anime Upscaler 2x")

    input_files = gr.File(
        label="Upload Images",
        file_count="multiple",
        file_types=["image"]
    )

    run_button = gr.Button("Upscale")

    output_zip = gr.File(
        label="Download ZIP"
    )

    run_button.click(
        fn=upscale,
        inputs=input_files,
        outputs=output_zip
    )


# =========================================
# LAUNCH
# =========================================

demo.queue()

demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    share=False
)