| import glob |
| import os |
| from configs import global_config, paths_config, hyperparameters |
| from scripts.latent_creators.sg2_plus_latent_creator import SG2PlusLatentCreator |
| from scripts.latent_creators.e4e_latent_creator import E4ELatentCreator |
| from scripts.run_pti import run_PTI |
| import pickle |
| import torch |
| from utils.models_utils import toogle_grad, load_old_G |
|
|
|
|
| class ExperimentRunner: |
|
|
| def __init__(self, run_id=''): |
| self.images_paths = glob.glob(f'{paths_config.input_data_path}/*') |
| self.target_paths = glob.glob(f'{paths_config.input_data_path}/*') |
| self.run_id = run_id |
| self.sampled_ws = None |
|
|
| self.old_G = load_old_G() |
|
|
| toogle_grad(self.old_G, False) |
|
|
| def run_experiment(self, run_pt, create_other_latents, use_multi_id_training, use_wandb=False): |
| if run_pt: |
| self.run_id = run_PTI(self.run_id, use_wandb=use_wandb, use_multi_id_training=use_multi_id_training) |
| if create_other_latents: |
| sg2_plus_latent_creator = SG2PlusLatentCreator(use_wandb=use_wandb) |
| sg2_plus_latent_creator.create_latents() |
| e4e_latent_creator = E4ELatentCreator(use_wandb=use_wandb) |
| e4e_latent_creator.create_latents() |
|
|
| torch.cuda.empty_cache() |
|
|
| return self.run_id |
|
|
|
|
| if __name__ == '__main__': |
| os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID' |
| os.environ['CUDA_VISIBLE_DEVICES'] = global_config.cuda_visible_devices |
|
|
| runner = ExperimentRunner() |
| runner.run_experiment(True, False, False) |
|
|
