From 49bc9493d78a6cd2132a31ed9952d89d8956d67e Mon Sep 17 00:00:00 2001 From: dnth Date: Thu, 10 Oct 2024 12:31:28 +0800 Subject: [PATCH] update --- InferX/transformers/blip2.py | 4 ++-- README.md | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/InferX/transformers/blip2.py b/InferX/transformers/blip2.py index b2bbf2e..b37767a 100644 --- a/InferX/transformers/blip2.py +++ b/InferX/transformers/blip2.py @@ -36,11 +36,11 @@ def preprocess(self, image: str | Image.Image, prompt: str = None): self.device ) - def predict(self, processed_data, **generate_kwargs): + def predict(self, preprocessed_input, **generate_kwargs): with torch.inference_mode(), torch.amp.autocast( device_type=self.device, dtype=torch.bfloat16 ): - return self.model.generate(**processed_data, **generate_kwargs) + return self.model.generate(**preprocessed_input, **generate_kwargs) def postprocess(self, prediction): return self.processor.batch_decode(prediction, skip_special_tokens=True)[0] diff --git a/README.md b/README.md index 66f8013..f2e171c 100644 --- a/README.md +++ b/README.md @@ -59,4 +59,18 @@ output = model.postprocess(prediction) print(output) >>> A cat on a yellow background + + +image = "https://img.freepik.com/free-photo/adorable-black-white-kitty-with-monochrome-wall-her_23-2148955182.jpg" +prompt = "Describe this image in concise detail. Answer:" + + +processed_input = model.preprocess(image, prompt) + +prediction = model.predict(processed_input, max_new_tokens=200) +output = model.postprocess(prediction) + +print(output) +>>> a black and white cat sitting on a table looking up at the camera + ```