Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import gradio as gr
 from prediction import run_image_prediction
 import torch
@@ -5,54 +6,84 @@ import torchvision.transforms as T
 from celle.utils import process_image
 from PIL import Image
 from matplotlib import pyplot as plt
+from celle_main import instantiate_from_config
+from omegaconf import OmegaConf
+
+
+class model:
+    def __init__(self):
+        self.model = None
+        self.model_name = None
+
+    def gradio_demo(self, model_name, sequence_input, nucleus_image, protein_image):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        if self.model_name != model_name:
+            self.model_name = model_name
+            model_ckpt_path = f"CELL-E_2-Image_Prediction/models/{model_name}.ckpt"
+            model_config_path = f"CELL-E_2-Image_Prediction/models/{model_name}.yaml"
+            
+            
+            # 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
+            
+            base_path = os.getcwd()
+
+            os.chdir(os.path.dirname(model_ckpt_path))
+
+            # Instantiate model from config and move to device
+            self.model = instantiate_from_config(config.model).to(device)
+            self.model = torch.compile(self.model,mode='reduce-overhead')
+            
+            os.chdir(base_path)
+
+
+        if "Finetuned" in model_name:
+            dataset = "OpenCell"
+
+        else:
+            dataset = "HPA"
+
+        nucleus_image = process_image(nucleus_image, dataset, "nucleus")
+        if protein_image:
+            protein_image = process_image(protein_image, dataset, "protein")
+            protein_image = protein_image > torch.median(protein_image)
+            protein_image = protein_image[0, 0]
+            protein_image = protein_image * 1.0
+        else:
+            protein_image = torch.ones((256, 256))
+
+        threshold, heatmap = run_image_prediction(
+            sequence_input=sequence_input,
+            nucleus_image=nucleus_image,
+            model=self.model,
+            device=device,
+        )
 
+        # Plot the heatmap
+        plt.imshow(heatmap.cpu(), cmap="rainbow", interpolation="bicubic")
+        plt.axis("off")
 
-def gradio_demo(model_name, sequence_input, nucleus_image, protein_image):
-    model = f"CELL-E_2-Image_Prediction/models/{model_name}.ckpt"
-    config = f"CELL-E_2-Image_Prediction/models/{model_name}.yaml"
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    if "Finetuned" in model_name:
-        dataset = "OpenCell"
-
-    else:
-        dataset = "HPA"
-
-    nucleus_image = process_image(nucleus_image, dataset, "nucleus")
-    if protein_image:
-        protein_image = process_image(protein_image, dataset, "protein")
-        protein_image = protein_image > torch.median(protein_image)
-        protein_image = protein_image[0, 0]
-        protein_image = protein_image * 1.0
-    else:
-        protein_image = torch.ones((256, 256))
-
-    threshold, heatmap = run_image_prediction(
-        sequence_input=sequence_input,
-        nucleus_image=nucleus_image,
-        model_ckpt_path=model,
-        model_config_path=config,
-        device=device,
-    )
-
-    # Plot the heatmap
-    plt.imshow(heatmap.cpu(), cmap="rainbow", interpolation="bicubic")
-    plt.axis("off")
-
-    # Save the plot to a temporary file
-    plt.savefig("temp.png", bbox_inches="tight", dpi=256)
+        # Save the plot to a temporary file
+        plt.savefig("temp.png", bbox_inches="tight", dpi=256)
 
-    # Open the temporary file as a PIL image
-    heatmap = Image.open("temp.png")
+        # Open the temporary file as a PIL image
+        heatmap = Image.open("temp.png")
 
-    return (
-        T.ToPILImage()(nucleus_image[0, 0]),
-        T.ToPILImage()(protein_image),
-        T.ToPILImage()(threshold),
-        heatmap,
-    )
+        return (
+            T.ToPILImage()(nucleus_image[0, 0]),
+            T.ToPILImage()(protein_image),
+            T.ToPILImage()(threshold),
+            heatmap,
+        )
 
+base_class = model()
 
 with gr.Blocks(theme='gradio/soft') as demo:
     gr.Markdown("Select the prediction model.")
@@ -64,8 +95,8 @@ with gr.Blocks(theme='gradio/soft') as demo:
     )
     with gr.Row():
         model_name = gr.Dropdown(
-            ["CELL-E_2-HPA_480", "CELL-E_2-HPA_Finetuned_480"],
-            value="CELL-E_2-HPA_480",
+            ["CELL-E_2_HPA_480", "CELL-E_2_HPA_Finetuned_480"],
+            value="CELL-E_2_HPA_480",
             label="Model Name",
         )
     with gr.Row():
@@ -120,6 +151,6 @@ with gr.Blocks(theme='gradio/soft') as demo:
             predicted_heatmap,
         ]
 
-        button.click(gradio_demo, inputs, outputs)
+        button.click(base_class.gradio_demo, inputs, outputs)
 
 demo.launch(enable_queue=True)