Upload folder using huggingface_hub

omg_llava/configs/finetune/ablation_multi_seg_states/ablation_multi_seg_states_baseline.py

@@ -0,0 +1,951 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = '/mnt/workspace/taozhang/chekpoints/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './work_dirs/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='/mnt/workspace/taozhang/chekpoints/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = '/mnt/workspace/taozhang/chekpoints/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 0
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/ablation_multi_seg_states/ablation_multi_seg_states_cat.py

@@ -0,0 +1,954 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = '/mnt/workspace/taozhang/chekpoints/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './work_dirs/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='/mnt/workspace/taozhang/chekpoints/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = '/mnt/workspace/taozhang/chekpoints/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 0
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    using_multilayer_states=True,
+    seg_token_merge_type='cat',
+    selected_layers=32,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/ablation_multi_seg_states/ablation_multi_seg_states_cat_debug.py

@@ -0,0 +1,927 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = '/mnt/workspace/taozhang/chekpoints/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './work_dirs/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='/mnt/workspace/taozhang/chekpoints/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = '/mnt/workspace/taozhang/chekpoints/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 0
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    using_multilayer_states=True,
+    seg_token_merge_type='cat',
+    selected_layers=32,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[glamm_refcocog_dataset, ],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/ablation_multi_seg_states/ablation_multi_seg_states_linear_cat.py

@@ -0,0 +1,954 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = '/mnt/workspace/taozhang/chekpoints/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './work_dirs/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='/mnt/workspace/taozhang/chekpoints/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = '/mnt/workspace/taozhang/chekpoints/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 0
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    using_multilayer_states=True,
+    seg_token_merge_type='linear_cat',
+    selected_layers=32,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/ablation_multi_seg_states/ablation_multi_seg_states_linearcat_debug.py

@@ -0,0 +1,927 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = '/mnt/workspace/taozhang/chekpoints/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './work_dirs/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='/mnt/workspace/taozhang/chekpoints/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = '/mnt/workspace/taozhang/chekpoints/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 0
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    using_multilayer_states=True,
+    seg_token_merge_type='linear_cat',
+    selected_layers=32,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[glamm_refcocog_dataset, ],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/ablation_multi_seg_states/ablation_multi_seg_states_mean.py

@@ -0,0 +1,954 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = '/mnt/workspace/taozhang/chekpoints/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './work_dirs/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='/mnt/workspace/taozhang/chekpoints/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = '/mnt/workspace/taozhang/chekpoints/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 0
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    using_multilayer_states=True,
+    seg_token_merge_type='mean',
+    selected_layers=32,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/ablation_multi_seg_states/debug.py

