class Predictor(cog.Predictor):
def setup(self):
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
self.generator = StyledGenerator(512).to(self.device)
print("Loading checkpoint")
self.generator.load_state_dict(
torch.load(
"stylegan-1024px-new.model",
map_location=self.device,
)["g_running"], )
self.generator.eval()
@cog.input("seed", type=int, default=-1, help="Random seed, -1 for random")
def predict(self, seed):
if seed < 0:
seed = int.from_bytes(os.urandom(2), "big")
torch.manual_seed(seed)
print(f"seed: {seed}")
mean_style = get_mean_style(self.generator, self.device)
step = int(math.log(SIZE, 2)) - 2
img = sample(self.generator, step, mean_style, 1, self.device)
output_path = Path(tempfile.mkdtemp()) / "output.png"
utils.save_image(img, output_path, normalize=True)
return output_path
def loadStyleGAN():
sys.path.append(StyleGAN1_root)
ckpt_root = join(StyleGAN1_root, 'checkpoint')
from model import StyledGenerator
from generate import get_mean_style
import math
generator = StyledGenerator(512).to("cuda")
# generator.load_state_dict(torch.load(r"E:\Github_Projects\style-based-gan-pytorch\checkpoint\stylegan-256px-new.model")['g_running'])
generator.load_state_dict(
torch.load(join(StyleGAN1_root,
"checkpoint\stylegan-256px-new.model"))['g_running'])
generator.eval()
for param in generator.parameters():
param.requires_grad_(False)
return generator
type=int,
default=3,
help='number of rows of sample matrix')
parser.add_argument('--n_col',
type=int,
default=5,
help='number of columns of sample matrix')
parser.add_argument('path', type=str, help='path to checkpoint file')
args = parser.parse_args()
device = 'cuda'
generator = StyledGenerator(512).to(device)
generator.load_state_dict(torch.load(args.path)['g_running'])
generator.eval()
mean_style = get_mean_style(generator, device)
step = int(math.log(args.size, 2)) - 2
img = sample(generator, step, mean_style, args.n_row * args.n_col, device)
utils.save_image(img,
'sample.png',
nrow=args.n_col,
normalize=True,
range=(-1, 1))
for j in range(20):
img = style_mixing(generator, step, mean_style, args.n_col, args.n_row,
device)
transforms.CenterCrop(resize),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
]
)
imgs = []
for imgfile in args.files:
img = transform(Image.open(imgfile).convert("RGB"))
imgs.append(img)
imgs = torch.stack(imgs, 0).to(device)
g_ema = StyledGenerator(512)
g_ema.load_state_dict(torch.load(args.ckpt)["g_running"], strict=False)
g_ema.eval()
g_ema = g_ema.to(device)
step = int(math.log(args.size, 2)) - 2
with torch.no_grad():
noise_sample = torch.randn(n_mean_latent, 512, device=device)
latent_out = g_ema.style(noise_sample)
latent_mean = latent_out.mean(0)
latent_std = ((latent_out - latent_mean).pow(2).sum() / n_mean_latent) ** 0.5
percept = lpips.PerceptualLoss(
model="net-lin", net="vgg", use_gpu=device.startswith("cuda")
)
noises_single = make_noise(device,args.size)
noises = []
def optimize_latents():
print("Optimizing Latents.")
generator = StyledGenerator(512).to(device)
generator.load_state_dict(torch.load(args.path)['generator'])
generator.eval()
latent_optimizer = LatentOptimizer(generator, args.vgg_layer)
mean_style = get_mean_style(generator, device)
total = np.zeros((83 * 3, 512))
# Optimize only the dlatents.
for param in latent_optimizer.parameters():
param.requires_grad_(False)
if args.video or args.save_optimized_image:
# Hook, saves an image during optimization to be used to create video.
generated_image_hook = GeneratedImageHook(
latent_optimizer.post_synthesis_processing, args.save_frequency)
for i in range(3 * 83): #3 for each pictrue
iid = i % 3
path = int(i / 3)
iterations = int(200 * iid + 300)
image_path = './data/' + str(path) + '.jpg'
print(image_path)
reference_image = load_images([image_path])
reference_image = torch.from_numpy(reference_image).to(device)
reference_image = latent_optimizer.vgg_processing(
reference_image) #normalize
reference_features = latent_optimizer.vgg16(
reference_image).detach() #vgg
reference_image = reference_image.detach()
if args.use_latent_finder:
image_to_latent = ImageToLatent().cuda()
image_to_latent.load_state_dict(
torch.load(args.image_to_latent_path))
image_to_latent.eval()
latents_to_be_optimized = image_to_latent(reference_image)
latents_to_be_optimized = latents_to_be_optimized.detach().cuda(
).requires_grad_(True)
else:
latents_to_be_optimized = torch.zeros(
(1, 512)).cuda().requires_grad_(True)
criterion = LatentLoss()
optimizer = torch.optim.SGD([latents_to_be_optimized],
lr=args.learning_rate)
progress_bar = tqdm(range(iterations))
for step in progress_bar:
optimizer.zero_grad()
generated_image_features = latent_optimizer(
latents_to_be_optimized, mean_style, i)
#print(latents_to_be_optimized)
loss = criterion(generated_image_features, reference_features)
loss.backward()
loss = loss.item()
optimizer.step()
progress_bar.set_description("Step: {}, Loss: {}".format(
step, loss))
optimized_dlatents = latents_to_be_optimized.detach().cpu().numpy()
total[i] = optimized_dlatents[0]
np.save(args.dlatent_path, total)
def optimize_latents():
print("Optimizing Latents.")
generator = StyledGenerator(512).to(device)
generator.load_state_dict(torch.load(args.path)['generator'])
generator.eval()
latent_optimizer = LatentOptimizer(generator, args.vgg_layer)
mean_style = get_mean_style(generator, device)
# Optimize only the dlatents.
for param in latent_optimizer.parameters():
param.requires_grad_(False)
if args.video or args.save_optimized_image:
# Hook, saves an image during optimization to be used to create video.
generated_image_hook = GeneratedImageHook(
latent_optimizer.post_synthesis_processing, args.save_frequency)
reference_image = load_images([args.image_path])
reference_image = torch.from_numpy(reference_image).to(device)
reference_image = latent_optimizer.vgg_processing(
reference_image) #normalize
utils.save_image(reference_image,
'./reference.png',
nrow=1,
normalize=True,
range=(-1, 1))
reference_features = latent_optimizer.vgg16(reference_image).detach() #vgg
reference_image = reference_image.detach()
if args.use_latent_finder:
image_to_latent = ImageToLatent().cuda()
image_to_latent.load_state_dict(torch.load(args.image_to_latent_path))
image_to_latent.eval()
latents_to_be_optimized = image_to_latent(reference_image)
latents_to_be_optimized = latents_to_be_optimized.detach().cuda(
).requires_grad_(True)
else:
latents_to_be_optimized = torch.zeros(
(1, 512)).cuda().requires_grad_(True)
#print(latents_to_be_optimized.mean(),latents_to_be_optimized.std())
criterion = LatentLoss()
optimizer = torch.optim.RMSprop([latents_to_be_optimized],
lr=args.learning_rate,
weight_decay=0.02)
#print(latents_to_be_optimized.mean(),latents_to_be_optimized.std())
progress_bar = tqdm(range(args.iterations))
for step in progress_bar:
#print(latents_to_be_optimized)
optimizer.zero_grad()
generated_image_features = latent_optimizer(latents_to_be_optimized,
mean_style)
loss = criterion(generated_image_features, reference_features)
loss.backward()
loss = loss.item()
optimizer.step()
# if step==args.iterations:
# break
# with torch.no_grad():
# latents_to_be_optimized.add_(-latents_to_be_optimized.mean()+3e-2*torch.randn(1).to('cuda'))
# latents_to_be_optimized.div_(latents_to_be_optimized.std()+3e-2*torch.randn(1).to('cuda'))
#print(latents_to_be_optimized.mean(),latents_to_be_optimized.std())
progress_bar.set_description("Step: {}, Loss: {}".format(step, loss))
print(latents_to_be_optimized)
#latents_to_be_optimized=latent_optimizer.normalize(latents_to_be_optimized)
#print(latents_to_be_optimized.mean(),latents_to_be_optimized.std())
optimized_dlatents = latents_to_be_optimized.detach().cpu().numpy()
np.save(args.dlatent_path, optimized_dlatents)
if args.video:
images_to_video(generated_image_hook.get_images(), args.video_path)
if args.save_optimized_image:
save_image(generated_image_hook.last_image, args.optimized_image_path)
batch_default = 32
phase = 150_000
max_iter = 100_000
transform = transform.Compose([
transform.ToPILImage(),
transform.RandomHorizontalFlip(),
transform.ToTensor(),
transform.ImageNormalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
netG = StyledGenerator(code_dim=code_size)
netD = Discriminator(from_rgb_activate=True)
g_running = StyledGenerator(code_size)
g_running.eval()
d_optimizer = jt.optim.Adam(netD.parameters(), lr=lr, betas=(0.0, 0.99))
g_optimizer = jt.optim.Adam(netG.generator.parameters(),
lr=lr,
betas=(0.0, 0.99))
g_optimizer.add_param_group({
'params': netG.style.parameters(),
'lr': lr * 0.01,
'mult': 0.01,
})
accumulate(g_running, netG, 0)
if args.ckpt is not None:
ckpt = jt.load(args.ckpt)