python -m venv venv
source ./venv/bin/activate # linux
call ./venv/scripts/Activate.bat  # windows?

pip install transformers peft datasets

from transformers import AutoModelForVision2Seq, AutoProcessor
from peft import LoraConfig, get_peft_model

config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    # target specific modules (look for Linear in the model)
    # print(model) to see the architecture of the model
    target_modules=[
        "self.query",
        "self.key",
        "self.value",
        "output.dense",
        "self_attn.qkv",
        "self_attn.projection",
        "mlp.fc1",
        "mlp.fc2",
    ],
)

# loading the model in float16
model = AutoModelForVision2Seq.from_pretrained("Salesforce/blip-image-captioning-base", torch_dtype=torch.float16)

# Load the blip model processor (just sets up the associated image processor so no additional model)
processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")

# We layer our PEFT on top of our model using the PEFT config
model = get_peft_model(model, config)
model.print_trainable_parameters()

from datasets import load_dataset

# We are extracting the train dataset
dataset = load_dataset("ybelkada/football-dataset", split="train")

# Sample the dataset
print("Dataset images: ", len(dataset))
print(next(iter(dataset)))

from torch.utils.data import DataLoader, Dataset


# Each dataset may have different names and format and this function should be
# adjusted for each dataset. Creating `image` and `text` keys to correlate with
# the collator
class FootballImageCaptioningDataset(Dataset):
    def __init__(self, dataset, processor):
        self.dataset = dataset
        self.processor = processor

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        item = self.dataset[idx]
        encoding = self.processor(
            images=item["image"], padding="max_length", return_tensors="pt"
        )
        # remove batch dimension
        encoding = {k: v.squeeze() for k, v in encoding.items()}
        encoding["text"] = item["text"]
        return encoding


def collator(batch):
    # pad the input_ids and attention_mask
    processed_batch = {}
    for key in batch[0].keys():
        if key != "text":
            processed_batch[key] = torch.stack([example[key] for example in batch])
        else:
            text_inputs = processor.tokenizer(
                [example["text"] for example in batch], padding=True, return_tensors="pt"
            )
            processed_batch["input_ids"] = text_inputs["input_ids"]
            processed_batch["attention_mask"] = text_inputs["attention_mask"]
    return processed_batch


train_dataset = FootballImageCaptioningDataset(dataset, processor)

# Set batch_size to the batch that works for you.
train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=2, collate_fn=collator)

import torch.optim as optim

# Setup AdamW
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)

# Set the device which works for you or use these defaults
device = "cuda" if torch.cuda.is_available() else "cpu"

# Set the model as ready for training, makes sure the gradient are on
model.train()

# 50 is the number of epochs. adjust to how many you may need
for epoch in range(50):
    print("Epoch:", epoch)
    for idx, batch in enumerate(train_dataloader):
        input_ids = batch.pop("input_ids").to(device)
        pixel_values = batch.pop("pixel_values").to(device)
        attention_mask = batch.pop("attention_mask").to(device)

        outputs = model(
            input_ids=input_ids,
            pixel_values=pixel_values,
            labels=input_ids,
            attention_mask=attention_mask,
        )

        loss = outputs.loss

        print("Loss:", loss.item())

        loss.backward()

        optimizer.step()
        optimizer.zero_grad()

        if idx % 10 == 0:
            generated_output = model.generate(pixel_values=pixel_values)
            print(processor.batch_decode(generated_output, skip_special_tokens=True))

# hugging face models saved to the directory
model.save_pretrained('./training/caption')

from PIL import Image
from transformers import AutoModelForVision2Seq, AutoProcessor

# The model here is to get the appropriate processor
model_id = "Salesforce/blip-image-captioning-base"

processor = AutoProcessor.from_pretrained(model_id)

# Caption PEFT model has the metadata for the appropriate base model. So we can just load
# that model directly
model = AutoModelForVision2Seq.from_pretrained('./training/caption')

# Turn on evaluation mode so we don't calculate gradients
model.eval()

# https://upload.wikimedia.org/wikipedia/commons/5/52/Lionel_Messi_playing_in_Argentina_2022_FIFA_World_Cup.jpg
# Set the images to be to captioned
images = [Image.open("Lionel_Messi_playing_in_Argentina_2022_FIFA_World_Cup.jpg")]

processed = processor(images=images, padding="max_length", return_tensors="pt")
generated_output = model.generate(pixel_values=processed['pixel_values'], max_new_tokens=64)

print(processor.batch_decode(generated_output, skip_special_tokens=True))

$ python inference2.py
['messi with argentina at the 2018 world cup']