skycaptioner_v1/scripts/vllm_struct_caption.py

import torch
import decord
import argparse

import pandas as pd
import numpy as np

from tqdm import tqdm
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer, AutoProcessor

from torch.utils.data import DataLoader

SYSTEM_PROMPT = "I need you to generate a structured and detailed caption for the provided video. The structured output and the requirements for each field are as shown in the following JSON content: {\"subjects\": [{\"appearance\": \"Main subject appearance description\", \"action\": \"Main subject action\", \"expression\": \"Main subject expression  (Only for human/animal categories, empty otherwise)\", \"position\": \"Subject position in the video (Can be relative position to other objects or spatial description)\", \"TYPES\": {\"type\": \"Main category (e.g., Human)\", \"sub_type\": \"Sub-category (e.g., Man)\"}, \"is_main_subject\": true}, {\"appearance\": \"Non-main subject appearance description\", \"action\": \"Non-main subject action\", \"expression\": \"Non-main subject expression (Only for human/animal categories, empty otherwise)\", \"position\": \"Position of non-main subject 1\", \"TYPES\": {\"type\": \"Main category (e.g., Vehicles)\", \"sub_type\": \"Sub-category (e.g., Ship)\"}, \"is_main_subject\": false}], \"shot_type\": \"Shot type(Options: long_shot/full_shot/medium_shot/close_up/extreme_close_up/other)\", \"shot_angle\": \"Camera angle(Options: eye_level/high_angle/low_angle/other)\", \"shot_position\": \"Camera position(Options: front_view/back_view/side_view/over_the_shoulder/overhead_view/point_of_view/aerial_view/overlooking_view/other)\", \"camera_motion\": \"Camera movement description\", \"environment\": \"Video background/environment description\", \"lighting\": \"Lighting information in the video\"}"


class VideoTextDataset(torch.utils.data.Dataset):
    def __init__(self, csv_path, model_path):
        self.meta = pd.read_csv(csv_path)
        self._path = 'path'
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.processor = AutoProcessor.from_pretrained(model_path)
  
    def __getitem__(self, index):
        row = self.meta.iloc[index]
        path = row[self._path]
        real_index = self.meta.index[index]
        vr = decord.VideoReader(path, ctx=decord.cpu(0), width=360, height=420)
        start = 0
        end = len(vr)
        # avg_fps = vr.get_avg_fps()
        index = self.get_index(end-start, 16, st=start)
        frames = vr.get_batch(index).asnumpy() # n h w c
        video_inputs = [torch.from_numpy(frames).permute(0, 3, 1, 2)]
        conversation = {
                    "role": "user",
                    "content": [
                        {
                            "type": "video",
                            "video": row['path'],
                            "max_pixels": 360 * 420, # 460800
                            "fps": 2.0,
                        },
                        {   
                            "type": "text", 
                            "text": SYSTEM_PROMPT
                        },
                    ],
                }
                
        # 生成 user_input
        user_input = self.processor.apply_chat_template(
            [conversation],
            tokenize=False,
            add_generation_prompt=True
        )
        results = dict()
        inputs = {
            'prompt': user_input,
            'multi_modal_data': {'video': video_inputs}
        }
        results["index"] = real_index
        results['input'] = inputs
        return results

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

    def get_index(self, video_size, num_frames, st=0):
        seg_size = max(0., float(video_size - 1) / num_frames)
        max_frame = int(video_size) - 1
        seq = []
        # index from 1, must add 1
        for i in range(num_frames):
            start = int(np.round(seg_size * i))
            # end = int(np.round(seg_size * (i + 1)))
            idx = min(start, max_frame)
            seq.append(idx+st)
        return seq
    
def result_writer(indices_list: list, result_list: list, meta: pd.DataFrame, column):
    flat_indices = []
    for x in zip(indices_list):
        flat_indices.extend(x)
    flat_results = []
    for x in zip(result_list):
        flat_results.extend(x)
    
    flat_indices = np.array(flat_indices)
    flat_results = np.array(flat_results)

    unique_indices, unique_indices_idx = np.unique(flat_indices, return_index=True)
    meta.loc[unique_indices, column[0]] = flat_results[unique_indices_idx]

    meta = meta.loc[unique_indices]
    return meta


def worker_init_fn(worker_id):
    # Set different seed for each worker
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    # Prevent deadlocks by setting timeout
    torch.set_num_threads(1)

def main():
    parser = argparse.ArgumentParser(description="SkyCaptioner-V1 vllm batch inference")
    parser.add_argument("--input_csv", default="./examples/test.csv")
    parser.add_argument("--out_csv", default="./examples/test_result.csv")
    parser.add_argument("--bs", type=int, default=4)
    parser.add_argument("--tp", type=int, default=1)
    parser.add_argument("--model_path", required=True, type=str, help="skycaptioner-v1 model path")
    args = parser.parse_args()
    
    dataset = VideoTextDataset(csv_path=args.input_csv, model_path=args.model_path)
    dataloader = DataLoader(
        dataset,
        batch_size=args.bs,
        num_workers=4,
        worker_init_fn=worker_init_fn,
        persistent_workers=True,
        timeout=180,
    )

    sampling_params = SamplingParams(temperature=0.05, max_tokens=2048)
    
    llm = LLM(model=args.model_path,
        gpu_memory_utilization=0.6, 
        max_model_len=31920,
        tensor_parallel_size=args.tp)
    
    indices_list = []
    caption_save = []
    for video_batch in tqdm(dataloader):
        indices = video_batch["index"]
        inputs = video_batch["input"]
        batch_user_inputs = []
        for prompt, video in zip(inputs['prompt'], inputs['multi_modal_data']['video'][0]):
            usi={'prompt':prompt, 'multi_modal_data':{'video':video}}
            batch_user_inputs.append(usi)
        outputs = llm.generate(batch_user_inputs, sampling_params, use_tqdm=False)
        struct_outputs = [output.outputs[0].text for output in outputs]

        indices_list.extend(indices.tolist())
        caption_save.extend(struct_outputs)
    
    meta_new = result_writer(indices_list, caption_save, dataset.meta, column=["structural_caption"])
    meta_new.to_csv(args.out_csv, index=False)
    print(f'Saved structural_caption to {args.out_csv}')

if __name__ == '__main__':
    main()