Spaces:
Build error
Build error
File size: 1,808 Bytes
6d314be |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import argparse
import os
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
import yaml
from PIL import Image
from torchvision import transforms, utils
from tensorboard_logger import Logger
from tqdm import tqdm
from utils.functions import *
import sys
sys.path.append('pixel2style2pixel/')
from pixel2style2pixel.models.stylegan2.model import Generator, get_keys
torch.backends.cudnn.enabled = True
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True
torch.autograd.set_detect_anomaly(True)
Image.MAX_IMAGE_PIXELS = None
device = torch.device('cuda')
parser = argparse.ArgumentParser()
parser.add_argument('--config', type=str, default='002', help='Path to the config file.')
parser.add_argument('--dataset_path', type=str, default='./data/stylegan2-generate-images/', help='dataset path')
parser.add_argument('--stylegan_model_path', type=str, default='./pixel2style2pixel/pretrained_models/psp_ffhq_encode.pt', help='pretrained stylegan model')
opts = parser.parse_args()
StyleGAN = Generator(1024, 512, 8)
state_dict = torch.load(opts.stylegan_model_path, map_location='cpu')
StyleGAN.load_state_dict(get_keys(state_dict, 'decoder'), strict=True)
StyleGAN.to(device)
#seeds = np.array([torch.random.seed() for i in range(100000)])
seeds = np.load(opts.dataset_path + 'seeds_pytorch_1.8.1.npy')
with torch.no_grad():
os.makedirs(opts.dataset_path + 'ims/', exist_ok=True)
for i, seed in enumerate(tqdm(seeds)):
torch.manual_seed(seed)
z = torch.randn(1, 512).to(device)
n = StyleGAN.make_noise()
w = StyleGAN.get_latent(z)
x, _ = StyleGAN([w], input_is_latent=True, noise=n)
utils.save_image(clip_img(x), opts.dataset_path + 'ims/%06d.jpg'%i)
|