Spaces:
Build error
Build error
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() |