| import os, re, argparse, threading | |
| import mlx.core as mx | |
| import numpy as np | |
| from PIL import Image, PngImagePlugin | |
| from tqdm import tqdm | |
| from ESRGAN import ESRGAN | |
| def parse_args(): | |
| parser = argparse.ArgumentParser(description="Process tile size, padding, and file paths.") | |
| parser.add_argument('--model', metavar='file_path', type=str, default='4x_NMKD-YandereNeoXL_200k.safetensors', help='Path to the model file') | |
| parser.add_argument('--tile_size', metavar='256', type=int, default=256, help='Size of each tile (default: 256)') | |
| parser.add_argument('--tile_pad', metavar='10', type=int, default=10, help='Padding around each tile (default: 10)') | |
| parser.add_argument('files', metavar='in_file_path', type=str, nargs='+', help='List of file paths to process') | |
| return parser.parse_args() | |
| def load_model(model_path): | |
| model = ESRGAN(mx.load(model_path)) | |
| return mx.compile(model), model.scale | |
| def upscale_img(args, model, file_path, scale=4.0): | |
| ts, tp, s = (args.tile_size, args.tile_pad, scale) | |
| img_in = Image.open(file_path) | |
| png_info = PngImagePlugin.PngInfo() | |
| for k, v in (getattr(img_in, "text", None) or {}).items(): | |
| png_info.add_text(k, v) | |
| img_save_argv = { | |
| "icc_profile": img_in.info.get('icc_profile'), | |
| "pnginfo": png_info, | |
| } | |
| img_in = mx.array(np.array(img_in.convert("RGB"), dtype=np.float32))[None] / 255.0 | |
| _, H, W, C = img_in.shape | |
| mx.eval(img_in) | |
| img_out = mx.zeros((1, H*s, W*s, C), dtype=mx.uint8) | |
| mx.eval(img_out) | |
| for hi, wj in tqdm([(hi, wj) for hi in range(0, H, ts) for wj in range(0, W, ts)]): | |
| phs = min(hi, tp) | |
| pws = min(wj, tp) | |
| img_out[:, hi*4:(hi+ts)*s, wj*4:(wj+ts)*s, :] = ( | |
| model(img_in[:, max(0,hi-tp):hi+ts+tp, max(0, wj-tp):wj+ts+tp, :])[:, phs*s:(ts+phs)*s, pws*s:(ts+pws)*s, :] * 255.0 | |
| ).astype(mx.uint8) | |
| mx.eval(img_out) | |
| img_out = np.array(img_out[0], copy=False) | |
| img_out = Image.fromarray(img_out) | |
| img_out.save(re.sub(r'(\.\w+)$', r'_4x.png', file_path), **img_save_argv) | |
| def main(): | |
| print("\033[1;32mkaeru tiny mlx upscaler v0.1\033[0m") | |
| mx.metal.set_cache_limit(0) | |
| args = parse_args() | |
| model, scale = load_model(args.model) | |
| for file_path in args.files: | |
| print(f"Upscaling {file_path}") | |
| upscale_img(args, model, file_path, scale=scale) | |
| if __name__ == "__main__": | |
| th = threading.Thread(target=main) | |
| th.start() | |
| th.join() | |