Update my_model/captioner/image_captioning.py

my_model/captioner/image_captioning.py

@@ -71,6 +71,9 @@ class ImageCaptioningModel:
         # If it's a file path or file-like object, open it as a PIL Image
             image = Image.open(image_path)
             
+        elif isinstance(image_path, Image.Image):
+            image = image_path
+            
         image = self.resize_image(image)
         inputs = self.processor(image, self.prompt, return_tensors="pt").to("cuda", self.torch_dtype)
         outputs = self.model.generate(**inputs, min_length=self.min_length, max_new_tokens=self.max_new_tokens)