update format

app.py

@@ -64,7 +64,7 @@ def image_classifier(input_img, input_vid, scorer_type):
 
 title_markdown = ("""
 
-<h1 align="center">Q-Align: Teaching LMMs for Visual Scoring via Discrete Text-Defined Levels</h1>
+<h3 align="center">Q-Align: Teaching LMMs for Visual Scoring via Discrete Text-Defined Levels</h3>
 
 <h3 align="center"> One Unified Model for Visual scoring. </h3>
 

q_align/.ipynb_checkpoints/utils-checkpoint.py

@@ -1,128 +0,0 @@
-import datetime
-import logging
-import logging.handlers
-import os
-import sys
-
-import requests
-
-
-
-from q_align.constants import LOGDIR
-
-server_error_msg = "**NETWORK ERROR DUE TO HIGH TRAFFIC. PLEASE REGENERATE OR REFRESH THIS PAGE.**"
-moderation_msg = "YOUR INPUT VIOLATES OUR CONTENT MODERATION GUIDELINES. PLEASE TRY AGAIN."
-
-handler = None
-
-
-def build_logger(logger_name, logger_filename):
-    global handler
-
-    formatter = logging.Formatter(
-        fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
-        datefmt="%Y-%m-%d %H:%M:%S",
-    )
-
-    # Set the format of root handlers
-    if not logging.getLogger().handlers:
-        logging.basicConfig(level=logging.INFO)
-    logging.getLogger().handlers[0].setFormatter(formatter)
-
-    # Redirect stdout and stderr to loggers
-    stdout_logger = logging.getLogger("stdout")
-    stdout_logger.setLevel(logging.INFO)
-    sl = StreamToLogger(stdout_logger, logging.INFO)
-    sys.stdout = sl
-
-    stderr_logger = logging.getLogger("stderr")
-    stderr_logger.setLevel(logging.ERROR)
-    sl = StreamToLogger(stderr_logger, logging.ERROR)
-    sys.stderr = sl
-
-    # Get logger
-    logger = logging.getLogger(logger_name)
-    logger.setLevel(logging.INFO)
-
-    # Add a file handler for all loggers
-    if handler is None:
-        os.makedirs(LOGDIR, exist_ok=True)
-        filename = os.path.join(LOGDIR, logger_filename)
-        handler = logging.handlers.TimedRotatingFileHandler(
-            filename, when='D', utc=True)
-        handler.setFormatter(formatter)
-
-        for name, item in logging.root.manager.loggerDict.items():
-            if isinstance(item, logging.Logger):
-                item.addHandler(handler)
-
-    return logger
-
-
-class StreamToLogger(object):
-    """
-    Fake file-like stream object that redirects writes to a logger instance.
-    """
-    def __init__(self, logger, log_level=logging.INFO):
-        self.terminal = sys.stdout
-        self.logger = logger
-        self.log_level = log_level
-        self.linebuf = ''
-
-    def __getattr__(self, attr):
-        return getattr(self.terminal, attr)
-
-    def write(self, buf):
-        temp_linebuf = self.linebuf + buf
-        self.linebuf = ''
-        for line in temp_linebuf.splitlines(True):
-            # From the io.TextIOWrapper docs:
-            #   On output, if newline is None, any '\n' characters written
-            #   are translated to the system default line separator.
-            # By default sys.stdout.write() expects '\n' newlines and then
-            # translates them so this is still cross platform.
-            if line[-1] == '\n':
-                self.logger.log(self.log_level, line.rstrip())
-            else:
-                self.linebuf += line
-
-    def flush(self):
-        if self.linebuf != '':
-            self.logger.log(self.log_level, self.linebuf.rstrip())
-        self.linebuf = ''
-
-
-def disable_torch_init():
-    """
-    Disable the redundant torch default initialization to accelerate model creation.
-    """
-    import torch
-    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
-    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
-
-
-def violates_moderation(text):
-    """
-    Check whether the text violates OpenAI moderation API.
-    """
-    url = "https://api.openai.com/v1/moderations"
-    headers = {"Content-Type": "application/json",
-               "Authorization": "Bearer " + os.environ["OPENAI_API_KEY"]}
-    text = text.replace("\n", "")
-    data = "{" + '"input": ' + f'"{text}"' + "}"
-    data = data.encode("utf-8")
-    try:
-        ret = requests.post(url, headers=headers, data=data, timeout=5)
-        flagged = ret.json()["results"][0]["flagged"]
-    except requests.exceptions.RequestException as e:
-        flagged = False
-    except KeyError as e:
-        flagged = False
-
-    return flagged
-
-
-def pretty_print_semaphore(semaphore):
-    if semaphore is None:
-        return "None"
-    return f"Semaphore(value={semaphore._value}, locked={semaphore.locked()})"

q_align/evaluate/.ipynb_checkpoints/iaa_eval-checkpoint.py

@@ -1,164 +0,0 @@
-import argparse
-import torch
-
-from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
-from q_align.conversation import conv_templates, SeparatorStyle
-from q_align.model.builder import load_pretrained_model
-from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
-
-from PIL import Image
-from PIL import ImageFile
-ImageFile.LOAD_TRUNCATED_IMAGES = True
-
-import requests
-from PIL import Image
-from io import BytesIO
-from transformers import TextStreamer
-
-from scipy.stats import spearmanr, pearsonr
-
-
-import json
-from tqdm import tqdm
-from collections import defaultdict
-
-import os
-
-def wa5(logits):
-    import numpy as np
-    logprobs = np.array([logits["excellent"], logits["good"], logits["fair"], logits["poor"], logits["bad"]])
-    probs = np.exp(logprobs) / np.sum(np.exp(logprobs))
-    return np.inner(probs, np.array([1,0.75,0.5,0.25,0.]))
-
-
-
-
-def disable_torch_init():
-    """
-    Disable the redundant torch default initialization to accelerate model creation.
-    """
-    import torch
-    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
-    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
-
-
-def load_image(image_file):
-    if image_file.startswith('http://') or image_file.startswith('https://'):
-        response = requests.get(image_file)
-        image = Image.open(BytesIO(response.content)).convert('RGB')
-    else:
-        image = Image.open(image_file).convert('RGB')
-    return image
-
-
-def main(args):
-    # Model
-    disable_torch_init()
-
-    model_name = get_model_name_from_path(args.model_path)
-    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
-    
-    
-    import json
-
-    
-    image_path = "playground/data/"
-                  
-
-    json_prefix = "playground/data/test_jsons/"
-    jsons = [
-        json_prefix + "test_ava.json",
-    ]
-
-    os.makedirs(f"results/{args.model_path}/", exist_ok=True)
-
-
-    conv_mode = "mplug_owl2"
-    
-    inp = "How would you rate the aesthetics of this image?"
-        
-    conv = conv_templates[conv_mode].copy()
-    inp = DEFAULT_IMAGE_TOKEN + inp
-    conv.append_message(conv.roles[0], inp)
-    image = None
-        
-    conv.append_message(conv.roles[1], None)
-    prompt = conv.get_prompt() + " The aesthetics of the image is"
-    
-    toks = ["good", "poor", "high", "fair", "low", "excellent", "bad", "fine", "moderate",  "decent", "average", "medium", "acceptable"]
-    print(toks)
-    ids_ = [id_[1] for id_ in tokenizer(toks)["input_ids"]]
-    print(ids_)
-
-    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(args.device)
-    
-    for json_ in jsons:
-        with open(json_) as f:
-            iqadata = json.load(f)  
-
-            image_tensors = []
-            batch_data = []
-            prs, gts = [], []
-            for i, llddata in enumerate(tqdm(iqadata, desc="Evaluating [{}]".format(json_.split("/")[-1]))):
-                filename = llddata["image"]
-                llddata["logits"] = defaultdict(float)
-                
-                
-                
-                image = load_image(image_path + filename)
-                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
-                image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
-                image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().to(args.device)
-
-                image_tensors.append(image_tensor)
-                batch_data.append(llddata)
-
-                if i % 8 == 7 or i == len(iqadata) - 1:                     
-                    with torch.inference_mode():
-                        output_logits = model(input_ids.repeat(len(image_tensors), 1),
-                            images=torch.cat(image_tensors, 0))["logits"][:,-1]
-    
-                    for j, xllddata in enumerate(batch_data):
-                        for tok, id_ in zip(toks, ids_):
-                            xllddata["logits"][tok] += output_logits[j,id_].item()
-                        xllddata["score"] = wa5(xllddata["logits"])
-                        # print(llddata)
-                        prs.append(xllddata["score"])
-                        gts.append(xllddata["gt_score"])
-                        json_ = json_.replace("combined/", "combined-")
-                        with open(f"results/{args.model_path}/{json_.split('/')[-1]}", "a") as wf:
-                            json.dump(xllddata, wf)
-
-                    image_tensors = []
-                    batch_data = []
-                
-                #if i > 0 and i % 200 == 0:
-                #    print(spearmanr(prs,gts)[0], pearsonr(prs,gts)[0])
-            print("Spearmanr", spearmanr(prs,gts)[0], "Pearson", pearsonr(prs,gts)[0])
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model-path", type=str, default="q-future/one-align")
-    parser.add_argument("--model-base", type=str, default=None)
-    parser.add_argument("--device", type=str, default="cuda:0")
-    parser.add_argument("--conv-mode", type=str, default=None)
-    parser.add_argument("--temperature", type=float, default=0.2)
-    parser.add_argument("--max-new-tokens", type=int, default=512)
-    parser.add_argument("--load-8bit", action="store_true")
-    parser.add_argument("--load-4bit", action="store_true")
-    parser.add_argument("--debug", action="store_true")
-    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
-    args = parser.parse_args()
-    main(args)

q_align/evaluate/.ipynb_checkpoints/iqa4vqa_eval-checkpoint.py

@@ -1,150 +0,0 @@
-import argparse
-import torch
-
-from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
-from q_align.conversation import conv_templates, SeparatorStyle
-from q_align.model.builder import load_pretrained_model
-from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
-
-from PIL import Image
-
-import requests
-from PIL import Image
-from io import BytesIO
-from transformers import TextStreamer
-
-from decord import VideoReader
-
-
-import json
-from tqdm import tqdm
-from collections import defaultdict
-
-import os
-
-
-
-
-def disable_torch_init():
-    """
-    Disable the redundant torch default initialization to accelerate model creation.
-    """
-    import torch
-    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
-    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
-
-
-def load_video(video_file):
-    vr = VideoReader(video_file)
-
-    # Get video frame rate
-    fps = vr.get_avg_fps()
-
-    # Calculate frame indices for 1fps
-    frame_indices = [int(fps * i) for i in range(int(len(vr) / fps))]
-    frames = vr.get_batch(frame_indices).asnumpy()
-    return [Image.fromarray(frames[i]) for i in range(int(len(vr) / fps))]
-
-
-def main(args):
-    # Model
-    disable_torch_init()
-
-    model_name = get_model_name_from_path(args.model_path)
-    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
-    
-    
-    import json
-
-    
-    image_paths = [
-        "playground/data/",
-        "playground/data/",
-        "playground/data/KoNViD_1k_videos/",
-        "playground/data/maxwell/",
-    ]
-
-    json_prefix = "playground/data/test_jsons/"
-    jsons = [
-        json_prefix + "test_lsvq.json",
-        json_prefix + "test_lsvq_1080p.json",
-        json_prefix + "konvid.json",
-        json_prefix + "maxwell_test.json",
-    ]
-
-    os.makedirs(f"results/{args.model_path}/", exist_ok=True)
-
-
-    conv_mode = "mplug_owl2"
-    
-    inp = "How would you rate the quality of this image?"
-        
-    conv = conv_templates[conv_mode].copy()
-    inp =  inp + "\n" + DEFAULT_IMAGE_TOKEN
-    conv.append_message(conv.roles[0], inp)
-    image = None
-        
-    conv.append_message(conv.roles[1], None)
-    prompt = conv.get_prompt() + " The quality of the image is"
-    
-    toks = ["good", "poor", "high", "fair", "low", "excellent", "bad", "fine", "moderate",  "decent", "average", "medium", "acceptable"]
-    print(toks)
-    ids_ = [id_[1] for id_ in tokenizer(toks)["input_ids"]]
-    print(ids_)
-
-    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(args.device)
-    
-    for image_path, json_ in zip(image_paths, jsons):
-        with open(json_) as f:
-            iqadata = json.load(f) 
-            try:
-                for i, llddata in enumerate(tqdm(iqadata, desc="Evaluating [{}]".format(json_.split("/")[-1]))):
-                    filename = llddata["img_path"]
-                    llddata["logits"] = defaultdict(float)
-
-                    image = load_video(image_path + filename)
-                    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
-                    image = [expand2square(img, tuple(int(x*255) for x in image_processor.image_mean)) for img in image]
-                    image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().to(args.device)
-
-
-                    if True:
-                        with torch.inference_mode():
-                            output_logits = model(input_ids.repeat(image_tensor.shape[0], 1),
-                                images=image_tensor)["logits"][:,-1]
-
-                            for tok, id_ in zip(toks, ids_):
-                                llddata["logits"][tok] += output_logits.mean(0)[id_].item()
-                            # print(llddata)
-                            json_ = json_.replace("combined/", "combined-")
-                            with open(f"results/{args.model_path}/{json_.split('/')[-1]}", "a") as wf:
-                                json.dump(llddata, wf)
-            except:
-                continue
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model-path", type=str, default="q-future/q-align-image")
-    parser.add_argument("--model-base", type=str, default=None)
-    parser.add_argument("--device", type=str, default="cuda:0")
-    parser.add_argument("--conv-mode", type=str, default=None)
-    parser.add_argument("--temperature", type=float, default=0.2)
-    parser.add_argument("--max-new-tokens", type=int, default=512)
-    parser.add_argument("--load-8bit", action="store_true")
-    parser.add_argument("--load-4bit", action="store_true")
-    parser.add_argument("--debug", action="store_true")
-    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
-    args = parser.parse_args()
-    main(args)

q_align/evaluate/.ipynb_checkpoints/iqa_eval-checkpoint.py

@@ -1,156 +0,0 @@
-import argparse
-import torch
-
-from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
-from q_align.conversation import conv_templates, SeparatorStyle
-from q_align.model.builder import load_pretrained_model
-from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
-
-from PIL import Image
-
-import requests
-from PIL import Image
-from io import BytesIO
-from transformers import TextStreamer
-
-import json
-from tqdm import tqdm
-from collections import defaultdict
-
-import os
-
-
-
-
-def disable_torch_init():
-    """
-    Disable the redundant torch default initialization to accelerate model creation.
-    """
-    import torch
-    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
-    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
-
-
-def load_image(image_file):
-    if image_file.startswith('http://') or image_file.startswith('https://'):
-        response = requests.get(image_file)
-        image = Image.open(BytesIO(response.content)).convert('RGB')
-    else:
-        image = Image.open(image_file).convert('RGB')
-    return image
-
-
-def main(args):
-    # Model
-    disable_torch_init()
-
-    model_name = get_model_name_from_path(args.model_path)
-    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
-    
-    
-    import json
-
-    
-    image_path = "playground/data/"
-                  
-
-    json_prefix = "playground/data/test_jsons/"
-    jsons = [
-        json_prefix + "test_imagerewarddb.json",
-        json_prefix + "test_koniq.json",
-        json_prefix + "test_spaq.json",
-        json_prefix + "test_kadid.json",
-        json_prefix + "livec.json",
-        json_prefix + "agi.json",
-        json_prefix + "live.json",
-        json_prefix + "csiq.json",
-    ]
-
-    os.makedirs(f"results/{args.model_path}/", exist_ok=True)
-
-
-    conv_mode = "mplug_owl2"
-    
-    inp = "Evaluate the image quality of the following image."#"How would you rate the quality of this image?"
-        
-    conv = conv_templates[conv_mode].copy()
-    inp =  inp + "\n" + DEFAULT_IMAGE_TOKEN
-    conv.append_message(conv.roles[0], inp)
-    image = None
-        
-    conv.append_message(conv.roles[1], None)
-    prompt = conv.get_prompt() + " The quality of the image is"
-    
-    toks = ["good", "poor", "high", "fair", "low", "excellent", "bad", "fine", "moderate",  "decent", "average", "medium", "acceptable"]
-    print(toks)
-    ids_ = [id_[1] for id_ in tokenizer(toks)["input_ids"]]
-    print(ids_)
-
-    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(args.device)
-    
-    for json_ in jsons:
-        with open(json_) as f:
-            iqadata = json.load(f)  
-
-            image_tensors = []
-            batch_data = []
-            
-            for i, llddata in enumerate(tqdm(iqadata, desc="Evaluating [{}]".format(json_.split("/")[-1]))):
-                if True:
-                    try:
-                        filename = llddata["image"]
-                    except:
-                        filename = llddata["img_path"]
-                    llddata["logits"] = defaultdict(float)
-
-                    image = load_image(image_path + filename)
-                    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
-                    image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
-                    image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().to(args.device)
-
-                    image_tensors.append(image_tensor)
-                    batch_data.append(llddata)
-
-                    if i % 8 == 7 or i == len(iqadata) - 1:                     
-                        with torch.inference_mode():
-                            output_logits = model(input_ids.repeat(len(image_tensors), 1),
-                                images=torch.cat(image_tensors, 0))["logits"][:,-1]
-
-                        for j, xllddata in enumerate(batch_data):
-                            for tok, id_ in zip(toks, ids_):
-                                xllddata["logits"][tok] += output_logits[j,id_].item()
-                            # print(llddata)
-                            json_ = json_.replace("combined/", "combined-")
-                            with open(f"results/{args.model_path}/2{json_.split('/')[-1]}", "a") as wf:
-                                json.dump(xllddata, wf)
-
-                        image_tensors = []
-                        batch_data = []
-
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model-path", type=str, default="q-future/one-align")
-    parser.add_argument("--model-base", type=str, default=None)
-    parser.add_argument("--device", type=str, default="cuda:0")
-    parser.add_argument("--conv-mode", type=str, default=None)
-    parser.add_argument("--temperature", type=float, default=0.2)
-    parser.add_argument("--max-new-tokens", type=int, default=512)
-    parser.add_argument("--load-8bit", action="store_true")
-    parser.add_argument("--load-4bit", action="store_true")
-    parser.add_argument("--debug", action="store_true")
-    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
-    args = parser.parse_args()
-    main(args)

q_align/evaluate/.ipynb_checkpoints/scorer-checkpoint.py

@@ -1,155 +0,0 @@
-from PIL import Image
-
-import torch.nn as nn
-import torch
-
-from typing import List
-
-from q_align.model.builder import load_pretrained_model
-
-from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
-from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
-
-def load_video(video_file):
-    from decord import VideoReader
-    vr = VideoReader(video_file)
-
-    # Get video frame rate
-    fps = vr.get_avg_fps()
-
-    # Calculate frame indices for 1fps
-    frame_indices = [int(fps * i) for i in range(int(len(vr) / fps))]
-    frames = vr.get_batch(frame_indices).asnumpy()
-    return [Image.fromarray(frames[i]) for i in range(int(len(vr) / fps))]
-
-
-class QAlignScorer(nn.Module):
-    def __init__(self, pretrained="q-future/one-align", device="cuda:0", tokenizer=None, model=None, image_processor=None):
-        super().__init__()
-        if model is None:
-            tokenizer, model, image_processor, _ = load_pretrained_model(pretrained, None, "mplug_owl2", device=device)
-        prompt = "USER: How would you rate the quality of this image?\n<|image|>\nASSISTANT: The quality of the image is"
-        
-        self.preferential_ids_ = [id_[1] for id_ in tokenizer(["excellent","good","fair","poor","bad"])["input_ids"]]
-        self.weight_tensor = torch.Tensor([1,0.75,0.5,0.25,0.]).half().to(model.device)
-    
-        self.tokenizer = tokenizer
-        self.model = model
-        self.image_processor = image_processor
-        self.input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
-        
-    def expand2square(self, 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 forward(self, image: List[Image.Image]):
-        image = [self.expand2square(img, tuple(int(x*255) for x in self.image_processor.image_mean)) for img in image]
-        with torch.inference_mode():
-            image_tensor = self.image_processor.preprocess(image, return_tensors="pt")["pixel_values"].half().to(self.model.device)
-            output_logits = self.model(self.input_ids.repeat(image_tensor.shape[0], 1),
-                        images=image_tensor)["logits"][:,-1, self.preferential_ids_]
-
-            return torch.softmax(output_logits, -1) #@ self.weight_tensor
-        
-        
-class QAlignAestheticScorer(nn.Module):
-    def __init__(self, pretrained="q-future/one-align", device="cuda:0", tokenizer=None, model=None, image_processor=None):
-        super().__init__()
-        if model is None:
-            tokenizer, model, image_processor, _ = load_pretrained_model(pretrained, None, "mplug_owl2", device=device)
-        prompt = "USER: How would you rate the aesthetics of this image?\n<|image|>\nASSISTANT: The aesthetics of the image is"
-        
-        self.preferential_ids_ = [id_[1] for id_ in tokenizer(["excellent","good","fair","poor","bad"])["input_ids"]]
-        self.weight_tensor = torch.Tensor([1,0.75,0.5,0.25,0.]).half().to(model.device)
-    
-        self.tokenizer = tokenizer
-        self.model = model
-        self.image_processor = image_processor
-        self.input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
-        
-    def expand2square(self, 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 forward(self, image: List[Image.Image]):
-        image = [self.expand2square(img, tuple(int(x*255) for x in self.image_processor.image_mean)) for img in image]
-        with torch.inference_mode():
-            image_tensor = self.image_processor.preprocess(image, return_tensors="pt")["pixel_values"].half().to(self.model.device)
-            output_logits = self.model(self.input_ids.repeat(image_tensor.shape[0], 1),
-                        images=image_tensor)["logits"][:,-1, self.preferential_ids_]
-
-            return torch.softmax(output_logits, -1) #@ self.weight_tensor
-        
-class QAlignVideoScorer(nn.Module):
-    def __init__(self, pretrained="q-future/one-align", device="cuda:0", tokenizer=None, model=None, image_processor=None):
-        super().__init__()
-        if model is None:
-            tokenizer, model, image_processor, _ = load_pretrained_model(pretrained, None, "mplug_owl2", device=device)
-        prompt = "USER: How would you rate the quality of this video?\n<|image|>\nASSISTANT: The quality of the video is"
-        
-        self.preferential_ids_ = [id_[1] for id_ in tokenizer(["excellent","good","fair","poor","bad"])["input_ids"]]
-        self.weight_tensor = torch.Tensor([1,0.75,0.5,0.25,0.]).half().to(model.device)
-    
-        self.tokenizer = tokenizer
-        self.model = model
-        self.image_processor = image_processor
-        self.input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
-        
-    def expand2square(self, 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 forward(self, video: List[List[Image.Image]]):
-        video = [[self.expand2square(frame, tuple(int(x*255) for x in self.image_processor.image_mean)) for frame in vid] for vid in video]
-        with torch.inference_mode():
-            video_tensors = [self.image_processor.preprocess(vid, return_tensors="pt")["pixel_values"].half().to(self.model.device) for vid in video]
-            output_logits = self.model(self.input_ids.repeat(len(video_tensors), 1),
-                        images=video_tensors)["logits"][:,-1, self.preferential_ids_]
-            return torch.softmax(output_logits, -1) #@ self.weight_tensor
-    
-
-if __name__ == "__main__":
-    import argparse
-    
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model-path", type=str, default="q-future/one-align")
-    parser.add_argument("--device", type=str, default="cuda:0")
-    parser.add_argument("--img_path", type=str, default="fig/singapore_flyer.jpg")
-    parser.add_argument("--aesthetic", action="store_true")
-    parser.add_argument("--video", action="store_true")
-    args = parser.parse_args()
-
-    if args.video:
-        scorer = QAlignVideoScorer(pretrained=args.model_path, device=args.device)
-        print(scorer([load_video(args.img_path)]).tolist())
-    else:
-        scorer = QAlignScorer(pretrained=args.model_path, device=args.device) if not args.aesthetic else QAlignAestheticScorer(pretrained=args.model_path, device=args.device) 
-        print(scorer([Image.open(args.img_path)]).tolist())
-

q_align/evaluate/.ipynb_checkpoints/vqa_eval-checkpoint.py

@@ -1,167 +0,0 @@
-import argparse
-import torch
-
-from q_align.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
-from q_align.conversation import conv_templates, SeparatorStyle
-from q_align.model.builder import load_pretrained_model
-from q_align.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
-
-from PIL import Image
-
-import requests
-from PIL import Image
-from io import BytesIO
-from transformers import TextStreamer
-
-
-from scipy.stats import spearmanr, pearsonr
-
-import json
-from tqdm import tqdm
-from collections import defaultdict
-
-import os
-
-def wa5(logits):
-    import numpy as np
-    logprobs = np.array([logits["excellent"], logits["good"], logits["fair"], logits["poor"], logits["bad"]])
-    probs = np.exp(logprobs) / np.sum(np.exp(logprobs))
-    return np.inner(probs, np.array([1,0.75,0.5,0.25,0.]))
-
-
-
-def disable_torch_init():
-    """
-    Disable the redundant torch default initialization to accelerate model creation.
-    """
-    import torch
-    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
-    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
-
-
-def load_video(video_file):
-    from decord import VideoReader
-    vr = VideoReader(video_file)
-
-    # Get video frame rate
-    fps = vr.get_avg_fps()
-
-    # Calculate frame indices for 1fps
-    frame_indices = [int(fps * i) for i in range(int(len(vr) / fps))]
-    frames = vr.get_batch(frame_indices).asnumpy()
-    return [Image.fromarray(frames[i]) for i in range(int(len(vr) / fps))]
-
-
-def main(args):
-    # Model
-    disable_torch_init()
-
-    model_name = get_model_name_from_path(args.model_path)
-    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
-    
-    
-    import json
-
-    
-    image_paths = [
-        #"playground/data/",
-        #"playground/data/",
-        "playground/data/KoNViD_1k_videos/",
-        "playground/data/maxwell/",
-
-    ]
-
-    json_prefix = "playground/data/test_jsons/"
-    jsons = [
-        #json_prefix + "test_lsvq.json",
-        #json_prefix + "test_lsvq_1080p.json",
-        json_prefix + "konvid.json",
-        json_prefix + "maxwell_test.json",
-    ]
-
-    os.makedirs(f"results/{args.model_path}/", exist_ok=True)
-
-
-    conv_mode = "mplug_owl2"
-    
-    inp = "How would you rate the quality of this video?"
-        
-    conv = conv_templates[conv_mode].copy()
-    inp =  inp + "\n" + DEFAULT_IMAGE_TOKEN
-    conv.append_message(conv.roles[0], inp)
-    image = None
-        
-    conv.append_message(conv.roles[1], None)
-    prompt = conv.get_prompt() + " The quality of the video is"
-    
-    toks = ["good", "poor", "high", "fair", "low", "excellent", "bad", "fine", "moderate",  "decent", "average", "medium", "acceptable"]
-    print(toks)
-    ids_ = [id_[1] for id_ in tokenizer(toks)["input_ids"]]
-    print(ids_)
-
-    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(args.device)
-    
-    for image_path, json_ in zip(image_paths, jsons):
-        with open(json_) as f:
-            iqadata = json.load(f) 
-            prs, gts = [], []
-            for i, llddata in enumerate(tqdm(iqadata, desc="Evaluating [{}]".format(json_.split("/")[-1]))):
-                try:
-                    try:
-                        filename = llddata["img_path"]
-                    except:
-                        filename = llddata["image"]
-                    llddata["logits"] = defaultdict(float)
-
-                    image = load_video(image_path + filename)
-                    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
-                    image = [expand2square(img, tuple(int(x*255) for x in image_processor.image_mean)) for img in image]
-                    image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().to(args.device)
-
-                    if True:
-                        with torch.inference_mode():
-                            output_logits = model(input_ids,
-                                images=[image_tensor])["logits"][:,-1]
-                            for tok, id_ in zip(toks, ids_):
-                                llddata["logits"][tok] += output_logits.mean(0)[id_].item()
-                            llddata["score"] = wa5(llddata["logits"])
-                            # print(llddata)
-                            prs.append(llddata["score"])
-                            gts.append(llddata["gt_score"])
-                            # print(llddata)
-                            json_ = json_.replace("combined/", "combined-")
-                            with open(f"results/{args.model_path}/2{json_.split('/')[-1]}", "a") as wf:
-                                json.dump(llddata, wf)
-
-                    if i > 0 and i % 200 == 0:
-                        print(spearmanr(prs,gts)[0], pearsonr(prs,gts)[0])
-                except:
-                    continue
-            print("Spearmanr", spearmanr(prs,gts)[0], "Pearson", pearsonr(prs,gts)[0])
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model-path", type=str, default="q-future/one-align")
-    parser.add_argument("--model-base", type=str, default=None)
-    parser.add_argument("--device", type=str, default="cuda:0")
-    parser.add_argument("--conv-mode", type=str, default=None)
-    parser.add_argument("--temperature", type=float, default=0.2)
-    parser.add_argument("--max-new-tokens", type=int, default=512)
-    parser.add_argument("--load-8bit", action="store_true")
-    parser.add_argument("--load-4bit", action="store_true")
-    parser.add_argument("--debug", action="store_true")
-    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
-    args = parser.parse_args()
-    main(args)

q_align/train/.ipynb_checkpoints/train-checkpoint.py

@@ -1,844 +0,0 @@
-# Adopted from https://github.com/lm-sys/FastChat. Below is the original copyright:
-# Adopted from tatsu-lab@stanford_alpaca. Below is the original copyright:
-#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
-#
-#    Licensed under the Apache License, Version 2.0 (the "License");
-#    you may not use this file except in compliance with the License.
-#    You may obtain a copy of the License at
-#
-#        http://www.apache.org/licenses/LICENSE-2.0
-#
-#    Unless required by applicable law or agreed to in writing, software
-#    distributed under the License is distributed on an "AS IS" BASIS,
-#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-#    See the License for the specific language governing permissions and
-#    limitations under the License.
-
-import os
-import copy
-from dataclasses import dataclass, field
-import json
-import logging
-import pathlib
-from typing import Dict, Optional, Sequence, List
-
-from PIL import ImageFile
-ImageFile.LOAD_TRUNCATED_IMAGES = True 
-
-import torch
-
-import transformers
-from transformers.models.clip.image_processing_clip import CLIPImageProcessor
-
-from torch.utils.data import Dataset
-from q_align.train.mplug_owl2_trainer import MPLUGOwl2Trainer
-from q_align.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
-
-from q_align import conversation as conversation_lib
-from q_align.model import *
-from q_align.mm_utils import tokenizer_image_token
-
-from PIL import Image
-from icecream import ic
-
-local_rank = None
-
-
-def rank0_print(*args):
-    if local_rank == 0:
-        print(*args)
-
-
-@dataclass
-class ModelArguments:
-    model_name_or_path: Optional[str] = field(default="facebook/opt-125m")
-    version: Optional[str] = field(default="v0")
-    freeze_backbone: bool = field(default=False)
-
-@dataclass
-class DataArguments:
-    data_path: str = field(default=None,
-                           metadata={"help": "Path to the training data."})
-    lazy_preprocess: bool = False
-    is_multimodal: bool = False
-    image_folder: Optional[str] = field(default=None)
-    image_aspect_ratio: str = 'square'
-    image_grid_pinpoints: Optional[str] = field(default=None)
-
-
-@dataclass
-class TrainingArguments(transformers.TrainingArguments):
-    cache_dir: Optional[str] = field(default=None)
-    optim: str = field(default="adamw_torch")
-    remove_unused_columns: bool = field(default=False)
-    
-    tune_visual_abstractor: bool = field(default=True)
-    freeze_vision_model: bool = field(default=True)
-
-    model_max_length: int = field(
-        default=512,
-        metadata={
-            "help":
-            "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
-        },
-    )
-    double_quant: bool = field(
-        default=True,
-        metadata={"help": "Compress the quantization statistics through double quantization."}
-    )
-    quant_type: str = field(
-        default="nf4",
-        metadata={"help": "Quantization data type to use. Should be one of `fp4` or `nf4`."}
-    )
-    bits: int = field(
-        default=16,
-        metadata={"help": "How many bits to use."}
-    )
-    lora_enable: bool = False
-    lora_r: int = 64
-    lora_alpha: int = 16
-    lora_dropout: float = 0.05
-    lora_weight_path: str = ""
-    lora_bias: str = "none"
-    visual_abstractor_lr: Optional[float] = None
-    group_by_modality_length: bool = field(default=False)
-
-
-def maybe_zero_3(param, ignore_status=False, name=None):
-    from deepspeed import zero
-    from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
-    if hasattr(param, "ds_id"):
-        if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
-            if not ignore_status:
-                logging.warning(f"{name}: param.ds_status != ZeroParamStatus.NOT_AVAILABLE: {param.ds_status}")
-        with zero.GatheredParameters([param]):
-            param = param.data.detach().cpu().clone()
-    else:
-        param = param.detach().cpu().clone()
-    return param
-
-
-# Borrowed from peft.utils.get_peft_model_state_dict
-def get_peft_state_maybe_zero_3(named_params, bias):
-    if bias == "none":
-        to_return = {k: t for k, t in named_params if "lora_" in k}
-    elif bias == "all":
-        to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k}
-    elif bias == "lora_only":
-        to_return = {}
-        maybe_lora_bias = {}
-        lora_bias_names = set()
-        for k, t in named_params:
-            if "lora_" in k:
-                to_return[k] = t
-                bias_name = k.split("lora_")[0] + "bias"
-                lora_bias_names.add(bias_name)
-            elif "bias" in k:
-                maybe_lora_bias[k] = t
-        for k, t in maybe_lora_bias:
-            if bias_name in lora_bias_names:
-                to_return[bias_name] = t
-    else:
-        raise NotImplementedError
-    to_return = {k: maybe_zero_3(v, ignore_status=True) for k, v in to_return.items()}
-    return to_return
-
-
-def get_peft_state_non_lora_maybe_zero_3(named_params, require_grad_only=True):
-    to_return = {k: t for k, t in named_params if "lora_" not in k}
-    if require_grad_only:
-        to_return = {k: t for k, t in to_return.items() if t.requires_grad}
-    to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}
-    return to_return
-
-
-def get_mm_adapter_state_maybe_zero_3(named_params, keys_to_match):
-    to_return = {k: t for k, t in named_params if any(key_match in k for key_match in keys_to_match)}
-    to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}
-    return to_return
-
-
-def find_all_linear_names(model):
-    cls = torch.nn.Linear
-    lora_module_names = set()
-    multimodal_keywords = ['vision_model', 'visual_abstractor']
-    for name, module in model.named_modules():
-        if any(mm_keyword in name for mm_keyword in multimodal_keywords):
-            continue
-        if isinstance(module, cls):
-            lora_module_names.add(name)
-
-    if 'lm_head' in lora_module_names: # needed for 16-bit
-        lora_module_names.remove('lm_head')
-    return list(lora_module_names)
-
-
-def safe_save_model_for_hf_trainer(trainer: transformers.Trainer,
-                                   output_dir: str):
-    """Collects the state dict and dump to disk."""
-
-    if trainer.deepspeed:
-        torch.cuda.synchronize()
-        trainer.save_model(output_dir)
-        return
-
-    state_dict = trainer.model.state_dict()
-    if trainer.args.should_save:
-        cpu_state_dict = {
-            key: value.cpu()
-            for key, value in state_dict.items()
-        }
-        del state_dict
-        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa
-
-
-def smart_tokenizer_and_embedding_resize(
-    special_tokens_dict: Dict,
-    tokenizer: transformers.PreTrainedTokenizer,
-    model: transformers.PreTrainedModel,
-):
-    """Resize tokenizer and embedding.
-
-    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
-    """
-    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
-    model.resize_token_embeddings(len(tokenizer))
-
-    if num_new_tokens > 0:
-        input_embeddings = model.get_input_embeddings().weight.data
-        output_embeddings = model.get_output_embeddings().weight.data
-
-        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
-            dim=0, keepdim=True)
-        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
-            dim=0, keepdim=True)
-
-        input_embeddings[-num_new_tokens:] = input_embeddings_avg
-        output_embeddings[-num_new_tokens:] = output_embeddings_avg
-
-
-def _tokenize_fn(strings: Sequence[str],
-                 tokenizer: transformers.PreTrainedTokenizer) -> Dict:
-    """Tokenize a list of strings."""
-    tokenized_list = [
-        tokenizer(
-            text,
-            return_tensors="pt",
-            padding="longest",
-            max_length=tokenizer.model_max_length,
-            truncation=True,
-        ) for text in strings
-    ]
-    input_ids = labels = [
-        tokenized.input_ids[0] for tokenized in tokenized_list
-    ]
-    input_ids_lens = labels_lens = [
-        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item()
-        for tokenized in tokenized_list
-    ]
-    return dict(
-        input_ids=input_ids,
-        labels=labels,
-        input_ids_lens=input_ids_lens,
-        labels_lens=labels_lens,
-    )
-
-
-def _mask_targets(target, tokenized_lens, speakers):
-    # cur_idx = 0
-    cur_idx = tokenized_lens[0]
-    tokenized_lens = tokenized_lens[1:]
-    target[:cur_idx] = IGNORE_INDEX
-    for tokenized_len, speaker in zip(tokenized_lens, speakers):
-        if speaker == "human":
-            target[cur_idx+2:cur_idx + tokenized_len] = IGNORE_INDEX
-        cur_idx += tokenized_len
-
-
-def _add_speaker_and_signal(header, source, get_conversation=True):
-    """Add speaker and start/end signal on each round."""
-    BEGIN_SIGNAL = "### "
-    END_SIGNAL = "\n"
-    conversation = header
-    for sentence in source:
-        from_str = sentence["from"]
-        if from_str.lower() == "human":
-            from_str = conversation_lib.default_conversation.roles[0]
-        elif from_str.lower() == "gpt":
-            from_str = conversation_lib.default_conversation.roles[1]
-        else:
-            from_str = 'unknown'
-        sentence["value"] = (BEGIN_SIGNAL + from_str + ": " +
-                             sentence["value"] + END_SIGNAL)
-        if get_conversation:
-            conversation += sentence["value"]
-    conversation += BEGIN_SIGNAL
-    return conversation
-
-
-def preprocess_multimodal(
-    sources: Sequence[str],
-    data_args: DataArguments
-) -> Dict:
-    is_multimodal = data_args.is_multimodal
-    if not is_multimodal:
-        return sources
-
-    for source in sources:
-        for sentence in source:
-            if DEFAULT_IMAGE_TOKEN in sentence['value']:
-                sentence['value'] = sentence['value'].replace(DEFAULT_IMAGE_TOKEN, '').strip()
-                sentence['value'] = DEFAULT_IMAGE_TOKEN + '\n' + sentence['value']
-                sentence['value'] = sentence['value'].strip()
-               
-            replace_token = DEFAULT_IMAGE_TOKEN
-            sentence["value"] = sentence["value"].replace(DEFAULT_IMAGE_TOKEN, replace_token)
-        
-    return sources
-
-
-def preprocess_v1(
-    sources,
-    tokenizer: transformers.PreTrainedTokenizer,
-    has_image: bool = False
-) -> Dict:
-    conv = conversation_lib.default_conversation.copy()
-    roles = {"human": conv.roles[0], "gpt": conv.roles[1]}
-
-    # Apply prompt templates
-    conversations = []
-    for i, source in enumerate(sources):
-        if roles[source[0]["from"]] != conv.roles[0]:
-            # Skip the first one if it is not from human
-            source = source[1:]
-
-        conv.messages = []
-        for j, sentence in enumerate(source):
-            role = roles[sentence["from"]]
-            assert role == conv.roles[j % 2], f"{i}"
-            conv.append_message(role, sentence["value"])
-        conversations.append(conv.get_prompt())
-
-    # Tokenize conversations
-
-    if has_image:
-        input_ids = torch.stack([tokenizer_image_token(prompt, tokenizer, return_tensors='pt') for prompt in conversations], dim=0)
-    else:
-        input_ids = tokenizer(
-            conversations,
-            return_tensors="pt",
-            padding="longest",
-            max_length=tokenizer.model_max_length,
-            truncation=True,
-        ).input_ids
-
-    targets = input_ids.clone()
-
-    assert conv.sep_style == conversation_lib.SeparatorStyle.TWO or conv.sep_style == conversation_lib.SeparatorStyle.TWO_NO_SYS
-
-    # Mask targets
-    sep = conv.sep + conv.roles[1] + ": "
-    for conversation, target in zip(conversations, targets):
-        total_len = int(target.ne(tokenizer.pad_token_id).sum())
-
-        rounds = conversation.split(conv.sep2)
-        cur_len = 1
-        target[:cur_len] = IGNORE_INDEX
-        for i, rou in enumerate(rounds):
-            if rou == "":
-                break
-
-            parts = rou.split(sep)
-            if len(parts) != 2:
-                break
-            parts[0] += sep
-
-            if has_image:
-                round_len = len(tokenizer_image_token(rou, tokenizer))
-                instruction_len = len(tokenizer_image_token(parts[0], tokenizer)) - 2
-            else:
-                round_len = len(tokenizer(rou).input_ids)
-                instruction_len = len(tokenizer(parts[0]).input_ids) - 2
-
-            target[cur_len : cur_len + instruction_len] = IGNORE_INDEX
-
-            cur_len += round_len
-        target[cur_len:] = IGNORE_INDEX
-
-        if cur_len < tokenizer.model_max_length:
-            if cur_len != total_len:
-                target[:] = IGNORE_INDEX
-                print(
-                    f"WARNING: tokenization mismatch: {cur_len} vs. {total_len}."
-                    f" (ignored)"
-                )
-
-    return dict(
-        input_ids=input_ids,
-        labels=targets,
-    )
-
-
-def preprocess_plain(
-    sources: Sequence[str],
-    tokenizer: transformers.PreTrainedTokenizer,
-) -> Dict:
-    # add end signal and concatenate together
-    conversations = []
-    for source in sources:
-        assert len(source) == 2
-        assert DEFAULT_IMAGE_TOKEN in source[0]['value']
-        source[0]['value'] = DEFAULT_IMAGE_TOKEN
-        conversation = source[0]['value'] + source[1]['value'] + conversation_lib.default_conversation.sep
-        conversations.append(conversation)
-    # tokenize conversations
-    input_ids = [tokenizer_image_token(prompt, tokenizer, return_tensors='pt') for prompt in conversations]
-    targets = copy.deepcopy(input_ids)
-    for target, source in zip(targets, sources):
-        tokenized_len = len(tokenizer_image_token(source[0]['value'], tokenizer))
-        target[:tokenized_len] = IGNORE_INDEX
-
-    return dict(input_ids=input_ids, labels=targets)
-
-
-def preprocess(
-    sources: Sequence[str],
-    tokenizer: transformers.PreTrainedTokenizer,
-    has_image: bool = False
-) -> Dict:
-    """
-    Given a list of sources, each is a conversation list. This transform:
-    1. Add signal '### ' at the beginning each sentence, with end signal '\n';
-    2. Concatenate conversations together;
-    3. Tokenize the concatenated conversation;
-    4. Make a deepcopy as the target. Mask human words with IGNORE_INDEX.
-    """
-    if conversation_lib.default_conversation.sep_style == conversation_lib.SeparatorStyle.PLAIN:
-        return preprocess_plain(sources, tokenizer)
-    if conversation_lib.default_conversation.version.startswith("v1"):
-        return preprocess_v1(sources, tokenizer, has_image=has_image)
-    # add end signal and concatenate together
-    conversations = []
-    for source in sources:
-        header = f"{conversation_lib.default_conversation.system}\n\n"
-        conversation = _add_speaker_and_signal(header, source)
-        conversations.append(conversation)
-    # tokenize conversations
-    def get_tokenize_len(prompts):
-        return [len(tokenizer_image_token(prompt, tokenizer)) for prompt in prompts]
-    if has_image:
-        input_ids = [tokenizer_image_token(prompt, tokenizer, return_tensors='pt') for prompt in conversations]
-    else:
-        conversations_tokenized = _tokenize_fn(conversations, tokenizer)
-        input_ids = conversations_tokenized["input_ids"]
-
-    targets = copy.deepcopy(input_ids)
-    for target, source in zip(targets, sources):
-        if has_image:
-            tokenized_lens = get_tokenize_len([header] + [s["value"] for s in source])
-        else:
-            tokenized_lens = _tokenize_fn([header] + [s["value"] for s in source], tokenizer)["input_ids_lens"]
-        speakers = [sentence["from"] for sentence in source]
-        _mask_targets(target, tokenized_lens, speakers)
-
-    return dict(input_ids=input_ids, labels=targets)
-
-
-def load_video(video_file):
-    from decord import VideoReader
-    vr = VideoReader(video_file)
-
-    # Get video frame rate
-    fps = vr.get_avg_fps()
-
-    # Calculate frame indices for 1fps
-    frame_indices = [int(fps * i) for i in range(int(len(vr) / fps))]
-    frames = vr.get_batch(frame_indices).asnumpy()
-    return [Image.fromarray(frames[i]) for i in range(int(len(vr) / fps))]
-
-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 LazySupervisedDataset(Dataset):
-    """Dataset for supervised fine-tuning."""
-
-    def __init__(self, data_path: str,
-                 tokenizer: transformers.PreTrainedTokenizer,
-                 data_args: DataArguments):
-        super(LazySupervisedDataset, self).__init__()
-        list_data_dict = json.load(open(data_path, "r"))
-
-        rank0_print("Formatting inputs...Skip in lazy mode")
-        self.tokenizer = tokenizer
-        self.list_data_dict = list_data_dict
-        self.data_args = data_args
-
-    def __len__(self):
-        return len(self.list_data_dict)
-
-    @property
-    def lengths(self):
-        length_list = []
-        for sample in self.list_data_dict:
-            img_tokens = 128 if 'image' in sample else 0
-            length_list.append(sum(len(conv['value'].split()) for conv in sample['conversations']) + img_tokens)
-        return length_list
-
-
-    @property
-    def modality_lengths(self):
-        length_list = []
-        for sample in self.list_data_dict:
-            cur_len = sum(len(conv['value'].split()) for conv in sample['conversations'])
-            cur_len = cur_len if 'image' in sample else -cur_len
-            length_list.append(cur_len)
-        return length_list
-
-#     def __getitem__(self, i) -> Dict[str, torch.Tensor]:
-#         sources = self.list_data_dict[i]
-#         if isinstance(i, int):
-#             sources = [sources]
-#         assert len(sources) == 1, "Don't know why it is wrapped to a list"  # FIXME
-#         if 'image' in sources[0]:
-#             image_file = self.list_data_dict[i]['image']
-#             image_folder = self.data_args.image_folder
-#             processor = self.data_args.image_processor
-#             image = Image.open(os.path.join(image_folder, image_file)).convert('RGB')
-#             if self.data_args.image_aspect_ratio == 'pad':
-#                 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
-#                 image = expand2square(image, tuple(int(x*255) for x in processor.image_mean))
-#                 image = processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
-#             else:
-#                 image = processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
-#             sources = preprocess_multimodal(
-#                 copy.deepcopy([e["conversations"] for e in sources]),
-#                 self.data_args)
-#         else:
-#             sources = copy.deepcopy([e["conversations"] for e in sources])
-#         data_dict = preprocess(
-#             sources,
-#             self.tokenizer,
-#             has_image=('image' in self.list_data_dict[i]))
-#         if isinstance(i, int):
-#             data_dict = dict(input_ids=data_dict["input_ids"][0],
-#                              labels=data_dict["labels"][0])
-
-#         # image exist in the data
-#         if 'image' in self.list_data_dict[i]:
-#             data_dict['image'] = image
-#         elif self.data_args.is_multimodal:
-#             # image does not exist in the data, but the model is multimodal
-#             crop_size = self.data_args.image_processor.crop_size
-#             data_dict['image'] = torch.zeros(3, crop_size['height'], crop_size['width'])
-#         return data_dict
-    
-    def next_rand(self):
-        import random
-        return random.randint(0,len(self)-1)
-
-    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
-        while True:
-            sources = self.list_data_dict[i]
-            if isinstance(i, int):
-                sources = [sources]
-            assert len(sources) == 1, "Don't know why it is wrapped to a list"  # FIXME
-            if 'image' in sources[0]:
-                image_file = self.list_data_dict[i]['image']
-                    
-                image_folder = self.data_args.image_folder
-                processor = self.data_args.image_processor
-                from pathlib import Path
-                #if not Path(os.path.join(image_folder, image_file)).exists():
-                #    i = self.next_rand()
-                #    continue
-                if isinstance(image_file, list):
-                    # Multiple Images as Input
-                    try:
-                        image = [Image.open(os.path.join(image_folder, imfile)).convert('RGB') for imfile in image_file]
-                    except Exception as ex:
-                        print(ex)
-                        i = self.next_rand()
-                        continue
-                    if self.data_args.image_aspect_ratio == 'pad':
-                        image = [expand2square(img, tuple(int(x*255) for x in processor.image_mean)) for img in image]
-                        image = processor.preprocess(image, return_tensors='pt')['pixel_values']
-                    else:
-                        image = processor.preprocess(image, return_tensors='pt')['pixel_values']
-                elif os.path.join(image_folder, image_file).endswith("mp4"):
-                    # Video as Input
-                    image = load_video(os.path.join(image_folder, image_file))
-                    if self.data_args.image_aspect_ratio == 'pad':
-                        image = [expand2square(img, tuple(int(x*255) for x in processor.image_mean)) for img in image]
-                        image = processor.preprocess(image, return_tensors='pt')['pixel_values']
-                    else:
-                        image = processor.preprocess(image, return_tensors='pt')['pixel_values']
-                else:
-                    try:
-                        image = Image.open(os.path.join(image_folder, image_file)).convert('RGB')
-                    except Exception as ex:
-                        print(ex)
-                        i = self.next_rand()
-                        continue
-                    if self.data_args.image_aspect_ratio == 'pad':
-                        image = expand2square(image, tuple(int(x*255) for x in processor.image_mean))
-                        image = processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
-                    else:
-                        image = processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
-                sources = preprocess_multimodal(
-                    copy.deepcopy([e["conversations"] for e in sources]),
-                    self.data_args)
-            else:
-                
-                sources = copy.deepcopy([e["conversations"] for e in sources])
-            data_dict = preprocess(
-                sources,
-                self.tokenizer,
-                has_image=('image' in self.list_data_dict[i]))
-            if isinstance(i, int):
-                data_dict = dict(input_ids=data_dict["input_ids"][0],
-                                labels=data_dict["labels"][0])
-
-            # image exist in the data
-            if 'image' in self.list_data_dict[i]:
-                data_dict['image'] = image
-            elif self.data_args.is_multimodal:
-                # image does not exist in the data, but the model is multimodal
-                crop_size = self.data_args.image_processor.crop_size
-                data_dict['image'] = torch.zeros(3, crop_size['height'], crop_size['width'])
-            return data_dict
-
-
-@dataclass
-class DataCollatorForSupervisedDataset(object):
-    """Collate examples for supervised fine-tuning."""
-
-    tokenizer: transformers.PreTrainedTokenizer
-
-    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
-        input_ids, labels = tuple([instance[key] for instance in instances]
-                                  for key in ("input_ids", "labels"))
-        input_ids = torch.nn.utils.rnn.pad_sequence(
-            input_ids,
-            batch_first=True,
-            padding_value=self.tokenizer.pad_token_id)
-        labels = torch.nn.utils.rnn.pad_sequence(labels,
-                                                 batch_first=True,
-                                                 padding_value=IGNORE_INDEX)
-        input_ids = input_ids[:, :self.tokenizer.model_max_length]
-        labels = labels[:, :self.tokenizer.model_max_length]
-        batch = dict(
-            input_ids=input_ids,
-            labels=labels,
-            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
-        )
-
-        if 'image' in instances[0]:
-            images = [instance['image'] for instance in instances]
-            if all(x is not None and x.shape == images[0].shape for x in images):
-                batch['images'] = torch.stack(images)
-            else:
-                batch['images'] = images
-
-        return batch
-
-
-def make_supervised_data_module(tokenizer: transformers.PreTrainedTokenizer,
-                                data_args) -> Dict:
-    """Make dataset and collator for supervised fine-tuning."""
-    train_dataset = LazySupervisedDataset(tokenizer=tokenizer,
-                                data_path=data_args.data_path,
-                                data_args=data_args)
-    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
-    return dict(train_dataset=train_dataset,
-                eval_dataset=None,
-                data_collator=data_collator)
-
-
-def train():
-    global local_rank
-
-    parser = transformers.HfArgumentParser(
-        (ModelArguments, DataArguments, TrainingArguments))
-    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
-    local_rank = training_args.local_rank
-    compute_dtype = (torch.float16 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32))
-
-    bnb_model_from_pretrained_args = {}
-    if training_args.bits in [4, 8]:
-        from transformers import BitsAndBytesConfig
-        bnb_model_from_pretrained_args.update(dict(
-            device_map={"": training_args.device},
-            load_in_4bit=training_args.bits == 4,
-            load_in_8bit=training_args.bits == 8,
-            quantization_config=BitsAndBytesConfig(
-                load_in_4bit=training_args.bits == 4,
-                load_in_8bit=training_args.bits == 8,
-                llm_int8_threshold=6.0,
-                llm_int8_has_fp16_weight=False,
-                bnb_4bit_compute_dtype=compute_dtype,
-                bnb_4bit_use_double_quant=training_args.double_quant,
-                bnb_4bit_quant_type=training_args.quant_type # {'fp4', 'nf4'}
-            )
-        ))
-
-    model = MPLUGOwl2LlamaForCausalLM.from_pretrained(
-        model_args.model_name_or_path,
-        cache_dir=training_args.cache_dir,
-        **bnb_model_from_pretrained_args
-    )
-    model.config.use_cache = False
-
-    if model_args.freeze_backbone:
-        model.model.requires_grad_(False)
-
-    if training_args.bits in [4, 8]:
-        from peft import prepare_model_for_kbit_training
-        model.config.torch_dtype=(torch.float32 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32))
-        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=training_args.gradient_checkpointing)
-
-    if training_args.gradient_checkpointing:
-        if hasattr(model, "enable_input_require_grads"):
-            model.enable_input_require_grads()
-        else:
-            def make_inputs_require_grad(module, input, output):
-                output.requires_grad_(True)
-            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
-
-    if training_args.lora_enable:
-        from peft import LoraConfig, get_peft_model
-        lora_config = LoraConfig(
-            r=training_args.lora_r,
-            lora_alpha=training_args.lora_alpha,
-            target_modules=find_all_linear_names(model),
-            lora_dropout=training_args.lora_dropout,
-            bias=training_args.lora_bias,
-            task_type="CAUSAL_LM",
-        )
-        if training_args.bits == 16:
-            if training_args.bf16:
-                model.to(torch.bfloat16)
-            if training_args.fp16:
-                model.to(torch.float16)
-        rank0_print("Adding LoRA adapters...")
-        model = get_peft_model(model, lora_config)
-
-    tokenizer = transformers.AutoTokenizer.from_pretrained(
-        model_args.model_name_or_path,
-        cache_dir=training_args.cache_dir,
-        model_max_length=training_args.model_max_length,
-        padding_side="right",
-        use_fast=False,
-    )
-
-    
-    tokenizer.pad_token = tokenizer.unk_token
-    if model_args.version in conversation_lib.conv_templates:
-        conversation_lib.default_conversation = conversation_lib.conv_templates[model_args.version]
-    else:
-        conversation_lib.default_conversation = conversation_lib.conv_templates["vicuna_v1"]
-    
-    if not training_args.freeze_vision_model and training_args.bits in [4, 8]:
-        model.get_model().vision_model.to(dtype=compute_dtype, device=training_args.device)
-    else:
-        vision_tower = model.get_model().vision_model
-        vision_tower.to(dtype=torch.bfloat16 if training_args.bf16 else torch.float16, device=training_args.device)
-    
-    if training_args.tune_visual_abstractor and training_args.bits in [4, 8]:
-        model.get_model().visual_abstractor.to(dtype=compute_dtype, device=training_args.device)
-    else:
-        visual_abstractor = model.get_model().visual_abstractor
-        visual_abstractor.to(dtype=torch.bfloat16 if training_args.bf16 else torch.float16, device=training_args.device)
-    
-    data_args.image_processor = CLIPImageProcessor.from_pretrained(model_args.model_name_or_path)
-    data_args.is_multimodal = True
-
-    model.config.image_aspect_ratio = data_args.image_aspect_ratio
-    model.config.image_grid_pinpoints = data_args.image_grid_pinpoints
-    model.config.tune_visual_abstractor = model_args.tune_visual_abstractor = training_args.tune_visual_abstractor
-    ic(training_args.tune_visual_abstractor)
-    model.requires_grad_(True)
-    if training_args.tune_visual_abstractor:
-        # model.requires_grad_(False)
-        for p in model.get_model().visual_abstractor.parameters():
-            p.requires_grad = True
-            
-    model.config.freeze_vision_model = training_args.freeze_vision_model
-    ic(training_args.freeze_vision_model)
-    if training_args.freeze_vision_model:
-        for p in model.get_model().vision_model.parameters():
-            p.requires_grad = False
-            
-    model.config.visual_abstractor_lr = training_args.visual_abstractor_lr
-
-
-    if training_args.bits in [4, 8]:
-        from peft.tuners.lora import LoraLayer
-        for name, module in model.named_modules():
-            if isinstance(module, LoraLayer):
-                if training_args.bf16:
-                    module = module.to(torch.bfloat16)
-            if 'norm' in name:
-                module = module.to(torch.float32)
-            if 'lm_head' in name or 'embed_tokens' in name:
-                if hasattr(module, 'weight'):
-                    if training_args.bf16 and module.weight.dtype == torch.float32:
-                        module = module.to(torch.bfloat16)
-
-    data_module = make_supervised_data_module(tokenizer=tokenizer,
-                                              data_args=data_args)
-    trainer = MPLUGOwl2Trainer(model=model,
-                    tokenizer=tokenizer,
-                    args=training_args,
-                    **data_module)
-
-    # if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
-    #     trainer.train(resume_from_checkpoint=True)
-    # else:
-    #     trainer.train()
-    
-    # TODO I dont like auto resume << REMOVE IT AND UNCOMMENT THE ABOVE CODE
-    trainer.train()
-
-    trainer.save_state()
-
-    model.config.use_cache = True
-
-    if training_args.lora_enable:
-        state_dict = get_peft_state_maybe_zero_3(
-            model.named_parameters(), training_args.lora_bias
-        )
-        non_lora_state_dict = get_peft_state_non_lora_maybe_zero_3(
-            model.named_parameters()
-        )
-        if training_args.local_rank == 0 or training_args.local_rank == -1:
-            model.config.save_pretrained(training_args.output_dir)
-            model.save_pretrained(training_args.output_dir, state_dict=state_dict)
-            torch.save(non_lora_state_dict, os.path.join(training_args.output_dir, 'non_lora_trainables.bin'))
-    else:
-        safe_save_model_for_hf_trainer(trainer=trainer,
-                                       output_dir=training_args.output_dir)
-
-
-if __name__ == "__main__":
-    train()