| import torch |
| import numpy as np |
| from tqdm import tqdm |
|
|
|
|
| def edit(generator, latents, indices, semantics=1, start_distance=-15.0, end_distance=15.0, num_samples=1, step=11): |
|
|
| layers, boundaries, values = factorize_weight(generator, indices) |
| codes = latents.detach().cpu().numpy() |
|
|
| |
| distances = np.linspace(start_distance, end_distance, step) |
| num_sam = num_samples |
| num_sem = semantics |
|
|
| edited_latents = [] |
| for sem_id in tqdm(range(num_sem), desc='Semantic ', leave=False): |
| boundary = boundaries[sem_id:sem_id + 1] |
| for sam_id in tqdm(range(num_sam), desc='Sample ', leave=False): |
| code = codes[sam_id:sam_id + 1] |
| for col_id, d in enumerate(distances, start=1): |
| temp_code = code.copy() |
| temp_code[:, layers, :] += boundary * d |
| edited_latents.append(torch.from_numpy(temp_code).float().cuda()) |
| return torch.cat(edited_latents) |
|
|
|
|
| def factorize_weight(g_ema, layers='all'): |
|
|
| weights = [] |
| if layers == 'all' or 0 in layers: |
| weight = g_ema.conv1.conv.modulation.weight.T |
| weights.append(weight.cpu().detach().numpy()) |
|
|
| if layers == 'all': |
| layers = list(range(g_ema.num_layers - 1)) |
| else: |
| layers = [l - 1 for l in layers if l != 0] |
|
|
| for idx in layers: |
| weight = g_ema.convs[idx].conv.modulation.weight.T |
| weights.append(weight.cpu().detach().numpy()) |
| weight = np.concatenate(weights, axis=1).astype(np.float32) |
| weight = weight / np.linalg.norm(weight, axis=0, keepdims=True) |
| eigen_values, eigen_vectors = np.linalg.eig(weight.dot(weight.T)) |
| return layers, eigen_vectors.T, eigen_values |
|
|
