From c136b91fbad74fab24f7d8bef29b45ad5eb60e16 Mon Sep 17 00:00:00 2001
From: Petr Andriushchenko <pitandmind@gmail.com>
Date: Thu, 23 Feb 2023 14:09:19 +0100
Subject: [PATCH] add model change function to summary

---
 misinformation/summary.py | 58 +++++++++++++++++++++++++++------------
 1 file changed, 41 insertions(+), 17 deletions(-)

diff --git a/misinformation/summary.py b/misinformation/summary.py
index 4448a0c..bc1d6b4 100644
--- a/misinformation/summary.py
+++ b/misinformation/summary.py
@@ -15,9 +15,23 @@ class SummaryDetector(AnalysisMethod):
 
     summary_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    summary_model, summary_vis_processors, _ = load_model_and_preprocess(
-        name="blip_caption", model_type="base_coco", is_eval=True, device=summary_device
-    )
+    def load_model_base(self):
+        summary_model, summary_vis_processors, _ = load_model_and_preprocess(
+            name="blip_caption",
+            model_type="base_coco",
+            is_eval=True,
+            device=SummaryDetector.summary_device,
+        )
+        return summary_model, summary_vis_processors
+
+    def load_model_large(self):
+        summary_model, summary_vis_processors, _ = load_model_and_preprocess(
+            name="blip_caption",
+            model_type="large_coco",
+            is_eval=True,
+            device=SummaryDetector.summary_device,
+        )
+        return summary_model, summary_vis_processors
 
     def set_keys(self) -> dict:
         params = {
@@ -26,21 +40,29 @@ class SummaryDetector(AnalysisMethod):
         }
         return params
 
-    def analyse_image(self):
+    def analyse_image(self, model_type):
 
+        select_model = {
+            "base": self.load_model_base,
+            "large": self.load_model_large,
+        }
+        summary_model, summary_vis_processors = select_model[model_type]()
         path = self.subdict["filename"]
         raw_image = Image.open(path).convert("RGB")
         image = (
-            self.summary_vis_processors["eval"](raw_image)
+            summary_vis_processors["eval"](raw_image)
             .unsqueeze(0)
             .to(self.summary_device)
         )
-        self.image_summary["const_image_summary"] = self.summary_model.generate(
-            {"image": image}
-        )[0]
-        self.image_summary["3_non-deterministic summary"] = self.summary_model.generate(
-            {"image": image}, use_nucleus_sampling=True, num_captions=3
-        )
+        with torch.no_grad():
+            self.image_summary["const_image_summary"] = summary_model.generate(
+                {"image": image}
+            )[0]
+            self.image_summary[
+                "3_non-deterministic summary"
+            ] = summary_model.generate(
+                {"image": image}, use_nucleus_sampling=True, num_captions=3
+            )
         for key in self.image_summary:
             self.subdict[key] = self.image_summary[key]
         return self.subdict
@@ -53,12 +75,13 @@ class SummaryDetector(AnalysisMethod):
         name="blip_vqa", model_type="vqav2", is_eval=True, device=summary_device
     )
 
-    def analyse_questions(self, list_of_questions):
+    def analyse_questions(self, model_type, list_of_questions):
+
         if len(list_of_questions) > 0:
             path = self.subdict["filename"]
             raw_image = Image.open(path).convert("RGB")
             image = (
-                self.summary_VQA_vis_processors["eval"](raw_image)
+                summary_VQA_vis_processors["eval"](raw_image)
                 .unsqueeze(0)
                 .to(self.summary_device)
             )
@@ -68,10 +91,11 @@ class SummaryDetector(AnalysisMethod):
             batch_size = len(list_of_questions)
             image_batch = image.repeat(batch_size, 1, 1, 1)
 
-            answers_batch = self.summary_VQA_model.predict_answers(
-                samples={"image": image_batch, "text_input": question_batch},
-                inference_method="generate",
-            )
+            with torch.no_grad():
+                answers_batch = self.summary_VQA_model.predict_answers(
+                    samples={"image": image_batch, "text_input": question_batch},
+                    inference_method="generate",
+                )
 
             for q, a in zip(question_batch, answers_batch):
                 self.image_summary[q] = a