| import os | |
| import torch | |
| import modules.core as core | |
| from ldm_patched.pfn.architecture.RRDB import RRDBNet as ESRGAN | |
| from ldm_patched.contrib.external_upscale_model import ImageUpscaleWithModel | |
| from collections import OrderedDict | |
| from modules.config import path_upscale_models | |
| model_filename = os.path.join(path_upscale_models, 'fooocus_upscaler_s409985e5.bin') | |
| opImageUpscaleWithModel = ImageUpscaleWithModel() | |
| model = None | |
| def perform_upscale(img): | |
| global model | |
| print(f'Upscaling image with shape {str(img.shape)} ...') | |
| if model is None: | |
| sd = torch.load(model_filename) | |
| sdo = OrderedDict() | |
| for k, v in sd.items(): | |
| sdo[k.replace('residual_block_', 'RDB')] = v | |
| del sd | |
| model = ESRGAN(sdo) | |
| model.cpu() | |
| model.eval() | |
| img = core.numpy_to_pytorch(img) | |
| img = opImageUpscaleWithModel.upscale(model, img)[0] | |
| img = core.pytorch_to_numpy(img)[0] | |
| return img | |