skycaptioner_v1/scripts/vllm_fusion_caption.py

import os
from pathlib import Path
import argparse
import glob
import time
import gc
from tqdm import tqdm
import torch
from transformers import AutoTokenizer
import pandas as pd
from vllm import LLM, SamplingParams
from torch.utils.data import DataLoader
import json
import random
from utils import result_writer

SYSTEM_PROMPT_I2V = """
You are an expert in video captioning. You are given a structured video caption and you need to compose it to be more natural and fluent in English.

## Structured Input
{structured_input}

## Notes
1. If there has an empty field, just ignore it and do not mention it in the output.
2. Do not make any semantic changes to the original fields. Please be sure to follow the original meaning.
3. If the action field is not empty, eliminate the irrelevant information in the action field that is not related to the timing action(such as wearings, background and environment information) to make a pure action field.

## Output Principles and Orders
1. First, eliminate the static information in the action field that is not related to the timing action, such as background or environment information.
2. Second, describe each subject with its pure action and expression if these fields exist.

## Output
Please directly output the final composed caption without any additional information.
"""

SYSTEM_PROMPT_T2V = """
You are an expert in video captioning. You are given a structured video caption and you need to compose it to be more natural and fluent in English.

## Structured Input
{structured_input}

## Notes
1. According to the action field information, change its name field to the subject pronoun in the action.
2. If there has an empty field, just ignore it and do not mention it in the output.
3. Do not make any semantic changes to the original fields. Please be sure to follow the original meaning.

## Output Principles and Orders
1. First, declare the shot_type, then declare the shot_angle and the shot_position fields.
2. Second, eliminate information in the action field that is not related to the timing action, such as background or environment information if action is not empty.
3. Third, describe each subject with its pure action, appearance, expression, position if these fields exist.
4. Finally, declare the environment and lighting if the environment and lighting fields are not empty.

## Output
Please directly output the final composed caption without any additional information.
"""

SHOT_TYPE_LIST = [
    'close-up shot',
    'extreme close-up shot',
    'medium shot',
    'long shot',
    'full shot',
]


class StructuralCaptionDataset(torch.utils.data.Dataset):
    def __init__(self, input_csv, model_path):
        self.meta = pd.read_csv(input_csv)
        self.task = args.task
        self.system_prompt = SYSTEM_PROMPT_T2V if self.task == 't2v' else SYSTEM_PROMPT_I2V
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        
    
    def __len__(self):
        return len(self.meta)
    
    def __getitem__(self, index):
        row = self.meta.iloc[index]
        real_index = self.meta.index[index]

        struct_caption = json.loads(row["structural_caption"])

        camera_movement = struct_caption.get('camera_motion', '')
        if camera_movement != '':
            camera_movement += '.'
        camera_movement = camera_movement.capitalize()
        
        fusion_by_llm = False
        cleaned_struct_caption = self.clean_struct_caption(struct_caption, self.task)
        if cleaned_struct_caption.get('num_subjects', 0) > 0:
            new_struct_caption = json.dumps(cleaned_struct_caption, indent=4, ensure_ascii=False)
            conversation = [
                {
                    "role": "system",
                    "content": self.system_prompt.format(structured_input=new_struct_caption),
                },
            ]
            text = self.tokenizer.apply_chat_template(
                conversation,
                tokenize=False,
                add_generation_prompt=True
            )
            fusion_by_llm = True
        else:
            text = '-'
        return real_index, fusion_by_llm, text, '-', camera_movement
    
    def clean_struct_caption(self, struct_caption, task):
        raw_subjects = struct_caption.get('subjects', [])
        subjects = []
        for subject in raw_subjects:
            subject_type = subject.get("TYPES", {}).get('type', '')
            subject_sub_type = subject.get("TYPES", {}).get('sub_type', '')
            if subject_type not in ["Human", "Animal"]:
                subject['expression'] = ''
            if subject_type == 'Human' and subject_sub_type == 'Accessory':
                subject['expression'] = ''
            if subject_sub_type != '':
                subject['name'] = subject_sub_type
            if 'TYPES' in subject:
                del subject['TYPES']
            if 'is_main_subject' in subject:
                del subject['is_main_subject']
            subjects.append(subject)

        to_del_subject_ids = []
        for idx, subject in enumerate(subjects):
            action = subject.get('action', '').strip()
            subject['action'] = action
            if random.random() > 0.9 and 'appearance' in subject:
                del subject['appearance']
            if random.random() > 0.9 and 'position' in subject:
                del subject['position']
            if task == 'i2v':
                # just keep name and action, expression in subjects
                dropped_keys = ['appearance', 'position']
                for key in dropped_keys:
                    if key in subject:
                        del subject[key]
                if subject['action'] == '' and ('expression' not in subject or subject['expression'] == ''):
                    to_del_subject_ids.append(idx)
        
        # delete the subjects according to the to_del_subject_ids
        for idx in sorted(to_del_subject_ids, reverse=True):
            del subjects[idx]


        shot_type = struct_caption.get('shot_type', '').replace('_', ' ')
        if shot_type not in SHOT_TYPE_LIST:
            struct_caption['shot_type'] = ''
        
        new_struct_caption = {
            'num_subjects': len(subjects),
            'subjects': subjects,
            'shot_type': struct_caption.get('shot_type', ''),
            'shot_angle': struct_caption.get('shot_angle', ''),
            'shot_position': struct_caption.get('shot_position', ''),
            'environment': struct_caption.get('environment', ''),
            'lighting': struct_caption.get('lighting', ''),
        }

        if task == 't2v' and random.random() > 0.9:
            del new_struct_caption['lighting']

        if task == 'i2v':
            drop_keys = ['environment', 'lighting', 'shot_type', 'shot_angle', 'shot_position']
            for drop_key in drop_keys:
                del new_struct_caption[drop_key]
        return new_struct_caption

