Update prediction.py

prediction.py

@@ -1,15 +1,9 @@
-import os
-os.chdir('..')
-base_dir = os.getcwd()
 from dataloader import CellLoader
-from celle_main import instantiate_from_config
-from omegaconf import OmegaConf
 
 def run_image_prediction(
     sequence_input,
     nucleus_image,
-    model_ckpt_path,
-    model_config_path,
+    model,
     device
 ):
     """
@@ -37,22 +31,6 @@ def run_image_prediction(
     # Convert SEQUENCE to sequence using dataset.tokenize_sequence()
     sequence = dataset.tokenize_sequence(sequence_input)
 
-    # Load model config and set ckpt_path if not provided in config
-    config = OmegaConf.load(model_config_path)
-    if config["model"]["params"]["ckpt_path"] is None:
-        config["model"]["params"]["ckpt_path"] = model_ckpt_path
-
-    # 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.model).to(device)
-
-    os.chdir(base_dir)
-
     # Sample from model using provided sequence and nucleus image
     _, _, _, predicted_threshold, predicted_heatmap = model.celle.sample(
         text=sequence.to(device),