| import torch |
| from configs import paths_config |
| from editings.latent_editor import LatentEditor |
|
|
|
|
| class LatentEditorWrapper: |
|
|
| def __init__(self): |
|
|
| self.interfacegan_directions = {'age': f'{paths_config.interfacegan_age}', |
| 'smile': f'{paths_config.interfacegan_smile}', |
| 'rotation': f'{paths_config.interfacegan_rotation}'} |
| self.interfacegan_directions_tensors = {name: torch.load(path).cuda() for name, path in |
| self.interfacegan_directions.items()} |
| self.ganspace_pca = torch.load(f'{paths_config.ffhq_pca}') |
|
|
| |
| self.ganspace_directions = { |
| 'eye_openness': (54, 7, 8, 5), |
| 'smile': (46, 4, 5, -6), |
| 'trimmed_beard': (58, 7, 9, 7), |
| } |
|
|
| self.latent_editor = LatentEditor() |
|
|
| def get_single_ganspace_edits(self, start_w, factors): |
| latents_to_display = [] |
| for ganspace_direction in self.ganspace_directions.values(): |
| for factor in factors: |
| edit_direction = list(ganspace_direction) |
| edit_direction[-1] = factor |
| edit_direction = tuple(edit_direction) |
| new_w = self.latent_editor.apply_ganspace(start_w, self.ganspace_pca, [edit_direction]) |
| latents_to_display.append(new_w) |
| return latents_to_display |
|
|
| def get_single_interface_gan_edits(self, start_w, factors): |
| latents_to_display = {} |
| for direction in ['rotation', 'smile', 'age']: |
| for factor in factors: |
| if direction not in latents_to_display: |
| latents_to_display[direction] = {} |
| latents_to_display[direction][factor] = self.latent_editor.apply_interfacegan( |
| start_w, self.interfacegan_directions_tensors[direction], factor / 2) |
|
|
| return latents_to_display |
|
|
