Official implementation of paper "MiniGPT-5: Interleaved Vision-and-Language Generation via Generative Vokens"
import torch
import os

CC3M_FOLDER = os.getenv("CC3M_FOLDER", "datasets/CC3M")
DATAFOLDER = os.getenv("DATAFOLDER", "datasets/CC3M")
PREPROCESS_FEATURE_FOLDER = os.getenv("PREPROCESS_FEATURE_FOLDER", None)

WEIGHTFOLDER = os.getenv("WEIGHTFOLDER", "weights")
OUTPUT_FOLDER = os.getenv("OUTPUT_FOLDER", "outputs")
OUTPUT_SUFFIX = os.getenv("OUTPUT_SUFFIX", "")

IS_STAGE2 = os.getenv("IS_STAGE2", False)
IS_STAGE2 = True if IS_STAGE2 in ["True", "true"] else False

IMG_TOKEN_NUM = 8
ALL_IMG_TOKENS = [f"[IMG{i}]" for i in range(IMG_TOKEN_NUM)]
ALL_IMG_TOKENS_STR = "".join(ALL_IMG_TOKENS)

USE_PREFIX_TUNING = False
USE_LORA = False
USE_CFG = True

if IS_STAGE2:
    USE_LORA = True

PRECISION = torch.bfloat16
TRAINABLE_PRECISION = torch.float32
constants.py
eric-ai-lab/MiniGPT-5
from torch.utils.data import Dataset
from typing import Optional, Dict, Sequence
import transformers
from tqdm import tqdm
from dataclasses import dataclass
import pandas as pd
import re
import os
import numpy as np
from PIL import Image
import random
from pathlib import Path
import json
import copy
from constants import *

