File size: 2,771 Bytes
f62cb15
 
 
 
 
 
 
 
 
 
 
 
 
 
9e228c1
f62cb15
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import gradio as gr
import torch
import os
from PIL import Image
import numpy as np
from huggingface_hub import hf_hub_download, list_repo_files

# -----------------------------
# CONFIG
# -----------------------------
REPO_ID = "easygoing0114/AI_upscalers"
MODEL_DIR = "/tmp/upscalers"
os.makedirs(MODEL_DIR, exist_ok=True)

DEVICE = "cpu"

# -----------------------------
# LOAD MODEL LIST
# -----------------------------
def get_models():
    files = list_repo_files(REPO_ID)
    models = [f for f in files if f.endswith(".pth")]
    return models

MODEL_LIST = get_models()

# -----------------------------
# DOWNLOAD MODEL
# -----------------------------
def download_model(model_name):
    path = hf_hub_download(
        repo_id=REPO_ID,
        filename=model_name,
        local_dir=MODEL_DIR
    )
    return path

# -----------------------------
# LOAD UPSCALER (GENERIC ESRGAN STYLE)
# -----------------------------
def load_model(model_path):
    # Lazy import to avoid heavy startup
    from basicsr.archs.rrdbnet_arch import RRDBNet
    from realesrgan import RealESRGANer

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

    upsampler = RealESRGANer(
        scale=4,
        model_path=model_path,
        model=model,
        tile=0,
        tile_pad=10,
        pre_pad=0,
        half=(DEVICE == "cuda"),
        device=DEVICE
    )

    return upsampler

# -----------------------------
# UPSCALE FUNCTION
# -----------------------------
def upscale_image(image, model_name):
    if image is None:
        return None

    # Download model
    model_path = download_model(model_name)

    # Load model
    upsampler = load_model(model_path)

    # Convert image
    img = np.array(image)

    # Upscale
    output, _ = upsampler.enhance(img, outscale=4)

    return Image.fromarray(output)

# -----------------------------
# UI
# -----------------------------
with gr.Blocks(title="AI Image Upscaler") as app:
    gr.Markdown("# 🔍 AI Image Upscaler (Multi-Model)")
    gr.Markdown("Select any model from the repository and upscale your image.")

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Image")

    model_dropdown = gr.Dropdown(
        choices=MODEL_LIST,
        value=MODEL_LIST[0] if MODEL_LIST else None,
        label="Select Upscaler Model"
    )

    upscale_btn = gr.Button("✨ Upscale Image")

    output_image = gr.Image(label="Upscaled Image")

    upscale_btn.click(
        fn=upscale_image,
        inputs=[image_input, model_dropdown],
        outputs=output_image
    )

# -----------------------------
# LAUNCH
# -----------------------------
if __name__ == "__main__":
    app.launch()