HairFastGAN

1
import argparse
2
import os
3
import numpy as np
4
import torch
5
import torch.nn as nn
6
import torch.nn.functional as F
7
import torch.utils.data as data
8
import yaml
9

10
from PIL import Image
11
from torchvision import transforms, utils
12
from tensorboard_logger import Logger
13
from tqdm import tqdm
14
from utils.functions import *
15

16
import sys
17
sys.path.append('pixel2style2pixel/')
18
from pixel2style2pixel.models.stylegan2.model import Generator, get_keys
19

20
torch.backends.cudnn.enabled = True
21
torch.backends.cudnn.deterministic = True
22
torch.backends.cudnn.benchmark = True
23
torch.autograd.set_detect_anomaly(True)
24
Image.MAX_IMAGE_PIXELS = None
25
device = torch.device('cuda')
26

27
parser = argparse.ArgumentParser()
28
parser.add_argument('--config', type=str, default='002', help='Path to the config file.')
29
parser.add_argument('--dataset_path', type=str, default='./data/stylegan2-generate-images/', help='dataset path')
30
parser.add_argument('--stylegan_model_path', type=str, default='./pixel2style2pixel/pretrained_models/psp_ffhq_encode.pt', help='pretrained stylegan model')
31
opts = parser.parse_args()
32

33

34
StyleGAN = Generator(1024, 512, 8)
35
state_dict = torch.load(opts.stylegan_model_path, map_location='cpu')
36
StyleGAN.load_state_dict(get_keys(state_dict, 'decoder'), strict=True)
37
StyleGAN.to(device)
38

39
#seeds = np.array([torch.random.seed() for i in range(100000)])
40
seeds = np.load(opts.dataset_path + 'seeds_pytorch_1.8.1.npy')
41

42
with torch.no_grad():
43
    os.makedirs(opts.dataset_path + 'ims/', exist_ok=True)
44

45
    for i, seed in enumerate(tqdm(seeds)):
46

47
        torch.manual_seed(seed)
48
        z = torch.randn(1, 512).to(device)
49
        n = StyleGAN.make_noise()
50
        w = StyleGAN.get_latent(z)
51
        x, _ = StyleGAN([w], input_is_latent=True, noise=n)
52
        utils.save_image(clip_img(x), opts.dataset_path + 'ims/%06d.jpg'%i)
53