initial fastfold

inference.py

@@ -227,9 +227,11 @@ def inference_multimer_model(args):
     )
 
     output_dir_base = args.output_dir
+    
     random_seed = args.data_random_seed
     if random_seed is None:
         random_seed = random.randrange(sys.maxsize)
+    # seed_torch(seed=1029)
     
     feature_processor = feature_pipeline.FeaturePipeline(
         config.data
@@ -312,6 +314,7 @@ def inference_multimer_model(args):
     with open(unrelaxed_output_path, 'w') as f:
         f.write(protein.to_pdb(unrelaxed_protein))
 
+    if(args.relaxation):
         amber_relaxer = relax.AmberRelaxation(
             use_gpu=True,
             **config.relax,
@@ -328,6 +331,15 @@ def inference_multimer_model(args):
         with open(relaxed_output_path, 'w') as f:
             f.write(relaxed_pdb_str)
         
+    if(args.save_outputs):
+            output_dict_path = os.path.join(
+                args.output_dir, f'{tag}_{args.model_name}_output_dict.pkl'
+            )
+            with open(output_dict_path, "wb") as fp:
+                pickle.dump(out, fp, protocol=pickle.HIGHEST_PROTOCOL)
+
+            logger.info(f"Model output written to {output_dict_path}...")
+
 
 def inference_monomer_model(args):
     print("running in monomer mode...")
@@ -454,6 +466,7 @@ def inference_monomer_model(args):
     with open(unrelaxed_output_path, 'w') as f:
         f.write(protein.to_pdb(unrelaxed_protein))
 
+    if(args.relaxation):
         amber_relaxer = relax.AmberRelaxation(
             use_gpu=True,
             **config.relax,
@@ -512,6 +525,9 @@ if __name__ == "__main__":
                         help="""Path to model parameters. If None, parameters are selected
              automatically according to the model name from 
              ./data/params""")
+    parser.add_argument(
+        "--relaxation", action="store_false", default=False,
+    )
     parser.add_argument(
         "--save_outputs", action="store_true", default=False,
         help="Whether to save all model outputs, including embeddings, etc."