Update prediction.py

prediction.py

@@ -1,5 +1,6 @@
 import os
 os.chdir('..')
+base_dir = os.getcwd()
 from dataloader import CellLoader
 from celle_main import instantiate_from_config
 from omegaconf import OmegaConf
@@ -53,30 +54,22 @@ def run_sequence_prediction(
     # Set condition_model_path and vqgan_model_path to None
     config["model"]["params"]["condition_model_path"] = None
     config["model"]["params"]["vqgan_model_path"] = None
+    
+    os.chdir(os.path.dirname(model_ckpt_path))
 
     # Instantiate model from config and move to device
-    model = instantiate_from_config(config).to(device)
+    model = instantiate_from_config(config.model).to(device)
 
     # Sample from model using provided sequence and nucleus image
     _, predicted_sequence, _ = model.celle.sample_text(
-        text=sequence,
-        condition=nucleus_image,
-        image=protein_image,
+        text=sequence.to(device),
+        condition=nucleus_image.to(device),
+        image=protein_image.to(device),
         force_aas=True,
-        timesteps=1,
         temperature=1,
-        progress=True,
+        progress=False,
     )
+    
+    os.chdir(base_dir)
 
-    formatted_predicted_sequence = ""
-
-    for i in range(min(len(predicted_sequence), len(sequence))):
-        if predicted_sequence[i] != sequence[i]:
-            formatted_predicted_sequence += f"**{predicted_sequence[i]}**"
-        else:
-            formatted_predicted_sequence += predicted_sequence[i]
-
-    if len(predicted_sequence) > len(sequence):
-        formatted_predicted_sequence += f"**{predicted_sequence[len(sequence):]}**"
-
-    return formatted_predicted_sequence
+    return predicted_sequence