| import torch | |
| def edit(latents, pca, edit_directions): | |
| edit_latents = [] | |
| for latent in latents: | |
| for pca_idx, start, end, strength in edit_directions: | |
| delta = get_delta(pca, latent, pca_idx, strength) | |
| delta_padded = torch.zeros(latent.shape).to("cuda") | |
| delta_padded[start:end] += delta.repeat(end - start, 1) | |
| edit_latents.append(latent + delta_padded) | |
| return torch.stack(edit_latents) | |
| def get_delta(pca, latent, idx, strength): | |
| w_centered = latent - pca["mean"].to("cuda") | |
| lat_comp = pca["comp"].to("cuda") | |
| lat_std = pca["std"].to("cuda") | |
| w_coord = ( | |
| torch.sum(w_centered[0].reshape(-1) * lat_comp[idx].reshape(-1)) / lat_std[idx] | |
| ) | |
| delta = (strength - w_coord) * lat_comp[idx] * lat_std[idx] | |
| return delta | |
| def edit_latent(latent, pca, edit_direction): | |
| pca_idx, start, end, strength = edit_direction | |
| delta = get_delta(pca, latent, pca_idx, strength) | |
| delta_padded = torch.zeros(latent.shape).to("cuda") | |
| delta_padded[start:end] += delta.repeat(end - start, 1) | |
| edit_latent = latent + delta_padded | |
| return edit_latent | |