def main():
    args = parse_args()
    cfg = Config.fromfile(args.config)
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True
    if (args.seed is not None):
        print('set random seed to', args.seed)
        set_random_seed(args.seed, deterministic=args.deterministic)
    model = build_model(cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg)
    _ = load_checkpoint(model, args.checkpoint, map_location='cpu')
    if (not model.use_ema):
        args.sample_model = 'orig'
    if (args.sample_model == 'ema'):
        generator = model.generator_ema
    else:
        generator = model.generator
    mmcv.print_log(f'Sampling model: {args.sample_model}', 'mmgen')
    generator.eval()
    device = 'cpu'
    if (not args.use_cpu):
        generator = generator.cuda()
        device = 'cuda'
    img_size = min(generator.out_size, 256)
    transform = transforms.Compose([transforms.Resize(img_size), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
    imgs = []
    for imgfile in args.files:
        img = Image.open(imgfile).convert('RGB')
        img = transform(img)
        img = img[([2, 1, 0], ...)]
        imgs.append(img)
    imgs = torch.stack(imgs, 0).to(device)
    with torch.no_grad():
        noise_sample = torch.randn(args.n_mean_latent, generator.style_channels, device=device)
        latent_out = generator.style_mapping(noise_sample)
        latent_mean = latent_out.mean(0)
        latent_std = (((latent_out - latent_mean).pow(2).sum() / args.n_mean_latent) ** 0.5)
    latent_in = latent_mean.detach().clone().unsqueeze(0).repeat(imgs.shape[0], 1)
    if args.w_plus:
        latent_in = latent_in.unsqueeze(1).repeat(1, generator.num_latents, 1)
    latent_in.requires_grad = True
    percept = PerceptualLoss(use_gpu=device.startswith('cuda'))
    noises_single = generator.make_injected_noise()
    noises = []
    for noise in noises_single:
        noises.append(noise.repeat(imgs.shape[0], 1, 1, 1).normal_())
    for noise in noises:
        noise.requires_grad = True
    optimizer = optim.Adam(([latent_in] + noises), lr=args.lr)
    pbar = tqdm(range(args.total_iters))
    for i in pbar:
        t = (i / args.total_iters)
        lr = get_lr(t, args.lr, args.lr_rampdown, args.lr_rampup)
        optimizer.param_groups[0]['lr'] = lr
        noise_strength = ((latent_std * args.noise) * (max(0, (1 - (t / args.noise_ramp))) ** 2))
        latent_n = latent_noise(latent_in, noise_strength.item())
        img_gen = generator([latent_n], input_is_latent=True, injected_noise=noises)
        (batch, channel, height, width) = img_gen.shape
        if (height > 256):
            factor = (height // 256)
            img_gen = img_gen.reshape(batch, channel, (height // factor), factor, (width // factor), factor)
            img_gen = img_gen.mean([3, 5])
        p_loss = percept(img_gen, imgs).sum()
        n_loss = noise_regularize(noises)
        mse_loss = F.mse_loss(img_gen, imgs)
        loss = ((p_loss + (args.noise_regularize * n_loss)) + (args.mse * mse_loss))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        noise_normalize_(noises)
        pbar.set_description(f' perceptual: {p_loss.item():.4f}, noise regularize:{n_loss.item():.4f}, mse: {mse_loss.item():.4f}, lr: {lr:.4f}')
    results = generator([latent_in.detach().clone()], input_is_latent=True, injected_noise=noises)
    results = ((results + 1) / 2)
    results = results[:, [2, 1, 0], ...]
    results = results.clamp_(0, 1)
    mmcv.mkdir_or_exist(args.results_path)
    result_file = {}
    for (i, input_name) in enumerate(args.files):
        noise_single = []
        for noise in noises:
            noise_single.append(noise[i:(i + 1)])
        result_file[input_name] = {'img': img_gen[i], 'latent': latent_in[i], 'injected_noise': noise_single}
        img_name = (os.path.splitext(os.path.basename(input_name))[0] + '-project.png')
        save_image(results[i], os.path.join(args.results_path, img_name))
    torch.save(result_file, os.path.join(args.results_path, 'project_result.pt'))
def main():
    args = parse_args()
    cfg = Config.fromfile(args.config)
    if cfg.get('cudnn_benchmark', False):
        torch.backends.cudnn.benchmark = True
    if (args.seed is not None):
        print('set random seed to', args.seed)
        set_random_seed(args.seed, deterministic=args.deterministic)
    model = build_model(cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg)
    _ = load_checkpoint(model, args.checkpoint, map_location='cpu')
    if (not model.use_ema):
        args.sample_model = 'orig'
    if (args.sample_model == 'ema'):
        generator = model.generator_ema
    else:
        generator = model.generator
    mmcv.print_log(f'Sampling model: {args.sample_model}', 'mmgen')
    generator.eval()
    device = 'cpu'
    if (not args.use_cpu):
        generator = generator.cuda()
        device = 'cuda'
    img_size = min(generator.out_size, 256)
    transform = transforms.Compose([transforms.Resize(img_size), transforms.ToTensor(), transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
    imgs = []
    for imgfile in args.files:
        img = Image.open(imgfile).convert('RGB')
        img = transform(img)
        img = img[([2, 1, 0], ...)]
        imgs.append(img)
    imgs = torch.stack(imgs, 0).to(device)
    with torch.no_grad():
        noise_sample = torch.randn(args.n_mean_latent, generator.style_channels, device=device)
        latent_out = generator.style_mapping(noise_sample)
        latent_mean = latent_out.mean(0)
        latent_std = (((latent_out - latent_mean).pow(2).sum() / args.n_mean_latent) ** 0.5)
    latent_in = latent_mean.detach().clone().unsqueeze(0).repeat(imgs.shape[0], 1)
    if args.w_plus:
        latent_in = latent_in.unsqueeze(1).repeat(1, generator.num_latents, 1)
    latent_in.requires_grad = True
    percept = PerceptualLoss(use_gpu=device.startswith('cuda'))
    noises_single = generator.make_injected_noise()
    noises = []
    for noise in noises_single:
        noises.append(noise.repeat(imgs.shape[0], 1, 1, 1).normal_())
    for noise in noises:
        noise.requires_grad = True
    optimizer = optim.Adam(([latent_in] + noises), lr=args.lr)
    pbar = tqdm(range(args.total_iters))
    for i in pbar:
        t = (i / args.total_iters)
        lr = get_lr(t, args.lr, args.lr_rampdown, args.lr_rampup)
        optimizer.param_groups[0]['lr'] = lr
        noise_strength = ((latent_std * args.noise) * (max(0, (1 - (t / args.noise_ramp))) ** 2))
        latent_n = latent_noise(latent_in, noise_strength.item())
        img_gen = generator([latent_n], input_is_latent=True, injected_noise=noises)
        (batch, channel, height, width) = img_gen.shape
        if (height > 256):
            factor = (height // 256)
            img_gen = img_gen.reshape(batch, channel, (height // factor), factor, (width // factor), factor)
            img_gen = img_gen.mean([3, 5])
        p_loss = percept(img_gen, imgs).sum()
        n_loss = noise_regularize(noises)
        mse_loss = F.mse_loss(img_gen, imgs)
        loss = ((p_loss + (args.noise_regularize * n_loss)) + (args.mse * mse_loss))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        noise_normalize_(noises)
        pbar.set_description(f' perceptual: {p_loss.item():.4f}, noise regularize:{n_loss.item():.4f}, mse: {mse_loss.item():.4f}, lr: {lr:.4f}')
    results = generator([latent_in.detach().clone()], input_is_latent=True, injected_noise=noises)
    results = ((results + 1) / 2)
    results = results[:, [2, 1, 0], ...]
    results = results.clamp_(0, 1)
    mmcv.mkdir_or_exist(args.results_path)
    result_file = {}
    for (i, input_name) in enumerate(args.files):
        noise_single = []
        for noise in noises:
            noise_single.append(noise[i:(i + 1)])
        result_file[input_name] = {'img': img_gen[i], 'latent': latent_in[i], 'injected_noise': noise_single}
        img_name = (os.path.splitext(os.path.basename(input_name))[0] + '-project.png')
        save_image(results[i], os.path.join(args.results_path, img_name))
    torch.save(result_file, os.path.join(args.results_path, 'project_result.pt'))