| import glob |
| from argparse import Namespace |
| from configs import paths_config |
| from models.StyleCLIP.mapper.scripts.inference import run |
| from scripts.run_pti import run_PTI |
|
|
| meta_data = { |
| 'afro': ['afro', False, False, True], |
| 'angry': ['angry', False, False, True], |
| 'Beyonce': ['beyonce', False, False, False], |
| 'bobcut': ['bobcut', False, False, True], |
| 'bowlcut': ['bowlcut', False, False, True], |
| 'curly hair': ['curly_hair', False, False, True], |
| 'Hilary Clinton': ['hilary_clinton', False, False, False], |
| 'Jhonny Depp': ['depp', False, False, False], |
| 'mohawk': ['mohawk', False, False, True], |
| 'purple hair': ['purple_hair', False, False, False], |
| 'surprised': ['surprised', False, False, True], |
| 'Taylor Swift': ['taylor_swift', False, False, False], |
| 'trump': ['trump', False, False, False], |
| 'Mark Zuckerberg': ['zuckerberg', False, False, False] |
| } |
|
|
|
|
| def styleclip_edit(use_multi_id_G, run_id, use_wandb, edit_types): |
| images_dir = paths_config.input_data_path |
| pretrained_mappers = paths_config.style_clip_pretrained_mappers |
| data_dir_name = paths_config.input_data_id |
| if run_id == '': |
| run_id = run_PTI(run_name='', use_wandb=use_wandb, use_multi_id_training=False) |
| images = glob.glob(f"{images_dir}/*.jpeg") |
| w_path_dir = f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}' |
| for image_name in images: |
| image_name = image_name.split(".")[0].split("/")[-1] |
| embedding_dir = f'{w_path_dir}/{paths_config.pti_results_keyword}/{image_name}' |
| latent_path = f'{embedding_dir}/0.pt' |
| for edit_type in set(meta_data.keys()).intersection(edit_types): |
| edit_id = meta_data[edit_type][0] |
| args = { |
| "exp_dir": f'{paths_config.styleclip_output_dir}', |
| "checkpoint_path": f"{pretrained_mappers}/{edit_id}.pt", |
| "couple_outputs": False, |
| "mapper_type": "LevelsMapper", |
| "no_coarse_mapper": meta_data[edit_type][1], |
| "no_medium_mapper": meta_data[edit_type][2], |
| "no_fine_mapper": meta_data[edit_type][3], |
| "stylegan_size": 1024, |
| "test_batch_size": 1, |
| "latents_test_path": latent_path, |
| "test_workers": 1, |
| "run_id": run_id, |
| "image_name": image_name, |
| 'edit_name': edit_type, |
| "data_dir_name": data_dir_name |
| } |
|
|
| run(Namespace(**args), run_id, image_name, use_multi_id_G) |
|
|
