| import argparse |
| import copy |
| import os.path as osp |
|
|
| import torch |
| from torch.utils.data import DataLoader |
|
|
| from mmengine import Config |
| from mmengine.dist import init_dist, get_dist_info |
| from mmengine.utils.dl_utils import set_multi_processing |
| from transformers import GenerationConfig |
| from xtuner.configs import cfgs_name_path |
| from xtuner.registry import BUILDER |
| from xtuner.tools.chat import TORCH_DTYPE_MAP |
| from xtuner.tools.utils import get_stop_criteria |
| from xtuner.utils import PROMPT_TEMPLATE |
|
|
|
|
| from projects.llava_sam2.datasets import video_lisa_collate_fn |
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser(description='RefCocoSeg') |
| parser.add_argument('config', help='config file name or path.') |
| parser.add_argument('--pth_model', default=None, help='pth model file') |
| parser.add_argument( |
| '--split', |
| default='val', |
| help='Specify a split') |
| parser.add_argument( |
| '--prompt-template', |
| choices=PROMPT_TEMPLATE.keys(), |
| default='internlm2_chat', |
| help='Specify a prompt template') |
| parser.add_argument( |
| '--stop-words', nargs='+', type=str, default=[], help='Stop words') |
| parser.add_argument( |
| '--torch-dtype', |
| default='fp16', |
| choices=TORCH_DTYPE_MAP.keys(), |
| help='Override the default `torch.dtype` and load the model under ' |
| 'a specific `dtype`.') |
| parser.add_argument( |
| '--bits', |
| type=int, |
| choices=[4, 8, None], |
| default=None, |
| help='LLM bits') |
| parser.add_argument( |
| '--bot-name', type=str, default='BOT', help='Name for Bot') |
| parser.add_argument( |
| '--offload-folder', |
| default=None, |
| help='The folder in which to offload the model weights (or where the ' |
| 'model weights are already offloaded).') |
| parser.add_argument( |
| '--max-new-tokens', |
| type=int, |
| default=100, |
| help='Maximum number of new tokens allowed in generated text') |
| parser.add_argument( |
| '--seed', |
| type=int, |
| default=0, |
| help='Random seed for reproducible text generation') |
| parser.add_argument( |
| '--launcher', |
| choices=['none', 'pytorch', 'slurm', 'mpi'], |
| default='none', |
| help='job launcher') |
| args = parser.parse_args() |
| return args |
|
|
| def main(): |
| args = parse_args() |
| if args.launcher != 'none': |
| set_multi_processing(distributed=True) |
| init_dist(args.launcher) |
|
|
| rank, world_size = get_dist_info() |
| torch.cuda.set_device(rank) |
| else: |
| rank = 0 |
| world_size = 1 |
| print(f'Rank: {rank} / World size: {world_size}') |
|
|
| if not osp.isfile(args.config): |
| try: |
| args.config = cfgs_name_path[args.config] |
| except KeyError: |
| raise FileNotFoundError(f'Cannot find {args.config}') |
|
|
| cfg = Config.fromfile(args.config) |
| model_name = cfg.model.type if isinstance(cfg.model.type, str) else cfg.model.type.__name__ |
| assert model_name in ('VideoLLaVASAMModel',) |
|
|
| model = BUILDER.build(cfg.model) |
|
|
| if args.pth_model is not None: |
| state_dict_pth = args.pth_model |
| state_dict = torch.load(state_dict_pth, map_location='cpu') |
| model.load_state_dict(state_dict, strict=False) |
|
|
| model.to('cuda:0') |
| model.eval() |
|
|
| |
| tokenizer = model.tokenizer |
|
|
| |
| gen_config = GenerationConfig( |
| max_new_tokens=args.max_new_tokens, |
| do_sample=False, |
| eos_token_id=tokenizer.eos_token_id, |
| pad_token_id=tokenizer.pad_token_id |
| if tokenizer.pad_token_id is not None else tokenizer.eos_token_id, |
| ) |
| stop_words = args.stop_words |
| if args.prompt_template: |
| template = PROMPT_TEMPLATE[args.prompt_template] |
| stop_words += template.get('STOP_WORDS', []) |
| stop_criteria = get_stop_criteria( |
| tokenizer=tokenizer, stop_words=stop_words) |
|
|
| data_cfg = copy.deepcopy(cfg.video_revos_dataset) |
| data_cfg.update(expression_file=data_cfg.expression_file.replace('train', 'val')) |
| dataset = BUILDER.build(cfg.video_revos_dataset) |
| dataloader = DataLoader( |
| dataset, |
| batch_size=1, |
| num_workers=1, |
| shuffle=False, |
| collate_fn=video_lisa_collate_fn, |
| ) |
|
|
| for data_item in dataloader: |
| data_item = model.data_preprocessor(data_item) |
| inputs, data_samples = data_item['data'], data_item['data_samples'] |
| g_pixel_values = inputs.pop('g_pixel_values', None) |
| gt_masks = inputs.pop('masks', None) |
|
|
| output = model.mllm.generate( |
| pixel_values=inputs['pixel_values'], |
| input_ids=inputs['input_ids'], |
| attention_mask=inputs['attention_mask'], |
| visual_features=None, |
| generation_config=gen_config, |
| streamer=None, |
| bos_token_id=tokenizer.bos_token_id, |
| stopping_criteria=stop_criteria, |
| output_hidden_states=True, |
| return_dict_in_generate=True |
| ) |
| print(1) |
| print(1) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
