TSE1966/Super-Resolution-Anime-Diffusion
0
1import cv22import numpy as np3from PIL import Image4import glob5import os6from basicsr.archs.rrdbnet_arch import RRDBNet7from basicsr.utils.download_util import load_file_from_url8 9from realesrgan import RealESRGANer10from realesrgan.archs.srvgg_arch import SRVGGNetCompact11 12 13def realEsrgan(14 model_name="RealESRGAN_x4plus_anime_6B",15 model_path=None,16 input_dir="inputs",17 output_dir="results",18 denoise_strength=0.5,19 outscale=4,20 suffix="out",21 tile=200,22 tile_pad=10,23 pre_pad=0,24 face_enhance=True,25 alpha_upsampler="realsrgan",26 out_ext="auto",27 fp32=True,28 gpu_id=None,29):30 31 # determine models according to model names32 model_name = model_name.split(".")[0]33 if model_name == "RealESRGAN_x4plus": # x4 RRDBNet model34 model = RRDBNet(35 num_in_ch=3,36 num_out_ch=3,37 num_feat=64,38 num_block=23,39 num_grow_ch=32,40 scale=4,41 )42 netscale = 443 file_url = [44 "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"45 ]46 elif model_name == "RealESRNet_x4plus": # x4 RRDBNet model47 model = RRDBNet(48 num_in_ch=3,49 num_out_ch=3,50 num_feat=64,51 num_block=23,52 num_grow_ch=32,53 scale=4,54 )55 netscale = 456 file_url = [57 "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth"58 ]59 elif model_name == "RealESRGAN_x4plus_anime_6B": # x4 RRDBNet model with 6 blocks60 model = RRDBNet(61 num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=462 )63 netscale = 464 file_url = [65 "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth"66 ]67 elif model_name == "RealESRGAN_x2plus": # x2 RRDBNet model68 model = RRDBNet(69 num_in_ch=3,70 num_out_ch=3,71 num_feat=64,72 num_block=23,73 num_grow_ch=32,74 scale=2,75 )76 netscale = 277 file_url = [78 "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"79 ]80 elif model_name == "realesr-animevideov3": # x4 VGG-style model (XS size)81 model = SRVGGNetCompact(82 num_in_ch=3,83 num_out_ch=3,84 num_feat=64,85 num_conv=16,86 upscale=4,87 act_type="prelu",88 )89 netscale = 490 file_url = [91 "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth"92 ]93 elif model_name == "realesr-general-x4v3": # x4 VGG-style model (S size)94 model = SRVGGNetCompact(95 num_in_ch=3,96 num_out_ch=3,97 num_feat=64,98 num_conv=32,99 upscale=4,100 act_type="prelu",101 )102 netscale = 4103 file_url = [104 "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth",105 "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",106 ]107 108 # determine model paths109 if model_path is None:110 model_path = os.path.join("weights", model_name + ".pth")111 if not os.path.isfile(model_path):112 ROOT_DIR = os.path.dirname(os.path.abspath(__file__))113 for url in file_url:114 # model_path will be updated115 model_path = load_file_from_url(116 url=url,117 model_dir=os.path.join(ROOT_DIR, "weights"),118 progress=True,119 file_name=None,120 )121 122 # use dni to control the denoise strength123 dni_weight = None124 if model_name == "realesr-general-x4v3" and denoise_strength != 1:125 wdn_model_path = model_path.replace(126 "realesr-general-x4v3", "realesr-general-wdn-x4v3"127 )128 model_path = [model_path, wdn_model_path]129 dni_weight = [denoise_strength, 1 - denoise_strength]130 131 # restorer132 upsampler = RealESRGANer(133 scale=netscale,134 model_path=model_path,135 dni_weight=dni_weight,136 model=model,137 tile=tile,138 tile_pad=tile_pad,139 pre_pad=pre_pad,140 half=not fp32,141 gpu_id=gpu_id,142 )143 144 if face_enhance: # Use GFPGAN for face enhancement145 from gfpgan import GFPGANer146 147 face_enhancer = GFPGANer(148 model_path="https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth",149 upscale=outscale,150 arch="clean",151 channel_multiplier=2,152 bg_upsampler=upsampler,153 )154 os.makedirs(output_dir, exist_ok=True)155 156 if not isinstance(input_dir, list):157 paths = [input_dir]158 else:159 paths = sorted(glob.glob(os.path.join(input_dir, "*")))160 161 Imgs = []162 for idx, path in enumerate(paths):163 print(f"Scaling x{outscale}:", path)164 if isinstance(path, Image.Image):165 img = path166 img = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)167 imgname = f"img_{idx}"168 else:169 imgname, extension = os.path.splitext(os.path.basename(path))170 img = cv2.imread(path, cv2.IMREAD_UNCHANGED)171 if len(img.shape) == 3 and img.shape[2] == 4:172 img_mode = "RGBA"173 else:174 img_mode = None175 176 try:177 if face_enhance:178 _, _, output = face_enhancer.enhance(179 img, has_aligned=False, only_center_face=False, paste_back=True180 )181 else:182 output, _ = upsampler.enhance(img, outscale=outscale)183 except RuntimeError as error:184 print("Error", error)185 print(186 "If you encounter CUDA or RAM out of memory, try to set --tile with a smaller number."187 )188 else:189 # if out_ext == "auto":190 # extension = extension[1:]191 # else:192 # extension = out_ext193 # if img_mode == "RGBA": # RGBA images should be saved in png format194 # extension = "png"195 # if suffix == "":196 # save_path = os.path.join(output_dir, f"{imgname}.{extension}")197 # else:198 # save_path = os.path.join(output_dir, f"{imgname}_{suffix}.{extension}")199 #200 # cv2.imwrite(save_path, output)201 202 img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)203 img = Image.fromarray(img)204 Imgs.append(img)205 206 return Imgs207