class CC3MDataset(Dataset):
    def __init__(self, data_path: str, input_processor=None, output_vis_processor=None, test=False):
        self.test = test
        self.input_processor = input_processor
        self.output_vis_processor = output_vis_processor
        self.output_img_id = input_processor.tokenizer.convert_tokens_to_ids(ALL_IMG_TOKENS[0])
        self.load_preprocessed_image_features = not test
        saved_data_path = data_path.replace('.tsv', '_8-token_stage1.pkl')
        if os.path.exists(saved_data_path):
            print("Loading saved data...")
            self.recover_data(saved_data_path)
            print("Loaded saved data for CC3M")
        else:
            list_data_table = pd.read_csv(data_path, delimiter='\t')
            self.sources, self.targets, self.input_image_path, self.output_image_path = [], [], [], []
            self.caption, self.task_names = [], []

            system_prompt="You will be able to generate image according to command."
            generation_prompts = [
                "generate image with caption:",
                "can you give me the image with caption:",
                "help me to generate this image:",
                "generate image with according to caption:",
                "according to caption, generate image:",
                "a image with caption:",
                "can you visualize this caption:",
            ]

            for i in tqdm(range(len(list_data_table))):
                data = list_data_table.iloc[i]
                step_image = data['image_path']
                step_caption = data['caption']
                path = Path(step_image)
                step_image = Path(DATAFOLDER).joinpath(path)
                step_image = str(step_image)
                step_caption = self.pre_caption(step_caption)

                caption_source = f"{step_caption}"
                caption_target = f'{ALL_IMG_TOKENS_STR} ###'
                self.sources.append(caption_source)
                self.targets.append(caption_target)
                self.caption.append(step_caption)
                self.task_names.append(f'cc3m_{i}')
                self.input_image_path.append([None])
                self.output_image_path.append(step_image)
                
                if i%100 == 0 and not test:
                    caption_source = f"###Human: {random.choice(generation_prompts)} {step_caption} ###Assistant:"
                    caption_source = system_prompt + caption_source
                    caption_target = f'{ALL_IMG_TOKENS_STR} ###'
                    self.sources.append(caption_source)
                    self.targets.append(caption_target)
                    self.caption.append(step_caption)
                    self.task_names.append(f'cc3m_{i}_instruction')
                    self.input_image_path.append([None])
                    self.output_image_path.append(step_image)
            self.valid_idx = list(range(len(self.sources)))
            print("Saving data...")
            self.save_process_data(saved_data_path)
            print("Saved data for cc3m!")
        if test:
            self.targets = self.caption

    def recover_data(self, saved_file):
        all_data = torch.load(saved_file)
        self.sources = all_data['sources']
        self.targets = all_data['targets']
        self.input_image_path = all_data['input_image_path']
        self.output_image_path = all_data['output_image_path']
        self.caption = all_data['caption']
        self.task_names = all_data['task_names']
        del all_data
        if self.test:
            self.valid_idx = []
            for i in range(len(self.targets)):
                if self.output_image_path[i] is not None:
                    self.valid_idx.append(i)
            

    def save_process_data(self, saved_file):
        all_data = {'sources': self.sources,
                    'targets': self.targets,
                    'input_image_path': self.input_image_path,
                    'output_image_path': self.output_image_path,
                    'caption': self.caption,
                    'task_names': self.task_names,
                    }
        torch.save(all_data, saved_file)
    
    def __len__(self):
        if self.test:
            return len(self.valid_idx)
        return len(self.sources)

    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        if self.test:
            i = self.valid_idx[i]
        input_image_path = self.input_image_path[i]
        output_image_path = self.output_image_path[i]
        input_text = self.sources[i]
        output_text = self.targets[i]
        if self.load_preprocessed_image_features and PREPROCESS_FEATURE_FOLDER is not None and os.path.isdir(PREPROCESS_FEATURE_FOLDER):
            if output_image_path is not None:
                output_feature_name = Path(output_image_path).name
                output_feature_name = output_feature_name.replace('.jpg', '_output.pt')
                if 'val' in output_image_path:
                    output_feature_path = Path(PREPROCESS_FEATURE_FOLDER).joinpath('val', output_feature_name)
                elif 'train' in output_image_path:
                    output_feature_path = Path(PREPROCESS_FEATURE_FOLDER).joinpath('train', output_feature_name)
                output_image_feature = torch.load(output_feature_path).unsqueeze(0)
            else:
                output_image_path = 'none'
                output_image_feature = torch.zeros((1, 8, 64, 64))
            
            input_images_feature = []
            for in_img_path in input_image_path:
                if in_img_path is not None:
                    input_feature_name = Path(in_img_path).name
                    input_feature_name = input_feature_name.replace('.jpg', '_input.pt')
                    if 'val' in in_img_path:
                        input_feature_path = Path(PREPROCESS_FEATURE_FOLDER).joinpath('val', input_feature_name)
                    elif 'train' in in_img_path:
                        input_feature_path = Path(PREPROCESS_FEATURE_FOLDER).joinpath('train', input_feature_name)
                    input_image_feature = torch.load(input_feature_path).unsqueeze(0)
                else:
                    input_image_feature = torch.zeros((1, 32, 4096))
                input_images_feature.append(input_image_feature)
            input_images_feature = torch.cat(input_images_feature, dim=0)
            input_dict = self.input_processor(text = input_text, add_special_tokens=False)
            input_dict['input_images_feature'] = input_images_feature
            input_dict['output_image_feature'] = output_image_feature
        else:
            input_images = []
            for in_img_path in input_image_path:
                if in_img_path is not None:
                    input_image = Image.open(in_img_path).convert("RGB")
                else:
                    input_image = Image.fromarray(np.zeros((224, 224, 3), dtype=np.uint8))
                input_images.append(input_image)
            input_dict = self.input_processor(text = input_text, images = input_images, add_special_tokens=False)
            input_dict['original_images'] = input_images
            
            if output_image_path is not None:
                output_image = Image.open(output_image_path).convert("RGB")
                output_image = self.expand2square(output_image, (255, 255, 255))
                output_image = self.output_vis_processor(output_image)
                output_image = output_image.unsqueeze(0)
            else:
                output_image_path = 'none'
                output_image = torch.zeros((1, 3, 512, 512))
            input_dict["output_image"] = output_image

        input_dict["caption"] = self.caption[i]
        input_dict["task_name"] = self.task_names[i]
        target_ids = self.input_processor(text = output_text, add_special_tokens=False)['input_ids']
        label = torch.ones_like(input_dict["input_ids"])*-100
        label = torch.cat((label, target_ids), dim=1)
        index = torch.nonzero(label == self.output_img_id)
        if len(index):
            index = index[0,1]
            label[:, index+1:index+IMG_TOKEN_NUM-1] = -100
        input_dict["labels"] = label
        input_dict["input_ids"] = torch.cat((input_dict["input_ids"], target_ids), dim=1)
        input_dict["attention_mask"] = torch.cat((input_dict["attention_mask"], torch.ones_like(target_ids)), dim=1)
        input_dict["source"] = input_text
        input_dict["target"] = output_text

        return input_dict

    def pre_caption(self, caption):
        
        caption = re.sub(
            r"([.!\"()*#:;~])",
            " ",
            caption.lower(),
        )
        caption = re.sub(
            r"\s{2,}",
            " ",
            caption,
        )
        caption = caption.rstrip("\n")
        caption = caption.strip(" ")

        # truncate caption
        # max_words = 100
        # caption_words = caption.split(" ")
        # if len(caption_words) > max_words:
        #     caption = " ".join(caption_words[: max_words])

        return caption
    
    @staticmethod
    def expand2square(pil_img, background_color):
        width, height = pil_img.size
        if width == height:
            return pil_img
        elif width > height:
            result = Image.new(pil_img.mode, (width, width), background_color)
            result.paste(pil_img, (0, (width - height) // 2))
            return result
        else:
            result = Image.new(pil_img.mode, (height, height), background_color)
            result.paste(pil_img, ((height - width) // 2, 0))
            return result

class VISTDataset(CC3MDataset):
    def __init__(self, data_path: str, input_processor=None, output_vis_processor=None, test=False):
        self.test = test
        self.input_processor = input_processor
        self.output_vis_processor = output_vis_processor
        self.output_img_id = input_processor.tokenizer.convert_tokens_to_ids(ALL_IMG_TOKENS[0])
        eos_token = input_processor.tokenizer.eos_token
        self.load_preprocessed_image_features = False

        self.sources, self.targets, self.input_image_path, self.output_image_path = [], [], [], []
        self.caption, self.task_names = [], []
        all_tasks = json.load(open(data_path, 'r'))
        image_id_mapping = json.load(open(data_path.replace('cleaned', 'image_mapping'), 'r'))
        system_prompt1="Give the following images in <Img>ImageContent</Img> format. "\
           "You will be able to see the images once I provide it to you. Please understanding images and generate story."
        human_prompts1 = [
            "###Human:{prompt} Generate an image with the scene description: {step_text} ###Assistant:",
            "###Human:{prompt} the scene description: {step_text} ###Assistant:",
        ]
        human_prompts2 = [
            "###Human:{prompt} Tell me the next scene with image. ###Assistant:",
            "###Human:{prompt} Generate the next scene with image. ###Assistant:",
            "###Human:{prompt} What should happen then? ###Assistant:"
        ]
        human_prompts3 = [
            "###Human:{prompt} Tell me the next scene description by this image: <Img><ImageHere></Img> ###Assistant:",
            "###Human:{prompt} What happen in the next scene image: <Img><ImageHere></Img> ###Assistant:"
        ]
        for task_name, task in tqdm(all_tasks.items()):
            task_prompts = []
            task_input_image_path = []
            for step in task:
                step_text = step['caption']
                step_image = step['image_id']
                sequence_index = step['sequence_index']
                step_image = os.path.join(DATAFOLDER, image_id_mapping[step_image])

                prompt = "<Img><ImageHere></Img>".join(task_prompts)
                if len(task_prompts):
                    prompt = f"{prompt}<Img><ImageHere></Img>\n"

                step_input_image = copy.deepcopy(task_input_image_path)
                step_name = f"{task_name}_{sequence_index}"
                #image generation
                step_source = random.choice(human_prompts1).format(prompt=prompt, step_text=step_text)
                step_source = system_prompt1 + step_source
                step_target = f"{ALL_IMG_TOKENS_STR} ###"
                self.sources.append(step_source)
                self.caption.append(None)
                self.targets.append(step_target)
                self.task_names.append(step_name+"-gen")
                if len(step_input_image):
                    self.input_image_path.append(step_input_image)
                else:
                    self.input_image_path.append([None])
                self.output_image_path.append(step_image)

                if len(task_prompts) and not test:
                    #image and text generation
                    step_source = random.choice(human_prompts2).format(prompt=prompt)
                    step_source = system_prompt1 + step_source
                    step_target = f"{step_text} {ALL_IMG_TOKENS_STR} ###"
                    self.sources.append(step_source)
                    self.targets.append(step_target)
                    self.caption.append(None)
                    self.task_names.append(step_name+"-multimodal")
                    self.input_image_path.append(step_input_image)
                    self.output_image_path.append(step_image)
                    
                    #image understanding
                    step_source = random.choice(human_prompts3).format(prompt=prompt)
                    step_source = system_prompt1 + step_source
                    step_target = f"{step_text} ###"
                    self.sources.append(step_source)
                    self.targets.append(step_target)
                    self.caption.append(None)
                    self.task_names.append(step_name+"-understanding")
                    self.input_image_path.append(step_input_image+[step_image])
                    self.output_image_path.append(None)

                task_prompts.append(step_text)
                task_input_image_path.append(step_image)

        self.valid_idx = list(range(len(self.sources)))
        print('Load data done!')

class MMDialogDataset(CC3MDataset):
    def __init__(self, data_path: str, input_processor=None, output_vis_processor=None, test=False):
        self.test = test
        self.input_processor = input_processor
        self.output_vis_processor = output_vis_processor
        self.output_img_id = input_processor.tokenizer.convert_tokens_to_ids(ALL_IMG_TOKENS[0])
        eos_token = input_processor.tokenizer.eos_token
        self.load_preprocessed_image_features = False
        
        system_prompt="Give the following images in <Img>ImageContent</Img> format. "\
           "You will be able to see the images once I provide it to you. Please generate conversation with appropriate image."
        
        self.sources, self.targets, self.input_image_path, self.output_image_path = [], [], [], []
        self.caption, self.task_names = [], []
        data_folder = os.path.dirname(data_path)
        with open(data_path, 'r') as f:
            all_data = f.readlines()
        for data in tqdm(all_data): 
            data = json.loads(data)
            data_num = data['conversation_id']
            conversation = data['conversation']
            if len(conversation)<2:
                continue
            history_prompt = system_prompt
            history_images = []
            for i, conv in enumerate(conversation):
                turn = conv['turn']
                turn_text = turn[0]['__TEXT__']
                if len(turn)==1:
                    turn_image_path = None
                else:
                    turn_image_path = os.path.join(data_folder, f"{turn[1]['__MEDIA__']}.jpg")
                    if not os.path.exists(turn_image_path):
                        # print(f'Cannot Find: {turn_image_path}')
                        break
                if len(turn_text.split(' '))>20 and not test:
                    break
                if i%2==0:
                    source = history_prompt + "###Human:"
                else:
                    source = history_prompt + "###Assistant:"

                if i>0:
                    self.sources.append(source)
                    if turn_image_path is not None:
                        target = f"{turn_text} {ALL_IMG_TOKENS_STR} ###"
                    else:
                        target = f"{turn_text} ###"
                    self.targets.append(target)
                    self.caption.append(None)
                    self.task_names.append(f'mmdialog{data_num}_{i}')
                    if len(history_images):
                        self.input_image_path.append(copy.deepcopy(history_images))
                    else:
                        self.input_image_path.append([None])
                    self.output_image_path.append(turn_image_path)

                if turn_image_path is not None:
                    history_prompt = source + f" {turn_text} <Img><ImageHere></Img>\n"
                    history_images.append(turn_image_path)
                else:
                    history_prompt = source + f" {turn_text}\n"
                
                if (i%2==1 and i>2):
                    pattern = "###Human:(.*?)###Human:"
                    match = re.search(pattern, history_prompt, re.DOTALL)
                    match_text = match.group(0)
                    history_prompt = history_prompt.replace(match_text, "###Human:")
                    pattern2 = '<Img><ImageHere></Img>'
                    match_num = len(re.findall(pattern2, match_text))
                    if match_num>0:
                        history_images = history_images[match_num:]

        self.valid_idx = list(range(len(self.sources)))
        print('Load data done with {} samples!'.format(len(self.sources)))

    def __getitem__(self, i):
        for _ in range(10):
            try:
                item = super().__getitem__(i)
                break
            except:
                i = random.choice(self.valid_idx)
        return item
    
@dataclass
class DataCollator(object):
    """Collate examples for supervised fine-tuning."""

    tokenizer: transformers.PreTrainedTokenizer
    sd_tokenizer: transformers.PreTrainedTokenizer

    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
        key_list = instances[0].keys()
        output_dict = {}
        for key in key_list:
            # Need to remove the batch dimension
            if key in ['input_ids', 'attention_mask', 'labels']:
                output_value = [instance[key][0] for instance in instances]
            else:
                output_value = [instance[key] for instance in instances]
            if key == "input_ids":
                output_value = torch.nn.utils.rnn.pad_sequence(output_value, batch_first=True, padding_value=self.tokenizer.pad_token_id)
            elif key == "labels":
                output_value = torch.nn.utils.rnn.pad_sequence(output_value, batch_first=True, padding_value=-100)
            elif key == "attention_mask":
                output_value = torch.nn.utils.rnn.pad_sequence(output_value, batch_first=True, padding_value=0)
            elif key == 'input_images':
                output_value = [v.to(PRECISION) for v in output_value]
            elif key == 'output_image':
                output_value = torch.concat(output_value).to(PRECISION)
            elif key == 'output_image_feature':
                output_value = torch.concat(output_value)
            output_dict[key] = output_value
        return output_dict

if 'CC3M' in DATAFOLDER:
    SupervisedDataset = CC3MDataset
elif 'MMDialog' in DATAFOLDER:
    SupervisedDataset = MMDialogDataset
else:
    SupervisedDataset = VISTDataset

dataloader.py
eric-ai-lab/MiniGPT-5
from typing import Any
import torch
import os
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import numpy as np
import pandas as pd
from lightning.pytorch import LightningModule, Trainer, seed_everything
from diffusers import AutoencoderKL
from pathlib import Path
from joblib import Parallel, delayed
from constants import *


class CC3MProcessDataset(Dataset):
    def __init__(self, data_path: str, output_vis_processor=None):
        self.output_vis_processor = output_vis_processor
        list_data_table = pd.read_csv(data_path, delimiter='\t')
        self.all_images = list_data_table['image_path'].tolist()

    def __len__(self):
        return len(self.all_images)
    
    def __getitem__(self, i):
        image_path = self.all_images[i]
        path = Path(image_path)
        image_path = Path(CC3M_FOLDER).joinpath(path)
        image_path = str(image_path)
        image = Image.open(image_path).convert("RGB")
        image = self.expand2square(image, (255, 255, 255))

        output_image = self.output_vis_processor(image)

        return {"output_processed_image": output_image, 'image_path': image_path}
    
    @staticmethod
    def expand2square(pil_img, background_color):
        width, height = pil_img.size
        if width == height:
            return pil_img
        elif width > height:
            result = Image.new(pil_img.mode, (width, width), background_color)
            result.paste(pil_img, (0, (width - height) // 2))
            return result
        else:
            result = Image.new(pil_img.mode, (height, height), background_color)
            result.paste(pil_img, ((height - width) // 2, 0))
            return result

def save_features(output_image_moment, image_path, save_folder):
    image_name = Path(image_path).name
    output_feature_name = image_name.replace('.jpg', '_output.pt')
    output_feature_path = os.path.join(save_folder, output_feature_name)
    torch.save(output_image_moment.detach().cpu(), output_feature_path, pickle_protocol=5)


class Processor(LightningModule):
    def __init__(self, save_folder) -> None:
        super().__init__()
        self.save_folder = save_folder
        self.output_vis_processor = transforms.Compose(
                [
                    transforms.Resize(512, interpolation=transforms.InterpolationMode.BILINEAR),
                    transforms.CenterCrop(512),
                    transforms.ToTensor(),
                ]
            )
        

        sd_model_name = "stabilityai/stable-diffusion-2-1-base"
        self.output_model = AutoencoderKL.from_pretrained(sd_model_name, subfolder="vae").to(PRECISION)
    
    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
        output_processed_image = batch["output_processed_image"].to(self.device,PRECISION)
        image_path = batch["image_path"]

        with torch.no_grad():
            h = self.output_model.encoder(output_processed_image)
            output_image_moments = self.output_model.quant_conv(h)
        
        Parallel(n_jobs=25)(delayed(save_features)(output_image_moments[i], image_path[i], self.save_folder) for i in range(len(image_path)))

if __name__ == "__main__":
    seed_everything(42)
    process_file = os.path.join(CC3M_FOLDER, 'cc3m_train.tsv')
    save_folder = os.path.join(CC3M_FOLDER, 'processed_features', 'train')

    if not os.path.exists(save_folder):
        os.makedirs(save_folder)
    processor = Processor(save_folder)
    processor.eval()
    processor.freeze()

    devices = [0,1,2,3]
    prediction_manager = Trainer(
                        max_epochs=1,
                        accelerator="gpu", devices=devices, 
                        strategy='ddp', 
                        logger = None,
                        precision='bf16-mixed')
    output_vis_processor = processor.output_vis_processor

    dataset = CC3MProcessDataset(process_file, output_vis_processor)
    dataloader = DataLoader(dataset, batch_size=50, num_workers=32, shuffle=False)
    prediction_manager.predict(processor, dataloaders=dataloader, return_predictions=False)
datasets/CC3M/preprocess_images.py
eric-ai-lab/MiniGPT-5
import os
import sys
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)

from torchvision import transforms
from model import MiniGPT5_Model
from train_eval import ModelArguments, DataArguments, TrainingArguments
from PIL import Image
import transformers
import torch
import matplotlib.pyplot as plt
import textwrap
from lightning.pytorch import seed_everything

def expand2square(pil_img, background_color):
    width, height = pil_img.size
    if width == height:
        return pil_img
    elif width > height:
        result = Image.new(pil_img.mode, (width, width), background_color)
        result.paste(pil_img, (0, (width - height) // 2))
        return result
    else:
        result = Image.new(pil_img.mode, (height, height), background_color)
        result.paste(pil_img, ((height - width) // 2, 0))
        return result
    
if __name__ == "__main__":
    seed_everything(42)
    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    if isinstance(training_args.gpus, str):
        training_args.gpus = [int(x) for x in training_args.gpus.split(',')]

    stage1_ckpt = model_args.stage1_weight
    stage2_ckpt = training_args.test_weight

    minigpt5 = MiniGPT5_Model.load_from_checkpoint(stage1_ckpt, strict=False, map_location="cpu", encoder_model_config=model_args, **vars(training_args))
    finetuned_state_dict = torch.load(stage2_ckpt, map_location="cpu")['state_dict']
    minigpt5.load_state_dict(finetuned_state_dict, strict=False)
    minigpt5.to("cuda:0")
    minigpt5.eval()
    input_images = None
    input_image_path = [os.path.join(current_dir, '000000005.jpg'), os.path.join(current_dir, '000000007.jpg')]

    if input_image_path:
        input_vis_processor = transforms.Compose(
            [
                transforms.Resize(224, interpolation=transforms.InterpolationMode.BILINEAR),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
            ]
        )
        input_images = []
        for img_path in input_image_path:
            input_image = Image.open(img_path).convert("RGB")
            input_image = expand2square(input_image, (255, 255, 255))
            input_image = input_vis_processor(input_image)
            input_image = input_image.unsqueeze(0).to("cuda:0")
            input_images.append(input_image)
        input_images = torch.cat(input_images, dim=0)

    system_prompt="Give the following images in <Img>ImageContent</Img> format. "\
           "You will be able to see the images once I provide it to you. Please understanding images and generate story."
    utterance = "my sister arrived early to help me with the family bar bq.<Img><ImageHere></Img>every one else arrived soon after.<Img><ImageHere>\n"
    utterance = system_prompt + f"###Human:{utterance} Tell me the next scene with image. ###Assistant:"

    for i in range(5):
        with torch.inference_mode():
            with torch.autocast("cuda"):
                text_out, image_out = minigpt5.generate(utterance, input_images)
        fig, ax = plt.subplots()
        ax.imshow(image_out)
        generated_text = text_out.replace("###", "").replace("[IMG0]", "")
        wrapped_generated_text = textwrap.fill(generated_text, 50)
        ax.set_title(wrapped_generated_text, pad=20)
        ax.axis('off')
        plt.savefig(os.path.join(current_dir, f'test_{i}.png'), bbox_inches='tight')
        plt.close(fig)
examples/playground.py
eric-ai-lab/MiniGPT-5
import os
import copy
from pathlib import Path
import json
import torch
import open_clip
from torchvision import transforms
from torchmetrics.image.inception import InceptionScore
from torchmetrics.image.fid import FrechetInceptionDistance
from PIL import Image
from tqdm import tqdm
from torchvision.transforms import ToPILImage, ToTensor
from sentence_transformers import SentenceTransformer, util
from rouge import Rouge 
from nltk.translate import bleu
from nltk.translate import meteor_score
from nltk import word_tokenize
from collections import OrderedDict
import numpy as np
from lightning.pytorch import seed_everything
import string
import argparse

from constants import *

class CLIPEvaluator(object):
    def __init__(self, device, name = 'ViT-L-14') -> None:
        self.device = device
        # self.model, self.preprocess = clip.load(clip_model, device=self.device)
        self.model, _, self.preprocess = open_clip.create_model_and_transforms(name, pretrained='openai', device=device)
        self.tokenizer = open_clip.get_tokenizer(name)                                    # + skip convert PIL to tensor

    @torch.no_grad()
    def encode_text(self, tokens: list) -> torch.Tensor:
        return self.model.encode_text(tokens)

    @torch.no_grad()
    def encode_images(self, images: torch.Tensor) -> torch.Tensor:
        images = self.preprocess(images).unsqueeze(0).to(self.device)
        return self.model.encode_image(images)

    def get_text_features(self, text: str, norm: bool = True) -> torch.Tensor:

        tokens = self.tokenizer(text).to(self.device)

        text_features = self.encode_text(tokens).detach()

        if norm:
            text_features /= text_features.norm(dim=-1, keepdim=True)

        return text_features

    def get_image_features(self, img: torch.Tensor, norm: bool = True) -> torch.Tensor:
        image_features = self.encode_images(img)
        
        if norm:
            image_features /= image_features.clone().norm(dim=-1, keepdim=True)

        return image_features

    def img_to_img_similarity(self, src_images, generated_images):
        src_img_features = self.get_image_features(src_images)
        gen_img_features = self.get_image_features(generated_images)

        return (src_img_features @ gen_img_features.T).mean()

    def txt_to_img_similarity(self, text, generated_images):
        text_features    = self.get_text_features(text)
        gen_img_features = self.get_image_features(generated_images)

        return (text_features @ gen_img_features.T).mean()

def fid_preprocess_image(image):
    image = torch.tensor(image).unsqueeze(0)
    image = image.permute(0, 3, 1, 2) / 255.0
    return transforms.Resize(256, interpolation=transforms.InterpolationMode.BILINEAR)(image)

def cc3m_calculate_metrics(pred_folder):
    all_prediction_files = Path(pred_folder).glob("predictions-*.pt")
    predictions = []
    for file in all_prediction_files:
        predictions.extend(torch.load(file))

    clip_evaluator = CLIPEvaluator(device="cuda" if torch.cuda.is_available() else "cpu")
    inception = InceptionScore(normalize=True).cuda()
    basline_inception = InceptionScore(normalize=True).cuda()
    fid = FrechetInceptionDistance(normalize=True)
    to_pil = ToPILImage()

    fid_vis_processor = transforms.Compose(
                [
                    transforms.Resize(256, interpolation=transforms.InterpolationMode.BILINEAR),
                    transforms.ToTensor(),
                ]
            )
    
    # save true features
    fid_real_feaures_file = "fid_real_feaures.pt"
    if os.path.exists(fid_real_feaures_file):
        real_feaures = torch.load(fid_real_feaures_file)
        fid.real_features_sum = real_feaures[0]
        fid.real_features_cov_sum = real_feaures[1]
        fid.real_features_num_samples = real_feaures[2]
    else:
        image_paths = sorted([os.path.join(CC3M_FOLDER, x) for x in Path(CC3M_FOLDER).glob('val/*/*.jpg')])
        for path in tqdm(image_paths):
            image = Image.open(path).convert("RGB")
            image = fid_vis_processor(image).unsqueeze(0)
            fid.update(image, real=True)

        real_feaures = [fid.real_features_sum, fid.real_features_cov_sum, fid.real_features_num_samples]
        torch.save(real_feaures, "fid_real_feaures.pt")
    baseline_fid = copy.deepcopy(fid)

    clip_similarities = []
    clip_t_similarities = []
    basline_clip_similarities = []
    basline_clip_t_similarities = []

    for prediction in tqdm(predictions):
        _, _, gt_out, predicted_images_ft, predicted_images_nl, gt_image, _, _ = prediction
        if predicted_images_ft is not None:
            gt_image = to_pil(gt_image)
            gt_image_features = clip_evaluator.get_image_features(gt_image)
            gt_text_features = clip_evaluator.get_text_features(gt_out)
            predicted_images_ft_features = clip_evaluator.get_image_features(predicted_images_ft)
            predicted_images_nl_features = clip_evaluator.get_image_features(predicted_images_nl)

            clip_similarity = (gt_image_features @ predicted_images_ft_features.T).mean()
            clip_similarities.append(clip_similarity)

            clip_t_similarity = (gt_text_features @ predicted_images_ft_features.T).mean()
            clip_t_similarities.append(clip_t_similarity)

            basline_clip_similaritie = (gt_image_features @ predicted_images_nl_features.T).mean()
            basline_clip_similarities.append(basline_clip_similaritie)

            basline_clip_t_similarity = (gt_text_features @ predicted_images_nl_features.T).mean()
            basline_clip_t_similarities.append(basline_clip_t_similarity)

            totensor = ToTensor()
            inception.update(totensor(predicted_images_ft).unsqueeze(0).cuda())
            basline_inception.update(totensor(predicted_images_nl).unsqueeze(0).cuda())
            

            fid.update(fid_vis_processor(predicted_images_ft).unsqueeze(0), real=False)
            baseline_fid.update(fid_vis_processor(predicted_images_nl).unsqueeze(0), real=False)


    if len(clip_similarities) == 0:
        overall_clip_similarity = 0
        overall_basline_clip_similarity = 0
        overall_clip_t_similarity = 0
        overall_basline_clip_t_similarity = 0
    else:
        overall_clip_similarity = sum(clip_similarities) / len(clip_similarities)
        overall_basline_clip_similarity = sum(basline_clip_similarities) / len(basline_clip_similarities)
        overall_clip_t_similarity = sum(clip_t_similarities) / len(clip_t_similarities)
        overall_basline_clip_t_similarity = sum(basline_clip_t_similarities) / len(basline_clip_t_similarities)
    overall_fid_score = fid.compute().item()
    overall_basline_fid_score = baseline_fid.compute().item()
    overall_inception_score = inception.compute()[0].item()
    overall_basline_inception_score = basline_inception.compute()[0].item()

    if type(overall_clip_similarity) == torch.Tensor:
        overall_clip_similarity = overall_clip_similarity.item()
    if type(overall_inception_score) == torch.Tensor:
        overall_inception_score = overall_inception_score.item()
    if type(overall_basline_clip_similarity) == torch.Tensor:
        overall_basline_clip_similarity = overall_basline_clip_similarity.item()
    if type(overall_clip_t_similarity) == torch.Tensor:
        overall_clip_t_similarity = overall_clip_t_similarity.item()
    if type(overall_basline_clip_t_similarity) == torch.Tensor:
        overall_basline_clip_t_similarity = overall_basline_clip_t_similarity.item()

    print(f"Overall CLIP similarity: {overall_clip_similarity}")
    print(f"Overall CLIP text similarity: {overall_clip_t_similarity}")
    print(f"Overall Baseline CLIP similarity: {overall_basline_clip_similarity}")
    print(f"Overall Baseline CLIP text similarity: {overall_basline_clip_t_similarity}")
    print(f"Overall Inception Score: {overall_inception_score}")
    print(f"Overall Baseline Inception Score: {overall_basline_inception_score}")
    print(f"Overall FID score: {overall_fid_score}")
    print(f"Overall Baseline FID score: {overall_basline_fid_score}")


    results_dict = {
        'overall_clip_similarity': overall_clip_similarity,
        'overall_clip_t_similarity': overall_clip_t_similarity,
        'overall_basline_clip_similarity': overall_basline_clip_similarity,
        'overall_basline_clip_t_similarity': overall_basline_clip_t_similarity,
        'overall_inception_score': overall_inception_score,
        'overall_basline_inception_score': overall_basline_inception_score,
        'overall_fid_score': overall_fid_score,
        'overall_basline_fid_score': overall_basline_fid_score,
    }

    with open(os.path.join(pred_folder,'results_dict.json'), 'w') as f:
        json.dump(results_dict, f)

def vist_calculate_metrics(pred_folder, calculate_instace_level=True):
    all_prediction_files = Path(pred_folder).glob("predictions-*.pt")
    predictions = []
    for file in all_prediction_files:
        predictions.extend(torch.load(file))

    clip_evaluator = CLIPEvaluator(device="cuda" if torch.cuda.is_available() else "cpu")
    fid = FrechetInceptionDistance(normalize=True)
    baseline_fid = copy.deepcopy(fid)
    inception = InceptionScore(normalize=True).cuda()
    basline_inception = InceptionScore(normalize=True).cuda()

    to_pil = ToPILImage()
    totensor = ToTensor()

    fid_vis_processor = transforms.Compose(
                [
                    transforms.Resize(256, interpolation=transforms.InterpolationMode.BILINEAR),
                    transforms.ToTensor(),
                ]
            )
    with_basline = False

    clip_similarities = []
    clip_t_similarities = []
    basline_clip_similarities = []
    basline_clip_t_similarities = []
    fid_scores = []
    inception_scores = []
    task_dicts = {}
    for prediction in tqdm(predictions):
        _, _, gt_out, predicted_images_ft, predicted_images_nl, gt_image, _, task_name = prediction
        gt_image = to_pil(gt_image)
        task_name, step_id = task_name.split('_')
        if "-" in step_id:
            step_id = step_id.split('-')[0]
        step_id = int(step_id)
        if task_name not in task_dicts:
            task_dicts[task_name] = OrderedDict()
        if step_id not in task_dicts[task_name]:
            task_dicts[task_name][step_id] = {}
        task_dicts[task_name][step_id]['pred_out'] = predicted_images_ft
        task_dicts[task_name][step_id]['gt_out'] = gt_image
        if predicted_images_ft is not None:
            gt_image_features = clip_evaluator.get_image_features(gt_image)
            gt_text_features = clip_evaluator.get_text_features(gt_out)
            predicted_images_ft_features = clip_evaluator.get_image_features(predicted_images_ft)

            clip_similarity = (gt_image_features @ predicted_images_ft_features.T).mean()
            clip_similarities.append(clip_similarity)

            clip_t_similarity = (gt_text_features @ predicted_images_ft_features.T).mean()
            clip_t_similarities.append(clip_t_similarity)

            if calculate_instace_level:
                inception.update(totensor(predicted_images_ft).unsqueeze(0).cuda())
                

                fid.update(fid_vis_processor(predicted_images_ft).unsqueeze(0), real=False)
                fid.update(fid_vis_processor(gt_image).unsqueeze(0), real=True)
            
            if predicted_images_nl is not None:
                with_basline = True
                predicted_images_nl_features = clip_evaluator.get_image_features(predicted_images_nl)
                basline_clip_similaritie = (gt_image_features @ predicted_images_nl_features.T).mean()
                basline_clip_similarities.append(basline_clip_similaritie)

                basline_clip_t_similarity = (gt_text_features @ predicted_images_nl_features.T).mean()
                basline_clip_t_similarities.append(basline_clip_t_similarity)

                task_dicts[task_name][step_id]['pred_out_nl'] = predicted_images_nl

                if calculate_instace_level:
                    basline_inception.update(totensor(predicted_images_nl).unsqueeze(0).cuda())
                    baseline_fid.update(fid_vis_processor(predicted_images_nl).unsqueeze(0), real=False)

    if with_basline and calculate_instace_level:
        baseline_fid.real_features_sum = fid.real_features_sum
        baseline_fid.real_features_cov_sum = fid.real_features_cov_sum
        baseline_fid.real_features_num_samples = fid.real_features_num_samples
        overall_instance_basline_fid_score = baseline_fid.compute().item()
        overall_instance_basline_inception_score = basline_inception.compute()[0].item()
    else:
        overall_instance_basline_fid_score = 0
        overall_instance_basline_inception_score = 0

    if calculate_instace_level:
        overall_instance_fid_score = fid.compute().item()
        overall_instance_inception_score = inception.compute()[0].item()
    else:
        overall_instance_fid_score = 0
        overall_instance_inception_score = 0

    fid.reset()
    baseline_fid.reset()

    if len(clip_similarities) == 0:
        overall_clip_similarity = 0
        # overall_inception_score = 0
        overall_clip_t_similarity = 0
        overall_basline_clip_t_similarity = 0
    else:
        overall_clip_similarity = sum(clip_similarities) / len(clip_similarities)
        # overall_inception_score = sum(inception_scores) / len(inception_scores) 
        overall_clip_t_similarity = sum(clip_t_similarities) / len(clip_t_similarities)
    
    if len(basline_clip_similarities)==0:
        overall_basline_clip_similarity = 0
        overall_basline_clip_t_similarity = 0
    else:
        overall_basline_clip_similarity = sum(basline_clip_similarities) / len(basline_clip_similarities)
        overall_basline_clip_t_similarity = sum(basline_clip_t_similarities) / len(basline_clip_t_similarities)
    
    print(f"Overall CLIP similarity: {overall_clip_similarity}")
    print(f"Overall CLIP text similarity: {overall_clip_t_similarity}")
    print(f"Overall Baseline CLIP similarity: {overall_basline_clip_similarity}")
    print(f"Overall Baseline CLIP text similarity: {overall_basline_clip_t_similarity}")
    print(f"Overall Instance Inception Score: {overall_instance_inception_score}")
    print(f"Overall Instance Baseline Inception Score: {overall_instance_basline_inception_score}")
    print(f"Overall Instance FID score: {overall_instance_fid_score}")
    print(f"Overall Instance Baseline FID score: {overall_instance_basline_fid_score}")

    overall_fid_score = []
    overall_basline_fid_score = []
    for task in tqdm(task_dicts.values(), desc="Calculating task-level FID"):
        if len(task)==1:
            continue
        task = dict(sorted(task.items()))

        gt_images = [fid_vis_processor(t['gt_out']) for t in task.values()]
        gt_images = torch.stack(gt_images)
        fid.update(gt_images, real=True)

        predicted_images_ft = [fid_vis_processor(t['pred_out']) for t in task.values()]
        predicted_images_ft = torch.stack(predicted_images_ft)
        fid.update(predicted_images_ft, real=False)

        overall_fid_score.append(fid.compute().item())
        if 'pred_out_nl' in list(task.values())[0]:
            predicted_images_nl = [fid_vis_processor(t['pred_out_nl']) for t in task.values()]
            predicted_images_nl = torch.stack(predicted_images_nl)
            baseline_fid.update(predicted_images_nl, real=False)
            baseline_fid.real_features_sum = fid.real_features_sum
            baseline_fid.real_features_cov_sum = fid.real_features_cov_sum
            baseline_fid.real_features_num_samples = fid.real_features_num_samples
            overall_basline_fid_score.append(baseline_fid.compute().item())
        fid.reset()
        baseline_fid.reset()

    overall_fid_score = np.mean(overall_fid_score)
    overall_basline_fid_score = np.mean(overall_basline_fid_score)

    if type(overall_clip_similarity) == torch.Tensor:
        overall_clip_similarity = overall_clip_similarity.item()
    if type(overall_basline_clip_similarity) == torch.Tensor:
        overall_basline_clip_similarity = overall_basline_clip_similarity.item()
    if type(overall_clip_t_similarity) == torch.Tensor:
        overall_clip_t_similarity = overall_clip_t_similarity.item()
    if type(overall_basline_clip_t_similarity) == torch.Tensor:
        overall_basline_clip_t_similarity = overall_basline_clip_t_similarity.item()

    print(f"Overall FID score: {overall_fid_score}")
    print(f"Overall Baseline FID score: {overall_basline_fid_score}")


    results_dict = {
        'overall_clip_similarity': overall_clip_similarity,
        'overall_clip_t_similarity': overall_clip_t_similarity,
        'overall_basline_clip_similarity': overall_basline_clip_similarity,
        'overall_basline_clip_t_similarity': overall_basline_clip_t_similarity,
        'overall_fid_score': overall_fid_score,
        'overall_basline_fid_score': overall_basline_fid_score,
        'overall_instance_inception_score': overall_instance_inception_score,
        'overall_instance_basline_inception_score': overall_instance_basline_inception_score,
        'overall_instance_fid_score': overall_instance_fid_score,
        'overall_instance_basline_fid_score': overall_instance_basline_fid_score,
    }

    with open(os.path.join(pred_folder,'results_dict.json'), 'w') as f:
        json.dump(results_dict, f)

def mmdialog_calculate_metrics(pred_folder, calculate_instace_level=True):
    all_prediction_files = Path(pred_folder).glob("predictions-*.pt")
    predictions = []
    for file in all_prediction_files:
        predictions.extend(torch.load(file))

    clip_evaluator = CLIPEvaluator(device="cuda" if torch.cuda.is_available() else "cpu", name="ViT-B-32")
    
    inception = InceptionScore(normalize=True).cuda()
    rouge_model = Rouge()
    sentence_transformer_model = SentenceTransformer('all-MiniLM-L6-v2')

    to_pil = ToPILImage()
    totensor = ToTensor()

    bleu_weights = [(1.0,), (0.5, 0.5)]
    clip_similarities = []
    clip_t_similarities = []
    task_dicts = {}
    for prediction in tqdm(predictions):
        _, pred_out, gt_out, predicted_images_ft, predicted_images_nl, gt_image, _, task_name = prediction
        if '[IMG0]' in gt_out:
            need_image = True
        else:
            need_image = False
        
        if '[IMG0]' in pred_out:
            generate_image = True
        else:
            generate_image = False

        gt_out = gt_out.replace('###', '').split('[IMG0]')[0].strip()
        pred_out = pred_out.replace('###', '').split('[IMG0]')[0].strip()
        if all(c in string.punctuation for c in pred_out):
            pred_out = ''
        if all(c in string.punctuation for c in gt_out):
            gt_out = ''
        gt_image = to_pil(gt_image)
        task_name, step_id = task_name.split('_')
        step_id = int(step_id)
        if task_name not in task_dicts:
            task_dicts[task_name] = OrderedDict()
        if step_id not in task_dicts[task_name]:
            task_dicts[task_name][step_id] = {}
        if len(gt_out) and len(pred_out):
            bleu_score_1, bleu_score_2 = bleu([gt_out], pred_out, bleu_weights)
            rouge_scores = rouge_model.get_scores(pred_out, gt_out)
            rouge_l_score = rouge_scores[0]['rouge-l']['f']

            embeddings1 = sentence_transformer_model.encode(pred_out, convert_to_tensor=True, show_progress_bar=False)
            embeddings2 = sentence_transformer_model.encode(gt_out, convert_to_tensor=True, show_progress_bar=False)
            sbert_score = util.pytorch_cos_sim(embeddings1, embeddings2).item()

            task_dicts[task_name][step_id]['bleu_score_1'] = bleu_score_1
            task_dicts[task_name][step_id]['bleu_score_2'] = bleu_score_2
            task_dicts[task_name][step_id]['rouge_l_score'] = rouge_l_score
            task_dicts[task_name][step_id]['sbert_score'] = sbert_score

            gt_text_features = clip_evaluator.get_text_features(gt_out)
            pred_text_features = clip_evaluator.get_text_features(pred_out)
            clip_t_similarity = (gt_text_features @ pred_text_features.T).mean().item()

            task_dicts[task_name][step_id]['clip_t_similarity'] = clip_t_similarity
        
        elif len(gt_out):
            task_dicts[task_name][step_id]['bleu_score_1'] = 0
            task_dicts[task_name][step_id]['bleu_score_2'] = 0
            task_dicts[task_name][step_id]['rouge_l_score'] = 0
            task_dicts[task_name][step_id]['sbert_score'] = 0
            task_dicts[task_name][step_id]['clip_t_similarity'] = -1
        elif len(pred_out):
            task_dicts[task_name][step_id]['clip_t_similarity'] = -2
        
        if generate_image:
            inception.update(totensor(predicted_images_ft).unsqueeze(0).cuda())

        if need_image and generate_image:
            gt_image_features = clip_evaluator.get_image_features(gt_image)
            predicted_images_ft_features = clip_evaluator.get_image_features(predicted_images_ft)
            clip_i_similarity = (gt_image_features @ predicted_images_ft_features.T).mean().item()

            task_dicts[task_name][step_id]['clip_i_similarity'] = clip_i_similarity
        elif need_image:
            task_dicts[task_name][step_id]['clip_i_similarity'] = -1
        elif generate_image:
            task_dicts[task_name][step_id]['clip_i_similarity'] = -2

    if calculate_instace_level:
        overall_instance_inception_score = inception.compute()[0].item()
    else:
        overall_instance_fid_score = 0
        overall_instance_inception_score = 0

    print(f"Overall Instance Inception Score: {overall_instance_inception_score}")

    overall_bleu_score_1 = []
    overall_bleu_score_2 = []
    overall_rouge_l_score = []
    overall_sbert_score = []
    overall_clip_f1 = []
    for task in tqdm(task_dicts.values(), desc="Calculating task-level FID"):
        task = dict(sorted(task.items()))

        task_bleu_score_1 = [t['bleu_score_1'] for t in task.values() if 'bleu_score_1' in t]
        task_bleu_score_2 = [t['bleu_score_2'] for t in task.values() if 'bleu_score_2' in t]
        task_rouge_l_score = [t['rouge_l_score'] for t in task.values() if 'rouge_l_score' in t]
        task_sbert_score = [t['sbert_score'] for t in task.values() if 'sbert_score' in t]

        task_clip_t_similarity = [t['clip_t_similarity'] for t in task.values() if 'clip_t_similarity' in t]
        task_clip_i_similarity = [t['clip_i_similarity'] for t in task.values() if 'clip_i_similarity' in t]

        task_clip_similarity = task_clip_t_similarity+task_clip_i_similarity
        valid_task_clip_similarity = [t for t in task_clip_similarity if t!=-1 and t!=-2]
        if len(valid_task_clip_similarity) == 0:
            clip_f1 = 0
        else:
            valid_clip_sum = sum(valid_task_clip_similarity)

            clip_precision = valid_clip_sum / len([t for t in task_clip_similarity if t!=-1])
            clip_recall = valid_clip_sum / len([t for t in task_clip_similarity if t!=-2])

            clip_f1 = 2 * clip_precision * clip_recall / (clip_precision + clip_recall)

        if len(task_bleu_score_1):
            overall_bleu_score_1.append(np.mean(task_bleu_score_1))
            overall_bleu_score_2.append(np.mean(task_bleu_score_2))
            overall_rouge_l_score.append(np.mean(task_rouge_l_score))
            overall_sbert_score.append(np.mean(task_sbert_score))
        overall_clip_f1.append(clip_f1)
    
    overall_bleu_score_1 = np.mean(overall_bleu_score_1)
    overall_bleu_score_2 = np.mean(overall_bleu_score_2)
    overall_rouge_l_score = np.mean(overall_rouge_l_score)
    overall_sbert_score = np.mean(overall_sbert_score)
    overall_clip_f1 = np.mean(overall_clip_f1)

    print(f"Overall BLEU-1 score: {overall_bleu_score_1}")
    print(f"Overall BLEU-2 score: {overall_bleu_score_2}")
    print(f"Overall ROUGE-L score: {overall_rouge_l_score}")
    print(f"Overall Sentence-BERT score: {overall_sbert_score}")
    print(f"Overall CLIP F1 score: {overall_clip_f1}")

    results_dict = {
        'overall_bleu_score_1': overall_bleu_score_1,
        'overall_bleu_score_2': overall_bleu_score_2,
        'overall_rouge_l_score': overall_rouge_l_score,
        'overall_sbert_score': overall_sbert_score,
        'overall_clip_f1': overall_clip_f1,
        'overall_instance_inception_score': overall_instance_inception_score,
    }

    with open(os.path.join(pred_folder,'results_dict.json'), 'w') as f:
        json.dump(results_dict, f)

if __name__=="__main__":
    seed_everything(42)
    parser = argparse.ArgumentParser()
    parser.add_argument('--test_weight', type=str, help='an integer for the accumulator')
    args = parser.parse_args()
    test_weight = args.test_weight
    output_folder = os.path.join(OUTPUT_FOLDER, test_weight.split(".")[0]+OUTPUT_SUFFIX)
    if "CC3M" in DATAFOLDER:
        cc3m_calculate_metrics(output_folder)
    elif "MMDialog" in DATAFOLDER:
        mmdialog_calculate_metrics(output_folder)
    else:
        vist_calculate_metrics(output_folder)
    exit()
metric.py
eric-ai-lab/MiniGPT-5
"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import logging
import json
from typing import Dict

from omegaconf import OmegaConf
from minigpt4.common.registry import registry


class Config:
    def __init__(self, args):
        self.config = {}

        self.args = args

        # Register the config and configuration for setup
        registry.register("configuration", self)

        user_config = self._build_opt_list(self.args.options)

        config = OmegaConf.load(self.args.cfg_path)

        runner_config = self.build_runner_config(config)
        model_config = self.build_model_config(config, **user_config)
        dataset_config = self.build_dataset_config(config)

        # Validate the user-provided runner configuration
        # model and dataset configuration are supposed to be validated by the respective classes
        # [TODO] validate the model/dataset configuration
        # self._validate_runner_config(runner_config)

        # Override the default configuration with user options.
        self.config = OmegaConf.merge(
            runner_config, model_config, dataset_config, user_config
        )

    def _validate_runner_config(self, runner_config):
        """
        This method validates the configuration, such that
            1) all the user specified options are valid;
            2) no type mismatches between the user specified options and the config.
        """
        runner_config_validator = create_runner_config_validator()
        runner_config_validator.validate(runner_config)

    def _build_opt_list(self, opts):
        opts_dot_list = self._convert_to_dot_list(opts)
        return OmegaConf.from_dotlist(opts_dot_list)

    @staticmethod
    def build_model_config(config, **kwargs):
        model = config.get("model", None)
        assert model is not None, "Missing model configuration file."

        model_cls = registry.get_model_class(model.arch)
        assert model_cls is not None, f"Model '{model.arch}' has not been registered."

        model_type = kwargs.get("model.model_type", None)
        if not model_type:
            model_type = model.get("model_type", None)
        # else use the model type selected by user.

        assert model_type is not None, "Missing model_type."

        model_config_path = model_cls.default_config_path(model_type=model_type)

        model_config = OmegaConf.create()
        # hierarchy override, customized config > default config
        model_config = OmegaConf.merge(
            model_config,
            OmegaConf.load(model_config_path),
            {"model": config["model"]},
        )

        return model_config

    @staticmethod
    def build_runner_config(config):
        return {"run": config.run}

    @staticmethod
    def build_dataset_config(config):
        datasets = config.get("datasets", None)
        if datasets is None:
            raise KeyError(
                "Expecting 'datasets' as the root key for dataset configuration."
            )

        dataset_config = OmegaConf.create()

        for dataset_name in datasets:
            builder_cls = registry.get_builder_class(dataset_name)

            dataset_config_type = datasets[dataset_name].get("type", "default")
            dataset_config_path = builder_cls.default_config_path(
                type=dataset_config_type
            )

            # hierarchy override, customized config > default config
            dataset_config = OmegaConf.merge(
                dataset_config,
                OmegaConf.load(dataset_config_path),
                {"datasets": {dataset_name: config["datasets"][dataset_name]}},
            )

        return dataset_config

    def _convert_to_dot_list(self, opts):
        if opts is None:
            opts = []

        if len(opts) == 0:
            return opts

        has_equal = opts[0].find("=") != -1

        if has_equal:
            return opts

        return [(opt + "=" + value) for opt, value in zip(opts[0::2], opts[1::2])]

    def get_config(self):
        return self.config

    @property
    def run_cfg(self):
        return self.config.run

    @property
    def datasets_cfg(self):
        return self.config.datasets

    @property
    def model_cfg(self):
        return self.config.model

    def pretty_print(self):
        logging.info("\n=====  Running Parameters    =====")
        logging.info(self._convert_node_to_json(self.config.run))

        logging.info("\n======  Dataset Attributes  ======")
        datasets = self.config.datasets

        for dataset in datasets:
            if dataset in self.config.datasets:
                logging.info(f"\n======== {dataset} =======")
                dataset_config = self.config.datasets[dataset]
                logging.info(self._convert_node_to_json(dataset_config))
            else:
                logging.warning(f"No dataset named '{dataset}' in config. Skipping")

        logging.info(f"\n======  Model Attributes  ======")
        logging.info(self._convert_node_to_json(self.config.model))

    def _convert_node_to_json(self, node):
        container = OmegaConf.to_container(node, resolve=True)
        return json.dumps(container, indent=4, sort_keys=True)

    def to_dict(self):
        return OmegaConf.to_container(self.config)


def node_to_dict(node):
    return OmegaConf.to_container(node)


class ConfigValidator:
    """
    This is a preliminary implementation to centralize and validate the configuration.
    May be altered in the future.

    A helper class to validate configurations from yaml file.

    This serves the following purposes:
        1. Ensure all the options in the yaml are defined, raise error if not.
        2. when type mismatches are found, the validator will raise an error.
        3. a central place to store and display helpful messages for supported configurations.

    """

    class _Argument:
        def __init__(self, name, choices=None, type=None, help=None):
            self.name = name
            self.val = None
            self.choices = choices
            self.type = type
            self.help = help

        def __str__(self):
            s = f"{self.name}={self.val}"
            if self.type is not None:
                s += f", ({self.type})"
            if self.choices is not None:
                s += f", choices: {self.choices}"
            if self.help is not None:
                s += f", ({self.help})"
            return s

    def __init__(self, description):
        self.description = description

        self.arguments = dict()

        self.parsed_args = None

    def __getitem__(self, key):
        assert self.parsed_args is not None, "No arguments parsed yet."

        return self.parsed_args[key]

    def __str__(self) -> str:
        return self.format_help()

    def add_argument(self, *args, **kwargs):
        """
        Assume the first argument is the name of the argument.
        """
        self.arguments[args[0]] = self._Argument(*args, **kwargs)

    def validate(self, config=None):
        """
        Convert yaml config (dict-like) to list, required by argparse.
        """
        for k, v in config.items():
            assert (
                k in self.arguments
            ), f"""{k} is not a valid argument. Support arguments are {self.format_arguments()}."""

            if self.arguments[k].type is not None:
                try:
                    self.arguments[k].val = self.arguments[k].type(v)
                except ValueError:
                    raise ValueError(f"{k} is not a valid {self.arguments[k].type}.")

            if self.arguments[k].choices is not None:
                assert (
                    v in self.arguments[k].choices
                ), f"""{k} must be one of {self.arguments[k].choices}."""

        return config

    def format_arguments(self):
        return str([f"{k}" for k in sorted(self.arguments.keys())])

    def format_help(self):
        # description + key-value pair string for each argument
        help_msg = str(self.description)
        return help_msg + ", available arguments: " + self.format_arguments()

    def print_help(self):
        # display help message
        print(self.format_help())


def create_runner_config_validator():
    validator = ConfigValidator(description="Runner configurations")

    validator.add_argument(
        "runner",
        type=str,
        choices=["runner_base", "runner_iter"],
        help="""Runner to use. The "runner_base" uses epoch-based training while iter-based
            runner runs based on iters. Default: runner_base""",
    )
    # add argumetns for training dataset ratios
    validator.add_argument(
        "train_dataset_ratios",
        type=Dict[str, float],
        help="""Ratios of training dataset. This is used in iteration-based runner.
        Do not support for epoch-based runner because how to define an epoch becomes tricky.
        Default: None""",
    )
    validator.add_argument(
        "max_iters",
        type=float,
        help="Maximum number of iterations to run.",
    )
    validator.add_argument(
        "max_epoch",
        type=int,
        help="Maximum number of epochs to run.",
    )
    # add arguments for iters_per_inner_epoch
    validator.add_argument(
        "iters_per_inner_epoch",
        type=float,
        help="Number of iterations per inner epoch. This is required when runner is runner_iter.",
    )
    lr_scheds_choices = registry.list_lr_schedulers()
    validator.add_argument(
        "lr_sched",
        type=str,
        choices=lr_scheds_choices,
        help="Learning rate scheduler to use, from {}".format(lr_scheds_choices),
    )
    task_choices = registry.list_tasks()
    validator.add_argument(
        "task",
        type=str,
        choices=task_choices,
        help="Task to use, from {}".format(task_choices),
    )
    # add arguments for init_lr
    validator.add_argument(
        "init_lr",
        type=float,
        help="Initial learning rate. This will be the learning rate after warmup and before decay.",
    )
    # add arguments for min_lr
    validator.add_argument(
        "min_lr",
        type=float,
        help="Minimum learning rate (after decay).",
    )
    # add arguments for warmup_lr
    validator.add_argument(
        "warmup_lr",
        type=float,
        help="Starting learning rate for warmup.",
    )
    # add arguments for learning rate decay rate
    validator.add_argument(
        "lr_decay_rate",
        type=float,
        help="Learning rate decay rate. Required if using a decaying learning rate scheduler.",
    )
    # add arguments for weight decay
    validator.add_argument(
        "weight_decay",
        type=float,
        help="Weight decay rate.",
    )
    # add arguments for training batch size
    validator.add_argument(
        "batch_size_train",
        type=int,
        help="Training batch size.",
    )
    # add arguments for evaluation batch size
    validator.add_argument(
        "batch_size_eval",
        type=int,
        help="Evaluation batch size, including validation and testing.",
    )
    # add arguments for number of workers for data loading
    validator.add_argument(
        "num_workers",
        help="Number of workers for data loading.",
    )
    # add arguments for warm up steps
    validator.add_argument(
        "warmup_steps",
        type=int,
        help="Number of warmup steps. Required if a warmup schedule is used.",
    )
    # add arguments for random seed
    validator.add_argument(
        "seed",
        type=int,
        help="Random seed.",
    )
    # add arguments for output directory
    validator.add_argument(
        "output_dir",
        type=str,
        help="Output directory to save checkpoints and logs.",
    )
    # add arguments for whether only use evaluation
    validator.add_argument(
        "evaluate",
        help="Whether to only evaluate the model. If true, training will not be performed.",
    )
    # add arguments for splits used for training, e.g. ["train", "val"]
    validator.add_argument(
        "train_splits",
        type=list,
        help="Splits to use for training.",
    )
    # add arguments for splits used for validation, e.g. ["val"]
    validator.add_argument(
        "valid_splits",
        type=list,
        help="Splits to use for validation. If not provided, will skip the validation.",
    )
    # add arguments for splits used for testing, e.g. ["test"]
    validator.add_argument(
        "test_splits",
        type=list,
        help="Splits to use for testing. If not provided, will skip the testing.",
    )
    # add arguments for accumulating gradient for iterations
    validator.add_argument(
        "accum_grad_iters",
        type=int,
        help="Number of iterations to accumulate gradient for.",
    )

    # ====== distributed training ======
    validator.add_argument(
        "device",
        type=str,
        choices=["cpu", "cuda"],
        help="Device to use. Support 'cuda' or 'cpu' as for now.",
    )
    validator.add_argument(
        "world_size",
        type=int,
        help="Number of processes participating in the job.",
    )
    validator.add_argument("dist_url", type=str)
    validator.add_argument("distributed", type=bool)
    # add arguments to opt using distributed sampler during evaluation or not
    validator.add_argument(
        "use_dist_eval_sampler",
        type=bool,
        help="Whether to use distributed sampler during evaluation or not.",
    )

    # ====== task specific ======
    # generation task specific arguments
    # add arguments for maximal length of text output
    validator.add_argument(
        "max_len",
        type=int,
        help="Maximal length of text output.",
    )
    # add arguments for minimal length of text output
    validator.add_argument(
        "min_len",
        type=int,
        help="Minimal length of text output.",
    )
    # add arguments number of beams
    validator.add_argument(
        "num_beams",
        type=int,
        help="Number of beams used for beam search.",
    )

    # vqa task specific arguments
    # add arguments for number of answer candidates
    validator.add_argument(
        "num_ans_candidates",
        type=int,
        help="""For ALBEF and BLIP, these models first rank answers according to likelihood to select answer candidates.""",
    )
    # add arguments for inference method
    validator.add_argument(
        "inference_method",
        type=str,
        choices=["generate", "rank"],
        help="""Inference method to use for question answering. If rank, requires an answer list.""",
    )

    # ====== model specific ======
    validator.add_argument(
        "k_test",
        type=int,
        help="Number of top k most similar samples from ITC/VTC selection to be tested.",
    )

    return validator

minigpt4/common/config.py
eric-ai-lab/MiniGPT-5
import argparse
import time
from PIL import Image

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaTokenizer
from transformers import StoppingCriteria, StoppingCriteriaList

import dataclasses
from enum import auto, Enum
from typing import List, Tuple, Any

from minigpt4.common.registry import registry


class SeparatorStyle(Enum):
    """Different separator style."""
    SINGLE = auto()
    TWO = auto()


@dataclasses.dataclass
class Conversation:
    """A class that keeps all conversation history."""
    system: str
    roles: List[str]
    messages: List[List[str]]
    offset: int
    # system_img: List[Image.Image] = []
    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
    sep: str = "###"
    sep2: str = None

    skip_next: bool = False
    conv_id: Any = None

    def get_prompt(self):
        if self.sep_style == SeparatorStyle.SINGLE:
            ret = self.system + self.sep
            for role, message in self.messages:
                if message:
                    ret += role + ": " + message + self.sep
                else:
                    ret += role + ":"
            return ret
        elif self.sep_style == SeparatorStyle.TWO:
            seps = [self.sep, self.sep2]
            ret = self.system + seps[0]
            for i, (role, message) in enumerate(self.messages):
                if message:
                    ret += role + ": " + message + seps[i % 2]
                else:
                    ret += role + ":"
            return ret
        else:
            raise ValueError(f"Invalid style: {self.sep_style}")

    def append_message(self, role, message):
        self.messages.append([role, message])

    def to_gradio_chatbot(self):
        ret = []
        for i, (role, msg) in enumerate(self.messages[self.offset:]):
            if i % 2 == 0:
                ret.append([msg, None])
            else:
                ret[-1][-1] = msg
        return ret

    def copy(self):
        return Conversation(
            system=self.system,
            # system_img=self.system_img,
            roles=self.roles,
            messages=[[x, y] for x, y in self.messages],
            offset=self.offset,
            sep_style=self.sep_style,
            sep=self.sep,
            sep2=self.sep2,
            conv_id=self.conv_id)

    def dict(self):
        return {
            "system": self.system,
            # "system_img": self.system_img,
            "roles": self.roles,
            "messages": self.messages,
            "offset": self.offset,
            "sep": self.sep,
            "sep2": self.sep2,
            "conv_id": self.conv_id,
        }


class StoppingCriteriaSub(StoppingCriteria):

    def __init__(self, stops=[], encounters=1):
        super().__init__()
        self.stops = stops

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
        for stop in self.stops:
            if torch.all((stop == input_ids[0][-len(stop):])).item():
                return True

        return False


CONV_VISION = Conversation(
    system="Give the following image: <Img>ImageContent</Img>. "
           "You will be able to see the image once I provide it to you. Please answer my questions.",
    roles=("Human", "Assistant"),
    messages=[],
    offset=2,
    sep_style=SeparatorStyle.SINGLE,
    sep="###",
)



class Chat:
    def __init__(self, model, vis_processor, device='cuda:0'):
        self.device = device
        self.model = model
        self.vis_processor = vis_processor
        stop_words_ids = [torch.tensor([835]).to(self.device),
                          torch.tensor([2277, 29937]).to(self.device)]  # '###' can be encoded in two different ways.
        self.stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])

    def ask(self, text, conv):
        if len(conv.messages) > 0 and conv.messages[-1][0] == conv.roles[0] \
                and conv.messages[-1][1][-6:] == '</Img>':  # last message is image.
            conv.messages[-1][1] = ' '.join([conv.messages[-1][1], text])
        else:
            conv.append_message(conv.roles[0], text)

    def answer(self, conv, img_list, max_new_tokens=300, num_beams=1, min_length=1, top_p=0.9,
               repetition_penalty=1.0, length_penalty=1, temperature=1.0, max_length=2000):
        conv.append_message(conv.roles[1], None)
        embs = self.get_context_emb(conv, img_list)

        current_max_len = embs.shape[1] + max_new_tokens
        if current_max_len - max_length > 0:
            print('Warning: The number of tokens in current conversation exceeds the max length. '
                  'The model will not see the contexts outside the range.')
        begin_idx = max(0, current_max_len - max_length)

        embs = embs[:, begin_idx:]

        outputs = self.model.llama_model.generate(
            inputs_embeds=embs,
            max_new_tokens=max_new_tokens,
            stopping_criteria=self.stopping_criteria,
            num_beams=num_beams,
            do_sample=True,
            min_length=min_length,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            length_penalty=length_penalty,
            temperature=temperature,
        )
        output_token = outputs[0]
        if output_token[0] == 0:  # the model might output an unknown token <unk> at the beginning. remove it
            output_token = output_token[1:]
        if output_token[0] == 1:  # some users find that there is a start token <s> at the beginning. remove it
            output_token = output_token[1:]
        output_text = self.model.llama_tokenizer.decode(output_token, add_special_tokens=False)
        output_text = output_text.split('###')[0]  # remove the stop sign '###'
        output_text = output_text.split('Assistant:')[-1].strip()
        conv.messages[-1][1] = output_text
        return output_text, output_token.cpu().numpy()

    def upload_img(self, image, conv, img_list):
        if isinstance(image, str):  # is a image path
            raw_image = Image.open(image).convert('RGB')
            image = self.vis_processor(raw_image).unsqueeze(0).to(self.device)
        elif isinstance(image, Image.Image):
            raw_image = image
            image = self.vis_processor(raw_image).unsqueeze(0).to(self.device)
        elif isinstance(image, torch.Tensor):
            if len(image.shape) == 3:
                image = image.unsqueeze(0)
            image = image.to(self.device)

        image_emb, _ = self.model.encode_img(image)
        img_list.append(image_emb)
        conv.append_message(conv.roles[0], "<Img><ImageHere></Img>")
        msg = "Received."
        # self.conv.append_message(self.conv.roles[1], msg)
        return msg

    def get_context_emb(self, conv, img_list):
        prompt = conv.get_prompt()
        prompt_segs = prompt.split('<ImageHere>')
        assert len(prompt_segs) == len(img_list) + 1, "Unmatched numbers of image placeholders and images."
        seg_tokens = [
            self.model.llama_tokenizer(
                seg, return_tensors="pt", add_special_tokens=i == 0).to(self.device).input_ids
            # only add bos to the first seg
            for i, seg in enumerate(prompt_segs)
        ]
        seg_embs = [self.model.llama_model.model.embed_tokens(seg_t) for seg_t in seg_tokens]
        mixed_embs = [emb for pair in zip(seg_embs[:-1], img_list) for emb in pair] + [seg_embs[-1]]
        mixed_embs = torch.cat(mixed_embs, dim=1)
        return mixed_embs



minigpt4/conversation/conversation.py
eric-ai-lab/MiniGPT-5
"""
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import logging
import torch
from omegaconf import OmegaConf

from minigpt4.common.registry import registry
from minigpt4.models.base_model import BaseModel
from minigpt4.models.blip2 import Blip2Base
from minigpt4.models.mini_gpt4 import MiniGPT4
from minigpt4.models.mini_gpt5 import MiniGPT5
from minigpt4.processors.base_processor import BaseProcessor


__all__ = [
    "load_model",
    "BaseModel",
    "Blip2Base",
    "MiniGPT4",
    "MiniGPT5"
]


def load_model(name, model_type, is_eval=False, device="cpu", checkpoint=None):
    """
    Load supported models.

    To list all available models and types in registry:
    >>> from minigpt4.models import model_zoo
    >>> print(model_zoo)

    Args:
        name (str): name of the model.
        model_type (str): type of the model.
        is_eval (bool): whether the model is in eval mode. Default: False.
        device (str): device to use. Default: "cpu".
        checkpoint (str): path or to checkpoint. Default: None.
            Note that expecting the checkpoint to have the same keys in state_dict as the model.

    Returns:
        model (torch.nn.Module): model.
    """

    model = registry.get_model_class(name).from_pretrained(model_type=model_type)

    if checkpoint is not None:
        model.load_checkpoint(checkpoint)

    if is_eval:
        model.eval()

    if device == "cpu":
        model = model.float()

    return model.to(device)


def load_preprocess(config):
    """
    Load preprocessor configs and construct preprocessors.

    If no preprocessor is specified, return BaseProcessor, which does not do any preprocessing.

    Args:
        config (dict): preprocessor configs.

    Returns:
        vis_processors (dict): preprocessors for visual inputs.
        txt_processors (dict): preprocessors for text inputs.

        Key is "train" or "eval" for processors used in training and evaluation respectively.
    """

    def _build_proc_from_cfg(cfg):
        return (
            registry.get_processor_class(cfg.name).from_config(cfg)
            if cfg is not None
            else BaseProcessor()
        )

    vis_processors = dict()
    txt_processors = dict()

    vis_proc_cfg = config.get("vis_processor")
    txt_proc_cfg = config.get("text_processor")

    if vis_proc_cfg is not None:
        vis_train_cfg = vis_proc_cfg.get("train")
        vis_eval_cfg = vis_proc_cfg.get("eval")
    else:
        vis_train_cfg = None
        vis_eval_cfg = None

    vis_processors["train"] = _build_proc_from_cfg(vis_train_cfg)
    vis_processors["eval"] = _build_proc_from_cfg(vis_eval_cfg)

    if txt_proc_cfg is not None:
        txt_train_cfg = txt_proc_cfg.get("train")
        txt_eval_cfg = txt_proc_cfg.get("eval")
    else:
        txt_train_cfg = None
        txt_eval_cfg = None

    txt_processors["train"] = _build_proc_from_cfg(txt_train_cfg)
    txt_processors["eval"] = _build_proc_from_cfg(txt_eval_cfg)

    return vis_processors, txt_processors


def load_model_and_preprocess(name, model_type, is_eval=False, device="cpu"):
    """
    Load model and its related preprocessors.

    List all available models and types in registry:
    >>> from minigpt4.models import model_zoo
    >>> print(model_zoo)

    Args:
        name (str): name of the model.
        model_type (str): type of the model.
        is_eval (bool): whether the model is in eval mode. Default: False.
        device (str): device to use. Default: "cpu".

    Returns:
        model (torch.nn.Module): model.
        vis_processors (dict): preprocessors for visual inputs.
        txt_processors (dict): preprocessors for text inputs.
    """
    model_cls = registry.get_model_class(name)

    # load model
    model = model_cls.from_pretrained(model_type=model_type)

    if is_eval:
        model.eval()

    # load preprocess
    cfg = OmegaConf.load(model_cls.default_config_path(model_type))
    if cfg is not None:
        preprocess_cfg = cfg.preprocess

        vis_processors, txt_processors = load_preprocess(preprocess_cfg)
    else:
        vis_processors, txt_processors = None, None
        logging.info(
            f"""No default preprocess for model {name} ({model_type}).
                This can happen if the model is not finetuned on downstream datasets,
                or it is not intended for direct use without finetuning.
            """
        )

    if device == "cpu" or device == torch.device("cpu"):
        model = model.float()

    return model.to(device), vis_processors, txt_processors


class ModelZoo:
    """
    A utility class to create string representation of available model architectures and types.

    >>> from minigpt4.models import model_zoo
    >>> # list all available models
    >>> print(model_zoo)
    >>> # show total number of models
    >>> print(len(model_zoo))
    """

    def __init__(self) -> None:
        self.model_zoo = {
            k: list(v.PRETRAINED_MODEL_CONFIG_DICT.keys())
            for k, v in registry.mapping["model_name_mapping"].items()
        }

    def __str__(self) -> str:
        return (
            "=" * 50
            + "\n"
            + f"{'Architectures':<30} {'Types'}\n"
            + "=" * 50
            + "\n"
            + "\n".join(
                [
                    f"{name:<30} {', '.join(types)}"
                    for name, types in self.model_zoo.items()
                ]
            )
        )

    def __iter__(self):
        return iter(self.model_zoo.items())

    def __len__(self):
        return sum([len(v) for v in self.model_zoo.values()])


model_zoo = ModelZoo()

minigpt4/models/__init__.py
eric-ai-lab/MiniGPT-5
from transformers.generation.utils import *

def sample(
        self,
        input_ids: torch.LongTensor,
        logits_processor: Optional[LogitsProcessorList] = None,
        stopping_criteria: Optional[StoppingCriteriaList] = None,
        logits_warper: Optional[LogitsProcessorList] = None,
        max_length: Optional[int] = None,
        pad_token_id: Optional[int] = None,
        eos_token_id: Optional[Union[int, List[int]]] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        output_scores: Optional[bool] = None,
        return_dict_in_generate: Optional[bool] = None,
        synced_gpus: bool = False,
        streamer: Optional["BaseStreamer"] = None,
        **model_kwargs,
    ) -> Union[SampleOutput, torch.LongTensor]:
        r"""
        Generates sequences of token ids for models with a language modeling head using **multinomial sampling** and
        can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models.

        <Tip warning={true}>

        In most cases, you do not need to call [`~generation.GenerationMixin.sample`] directly. Use generate() instead.
        For an overview of generation strategies and code examples, check the [following
        guide](../generation_strategies).

        </Tip>

        Parameters:
            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
                The sequence used as a prompt for the generation.
            logits_processor (`LogitsProcessorList`, *optional*):
                An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]
                used to modify the prediction scores of the language modeling head applied at each generation step.
            stopping_criteria (`StoppingCriteriaList`, *optional*):
                An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`]
                used to tell if the generation loop should stop.
            logits_warper (`LogitsProcessorList`, *optional*):
                An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsWarper`] used
                to warp the prediction score distribution of the language modeling head applied before multinomial
                sampling at each generation step.
            max_length (`int`, *optional*, defaults to 20):
                **DEPRECATED**. Use `logits_processor` or `stopping_criteria` directly to cap the number of generated
                tokens. The maximum length of the sequence to be generated.
            pad_token_id (`int`, *optional*):
                The id of the *padding* token.
            eos_token_id (`Union[int, List[int]]`, *optional*):
                The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens.
            output_attentions (`bool`, *optional*, defaults to `False`):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more details.
            output_hidden_states (`bool`, *optional*, defaults to `False`):
                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
                for more details.
            output_scores (`bool`, *optional*, defaults to `False`):
                Whether or not to return the prediction scores. See `scores` under returned tensors for more details.
            return_dict_in_generate (`bool`, *optional*, defaults to `False`):
                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
            synced_gpus (`bool`, *optional*, defaults to `False`):
                Whether to continue running the while loop until max_length (needed for ZeRO stage 3)
            streamer (`BaseStreamer`, *optional*):
                Streamer object that will be used to stream the generated sequences. Generated tokens are passed
                through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
            model_kwargs:
                Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is
                an encoder-decoder model the kwargs should include `encoder_outputs`.

        Return:
            [`~generation.SampleDecoderOnlyOutput`], [`~generation.SampleEncoderDecoderOutput`] or `torch.LongTensor`:
            A `torch.LongTensor` containing the generated tokens (default behaviour) or a
            [`~generation.SampleDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and
            `return_dict_in_generate=True` or a [`~generation.SampleEncoderDecoderOutput`] if
            `model.config.is_encoder_decoder=True`.

        Examples:

        ```python
        >>> from transformers import (
        ...     AutoTokenizer,
        ...     AutoModelForCausalLM,
        ...     LogitsProcessorList,
        ...     MinLengthLogitsProcessor,
        ...     TopKLogitsWarper,
        ...     TemperatureLogitsWarper,
        ...     StoppingCriteriaList,
        ...     MaxLengthCriteria,
        ... )
        >>> import torch

        >>> tokenizer = AutoTokenizer.from_pretrained("gpt2")
        >>> model = AutoModelForCausalLM.from_pretrained("gpt2")

        >>> # set pad_token_id to eos_token_id because GPT2 does not have a EOS token
        >>> model.config.pad_token_id = model.config.eos_token_id
        >>> model.generation_config.pad_token_id = model.config.eos_token_id

        >>> input_prompt = "Today is a beautiful day, and"
        >>> input_ids = tokenizer(input_prompt, return_tensors="pt").input_ids

        >>> # instantiate logits processors
        >>> logits_processor = LogitsProcessorList(
        ...     [
        ...         MinLengthLogitsProcessor(15, eos_token_id=model.generation_config.eos_token_id),
        ...     ]
        ... )
        >>> # instantiate logits processors
        >>> logits_warper = LogitsProcessorList(
        ...     [
        ...         TopKLogitsWarper(50),
        ...         TemperatureLogitsWarper(0.7),
        ...     ]
        ... )

        >>> stopping_criteria = StoppingCriteriaList([MaxLengthCriteria(max_length=20)])

        >>> torch.manual_seed(0)  # doctest: +IGNORE_RESULT
        >>> outputs = model.sample(
        ...     input_ids,
        ...     logits_processor=logits_processor,
        ...     logits_warper=logits_warper,
        ...     stopping_criteria=stopping_criteria,
        ... )

        >>> tokenizer.batch_decode(outputs, skip_special_tokens=True)
        ['Today is a beautiful day, and we must do everything possible to make it a day of celebration.']
        ```"""
        # init values
        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
        if max_length is not None:
            warnings.warn(
                "`max_length` is deprecated in this function, use"
                " `stopping_criteria=StoppingCriteriaList(MaxLengthCriteria(max_length=max_length))` instead.",
                UserWarning,
            )
            stopping_criteria = validate_stopping_criteria(stopping_criteria, max_length)
        logits_warper = logits_warper if logits_warper is not None else LogitsProcessorList()
        pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id
        eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id
        if isinstance(eos_token_id, int):
            eos_token_id = [eos_token_id]
        eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None
        output_scores = output_scores if output_scores is not None else self.generation_config.output_scores
        output_attentions = (
            output_attentions if output_attentions is not None else self.generation_config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states
        )
        return_dict_in_generate = (
            return_dict_in_generate
            if return_dict_in_generate is not None
            else self.generation_config.return_dict_in_generate
        )

        # init attention / hidden states / scores tuples
        scores = () if (return_dict_in_generate and output_scores) else None
        decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
        cross_attentions = () if (return_dict_in_generate and output_attentions) else None
        decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None

        # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
        if return_dict_in_generate and self.config.is_encoder_decoder:
            encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
            encoder_hidden_states = (
                model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
            )

        # keep track of which sequences are already finished
        unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)

        this_peer_finished = False  # used by synced_gpus only
        # auto-regressive generation
        while True:
            if synced_gpus:
                # Under synced_gpus the `forward` call must continue until all gpus complete their sequence.
                # The following logic allows an early break if all peers finished generating their sequence
                this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0).to(input_ids.device)
                # send 0.0 if we finished, 1.0 otherwise
                dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM)
                # did all peers finish? the reduced sum will be 0.0 then
                if this_peer_finished_flag.item() == 0.0:
                    break

            # prepare model inputs
            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
            model_kwargs["attention_mask"] = model_inputs["attention_mask"]
            # forward pass to get next token
            outputs = self(
                **model_inputs,
                return_dict=True,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
            )

            if synced_gpus and this_peer_finished:
                continue  # don't waste resources running the code we don't need

            next_token_logits = outputs.logits[:, -1, :]

            # pre-process distribution
            next_token_scores = logits_processor(input_ids, next_token_logits)
            next_token_scores = logits_warper(input_ids, next_token_scores)

            force_generation = model_kwargs.get('force_generation', False)
            if force_generation:
                if self.output_img_id is not None and self.output_img_id not in input_ids:
                    next_token_scores[:, self.output_img_id] = 10*next_token_scores.max()

            # Store scores, attentions and hidden_states when required
            if return_dict_in_generate:
                if output_scores:
                    scores += (next_token_scores,)
                if output_attentions:
                    decoder_attentions += (
                        (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
                    )
                    if self.config.is_encoder_decoder:
                        cross_attentions += (outputs.cross_attentions,)

                if output_hidden_states:
                    decoder_hidden_states += (
                        (outputs.decoder_hidden_states,)
                        if self.config.is_encoder_decoder
                        else (outputs.hidden_states,)
                    )

            # sample
            probs = nn.functional.softmax(next_token_scores, dim=-1)
            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)

            # finished sentences should have their next token be a padding token
            if eos_token_id is not None:
                if pad_token_id is None:
                    raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.")
                next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)

            # update generated ids, model inputs, and length for next step
            input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
            if streamer is not None:
                streamer.put(next_tokens.cpu())
            model_kwargs = self._update_model_kwargs_for_generation(
                outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
            )

            # if eos_token was found in one sentence, set sentence to finished
            if eos_token_id_tensor is not None:
                unfinished_sequences = unfinished_sequences.mul(
                    next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
                )

                # stop when each sentence is finished
                if unfinished_sequences.max() == 0:
                    this_peer_finished = True

            # stop if we exceed the maximum length
            if stopping_criteria(input_ids, scores):
                this_peer_finished = True

            if this_peer_finished and not synced_gpus:
                break

        if streamer is not None:
            streamer.end()

        if return_dict_in_generate:
            if self.config.is_encoder_decoder:
                return SampleEncoderDecoderOutput(
                    sequences=input_ids,
                    scores=scores,
                    encoder_attentions=encoder_attentions,
                    encoder_hidden_states=encoder_hidden_states,
                    decoder_attentions=decoder_attentions,
                    cross_attentions=cross_attentions,
                    decoder_hidden_states=decoder_hidden_states,
                )
            else:
                return SampleDecoderOnlyOutput(
                    sequences=input_ids,
                    scores=scores,
                    attentions=decoder_attentions,
                    hidden_states=decoder_hidden_states,
                )
        else:
            return input_ids
minigpt4/models/generation_helper.py
eric-ai-lab/MiniGPT-5
import logging
import random

import torch
from torch.cuda.amp import autocast as autocast
import torch.nn as nn

from minigpt4.common.registry import registry
from minigpt4.models.blip2 import Blip2Base, disabled_train
from minigpt4.models.modeling_llama import LlamaForCausalLM
from transformers import LlamaTokenizer
import transformers
assert transformers.__version__ >= '4.31.0', "Please update transformers to >= 4.31.0"
from constants import *


@registry.register_model("mini_gpt4")
class MiniGPT4(Blip2Base):
    """
    BLIP2 GPT-LLAMA model.
    """

    PRETRAINED_MODEL_CONFIG_DICT = {
        "pretrain_vicuna": "configs/models/minigpt4.yaml",
    }

    def __init__(
        self,
        vit_model="eva_clip_g",
        q_former_model="https://storage.googleapis.com/sfr-vision-language-research/LAVIS/models/BLIP2/blip2_pretrained_flant5xxl.pth",
        img_size=224,
        drop_path_rate=0,
        use_grad_checkpoint=False,
        vit_precision="fp16",
        freeze_vit=True,
        freeze_qformer=True,
        num_query_token=32,
        llama_model="",
        prompt_path="",
        prompt_template="",
        max_txt_len=32,
        end_sym='\n',
        low_resource=False,  # use 8 bit and put vit in cpu
        device_8bit=0,  # the device of 8bit model should be set when loading and cannot be changed anymore.
        **kwargs
    ):
        super().__init__()

        self.tokenizer = self.init_tokenizer()
        self.low_resource = low_resource

        print('Loading VIT')
        self.visual_encoder, self.ln_vision = self.init_vision_encoder(
            vit_model, img_size, drop_path_rate, use_grad_checkpoint, vit_precision
        )
        if freeze_vit:
            for name, param in self.visual_encoder.named_parameters():
                param.requires_grad = False
            self.visual_encoder = self.visual_encoder.eval()
            self.visual_encoder.train = disabled_train
            for name, param in self.ln_vision.named_parameters():
                param.requires_grad = False
            self.ln_vision = self.ln_vision.eval()
            self.ln_vision.train = disabled_train
            logging.info("freeze vision encoder")
        print('Loading VIT Done')

        print('Loading Q-Former')
        self.Qformer, self.query_tokens = self.init_Qformer(
            num_query_token, self.visual_encoder.num_features
        )
        self.Qformer.cls = None
        self.Qformer.bert.embeddings.word_embeddings = None
        self.Qformer.bert.embeddings.position_embeddings = None
        for layer in self.Qformer.bert.encoder.layer:
            layer.output = None
            layer.intermediate = None
        self.load_from_pretrained(url_or_filename=q_former_model)

        if freeze_qformer:
            for name, param in self.Qformer.named_parameters():
                param.requires_grad = False
            self.Qformer = self.Qformer.eval()
            self.Qformer.train = disabled_train
            self.query_tokens.requires_grad = False
            logging.info("freeze Qformer")
        print('Loading Q-Former Done')

        print('Loading LLAMA')
        self.llama_tokenizer = LlamaTokenizer.from_pretrained(llama_model, legacy=True, use_fast=False)
        self.llama_tokenizer.pad_token = self.llama_tokenizer.eos_token

        if self.low_resource:
            self.llama_model = LlamaForCausalLM.from_pretrained(
                llama_model,
                torch_dtype=PRECISION,
                load_in_8bit=True,
                device_map={'': device_8bit}
            )
        else:
            self.llama_model = LlamaForCausalLM.from_pretrained(
                llama_model,
                torch_dtype=PRECISION,
            )

        for name, param in self.llama_model.named_parameters():
            param.requires_grad = False
        print('Loading LLAMA Done')

        self.llama_proj = nn.Linear(
            self.Qformer.config.hidden_size, self.llama_model.config.hidden_size
        )
        self.max_txt_len = max_txt_len
        self.end_sym = end_sym

        if prompt_path:
            with open(prompt_path, 'r') as f:
                raw_prompts = f.read().splitlines()
            filted_prompts = [raw_prompt for raw_prompt in raw_prompts if "<ImageHere>" in raw_prompt]
            self.prompt_list = [prompt_template.format(p) for p in filted_prompts]
            print('Load {} training prompts'.format(len(self.prompt_list)))
            print('Prompt Example \n{}'.format(random.choice(self.prompt_list)))
        else:
            self.prompt_list = []

    def vit_to_cpu(self):
        self.ln_vision.to("cpu")
        self.ln_vision.float()
        self.visual_encoder.to("cpu")
        self.visual_encoder.float()

    def encode_img(self, image):
        device = image.device
        if self.low_resource:
            self.vit_to_cpu()
            image = image.to("cpu")

        with self.maybe_autocast():
            image_embeds = self.ln_vision(self.visual_encoder(image)).to(device)
            image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)

            query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)
            query_output = self.Qformer.bert(
                query_embeds=query_tokens,
                encoder_hidden_states=image_embeds,
                encoder_attention_mask=image_atts,
                return_dict=True,
            )

            inputs_llama = self.llama_proj(query_output.last_hidden_state)
            atts_llama = torch.ones(inputs_llama.size()[:-1], dtype=torch.long).to(image.device)
        return inputs_llama, atts_llama

    def prompt_wrap(self, img_embeds, atts_img, prompt):
        if prompt:
            batch_size = img_embeds.shape[0]
            p_before, p_after = prompt.split('<ImageHere>')
            p_before_tokens = self.llama_tokenizer(
                p_before, return_tensors="pt", add_special_tokens=False).to(img_embeds.device)
            p_after_tokens = self.llama_tokenizer(
                p_after, return_tensors="pt", add_special_tokens=False).to(img_embeds.device)
            p_before_embeds = self.llama_model.model.embed_tokens(p_before_tokens.input_ids).expand(batch_size, -1, -1)
            p_after_embeds = self.llama_model.model.embed_tokens(p_after_tokens.input_ids).expand(batch_size, -1, -1)
            wrapped_img_embeds = torch.cat([p_before_embeds, img_embeds, p_after_embeds], dim=1)
            wrapped_atts_img = atts_img[:, :1].expand(-1, wrapped_img_embeds.shape[1])
            return wrapped_img_embeds, wrapped_atts_img
        else:
            return img_embeds, atts_img

    def forward(self, samples):
        image = samples["image"]
        img_embeds, atts_img = self.encode_img(image)
        if hasattr(samples, 'question_split'):  # VQA dataset
            print('VQA Batch')
            vqa_prompt = '###Human: <Img><ImageHere></Img> '
            img_embeds, atts_img = self.prompt_wrap(img_embeds, atts_img, vqa_prompt)
        elif self.prompt_list:
            prompt = random.choice(self.prompt_list)
            img_embeds, atts_img = self.prompt_wrap(img_embeds, atts_img, prompt)

        self.llama_tokenizer.padding_side = "right"

        text = [t + self.end_sym for t in samples["text_input"]]

        to_regress_tokens = self.llama_tokenizer(
            text,
            return_tensors="pt",
            padding="longest",
            truncation=True,
            max_length=self.max_txt_len,
            add_special_tokens=False
        ).to(image.device)

        targets = to_regress_tokens.input_ids.masked_fill(
            to_regress_tokens.input_ids == self.llama_tokenizer.pad_token_id, -100
        )

        empty_targets = (
            torch.ones([atts_img.shape[0], atts_img.shape[1]+1],
                       dtype=torch.long).to(image.device).fill_(-100)  # plus one for bos
        )
        targets = torch.cat([empty_targets, targets], dim=1)

        batch_size = img_embeds.shape[0]
        bos = torch.ones([batch_size, 1],
                         dtype=to_regress_tokens.input_ids.dtype,
                         device=to_regress_tokens.input_ids.device) * self.llama_tokenizer.bos_token_id
        bos_embeds = self.llama_model.model.embed_tokens(bos)
        atts_bos = atts_img[:, :1]

        to_regress_embeds = self.llama_model.model.embed_tokens(to_regress_tokens.input_ids)
        inputs_embeds = torch.cat([bos_embeds, img_embeds, to_regress_embeds], dim=1)
        attention_mask = torch.cat([atts_bos, atts_img, to_regress_tokens.attention_mask], dim=1)

        with self.maybe_autocast():
            outputs = self.llama_model(
                inputs_embeds=inputs_embeds,
                attention_mask=attention_mask,
                return_dict=True,
                labels=targets,
            )
        loss = outputs.loss

        return {"loss": loss}

    @classmethod
    def from_config(cls, cfg):
        vit_model = cfg.get("vit_model", "eva_clip_g")
        q_former_model = cfg.get("q_former_model", "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/models/BLIP2/blip2_pretrained_flant5xxl.pth")
        img_size = cfg.get("image_size")
        num_query_token = cfg.get("num_query_token")
        llama_model = cfg.get("llama_model")

        drop_path_rate = cfg.get("drop_path_rate", 0)
        use_grad_checkpoint = cfg.get("use_grad_checkpoint", False)
        vit_precision = cfg.get("vit_precision", "fp16")
        freeze_vit = cfg.get("freeze_vit", True)
        freeze_qformer = cfg.get("freeze_qformer", True)
        low_resource = cfg.get("low_resource", False)
        device_8bit = cfg.get("device_8bit", 0)

        prompt_path = cfg.get("prompt_path", "")
        prompt_template = cfg.get("prompt_template", "")
        max_txt_len = cfg.get("max_txt_len", 32)
        end_sym = cfg.get("end_sym", '\n')

        model = cls(
            vit_model=vit_model,
            q_former_model=q_former_model,
            img_size=img_size,
            drop_path_rate=drop_path_rate,
            use_grad_checkpoint=use_grad_checkpoint,
            vit_precision=vit_precision,
            freeze_vit=freeze_vit,
            freeze_qformer=freeze_qformer,
            num_query_token=num_query_token,
            llama_model=llama_model,
            prompt_path=prompt_path,
            prompt_template=prompt_template,
            max_txt_len=max_txt_len,
            end_sym=end_sym,
            low_resource=low_resource,
            device_8bit=device_8bit,
        )

        ckpt_path = cfg.get("ckpt", "")  # load weights of MiniGPT-4
        if ckpt_path:
            print("Load BLIP2-LLM Checkpoint: {}".format(ckpt_path))
            ckpt = torch.load(ckpt_path, map_location="cpu")
            msg = model.load_state_dict(ckpt['model'], strict=False)

        return model

minigpt4/models/mini_gpt4.py
eric-ai-lab/MiniGPT-5
import torch.nn as nn
import torch

from minigpt4.common.registry import registry
from minigpt4.models.mini_gpt4 import MiniGPT4
from transformers import StoppingCriteria, StoppingCriteriaList
from constants import *

from peft import (
    LoraConfig,
    PeftType,
    get_peft_model,
    PrefixTuningConfig,
    PromptLearningConfig,
    PrefixEncoder,
    TaskType
)

import types
from .generation_helper import sample

IGNORE_INDEX = -100
DEFAULT_PAD_TOKEN = "[PAD]"
DEFAULT_EOS_TOKEN = "</s>"
DEFAULT_BOS_TOKEN = "</s>"
DEFAULT_UNK_TOKEN = "</s>"


class StoppingCriteriaSub(StoppingCriteria):

    def __init__(self, stops=[], encounters=1):
        super().__init__()
        self.stops = stops

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
        for stop in self.stops:
            if torch.all((stop == input_ids[0][-len(stop):])).item():
                return True

        return False
    
@registry.register_model("minigpt5")
class MiniGPT5(MiniGPT4):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.to(PRECISION)
        all_img_tokens = ALL_IMG_TOKENS + ["<ImageHere>"]
        self.num_new_tokens = self.llama_tokenizer.add_special_tokens(
            {   
                # "pad_token": DEFAULT_PAD_TOKEN,
                "eos_token": DEFAULT_EOS_TOKEN,
                "bos_token": DEFAULT_BOS_TOKEN,
                # "unk_token": DEFAULT_UNK_TOKEN,
                "additional_special_tokens": all_img_tokens
            }
        )
        if self.num_new_tokens > 0:
            self.llama_model.resize_token_embeddings(len(self.llama_tokenizer))
            input_embeddings = self.llama_model.get_input_embeddings().weight.data
            output_embeddings = self.llama_model.get_output_embeddings().weight.data

            input_embeddings_avg = input_embeddings[:-self.num_new_tokens].mean(dim=0, keepdim=True)
            output_embeddings_avg = output_embeddings[:-self.num_new_tokens].mean(dim=0, keepdim=True)

            input_embeddings[-self.num_new_tokens:] = input_embeddings_avg
            output_embeddings[-self.num_new_tokens:] = output_embeddings_avg
        
        if len(ALL_IMG_TOKENS):
            self.output_img_id = self.llama_tokenizer.convert_tokens_to_ids(ALL_IMG_TOKENS[0])
        self.input_img_id = self.llama_tokenizer.convert_tokens_to_ids('<ImageHere>')

        self.llama_model.get_input_embeddings().to(TRAINABLE_PRECISION)
        self.llama_model.get_output_embeddings().to(TRAINABLE_PRECISION)

        # original_input_embeddings = self.llama_model.get_input_embeddings().weight.data.clone()
        # original_output_embeddings = self.llama_model.get_output_embeddings().weight.data.clone()
        # self.register_buffer("original_input_embeddings", original_input_embeddings, persistent=False)
        # self.register_buffer("original_output_embeddings", original_output_embeddings, persistent=False)
        input_embed_grad_mask = torch.ones_like(self.llama_model.get_input_embeddings().weight.data)
        output_embed_grad_mask = torch.ones_like(self.llama_model.get_output_embeddings().weight.data)
        input_embed_grad_mask[:-self.num_new_tokens] = 0
        output_embed_grad_mask[:-self.num_new_tokens] = 0
        self.register_buffer("input_embed_grad_mask", input_embed_grad_mask, persistent=False)
        self.register_buffer("output_embed_grad_mask", output_embed_grad_mask, persistent=False)

        self.base_model_prepare_inputs_for_generation = self.llama_model.prepare_inputs_for_generation
        self.base_model_torch_dtype = self.llama_model.dtype
        self.llama_model.sample = types.MethodType(sample, self.llama_model)
        self.llama_model.output_img_id = self.output_img_id

        self.llama_model.prepare_inputs_for_generation = self.prepare_inputs_for_generation     

        self.using_lora = USE_LORA
        self.using_prefix_tuning = USE_PREFIX_TUNING
        if self.using_lora:
            print("Using LoRA")
            lora_target_modules = [
                "q_proj",
                "v_proj",
            ]
            lora_r = 8
            lora_alpha = 16
            lora_dropout = 0.05
            self.lora_config = LoraConfig(
                r=lora_r,
                lora_alpha=lora_alpha,
                target_modules=lora_target_modules,
                lora_dropout=lora_dropout,
                bias="none",
                task_type=TaskType.CAUSAL_LM,
                modules_to_save=['lm_head','embed_tokens']
            )
            self.llama_model = get_peft_model(self.llama_model, self.lora_config)
            self.llama_model.base_model.model.model.embed_tokens.original_module.weight.requires_grad = False
            self.llama_model.base_model.model.lm_head.original_module.weight.requires_grad = False
        else:
            for name, param in self.llama_model.named_parameters():
                if "lm_head" in name or "embed_tokens" in name:
                    param.requires_grad = True
                else:
                    param.requires_grad = False

        llama_config = self.llama_model.config
        if self.using_prefix_tuning:
            num_virtual_tokens = 5
            self.peft_config = PrefixTuningConfig(peft_type="P_TUNING",
                                              task_type=TaskType.CAUSAL_LM, 
                                              num_virtual_tokens=num_virtual_tokens,
                                              num_layers=llama_config.num_hidden_layers,
                                              num_attention_heads = llama_config.num_attention_heads,
                                              token_dim=llama_config.hidden_size,
                                              encoder_hidden_size=1024,
                                              prefix_projection=True)
            self.prefix_encoder = PrefixEncoder(self.peft_config).to(TRAINABLE_PRECISION)
            self.prefix_tokens = torch.arange(num_virtual_tokens).long()

        for name, param in self.llama_proj.named_parameters():
            param.requires_grad = False

    def get_prompt(self, batch_size, device, inference_mode=False):
        prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device)
        past_key_values = self.prefix_encoder(prefix_tokens)
        past_key_values = past_key_values.view(
            batch_size,
            self.peft_config.num_virtual_tokens,
            self.peft_config.num_layers * 2,
            self.peft_config.num_attention_heads,
            self.peft_config.token_dim//self.peft_config.num_attention_heads
                )
        past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2)
        return past_key_values
    
    def encode_img(self, image):
        device = image.device
        if self.low_resource:
            self.vit_to_cpu()
            image = image.to("cpu")
        with torch.no_grad():
            with torch.autocast('cuda'):
                image_embeds = self.ln_vision(self.visual_encoder(image)).to(device)
                image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)

                query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)
                query_output = self.Qformer.bert(
                    query_embeds=query_tokens,
                    encoder_hidden_states=image_embeds,
                    encoder_attention_mask=image_atts,
                    return_dict=True,
                )

                inputs_llama = self.llama_proj(query_output.last_hidden_state)
                atts_llama = torch.ones(inputs_llama.size()[:-1], dtype=torch.long).to(image.device)
        return inputs_llama, atts_llama
    
    def input_warp(self, input_ids, attention_mask, labels=None, input_image=None, input_image_feature=None):
        assert input_ids.shape[0] == 1, "warping each sample individually"

        bos = torch.ones([1, 1],
                        dtype=input_ids.dtype,
                        device=input_ids.device) * self.llama_tokenizer.bos_token_id
        if self.using_lora:
            bos_embeds = self.llama_model.base_model.model.model.embed_tokens(bos)
        else:
            bos_embeds = self.llama_model.model.embed_tokens(bos)
        atts_bos = torch.ones([1, 1],dtype=attention_mask.dtype,device=attention_mask.device)
        if labels is not None:
            labels_bos = torch.ones([1, 1],dtype=labels.dtype,device=labels.device) * -100
            wrapped_labels = labels_bos
        else:
            wrapped_labels = None
        
        wrapped_img_embeds, wrapped_atts_img = bos_embeds, atts_bos
        input_img_idx = (input_ids == self.input_img_id).nonzero(as_tuple=True)
        start_idx = 0
        if len(input_img_idx[0]) > 0:
            assert input_image is not None or input_image_feature is not None, 'input_image or input_image_feature should be provided'

            if input_image_feature is not None:
                img_embeds = input_image_feature
                atts_img = torch.ones(img_embeds.size()[:-1], dtype=torch.long).to(img_embeds.device)
            else:
                img_embeds, atts_img = self.encode_img(input_image)
            
            if labels is not None:
                img_label = torch.ones_like(atts_img, dtype=torch.long).to(img_embeds.device) * -100

            for i in range(len(input_img_idx[1])):
                p_before = input_ids[:, start_idx:input_img_idx[1][i]]
                p_before_attention_mask = attention_mask[:, start_idx:input_img_idx[1][i]]
                p_before_embeds = self.get_input_embeddings(p_before)
                wrapped_img_embeds = torch.cat([wrapped_img_embeds, p_before_embeds, img_embeds[i:i+1]], dim=1)
                wrapped_atts_img = torch.cat([wrapped_atts_img, p_before_attention_mask, atts_img[i:i+1]], dim=1)
                if labels is not None:
                    p_before_labels = labels[:, start_idx:input_img_idx[1][i]]
                    wrapped_labels = torch.cat([wrapped_labels, p_before_labels, img_label[i:i+1]], dim=1)
                start_idx = input_img_idx[1][i] + 1
            
        p_before = input_ids[:, start_idx:]
        p_before_attention_mask = attention_mask[:, start_idx:]
        p_before_embeds = self.get_input_embeddings(p_before)
        wrapped_img_embeds = torch.cat([wrapped_img_embeds, p_before_embeds], dim=1)
        wrapped_atts_img = torch.cat([wrapped_atts_img, p_before_attention_mask], dim=1)
        if labels is not None:
            p_before_labels = labels[:, start_idx:]
            wrapped_labels = torch.cat([wrapped_labels, p_before_labels], dim=1)
        return wrapped_img_embeds, wrapped_atts_img, wrapped_labels
            
    def forward(self, input_ids, labels, attention_mask, input_images=None, input_img_features=None, output_hidden_states=True):
        batch_size = input_ids.shape[0]
        all_input_embeds, all_attention, all_labels = [], [], []
        for b in range(batch_size):
            if input_img_features is not None:
                wrapped_img_embeds, wrapped_atts_img, wrapped_labels = self.input_warp(input_ids[b:b+1], attention_mask[b:b+1], labels[b:b+1], input_image_feature=input_img_features[b])
            elif input_images is not None:
                wrapped_img_embeds, wrapped_atts_img, wrapped_labels = self.input_warp(input_ids[b:b+1], attention_mask[b:b+1], labels[b:b+1], input_images[b])

            all_input_embeds.append(wrapped_img_embeds)
            all_attention.append(wrapped_atts_img)
            all_labels.append(wrapped_labels)

        #add padding features for batch
        max_len = max([x.shape[1] for x in all_input_embeds])
        for i in range(len(all_input_embeds)):
            if all_input_embeds[i].shape[1] < max_len:
                pad_len = max_len - all_input_embeds[i].shape[1]
                pad_embeds = torch.zeros([all_input_embeds[i].shape[0], pad_len, all_input_embeds[i].shape[2]]).to(all_input_embeds[i].device)
                pad_atts = torch.zeros([all_attention[i].shape[0], pad_len]).to(all_attention[i].device)
                pad_labels = torch.ones([all_labels[i].shape[0], pad_len], dtype=torch.long).to(all_labels[i].device) * -100
                all_input_embeds[i] = torch.cat([all_input_embeds[i], pad_embeds], dim=1)
                all_attention[i] = torch.cat([all_attention[i], pad_atts], dim=1)
                all_labels[i] = torch.cat([all_labels[i], pad_labels], dim=1)
        
        all_input_embeds = torch.cat(all_input_embeds, dim=0)
        all_attention = torch.cat(all_attention, dim=0)
        all_labels = torch.cat(all_labels, dim=0)

        past_key_values = None
        if self.using_prefix_tuning:
            device = all_input_embeds.device
            past_key_values = self.get_prompt(batch_size=batch_size, device=device)
            prefix_attention_mask = torch.ones(batch_size, self.peft_config.num_virtual_tokens).to(device)
            all_attention = torch.cat([prefix_attention_mask, all_attention], dim=1)
            # prefix_labels = torch.ones(batch_size, self.peft_config.num_virtual_tokens, dtype=wrapped_labels.dtype).to(device) * -100
            # wrapped_labels = torch.cat([prefix_labels, wrapped_labels], dim=1)

        outputs = self.llama_model(
                inputs_embeds=all_input_embeds,
                attention_mask=all_attention,
                return_dict=True,
                labels=all_labels,
                output_hidden_states=output_hidden_states,
                past_key_values=past_key_values,
            )
        output_token_index = (all_labels == self.output_img_id).nonzero()

        if len(output_token_index):
            addon_index = torch.ones_like(output_token_index)*(-1)
            addon_index[:, 0] = 0
            output_token_index += addon_index
        
        return outputs, output_token_index
    
    def predict(self, instruction, input_image, max_new_tokens=300, num_beams=1, min_length=1, top_p=0.9,
               repetition_penalty=1.0, length_penalty=1, temperature=1.0, task_name=None, output_hidden_states=True, force_generation=False):

        sample_inputs = self.llama_tokenizer(instruction, return_tensors="pt", add_special_tokens = False).to(self.device)
        input_ids = sample_inputs.input_ids
        attention_mask = sample_inputs.attention_mask

        wrapped_img_embeds, wrapped_atts_img, _ = self.input_warp(input_ids, attention_mask, input_image=input_image)

        stop_words_ids = [torch.tensor([835]).to(self.device),
                          torch.tensor([2277, 29937]).to(self.device)]  # '###' can be encoded in two different ways.
        self.stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])

        # if self.using_prefix_tuning:

        sample_outputs = self.llama_model.generate(
            inputs_embeds=wrapped_img_embeds,
            attention_mask=wrapped_atts_img,
            max_new_tokens=max_new_tokens,
            stopping_criteria=self.stopping_criteria,
            num_beams=num_beams,
            do_sample=True,
            min_length=min_length,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            length_penalty=length_penalty,
            temperature=temperature,
            output_hidden_states=output_hidden_states,
            return_dict_in_generate=True,
            force_generation=force_generation
        )

        return sample_outputs

    def prepare_inputs_for_generation(self, input_ids, inputs_embeds=None, force_generation=None, *args, **kwargs):
        model_kwargs = self.base_model_prepare_inputs_for_generation(input_ids, inputs_embeds=inputs_embeds, *args, **kwargs)
        bs = inputs_embeds.shape[0]
        device = inputs_embeds.device
        if 'input_ids' in model_kwargs:
            new_token_ids = model_kwargs['input_ids'][:, -1:]
            if new_token_ids == self.output_img_id:
                #Generated the image token, force add all the image tokens
                current_position_ids = model_kwargs['position_ids'][0, -1] #TODO: Only support batch size 1
                all_img_tokens = torch.tensor(self.llama_tokenizer.convert_tokens_to_ids(ALL_IMG_TOKENS)).unsqueeze(0).to(device)
                all_img_tokens_mask = torch.ones_like(all_img_tokens)[:, :IMG_TOKEN_NUM-1].to(device)
                all_img_position_ids = torch.arange(current_position_ids, current_position_ids + IMG_TOKEN_NUM).unsqueeze(0).to(device)
                
                model_kwargs['attention_mask'] = torch.cat([model_kwargs['attention_mask'], all_img_tokens_mask], dim=1)
                model_kwargs['position_ids'] = all_img_position_ids
                inputs_embeds = self.get_input_embeddings(all_img_tokens)
                model_kwargs['input_ids'] = None
                model_kwargs['inputs_embeds'] = inputs_embeds
                # input_ids = torch.cat([input_ids[:, :-1], all_img_tokens], dim=1)

        if self.using_prefix_tuning:
            if isinstance(self.peft_config, PromptLearningConfig):
                if self.peft_config.peft_type == PeftType.PREFIX_TUNING and model_kwargs["past_key_values"] is None:
                    prefix_attention_mask = torch.ones(bs, self.peft_config.num_virtual_tokens).to(device)
                    model_kwargs["attention_mask"] = torch.cat(
                        (prefix_attention_mask, model_kwargs["attention_mask"]), dim=1
                    )

                if model_kwargs["past_key_values"] is None and self.peft_config.peft_type == PeftType.PREFIX_TUNING:
                    past_key_values = self.get_prompt(batch_size=bs, device=device)

                    if self.base_model_torch_dtype is not None:
                        # handle the case for Bloom where it outputs tuple of tuples
                        if isinstance(past_key_values[0], tuple):
                            past_key_values = tuple(
                                tuple(
                                    past_key_value.to(self.base_model_torch_dtype)
                                    for past_key_value in past_key_value_tuple
                                )
                                for past_key_value_tuple in past_key_values
                            )
                        else:
                            past_key_values = tuple(
                                past_key_value.to(self.base_model_torch_dtype) for past_key_value in past_key_values
                            )

                    model_kwargs["past_key_values"] = past_key_values
                else:
                    if model_kwargs["past_key_values"] is None:
                        inputs_embeds = self.word_embeddings(model_kwargs["input_ids"])
                        prompts = self.get_prompt(batch_size=model_kwargs["input_ids"].shape[0])
                        prompts = prompts.to(inputs_embeds.dtype)
                        model_kwargs["inputs_embeds"] = torch.cat((prompts, inputs_embeds), dim=1)
                        model_kwargs["input_ids"] = None

        return model_kwargs
    
    def get_input_embeddings(self, input_ids):
        if self.using_lora:
            embed_tokens = self.llama_model.base_model.model.model.embed_tokens
        else:
            embed_tokens = self.llama_model.model.embed_tokens
        inputs_embeds = embed_tokens(input_ids)

        return inputs_embeds
    
    def reset_embeddings(self):
        with torch.no_grad():
            if self.using_lora:
                for n, p in self.llama_model.named_parameters():
                    if p.grad is None:
                        continue
                    if "lm_head" in n:
                        p.grad = p.grad*self.output_embed_grad_mask
                    elif "embed_tokens" in n:
                        p.grad = p.grad*self.input_embed_grad_mask
            else:
                self.llama_model.get_input_embeddings().weight.grad = self.llama_model.get_input_embeddings().weight.grad*self.input_embed_grad_mask
                self.llama_model.get_output_embeddings().weight.grad = self.llama_model.get_output_embeddings().weight.grad*self.output_embed_grad_mask

minigpt4/models/mini_gpt5.py
eric-ai-lab/MiniGPT-5
import os
import random
from typing import Any, Optional, Dict, List

import torch
from lightning.pytorch import LightningModule
from transformers import get_linear_schedule_with_warmup, CLIPTextModel, CLIPTokenizer, PreTrainedTokenizer
from torch.optim import AdamW
import torch.nn as nn

from minigpt4.models.mini_gpt5 import MiniGPT5
from minigpt4.common.config import Config

from diffusers import AutoencoderKL, UNet2DConditionModel
import wandb
import torch.nn.functional as F
from utils import plot_images_and_text
from constants import *
from diffusers import StableDiffusionPipeline
from diffusers.models.vae import DiagonalGaussianDistribution

class MiniGPT5_InputProcessor(object):

    def __init__(self, tokenizer: PreTrainedTokenizer, image_processor: Any):
        self.tokenizer = tokenizer
        self.image_processor = image_processor
        
    def __call__(self, text = None, images = None, **kwargs) -> Any:
        output_dict = {}
        if text is not None:
            text_output = self.tokenizer(text, 
                                  return_tensors="pt", 
                                  padding=True,
                                  max_length=self.tokenizer.model_max_length,
                                  truncation=True,
                                  **kwargs)
            output_dict.update(text_output)

        if images is not None:
            all_images = []
            if isinstance(images, list):
                for img in images:
                    image_output = self.image_processor(img)
                    all_images.append(image_output)
                input_images = torch.stack(all_images, dim=0)
            else:
                input_images = self.image_processor(images)
            output_dict['input_images'] = input_images
        return output_dict

class MiniGPT4Args:
    cfg_path = "config/minigpt4.yaml"
    options = []

# define the LightningModule
class MiniGPT5_Model(LightningModule):
    def __init__(self,
        encoder_model_config,
        **kwargs,
    ):
        super().__init__()
        self.save_hyperparameters(ignore=['encoder_model_config'])
        self.encoder_model_config = encoder_model_config
        self.input_vis_processor = None

        if encoder_model_config.model_type == 'multimodal_encoder':
            minigpt4_config = Config(MiniGPT4Args)
            self.model = MiniGPT5.from_config(minigpt4_config.model_cfg)
            self.tokenizer = self.model.llama_tokenizer

            hidden_size = self.model.llama_model.config.hidden_size

        sd_model_name = "stabilityai/stable-diffusion-2-1-base"

        self.sd_text_encoder = CLIPTextModel.from_pretrained(sd_model_name, subfolder="text_encoder")
        self.sd_tokenizer = CLIPTokenizer.from_pretrained(sd_model_name, subfolder="tokenizer")
        self.vae = AutoencoderKL.from_pretrained(sd_model_name, subfolder="vae").to(PRECISION)
        
        self.unet = UNet2DConditionModel.from_pretrained(sd_model_name, subfolder="unet").to(PRECISION)
        # Freeze vae and text_encoder
        self.vae.requires_grad_(False)
        self.sd_text_encoder.requires_grad_(False)
        self.unet.requires_grad_(False)

        sd_hidden_size = self.sd_text_encoder.config.hidden_size
        self.t2i_decoder_prompt = torch.nn.Parameter(torch.randn((1,77, sd_hidden_size), dtype=TRAINABLE_PRECISION))
        self.llm_to_t2i_mapping = nn.Transformer(batch_first=True, norm_first=True, d_model = sd_hidden_size, num_encoder_layers=4, num_decoder_layers=4, dim_feedforward=sd_hidden_size*4, dropout=0.0, dtype=TRAINABLE_PRECISION)
        
        if len(ALL_IMG_TOKENS):
            self.output_img_id = self.tokenizer.convert_tokens_to_ids(ALL_IMG_TOKENS[0])
        self.img_token_num = IMG_TOKEN_NUM

        self.image_pipeline = StableDiffusionPipeline.from_pretrained(
            sd_model_name,
            vae = self.vae,
            unet = self.unet,
            safety_checker = None,
        )

        self.noise_scheduler = self.image_pipeline.scheduler

        self.fc = nn.Sequential(
                    nn.Linear(hidden_size, sd_hidden_size),
                    nn.GELU(),
                    nn.Linear(sd_hidden_size, sd_hidden_size),
                ).to(TRAINABLE_PRECISION)

        empty_text_feature = self.encode_caption('', self.sd_tokenizer.model_max_length, inference=True)
        self.register_buffer('empty_text_feature', empty_text_feature, persistent=False)

        zero_img_feature = torch.zeros((1, self.img_token_num, hidden_size), dtype=TRAINABLE_PRECISION)
        self.register_buffer('zero_img_feature', zero_img_feature, persistent=False)

        self.sd_text_encoder.to(PRECISION)

        if IS_STAGE2:
            for n, p in self.fc.named_parameters():
                p.requires_grad = False
            # for n, p in self.llm_to_t2i_mapping.named_parameters():
            #     p.requires_grad = False
            self.t2i_decoder_prompt.requires_grad = False

    def training_step(self, batch, batch_idx):
        for key in batch.keys():
            if type(batch[key]) == list:
                batch[key] = batch[key]
            else:
                batch[key] = batch[key].to(self.device)

        input_ids = batch['input_ids']
        labels = batch['labels']
        attention_mask = batch['attention_mask']
        source_text = batch['source']
        target_text = batch['target']
        captions = batch['caption']
        input_images = batch.get('input_images', None)
        output_image = batch.get('output_image', None)
        input_images_feature = batch.get('input_images_feature', None)
        output_image_feature = batch.get('output_image_feature', None)

        bs = len(source_text)
        loss_dict = self(input_ids, attention_mask, input_images, output_image, labels, captions, input_images_feature, output_image_feature)
        loss = loss_dict['loss']
        log_dict = {f'train_{k}': v for k, v in loss_dict.items()}
        self.log_dict(log_dict, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True, batch_size=bs)
        # check image generation for every 1000 steps
        if (self.global_step+1) % 500 == 0 and self.global_rank == 0:
            with torch.no_grad():
                self.eval()
                # utterance = "generate image with caption: a man on the sofa."
                utterance = source_text[0]
                gt_text = target_text[0]
                i_image = None
                if "<ImageHere>" in utterance:
                    i_image = input_images[0]
                text_out, image_out = self.generate(utterance, i_image)
                if image_out is not None:
                    if os.path.exists("train_eval") == False:
                        os.makedirs("train_eval")
                    if IS_STAGE2:
                        data = [[self.global_step, utterance, text_out, wandb.Image(image_out), gt_text]]
                        columns = ["step", "input_utterance", "text_out", "img_out", "gt_text"]
                    else:
                        predicted_images_nl = self.image_pipeline(prompt= captions[0], use_original=True).images[0]
                        data = [[self.global_step, utterance, text_out, wandb.Image(image_out), captions[0], wandb.Image(predicted_images_nl)]]
                        columns = ["step", "input_utterance", "text_out", "img_out", "caption", "caption_out"]
                        predicted_images_nl.save(os.path.join("train_eval", f'{self.global_step}_nl.png'))
                    self.logger.log_table(key="sample", data=data, columns=columns)
                    image_out.save(os.path.join("train_eval", f'{self.global_step}.png'))
                else:
                    data = [[self.global_step, utterance, text_out, None, gt_text]]
                    columns = ["step", "input_utterance", "text_out", "img_out", "gt_text"]
                    self.logger.log_table(key="sample", data=data, columns=columns)
                self.train()
        return loss
    
    def on_before_optimizer_step(self, optimizer) -> None:
        self.model.reset_embeddings()

    def validation_step(self, batch, batch_idx):
        for key in batch.keys():
            if type(batch[key]) == list:
                batch[key] = batch[key]
            else:
                batch[key] = batch[key].to(self.device)

        input_ids = batch['input_ids']
        labels = batch['labels']
        attention_mask = batch['attention_mask']
        source_text = batch['source']
        target_text = batch['target']
        captions = batch['caption']
        input_images = batch.get('input_images', None)
        output_image = batch.get('output_image', None)
        input_images_feature = batch.get('input_images_feature', None)
        output_image_feature = batch.get('output_image_feature', None)

        bs = len(source_text)
        loss_dict = self(input_ids, attention_mask, input_images, output_image, labels, captions, input_images_feature, output_image_feature)
        log_dict = {f'val_{k}': v for k, v in loss_dict.items()}
        self.log_dict(log_dict, batch_size=bs, logger=True, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)

    def configure_optimizers(self):
        """Prepare optimizer and schedule (linear warmup and decay)"""
        # no_decay = ["bias", "LayerNorm.weight"]
        if IS_STAGE2:
            optimizer_grouped_parameters = [
                {
                    "params": [p for n, p in self.model.named_parameters() if p.requires_grad],
                },
                {
                    "params": [p for n, p in self.llm_to_t2i_mapping.named_parameters() if p.requires_grad],
                }
            ]
        else:
            optimizer_grouped_parameters = [
                {
                    "params": [p for n, p in self.model.named_parameters() if p.requires_grad],
                },
                {
                    "params": [p for n, p in self.fc.named_parameters() if p.requires_grad],
                    "lr": self.hparams.learning_rate * 10,
                },
                {
                    "params": [p for n, p in self.llm_to_t2i_mapping.named_parameters() if p.requires_grad],
                    "lr": self.hparams.learning_rate * 10,
                },
                {
                    "params": self.t2i_decoder_prompt,
                    "lr": self.hparams.learning_rate * 10,
                }
            ]
        optimizer = AdamW(optimizer_grouped_parameters, lr=self.hparams.learning_rate, eps=self.hparams.adam_epsilon)

        scheduler = get_linear_schedule_with_warmup(
            optimizer,
            num_warmup_steps=self.hparams.warmup_steps,
            num_training_steps=self.trainer.estimated_stepping_batches,
        )
        scheduler = {"scheduler": scheduler, "interval": "step", "frequency": 1}
        return [optimizer], [scheduler]
    
    def forward(self, input_ids, attention_mask, input_images, output_image, labels, captions=None, input_images_feature=None, output_image_feature=None):
        if self.encoder_model_config.model_type=='multimodal_encoder':
            outputs, special_token_index = self.model(input_ids=input_ids, labels=labels, attention_mask=attention_mask, input_images=input_images, input_img_features=input_images_feature ,output_hidden_states=True)
        text_loss = outputs['loss']
        last_hidden_state = outputs['hidden_states'][-1]
        t2i_input_embedding = []
        caption_feature = []
        calculate_caption_loss = not any([c is None for c in captions])
        for i in range(len(special_token_index)):
            bs_id, seq_id = special_token_index[i]
            # random set 10% data with empty text feature
            if USE_CFG and random.random() < 0.1:
                t2i_input_embedding.append(self.zero_img_feature)
                if calculate_caption_loss:
                    caption_feature.append(self.empty_text_feature)
            else:
                t2i_input_embedding.append(last_hidden_state[bs_id:bs_id+1, seq_id:seq_id+self.img_token_num, :])
                if calculate_caption_loss:
                    caption_feature.append(self.encode_caption(captions[bs_id], self.sd_tokenizer.model_max_length, inference=True))

        if len(t2i_input_embedding) == 0:
            loss = 0.01 * text_loss
            if calculate_caption_loss:
                return {'loss': loss, 'text_loss': text_loss, 'image_loss': 0.0, 'caption_loss': 0.0}
            else:
                return {'loss': loss, 'text_loss': text_loss, 'image_loss': 0.0}
        
        else:
            t2i_input_embedding = torch.cat(t2i_input_embedding, dim=0)
            img_token_bs = t2i_input_embedding.shape[0]
            t2i_input_embedding = self.fc(t2i_input_embedding)
            mapping_feature = self.llm_to_t2i_mapping(src=t2i_input_embedding, tgt=self.t2i_decoder_prompt.repeat(img_token_bs, 1, 1))

            if output_image_feature is None:
                image_loss = self.compute_image_loss(mapping_feature, output_image[special_token_index[:, 0]])
            else:
                image_loss = self.compute_image_loss(mapping_feature, None, output_image_feature=output_image_feature[special_token_index[:, 0]])

            if calculate_caption_loss:
                caption_feature = torch.cat(caption_feature, dim=0)
                caption_loss = F.mse_loss(mapping_feature, caption_feature)

                loss = 0.01 * text_loss + image_loss + 0.1 * caption_loss

                return {'loss': loss, 'text_loss': text_loss, 'image_loss': image_loss, 'caption_loss': caption_loss}
            else:
                loss = 0.01 * text_loss + image_loss

                return {'loss': loss, 'text_loss': text_loss, 'image_loss': image_loss}

    def generate(self, utterance, input_image=None, task_name=None, max_new_tokens=256, force_generation=False, guidance_scale=7.5) -> Any:
        self.image_pipeline.to(self.device, PRECISION)
        if input_image is None:
            input_image = torch.zeros((1, 3, 224, 224), dtype=PRECISION).to(self.device)
        if type(utterance) == str:
            utterance = [utterance]
        llm_sample_outputs = self.model.predict(utterance, input_image, max_new_tokens=max_new_tokens, temperature=1.0, repetition_penalty=2.0, task_name=task_name, force_generation=force_generation)
        new_tokens = llm_sample_outputs['sequences'][0]
        pred_out = self.tokenizer.decode(new_tokens)
        print(f'Generated text: {pred_out}')

        last_hidden_state = llm_sample_outputs['hidden_states']
        special_token_index = (new_tokens == self.output_img_id).nonzero()

        predicted_images_ft = None
        if len(special_token_index):
            idx = special_token_index[0,0]
            t2i_input_embedding = last_hidden_state[idx][-1]
            assert t2i_input_embedding.shape[1] == self.img_token_num
            img0_output_feature = last_hidden_state[idx-1][-1][:, -1:]
            t2i_input_embedding = torch.cat([img0_output_feature, t2i_input_embedding[:, :-1]], dim=1)
            t2i_input_embedding = self.fc(t2i_input_embedding)
            mapping_feature = self.llm_to_t2i_mapping(src=t2i_input_embedding, tgt=self.t2i_decoder_prompt)

            if USE_CFG:
                empty_feature = self.fc(self.zero_img_feature)
                empty_feature = self.llm_to_t2i_mapping(src=empty_feature, tgt=self.t2i_decoder_prompt)
                predicted_images_ft = self.image_pipeline(prompt_embeds = mapping_feature, negative_prompt_embeds=empty_feature, guidance_scale=guidance_scale).images[0]
            else:
                predicted_images_ft = self.image_pipeline(prompt_embeds = mapping_feature, guidance_scale=guidance_scale, use_original=True).images[0]
            
        return pred_out, predicted_images_ft
  
    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
        self.image_pipeline.to(self.device, PRECISION)

        input_images = batch['input_images'][0]
        gt_image = batch['output_image'][0]
        input_utterance = batch['source'][0]
        gt_out = batch['target'][0]
        captions = batch['caption'][0]
        task_name = batch['task_name'][0]

        predicted_images_ft = None
        predicted_images_nl = None
        current_step_prompt = None

        save_dir_cpr = self.output_folder

        if self.encoder_model_config.model_type=='multimodal_encoder':
            pred_out, predicted_images_ft = self.generate(input_utterance, input_images)
            if predicted_images_ft is not None:
                if captions is not None:
                    predicted_images_nl = self.image_pipeline(prompt = captions).images[0]
                try:
                    if "###Human" in input_utterance and "###Assistant" in input_utterance:
                        input_texts = input_utterance.split("###Human:")[1].split("###Assistant:")[0].replace("/n","")
                        if "<Img><ImageHere></Img>" in input_texts:
                            input_texts = input_texts.split("<Img><ImageHere></Img>")
                        else:
                            input_texts = [input_texts]
                    else:
                        input_texts = [input_utterance]
                    # convert input images to PIL images
                    plot_images_and_text(predicted_images_ft, predicted_images_nl, gt_image, pred_out, gt_out, save_dir_cpr, task_name, input_texts, batch['original_images'][0])
                except:
                    print("Error in saving images")

        results = [input_utterance, pred_out, gt_out, predicted_images_ft, predicted_images_nl, gt_image, current_step_prompt, task_name]
        return results
    
    def compute_snr(self,timesteps):
        """
        Computes SNR as per https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849
        """
        alphas_cumprod = self.noise_scheduler.alphas_cumprod
        sqrt_alphas_cumprod = alphas_cumprod**0.5
        sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5

        # Expand the tensors.
        # Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026
        sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[timesteps].float()
        while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):
            sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]
        alpha = sqrt_alphas_cumprod.expand(timesteps.shape)

        sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(device=timesteps.device)[timesteps].float()
        while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape):
            sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None]
        sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape)

        # Compute SNR.
        snr = (alpha / sigma) ** 2
        return snr
    
    
    def compute_image_loss(self, mapping_feature, output_image, output_image_feature=None):
        if output_image_feature is not None:
            latents = DiagonalGaussianDistribution(output_image_feature).sample()
        else:
            if len(output_image.shape) == 3:
                output_image = output_image.unsqueeze(0)

            latents = self.vae.encode(output_image).latent_dist.sample()
        latents = latents * self.vae.config.scaling_factor

        noise = torch.randn_like(latents)
        bsz = latents.shape[0]
        # Sample a random timestep for each image
        timesteps = torch.randint(0, self.noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
        timesteps = timesteps.long()

        noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)

        target = noise

        model_pred = self.unet(noisy_latents, timesteps, mapping_feature).sample

        if self.encoder_model_config.snr_loss:
            snr = self.compute_snr(timesteps)
            mse_loss_weights = (
                torch.stack([snr, 5 * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr
            )
            # We first calculate the original loss. Then we mean over the non-batch dimensions and
            # rebalance the sample-wise losses with their respective loss weights.
            # Finally, we take the mean of the rebalanced loss.
            loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
            loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights
            loss = loss.mean()
        else:
            loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
        return loss
    
    def encode_caption(self, caption, length, inference=False):
        # add_special_tokens = False
        # if len(caption) == 0:
        add_special_tokens = True
        text_inputs = self.sd_tokenizer(
                caption,
                padding="max_length",
                max_length=length,
                truncation=True,
                return_tensors="pt",
                add_special_tokens=add_special_tokens
            ).to(self.device)
        # text_inputs = {k: v.to(self.device) for k, v in text_inputs.items()}
        prompt_embeds = self.sd_text_encoder(**text_inputs)[0]
        return prompt_embeds
    
    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
        trainable_param_names = [n for n, p in self.named_parameters() if p.requires_grad]
        # remove untrainable params
        for k in list(checkpoint["state_dict"].keys()):
            if k not in trainable_param_names:
                del checkpoint["state_dict"][k]

    def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
        #use pretrained weights for unsaved params
        current_state_dict = self.state_dict()
        state_dict = checkpoint["state_dict"]
        if self.model.using_lora:
            # load lm_head and embed_tokens from pretrained model
            for name in state_dict.keys():
                if "lm_head" in name:
                    for key in current_state_dict.keys():
                        if "lm_head" in key and key != name:
                            current_state_dict[key] = state_dict[name]
                elif "embed_tokens" in name:
                    for key in current_state_dict.keys():
                        if "embed_tokens" in key and key != name:
                            current_state_dict[key] = state_dict[name]
        current_state_dict.update(state_dict)
        checkpoint["state_dict"] = current_state_dict

model.py
eric-ai-lab/MiniGPT-5
import os
from typing import Any, Optional, Dict, List
from dataclasses import dataclass, field

import torch
from lightning.pytorch import LightningModule, Trainer, seed_everything
from lightning.pytorch.callbacks import ModelCheckpoint
import transformers
from torch.utils.data import DataLoader
from typing import Optional, Dict
from dataclasses import dataclass, field
from dataloader import SupervisedDataset, DataCollator
from lightning.pytorch.loggers import WandbLogger

from lightning.pytorch import seed_everything
from torchvision import transforms
from constants import *

from lightning.pytorch.callbacks import BasePredictionWriter
from model import MiniGPT5_Model, MiniGPT5_InputProcessor
from metric import *

class PredWriter(BasePredictionWriter):
    def write_on_epoch_end(
        self,
        trainer: Trainer,
        pl_module: LightningModule,
        predictions: Any,  # complex variables is ok
        batch_indices: list[list[list[int]]],
    ) -> None:
        output_folder = pl_module.output_folder
        torch.save(predictions, os.path.join(output_folder, f"predictions-{trainer.local_rank}.pt"))
        print(f'rank {trainer.local_rank} predictions saved')

def default_gpus():
    return [0,1,2,3]

@dataclass
class ModelArguments:
    model_type: Optional[str] = field(default="multimodal_encoder")
    snr_loss: Optional[bool] = field(default=True)
    model_save_name: Optional[str] = field(default="model_{epoch}-{step}")
    stage1_weight: Optional[str] = field(default=None)

@dataclass
class DataArguments:
    train_data_path: str = field(default=None, metadata={"help": "Path to the training data."})
    val_data_path: str = field(default=None, metadata={"help": "Path to the validation data."})
    test_data_path: str = field(default=None, metadata={"help": "Path to the test data."})

@dataclass
class TrainingArguments:
    cache_dir: Optional[str] = field(default=None)
    output_dir: str = field(default=WEIGHTFOLDER)
    num_train_epochs:int = field(default=2)
    per_device_train_batch_size:int = field(default=2)
    per_device_eval_batch_size:int = field(default=2)
    real_batch_size:int = field(default=48)
    save_total_limit:int = field(default=1)
    learning_rate:float = field(default=2e-5)
    warmup_ratio:float = field(default=0.03)
    warmup_steps:int = field(default=1000)
    adam_epsilon:float = field(default=1e-8)

    num_workers:int = field(default=16)

    gpus: List[int] = field(default_factory=default_gpus)
    resume: Optional[str] = field(default=None)
    is_training: Optional[bool] = field(default=False)
    test_weight: Optional[str] = field(default=None)
        
def make_supervised_data_module(data_args, training_args, data_collator, input_processor=None, output_vis_processor=None) -> Dict:
    """Make dataset and collator for supervised fine-tuning."""

    train_dataset = SupervisedDataset(data_path=data_args.train_data_path, input_processor=input_processor, output_vis_processor=output_vis_processor)
    eval_dataset = SupervisedDataset(data_path=data_args.val_data_path, input_processor=input_processor, output_vis_processor=output_vis_processor)

    train_dataloader = DataLoader(train_dataset, 
                                  batch_size=training_args.per_device_train_batch_size,
                                    num_workers=training_args.num_workers,
                                  collate_fn=data_collator,
                                  prefetch_factor=4,
                                  pin_memory=True)
    val_dataloader = DataLoader(eval_dataset, 
                                batch_size=training_args.per_device_eval_batch_size,
                                num_workers=training_args.num_workers,
                                collate_fn=data_collator,
                                prefetch_factor=4,
                                pin_memory=True)
    return train_dataloader, val_dataloader


def make_eval_data_module(data_args, training_args, data_collator, input_processor=None, output_vis_processor=None) -> Dict:
    eval_dataset = SupervisedDataset(data_path=data_args.test_data_path, input_processor=input_processor, output_vis_processor=output_vis_processor, test=True)
    val_dataloader = DataLoader(eval_dataset, batch_size=1, collate_fn=data_collator)
    return val_dataloader

if __name__ == "__main__":
    seed_everything(42)
    torch.backends.cuda.matmul.allow_tf32 = True
    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    if isinstance(training_args.gpus, str):
        training_args.gpus = [int(x) for x in training_args.gpus.split(',')]
    if isinstance(data_args.train_data_path, str):
        data_args.train_data_path = os.path.join(DATAFOLDER, data_args.train_data_path)
    if isinstance(data_args.val_data_path, str):
        data_args.val_data_path = os.path.join(DATAFOLDER, data_args.val_data_path)
    if isinstance(data_args.test_data_path, str):
        data_args.test_data_path = os.path.join(DATAFOLDER, data_args.test_data_path)

    batch_size = training_args.real_batch_size
    devices = training_args.gpus
    num_devices = len(devices)
    gradient_accumulation_steps = max(1,batch_size // (training_args.per_device_train_batch_size*num_devices))

    if IS_STAGE2:
        stage1_weight = model_args.stage1_weight
        assert stage1_weight is not None, "stage2 weight needs stage1 weight, but stage1 weight path is None"
        stage1_weight = os.path.join(WEIGHTFOLDER, stage1_weight)
        model = MiniGPT5_Model.load_from_checkpoint(stage1_weight, strict=False, map_location="cpu", encoder_model_config=model_args, **vars(training_args))
    else:
        model = MiniGPT5_Model(encoder_model_config=model_args, **vars(training_args))
    
    tokenizer = model.tokenizer
    sd_tokenizer = model.sd_tokenizer

    if training_args.is_training:

        output_vis_processor = transforms.Compose(
                [
                    transforms.Resize(512, interpolation=transforms.InterpolationMode.BILINEAR),
                    transforms.CenterCrop(512),
                    transforms.RandomHorizontalFlip(),
                    transforms.ToTensor(),
                    transforms.Normalize([0.5], [0.5]),
                ]
            )
        input_vis_processor = transforms.Compose(
                [
                    transforms.Resize(224, interpolation=transforms.InterpolationMode.BILINEAR),
                    transforms.CenterCrop(224),
                    transforms.RandomHorizontalFlip(),
                    transforms.ToTensor(),
                    transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
                ]
            )
        input_processor = MiniGPT5_InputProcessor(tokenizer=tokenizer, image_processor=input_vis_processor)

        data_collator = DataCollator(tokenizer=tokenizer, sd_tokenizer=sd_tokenizer)

        train_dataloader, val_dataloader = make_supervised_data_module(data_args, training_args, data_collator, input_processor, output_vis_processor)
        
        checkpoint_callback = ModelCheckpoint(
                dirpath=training_args.output_dir,
                filename=model_args.model_save_name,
                monitor="val_loss",
                save_top_k=1,
                # save_last=True,
            )
        
        strategy = 'ddp'
        if "CC3M" in DATAFOLDER:
            val_check_interval = 0.25
        else:
            val_check_interval = 0.5
            strategy = 'ddp_find_unused_parameters_true'

        wandb_logger = WandbLogger(save_dir=training_args.output_dir, project="MiniGPT5_Model", offline=False, name=model_args.model_save_name)
        trainer = Trainer(default_root_dir=training_args.output_dir, max_epochs=training_args.num_train_epochs, 
                        accumulate_grad_batches=gradient_accumulation_steps,
                        accelerator="gpu", devices=devices, 
                        strategy=strategy,
                        logger = wandb_logger, 
                        precision='bf16-mixed',
                        val_check_interval=val_check_interval,
                        callbacks=[checkpoint_callback])
        resume = training_args.resume
        trainer.fit(model, train_dataloader, val_dataloader, ckpt_path=resume)
    else:
        model.image_pipeline.enable_xformers_memory_efficient_attention()
        output_vis_processor = transforms.Compose(
                [
                    transforms.Resize(512, interpolation=transforms.InterpolationMode.BILINEAR),
                    transforms.CenterCrop(512),
                    transforms.ToTensor(),
                ]
            )
        input_vis_processor = transforms.Compose(
                [
                    transforms.Resize(224, interpolation=transforms.InterpolationMode.BILINEAR),
                    transforms.CenterCrop(224),
                    transforms.ToTensor(),
                    transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
                ]
            )

        input_processor = MiniGPT5_InputProcessor(tokenizer=tokenizer, image_processor=input_vis_processor)

        data_collator = DataCollator(tokenizer=tokenizer, sd_tokenizer=sd_tokenizer)

        pred_writer = PredWriter(write_interval="epoch")

        wandb_logger = WandbLogger(save_dir=training_args.output_dir, project="MiniGPT5_Model", offline=True, name=model_args.model_save_name)
        trainer = Trainer(default_root_dir=training_args.output_dir, max_epochs=training_args.num_train_epochs, 
                        accelerator="gpu", devices=devices, 
                        strategy='ddp', 
                        logger = wandb_logger, 
                        precision='bf16-mixed',
                        callbacks=[pred_writer])
        
        val_dataloader = make_eval_data_module(data_args, training_args, data_collator, input_processor, output_vis_processor)
        assert training_args.test_weight is not None, "test weight path is None"
        ckpt_path = os.path.join(WEIGHTFOLDER, training_args.test_weight)
        model.output_folder = os.path.join(OUTPUT_FOLDER, training_args.test_weight.split(".")[0]+OUTPUT_SUFFIX)
        if not os.path.exists(model.output_folder):
            os.makedirs(model.output_folder)
        trainer.predict(model, val_dataloader, return_predictions=False, ckpt_path=ckpt_path)

train_eval.py
eric-ai-lab/MiniGPT-5

import os
import numpy as np
import matplotlib.pyplot as plt
import re
import textwrap

def sanitize_filename(filename):
    return re.sub('[^0-9a-zA-Z]+', '_', filename)

def plot_images_and_text(predicted_image1, predicted_image2, groundtruth_image, generated_text, gt_text, save_dir, task_name, input_texts, input_images):
    task_path = os.path.join(save_dir, task_name)
    if not os.path.exists(task_path):
        os.makedirs(task_path)
    max_width = 50  # adjust this value based on your needs

    fig, ax = plt.subplots()
    ax.imshow(predicted_image1)
    generated_text = generated_text.replace("###", "").replace("[IMG0]", "")
    wrapped_generated_text = textwrap.fill(generated_text, max_width)
    ax.set_title(wrapped_generated_text, pad=20)
    ax.axis('off')
    plt.savefig(os.path.join(task_path, f"generated.jpg"), bbox_inches='tight')
    plt.close(fig)

    gt_text = gt_text.replace("$", "\$")
    wrapped_gt = textwrap.fill(gt_text, max_width)
    if predicted_image2 is not None:
        fig, ax = plt.subplots()
        ax.imshow(predicted_image2)
        ax.set_title(wrapped_gt, pad=20)
        ax.axis('off')
        plt.savefig(os.path.join(task_path, f"sd_baseline.jpg"), bbox_inches='tight')
        plt.close(fig)

    if groundtruth_image is not None:
        fig, ax = plt.subplots()
        groundtruth_image = groundtruth_image.float().cpu().numpy().squeeze()
        groundtruth_image = np.transpose(groundtruth_image, (1, 2, 0))
        groundtruth_image = np.uint8(groundtruth_image*255)
        ax.imshow(groundtruth_image)
        ax.set_title(wrapped_gt, pad=20)
        ax.axis('off')
        plt.savefig(os.path.join(task_path, f"gt.jpg"), bbox_inches='tight')
        plt.close(fig)

    if len(input_texts):
        max_width = 30
        length = len(input_texts)
        if length > 1:
            fig, ax = plt.subplots(1, length, figsize=(10*length, 10))
            for i in range(length):
                if i < len(input_images):
                    ax[i].imshow(input_images[i])
                    ax[i].set_title(textwrap.fill(input_texts[i], max_width), fontsize=28)
                    ax[i].axis('off')
                else:
                    ax[i].text(0.5, 0.5, textwrap.fill(input_texts[i], max_width), horizontalalignment='center', verticalalignment='center', fontsize=28)
                    ax[i].axis('off')
        else:
            fig, ax = plt.subplots()
            ax.imshow(input_images[0])
            ax.set_title(textwrap.fill(input_texts[0], max_width), fontsize=28)
            ax.axis('off')
        plt.savefig(os.path.join(task_path, f"input.jpg"), bbox_inches='tight')
        plt.close(fig)

    return None

utils.py
eric-ai-lab/MiniGPT-5