|
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 |
|
|
