Testing

handler.py

@@ -1,7 +1,9 @@
 from typing import Any, Dict, List
-from transformers import Idefics2Processor, Idefics2Model
+from transformers import Idefics2Processor, Idefics2ForConditionalGeneration
 import torch
 import logging
+from PIL import Image
+import requests
 
 
 class EndpointHandler:
@@ -11,7 +13,7 @@ class EndpointHandler:
         self.logger.addHandler(logging.StreamHandler())
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.processor = Idefics2Processor.from_pretrained(path)
-        self.model = Idefics2Model.from_pretrained(path)
+        self.model = Idefics2ForConditionalGeneration.from_pretrained(path)
         self.model.to(self.device)
         self.logger.info("Initialisation finished!")
 
@@ -23,20 +25,65 @@ class EndpointHandler:
         Return:
               A :obj:`list` | `dict`: will be serialized and returned
         """
-        image = data.pop("inputs", data)
+        """image = data.pop("inputs", data)
         self.logger.info("image")
 
         # process image
         inputs = self.processor(images=image, return_tensors="pt").to(self.device)
         self.logger.info("inputs")
-        generated_ids = self.model.forward(input_ids=inputs)
+        self.logger.info(f"{inputs.input_ids}")
+        generated_ids = self.model.generate(**inputs)
         self.logger.info("generated")
 
         # run prediction
         generated_text = self.processor.batch_decode(
             generated_ids, skip_special_tokens=True
         )
-        self.logger.info("decoded")
+        self.logger.info("decoded")"""
+
+        url_1 = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        url_2 = "http://images.cocodataset.org/val2017/000000219578.jpg"
+
+        image_1 = Image.open(requests.get(url_1, stream=True).raw)
+        image_2 = Image.open(requests.get(url_2, stream=True).raw)
+        images = [image_1, image_2]
+
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "What’s the difference between these two images?",
+                    },
+                    {"type": "image"},
+                    {"type": "image"},
+                ],
+            }
+        ]
+
+        processor = Idefics2Processor.from_pretrained("HuggingFaceM4/idefics2-8b")
+        model = Idefics2ForConditionalGeneration.from_pretrained(
+            "HuggingFaceM4/idefics2-8b"
+        )
+        model.to(self.device)
+
+        # at inference time, one needs to pass `add_generation_prompt=True` in order to make sure the model completes the prompt
+        text = processor.apply_chat_template(messages, add_generation_prompt=True)
+        self.logger.info(text)
+        # 'User: What’s the difference between these two images?<image><image><end_of_utterance>\nAssistant:'
+
+        inputs = processor(images=images, text=text, return_tensors="pt").to(
+            self.device
+        )
+        self.logger.info("inputs")
+
+        generated_text = model.generate(**inputs, max_new_tokens=500)
+        self.logger.info("generated")
+        generated_text = processor.batch_decode(
+            generated_text, skip_special_tokens=True
+        )[0]
+        self.logger.info(f"Generated text: {generated_text}")
 
         # decode output
         return generated_text