@@ -0,0 +1,924 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = '/mnt/workspace/taozhang/chekpoints/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './work_dirs/omg_llava_7b_pretrain_1024image_8gpus/iter_4361.pth'
+omg_ov_class_embed_path='/mnt/workspace/taozhang/chekpoints/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = '/mnt/workspace/taozhang/chekpoints/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[glamm_refcocog_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/ablation_projector/finetune_ablation_projector_cross.py

@@ -0,0 +1,953 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_7b_pretrain_1024image_seperate_cross_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    visual_prompt_proj=False,
+    add_cross_attn_layer=True,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/ablation_projector/finetune_ablation_projector_seperate.py

@@ -0,0 +1,953 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    visual_prompt_proj=True,
+    add_cross_attn_layer=False,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/ablation_projector/finetune_ablation_projector_seperate_cross.py

@@ -0,0 +1,953 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_7b_pretrain_1024image_seperate_cross_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    visual_prompt_proj=True,
+    add_cross_attn_layer=True,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/ablation_projector/finetune_ablation_projector_seperate_cross_debug.py

@@ -0,0 +1,926 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_7b_pretrain_1024image_seperate_cross_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 0
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    visual_prompt_proj=True,
+    add_cross_attn_layer=True,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[mdpv_brief_description_lvis_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/debug.py

@@ -0,0 +1,967 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-7b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_mapillary_dataset = dict(
+    type=MapillarySemanticSegDataset,
+    data_path=mapillary_class_file,
+    image_folder=mapillary_image_path,
+    label_path=mapillary_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_pascal_part_dataset = dict(
+    type=PascalPartSemanticSegDataset,
+    data_path=pascal_file,
+    image_folder=pascal_part_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=pascal_part_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_paco_dataset = dict(
+    type=PacoSemanticSegDataset,
+    data_path=paco_file,
+    image_folder=paco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=pascal_part_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[glamm_refcocog_dataset,
+                   ],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/fix_unfrozen_bug_omg_llava_7b_finetune_stage1_1024image_uniSegFormat_8gpus.py

@@ -0,0 +1,951 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 0
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/hf_app.py

@@ -0,0 +1,951 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/omg_llava_20b_finetune_stage1_1024image_8gpus.py

@@ -0,0 +1,993 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-20b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_20b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 4
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_mapillary_dataset = dict(
+    type=MapillarySemanticSegDataset,
+    data_path=mapillary_class_file,
+    image_folder=mapillary_image_path,
+    label_path=mapillary_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_pascal_part_dataset = dict(
+    type=PascalPartSemanticSegDataset,
+    data_path=pascal_file,
+    image_folder=pascal_part_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=pascal_part_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_paco_dataset = dict(
+    type=PacoSemanticSegDataset,
+    data_path=paco_file,
+    image_folder=paco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=pascal_part_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/omg_llava_7b_convnextXXL_finetune_stage1_1024image_uniSegFormat_8gpus.py

@@ -0,0 +1,952 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_7b_pretrain_1024image_convnextXXL.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_xxlarge_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convxxl.pth' # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_xxlarge',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s34b_b82k_augreg_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[384, 768, 1536, 3072],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    clip_feat_channel=3072,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/omg_llava_7b_finetune_stage1_1024image_8gpus.py

@@ -0,0 +1,993 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_mapillary_dataset = dict(
+    type=MapillarySemanticSegDataset,
+    data_path=mapillary_class_file,
+    image_folder=mapillary_image_path,
+    label_path=mapillary_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_pascal_part_dataset = dict(
+    type=PascalPartSemanticSegDataset,
+    data_path=pascal_file,
+    image_folder=pascal_part_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=pascal_part_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_paco_dataset = dict(
+    type=PacoSemanticSegDataset,
+    data_path=paco_file,
+    image_folder=paco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=pascal_part_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/omg_llava_7b_finetune_stage1_1024image_8gpus_01.py

@@ -0,0 +1,1007 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-7b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+    num_proc=32
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+    num_proc=32
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+    num_proc=32
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    num_proc=32
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    num_proc=32
+)
+
+semantic_seg_mapillary_dataset = dict(
+    type=MapillarySemanticSegDataset,
+    data_path=mapillary_class_file,
+    image_folder=mapillary_image_path,
+    label_path=mapillary_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    num_proc=32
+)
+
+semantic_seg_pascal_part_dataset = dict(
+    type=PascalPartSemanticSegDataset,
+    data_path=pascal_file,
+    image_folder=pascal_part_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=pascal_part_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    num_proc=32
+)
+
+semantic_seg_paco_dataset = dict(
+    type=PacoSemanticSegDataset,
+    data_path=paco_file,
+    image_folder=paco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=pascal_part_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    num_proc=32
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    num_proc=32
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    num_proc=32
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    num_proc=32
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    num_proc=32
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    num_proc=32
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    num_proc=32
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/omg_llava_7b_finetune_stage1_1024image_decoupleGCG_8gpus.py

@@ -0,0 +1,1028 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn,\
+    DecoupledGranDfGCGDataset, DecoupledOpenPsgGCGDataset, DecoupledRefCOCOgGCGDataset, DecoupledFlickrGCGDataset,\
+    glamm_openpsg_decoupled_given_description_map_fn, glamm_openpsg_decoupled_given_objects_map_fn,\
+    glamm_flickr_decoupled_given_objects_map_fn, glamm_flickr_decoupled_given_description_map_fn,\
+    glamm_granf_decoupled_given_objects_map_fn, glamm_granf_decoupled_given_description_map_fn,\
+    glamm_refcocog_decoupled_given_objects_map_fn, glamm_refcocog_decoupled_given_description_map_fn
+
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 0
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset_given_description = dict(
+    type=DecoupledRefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_decoupled_given_description_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    mode='given_description'
+)
+
+glamm_refcocog_dataset_given_objects = dict(
+    type=DecoupledRefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_decoupled_given_objects_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    mode='given_objects'
+)
+
+glamm_grandf_dataset_given_description = dict(
+    type=DecoupledGranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_decoupled_given_description_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+    mode='given_description'
+)
+
+glamm_grandf_dataset_given_objects = dict(
+    type=DecoupledGranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_decoupled_given_objects_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+    mode='given_objects'
+)
+
+glamm_psg_dataset_given_description = dict(
+    type=DecoupledOpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_decoupled_given_description_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+    mode='given_description'
+)
+
+glamm_psg_dataset_given_objects = dict(
+    type=DecoupledOpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_decoupled_given_objects_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+    mode='given_objects'
+)
+
+glamm_flickr_dataset_given_description = dict(
+    type=DecoupledFlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_decoupled_given_description_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+    mode='given_description'
+)
+
+glamm_flickr_dataset_given_objects = dict(
+    type=DecoupledFlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_decoupled_given_objects_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+    mode='given_objects'
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset,
+                   glamm_flickr_dataset_given_description, glamm_flickr_dataset_given_objects,
+                   glamm_refcocog_dataset_given_objects, glamm_refcocog_dataset_given_description,
+                   glamm_psg_dataset_given_description, glamm_psg_dataset_given_objects,
+                   glamm_grandf_dataset_given_description, glamm_grandf_dataset_given_objects,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/omg_llava_7b_finetune_stage1_1024image_decoupleGCG_8gpus_debug.py

@@ -0,0 +1,1000 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn,\
+    DecoupledGranDfGCGDataset, DecoupledOpenPsgGCGDataset, DecoupledRefCOCOgGCGDataset, DecoupledFlickrGCGDataset,\
+    glamm_openpsg_decoupled_given_description_map_fn, glamm_openpsg_decoupled_given_objects_map_fn,\
+    glamm_flickr_decoupled_given_objects_map_fn, glamm_flickr_decoupled_given_description_map_fn,\
+    glamm_granf_decoupled_given_objects_map_fn, glamm_granf_decoupled_given_description_map_fn,\
+    glamm_refcocog_decoupled_given_objects_map_fn, glamm_refcocog_decoupled_given_description_map_fn
+
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 0
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset_given_description = dict(
+    type=DecoupledRefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_decoupled_given_description_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    mode='given_description'
+)
+
+glamm_refcocog_dataset_given_objects = dict(
+    type=DecoupledRefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_decoupled_given_objects_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    mode='given_objects'
+)
+
+glamm_grandf_dataset_given_description = dict(
+    type=DecoupledGranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_decoupled_given_description_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+    mode='given_description'
+)
+
+glamm_grandf_dataset_given_objects = dict(
+    type=DecoupledGranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_decoupled_given_objects_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+    mode='given_objects'
+)
+
+glamm_psg_dataset_given_description = dict(
+    type=DecoupledOpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_decoupled_given_description_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+    mode='given_description'
+)
+
+glamm_psg_dataset_given_objects = dict(
+    type=DecoupledOpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_decoupled_given_objects_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+    mode='given_objects'
+)
+
+glamm_flickr_dataset_given_description = dict(
+    type=DecoupledFlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_decoupled_given_description_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+    mode='given_description'
+)
+
+glamm_flickr_dataset_given_objects = dict(
+    type=DecoupledFlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_decoupled_given_objects_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+    mode='given_objects'
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[
+                   glamm_refcocog_dataset_given_objects, glamm_refcocog_dataset_given_description,
+                   ],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/omg_llava_7b_finetune_stage1_1024image_uniSegFormat_8gpus.py

@@ -0,0 +1,951 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './pretrained/omg_llava/omg_llava_7b_pretrain_1024image_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/omg_llava_7b_finetune_stage2_1024image_8gpus.py

@@ -0,0 +1,994 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './work_dirs/omg_llava_7b_finetune_stage1_1024image_8gpus/iter_27600.pth' # noqa: E501
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=True,
+    freeze_llm_with_lora=False,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_mapillary_dataset = dict(
+    type=MapillarySemanticSegDataset,
+    data_path=mapillary_class_file,
+    image_folder=mapillary_image_path,
+    label_path=mapillary_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_pascal_part_dataset = dict(
+    type=PascalPartSemanticSegDataset,
+    data_path=pascal_file,
+    image_folder=pascal_part_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=pascal_part_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+semantic_seg_paco_dataset = dict(
+    type=PacoSemanticSegDataset,
+    data_path=paco_file,
+    image_folder=paco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=pascal_part_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[llava_dataset, glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset, # repeat 3x
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   semantic_seg_ade20k_dataset, semantic_seg_cocostuff_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   region_cap_osprey_dataset, region_conversation_osprey_dataset,
+                   mdpv_detailed_description_ade20k_dataset,
+                   mdpv_detailed_description_cocostuff_10k_dataset,
+                   mdpv_detailed_description_cocostuff_164k_dataset,
+                   mdpv_detailed_description_vg_dataset,
+                   mdpv_brief_description_lvis_dataset,
+                   mdpv_brief_description_vg_dataset,
+                   mdpv_brief_description_ade20k_dataset,
+                   mdpv_brief_description_cocostuff10k_dataset,
+                   mdpv_brief_description_cocostuff164k_dataset,
+                   mdpv_qa_vg_dataset,
+                   mdpv_qa_lvis_dataset,
+                   mdpv_qa_ade20k_dataset,
+                   mdpv_qa_cocostuff10k_dataset,
+                   mdpv_qa_cocostuff164k_dataset,
+                   mdpv_multi_points_flicker30k_dataset,
+                   mdpv_multi_points_openpsg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/specific_tasks_finetune/finetune_gcg.py

@@ -0,0 +1,925 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './work_dirs/omg_llava_1024x_2stage_finetune_1_clear_reratio_rmqcache_uniformSegFormat_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=True,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[glamm_flickr_dataset, glamm_refcocog_dataset,
+                   glamm_grandf_dataset, glamm_psg_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/finetune/specific_tasks_finetune/finetune_refseg.py

@@ -0,0 +1,929 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from peft import LoraConfig
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.dataset import GranDfGCGDataset, FlickrGCGDataset, OpenPsgGCGDataset, RefCOCOgGCGDataset,\
+    CombineDataset, glamm_refcocog_map_fn, glamm_openpsg_map_fn, glamm_flickr_map_fn, glamm_granf_map_fn,\
+    ADE20kSemanticSegDataset, COCOStuffSemanticSegDataset, semantic_seg_map_fn, MapillarySemanticSegDataset,\
+    PascalPartSemanticSegDataset, pascal_part_map_fn, PacoSemanticSegDataset,\
+    RefcocoReferringSegDataset, referring_seg_map_fn, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset,\
+    OspreyRegionCaptionDataset, osprey_region_caption_map_fn,\
+    OspreyRegionConversationDataset, osprey_region_conversation_map_fn,\
+    MDPVPointDetailedCaptionDataset, mdpv_points_map_fn, MDPVPointBriefCaptionDataset,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_gcg_format_map_fn, osprey_region_caption_gcg_format_map_fn
+from xtuner.dataset.samplers import LengthGroupedSampler
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model import OpenCLIPBackbone_omgseg
+from omg_llava.model import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b'  # Please change to your own path
+pretrained_pth = './work_dirs/omg_llava_1024x_2stage_finetune_1_clear_reratio_rmqcache_uniformSegFormat_8gpus.pth'
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Instruct-150K/llava_v1_5_mix665k.json'
+image_folder = data_root + 'llava_images'
+
+glamm_data_root = './data/glamm_data/'
+
+refcocog_image_path = glamm_data_root + 'images/coco2014/train2014/'
+refcocog_ann_file = glamm_data_root + 'annotations/RefCOCOg_GCG_train.json'
+
+grandf_image_path = glamm_data_root + 'images/grandf/train/'
+grandf_ann_file = glamm_data_root + 'annotations/GranDf_HA_GCG_train.json'
+
+flickr_image_path = glamm_data_root + 'images/flickr30k/Flickr30K/'
+flickr_ann_file = glamm_data_root + 'annotations/flickr_mergedGT_GCG_train.json'
+
+psg_image_path = glamm_data_root + 'images/coco2017/'
+psg_ann_file = glamm_data_root + 'annotations/OpenPsgGCG_train.json'
+
+ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+ade20k_class_file = './omg_llava/dataset/utils/ade20k_classes.json'
+
+cocostuff_image_path = './data/glamm_data/images/coco2017/train2017/'
+cocostuff_class_file = './omg_llava/dataset/utils/cocostuff_classes.txt'
+cocostuff_label_path = './data/semantic_seg/coco_stuff/stuffthingmaps_trainval2017/train2017/'
+
+mapillary_image_path = './data/semantic_seg/mapillary/training/images/'
+mapillary_class_file = './data/semantic_seg/mapillary/config_v2.0.json'
+mapillary_label_path = './data/semantic_seg/mapillary/training/v2.0/labels/'
+
+pascal_part_image_path = './data/semantic_seg/pascal_part/VOCdevkit/VOC2010/JPEGImages/'
+pascal_file = './data/semantic_seg/pascal_part/train.json'
+
+paco_image_path = './data/glamm_data/images/coco2017/'
+paco_file = './data/semantic_seg/paco_lvis/paco_lvis_v1_train.json'
+
+referring_refcoco_image_path = refcocog_image_path
+referring_refcoco_data_path = "./data/ref_seg/"
+
+referring_refcoco_plus_image_path = refcocog_image_path
+referring_refcoco_plus_data_path = "./data/ref_seg/"
+
+referring_refcocog_image_path = refcocog_image_path
+referring_refcocog_data_path = "./data/ref_seg/"
+
+referring_refclef_image_path = "./data/ref_seg/saiapr_tc-12/"
+referring_refclef_data_path = "./data/ref_seg/"
+
+region_cap_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_cap_osprey_data_path = "./data/region_caption/osprey/osprey_detail_description.json"
+
+region_conversation_osprey_image_path = glamm_data_root + 'images/coco2014/train2014/'
+region_conversation_osprey_data_path = "./data/region_caption/osprey/osprey_conversation.json"
+
+mdpv_detailed_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_detailed_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_detailed_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_detailed_caption_point.json'
+
+mdpv_detailed_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_detailed_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_detailed_caption_point.json'
+
+mdpv_detailed_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_detailed_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_detailed_caption_point.json'
+
+mdpv_brief_caption_cocostuff_10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_brief_caption_cocostuff_10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_brief_caption_point.json'
+
+mdpv_brief_caption_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_brief_caption_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_brief_caption_point.json'
+
+mdpv_brief_caption_cocostuff_164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_cocostuff_164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_brief_caption_point.json'
+
+mdpv_brief_caption_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_brief_caption_vg_data_path = './data/mdpv_point/gpt4v_vg_brief_caption_point.json'
+
+mdpv_brief_caption_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_brief_caption_lvis_data_path = './data/mdpv_point/gpt4v_lvis_brief_caption_point.json'
+
+mdpv_qa_vg_image_path = './data/llava_data/llava_images/vg/VG_100K'
+mdpv_qa_vg_data_path = './data/mdpv_point/gpt4v_vg_QA_point.json'
+
+mdpv_qa_ade20k_image_path = './data/semantic_seg/ADEChallengeData2016/images/training/'
+mdpv_qa_ade20k_data_path = './data/mdpv_point/gpt4v_ade20k_QA_point.json'
+
+mdpv_qa_cocostuff164k_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_cocostuff164k_data_path = './data/mdpv_point/gpt4v_cocostuff_164k_QA_point.json'
+
+mdpv_qa_lvis_image_path = './data/glamm_data/images/coco2017/train2017'
+mdpv_qa_lvis_data_path = './data/mdpv_point/gpt4v_lvis_QA_point.json'
+
+mdpv_qa_cocostuff10k_image_path = glamm_data_root + 'images/coco2014/train2014/'
+mdpv_qa_cocostuff10k_data_path = './data/mdpv_point/gpt4v_cocostuff_10k_QA_point.json'
+
+mdpv_multi_points_flicker30k_image_path = './data/glamm_data/images/flickr30k/Flickr30K/'
+mdpv_multi_points_flicker30k_data_path = './data/mdpv_point/Flicker30K_multi_points_to_caption.json'
+
+mdpv_multi_points_openpsg_image_path = glamm_data_root + 'images/coco2017/train2017'
+mdpv_multi_points_openpsg_data_path = './data/mdpv_point/OpenPsgGCG_train_multi_points_to_caption.json'
+
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2 - 100)
+
+# Scheduler & Optimizer
+batch_size = 8  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 2e-4
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+
+# Save
+save_steps = 2000
+save_total_limit = 4  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 2000
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture',
+                     'Could you please give me a detailed description of the image? Please respond with interleaved segmentation masks for the corresponding parts of the answer.']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    require_omg_decoder=True,
+    pretrained_pth=pretrained_pth,
+    text2vision_projector=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    llm_lora=dict(
+        type=LoraConfig,
+        r=512,
+        lora_alpha=256,
+        lora_dropout=0.05,
+        bias='none',
+        task_type='CAUSAL_LM'),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+debug=False
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True)
+
+glamm_refcocog_dataset = dict(
+    type=RefCOCOgGCGDataset,
+    data_path=refcocog_ann_file,
+    image_folder=refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_refcocog_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+glamm_grandf_dataset = dict(
+    type=GranDfGCGDataset,
+    data_path=grandf_ann_file,
+    image_folder=grandf_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_granf_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=10,
+)
+
+glamm_psg_dataset = dict(
+    type=OpenPsgGCGDataset,
+    data_path=psg_ann_file,
+    image_folder=psg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_openpsg_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+glamm_flickr_dataset = dict(
+    type=FlickrGCGDataset,
+    data_path=flickr_ann_file,
+    image_folder=flickr_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=glamm_flickr_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=debug,
+    repeats=1,
+)
+
+semantic_seg_ade20k_dataset = dict(
+    type=ADE20kSemanticSegDataset,
+    data_path=ade20k_class_file,
+    image_folder=ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+semantic_seg_cocostuff_dataset = dict(
+    type=COCOStuffSemanticSegDataset,
+    data_path=cocostuff_class_file,
+    image_folder=cocostuff_image_path,
+    label_path=cocostuff_label_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=semantic_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+    gcg_format=True,
+)
+
+referring_seg_refcoco_dataset = dict(
+    type=RefcocoReferringSegDataset,
+    data_path=referring_refcoco_data_path,
+    image_folder=referring_refcoco_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcoco_plus_dataset = dict(
+    type=Refcoco_plus_ReferringSegDataset,
+    data_path=referring_refcoco_plus_data_path,
+    image_folder=referring_refcoco_plus_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refcocog_dataset = dict(
+    type=Refcocog_ReferringSegDataset,
+    data_path=referring_refcocog_data_path,
+    image_folder=referring_refcocog_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+referring_seg_refclef_dataset = dict(
+    type=Refclef_ReferringSegDataset,
+    data_path=referring_refclef_data_path,
+    image_folder=referring_refclef_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=referring_seg_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_cap_osprey_dataset = dict(
+    type=OspreyRegionCaptionDataset,
+    data_path=region_cap_osprey_data_path,
+    image_folder=region_cap_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_caption_gcg_format_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+region_conversation_osprey_dataset = dict(
+    type=OspreyRegionConversationDataset,
+    data_path=region_conversation_osprey_data_path,
+    image_folder=region_conversation_osprey_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=osprey_region_conversation_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_ade20k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_ade20k_data_path,
+    image_folder=mdpv_detailed_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_10k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_cocostuff_164k_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_detailed_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_detailed_description_vg_dataset = dict(
+    type=MDPVPointDetailedCaptionDataset,
+    data_path=mdpv_detailed_caption_vg_data_path,
+    image_folder=mdpv_detailed_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_vg_data_path,
+    image_folder=mdpv_brief_caption_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_10k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_cocostuff_164k_data_path,
+    image_folder=mdpv_brief_caption_cocostuff_164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_ade20k_data_path,
+    image_folder=mdpv_brief_caption_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_brief_description_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_brief_caption_lvis_data_path,
+    image_folder=mdpv_brief_caption_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_vg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_vg_data_path,
+    image_folder=mdpv_qa_vg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_ade20k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_ade20k_data_path,
+    image_folder=mdpv_qa_ade20k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_lvis_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_lvis_data_path,
+    image_folder=mdpv_qa_lvis_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff10k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff10k_data_path,
+    image_folder=mdpv_qa_cocostuff10k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_qa_cocostuff164k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_qa_cocostuff164k_data_path,
+    image_folder=mdpv_qa_cocostuff164k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_openpsg_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_openpsg_data_path,
+    image_folder=mdpv_multi_points_openpsg_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+mdpv_multi_points_flicker30k_dataset = dict(
+    type=MDPVPointBriefCaptionDataset,
+    data_path=mdpv_multi_points_flicker30k_data_path,
+    image_folder=mdpv_multi_points_flicker30k_image_path,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=mdpv_points_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+    repeats=1,
+)
+
+train_dataset = dict(
+    type=CombineDataset,
+    datasets_cfgs=[referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,  # repeat 3x
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,
+                   referring_seg_refcoco_dataset, referring_seg_refcoco_plus_dataset,
+                   referring_seg_refcocog_dataset, referring_seg_refclef_dataset,],
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=train_dataset,
+    sampler=dict(
+        type=LengthGroupedSampler,
+        length_property='modality_length',
+        per_device_batch_size=batch_size * accumulative_counts),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/pretrain/ablation_projector/ablation_projector_baseline.py

@@ -0,0 +1,377 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.dataset import DefaultSampler
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model.convnext_clip import OpenCLIPBackbone_omgseg
+from omg_llava.model.omg_seg import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model or model paths
+llm_name_or_path = '/mnt/workspace/taozhang/chekpoints/internlm2-chat-7b' # Please change to your own path
+omg_ov_class_embed_path='/mnt/workspace/taozhang/chekpoints/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = '/mnt/workspace/taozhang/chekpoints/omg_seg_convl.pth'  # Please change to your own path
+
+# Data paths
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Pretrain/blip_laion_cc_sbu_558k.json'
+image_folder = data_root + 'LLaVA-Pretrain/images'
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2)
+
+# Scheduler & Optimizer
+batch_size = 16  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 1e-3
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+# Save
+save_steps = 500
+save_total_limit = 2  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 200
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+# using coco class as the class classifier
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    text2vision_projector=True,
+    keep_omg_decoder_frozen=True,
+    add_seg_pretrain=True,
+    pixel_shuffle_ratio=2,
+    visual_prompt_proj=False,
+    add_cross_attn_layer=False,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=llava_dataset,
+    sampler=dict(type=DefaultSampler, shuffle=True),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/pretrain/ablation_projector/ablation_projector_seperate_cross.py

@@ -0,0 +1,377 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.dataset import DefaultSampler
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model.convnext_clip import OpenCLIPBackbone_omgseg
+from omg_llava.model.omg_seg import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model or model paths
+llm_name_or_path = '/mnt/workspace/taozhang/chekpoints/internlm2-chat-7b' # Please change to your own path
+omg_ov_class_embed_path='/mnt/workspace/taozhang/chekpoints/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = '/mnt/workspace/taozhang/chekpoints/omg_seg_convl.pth'  # Please change to your own path
+
+# Data paths
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Pretrain/blip_laion_cc_sbu_558k.json'
+image_folder = data_root + 'LLaVA-Pretrain/images'
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2)
+
+# Scheduler & Optimizer
+batch_size = 16  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 1e-3
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+# Save
+save_steps = 500
+save_total_limit = 2  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 200
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+# using coco class as the class classifier
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    text2vision_projector=True,
+    keep_omg_decoder_frozen=True,
+    add_seg_pretrain=True,
+    pixel_shuffle_ratio=2,
+    visual_prompt_proj=False,
+    add_cross_attn_layer=True,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=llava_dataset,
+    sampler=dict(type=DefaultSampler, shuffle=True),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/pretrain/ablation_projector/ablation_projector_seperate_cross_rmProjloss.py

@@ -0,0 +1,377 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.dataset import DefaultSampler
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model.convnext_clip import OpenCLIPBackbone_omgseg
+from omg_llava.model.omg_seg import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model or model paths
+llm_name_or_path = '/mnt/workspace/taozhang/chekpoints/internlm2-chat-7b' # Please change to your own path
+omg_ov_class_embed_path='/mnt/workspace/taozhang/chekpoints/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = '/mnt/workspace/taozhang/chekpoints/omg_seg_convl.pth'  # Please change to your own path
+
+# Data paths
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Pretrain/blip_laion_cc_sbu_558k.json'
+image_folder = data_root + 'LLaVA-Pretrain/images'
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2)
+
+# Scheduler & Optimizer
+batch_size = 16  # per_device
+accumulative_counts = 4
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 1e-3
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+# Save
+save_steps = 500
+save_total_limit = 2  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 200
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+# using coco class as the class classifier
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    text2vision_projector=True,
+    keep_omg_decoder_frozen=True,
+    add_seg_pretrain=True,
+    pixel_shuffle_ratio=2,
+    visual_prompt_proj=False,
+    add_cross_attn_layer=True,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=llava_dataset,
+    sampler=dict(type=DefaultSampler, shuffle=True),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/pretrain/omg_llava_20b_pretrain_1024image_8gpus.py

@@ -0,0 +1,379 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.dataset import DefaultSampler
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model.convnext_clip import OpenCLIPBackbone_omgseg
+from omg_llava.model.omg_seg import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model or model paths
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-20b' # Please change to your own path
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data paths
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Pretrain/blip_laion_cc_sbu_558k.json'
+image_folder = data_root + 'LLaVA-Pretrain/images'
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2)
+
+# Scheduler & Optimizer
+batch_size = 16  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 0
+max_epochs = 1
+optim_type = AdamW
+lr = 1e-3
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+# Save
+save_steps = 500
+save_total_limit = 2  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 200
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+# using coco class as the class classifier
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+
+
+
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    text2vision_projector=True,
+    keep_omg_decoder_frozen=True,
+    add_seg_pretrain=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=llava_dataset,
+    sampler=dict(type=DefaultSampler, shuffle=True),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/pretrain/omg_llava_7b_pretrain_1024image_8gpus.py

@@ -0,0 +1,375 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.dataset import DefaultSampler
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model.convnext_clip import OpenCLIPBackbone_omgseg
+from omg_llava.model.omg_seg import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model or model paths
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b' # Please change to your own path
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_large_d_320_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convl.pth'  # Please change to your own path
+
+# Data paths
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Pretrain/blip_laion_cc_sbu_558k.json'
+image_folder = data_root + 'LLaVA-Pretrain/images'
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2)
+
+# Scheduler & Optimizer
+batch_size = 16  # per_device
+accumulative_counts = 2
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 1e-3
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+# Save
+save_steps = 500
+save_total_limit = 2  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 200
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+# using coco class as the class classifier
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_large_d_320',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s29b_b131k_ft_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[192, 384, 768, 1536],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    text2vision_projector=True,
+    keep_omg_decoder_frozen=True,
+    add_seg_pretrain=True,
+    pixel_shuffle_ratio=2,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=llava_dataset,
+    sampler=dict(type=DefaultSampler, shuffle=True),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/configs/pretrain/omg_llava_7b_pretrain_1024image_convnextXXL_8gpus.py

@@ -0,0 +1,376 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+from mmengine.dataset import DefaultSampler
+from mmengine.hooks import (CheckpointHook, DistSamplerSeedHook, IterTimerHook,
+                            LoggerHook, ParamSchedulerHook)
+from mmengine.optim import AmpOptimWrapper, CosineAnnealingLR, LinearLR
+from torch.optim import AdamW
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, CLIPImageProcessor,
+                          CLIPVisionModel)
+
+from omg_llava.dataset import LLaVADataset
+from omg_llava.dataset.collect_fns import omg_llava_collate_fn
+from xtuner.dataset.map_fns import llava_map_fn, template_map_fn_factory
+from omg_llava.engine import DatasetInfoHook_withSpecoalTokens, EvaluateChatHook_withSpecialTokens
+from xtuner.engine.runner import TrainLoop
+from omg_llava.model import OMG_LLaVA
+from xtuner.utils import PROMPT_TEMPLATE
+from omg_llava.model.convnext_clip import OpenCLIPBackbone_omgseg
+from omg_llava.model.omg_seg import OMGSegVisualEncoder, Mask2FormerVideoSemSamHead
+
+from torch.nn import GroupNorm, ReLU
+
+from mmdet.models import BatchFixedSizePad, MSDeformAttnPixelDecoder, CrossEntropyLoss, \
+    DiceLoss, MaskFormerFusionHead, FocalLoss
+from mmdet.models.task_modules.assigners import HungarianAssigner, CrossEntropyLossCost, DiceCost
+from mmdet.models.task_modules.samplers import MaskPseudoSampler
+
+#######################################################################
+#                          PART 1  Settings                           #
+#######################################################################
+# Model or model paths
+llm_name_or_path = './pretrained/omg_llava/internlm2-chat-7b' # Please change to your own path
+omg_ov_class_embed_path='./pretrained/omg_llava/convnext_xxlarge_CocoPanopticOVDataset.pth' # Please change to your own path
+omg_head_pretrain_pth_path = './pretrained/omg_llava/omg_seg_convxxl.pth'  # Please change to your own path
+
+# Data paths
+data_root = './data/llava_data/'
+data_path = data_root + 'LLaVA-Pretrain/blip_laion_cc_sbu_558k.json'
+image_folder = data_root + 'LLaVA-Pretrain/images'
+prompt_template = PROMPT_TEMPLATE.internlm2_chat
+max_length = int(2048 - (1024 / 64)**2)
+
+# Scheduler & Optimizer
+batch_size = 16  # per_device
+accumulative_counts = 4
+dataloader_num_workers = 4
+max_epochs = 1
+optim_type = AdamW
+lr = 1e-3
+betas = (0.9, 0.999)
+weight_decay = 0
+max_norm = 1  # grad clip
+warmup_ratio = 0.03
+
+# Save
+save_steps = 500
+save_total_limit = 2  # Maximum checkpoints to keep (-1 means unlimited)
+
+# Evaluate the generation performance during the training
+evaluation_freq = 200
+SYSTEM = ''
+evaluation_images = './work_dirs/test.jpg'
+evaluation_inputs = ['请描述一下这张照片', 'Please describe this picture']
+
+#######################################################################
+#            PART 2  Model & Tokenizer & Image Processor              #
+#######################################################################
+tokenizer = dict(
+    type=AutoTokenizer.from_pretrained,
+    pretrained_model_name_or_path=llm_name_or_path,
+    trust_remote_code=True,
+    padding_side='right')
+
+image_processor = dict(
+    type=CLIPImageProcessor,
+    do_resize=True,
+    size=1024,
+    resample=3,
+    do_center_crop=True,
+    crop_size=1024,
+    do_rescale=True,
+    do_normalize=True,
+    image_mean=[0.4814, 0.4578, 0.4082],
+    image_std=[0.2686, 0.2613, 0.2757],
+    do_convert_rgb=True
+)
+
+# using coco class as the class classifier
+class_embed = 'convnext_large_d_320_CocoPanopticOVDataset'
+num_things_classes = 80
+num_stuff_classes = 53
+num_classes = num_things_classes + num_stuff_classes
+
+omgseg_model = dict(
+    type=OMGSegVisualEncoder,
+    data_preprocessor=None,
+    pixel_shuffle_down_ratio=2,
+    backbone=dict(
+        type=OpenCLIPBackbone_omgseg,
+        model_name='convnext_xxlarge',
+        fix=True,
+        init_cfg=dict(
+            type='clip_pretrain',
+            checkpoint='laion2b_s34b_b82k_augreg_soup'
+        )
+    ),
+    panoptic_head=dict(
+        type=Mask2FormerVideoSemSamHead,
+        sphere_cls=True,
+        ov_path=omg_ov_class_embed_path,
+        enable_box_query=False,
+        ov_classifier_name=class_embed,
+        logit=None,
+        in_channels=[384, 768, 1536, 3072],  # pass to pixel_decoder inside
+        strides=[4, 8, 16, 32],
+        feat_channels=256,
+        out_channels=256,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        num_queries=300,
+        num_transformer_feat_level=3,
+        pixel_decoder=dict(
+            type=MSDeformAttnPixelDecoder,
+            num_outs=3,
+            norm_cfg=dict(type=GroupNorm, num_groups=32),
+            act_cfg=dict(type=ReLU),
+            encoder=dict(  # DeformableDetrTransformerEncoder
+                num_layers=6,
+                layer_cfg=dict(  # DeformableDetrTransformerEncoderLayer
+                    self_attn_cfg=dict(  # MultiScaleDeformableAttention
+                        embed_dims=256,
+                        num_heads=8,
+                        num_levels=3,
+                        num_points=4,
+                        dropout=0.0,
+                        batch_first=True),
+                    ffn_cfg=dict(
+                        embed_dims=256,
+                        feedforward_channels=1024,
+                        num_fcs=2,
+                        ffn_drop=0.0,
+                        act_cfg=dict(type=ReLU, inplace=True)))),
+        positional_encoding=dict(num_feats=128, normalize=True)),
+        enforce_decoder_input_project=False,
+        positional_encoding=dict(num_feats=128, normalize=True),
+        transformer_decoder=dict(  # Mask2FormerTransformerDecoder
+            return_intermediate=True,
+            num_layers=9,
+            layer_cfg=dict(  # Mask2FormerTransformerDecoderLayer
+                self_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                cross_attn_cfg=dict(  # MultiheadAttention
+                    embed_dims=256,
+                    num_heads=8,
+                    dropout=0.0,
+                    batch_first=True),
+                ffn_cfg=dict(
+                    embed_dims=256,
+                    feedforward_channels=2048,
+                    num_fcs=2,
+                    ffn_drop=0.0,
+                    act_cfg=dict(type='ReLU', inplace=True))),
+            init_cfg=None),
+        loss_cls=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=False,
+            loss_weight=2.0,
+            reduction='mean',
+            class_weight=[1.0] * 240 + [0.1]),
+        loss_mask=dict(
+            type=CrossEntropyLoss,
+            use_sigmoid=True,
+            reduction='mean',
+            loss_weight=5.0),
+        loss_dice=dict(
+            type=DiceLoss,
+            use_sigmoid=True,
+            activate=True,
+            reduction='mean',
+            naive_dice=True,
+            eps=1.0,
+            loss_weight=5.0),
+        loss_iou=dict(
+            type=FocalLoss,
+            use_sigmoid=True,
+            loss_weight=2.0,
+            reduction='mean')
+    ),
+    panoptic_fusion_head=dict(
+        type=MaskFormerFusionHead,
+        num_things_classes=num_things_classes,
+        num_stuff_classes=num_stuff_classes,
+        loss_panoptic=None,
+        init_cfg=None),
+    train_cfg=dict(
+        num_points=12544,
+        oversample_ratio=3.0,
+        importance_sample_ratio=0.75,
+        assigner=dict(
+            type=HungarianAssigner,
+            match_costs=[
+                # dict(type=FlexibleClassificationCost, weight=2.0),
+                dict(type=CrossEntropyLossCost, weight=5.0, use_sigmoid=True),
+                dict(type=DiceCost, weight=5.0, pred_act=True, eps=1.0)
+            ]),
+        sampler=dict(type=MaskPseudoSampler)),
+    test_cfg=dict(
+        panoptic_on=True,
+        # For now, the dataset does not support
+        # evaluating semantic segmentation metric.
+        semantic_on=False,
+        instance_on=True,
+        # max_per_image is for instance segmentation.
+        max_per_image=100,
+        iou_thr=0.8,
+        # In Mask2Former's panoptic postprocessing,
+        # it will filter mask area where score is less than 0.5 .
+        filter_low_score=True),
+    init_cfg=dict(
+        type='Pretrained',
+        checkpoint=omg_head_pretrain_pth_path,
+    )
+)
+
+model = dict(
+    type=OMG_LLaVA,
+    freeze_llm=True,
+    freeze_visual_encoder=True,
+    text2vision_projector=True,
+    keep_omg_decoder_frozen=True,
+    add_seg_pretrain=True,
+    pixel_shuffle_ratio=2,
+    clip_feat_channel=3072,
+    llm=dict(
+        type=AutoModelForCausalLM.from_pretrained,
+        pretrained_model_name_or_path=llm_name_or_path,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        quantization_config=dict(
+            type=BitsAndBytesConfig,
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')),
+    visual_encoder=omgseg_model,
+    tokenizer=tokenizer,
+)
+
+#######################################################################
+#                      PART 3  Dataset & Dataloader                   #
+#######################################################################
+llava_dataset = dict(
+    type=LLaVADataset,
+    data_path=data_path,
+    image_folder=image_folder,
+    tokenizer=tokenizer,
+    image_processor=image_processor,
+    dataset_map_fn=llava_map_fn,
+    template_map_fn=dict(
+        type=template_map_fn_factory, template=prompt_template),
+    max_length=max_length,
+    pad_image_to_square=True,
+    debug=False,
+)
+
+train_dataloader = dict(
+    batch_size=batch_size,
+    num_workers=dataloader_num_workers,
+    dataset=llava_dataset,
+    sampler=dict(type=DefaultSampler, shuffle=True),
+    collate_fn=dict(type=omg_llava_collate_fn))
+
+#######################################################################
+#                    PART 4  Scheduler & Optimizer                    #
+#######################################################################
+# optimizer
+optim_wrapper = dict(
+    type=AmpOptimWrapper,
+    optimizer=dict(
+        type=optim_type, lr=lr, betas=betas, weight_decay=weight_decay),
+    clip_grad=dict(max_norm=max_norm, error_if_nonfinite=False),
+    accumulative_counts=accumulative_counts,
+    loss_scale='dynamic',
+    dtype='float16')
+
+# learning policy
+# More information: https://github.com/open-mmlab/mmengine/blob/main/docs/en/tutorials/param_scheduler.md  # noqa: E501
+param_scheduler = [
+    dict(
+        type=LinearLR,
+        start_factor=1e-5,
+        by_epoch=True,
+        begin=0,
+        end=warmup_ratio * max_epochs,
+        convert_to_iter_based=True),
+    dict(
+        type=CosineAnnealingLR,
+        eta_min=0.0,
+        by_epoch=True,
+        begin=warmup_ratio * max_epochs,
+        end=max_epochs,
+        convert_to_iter_based=True)
+]
+
+# train, val, test setting
+train_cfg = dict(type=TrainLoop, max_epochs=max_epochs)
+
+#######################################################################
+#                           PART 5  Runtime                           #
+#######################################################################
+# Log the dialogue periodically during the training process, optional
+custom_hooks = [
+    dict(type=DatasetInfoHook_withSpecoalTokens, tokenizer=tokenizer),
+    dict(
+        type=EvaluateChatHook_withSpecialTokens,
+        tokenizer=tokenizer,
+        image_processor=image_processor,
+        every_n_iters=evaluation_freq,
+        evaluation_inputs=evaluation_inputs,
+        evaluation_images=evaluation_images,
+        system=SYSTEM,
+        prompt_template=prompt_template)
+]
+
+# configure default hooks
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type=IterTimerHook),
+    # print log every 10 iterations.
+    logger=dict(type=LoggerHook, log_metric_by_epoch=False, interval=10),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type=ParamSchedulerHook),
+    # save checkpoint per `save_steps`.
+    checkpoint=dict(
+        type=CheckpointHook,
+        by_epoch=False,
+        interval=save_steps,
+        max_keep_ckpts=save_total_limit),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type=DistSamplerSeedHook),
+)
+
+# configure environment
+env_cfg = dict(
+    # whether to enable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+# set visualizer
+visualizer = None
+
+# set log level
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume training from the loaded checkpoint
+resume = False
+
+# Defaults to use random seed and disable `deterministic`
+randomness = dict(seed=None, deterministic=False)
+
+# set log processor
+log_processor = dict(by_epoch=False)

omg_llava/dataset/CombineDataset.py

@@ -0,0 +1,81 @@
+from torch.utils.data import Dataset
+import numpy as np
+
+class CombineDataset(Dataset):
+    def __init__(self,
+                 datasets_cfgs,
+                 ):
+        super().__init__()
+
+        self.datasets = []
+        self.datasets_length = []
+
+        self.tokenizer = datasets_cfgs[0].tokenizer
+        tokenizer_type = self.tokenizer['type']
+        del self.tokenizer['type']
+        self.tokenizer = tokenizer_type(**self.tokenizer)
+
+        self._add_special_tokens()
+
+        for i in range(len(datasets_cfgs)):
+            datasets_cfgs[i].tokenizer = self.tokenizer
+
+        for dataset_cfg in datasets_cfgs:
+            dataset = dataset_cfg['type']
+            del dataset_cfg['type']
+            dataset = dataset(**dataset_cfg)
+            self.datasets.append(dataset)
+            self.datasets_length.append(len(dataset))
+
+        self.dataset_threthold = []
+        for i, length in enumerate(self.datasets_length):
+            if i == 0:
+                self.dataset_threthold.append(length)
+            else:
+                self.dataset_threthold.append(length + self.dataset_threthold[i - 1])
+
+        np.random.seed(42)
+        self.shuffled_index = np.arange(self.dataset_threthold[-1])
+        np.random.shuffle(self.shuffled_index)
+
+    @property
+    def modality_length(self):
+        length_list = []
+        for dataset in self.datasets:
+            for data_dict in dataset.text_data:
+                cur_len = len(data_dict['input_ids'])
+                if data_dict.get('image', None) is None:
+                    cur_len = -cur_len
+                length_list.append(cur_len)
+        return length_list
+
+    def __len__(self):
+        return self.dataset_threthold[-1]
+
+    def __getitem__(self, index):
+        index = int(self.shuffled_index[index])
+        for i, thred in enumerate(self.dataset_threthold):
+            if index < thred:
+                break
+
+
+        if i == 0:
+            _index = index
+        else:
+            _index = index - self.dataset_threthold[i - 1]
+
+        return self.datasets[i][_index]
+
+    def _add_special_tokens(self):
+        assert hasattr(self, "tokenizer")
+        # Adding special tokens for pixel grounding
+        segmentation_tokens = ['[SEG]']
+        # Adding tokens for GCG
+        phrase_tokens = ['<p>', '</p>']
+        # add for visual prompt
+        region_tokens = ['<region>']
+        point_tokens = ['<mark>']
+        special_tokens = segmentation_tokens + phrase_tokens + region_tokens + point_tokens
+
+        self.tokenizer.add_tokens(special_tokens, special_tokens=True)
+        return

omg_llava/dataset/DecoupledGCGDataset.py

@@ -0,0 +1,381 @@
+import json
+import logging
+import os
+
+import torch
+from datasets import Dataset as HFDataset
+from datasets import DatasetDict, load_from_disk
+from mmengine import print_log
+from mmengine.config import Config, ConfigDict
+from PIL import Image
+from torch.utils.data import Dataset
+from pycocotools import mask
+import numpy as np
+import torch.nn.functional as F
+import copy
+
+from xtuner.registry import BUILDER
+from omg_llava.dataset.utils import expand2square, expand2square_mask
+from xtuner.dataset.huggingface import process_hf_dataset
+
+class DecoupledGCGDataset(Dataset):
+
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=32,
+                 debug=False,
+                 repeats=1,
+                 mode='given_description'):
+        super().__init__()
+
+        assert offline_processed_text_folder or (data_path and tokenizer)
+        self.debug = debug
+        if offline_processed_text_folder and data_path:
+            print_log(
+                'Both `offline_processed_text_folder` and '
+                '`data_path` are set, and we load dataset from'
+                '`offline_processed_text_folder` '
+                f'({offline_processed_text_folder})',
+                logger='current',
+                level=logging.WARNING)
+
+        if offline_processed_text_folder is not None:
+            raise NotImplementedError
+        else:
+            json_data = self.json_file_preprocess(data_path)
+            json_data = DatasetDict({'train': HFDataset.from_list(json_data)})
+            self.text_data = process_hf_dataset(
+                dataset=json_data,
+                tokenizer=tokenizer,
+                max_length=max_length,
+                dataset_map_fn=dataset_map_fn,
+                template_map_fn=template_map_fn,
+                split='train',
+                max_dataset_length=max_dataset_length,
+                remove_unused_columns=False,
+                pack_to_max_length=False,
+                with_image_token=True,
+                map_num_proc=num_proc,  # because limited mem
+            )
+
+        self.image_folder = image_folder
+        size = image_processor.crop_size
+        if isinstance(size, int):
+            self.image_h, self.image_w = size, size
+        else:
+            self.image_w, self.image_h = size
+
+        if isinstance(image_processor, dict) or isinstance(
+                image_processor, Config) or isinstance(image_processor,
+                                                       ConfigDict):
+            self.image_processor = BUILDER.build(image_processor)
+        else:
+            self.image_processor = image_processor
+        self.pad_image_to_square = pad_image_to_square
+        self.down_ratio = 1
+        self.repeats = repeats
+        self.mode = mode
+
+    def json_file_preprocess(self, data_path):
+        with open(data_path, 'r') as f:
+            json_data = json.load(f)
+
+        # for quickly debug with mini split
+        if self.debug:
+            json_data = json_data[:100]
+        return json_data
+
+    @property
+    def modality_length(self):
+        length_list = []
+        for data_dict in self.text_data:
+            cur_len = len(data_dict['input_ids'])
+            if data_dict.get('image', None) is None:
+                cur_len = -cur_len
+            length_list.append(cur_len)
+        length_list = length_list * self.repeats
+        return length_list
+
+    def __len__(self):
+        return len(self.text_data) * self.repeats
+
+    def real_len(self):
+        return len(self.text_data)
+
+    def decode_mask(self, object_masks, ori_height, ori_width):
+        binary_masks = []
+        for object_mask in object_masks:
+            binary_mask = np.zeros((ori_height, ori_width), dtype=np.uint8)
+            for seg in object_mask:
+                rles = mask.frPyObjects([seg], ori_height, ori_width)
+                m = mask.decode(rles)
+                m = m.astype(np.uint8)
+                binary_mask += m.squeeze()
+
+            binary_masks.append(binary_mask)
+        if len(binary_masks) == 0:
+            return None
+        masks = np.stack(binary_masks, axis=0)
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+        masks = torch.from_numpy(masks)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio, self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks
+
+    def __getitem__(self, index):
+        index = index % self.real_len()
+        data_dict = copy.deepcopy(self.text_data[index])
+
+        if data_dict.get('image', None) is not None:
+            image_file = data_dict['image']
+            image = Image.open(os.path.join(self.image_folder,
+                                            image_file)).convert('RGB')
+            ori_width, ori_height = image.size
+            if self.pad_image_to_square:
+                image = expand2square(
+                    image,
+                    tuple(
+                        int(x * 255) for x in self.image_processor.image_mean))
+            image = self.image_processor.preprocess(
+                image, return_tensors='pt')['pixel_values'][0]
+            data_dict['pixel_values'] = image
+
+            # process and get masks
+            data_dict['masks'] = self.decode_mask(data_dict['masks'], ori_height=ori_height, ori_width=ori_width)
+
+            assert self.mode in ['given_objects', 'given_description']
+            if self.mode == 'given_objects':
+                data_dict['regions'] = copy.deepcopy(data_dict['masks'])
+
+            # if data_dict['masks'] is None:
+            #     return self.__getitem__(0)
+        else:
+            if hasattr(self.image_processor, 'crop_size'):
+                crop_size = self.image_processor.crop_size
+            else:
+                crop_size = self.image_processor.size
+            data_dict['pixel_values'] = torch.zeros(3, crop_size['height'],
+                                                    crop_size['width'])
+            data_dict['masks'] = None
+        return data_dict
+
+class DecoupledRefCOCOgGCGDataset(DecoupledGCGDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 debug=False,
+                 repeats=1,
+                 mode='given_description',
+                 ):
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            debug=debug,
+            repeats=repeats,
+            mode=mode,
+        )
+
+    def json_file_preprocess(self, data_path):
+        json_data = json.load(open(data_path))
+        if self.debug:
+            json_data = json_data[:100]
+
+        # convert {id: dict} to dict(..., id=xx)
+        for idx in range(len(json_data)):
+            id = list(json_data[idx].keys())[0]
+            json_data[idx] = json_data[idx][id]
+            json_data[idx].update({'id': id})
+        return json_data
+
+class DecoupledGranDfGCGDataset(DecoupledGCGDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=4,
+                 debug=False,
+                 repeats=1,
+                 mode='given_description'):
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            num_proc=num_proc,
+            debug=debug,
+            repeats=repeats,
+            mode=mode
+        )
+
+    def decode_mask(self, object_masks, ori_height, ori_width):
+        binary_masks = []
+        for object_mask in object_masks:
+            binary_mask = np.zeros((ori_height, ori_width), dtype=np.uint8)
+
+            for rle in object_mask:
+                m = mask.decode(rle).astype(np.uint8)
+                binary_mask += m.squeeze()
+
+            binary_masks.append(binary_mask)
+        if len(binary_masks) == 0:
+            return None
+        masks = np.stack(binary_masks, axis=0)
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+        masks = torch.from_numpy(masks)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio, self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks
+
+class DecoupledOpenPsgGCGDataset(DecoupledGranDfGCGDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=4,
+                 debug=False,
+                 repeats=1,
+                 mode='given_description'):
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            num_proc=num_proc,
+            debug=debug,
+            repeats=repeats,
+            mode=mode
+        )
+
+class DecoupledFlickrGCGDataset(DecoupledGCGDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=4,
+                 debug=False,
+                 repeats=1,
+                 mode='given_description'
+                 ):
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            num_proc=num_proc,
+            debug=debug,
+            repeats=repeats,
+            mode=mode
+        )
+
+    def json_file_preprocess(self, data_path):
+        def filter_images(data_infos, min_size):
+            return [i for i, info in enumerate(data_infos) if min(info['width'], info['height']) >= min_size]
+
+        # convert {id: dict} to dict(..., id=xx)
+        from pycocotools.coco import COCO
+        self.coco = COCO(data_path)
+        self.image_ids = self.coco.getImgIds()
+        data_infos = []
+        total_ann_ids = []
+        removed_img_count = 0
+        for img_id in self.image_ids:
+            info = self.coco.loadImgs([img_id])[0]
+            if len(info['caption'].split(' ')) < 3:
+                removed_img_count += 1
+                continue
+            info['filename'] = info['file_name'].split('_')[-1]
+            info['height'] = int(info['height'])
+            info['width'] = int(info['width'])
+            data_infos.append(info)
+            ann_ids = self.coco.getAnnIds(imgIds=[img_id])
+            total_ann_ids.extend(ann_ids)
+        assert len(set(total_ann_ids)) == len(total_ann_ids), f"Non-unique annotation IDs in '{data_path}'!"
+        print(f'Removed {removed_img_count} images.')
+        data_infos = [data_infos[i] for i in filter_images(data_infos, min_size=32)]
+
+        # obtain_annotations
+        for data_info in data_infos:
+            ann_ids = self.coco.getAnnIds(imgIds=data_info['id'])
+            ann_info = self.coco.loadAnns(ann_ids)
+            data_info.update({'ann_info': ann_info})
+        if self.debug:
+            data_infos = data_infos[:32]
+        return data_infos
+
+    def decode_mask(self, object_masks, ori_height, ori_width):
+        binary_masks = []
+        for object_mask in object_masks:
+            binary_mask = mask.decode(object_mask).astype(np.uint8)
+            binary_masks.append(binary_mask)
+        if len(binary_masks) == 0:
+            return None
+        masks = np.stack(binary_masks, axis=0)
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+        masks = torch.from_numpy(masks)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio, self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks

omg_llava/dataset/GCGDataset.py

@@ -0,0 +1,364 @@
+import json
+import logging
+import os
+
+import torch
+from datasets import Dataset as HFDataset
+from datasets import DatasetDict, load_from_disk
+from mmengine import print_log
+from mmengine.config import Config, ConfigDict
+from PIL import Image
+from torch.utils.data import Dataset
+from pycocotools import mask
+import numpy as np
+import torch.nn.functional as F
+import copy
+
+from xtuner.registry import BUILDER
+from omg_llava.dataset.utils import expand2square, expand2square_mask
+from xtuner.dataset.huggingface import process_hf_dataset
+
+class GCGDataset(Dataset):
+
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=32,
+                 debug=False,
+                 repeats=1):
+        super().__init__()
+
+        assert offline_processed_text_folder or (data_path and tokenizer)
+        self.debug = debug
+        if offline_processed_text_folder and data_path:
+            print_log(
+                'Both `offline_processed_text_folder` and '
+                '`data_path` are set, and we load dataset from'
+                '`offline_processed_text_folder` '
+                f'({offline_processed_text_folder})',
+                logger='current',
+                level=logging.WARNING)
+
+        if offline_processed_text_folder is not None:
+            raise NotImplementedError
+        else:
+            json_data = self.json_file_preprocess(data_path)
+            json_data = DatasetDict({'train': HFDataset.from_list(json_data)})
+            self.text_data = process_hf_dataset(
+                dataset=json_data,
+                tokenizer=tokenizer,
+                max_length=max_length,
+                dataset_map_fn=dataset_map_fn,
+                template_map_fn=template_map_fn,
+                split='train',
+                max_dataset_length=max_dataset_length,
+                remove_unused_columns=False,
+                pack_to_max_length=False,
+                with_image_token=True,
+                map_num_proc=num_proc,  # because limited mem
+            )
+
+        self.image_folder = image_folder
+        size = image_processor.crop_size
+        if isinstance(size, int):
+            self.image_h, self.image_w = size, size
+        else:
+            self.image_w, self.image_h = size
+
+        if isinstance(image_processor, dict) or isinstance(
+                image_processor, Config) or isinstance(image_processor,
+                                                       ConfigDict):
+            self.image_processor = BUILDER.build(image_processor)
+        else:
+            self.image_processor = image_processor
+        self.pad_image_to_square = pad_image_to_square
+        self.down_ratio = 1
+        self.repeats = repeats
+
+    def json_file_preprocess(self, data_path):
+        with open(data_path, 'r') as f:
+            json_data = json.load(f)
+
+        # for quickly debug with mini split
+        if self.debug:
+            json_data = json_data[:100]
+        return json_data
+
+    @property
+    def modality_length(self):
+        length_list = []
+        for data_dict in self.text_data:
+            cur_len = len(data_dict['input_ids'])
+            if data_dict.get('image', None) is None:
+                cur_len = -cur_len
+            length_list.append(cur_len)
+        length_list = length_list * self.repeats
+        return length_list
+
+    def __len__(self):
+        return len(self.text_data) * self.repeats
+
+    def real_len(self):
+        return len(self.text_data)
+
+    def decode_mask(self, object_masks, ori_height, ori_width):
+        binary_masks = []
+        for object_mask in object_masks:
+            binary_mask = np.zeros((ori_height, ori_width), dtype=np.uint8)
+            for seg in object_mask:
+                rles = mask.frPyObjects([seg], ori_height, ori_width)
+                m = mask.decode(rles)
+                m = m.astype(np.uint8)
+                binary_mask += m.squeeze()
+
+            binary_masks.append(binary_mask)
+        if len(binary_masks) == 0:
+            return None
+        masks = np.stack(binary_masks, axis=0)
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+        masks = torch.from_numpy(masks)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio, self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks
+
+    def __getitem__(self, index):
+        index = index % self.real_len()
+        data_dict = copy.deepcopy(self.text_data[index])
+
+        if data_dict.get('image', None) is not None:
+            image_file = data_dict['image']
+            image = Image.open(os.path.join(self.image_folder,
+                                            image_file)).convert('RGB')
+            ori_width, ori_height = image.size
+            if self.pad_image_to_square:
+                image = expand2square(
+                    image,
+                    tuple(
+                        int(x * 255) for x in self.image_processor.image_mean))
+            image = self.image_processor.preprocess(
+                image, return_tensors='pt')['pixel_values'][0]
+            data_dict['pixel_values'] = image
+
+            # process and get masks
+            data_dict['masks'] = self.decode_mask(data_dict['masks'], ori_height=ori_height, ori_width=ori_width)
+            if data_dict['masks'] is None:
+                return self.__getitem__(0)
+        else:
+            if hasattr(self.image_processor, 'crop_size'):
+                crop_size = self.image_processor.crop_size
+            else:
+                crop_size = self.image_processor.size
+            data_dict['pixel_values'] = torch.zeros(3, crop_size['height'],
+                                                    crop_size['width'])
+            data_dict['masks'] = None
+        return data_dict
+
+class RefCOCOgGCGDataset(GCGDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 debug=False,
+                 repeats=1,):
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            debug=debug,
+            repeats=repeats,
+        )
+
+    def json_file_preprocess(self, data_path):
+        json_data = json.load(open(data_path))
+        if self.debug:
+            json_data = json_data[:100]
+
+        # convert {id: dict} to dict(..., id=xx)
+        for idx in range(len(json_data)):
+            id = list(json_data[idx].keys())[0]
+            json_data[idx] = json_data[idx][id]
+            json_data[idx].update({'id': id})
+        return json_data
+
+class GranDfGCGDataset(GCGDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=4,
+                 debug=False,
+                 repeats=1):
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            num_proc=num_proc,
+            debug=debug,
+            repeats=repeats,
+        )
+
+    def decode_mask(self, object_masks, ori_height, ori_width):
+        binary_masks = []
+        for object_mask in object_masks:
+            binary_mask = np.zeros((ori_height, ori_width), dtype=np.uint8)
+
+            for rle in object_mask:
+                m = mask.decode(rle).astype(np.uint8)
+                binary_mask += m.squeeze()
+
+            binary_masks.append(binary_mask)
+        if len(binary_masks) == 0:
+            return None
+        masks = np.stack(binary_masks, axis=0)
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+        masks = torch.from_numpy(masks)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio, self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks
+
+class OpenPsgGCGDataset(GranDfGCGDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=4,
+                 debug=False,
+                 repeats=1):
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            num_proc=num_proc,
+            debug=debug,
+            repeats=repeats,
+        )
+
+class FlickrGCGDataset(GCGDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=4,
+                 debug=False,
+                 repeats=1,):
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            num_proc=num_proc,
+            debug=debug,
+            repeats=repeats,
+        )
+
+    def json_file_preprocess(self, data_path):
+        def filter_images(data_infos, min_size):
+            return [i for i, info in enumerate(data_infos) if min(info['width'], info['height']) >= min_size]
+
+        # convert {id: dict} to dict(..., id=xx)
+        from pycocotools.coco import COCO
+        self.coco = COCO(data_path)
+        self.image_ids = self.coco.getImgIds()
+        data_infos = []
+        total_ann_ids = []
+        removed_img_count = 0
+        for img_id in self.image_ids:
+            info = self.coco.loadImgs([img_id])[0]
+            if len(info['caption'].split(' ')) < 3:
+                removed_img_count += 1
+                continue
+            info['filename'] = info['file_name'].split('_')[-1]
+            info['height'] = int(info['height'])
+            info['width'] = int(info['width'])
+            data_infos.append(info)
+            ann_ids = self.coco.getAnnIds(imgIds=[img_id])
+            total_ann_ids.extend(ann_ids)
+        assert len(set(total_ann_ids)) == len(total_ann_ids), f"Non-unique annotation IDs in '{data_path}'!"
+        print(f'Removed {removed_img_count} images.')
+        data_infos = [data_infos[i] for i in filter_images(data_infos, min_size=32)]
+
+        # obtain_annotations
+        for data_info in data_infos:
+            ann_ids = self.coco.getAnnIds(imgIds=data_info['id'])
+            ann_info = self.coco.loadAnns(ann_ids)
+            data_info.update({'ann_info': ann_info})
+        if self.debug:
+            data_infos = data_infos[:32]
+        return data_infos
+
+    def decode_mask(self, object_masks, ori_height, ori_width):
+        binary_masks = []
+        for object_mask in object_masks:
+            binary_mask = mask.decode(object_mask).astype(np.uint8)
+            binary_masks.append(binary_mask)
+        if len(binary_masks) == 0:
+            return None
+        masks = np.stack(binary_masks, axis=0)
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+        masks = torch.from_numpy(masks)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio, self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks

omg_llava/dataset/LlavaDataset.py

@@ -0,0 +1,134 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import json
+import logging
+import os
+
+import torch
+from datasets import Dataset as HFDataset
+from datasets import DatasetDict, load_from_disk
+from mmengine import print_log
+from mmengine.config import Config, ConfigDict
+from PIL import Image
+from torch.utils.data import Dataset
+
+from xtuner.registry import BUILDER
+from xtuner.dataset.huggingface import process_hf_dataset
+from .utils import expand2square
+import copy
+
+class LLaVADataset(Dataset):
+
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 debug=False):
+        super().__init__()
+
+        assert offline_processed_text_folder or (data_path and tokenizer)
+
+        self.tokenizer = tokenizer
+        if isinstance(tokenizer, dict) or isinstance(
+                tokenizer, Config) or isinstance(tokenizer, ConfigDict):
+            tokenizer_type = self.tokenizer['type']
+            del self.tokenizer['type']
+            self.tokenizer = tokenizer_type(**self.tokenizer)
+            self._add_special_tokens()
+
+        if offline_processed_text_folder and data_path:
+            print_log(
+                'Both `offline_processed_text_folder` and '
+                '`data_path` are set, and we load dataset from'
+                '`offline_processed_text_folder` '
+                f'({offline_processed_text_folder})',
+                logger='current',
+                level=logging.WARNING)
+
+        if offline_processed_text_folder is not None:
+            self.text_data = load_from_disk(offline_processed_text_folder)
+        else:
+            json_data = json.load(open(data_path))
+            if debug:
+                json_data = json_data[:10000]
+            for idx in range(len(json_data)):
+                if isinstance(json_data[idx]['id'], int):
+                    json_data[idx]['id'] = str(json_data[idx]['id'])
+            json_data = DatasetDict({'train': HFDataset.from_list(json_data)})
+            self.text_data = process_hf_dataset(
+                dataset=json_data,
+                tokenizer=self.tokenizer,
+                max_length=max_length,
+                dataset_map_fn=dataset_map_fn,
+                template_map_fn=template_map_fn,
+                split='train',
+                max_dataset_length=max_dataset_length,
+                remove_unused_columns=False,
+                pack_to_max_length=False,
+                with_image_token=True,
+                map_num_proc=32,  # because limited mem
+            )
+
+        self.image_folder = image_folder
+        if isinstance(image_processor, dict) or isinstance(
+                image_processor, Config) or isinstance(image_processor,
+                                                       ConfigDict):
+            self.image_processor = BUILDER.build(image_processor)
+        else:
+            self.image_processor = image_processor
+        self.pad_image_to_square = pad_image_to_square
+
+    @property
+    def modality_length(self):
+        length_list = []
+        for data_dict in self.text_data:
+            cur_len = len(data_dict['input_ids'])
+            if data_dict.get('image', None) is None:
+                cur_len = -cur_len
+            length_list.append(cur_len)
+        return length_list
+
+    def __len__(self):
+        return len(self.text_data)
+
+    def __getitem__(self, index):
+        data_dict = copy.deepcopy(self.text_data[index])
+        if data_dict.get('image', None) is not None:
+            image_file = data_dict['image']
+            image = Image.open(os.path.join(self.image_folder,
+                                            image_file)).convert('RGB')
+            if self.pad_image_to_square:
+                image = expand2square(
+                    image,
+                    tuple(
+                        int(x * 255) for x in self.image_processor.image_mean))
+            image = self.image_processor.preprocess(
+                image, return_tensors='pt')['pixel_values'][0]
+            data_dict['pixel_values'] = image
+        else:
+            if hasattr(self.image_processor, 'crop_size'):
+                crop_size = self.image_processor.crop_size
+            else:
+                crop_size = self.image_processor.size
+            data_dict['pixel_values'] = torch.zeros(3, crop_size['height'],
+                                                    crop_size['width'])
+        return data_dict
+
+    def _add_special_tokens(self):
+        assert hasattr(self, "tokenizer")
+        # Adding special tokens for pixel grounding
+        segmentation_tokens = ['[SEG]']
+        # Adding tokens for GCG
+        phrase_tokens = ['<p>', '</p>']
+        # add for visual prompt
+        region_tokens = ['<region>']
+        point_tokens = ['<mark>']
+        special_tokens = segmentation_tokens + phrase_tokens + region_tokens + point_tokens
+        self.tokenizer.add_tokens(special_tokens, special_tokens=True)
+        return

omg_llava/dataset/MDPVPointsDataset.py

@@ -0,0 +1,220 @@
+import json
+import logging
+import os
+import torch
+from datasets import Dataset as HFDataset
+from datasets import DatasetDict, load_from_disk
+from mmengine import print_log
+from mmengine.config import Config, ConfigDict
+from PIL import Image
+from torch.utils.data import Dataset
+from pycocotools import mask
+import numpy as np
+import torch.nn.functional as F
+
+from xtuner.registry import BUILDER
+from omg_llava.dataset.utils import expand2square, expand2square_mask, expand2square_points
+from xtuner.dataset.huggingface import process_hf_dataset
+import copy
+
+class MDPVPointDetailedCaptionDataset(Dataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=32,
+                 debug=False,
+                 repeats=1):
+        super().__init__()
+
+        assert offline_processed_text_folder or (data_path and tokenizer)
+        self.debug = debug
+        if offline_processed_text_folder and data_path:
+            print_log(
+                'Both `offline_processed_text_folder` and '
+                '`data_path` are set, and we load dataset from'
+                '`offline_processed_text_folder` '
+                f'({offline_processed_text_folder})',
+                logger='current',
+                level=logging.WARNING)
+
+        if offline_processed_text_folder is not None:
+            raise NotImplementedError
+        else:
+            json_data = self.json_file_preprocess(data_path)
+            self.json_data = json_data
+            hf_json_data = self.filter_hf_require_infos(json_data)
+            hf_json_data = DatasetDict({'train': HFDataset.from_list(hf_json_data)})
+            self.text_data = process_hf_dataset(
+                dataset=hf_json_data,
+                tokenizer=tokenizer,
+                max_length=max_length,
+                dataset_map_fn=dataset_map_fn,
+                template_map_fn=template_map_fn,
+                split='train',
+                max_dataset_length=max_dataset_length,
+                remove_unused_columns=False,
+                pack_to_max_length=False,
+                with_image_token=True,
+                map_num_proc=num_proc,  # because limited mem
+            )
+
+        self.image_folder = image_folder
+        size = image_processor.crop_size
+        if isinstance(size, int):
+            self.image_h, self.image_w = size, size
+        else:
+            self.image_w, self.image_h = size
+
+        if isinstance(image_processor, dict) or isinstance(
+                image_processor, Config) or isinstance(image_processor,
+                                                       ConfigDict):
+            self.image_processor = BUILDER.build(image_processor)
+        else:
+            self.image_processor = image_processor
+        self.pad_image_to_square = pad_image_to_square
+        self.down_ratio = 1
+        self.repeats = repeats
+
+    def filter_hf_require_infos(self, dataset_infos):
+        ret = []
+        for dataset_info in dataset_infos:
+            conversations = dataset_info["conversations"]
+            image = dataset_info['image'].split('/')[-1]
+            num_marks = len(dataset_info['points'])
+            required_info = {'image': image,
+                             'conversations': conversations,
+                             'num_marks': num_marks}
+            ret.append(required_info)
+        return ret
+
+    def json_file_preprocess(self, data_path):
+        with open(data_path, 'r') as f:
+            json_file = json.load(f)
+        if self.debug:
+            json_file = json_file[:10000]
+        return json_file
+
+    @property
+    def modality_length(self):
+        length_list = []
+        for data_dict in self.text_data:
+            cur_len = len(data_dict['input_ids'])
+            if data_dict.get('image', None) is None:
+                cur_len = -cur_len
+            length_list.append(cur_len)
+        length_list = length_list * self.repeats
+        return length_list
+
+    def __len__(self):
+        return len(self.text_data) * self.repeats
+
+    def real_len(self):
+        return len(self.text_data)
+
+    def decode_mask(self, object_masks, ori_height, ori_width):
+        binary_masks = []
+        for object_mask in object_masks:
+            binary_mask = np.zeros((ori_height, ori_width), dtype=np.uint8)
+            for seg in object_mask:
+                rles = mask.frPyObjects([seg], ori_height, ori_width)
+                m = mask.decode(rles)
+                m = m.astype(np.uint8)
+                binary_mask += m.squeeze()
+
+            binary_masks.append(binary_mask)
+        if len(binary_masks) == 0:
+            return None
+        masks = np.stack(binary_masks, axis=0)
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+        masks = torch.from_numpy(masks)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio, self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks
+
+    def __getitem__(self, index):
+        index = index % self.real_len()
+        data_dict = copy.deepcopy(self.json_data[index])
+        data_dict.update(self.text_data[index])
+
+        if data_dict.get('image', None) is not None:
+            image_file = data_dict['image']
+            image_path = os.path.join(self.image_folder, image_file)
+            if not os.path.exists(image_path) and "VG" in self.image_folder:
+                image_path =  os.path.join(self.image_folder + "_2", image_file)
+            image = Image.open(image_path).convert('RGB')
+            ori_width, ori_height = image.size
+            if self.pad_image_to_square:
+                image = expand2square(
+                    image,
+                    tuple(
+                        int(x * 255) for x in self.image_processor.image_mean))
+
+            image = self.image_processor.preprocess(
+                image, return_tensors='pt')['pixel_values'][0]
+            data_dict['pixel_values'] = image
+
+            # process and get masks
+            points = data_dict["points"]
+            points = np.array(points)
+            if self.pad_image_to_square:
+                points = expand2square_points(points, height=ori_height, width=ori_width)
+                points[:, 0] = points[:, 0] / max(ori_height, ori_width) * self.image_w
+                points[:, 1] = points[:, 1] / max(ori_height, ori_width) * self.image_h
+            else:
+                points[:, 0] = points[:, 0] / ori_width * self.image_w
+                points[:, 1] = points[:, 1] / ori_height * self.image_h
+            data_dict['points'] = torch.from_numpy(points)
+            if data_dict['points'] is None:
+                return self.__getitem__(0)
+            data_dict['masks'] = None
+            data_dict['regions'] = None
+        else:
+            if hasattr(self.image_processor, 'crop_size'):
+                crop_size = self.image_processor.crop_size
+            else:
+                crop_size = self.image_processor.size
+            data_dict['pixel_values'] = torch.zeros(3, crop_size['height'],
+                                                    crop_size['width'])
+            data_dict['masks'] = None
+            data_dict['regions'] = None
+            data_dict['points'] = None
+        return data_dict
+
+class MDPVPointBriefCaptionDataset(MDPVPointDetailedCaptionDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=32,
+                 debug=False,
+                 repeats=1):
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            num_proc=num_proc,
+            debug=debug,
+            repeats=repeats
+            )

omg_llava/dataset/ReferringSegDataset.py

@@ -0,0 +1,380 @@
+import logging
+import os
+import torch
+from datasets import Dataset as HFDataset
+from datasets import DatasetDict, load_from_disk
+from mmengine import print_log
+from mmengine.config import Config, ConfigDict
+from PIL import Image
+from torch.utils.data import Dataset
+from pycocotools import mask
+import numpy as np
+import torch.nn.functional as F
+
+from xtuner.registry import BUILDER
+from omg_llava.dataset.utils import expand2square, expand2square_mask
+from xtuner.dataset.huggingface import process_hf_dataset
+from omg_llava.dataset.utils.refcoco_refer import REFER
+import copy
+
+class RefcocoReferringSegDataset(Dataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=8,
+                 debug=False,
+                 repeats=1,):
+        self._set_attribute()
+        self.tokenizer = tokenizer
+        assert offline_processed_text_folder or (data_path and tokenizer)
+        self.debug = debug
+        if offline_processed_text_folder and data_path:
+            print_log(
+                'Both `offline_processed_text_folder` and '
+                '`data_path` are set, and we load dataset from'
+                '`offline_processed_text_folder` '
+                f'({offline_processed_text_folder})',
+                logger='current',
+                level=logging.WARNING)
+
+        if offline_processed_text_folder is not None:
+            raise NotImplementedError
+        else:
+            json_datas = self.json_file_preprocess(data_path)
+            self.json_datas = json_datas
+            json_datas = self.only_get_hf_map_infos()
+            json_data = DatasetDict({'train': HFDataset.from_list(json_datas)})
+            self.text_data = process_hf_dataset(
+                dataset=json_data,
+                tokenizer=tokenizer,
+                max_length=max_length,
+                dataset_map_fn=dataset_map_fn,
+                template_map_fn=template_map_fn,
+                split='train',
+                max_dataset_length=max_dataset_length,
+                remove_unused_columns=False,
+                pack_to_max_length=False,
+                with_image_token=True,
+                map_num_proc=num_proc,  # because limited mem
+            )
+
+        self.image_folder = image_folder
+        size = image_processor.crop_size
+        if isinstance(size, int):
+            self.image_h, self.image_w = size, size
+        else:
+            self.image_w, self.image_h = size
+
+        if isinstance(image_processor, dict) or isinstance(
+                image_processor, Config) or isinstance(image_processor,
+                                                       ConfigDict):
+            self.image_processor = BUILDER.build(image_processor)
+        else:
+            self.image_processor = image_processor
+        self.pad_image_to_square = pad_image_to_square
+        self.down_ratio = 1
+        self.repeats = repeats
+
+    def _set_attribute(self):
+        self.splitBy = "unc"
+        self.dataset_name = 'refcoco'
+
+    def only_get_hf_map_infos(self):
+        ret = []
+        for json_data in self.json_datas:
+            ret.append({'sampled_sents': json_data['selected_labels']})
+        return ret
+
+    def __len__(self):
+        return len(self.text_data) * self.repeats
+
+    @property
+    def modality_length(self):
+        length_list = []
+        for data_dict in self.text_data:
+            cur_len = len(data_dict['input_ids'])
+            if data_dict.get('image', None) is None:
+                cur_len = -cur_len
+            length_list.append(cur_len)
+        length_list = length_list * self.repeats
+        return length_list
+
+    def real_len(self):
+        return len(self.text_data)
+
+    def json_file_preprocess(self, data_path):
+        splitBy = self.splitBy
+        dataset_name = self.dataset_name
+        refer_api = REFER(data_path, dataset_name, splitBy)
+        ref_ids_train = refer_api.getRefIds(split='train')
+        images_ids_train = refer_api.getImgIds(ref_ids=ref_ids_train)
+        refs_train = refer_api.loadRefs(ref_ids=ref_ids_train)
+        self.img2refs = self.create_img_to_refs_mapping(refs_train)
+
+        image_infos = []
+        loaded_images = refer_api.loadImgs(image_ids=images_ids_train)
+        for item in loaded_images:
+            item = item.copy()
+            image_infos.append(item)
+
+        self.annotations = refer_api.Anns
+
+        refs = [self.img2refs[image_info['id']] for image_info in image_infos]
+
+        ret = []
+        for image_info, ref in zip(image_infos, refs):
+            if len(ref) == 0:
+                continue
+
+            sents = []
+            ann_ids = []
+            for _ref in ref:
+                for sent in _ref["sentences"]:
+                    text = sent["sent"]
+                    sents.append(text)
+                    ann_ids.append(_ref["ann_id"])
+            if len(sents) >= 3:
+                sampled_inds = np.random.choice(
+                    list(range(len(sents))), size=3, replace=False
+                )
+            else:
+                sampled_inds = list(range(len(sents)))
+            sampled_sents = np.vectorize(sents.__getitem__)(sampled_inds).tolist()
+            sampled_ann_ids = [ann_ids[ind] for ind in sampled_inds]
+            selected_labels = sampled_sents
+            ret.append(
+                {'image_info': image_info,
+                 'sampled_ann_id': sampled_ann_ids,
+                 'selected_labels': selected_labels,
+                 'image': image_info['file_name']
+                 }
+            )
+        if self.debug:
+            return ret[:1000]
+        return ret
+
+    def create_img_to_refs_mapping(self, refs_train):
+        img2refs = {}
+        for ref in refs_train:
+            img2refs[ref["image_id"]] = img2refs.get(ref["image_id"], []) + [ref, ]
+        return img2refs
+
+    def decode_mask(self, annotations_ids, image_info):
+        flag = False
+        masks = []
+
+        for ann_id in annotations_ids:
+            if isinstance(ann_id, list):
+                flag = True
+                if -1 in ann_id:
+                    assert len(ann_id) == 1
+                    m = np.zeros((image_info["height"], image_info["width"])).astype(
+                        np.uint8
+                    )
+                else:
+                    m_final = np.zeros(
+                        (image_info["height"], image_info["width"])
+                    ).astype(np.uint8)
+                    for ann_id_i in ann_id:
+                        ann = self.annotations[ann_id_i]
+
+                        if len(ann["segmentation"]) == 0:
+                            m = np.zeros(
+                                (image_info["height"], image_info["width"])
+                            ).astype(np.uint8)
+                        else:
+                            if type(ann["segmentation"][0]) == list:  # polygon
+                                rle = mask.frPyObjects(
+                                    ann["segmentation"], image_info["height"], image_info["width"], )
+                            else:
+                                rle = ann["segmentation"]
+                                for i in range(len(rle)):
+                                    if not isinstance(rle[i]["counts"], bytes):
+                                        rle[i]["counts"] = rle[i]["counts"].encode()
+                            m = mask.decode(rle)
+                            m = np.sum(
+                                m, axis=2
+                            )  # sometimes there are multiple binary map (corresponding to multiple segs)
+                            m = m.astype(np.uint8)  # convert to np.uint8
+                        m_final = m_final | m
+                    m = m_final
+                masks.append(m)
+                continue
+
+            ann = self.annotations[ann_id]
+
+            if len(ann["segmentation"]) == 0:
+                m = np.zeros((image_info["height"], image_info["width"])).astype(
+                    np.uint8
+                )
+                masks.append(m)
+                continue
+
+            if type(ann["segmentation"][0]) == list:  # polygon
+                rle = mask.frPyObjects(
+                    ann["segmentation"], image_info["height"], image_info["width"]
+                )
+            else:
+                rle = ann["segmentation"]
+                for i in range(len(rle)):
+                    if not isinstance(rle[i]["counts"], bytes):
+                        rle[i]["counts"] = rle[i]["counts"].encode()
+            m = mask.decode(rle)
+            m = np.sum(m, axis=2)  # sometimes there are multiple binary map (corresponding to multiple segs)
+            m = m.astype(np.uint8)  # convert to np.uint8
+            masks.append(m)
+        masks = np.stack(masks, axis=0)
+
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+        masks = torch.from_numpy(masks)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio,
+                                                        self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks
+
+    def __getitem__(self, index):
+        index = index % self.real_len()
+        data_dict = copy.deepcopy(self.text_data[index])
+        data_dict.update(self.json_datas[index])
+
+        if data_dict.get('image', None) is not None:
+            image_file = data_dict['image']
+            image_file = os.path.join(self.image_folder, image_file)
+            image = Image.open(image_file).convert('RGB')
+            ori_width, ori_height = image.size
+            if self.pad_image_to_square:
+                image = expand2square(
+                    image,
+                    tuple(
+                        int(x * 255) for x in self.image_processor.image_mean))
+            image = self.image_processor.preprocess(
+                image, return_tensors='pt')['pixel_values'][0]
+            data_dict['pixel_values'] = image
+
+            # process and get masks
+            masks = self.decode_mask(data_dict['sampled_ann_id'], data_dict['image_info'])
+            data_dict['masks'] = masks
+        else:
+            if hasattr(self.image_processor, 'crop_size'):
+                crop_size = self.image_processor.crop_size
+            else:
+                crop_size = self.image_processor.size
+            data_dict['pixel_values'] = torch.zeros(3, crop_size['height'],
+                                                    crop_size['width'])
+            data_dict['masks'] = None
+        return data_dict
+
+class Refcoco_plus_ReferringSegDataset(RefcocoReferringSegDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=8,
+                 debug=False,
+                 repeats=1,):
+
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            num_proc=num_proc,
+            debug=debug,
+            repeats=repeats,)
+
+    def _set_attribute(self):
+        self.splitBy = "unc"
+        self.dataset_name = 'refcoco+'
+
+class Refcocog_ReferringSegDataset(RefcocoReferringSegDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=8,
+                 debug=False,
+                 repeats=1,):
+
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            num_proc=num_proc,
+            debug=debug,
+            repeats=repeats,
+        )
+
+    def _set_attribute(self):
+        self.splitBy = "umd"
+        self.dataset_name = 'refcocog'
+
+class Refclef_ReferringSegDataset(RefcocoReferringSegDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=8,
+                 debug=False,
+                 repeats=1,):
+
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            num_proc=num_proc,
+            debug=debug,
+            repeats=repeats,
+        )
+
+    def _set_attribute(self):
+        self.splitBy = "unc"
+        self.dataset_name = 'refclef'

omg_llava/dataset/RegionCaptionDataset.py

@@ -0,0 +1,356 @@
+import json
+import logging
+import os
+import copy
+
+import torch
+from datasets import Dataset as HFDataset
+from datasets import DatasetDict, load_from_disk
+from mmengine import print_log
+from mmengine.config import Config, ConfigDict
+from PIL import Image, ImageDraw
+from torch.utils.data import Dataset
+from pycocotools import mask
+import numpy as np
+import torch.nn.functional as F
+
+from xtuner.registry import BUILDER
+from omg_llava.dataset.utils import expand2square, expand2square_mask
+from xtuner.dataset.huggingface import process_hf_dataset
+
+class OspreyRegionCaptionDataset(Dataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=32,
+                 debug=False,
+                 repeats=1):
+        super().__init__()
+
+        assert offline_processed_text_folder or (data_path and tokenizer)
+        self.debug = debug
+        if offline_processed_text_folder and data_path:
+            print_log(
+                'Both `offline_processed_text_folder` and '
+                '`data_path` are set, and we load dataset from'
+                '`offline_processed_text_folder` '
+                f'({offline_processed_text_folder})',
+                logger='current',
+                level=logging.WARNING)
+
+        if offline_processed_text_folder is not None:
+            raise NotImplementedError
+        else:
+            json_data = self.json_file_preprocess(data_path)
+            self.json_data = json_data
+            hf_json_data = self.filter_hf_require_infos(json_data)
+            hf_json_data = DatasetDict({'train': HFDataset.from_list(hf_json_data)})
+            self.text_data = process_hf_dataset(
+                dataset=hf_json_data,
+                tokenizer=tokenizer,
+                max_length=max_length,
+                dataset_map_fn=dataset_map_fn,
+                template_map_fn=template_map_fn,
+                split='train',
+                max_dataset_length=max_dataset_length,
+                remove_unused_columns=False,
+                pack_to_max_length=False,
+                with_image_token=True,
+                map_num_proc=num_proc,  # because limited mem
+            )
+
+        self.image_folder = image_folder
+        size = image_processor.crop_size
+        if isinstance(size, int):
+            self.image_h, self.image_w = size, size
+        else:
+            self.image_w, self.image_h = size
+
+        if isinstance(image_processor, dict) or isinstance(
+                image_processor, Config) or isinstance(image_processor,
+                                                       ConfigDict):
+            self.image_processor = BUILDER.build(image_processor)
+        else:
+            self.image_processor = image_processor
+        self.pad_image_to_square = pad_image_to_square
+        self.down_ratio = 1
+        self.repeats = repeats
+
+    def filter_hf_require_infos(self, dataset_infos):
+        ret = []
+        for dataset_info in dataset_infos:
+            description = dataset_info["description"]
+            image = dataset_info['file_name']
+            required_info = {'image': image, 'description': description}
+            ret.append(required_info)
+        return ret
+
+    def json_file_preprocess(self, data_path):
+        with open(data_path, 'r') as f:
+            json_file = json.load(f)
+
+        ret = []
+        for item in json_file:
+            if len(item["description"]) != len(item["annotation"]):
+                print("The number of description is not equal to seg !!!")
+            else:
+                ret.append(item)
+
+        if self.debug:
+            ret = ret[:10000]
+        return ret
+
+    @property
+    def modality_length(self):
+        length_list = []
+        for data_dict in self.text_data:
+            cur_len = len(data_dict['input_ids'])
+            if data_dict.get('image', None) is None:
+                cur_len = -cur_len
+            length_list.append(cur_len)
+        length_list = length_list * self.repeats
+        return length_list
+
+    def __len__(self):
+        return len(self.text_data) * self.repeats
+
+    def real_len(self):
+        return len(self.text_data)
+
+    def decode_mask(self, object_masks, ori_height, ori_width):
+        binary_masks = []
+        for object_mask in object_masks:
+            binary_mask = np.zeros((ori_height, ori_width), dtype=np.uint8)
+            for seg in object_mask:
+                rles = mask.frPyObjects([seg], ori_height, ori_width)
+                m = mask.decode(rles)
+                m = m.astype(np.uint8)
+                binary_mask += m.squeeze()
+
+            binary_masks.append(binary_mask)
+        if len(binary_masks) == 0:
+            return None
+        masks = np.stack(binary_masks, axis=0)
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+        masks = torch.from_numpy(masks)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio, self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks
+
+    def __getitem__(self, index):
+        index = index % self.real_len()
+        data_dict = copy.deepcopy(self.json_data[index])
+        data_dict.update(self.text_data[index])
+
+        if data_dict.get('image', None) is not None:
+            image_file = data_dict['image']
+            image = Image.open(os.path.join(self.image_folder,
+                                            image_file)).convert('RGB')
+            ori_width, ori_height = image.size
+            if self.pad_image_to_square:
+                image = expand2square(
+                    image,
+                    tuple(
+                        int(x * 255) for x in self.image_processor.image_mean))
+            image = self.image_processor.preprocess(
+                image, return_tensors='pt')['pixel_values'][0]
+            data_dict['pixel_values'] = image
+
+            # process and get masks
+            annotations = data_dict['annotation']
+            sampled_inds = data_dict['sampled_inds']
+            annotations = [annotations[idx]['segmentation'] for idx in sampled_inds]
+            data_dict['regions'] = self.decode_mask(annotations, ori_height=ori_height, ori_width=ori_width)
+
+            if data_dict['regions'] is None or len(data_dict['regions']) != len(sampled_inds):
+                print("Bad data item !!!")
+                return self.__getitem__(0)
+            seg_region_idx = data_dict['seg_region_idx']
+            if len(seg_region_idx) == 0:
+                data_dict['masks'] = None
+            else:
+                data_dict['masks'] = data_dict['regions'][seg_region_idx]
+        else:
+            if hasattr(self.image_processor, 'crop_size'):
+                crop_size = self.image_processor.crop_size
+            else:
+                crop_size = self.image_processor.size
+            data_dict['pixel_values'] = torch.zeros(3, crop_size['height'],
+                                                    crop_size['width'])
+            data_dict['masks'] = None
+            data_dict['regions'] = None
+        return data_dict
+
+class OspreyRegionConversationDataset(Dataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=32,
+                 debug=False,
+                 repeats=1):
+        super().__init__()
+
+        assert offline_processed_text_folder or (data_path and tokenizer)
+        self.debug = debug
+        if offline_processed_text_folder and data_path:
+            print_log(
+                'Both `offline_processed_text_folder` and '
+                '`data_path` are set, and we load dataset from'
+                '`offline_processed_text_folder` '
+                f'({offline_processed_text_folder})',
+                logger='current',
+                level=logging.WARNING)
+
+        if offline_processed_text_folder is not None:
+            raise NotImplementedError
+        else:
+            json_data = self.json_file_preprocess(data_path)
+            self.json_data = json_data
+            hf_json_data = self.filter_hf_require_infos(json_data)
+            hf_json_data = DatasetDict({'train': HFDataset.from_list(hf_json_data)})
+            self.text_data = process_hf_dataset(
+                dataset=hf_json_data,
+                tokenizer=tokenizer,
+                max_length=max_length,
+                dataset_map_fn=dataset_map_fn,
+                template_map_fn=template_map_fn,
+                split='train',
+                max_dataset_length=max_dataset_length,
+                remove_unused_columns=False,
+                pack_to_max_length=False,
+                with_image_token=True,
+                map_num_proc=num_proc,  # because limited mem
+            )
+
+        self.image_folder = image_folder
+        size = image_processor.crop_size
+        if isinstance(size, int):
+            self.image_h, self.image_w = size, size
+        else:
+            self.image_w, self.image_h = size
+
+        if isinstance(image_processor, dict) or isinstance(
+                image_processor, Config) or isinstance(image_processor,
+                                                       ConfigDict):
+            self.image_processor = BUILDER.build(image_processor)
+        else:
+            self.image_processor = image_processor
+        self.pad_image_to_square = pad_image_to_square
+        self.down_ratio = 1
+        self.repeats = repeats
+
+    def filter_hf_require_infos(self, dataset_infos):
+        ret = []
+        for dataset_info in dataset_infos:
+            conversations = dataset_info["conversations"]
+            image = dataset_info['file_name']
+            num_regions = len(dataset_info['annotation'])
+            required_info = {'image': image, 'conversations': conversations,
+                             'num_regions': num_regions}
+            ret.append(required_info)
+        return ret
+
+    def json_file_preprocess(self, data_path):
+        with open(data_path, 'r') as f:
+            json_file = json.load(f)
+
+        # filter
+        ret = []
+        for dataset_info in json_file:
+            if 'annotation' not in dataset_info or len(dataset_info['annotation']) == 0:
+                print("The annotation is not valid, filter out!!!")
+                continue
+            ret.append(dataset_info)
+
+        if self.debug:
+            ret = ret[:10000]
+        return ret
+
+    @property
+    def modality_length(self):
+        length_list = []
+        for data_dict in self.text_data:
+            cur_len = len(data_dict['input_ids'])
+            if data_dict.get('image', None) is None:
+                cur_len = -cur_len
+            length_list.append(cur_len)
+        length_list = length_list * self.repeats
+        return length_list
+
+    def __len__(self):
+        return len(self.text_data) * self.repeats
+
+    def real_len(self):
+        return len(self.text_data)
+
+    def decode_mask(self, object_masks, ori_height, ori_width):
+        binary_masks = []
+        for object_mask in object_masks:
+            binary_mask = np.zeros((ori_height, ori_width), dtype=np.uint8)
+            for seg in object_mask:
+                rles = mask.frPyObjects([seg], ori_height, ori_width)
+                m = mask.decode(rles)
+                m = m.astype(np.uint8)
+                binary_mask += m.squeeze()
+            binary_masks.append(binary_mask)
+        if len(binary_masks) == 0:
+            return None
+        masks = np.stack(binary_masks, axis=0)
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+        masks = torch.from_numpy(masks)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio, self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks
+
+    def __getitem__(self, index):
+        index = index % self.real_len()
+        data_dict = copy.deepcopy(self.json_data[index])
+        data_dict.update(self.text_data[index])
+
+        if data_dict.get('image', None) is not None:
+            image_file = data_dict['image']
+            image = Image.open(os.path.join(self.image_folder,
+                                            image_file)).convert('RGB')
+            ori_width, ori_height = image.size
+            if self.pad_image_to_square:
+                image = expand2square(
+                    image,
+                    tuple(
+                        int(x * 255) for x in self.image_processor.image_mean))
+            image = self.image_processor.preprocess(
+                image, return_tensors='pt')['pixel_values'][0]
+            data_dict['pixel_values'] = image
+
+            # process and get masks
+            annotations = data_dict['annotation']
+            annotations = [annotations[idx]['segmentation'] for idx in range(len(annotations))]
+            data_dict['regions'] = self.decode_mask(annotations, ori_height=ori_height, ori_width=ori_width)
+            if data_dict['regions'] is None:
+                return self.__getitem__(0)
+            data_dict['masks'] = None
+        else:
+            if hasattr(self.image_processor, 'crop_size'):
+                crop_size = self.image_processor.crop_size
+            else:
+                crop_size = self.image_processor.size
+            data_dict['pixel_values'] = torch.zeros(3, crop_size['height'],
+                                                    crop_size['width'])
+            data_dict['masks'] = None
+            data_dict['regions'] = None
+        return data_dict

omg_llava/dataset/SemanticSegDataset.py

@@ -0,0 +1,725 @@
+import random
+import glob
+import json
+import logging
+import os
+import torch
+from datasets import Dataset as HFDataset
+from datasets import DatasetDict, load_from_disk
+from mmengine import print_log
+from mmengine.config import Config, ConfigDict
+from PIL import Image
+from torch.utils.data import Dataset
+import numpy as np
+import torch.nn.functional as F
+from pycocotools.coco import COCO
+
+from xtuner.registry import BUILDER
+from omg_llava.dataset.utils import expand2square, expand2square_mask
+from xtuner.dataset.huggingface import process_hf_dataset
+from omg_llava.dataset.process_functions.semantic_seg_process import semantic_seg_conversations, semantic_seg_gcg_format_conversations
+import copy
+
+class SemanticSegDataset(Dataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=8,
+                 debug=False,
+                 repeats=1,
+                 gcg_format=False):
+        super().__init__()
+        self.tokenizer = tokenizer
+        assert offline_processed_text_folder or (data_path and tokenizer)
+        self.debug = debug
+        if offline_processed_text_folder and data_path:
+            print_log(
+                'Both `offline_processed_text_folder` and '
+                '`data_path` are set, and we load dataset from'
+                '`offline_processed_text_folder` '
+                f'({offline_processed_text_folder})',
+                logger='current',
+                level=logging.WARNING)
+
+        if offline_processed_text_folder is not None:
+            raise NotImplementedError
+        else:
+            self.image_label_datas = self.json_file_preprocess(data_path, image_folder)
+            if gcg_format:
+                conversations_datas = semantic_seg_gcg_format_conversations(self.classes)
+            else:
+                conversations_datas = semantic_seg_conversations(self.classes)
+            json_data = DatasetDict({'train': HFDataset.from_list(conversations_datas)})
+            self.text_data = process_hf_dataset(
+                dataset=json_data,
+                tokenizer=tokenizer,
+                max_length=max_length,
+                dataset_map_fn=dataset_map_fn,
+                template_map_fn=template_map_fn,
+                split='train',
+                max_dataset_length=max_dataset_length,
+                remove_unused_columns=False,
+                pack_to_max_length=False,
+                with_image_token=True,
+                map_num_proc=num_proc,  # because limited mem
+            )
+
+        self.clsid2convs = self.construct_cls2convs_dict()
+        self.image_folder = image_folder
+        size = image_processor.crop_size
+        if isinstance(size, int):
+            self.image_h, self.image_w = size, size
+        else:
+            self.image_w, self.image_h = size
+
+        if isinstance(image_processor, dict) or isinstance(
+                image_processor, Config) or isinstance(image_processor,
+                                                       ConfigDict):
+            self.image_processor = BUILDER.build(image_processor)
+        else:
+            self.image_processor = image_processor
+        self.pad_image_to_square = pad_image_to_square
+        self.down_ratio = 1
+        self.repeats = repeats
+
+    def construct_cls2convs_dict(self):
+        ret = {}
+        for conv_item in self.text_data:
+            cls_id = conv_item['class_id']
+            if cls_id in ret.keys():
+                ret[cls_id].append(conv_item)
+            else:
+                ret[cls_id] = [conv_item]
+        return ret
+
+    def json_file_preprocess(self, data_path, image_folder):
+        # ade20k
+        with open(data_path, 'r') as file:
+            ade20k_classes = json.load(file)
+        ade20k_image_dir = image_folder
+        ade20k_images = [os.path.join(ade20k_image_dir, img) for img in os.listdir(ade20k_image_dir) if
+                         img.endswith('.jpg')]
+        ade20k_labels = [img.replace(".jpg", ".png").replace("images", "annotations") for img in ade20k_images]
+        self.classes = np.array(ade20k_classes)
+
+        ret = []
+        for image, label in zip(ade20k_images, ade20k_labels):
+            ret.append({"image": image, "label": label})
+        if self.debug:
+            return ret[:1000]
+        return ret
+
+    def __len__(self):
+        return len(self.image_label_datas) * self.repeats
+
+    @property
+    def modality_length(self):
+        length_list = []
+        for data_dict in self.image_label_datas:
+            length_list.append(-100)
+        length_list = length_list * self.repeats
+        return length_list
+
+    def real_len(self):
+        return len(self.image_label_datas)
+
+    def decode_mask(self, label_path):
+        label = np.array(Image.open(label_path))
+
+        # ade 20k
+        label = np.where(label == 0, 255, label - 1)
+        unique_labels = [lbl for lbl in np.unique(label) if lbl != 255]
+        if not unique_labels:
+            return None, None
+
+        # only choose 1
+        selected_labels = np.random.choice(
+            unique_labels, 1, replace=False
+        )
+        label = torch.from_numpy(label).long()
+        masks = torch.stack([label == class_id for class_id in selected_labels], dim=0)
+
+        masks = masks.numpy()
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+
+        masks = torch.from_numpy(masks).to(torch.float32)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio,
+                                                        self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks, selected_labels[0]
+
+    def __getitem__(self, index):
+        index = index % self.real_len()
+        data_dict = copy.deepcopy(self.image_label_datas[index])
+
+        if data_dict.get('image', None) is not None:
+            image_file = data_dict['image']
+            image = Image.open(image_file).convert('RGB')
+            ori_width, ori_height = image.size
+            if self.pad_image_to_square:
+                image = expand2square(
+                    image,
+                    tuple(
+                        int(x * 255) for x in self.image_processor.image_mean))
+            image = self.image_processor.preprocess(
+                image, return_tensors='pt')['pixel_values'][0]
+            data_dict['pixel_values'] = image
+
+            # process and get masks
+            data_dict['masks'], class_id = self.decode_mask(data_dict['label'])
+            if class_id is None:
+                return self.__getitem__(0)
+            conv_datas = self.clsid2convs[class_id]
+            selected_idx = np.random.randint(0, len(conv_datas))
+            data_dict.update(conv_datas[selected_idx])
+        else:
+            if hasattr(self.image_processor, 'crop_size'):
+                crop_size = self.image_processor.crop_size
+            else:
+                crop_size = self.image_processor.size
+            data_dict['pixel_values'] = torch.zeros(3, crop_size['height'],
+                                                    crop_size['width'])
+            data_dict['masks'] = None
+        return data_dict
+
+class ADE20kSemanticSegDataset(SemanticSegDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=8,
+                 debug=False,
+                 repeats=1,
+                 gcg_format=False):
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            num_proc=num_proc,
+            debug=debug,
+            repeats=repeats,
+            gcg_format=gcg_format,
+        )
+
+class COCOStuffSemanticSegDataset(SemanticSegDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=8,
+                 debug=False,
+                 repeats=1,
+                 label_path=None,
+                 gcg_format=False,):
+        self.label_path = label_path
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            num_proc=num_proc,
+            debug=debug,
+            repeats=repeats,
+            gcg_format=gcg_format,
+        )
+        self.cocostuff_class2index = {c: i for i, c in enumerate(self.classes)}
+
+    def json_file_preprocess(self, data_path, image_folder):
+        # coco stuff
+        assert self.label_path is not None
+        with open(data_path, 'r') as file:
+            cocostuff_classes = [line.strip().split(": ")[-1] for line in file.readlines()[1:]]
+        coco_stuff_image_dir = image_folder
+        coco_stuff_label_dir = self.label_path
+        coco_stuff_labels = glob.glob(os.path.join(coco_stuff_label_dir, "*.png"))
+
+        coco_stuff_images = [label.replace(".png", ".jpg").replace(coco_stuff_label_dir, coco_stuff_image_dir)
+            for label in coco_stuff_labels]
+
+        self.classes = np.array(cocostuff_classes)
+
+        ret = []
+        for image, label in zip(coco_stuff_images, coco_stuff_labels):
+            ret.append({"image": image, "label": label})
+        if self.debug:
+            return ret[:1000]
+        return ret
+
+    def decode_mask(self, label_path):
+        label = np.array(Image.open(label_path))
+
+        # coco stuff
+        ignored_classes = [index for class_name, index in self.cocostuff_class2index.items() if
+                           "-" in class_name]
+        label = np.where(np.isin(label, ignored_classes), 255, label)
+
+        unique_labels = [lbl for lbl in np.unique(label) if lbl != 255]
+        if not unique_labels:
+            print("No valid label !!!")
+            return None, None
+
+        # only choose 1
+        selected_labels = np.random.choice(
+            unique_labels, 1, replace=False
+        )
+        label = torch.from_numpy(label).long()
+        masks = torch.stack([label == class_id for class_id in selected_labels], dim=0)
+
+        masks = masks.numpy()
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+
+        masks = torch.from_numpy(masks).to(torch.float32)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio,
+                                                        self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks, selected_labels[0]
+
+class MapillarySemanticSegDataset(SemanticSegDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=8,
+                 debug=False,
+                 repeats=1,
+                 label_path=None,
+                 gcg_format=False,):
+        self.label_path = label_path
+        super().__init__(
+            image_folder=image_folder,
+            image_processor=image_processor,
+            data_path=data_path,
+            tokenizer=tokenizer,
+            offline_processed_text_folder=offline_processed_text_folder,
+            max_dataset_length=max_dataset_length,
+            dataset_map_fn=dataset_map_fn,
+            template_map_fn=template_map_fn,
+            max_length=max_length,
+            pad_image_to_square=pad_image_to_square,
+            num_proc=num_proc,
+            debug=debug,
+            repeats=repeats,
+            gcg_format=gcg_format,
+        )
+
+    def json_file_preprocess(self, data_path, image_folder):
+        assert self.label_path is not None
+        # mapillary
+        with open(data_path, 'r') as file:
+            mapillary_classes = json.load(file)["labels"]
+        mapillary_classes = [cls["readable"].lower() for cls in mapillary_classes]
+
+        mapillary_labels = sorted(
+            glob.glob(os.path.join(self.label_path, "*.png")))
+        mapillary_images = [
+            label.replace(".png", ".jpg").replace(self.label_path, image_folder)
+            for label in mapillary_labels]
+
+        self.classes = np.array(mapillary_classes)
+
+        ret = []
+        for image, label in zip(mapillary_images, mapillary_labels):
+            ret.append({"image": image, "label": label})
+        if self.debug:
+            return ret[:1000]
+        return ret
+
+    def decode_mask(self, label_path):
+        label = np.array(Image.open(label_path))
+
+        ignored_classes = [index for index, class_name in enumerate(self.classes) if
+                           "-" in class_name or '(' in class_name or
+                           'unlabeled' in class_name]
+        label = np.where(np.isin(label, ignored_classes), 255, label)
+        unique_labels = [lbl for lbl in np.unique(label) if lbl != 255]
+        if not unique_labels:
+            print("No valid label !!!")
+            return None, None
+        # only choose 1
+        selected_labels = np.random.choice(
+            unique_labels, 1, replace=False
+        )
+        label = torch.from_numpy(label).long()
+        masks = torch.stack([label == class_id for class_id in selected_labels], dim=0)
+
+        masks = masks.numpy()
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+
+        masks = torch.from_numpy(masks).to(torch.float32)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio,
+                                                        self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks, selected_labels[0]
+
+class PascalPartSemanticSegDataset(Dataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=8,
+                 debug=False,
+                 repeats=1):
+        super().__init__()
+        self.tokenizer = tokenizer
+        assert offline_processed_text_folder or (data_path and tokenizer)
+        self.debug = debug
+        if offline_processed_text_folder and data_path:
+            print_log(
+                'Both `offline_processed_text_folder` and '
+                '`data_path` are set, and we load dataset from'
+                '`offline_processed_text_folder` '
+                f'({offline_processed_text_folder})',
+                logger='current',
+                level=logging.WARNING)
+
+        if offline_processed_text_folder is not None:
+            raise NotImplementedError
+        else:
+            json_datas = self.json_file_preprocess(data_path)
+            json_data = DatasetDict({'train': HFDataset.from_list(json_datas)})
+            self.text_data = process_hf_dataset(
+                dataset=json_data,
+                tokenizer=tokenizer,
+                max_length=max_length,
+                dataset_map_fn=dataset_map_fn,
+                template_map_fn=template_map_fn,
+                split='train',
+                max_dataset_length=max_dataset_length,
+                remove_unused_columns=False,
+                pack_to_max_length=False,
+                with_image_token=True,
+                map_num_proc=num_proc,  # because limited mem
+            )
+
+        self.image_folder = image_folder
+        size = image_processor.crop_size
+        if isinstance(size, int):
+            self.image_h, self.image_w = size, size
+        else:
+            self.image_w, self.image_h = size
+
+        if isinstance(image_processor, dict) or isinstance(
+                image_processor, Config) or isinstance(image_processor,
+                                                       ConfigDict):
+            self.image_processor = BUILDER.build(image_processor)
+        else:
+            self.image_processor = image_processor
+        self.pad_image_to_square = pad_image_to_square
+        self.down_ratio = 1
+        self.repeats = repeats
+
+    def json_file_preprocess(self, data_path):
+        pascal_part_api = COCO(data_path)
+        all_classes = pascal_part_api.loadCats(pascal_part_api.getCatIds())
+        class_map_pascal_part = {}
+        for cat in all_classes:
+            cat_main, cat_part = cat["name"].strip().split(":")
+            name = (cat_main, cat_part)
+            class_map_pascal_part[cat["id"]] = name
+        img_ids = pascal_part_api.getImgIds()
+        self.classes = class_map_pascal_part
+        self.coco_api = pascal_part_api
+
+        img_infos = [self.coco_api.loadImgs([img_id])[0] for img_id in img_ids]
+        valid_img_infos = []
+        for img_info in img_infos:
+            annotation_ids = self.coco_api.getAnnIds(imgIds=img_info["id"])
+            annotations = self.coco_api.loadAnns(annotation_ids)
+            if not annotations:
+                continue
+
+            # sampled to max number as 5
+            sampled_anns = np.random.choice(annotations, 5, replace=False) if len(
+                annotations
+            ) >= 5 else annotations
+
+            selected_labels = []
+            for ann in sampled_anns:
+                category_id = ann["category_id"]
+                sampled_cls = self.classes[category_id]
+                if isinstance(sampled_cls, tuple):
+                    obj, part = sampled_cls
+                    name = f"{obj} {part}" if random.random() < 0.5 else f"the {part} of the {obj}"
+                else:
+                    name = sampled_cls
+                selected_labels.append(name)
+
+            img_info.update({"annotations": sampled_anns,
+                             "selected_labels": selected_labels})
+            valid_img_infos.append(img_info)
+
+        if self.debug:
+            return valid_img_infos[:1000]
+        return valid_img_infos
+
+    def __len__(self):
+        return len(self.text_data) * self.repeats
+
+    @property
+    def modality_length(self):
+        length_list = []
+        for data_dict in self.text_data:
+            cur_len = len(data_dict['input_ids'])
+            if data_dict.get('image', None) is None:
+                cur_len = -cur_len
+            length_list.append(cur_len)
+        length_list = length_list * self.repeats
+        return length_list
+
+    def real_len(self):
+        return len(self.text_data)
+
+    def decode_mask(self, annotations):
+
+        try:
+            masks = [self.coco_api.annToMask(ann) for ann in annotations]
+        except Exception as e:
+            print(f"Error generating mask: {e}")
+            return None
+
+        masks = np.stack(masks, axis=0)
+        if self.pad_image_to_square:
+            masks = expand2square_mask(masks)
+        masks = torch.from_numpy(masks)
+        masks = F.interpolate(masks.unsqueeze(0), size=(self.image_h // self.down_ratio,
+                                                        self.image_w // self.down_ratio), mode='nearest').squeeze(0)
+        return masks
+
+    def __getitem__(self, index):
+        index = index % self.real_len()
+        data_dict = copy.deepcopy(self.text_data[index])
+
+        if data_dict.get('image', None) is not None:
+            image_file = data_dict['image']
+            image_file = os.path.join(self.image_folder, image_file)
+            image = Image.open(image_file).convert('RGB')
+            ori_width, ori_height = image.size
+            if self.pad_image_to_square:
+                image = expand2square(
+                    image,
+                    tuple(
+                        int(x * 255) for x in self.image_processor.image_mean))
+            image = self.image_processor.preprocess(
+                image, return_tensors='pt')['pixel_values'][0]
+            data_dict['pixel_values'] = image
+
+            # process and get masks
+            data_dict['masks'] = self.decode_mask(data_dict['annotations'])
+            if data_dict['masks'] is None:
+                return self.__getitem__(0)
+        else:
+            if hasattr(self.image_processor, 'crop_size'):
+                crop_size = self.image_processor.crop_size
+            else:
+                crop_size = self.image_processor.size
+            data_dict['pixel_values'] = torch.zeros(3, crop_size['height'],
+                                                    crop_size['width'])
+            data_dict['masks'] = None
+        return data_dict
+
+class PacoSemanticSegDataset(PascalPartSemanticSegDataset):
+    def __init__(self,
+                 image_folder,
+                 image_processor,
+                 data_path=None,
+                 tokenizer=None,
+                 offline_processed_text_folder=None,
+                 max_dataset_length=None,
+                 dataset_map_fn=None,
+                 template_map_fn=None,
+                 max_length=2048,
+                 pad_image_to_square=False,
+                 num_proc=8,
+                 debug=False,
+                 repeats=1,):
+        self.tokenizer = tokenizer
+        assert offline_processed_text_folder or (data_path and tokenizer)
+        self.debug = debug
+        if offline_processed_text_folder and data_path:
+            print_log(
+                'Both `offline_processed_text_folder` and '
+                '`data_path` are set, and we load dataset from'
+                '`offline_processed_text_folder` '
+                f'({offline_processed_text_folder})',
+                logger='current',
+                level=logging.WARNING)
+
+        if offline_processed_text_folder is not None:
+            raise NotImplementedError
+        else:
+            json_datas = self.json_file_preprocess(data_path)
+            self.json_datas = json_datas
+            json_datas = self.only_get_hf_map_infos()
+            json_data = DatasetDict({'train': HFDataset.from_list(json_datas)})
+            self.text_data = process_hf_dataset(
+                dataset=json_data,
+                tokenizer=tokenizer,
+                max_length=max_length,
+                dataset_map_fn=dataset_map_fn,
+                template_map_fn=template_map_fn,
+                split='train',
+                max_dataset_length=max_dataset_length,
+                remove_unused_columns=False,
+                pack_to_max_length=False,
+                with_image_token=True,
+                map_num_proc=num_proc,  # because limited mem
+            )
+
+        self.image_folder = image_folder
+        size = image_processor.crop_size
+        if isinstance(size, int):
+            self.image_h, self.image_w = size, size
+        else:
+            self.image_w, self.image_h = size
+
+        if isinstance(image_processor, dict) or isinstance(
+                image_processor, Config) or isinstance(image_processor,
+                                                       ConfigDict):
+            self.image_processor = BUILDER.build(image_processor)
+        else:
+            self.image_processor = image_processor
+        self.pad_image_to_square = pad_image_to_square
+        self.down_ratio = 1
+        self.repeats = repeats
+
+    def only_get_hf_map_infos(self):
+        ret = []
+        for json_data in self.json_datas:
+            ret.append({'file_name': json_data['file_name'],
+                        'selected_labels': json_data['selected_labels']})
+        return ret
+
+    def json_file_preprocess(self, data_path):
+        paco_api = COCO(data_path)
+        all_classes = paco_api.loadCats(paco_api.getCatIds())
+        class_map_paco = {}
+        for cat in all_classes:
+            cat_split = cat["name"].strip().split(":")
+            if len(cat_split) == 1:
+                name = cat_split[0].split("_(")[0]
+            else:
+                assert len(cat_split) == 2
+                obj, part = cat_split
+                obj = obj.split("_(")[0]
+                part = part.split("_(")[0]
+                name = (obj, part)
+            class_map_paco[cat["id"]] = name
+
+        img_ids = paco_api.getImgIds()
+        self.classes = class_map_paco
+        self.coco_api = paco_api
+
+        img_infos = [self.coco_api.loadImgs([img_id])[0] for img_id in img_ids]
+        valid_img_infos = []
+        for img_info in img_infos:
+            annotation_ids = self.coco_api.getAnnIds(imgIds=img_info["id"])
+            annotations = self.coco_api.loadAnns(annotation_ids)
+            if not annotations:
+                continue
+
+            # sampled to max number as 5
+            sampled_anns = np.random.choice(annotations, 5, replace=False) if len(
+                annotations
+            ) >= 5 else annotations
+
+            selected_labels = []
+            for ann in sampled_anns:
+                category_id = ann["category_id"]
+                sampled_cls = self.classes[category_id]
+                if isinstance(sampled_cls, tuple):
+                    obj, part = sampled_cls
+                    name = f"{obj} {part}" if random.random() < 0.5 else f"the {part} of the {obj}"
+                else:
+                    name = sampled_cls
+                selected_labels.append(name)
+
+            img_info.update({"annotations": sampled_anns,
+                             "selected_labels": selected_labels})
+            valid_img_infos.append(img_info)
+
+        if self.debug:
+            return valid_img_infos[:1000]
+        return valid_img_infos
+
+    def __getitem__(self, index):
+        index = index % self.real_len()
+        data_dict = copy.deepcopy(self.text_data[index])
+        data_dict.update(self.json_datas[index])
+
+        if data_dict.get('image', None) is not None:
+            image_file = data_dict['image']
+            image_file = os.path.join(self.image_folder, image_file)
+            image = Image.open(image_file).convert('RGB')
+            ori_width, ori_height = image.size
+            if self.pad_image_to_square:
+                image = expand2square(
+                    image,
+                    tuple(
+                        int(x * 255) for x in self.image_processor.image_mean))
+            image = self.image_processor.preprocess(
+                image, return_tensors='pt')['pixel_values'][0]
+            data_dict['pixel_values'] = image
+
+            # process and get masks
+            data_dict['masks'] = self.decode_mask(data_dict['annotations'])
+            if data_dict['masks'] is None:
+                return self.__getitem__(0)
+        else:
+            if hasattr(self.image_processor, 'crop_size'):
+                crop_size = self.image_processor.crop_size
+            else:
+                crop_size = self.image_processor.size
+            data_dict['pixel_values'] = torch.zeros(3, crop_size['height'],
+                                                    crop_size['width'])
+            data_dict['masks'] = None
+        return data_dict
+

omg_llava/dataset/__init__.py

@@ -0,0 +1,29 @@
+from .CombineDataset import CombineDataset
+from .GCGDataset import RefCOCOgGCGDataset, OpenPsgGCGDataset, GranDfGCGDataset, FlickrGCGDataset
+from .SemanticSegDataset import SemanticSegDataset, ADE20kSemanticSegDataset,\
+    COCOStuffSemanticSegDataset,MapillarySemanticSegDataset, PascalPartSemanticSegDataset,\
+    PacoSemanticSegDataset
+from .MDPVPointsDataset import MDPVPointDetailedCaptionDataset, MDPVPointBriefCaptionDataset
+from .ReferringSegDataset import RefcocoReferringSegDataset, Refcoco_plus_ReferringSegDataset,\
+    Refcocog_ReferringSegDataset, Refclef_ReferringSegDataset
+from .RegionCaptionDataset import OspreyRegionCaptionDataset, OspreyRegionConversationDataset
+from .LlavaDataset import LLaVADataset
+from .DecoupledGCGDataset import DecoupledRefCOCOgGCGDataset, DecoupledOpenPsgGCGDataset,\
+    DecoupledGranDfGCGDataset, DecoupledFlickrGCGDataset
+
+
+from .process_functions import glamm_openpsg_map_fn, glamm_refcocog_map_fn,\
+    glamm_granf_map_fn, glamm_flickr_map_fn,\
+    semantic_seg_map_fn, pascal_part_map_fn,\
+    semantic_seg_gcg_format_map_fn, pascal_part_gcg_format_map_fn,\
+    referring_seg_map_fn, referring_seg_gcg_format_map_fn,\
+    osprey_region_caption_map_fn, osprey_region_caption_gcg_format_map_fn,\
+    osprey_region_conversation_map_fn,\
+    mdpv_points_map_fn
+
+from .process_functions import glamm_refcocog_decoupled_given_objects_map_fn, glamm_refcocog_decoupled_given_description_map_fn,\
+    glamm_granf_decoupled_given_description_map_fn, glamm_granf_decoupled_given_objects_map_fn,\
+    glamm_flickr_decoupled_given_description_map_fn, glamm_flickr_decoupled_given_objects_map_fn,\
+    glamm_openpsg_decoupled_given_objects_map_fn, glamm_openpsg_decoupled_given_description_map_fn
+
+from .collect_fns import omg_llava_collate_fn