|
--- |
|
language: |
|
- en |
|
license: mit |
|
library_name: transformers |
|
tags: |
|
- vision |
|
- image-to-text |
|
- image-captioning |
|
pipeline_tag: image-to-text |
|
base_model: Salesforce/blip2-opt-2.7b |
|
--- |
|
# VLRM |
|
This repository contains the weights of BLIP-2 OPT-2.7B model fine-tuned by reinforcement learning method introduced in the paper [VLRM: Vision-Language Models act as |
|
Reward Models for Image Captioning](https://arxiv.org/abs/2404.01911). |
|
|
|
The RL-tuned model is able to generate longer and more comprehensive descriptions with zero computational overhead compared to the original model. |
|
|
|
You can find other details in the [GitHub Repository (to be done)](https://github.com/papermsucode). |
|
# Running the model |
|
## Option 1 |
|
<details> |
|
<summary> Load the whole model from this repo </summary> |
|
|
|
```python |
|
import torch |
|
import requests |
|
from PIL import Image |
|
from transformers import Blip2Processor, Blip2ForConditionalGeneration |
|
|
|
processor = Blip2Processor.from_pretrained("sashakunitsyn/vlrm-blip2-opt-2.7b") |
|
model = Blip2ForConditionalGeneration.from_pretrained("sashakunitsyn/vlrm-blip2-opt-2.7b", torch_dtype=torch.float16, device_map="auto") |
|
|
|
img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg' |
|
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB') |
|
|
|
inputs = processor(raw_image, return_tensors="pt").to("cuda", torch.float16) |
|
|
|
out = model.generate(**inputs, max_new_tokens=60) |
|
processor.decode(out[0], skip_special_tokens=True).strip() |
|
>>> 'a woman in a plaid shirt shaking hands with a yellow labrador retriever sitting on the ground at sunset on a beach in florida' |
|
``` |
|
</details> |
|
|
|
## Option 2 |
|
Since the fine-tuned layers take small part of the whole model, you can first load the original model, then load the RL-tuned weights. |
|
<details> |
|
<summary> Step 1. Load the original model </summary> |
|
|
|
```python |
|
import torch |
|
import requests |
|
from PIL import Image |
|
from transformers import Blip2Processor, Blip2ForConditionalGeneration |
|
|
|
processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b") |
|
model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16, device_map="auto") |
|
|
|
img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg' |
|
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB') |
|
|
|
inputs = processor(raw_image, return_tensors="pt").to("cuda", torch.float16) |
|
|
|
out = model.generate(**inputs, max_new_tokens=60) |
|
processor.decode(out[0], skip_special_tokens=True).strip() |
|
>>> 'a woman sitting on the beach with a dog' |
|
``` |
|
</details> |
|
|
|
<details> |
|
<summary> Step 2. Load the RL-tuned weights </summary> |
|
Available checkpoints: |
|
|
|
- `vlrm-blip2-opt-2.7b.pt` (VLRM in the paper) |
|
- `vlrm-rs-blip2-opt-2.7b.pt` (VLRM-RS in the paper) |
|
|
|
```python |
|
from huggingface_hub import hf_hub_download |
|
finetuned_weights_state_dict = torch.load(hf_hub_download(repo_id="sashakunitsyn/vlrm-blip2-opt-2.7b", filename="vlrm-blip2-opt-2.7b.pt")) |
|
model.load_state_dict(finetuned_weights_state_dict, strict=False) |
|
|
|
out = model.generate(**inputs, max_new_tokens=60) |
|
processor.decode(out[0], skip_special_tokens=True).strip() |
|
>>> 'a woman in a plaid shirt shaking hands with a yellow labrador retriever sitting on the ground at sunset on a beach in florida' |
|
``` |
|
</details> |
|
|