def custom_collate_fn(batch):
    real_indices, fusion_by_llm, texts, original_texts, camera_movements = zip(*batch)
    return list(real_indices), list(fusion_by_llm), list(texts), list(original_texts), list(camera_movements)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Caption Fusion by LLM")
    parser.add_argument("--input_csv", default="./examples/test_result.csv")
    parser.add_argument("--out_csv", default="./examples/test_result_caption.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="LLM model path")
    parser.add_argument("--task", default='t2v', help="t2v or i2v")
    
    args = parser.parse_args()

    sampling_params = SamplingParams(
        temperature=0.1,
        max_tokens=512,
        stop=['\n\n']
    )
    # model_path = "/maindata/data/shared/public/Common-Models/Qwen2.5-32B-Instruct/"

   
    llm = LLM(
        model=args.model_path,
        gpu_memory_utilization=0.9,
        max_model_len=4096,
        tensor_parallel_size = args.tp
    )
    

    dataset = StructuralCaptionDataset(input_csv=args.input_csv, model_path=args.model_path)
    
    dataloader = DataLoader(
        dataset,
        batch_size=args.bs,
        num_workers=8,
        collate_fn=custom_collate_fn,
        shuffle=False,
        drop_last=False,
    )

    indices_list = []
    result_list = []
    for indices, fusion_by_llms, texts, original_texts, camera_movements in tqdm(dataloader):
        llm_indices, llm_texts, llm_original_texts, llm_camera_movements = [], [], [], []
        for idx, fusion_by_llm, text, original_text, camera_movement in zip(indices, fusion_by_llms, texts, original_texts, camera_movements):
            if fusion_by_llm:
                llm_indices.append(idx)
                llm_texts.append(text)
                llm_original_texts.append(original_text)
                llm_camera_movements.append(camera_movement)    
            else:
                indices_list.append(idx)
                caption = original_text + " " + camera_movement
                result_list.append(caption)
        if len(llm_texts) > 0:
            try:
                outputs = llm.generate(llm_texts, sampling_params, use_tqdm=False)
                results = []
                for output in outputs:
                    result = output.outputs[0].text.strip()
                    results.append(result)
                indices_list.extend(llm_indices)
            except Exception as e:
                print(f"Error at {llm_indices}: {str(e)}")
                indices_list.extend(llm_indices)
                results = llm_original_texts
            
            for result, camera_movement in zip(results, llm_camera_movements):
                # concat camera movement to fusion_caption
                llm_caption = result + " " + camera_movement
                result_list.append(llm_caption)
    torch.cuda.empty_cache()
    gc.collect()
    gathered_list = [indices_list, result_list]
    meta_new = result_writer(indices_list, result_list, dataset.meta, column=[f"{args.task}_fusion_caption"])
    meta_new.to_csv(args.out_csv, index=False)