| import torch | |
| import argparse | |
| from models.psp import pSp | |
| from models.encoders.psp_encoders import Encoder4Editing | |
| def setup_model(checkpoint_path, device='cuda'): | |
| ckpt = torch.load(checkpoint_path, map_location='cpu') | |
| opts = ckpt['opts'] | |
| opts['checkpoint_path'] = checkpoint_path | |
| opts['device'] = device | |
| opts = argparse.Namespace(**opts) | |
| net = pSp(opts) | |
| net.eval() | |
| net = net.to(device) | |
| return net, opts | |
| def load_e4e_standalone(checkpoint_path, device='cuda'): | |
| ckpt = torch.load(checkpoint_path, map_location='cpu') | |
| opts = argparse.Namespace(**ckpt['opts']) | |
| e4e = Encoder4Editing(50, 'ir_se', opts) | |
| e4e_dict = {k.replace('encoder.', ''): v for k, v in ckpt['state_dict'].items() if k.startswith('encoder.')} | |
| e4e.load_state_dict(e4e_dict) | |
| e4e.eval() | |
| e4e = e4e.to(device) | |
| latent_avg = ckpt['latent_avg'].to(device) | |
| def add_latent_avg(model, inputs, outputs): | |
| return outputs + latent_avg.repeat(outputs.shape[0], 1, 1) | |
| e4e.register_forward_hook(add_latent_avg) | |
| return e4e | |