File size: 2,555 Bytes
2d7efb8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
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)
|