Update tune.py

tune.py

@@ -24,9 +24,26 @@ hyperparameters.lpips_type = 'squeeze'
 
 from scripts.run_pti import run_PTI
 
+def load_generator(model_id):
+  with open(f'{paths_config.checkpoints_dir}/model_{model_id}_file.pt', 'rb') as f_new: 
+    new_G = torch.load(f_new).cuda()
+  return new_G
+  
+def tensor_to_pil(img):
+  img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0] 
+  plt.axis('off') 
+  resized_image = Image.fromarray(img,mode='RGB').resize((256,256)) 
+  return resized_image
+
 def tune():
     model_id = run_PTI(run_name='',use_wandb=False, use_multi_id_training=False)
-
+    w_path_dir = f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}'
+    embedding_dir = f'{w_path_dir}/{paths_config.pti_results_keyword}/{image_name}'
+    w_pivot = torch.load(f'{embedding_dir}/0.pt')
+    new_G = load_generator(model_id)
+    new_image = new_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
+    tensor_to_pil(new_image).save("output/out.png")
+   
 #----------------------------------------------------------------------------
 if __name__ == '__main__':
     tune()