├── LICENSE
├── README.md
├── app.py
├── checkpoints
    └── PUT_CKPT_HERE
├── docs
    ├── eval.md
    ├── inference.md
    └── models.md
├── eval
    ├── eval_egoschema.py
    ├── eval_mlvu.py
    ├── eval_mvbench.py
    └── eval_videomme.py
├── examples
    ├── video1.mp4
    ├── video2.mp4
    └── video3.mp4
├── longvu
    ├── __init__.py
    ├── apply_delta.py
    ├── builder.py
    ├── cambrian_arch.py
    ├── consolidate.py
    ├── constants.py
    ├── conversation.py
    ├── file_io.py
    ├── language_model
    │   ├── cambrian_llama.py
    │   └── cambrian_qwen.py
    ├── make_delta.py
    ├── mm_datautils.py
    ├── mm_trainer.py
    ├── mm_utils.py
    ├── multimodal_encoder
    │   ├── base_encoder.py
    │   ├── builder.py
    │   ├── dino_encoder.py
    │   ├── drop.py
    │   ├── image.py
    │   ├── logging.py
    │   ├── loss.py
    │   ├── registry.py
    │   ├── siglip_encoder.py
    │   └── utils.py
    ├── multimodal_projector
    │   └── builder.py
    ├── train.py
    ├── utils.py
    └── vision_sampler.py
├── requirements.txt
└── scripts
    ├── train_image_llama3_2.sh
    ├── train_image_qwen.sh
    ├── train_video_llama3_2.sh
    └── train_video_qwen.sh


/LICENSE:
--------------------------------------------------------------------------------
  1 |                                  Apache License
  2 |                            Version 2.0, January 2004
  3 |                         http://www.apache.org/licenses/
  4 | 
  5 |    TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
  6 | 
  7 |    1. Definitions.
  8 | 
  9 |       "License" shall mean the terms and conditions for use, reproduction,
 10 |       and distribution as defined by Sections 1 through 9 of this document.
 11 | 
 12 |       "Licensor" shall mean the copyright owner or entity authorized by
 13 |       the copyright owner that is granting the License.
 14 | 
 15 |       "Legal Entity" shall mean the union of the acting entity and all
 16 |       other entities that control, are controlled by, or are under common
 17 |       control with that entity. For the purposes of this definition,
 18 |       "control" means (i) the power, direct or indirect, to cause the
 19 |       direction or management of such entity, whether by contract or
 20 |       otherwise, or (ii) ownership of fifty percent (50%) or more of the
 21 |       outstanding shares, or (iii) beneficial ownership of such entity.
 22 | 
 23 |       "You" (or "Your") shall mean an individual or Legal Entity
 24 |       exercising permissions granted by this License.
 25 | 
 26 |       "Source" form shall mean the preferred form for making modifications,
 27 |       including but not limited to software source code, documentation
 28 |       source, and configuration files.
 29 | 
 30 |       "Object" form shall mean any form resulting from mechanical
 31 |       transformation or translation of a Source form, including but
 32 |       not limited to compiled object code, generated documentation,
 33 |       and conversions to other media types.
 34 | 
 35 |       "Work" shall mean the work of authorship, whether in Source or
 36 |       Object form, made available under the License, as indicated by a
 37 |       copyright notice that is included in or attached to the work
 38 |       (an example is provided in the Appendix below).
 39 | 
 40 |       "Derivative Works" shall mean any work, whether in Source or Object
 41 |       form, that is based on (or derived from) the Work and for which the
 42 |       editorial revisions, annotations, elaborations, or other modifications
 43 |       represent, as a whole, an original work of authorship. For the purposes
 44 |       of this License, Derivative Works shall not include works that remain
 45 |       separable from, or merely link (or bind by name) to the interfaces of,
 46 |       the Work and Derivative Works thereof.
 47 | 
 48 |       "Contribution" shall mean any work of authorship, including
 49 |       the original version of the Work and any modifications or additions
 50 |       to that Work or Derivative Works thereof, that is intentionally
 51 |       submitted to Licensor for inclusion in the Work by the copyright owner
 52 |       or by an individual or Legal Entity authorized to submit on behalf of
 53 |       the copyright owner. For the purposes of this definition, "submitted"
 54 |       means any form of electronic, verbal, or written communication sent
 55 |       to the Licensor or its representatives, including but not limited to
 56 |       communication on electronic mailing lists, source code control systems,
 57 |       and issue tracking systems that are managed by, or on behalf of, the
 58 |       Licensor for the purpose of discussing and improving the Work, but
 59 |       excluding communication that is conspicuously marked or otherwise
 60 |       designated in writing by the copyright owner as "Not a Contribution."
 61 | 
 62 |       "Contributor" shall mean Licensor and any individual or Legal Entity
 63 |       on behalf of whom a Contribution has been received by Licensor and
 64 |       subsequently incorporated within the Work.
 65 | 
 66 |    2. Grant of Copyright License. Subject to the terms and conditions of
 67 |       this License, each Contributor hereby grants to You a perpetual,
 68 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 69 |       copyright license to reproduce, prepare Derivative Works of,
 70 |       publicly display, publicly perform, sublicense, and distribute the
 71 |       Work and such Derivative Works in Source or Object form.
 72 | 
 73 |    3. Grant of Patent License. Subject to the terms and conditions of
 74 |       this License, each Contributor hereby grants to You a perpetual,
 75 |       worldwide, non-exclusive, no-charge, royalty-free, irrevocable
 76 |       (except as stated in this section) patent license to make, have made,
 77 |       use, offer to sell, sell, import, and otherwise transfer the Work,
 78 |       where such license applies only to those patent claims licensable
 79 |       by such Contributor that are necessarily infringed by their
 80 |       Contribution(s) alone or by combination of their Contribution(s)
 81 |       with the Work to which such Contribution(s) was submitted. If You
 82 |       institute patent litigation against any entity (including a
 83 |       cross-claim or counterclaim in a lawsuit) alleging that the Work
 84 |       or a Contribution incorporated within the Work constitutes direct
 85 |       or contributory patent infringement, then any patent licenses
 86 |       granted to You under this License for that Work shall terminate
 87 |       as of the date such litigation is filed.
 88 | 
 89 |    4. Redistribution. You may reproduce and distribute copies of the
 90 |       Work or Derivative Works thereof in any medium, with or without
 91 |       modifications, and in Source or Object form, provided that You
 92 |       meet the following conditions:
 93 | 
 94 |       (a) You must give any other recipients of the Work or
 95 |           Derivative Works a copy of this License; and
 96 | 
 97 |       (b) You must cause any modified files to carry prominent notices
 98 |           stating that You changed the files; and
 99 | 
100 |       (c) You must retain, in the Source form of any Derivative Works
101 |           that You distribute, all copyright, patent, trademark, and
102 |           attribution notices from the Source form of the Work,
103 |           excluding those notices that do not pertain to any part of
104 |           the Derivative Works; and
105 | 
106 |       (d) If the Work includes a "NOTICE" text file as part of its
107 |           distribution, then any Derivative Works that You distribute must
108 |           include a readable copy of the attribution notices contained
109 |           within such NOTICE file, excluding those notices that do not
110 |           pertain to any part of the Derivative Works, in at least one
111 |           of the following places: within a NOTICE text file distributed
112 |           as part of the Derivative Works; within the Source form or
113 |           documentation, if provided along with the Derivative Works; or,
114 |           within a display generated by the Derivative Works, if and
115 |           wherever such third-party notices normally appear. The contents
116 |           of the NOTICE file are for informational purposes only and
117 |           do not modify the License. You may add Your own attribution
118 |           notices within Derivative Works that You distribute, alongside
119 |           or as an addendum to the NOTICE text from the Work, provided
120 |           that such additional attribution notices cannot be construed
121 |           as modifying the License.
122 | 
123 |       You may add Your own copyright statement to Your modifications and
124 |       may provide additional or different license terms and conditions
125 |       for use, reproduction, or distribution of Your modifications, or
126 |       for any such Derivative Works as a whole, provided Your use,
127 |       reproduction, and distribution of the Work otherwise complies with
128 |       the conditions stated in this License.
129 | 
130 |    5. Submission of Contributions. Unless You explicitly state otherwise,
131 |       any Contribution intentionally submitted for inclusion in the Work
132 |       by You to the Licensor shall be under the terms and conditions of
133 |       this License, without any additional terms or conditions.
134 |       Notwithstanding the above, nothing herein shall supersede or modify
135 |       the terms of any separate license agreement you may have executed
136 |       with Licensor regarding such Contributions.
137 | 
138 |    6. Trademarks. This License does not grant permission to use the trade
139 |       names, trademarks, service marks, or product names of the Licensor,
140 |       except as required for reasonable and customary use in describing the
141 |       origin of the Work and reproducing the content of the NOTICE file.
142 | 
143 |    7. Disclaimer of Warranty. Unless required by applicable law or
144 |       agreed to in writing, Licensor provides the Work (and each
145 |       Contributor provides its Contributions) on an "AS IS" BASIS,
146 |       WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147 |       implied, including, without limitation, any warranties or conditions
148 |       of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149 |       PARTICULAR PURPOSE. You are solely responsible for determining the
150 |       appropriateness of using or redistributing the Work and assume any
151 |       risks associated with Your exercise of permissions under this License.
152 | 
153 |    8. Limitation of Liability. In no event and under no legal theory,
154 |       whether in tort (including negligence), contract, or otherwise,
155 |       unless required by applicable law (such as deliberate and grossly
156 |       negligent acts) or agreed to in writing, shall any Contributor be
157 |       liable to You for damages, including any direct, indirect, special,
158 |       incidental, or consequential damages of any character arising as a
159 |       result of this License or out of the use or inability to use the
160 |       Work (including but not limited to damages for loss of goodwill,
161 |       work stoppage, computer failure or malfunction, or any and all
162 |       other commercial damages or losses), even if such Contributor
163 |       has been advised of the possibility of such damages.
164 | 
165 |    9. Accepting Warranty or Additional Liability. While redistributing
166 |       the Work or Derivative Works thereof, You may choose to offer,
167 |       and charge a fee for, acceptance of support, warranty, indemnity,
168 |       or other liability obligations and/or rights consistent with this
169 |       License. However, in accepting such obligations, You may act only
170 |       on Your own behalf and on Your sole responsibility, not on behalf
171 |       of any other Contributor, and only if You agree to indemnify,
172 |       defend, and hold each Contributor harmless for any liability
173 |       incurred by, or claims asserted against, such Contributor by reason
174 |       of your accepting any such warranty or additional liability.
175 | 
176 |    END OF TERMS AND CONDITIONS
177 | 
178 |    APPENDIX: How to apply the Apache License to your work.
179 | 
180 |       To apply the Apache License to your work, attach the following
181 |       boilerplate notice, with the fields enclosed by brackets "[]"
182 |       replaced with your own identifying information. (Don't include
183 |       the brackets!)  The text should be enclosed in the appropriate
184 |       comment syntax for the file format. We also recommend that a
185 |       file or class name and description of purpose be included on the
186 |       same "printed page" as the copyright notice for easier
187 |       identification within third-party archives.
188 | 
189 |    Copyright [yyyy] [name of copyright owner]
190 | 
191 |    Licensed under the Apache License, Version 2.0 (the "License");
192 |    you may not use this file except in compliance with the License.
193 |    You may obtain a copy of the License at
194 | 
195 |        http://www.apache.org/licenses/LICENSE-2.0
196 | 
197 |    Unless required by applicable law or agreed to in writing, software
198 |    distributed under the License is distributed on an "AS IS" BASIS,
199 |    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200 |    See the License for the specific language governing permissions and
201 |    limitations under the License.
202 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | # LongVU
  2 | 
  3 | > **LongVU: Spatiotemporal Adaptive Compression for Long Video-Language Understanding [ICML 2025]**
  4 | >
  5 | > <a href='https://arxiv.org/abs/2410.17434'><img src='https://img.shields.io/badge/arXiv-paper-red'></a> <a href='https://vision-cair.github.io/LongVU'><img src='https://img.shields.io/badge/project-LongVU-blue'></a> <a href='https://huggingface.co/spaces/Vision-CAIR/LongVU'><img src='https://img.shields.io/badge/demo-space-blue'></a> <a href='https://huggingface.co/collections/Vision-CAIR/longvu-67181d2debabfc1eb050c21d'><img src='https://img.shields.io/badge/model-checkpoints-yellow'></a> 
  6 | 
  7 | <div align="center">
  8 |     <a href='https://vision-cair.github.io/LongVU'><img src="https://longvu.s3.amazonaws.com/assets/demo.gif" alt="Demo GIF" style="width: 100%; max-width: 650px;"></a>
  9 | </div>
 10 | 
 11 | ## :rocket: Quick Start
 12 | 
 13 | Try our model on [HF 🤗 Demo](https://huggingface.co/spaces/Vision-CAIR/LongVU)
 14 | 
 15 | Or demploy from local
 16 | 
 17 | ### Linux
 18 | 
 19 | ```
 20 | git clone https://github.com/Vision-CAIR/LongVU
 21 | cd LongVU
 22 | conda create -n longvu python=3.10
 23 | conda activate longvu
 24 | pip install -r requirements.txt
 25 | ```
 26 | 
 27 | Download our checkpoints and put it under `./checkpoints`
 28 | 
 29 | | Modality | LongVU_Qwen2_7B | LongVU_Llama3_2_3B |
 30 | :--------------------------:| :--------------------------:|:--------------------------:
 31 | | Image | [Download](https://huggingface.co/Vision-CAIR/LongVU_Qwen2_7B_img) | [Download](https://huggingface.co/Vision-CAIR/LongVU_Llama3_2_3B_img) |
 32 | | Video | [Download](https://huggingface.co/Vision-CAIR/LongVU_Qwen2_7B) | [Download](https://huggingface.co/Vision-CAIR/LongVU_Llama3_2_3B) |
 33 | 
 34 | Run demo `python app.py` locally with minimum 40G GPU.
 35 | 
 36 | <details>
 37 |   <summary>Click for quick inference code</summary>
 38 |     
 39 | ```python
 40 | import numpy as np
 41 | import torch
 42 | from longvu.builder import load_pretrained_model
 43 | from longvu.constants import (
 44 |     DEFAULT_IMAGE_TOKEN,
 45 |     IMAGE_TOKEN_INDEX,
 46 | )
 47 | from longvu.conversation import conv_templates, SeparatorStyle
 48 | from longvu.mm_datautils import (
 49 |     KeywordsStoppingCriteria,
 50 |     process_images,
 51 |     tokenizer_image_token,
 52 | )
 53 | from decord import cpu, VideoReader
 54 | 
 55 | tokenizer, model, image_processor, context_len = load_pretrained_model(
 56 |     "./checkpoints/longvu_qwen", None, "cambrian_qwen",
 57 | )
 58 | 
 59 | model.eval()
 60 | video_path = "./examples/video1.mp4"
 61 | qs = "Describe this video in detail"
 62 | 
 63 | vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
 64 | fps = float(vr.get_avg_fps())
 65 | frame_indices = np.array([i for i in range(0, len(vr), round(fps),)])
 66 | video = []
 67 | for frame_index in frame_indices:
 68 |     img = vr[frame_index].asnumpy()
 69 |     video.append(img)
 70 | video = np.stack(video)
 71 | image_sizes = [video[0].shape[:2]]
 72 | video = process_images(video, image_processor, model.config)
 73 | video = [item.unsqueeze(0) for item in video]
 74 | 
 75 | qs = DEFAULT_IMAGE_TOKEN + "\n" + qs
 76 | conv = conv_templates["qwen"].copy()
 77 | conv.append_message(conv.roles[0], qs)
 78 | conv.append_message(conv.roles[1], None)
 79 | prompt = conv.get_prompt()
 80 | 
 81 | input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(model.device)
 82 | stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
 83 | keywords = [stop_str]
 84 | stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
 85 | with torch.inference_mode():
 86 |     output_ids = model.generate(
 87 |         input_ids,
 88 |         images=video,
 89 |         image_sizes=image_sizes,
 90 |         do_sample=False,
 91 |         temperature=0.2,
 92 |         max_new_tokens=128,
 93 |         use_cache=True,
 94 |         stopping_criteria=[stopping_criteria],
 95 |     )
 96 | pred = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
 97 | ```
 98 |     
 99 | </details>
100 | 
101 | ### Windows
102 | 
103 | Thanks the detailed instruction [(here)](https://github.com/Vision-CAIR/LongVU/issues/6) from @ipeevski for developing on Windows system with 24GB VRAM.
104 | 
105 | ## Training
106 | 
107 | ### Dataset
108 | 
109 | + image-text stage: [LLaVA-OneVision-Single](https://huggingface.co/datasets/lmms-lab/LLaVA-OneVision-Data)
110 | + video-text stage: [VideoChat2-IT](https://huggingface.co/datasets/OpenGVLab/VideoChat2-IT)
111 | 
112 | ### Scripts
113 | 
114 | Experiments are run on 64 H100-96G
115 | 
116 | Download [image_json_file](https://huggingface.co/datasets/shenxq/OneVision/blob/main/onevision.json) and [video_json_file](https://huggingface.co/datasets/shenxq/VideoChat2/blob/main/train_video_data.json).
117 | 
118 | We also provide [row_video_data](https://huggingface.co/datasets/shenxq/VideoChat2) for easy downloading.
119 | 
120 | Modify the PATH_TO_JSON and PATH_TO_FOLDER arguments in the training scripts to your save folder.
121 | 
122 | ```
123 | PATH_TO_JSON=""
124 | PATH_TO_FOLDER=""
125 | ```
126 | Training your own model
127 | ```
128 | # image sft
129 | sh scripts/train_image_qwen.sh
130 | sh scripts/train_image_llama3_2.sh
131 | ```
132 | 
133 | Modify PREV_STAGE_CHECKPOINT in the training scripts to your first stage model path
134 | 
135 | Change `image_token_len` and `query_num_list` in `config.json` to 144
136 | 
137 | ```
138 | # video sft
139 | sh scripts/train_video_qwen.sh
140 | sh scripts/train_video_llama3_2.sh
141 | ```
142 | 
143 | ## Evaluation
144 | 
145 | See detailed evaluation code in [eval.md](https://github.com/Vision-CAIR/LongVU/blob/main/docs/eval.md)
146 | 
147 | ## Acknowledgement
148 | 
149 | + The model architecture of LongVU follows [LLaVA](https://github.com/haotian-liu/LLaVA) and [Cambrian](https://github.com/cambrian-mllm/cambrian)
150 | + We base [Qwen2](https://huggingface.co/Qwen/Qwen2-7B-Instruct) and [Llama3.2](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) as the language backbone
151 | + We use [SigLIP](https://huggingface.co/google/siglip-so400m-patch14-384) and [DINOv2](https://huggingface.co/facebook/dinov2-giant) as the vision encoder
152 | 
153 | ## Citation
154 | 
155 | ```
156 | @article{shen2024longvu,
157 |   author ={Shen, Xiaoqian and Xiong, Yunyang and Zhao, Changsheng and Wu, Lemeng and Chen, Jun and Zhu, Chenchen and Liu, Zechun and Xiao, Fanyi and Varadarajan, Balakrishnan and Bordes, Florian and Liu, Zhuang and Xu, Hu and J. Kim, Hyunwoo and Soran, Bilge and Krishnamoorthi, Raghuraman and Elhoseiny, Mohamed and Chandra, Vikas},
158 |   title = {LongVU: Spatiotemporal Adaptive Compression for Long Video-Language Understanding},
159 |   journal = {arXiv preprint arXiv:2410.17434},
160 |   year = {2024},
161 | }
162 | ```
163 | 


--------------------------------------------------------------------------------
/app.py:
--------------------------------------------------------------------------------
  1 | # import spaces
  2 | 
  3 | import os
  4 | import re
  5 | import traceback
  6 | 
  7 | import torch
  8 | import gradio as gr
  9 | 
 10 | import sys
 11 | 
 12 | import numpy as np
 13 | 
 14 | from longvu.builder import load_pretrained_model
 15 | from longvu.constants import (
 16 |     DEFAULT_IMAGE_TOKEN,
 17 |     IMAGE_TOKEN_INDEX,
 18 | )
 19 | from longvu.conversation import conv_templates, SeparatorStyle
 20 | from longvu.mm_datautils import (
 21 |     KeywordsStoppingCriteria,
 22 |     process_images,
 23 |     tokenizer_image_token,
 24 | )
 25 | from decord import cpu, VideoReader
 26 | 
 27 | 
 28 | title_markdown = """
 29 | <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
 30 |   <div>
 31 |     <h1 >LongVU: Spatiotemporal Adaptive Compression for Long Video-Language Understanding</h1>
 32 |   </div>
 33 | </div>
 34 | <div align="center">
 35 |     <div style="display:flex; gap: 0.25rem; margin-top: 10px;" align="center">
 36 |         <a href='https://vision-cair.github.io/LongVU/'><img src='https://img.shields.io/badge/Project-LongVU-blue'></a>
 37 |         <a href='https://huggingface.co/Vision-CAIR/LongVU_Qwen2_7B'><img src='https://img.shields.io/badge/model-checkpoints-yellow'></a>
 38 |     </div>
 39 | </div>
 40 | """
 41 | 
 42 | block_css = """
 43 | #buttons button {
 44 |     min-width: min(120px,100%);
 45 |     color: #9C276A
 46 | }
 47 | """
 48 | 
 49 | plum_color = gr.themes.colors.Color(
 50 |     name='plum',
 51 |     c50='#F8E4EF',
 52 |     c100='#E9D0DE',
 53 |     c200='#DABCCD',
 54 |     c300='#CBA8BC',
 55 |     c400='#BC94AB',
 56 |     c500='#AD809A',
 57 |     c600='#9E6C89',
 58 |     c700='#8F5878',
 59 |     c800='#804467',
 60 |     c900='#713056',
 61 |     c950='#662647',
 62 | )
 63 | 
 64 | 
 65 | class Chat:
 66 | 
 67 |     def __init__(self):
 68 |         self.version = "qwen"
 69 |         model_name = "cambrian_qwen"
 70 |         model_path = "./checkpoints/longvu_qwen"
 71 |         device = "cuda:7"
 72 | 
 73 |         self.tokenizer, self.model, self.processor, _ = load_pretrained_model(model_path,  None, model_name, device=device)
 74 |         self.model.eval()
 75 | 
 76 |     def remove_after_last_dot(self, s):
 77 |         last_dot_index = s.rfind('.')
 78 |         if last_dot_index == -1:
 79 |             return s
 80 |         return s[:last_dot_index + 1]
 81 |     
 82 |     # @spaces.GPU(duration=120)
 83 |     @torch.inference_mode()
 84 |     def generate(self, data: list, message, temperature, top_p, max_output_tokens):
 85 |         # TODO: support multiple turns of conversation.
 86 |         assert len(data) == 1
 87 | 
 88 |         tensor, image_sizes, modal = data[0]
 89 | 
 90 |         conv = conv_templates[self.version].copy()
 91 | 
 92 |         if isinstance(message, str):
 93 |             conv.append_message("user", DEFAULT_IMAGE_TOKEN + '\n' + message)
 94 |         elif isinstance(message, list):
 95 |             if DEFAULT_IMAGE_TOKEN not in message[0]['content']:
 96 |                 message[0]['content'] = DEFAULT_IMAGE_TOKEN + '\n' + message[0]['content']
 97 |             for mes in message:
 98 |                 conv.append_message(mes["role"], mes["content"])
 99 | 
100 |         conv.append_message("assistant", None)
101 |         
102 |         prompt = conv.get_prompt()
103 | 
104 |         input_ids = (
105 |             tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt")
106 |             .unsqueeze(0)
107 |             .to(self.model.device)
108 |         )
109 | 
110 |         if "llama3" in self.version:
111 |             input_ids = input_ids[0][1:].unsqueeze(0)  # remove bos
112 |         
113 |         stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
114 |         keywords = [stop_str]
115 |         stopping_criteria = KeywordsStoppingCriteria(keywords, self.tokenizer, input_ids)
116 |         with torch.inference_mode():
117 |             output_ids = self.model.generate(
118 |                 input_ids,
119 |                 images=tensor,
120 |                 image_sizes=image_sizes,
121 |                 do_sample=True,
122 |                 temperature=temperature,
123 |                 max_new_tokens=max_output_tokens,
124 |                 use_cache=True,
125 |                 top_p=top_p,
126 |                 stopping_criteria=[stopping_criteria],
127 |             )
128 |         
129 |         pred = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
130 |         
131 |         return self.remove_after_last_dot(pred)
132 | 
133 | 
134 | # @spaces.GPU(duration=120)
135 | def generate(image, video, message, chatbot, textbox_in, temperature, top_p, max_output_tokens, dtype=torch.float16):
136 |     if textbox_in is None:
137 |         raise gr.Error("Chat messages cannot be empty")
138 |         return (
139 |             gr.update(value=image, interactive=True),
140 |             gr.update(value=video, interactive=True),
141 |             message,
142 |             chatbot,
143 |             None,
144 |         )
145 |     
146 |     data = []
147 | 
148 |     processor = handler.processor
149 |     try:
150 |         if image is not None:
151 |             data.append((processor['image'](image).to(handler.model.device, dtype=dtype), None, '<image>'))
152 |         elif video is not None:
153 |             vr = VideoReader(video, ctx=cpu(0), num_threads=1)
154 |             fps = float(vr.get_avg_fps())
155 |             frame_indices = np.array(
156 |                 [
157 |                     i
158 |                     for i in range(
159 |                         0,
160 |                         len(vr),
161 |                         round(fps),
162 |                     )
163 |                 ]
164 |             )
165 |             video_tensor = []
166 |             for frame_index in frame_indices:
167 |                 img = vr[frame_index].asnumpy()
168 |                 video_tensor.append(img)
169 |             video_tensor = np.stack(video_tensor)
170 |             image_sizes = [video_tensor[0].shape[:2]]
171 |             video_tensor = process_images(video_tensor, processor, handler.model.config)
172 |             video_tensor = [item.unsqueeze(0).to(handler.model.device, dtype=dtype) for item in video_tensor]
173 |             data.append((video_tensor, image_sizes, '<video>'))
174 |         elif image is None and video is None:
175 |             data.append((None, None, '<text>'))
176 |         else:
177 |             raise NotImplementedError("Not support image and video at the same time")
178 |     except Exception as e:
179 |         traceback.print_exc()
180 |         return gr.update(value=None, interactive=True), gr.update(value=None, interactive=True), message, chatbot, None
181 | 
182 |     assert len(message) % 2 == 0, "The message should be a pair of user and system message."
183 | 
184 |     show_images = ""
185 |     if image is not None:
186 |         show_images += f'<img src="./file={image}" style="display: inline-block;width: 250px;max-height: 400px;">'
187 |     if video is not None:
188 |         show_images += f'<video controls playsinline width="300" style="display: inline-block;"  src="./file={video}"></video>'
189 | 
190 |     one_turn_chat = [textbox_in, None]
191 | 
192 |     # 1. first run case
193 |     if len(chatbot) == 0:
194 |         one_turn_chat[0] += "\n" + show_images
195 |     # 2. not first run case
196 |     else:
197 |         # scanning the last image or video
198 |         length = len(chatbot)
199 |         for i in range(length - 1, -1, -1):
200 |             previous_image = re.findall(r'<img src="./file=(.+?)"', chatbot[i][0])
201 |             previous_video = re.findall(r'<video controls playsinline width="500" style="display: inline-block;"  src="./file=(.+?)"', chatbot[i][0])
202 | 
203 |             if len(previous_image) > 0:
204 |                 previous_image = previous_image[-1]
205 |                 # 2.1 new image append or pure text input will start a new conversation
206 |                 if (video is not None) or (image is not None and os.path.basename(previous_image) != os.path.basename(image)):
207 |                     message.clear()
208 |                     one_turn_chat[0] += "\n" + show_images
209 |                 break
210 |             elif len(previous_video) > 0:
211 |                 previous_video = previous_video[-1]
212 |                 # 2.2 new video append or pure text input will start a new conversation
213 |                 if image is not None or (video is not None and os.path.basename(previous_video) != os.path.basename(video)):
214 |                     message.clear()
215 |                     one_turn_chat[0] += "\n" + show_images
216 |                 break
217 | 
218 |     message.append({'role': 'user', 'content': textbox_in})
219 |     text_en_out = handler.generate(data, message, temperature=temperature, top_p=top_p, max_output_tokens=max_output_tokens)
220 |     message.append({'role': 'assistant', 'content': text_en_out})
221 | 
222 |     one_turn_chat[1] = text_en_out
223 |     chatbot.append(one_turn_chat)
224 | 
225 |     return gr.update(value=image, interactive=True), gr.update(value=video, interactive=True), message, chatbot, None
226 | 
227 | 
228 | def regenerate(message, chatbot):
229 |     message.pop(-1), message.pop(-1)
230 |     chatbot.pop(-1)
231 |     return message, chatbot
232 | 
233 | 
234 | def clear_history(message, chatbot):
235 |     message.clear(), chatbot.clear()
236 |     return (gr.update(value=None, interactive=True),
237 |             gr.update(value=None, interactive=True),
238 |             message, chatbot,
239 |             gr.update(value=None, interactive=True))
240 | 
241 | handler = Chat()
242 | 
243 | textbox = gr.Textbox(show_label=False, placeholder="Enter text and press ENTER", container=False)
244 | 
245 | theme = gr.themes.Default(primary_hue=plum_color)
246 | # theme.update_color("primary", plum_color.c500)
247 | theme.set(slider_color="#9C276A")
248 | theme.set(block_title_text_color="#9C276A")
249 | theme.set(block_label_text_color="#9C276A")
250 | theme.set(button_primary_text_color="#9C276A")
251 | 
252 | with gr.Blocks(title='LongVU', theme=theme, css=block_css) as demo:
253 |     gr.Markdown(title_markdown)
254 |     message = gr.State([])
255 | 
256 |     with gr.Row():
257 |         with gr.Column(scale=3):
258 |             image = gr.State(None)
259 |             video = gr.Video(label="Input Video")
260 | 
261 |             with gr.Accordion("Parameters", open=True) as parameter_row:
262 | 
263 |                 temperature = gr.Slider(
264 |                     minimum=0.1,
265 |                     maximum=1.0,
266 |                     value=0.2,
267 |                     step=0.1,
268 |                     interactive=True,
269 |                     label="Temperature",
270 |                 )
271 | 
272 |                 top_p = gr.Slider(
273 |                         minimum=0.0,
274 |                         maximum=1.0,
275 |                         value=0.7,
276 |                         step=0.1,
277 |                         interactive=True,
278 |                         label="Top P",
279 |                 )
280 | 
281 |                 max_output_tokens = gr.Slider(
282 |                     minimum=64,
283 |                     maximum=512,
284 |                     value=128,
285 |                     step=64,
286 |                     interactive=True,
287 |                     label="Max output tokens",
288 |                 )
289 | 
290 |         with gr.Column(scale=7):
291 |             chatbot = gr.Chatbot(label="LongVU", bubble_full_width=True, height=420)
292 |             with gr.Row():
293 |                 with gr.Column(scale=8):
294 |                     textbox.render()
295 |                 with gr.Column(scale=1, min_width=50):
296 |                     submit_btn = gr.Button(value="Send", variant="primary", interactive=True)
297 |             with gr.Row(elem_id="buttons") as button_row:
298 |                 upvote_btn     = gr.Button(value="👍  Upvote", interactive=True)
299 |                 downvote_btn   = gr.Button(value="👎  Downvote", interactive=True)
300 |                 regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=True)
301 |                 clear_btn      = gr.Button(value="🗑️  Clear history", interactive=True)
302 | 
303 |     with gr.Row():
304 |         with gr.Column():
305 |             gr.Examples(
306 |                 examples=[
307 |                     [
308 |                         f"./examples/video3.mp4",
309 |                         "What is the moving direction of the yellow ball?",
310 |                     ],
311 |                     [
312 |                         f"./examples/video1.mp4",
313 |                         "Describe this video in detail.",
314 |                     ],
315 |                     [
316 |                         f"./examples/video2.mp4",
317 |                         "What is the name of the store?",
318 |                     ],
319 |                 ],
320 |                 inputs=[video, textbox],
321 |             )
322 | 
323 |     submit_btn.click(
324 |         generate, 
325 |         [image, video, message, chatbot, textbox, temperature, top_p, max_output_tokens],
326 |         [image, video, message, chatbot])
327 | 
328 |     regenerate_btn.click(
329 |         regenerate, 
330 |         [message, chatbot], 
331 |         [message, chatbot]).then(
332 |         generate, 
333 |         [image, video, message, chatbot, textbox, temperature, top_p, max_output_tokens], 
334 |         [image, video, message, chatbot, textbox])
335 | 
336 |     textbox.submit(
337 |         generate,
338 |         [
339 |             image,
340 |             video,
341 |             message,
342 |             chatbot,
343 |             textbox,
344 |             temperature,
345 |             top_p,
346 |             max_output_tokens,
347 |         ],
348 |         [image, video, message, chatbot, textbox],
349 |     )
350 |     
351 |     clear_btn.click(
352 |         clear_history, 
353 |         [message, chatbot],
354 |         [image, video, message, chatbot, textbox])
355 | 
356 | demo.launch()
357 | 


--------------------------------------------------------------------------------
/checkpoints/PUT_CKPT_HERE:
--------------------------------------------------------------------------------
1 | 
2 | 


--------------------------------------------------------------------------------
/docs/eval.md:
--------------------------------------------------------------------------------
 1 | # Evaluation
 2 | 
 3 | To reproduce the results, please download our checkpoints to `./checkpoints`
 4 | + LongVU_Qwen2-7B: [checkpoint]()
 5 | + LongVU_Llama3_2-3B: [checkpoint]()
 6 | 
 7 | ## MVBench
 8 | 
 9 | Download [MVBench](https://huggingface.co/datasets/OpenGVLab/MVBench) to `./data/MVBench`
10 | 
11 | ```
12 | python -m torch.distributed.launch --nproc-per-node=8 eval/eval_mvbench.py --data_path ./data/MVBench --version qwen --model_path ./checkpoints/longvu_qwen
13 | python -m torch.distributed.launch --nproc-per-node=8 eval/eval_mvbench.py --data_path ./data/MVBench --version llama3 --model_path ./checkpoints/longvu_llama3_2
14 | ```
15 | 
16 | ## EgoSchema
17 | 
18 | Download [EgoSchema](https://github.com/egoschema/EgoSchema) to `./data/EgoSchema`
19 | 
20 | ```
21 | python -m torch.distributed.launch --nproc-per-node=8 eval/eval_egoschema.py --data_path ./data/EgoSchema --version qwen --model_path ./checkpoints/longvu_qwen
22 | python -m torch.distributed.launch --nproc-per-node=8 eval/eval_egoschema.py --data_path ./data/EgoSchema --version llama3 --model_path ./checkpoints/longvu_llama3_2
23 | ```
24 | 
25 | Then submit the result file `.csv` to [Kaggle](https://www.kaggle.com/competitions/egoschema-public/submissions)
26 | 
27 | 
28 | ## MLVU
29 | 
30 | Download [MVLU](https://huggingface.co/datasets/MLVU/MVLU) to `./data/MVLU`
31 | 
32 | ```
33 | python -m torch.distributed.launch --nproc-per-node=8 eval/eval_mlvu.py --data_path ./data/MVLU --version qwen --model_path ./checkpoints/longvu_qwen
34 | python -m torch.distributed.launch --nproc-per-node=8 eval/eval_mlvu.py --data_path ./data/MVLU --version llama3 --model_path ./checkpoints/longvu_llama3_2
35 | ```
36 | 
37 | ## VideoMME
38 | 
39 | Download [VideoMME](https://github.com/BradyFU/Video-MME?tab=readme-ov-file#-dataset) to `./data/VideoMME`
40 | 
41 | ```
42 | python -m torch.distributed.launch --nproc-per-node=8 eval/eval_videomme.py --data_path ./data/VideoMME --version qwen --model_path ./checkpoints/longvu_qwen
43 | python -m torch.distributed.launch --nproc-per-node=8 eval/eval_videomme.py --data_path ./data/VideoMME --version llama3 --model_path ./checkpoints/longvu_llama3_2
44 | ```
45 | 


--------------------------------------------------------------------------------
/docs/inference.md:
--------------------------------------------------------------------------------
 1 | ## Inference Code
 2 | 
 3 | Below is a quick inference code:
 4 | 
 5 | ```python
 6 | import numpy as np
 7 | import torch
 8 | from longvu.builder import load_pretrained_model
 9 | from longvu.constants import (
10 |     DEFAULT_IMAGE_TOKEN,
11 |     IMAGE_TOKEN_INDEX,
12 | )
13 | from longvu.conversation import conv_templates, SeparatorStyle
14 | from longvu.mm_datautils import (
15 |     KeywordsStoppingCriteria,
16 |     process_images,
17 |     tokenizer_image_token,
18 | )
19 | from decord import cpu, VideoReader
20 | 
21 | tokenizer, model, image_processor, context_len = load_pretrained_model(
22 |     "./checkpoints/longvu_qwen", None, "cambrian_qwen",
23 | )
24 | 
25 | model.eval()
26 | video_path = "./examples/video1.mp4"
27 | qs = "Describe this video in detail"
28 | 
29 | vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
30 | fps = float(vr.get_avg_fps())
31 | frame_indices = np.array([i for i in range(0, len(vr), round(fps),)])
32 | video = []
33 | for frame_index in frame_indices:
34 |     img = vr[frame_index].asnumpy()
35 |     video.append(img)
36 | video = np.stack(video)
37 | image_sizes = [video[0].shape[:2]]
38 | video = process_images(video, image_processor, model.config)
39 | video = [item.unsqueeze(0) for item in video]
40 | 
41 | qs = DEFAULT_IMAGE_TOKEN + "\n" + qs
42 | conv = conv_templates["qwen"].copy()
43 | conv.append_message(conv.roles[0], qs)
44 | conv.append_message(conv.roles[1], None)
45 | prompt = conv.get_prompt()
46 | 
47 | input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(model.device)
48 | stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
49 | keywords = [stop_str]
50 | stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
51 | with torch.inference_mode():
52 |     output_ids = model.generate(
53 |         input_ids,
54 |         images=video,
55 |         image_sizes=image_sizes,
56 |         do_sample=False,
57 |         temperature=0.2,
58 |         max_new_tokens=128,
59 |         use_cache=True,
60 |         stopping_criteria=[stopping_criteria],
61 |     )
62 | pred = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()
63 | ```
64 | 


--------------------------------------------------------------------------------
/docs/models.md:
--------------------------------------------------------------------------------
1 | # Model Zoo
2 | 
3 | | Modality | LongVU_Qwen2_7B | LongVU_Llama3_2_3B | LongVU_Llama3_2_1B
4 | :--------------------------:| :--------------------------:|:--------------------------: |:--------------------------:
5 | | Image | [Download](https://huggingface.co/Vision-CAIR/LongVU_Qwen2_7B_img) | [Download](https://huggingface.co/Vision-CAIR/LongVU_Llama3_2_3B_img) | [Download](https://huggingface.co/Vision-CAIR/LongVU_Llama3_2_1B_img) |
6 | | Video | [Download](https://huggingface.co/Vision-CAIR/LongVU_Qwen2_7B) | [Download](https://huggingface.co/Vision-CAIR/LongVU_Llama3_2_3B) | [Download](https://huggingface.co/Vision-CAIR/LongVU_Llama3_2_1B) |
7 | 


--------------------------------------------------------------------------------
/eval/eval_egoschema.py:
--------------------------------------------------------------------------------
  1 | # Adopted from https://github.com/lm-sys/FastChat. Below is the original copyright:
  2 | # Adopted from tatsu-lab@stanford_alpaca. Below is the original copyright:
  3 | # Make it more memory efficient by monkey patching the LLaMA model with FlashAttn.
  4 | 
  5 | # pyre-strict
  6 | 
  7 | # Need to call this before importing transformers.
  8 | 
  9 | 
 10 | import datetime
 11 | import json
 12 | import os
 13 | import re
 14 | import shutil
 15 | import uuid
 16 | from itertools import chain
 17 | import argparse
 18 | 
 19 | import sys
 20 | sys.path.append('./')
 21 | import numpy as np
 22 | from PIL import Image
 23 | import pandas as pd
 24 | 
 25 | import torch
 26 | 
 27 | from longvu.builder import load_pretrained_model
 28 | from longvu.constants import (
 29 |     DEFAULT_IM_END_TOKEN,
 30 |     DEFAULT_IM_START_TOKEN,
 31 |     DEFAULT_IMAGE_TOKEN,
 32 |     IMAGE_TOKEN_INDEX,
 33 | )
 34 | from longvu.conversation import conv_templates, SeparatorStyle
 35 | from longvu.mm_datautils import (
 36 |     KeywordsStoppingCriteria,
 37 |     process_images,
 38 |     tokenizer_image_token,
 39 | )
 40 | 
 41 | from decord import cpu, VideoReader  # @manual=fbsource//third-party/pypi/decord:decord
 42 | from torch import distributed as dist
 43 | from tqdm import tqdm
 44 | 
 45 | from transformers.trainer_pt_utils import IterableDatasetShard
 46 | 
 47 | class EvalDataset(torch.utils.data.IterableDataset):
 48 |     """Dataset for supervised fine-tuning."""
 49 | 
 50 |     def __init__(
 51 |         self,
 52 |         data_path: str,
 53 |     ) -> None:
 54 |         super(EvalDataset, self).__init__()
 55 | 
 56 |         # pyre-fixme[4]: Attribute must be annotated.
 57 |         self.data = json.load(open(data_path, "r"))
 58 | 
 59 |     def __len__(self) -> int:
 60 |         return len(self.data)
 61 | 
 62 |     # pyre-fixme[3]: Return type must be annotated.
 63 |     def __iter__(self):
 64 |         return iter(self.data)
 65 | 
 66 |     # pyre-fixme[3]: Return type must be annotated.
 67 |     # pyre-fixme[2]: Parameter must be annotated.
 68 |     def __getitem__(self, i):
 69 |         return self.data[i]
 70 | 
 71 | 
 72 | def train(args) -> None:
 73 |     dist.init_process_group(backend="nccl", timeout=datetime.timedelta(hours=8))
 74 |     
 75 |     version = args.version
 76 |     model_name = args.model_name
 77 |     model_path = args.model_path
 78 | 
 79 |     # torch.distributed.barrier()
 80 |     tokenizer, model, image_processor, context_len = load_pretrained_model(
 81 |         model_path,  # pyre-fixme
 82 |         None,
 83 |         model_name,
 84 |         device_map=None,
 85 |     )
 86 |     model.get_model().config.drop_threshold = 0.8
 87 |     model.config.use_cache = True
 88 |     model.cuda()
 89 |     dataset = EvalDataset(
 90 |         data_path=args.data_path,
 91 |     )
 92 |     world_size = torch.distributed.get_world_size()
 93 |     world_rank = torch.distributed.get_rank()
 94 |     shard_dataset = IterableDatasetShard(
 95 |         dataset,
 96 |         batch_size=1,
 97 |         num_processes=world_size,
 98 |         process_index=world_rank,
 99 |     )
100 |     torch.distributed.barrier()
101 |     output = []
102 |     final_output = [None] * world_size
103 | 
104 |     video_formats = [".mp4", ".avi", ".mov", ".mkv"]
105 | 
106 |     for line in tqdm(shard_dataset):
107 |         video_name = line["google_drive_id"]
108 |         idx = line["q_uid"]
109 |         question = line["question"]
110 |         a0 = line["option 0"]
111 |         a1 = line["option 1"]
112 |         a2 = line["option 2"]
113 |         a3 = line["option 3"]
114 |         a4 = line["option 4"]
115 |         qs = f"Question: {question}\nOptions:\n(A) {a0}\n(B) {a1}\n(C) {a2}\n(D) {a3}\n(E) {a4}\nRespond with only the letter (A, B, C, D or E) of the correct option."
116 |         video_path = os.path.join(
117 |             args.data_path,
118 |             "good_clips_git",
119 |             f"{idx}.mp4",
120 |         )
121 |         for fmt in video_formats:
122 |             temp_path = os.path.join(
123 |                 args.data_path,
124 |                 "good_clips_git",
125 |                 f"{idx}{fmt}",
126 |             )
127 |             if os.path.exists(temp_path):
128 |                 video_path = temp_path
129 |                 break
130 | 
131 |         if os.path.exists(video_path):
132 |             vr = VideoReader(video_path, ctx=cpu(0))
133 |             fps = round(vr.get_avg_fps())
134 |             frame_idx = [
135 |                     i
136 |                     for i in range(0, len(vr), round(fps / 0.5))
137 |                 ]
138 |             if len(frame_idx) > 1000:
139 |                 frame_idx = [
140 |                     frame_idx[i]
141 |                     for i in range(0, len(frame_idx), len(frame_idx) // 1000)
142 |                 ]
143 |             video = vr.get_batch(frame_idx).asnumpy()
144 |             image_sizes = [video[0].shape[:2]]
145 |             video = process_images(video, image_processor, model.config)
146 |             video = [item.unsqueeze(0) for item in video]
147 |         else:
148 |             video = np.zeros((1, 1024, 1024, 3)).astype(np.uint8)
149 |             image_sizes = [(1024, 1024)]
150 |             video = process_images(video, image_processor, model.config)
151 | 
152 |         if getattr(model.config, "mm_use_im_start_end", False):
153 |             qs = (
154 |                 DEFAULT_IM_START_TOKEN
155 |                 + DEFAULT_IMAGE_TOKEN
156 |                 + DEFAULT_IM_END_TOKEN
157 |                 + "\n"
158 |                 + qs
159 |             )
160 |         else:
161 |             qs = DEFAULT_IMAGE_TOKEN + "\n" + qs
162 | 
163 |         conv = conv_templates[version].copy()
164 |         conv.append_message(conv.roles[0], qs)
165 |         conv.append_message(conv.roles[1], None)
166 |         prompt = conv.get_prompt()
167 | 
168 |         input_ids = (
169 |             tokenizer_image_token(
170 |                 prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
171 |             )
172 |             .unsqueeze(0)
173 |             .cuda()
174 |         )
175 | 
176 |         if "llama3" in version:
177 |             input_ids = input_ids[0][1:].unsqueeze(0)  # remove bos
178 | 
179 |         stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
180 |         keywords = [stop_str]
181 |         stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
182 |             
183 |         with torch.inference_mode():
184 |             output_ids = model.generate(
185 |                 input_ids,
186 |                 images=video,
187 |                 image_sizes=image_sizes,
188 |                 do_sample=False,
189 |                 temperature=0.0,
190 |                 max_new_tokens=5,  
191 |                 use_cache=True,
192 |                 stopping_criteria=[stopping_criteria],
193 |             )
194 |         if isinstance(output_ids, tuple):
195 |             output_ids = output_ids[0]
196 |         pred = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[
197 |             0
198 |         ].strip()
199 |         if pred.endswith(stop_str):
200 |             pred = pred[: -len(stop_str)]
201 |             pred = pred.strip()
202 |         pred = pred.replace("Answer", "")
203 | 
204 |         letters = ["A", "B", "C", "D", "E"]
205 |         pred_answer = re.findall("[\(\ ]*[A-E][\)\ ]*", pred)
206 | 
207 |         pred_answer = pred_answer[0].strip()
208 |         pred_answer = pred_answer.strip("()")
209 |         if pred_answer in letters:
210 |             pred_idx = letters.index(pred_answer)
211 |             pred = letters[pred_idx]
212 |         else:
213 |             print("pred_answer: ", pred_answer, " pred: ", pred, flush=True)
214 |             pred_idx = 2
215 |             pred = letters[pred_idx]
216 | 
217 |         ans_id = uuid.uuid4()
218 |         output.append(
219 |             {
220 |                 "idx": idx,
221 |                 "prompt": qs,
222 |                 "pred": pred_idx,
223 |                 "answer_id": str(ans_id),
224 |                 "model_id": model_name,
225 |                 "video_name": video_name,
226 |                 "metadata": {},
227 |             }
228 |         )
229 | 
230 |     dist.barrier()
231 |     dist.all_gather_object(
232 |         final_output,
233 |         output,
234 |     )
235 |     all_output = list(chain(*final_output))
236 |     global_rank = dist.get_rank()
237 |     if global_rank == 0:
238 |         if os.path.exists("/tmp/generated_text"):
239 |             shutil.rmtree("/tmp/generated_text")
240 |         os.mkdir("/tmp/generated_text")
241 | 
242 |         with open(
243 |             os.path.join("/tmp/generated_text", "outputs.json"),
244 |             "w",
245 |         ) as f:
246 |             json.dump(all_output, f)
247 | 
248 |         answers = json.load(
249 |                 open(
250 |                     os.path.join(
251 |                         args.data_path, "subset_answers.json"
252 |                     ),
253 |                     "r",
254 |                 )
255 |             )
256 | 
257 |         correct = 0
258 |         total = 0
259 |         for output in all_output:
260 |             if output["idx"] in answers:
261 |                 total += 1
262 |                 if int(output["pred"]) == int(answers[output["idx"]]):
263 |                     correct += 1
264 | 
265 |         print(f"Accuracy: {correct / total}")
266 |         result = {"acc": correct / total}
267 | 
268 |         with open(os.path.join("/tmp/generated_text", "result.json"), "w") as f:
269 |             json.dump(result, f)
270 | 
271 |         submission = []
272 |         q_uids = []
273 |         for output in all_output:
274 |             item = {}
275 |             item["q_uid"] = output["idx"]
276 |             item["answer"] = int(output["pred"])
277 |             if output["idx"] not in q_uids:
278 |                 q_uids.append(output["idx"])
279 |                 submission.append(item)
280 |         df = pd.DataFrame(submission)
281 |         df.to_csv(
282 |             "/tmp/generated_text/submission.csv",
283 |             index=False,
284 |             columns=["q_uid", "answer"],
285 |         )
286 | 
287 | if __name__ == "__main__":
288 |     parser = argparse.ArgumentParser()
289 | 
290 |     parser.add_argument('--model_path', default="./checkpoints/longvu_qwen")
291 |     parser.add_argument('--model_name', default="cambrian_qwen")
292 |     parser.add_argument('--version', default="qwen")
293 |     parser.add_argument('--local-rank', default=0)
294 |     parser.add_argument('--data_path', required=True)
295 |     args = parser.parse_args()
296 | 
297 |     if "llama3" in args.version:
298 |         args.model_name = "cambrian_llama3"
299 | 
300 |     train(args)
301 | 


--------------------------------------------------------------------------------
/eval/eval_mlvu.py:
--------------------------------------------------------------------------------
  1 | # Adopted from https://github.com/lm-sys/FastChat. Below is the original copyright:
  2 | # Adopted from tatsu-lab@stanford_alpaca. Below is the original copyright:
  3 | # Make it more memory efficient by monkey patching the LLaMA model with FlashAttn.
  4 | 
  5 | # pyre-strict
  6 | 
  7 | # Need to call this before importing transformers.
  8 | 
  9 | 
 10 | import datetime
 11 | import json
 12 | import os
 13 | import re
 14 | import shutil
 15 | import uuid
 16 | from itertools import chain
 17 | import argparse
 18 | 
 19 | import sys
 20 | sys.path.append('./')
 21 | import numpy as np
 22 | from PIL import Image
 23 | 
 24 | import torch
 25 | 
 26 | from longvu.builder import load_pretrained_model
 27 | from longvu.constants import (
 28 |     DEFAULT_IM_END_TOKEN,
 29 |     DEFAULT_IM_START_TOKEN,
 30 |     DEFAULT_IMAGE_TOKEN,
 31 |     IMAGE_TOKEN_INDEX,
 32 | )
 33 | from longvu.conversation import conv_templates, SeparatorStyle
 34 | from longvu.mm_datautils import (
 35 |     KeywordsStoppingCriteria,
 36 |     process_images,
 37 |     tokenizer_image_token,
 38 | )
 39 | 
 40 | from decord import cpu, VideoReader  # @manual=fbsource//third-party/pypi/decord:decord
 41 | from torch import distributed as dist
 42 | from tqdm import tqdm
 43 | 
 44 | from transformers.trainer_pt_utils import IterableDatasetShard
 45 | 
 46 | class EvalDataset(torch.utils.data.IterableDataset):
 47 |     """Dataset for supervised fine-tuning."""
 48 | 
 49 |     def __init__(
 50 |         self,
 51 |         data_path: str,
 52 |     ) -> None:
 53 |         super(EvalDataset, self).__init__()
 54 | 
 55 |         self.data_path = data_path
 56 | 
 57 |         data_list = {
 58 |             "count": ("json/4_count.json", f"video/4_count", "video"),
 59 |             "ego": ("json/3_ego.json", f"video/3_ego", "video"),
 60 |             "needle": ("json/2_needle.json", f"video/2_needle", "video"),
 61 |             "order": ("json/5_order.json", f"video/5_order", "video"),
 62 |             "plotQA": ("json/1_plotQA.json", f"video/1_plotQA", "video"),
 63 |             "anomaly_reco": (
 64 |                 "json/6_anomaly_reco.json",
 65 |                 f"video/6_anomaly_reco",
 66 |                 "video",
 67 |             ),
 68 |             "topic_reasoning": (
 69 |                 "json/7_topic_reasoning.json",
 70 |                 f"video/7_topic_reasoning",
 71 |                 "video",
 72 |             ),
 73 |         }
 74 | 
 75 |         list_data_dict = []
 76 |         for k, v in data_list.items():
 77 |             with open(os.path.join(data_path, v[0]), "r") as f:
 78 |                 json_data = json.load(f)
 79 |             for data in json_data:
 80 |                 question, answer = self.qa_template(data)
 81 |                 list_data_dict.append(
 82 |                     {
 83 |                         "task_type": k,
 84 |                         "video": os.path.join(self.data_path, v[1], data["video"]),
 85 |                         "video_name": data["video"],
 86 |                         "question": data["question"],
 87 |                         "prompt": question,
 88 |                         "answer": answer,
 89 |                     }
 90 |                 )
 91 | 
 92 |         # pyre-fixme[4]: Attribute must be annotated.
 93 |         self.data = list_data_dict
 94 | 
 95 |     def __len__(self) -> int:
 96 |         return len(self.data)
 97 | 
 98 |     # pyre-fixme[3]: Return type must be annotated.
 99 |     # pyre-fixme[2]: Parameter must be annotated.
100 |     def qa_template(self, data):
101 |         question = f"Question: {data['question']}\n"
102 |         question += "Options:\n"
103 |         answer = data["answer"]
104 |         answer_idx = -1
105 |         for idx, c in enumerate(data["candidates"]):
106 |             question += f"({chr(ord('A') + idx)}) {c}\n"
107 |             if c == answer:
108 |                 answer_idx = idx
109 |         question += (
110 |             "Respond with only the letter (A, B, C or D) of the correct option.\n"
111 |         )
112 |         question = question.rstrip()
113 |         answer = f"{chr(ord('A') + answer_idx)}"
114 |         return question, answer
115 | 
116 |     # pyre-fixme[3]: Return type must be annotated.
117 |     def __iter__(self):
118 |         return iter(self.data)
119 | 
120 |     # pyre-fixme[3]: Return type must be annotated.
121 |     # pyre-fixme[2]: Parameter must be annotated.
122 |     def __getitem__(self, i):
123 |         return self.data[i]
124 | 
125 | 
126 | def train(args) -> None:
127 |     dist.init_process_group(backend="nccl", timeout=datetime.timedelta(hours=8))
128 |     
129 |     version = args.version
130 |     model_name = args.model_name
131 |     model_path = args.model_path
132 | 
133 |     # torch.distributed.barrier()
134 |     tokenizer, model, image_processor, context_len = load_pretrained_model(
135 |         model_path,  # pyre-fixme
136 |         None,
137 |         model_name,
138 |         device_map=None,
139 |     )
140 |     model.get_model().config.dino_threshold = 0.82
141 |     model.get_model().config.drop_threshold = 0.77
142 |     model.config.use_cache = True
143 |     model.cuda()
144 |     dataset = EvalDataset(
145 |         data_path=args.data_path,
146 |     )
147 |     world_size = torch.distributed.get_world_size()
148 |     world_rank = torch.distributed.get_rank()
149 |     shard_dataset = IterableDatasetShard(
150 |         dataset,
151 |         batch_size=1,
152 |         num_processes=world_size,
153 |         process_index=world_rank,
154 |     )
155 |     torch.distributed.barrier()
156 |     output = []
157 |     final_output = [None] * world_size
158 | 
159 |     for line in tqdm(shard_dataset):
160 |         video_name = line["video_name"]
161 |         answer = line["answer"]
162 |         qs = line["prompt"]
163 |         task_type = line["task_type"]
164 |         video_path = os.path.join(
165 |             args.data_path,
166 |             line["video"],
167 |         )
168 | 
169 |         if os.path.exists(video_path):
170 |             vr = VideoReader(video_path, ctx=cpu(0))
171 |             fps = round(vr.get_avg_fps())
172 |             frame_idx = [
173 |                     i
174 |                     for i in range(0, len(vr), round(fps / 0.5))
175 |                 ]
176 |             if len(frame_idx) > 1000:
177 |                 frame_idx = [
178 |                     frame_idx[i]
179 |                     for i in range(0, len(frame_idx), len(frame_idx) // 1000)
180 |                 ]
181 |             video = vr.get_batch(frame_idx).asnumpy()
182 |             image_sizes = [video[0].shape[:2]]
183 |             video = process_images(video, image_processor, model.config)
184 |             video = [item.unsqueeze(0) for item in video]
185 |         else:
186 |             video = np.zeros((1, 1024, 1024, 3)).astype(np.uint8)
187 |             image_sizes = [(1024, 1024)]
188 |             video = process_images(video, image_processor, model.config)
189 | 
190 |         if getattr(model.config, "mm_use_im_start_end", False):
191 |             qs = (
192 |                 DEFAULT_IM_START_TOKEN
193 |                 + DEFAULT_IMAGE_TOKEN
194 |                 + DEFAULT_IM_END_TOKEN
195 |                 + "\n"
196 |                 + qs
197 |             )
198 |         else:
199 |             qs = DEFAULT_IMAGE_TOKEN + "\n" + qs
200 | 
201 |         conv = conv_templates[version].copy()
202 |         conv.append_message(conv.roles[0], qs)
203 |         conv.append_message(conv.roles[1], None)
204 |         prompt = conv.get_prompt()
205 | 
206 |         input_ids = (
207 |             tokenizer_image_token(
208 |                 prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
209 |             )
210 |             .unsqueeze(0)
211 |             .cuda()
212 |         )
213 | 
214 |         if "llama3" in version:
215 |             input_ids = input_ids[0][1:].unsqueeze(0)  # remove bos
216 | 
217 |         stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
218 |         keywords = [stop_str]
219 |         stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
220 |             
221 |         with torch.inference_mode():
222 |             output_ids = model.generate(
223 |                 input_ids,
224 |                 images=video,
225 |                 image_sizes=image_sizes,
226 |                 do_sample=False,
227 |                 temperature=0.0,
228 |                 max_new_tokens=5,  
229 |                 use_cache=True,
230 |                 stopping_criteria=[stopping_criteria],
231 |             )
232 |         if isinstance(output_ids, tuple):
233 |             output_ids = output_ids[0]
234 |         pred = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[
235 |             0
236 |         ].strip()
237 |         if pred.endswith(stop_str):
238 |             pred = pred[: -len(stop_str)]
239 |             pred = pred.strip()
240 |         pred = pred.replace("Answer", "")
241 | 
242 |         letters = ["A", "B", "C", "D"]
243 | 
244 |         pred_answer = re.findall("[\(\ \[]*([A-D])[\)\.\ \]]*", pred)
245 | 
246 |         pred_answer = pred_answer[0].strip()
247 |         pred_answer = pred_answer.strip("()")
248 |         if pred_answer in letters:
249 |             pred_idx = letters.index(pred_answer)
250 |             pred = letters[pred_idx]
251 |         else:
252 |             print("pred_answer: ", pred_answer, " pred: ", pred, flush=True)
253 |             pred_idx = 2
254 |             pred = letters[pred_idx]
255 | 
256 |         ans_id = uuid.uuid4()
257 |         output.append(
258 |             {
259 |                 "question": line["question"],
260 |                 "prompt": qs,
261 |                 "answer": answer,
262 |                 "pred": pred,
263 |                 "task_type": task_type,
264 |                 "answer_id": str(ans_id),
265 |                 "model_id": model_name,
266 |                 "video_name": video_name,
267 |                 "metadata": {},
268 |             }
269 |         )
270 | 
271 |     dist.barrier()
272 |     dist.all_gather_object(
273 |         final_output,
274 |         output,
275 |     )
276 |     all_output = list(chain(*final_output))
277 |     global_rank = dist.get_rank()
278 |     if global_rank == 0:
279 |         if os.path.exists("/tmp/generated_text"):
280 |             shutil.rmtree("/tmp/generated_text")
281 |         os.mkdir("/tmp/generated_text")
282 | 
283 |         with open(
284 |             os.path.join("/tmp/generated_text", "outputs.json"),
285 |             "w",
286 |         ) as f:
287 |             json.dump(all_output, f)
288 | 
289 |         correct = 0
290 |         total = 0
291 |         acc_dict = {}
292 |         for output in all_output:
293 |             pred = output["pred"]
294 |             gt = output["answer"]
295 |             task_type = output["task_type"]
296 |             if task_type not in acc_dict:
297 |                 acc_dict[task_type] = [0, 0]
298 |             acc_dict[task_type][1] += 1
299 |             total += 1
300 | 
301 |             if pred == gt:
302 |                 acc_dict[task_type][0] += 1
303 |                 correct += 1
304 | 
305 |         final_res = dict()
306 |         total = 0
307 |         idx = 0
308 |         for k, v in acc_dict.items():
309 |             idx += 1
310 |             final_res[k] = v[0] / v[1] * 100
311 |             total += final_res[k]
312 |         final_res["Acc"] = total / idx
313 |         print(final_res, flush=True)
314 | 
315 |         with open(os.path.join("/tmp/generated_text", "result.json"), "w") as f:
316 |             json.dump(final_res, f)
317 | 
318 | if __name__ == "__main__":
319 |     parser = argparse.ArgumentParser()
320 | 
321 |     parser.add_argument('--model_path', default="./checkpoints/longvu_qwen")
322 |     parser.add_argument('--model_name', default="cambrian_qwen")
323 |     parser.add_argument('--version', default="qwen")
324 |     parser.add_argument('--local-rank', default=0)
325 |     parser.add_argument('--data_path', required=True)
326 |     args = parser.parse_args()
327 | 
328 |     if "llama3" in args.version:
329 |         args.model_name = "cambrian_llama3"
330 | 
331 |     train(args)
332 | 


--------------------------------------------------------------------------------
/eval/eval_mvbench.py:
--------------------------------------------------------------------------------
  1 | # Adopted from https://github.com/lm-sys/FastChat. Below is the original copyright:
  2 | # Adopted from tatsu-lab@stanford_alpaca. Below is the original copyright:
  3 | # Make it more memory efficient by monkey patching the LLaMA model with FlashAttn.
  4 | 
  5 | # pyre-strict
  6 | 
  7 | # Need to call this before importing transformers.
  8 | 
  9 | 
 10 | import datetime
 11 | import json
 12 | import os
 13 | import re
 14 | import shutil
 15 | import uuid
 16 | from itertools import chain
 17 | import argparse
 18 | 
 19 | import sys
 20 | sys.path.append('./')
 21 | import numpy as np
 22 | from PIL import Image
 23 | 
 24 | import torch
 25 | 
 26 | from longvu.builder import load_pretrained_model
 27 | from longvu.constants import (
 28 |     DEFAULT_IM_END_TOKEN,
 29 |     DEFAULT_IM_START_TOKEN,
 30 |     DEFAULT_IMAGE_TOKEN,
 31 |     IMAGE_TOKEN_INDEX,
 32 | )
 33 | from longvu.conversation import conv_templates, SeparatorStyle
 34 | from longvu.mm_datautils import (
 35 |     KeywordsStoppingCriteria,
 36 |     process_images,
 37 |     tokenizer_image_token,
 38 | )
 39 | 
 40 | from decord import cpu, VideoReader  
 41 | from torch import distributed as dist
 42 | from tqdm import tqdm
 43 | 
 44 | from transformers.trainer_pt_utils import IterableDatasetShard
 45 | 
 46 | tasks = {
 47 |     "Action Sequence": (
 48 |         "action_sequence.json",
 49 |         "star/Charades_v1_480/",
 50 |         "video",
 51 |         True,
 52 |     ),  # has start & end
 53 |     "Action Prediction": (
 54 |         "action_prediction.json",
 55 |         "star/Charades_v1_480/",
 56 |         "video",
 57 |         True,
 58 |     ),  # has start & end
 59 |     "Action Antonym": ("action_antonym.json", "ssv2_video/", "video", False),
 60 |     "Fine-grained Action": (
 61 |         "fine_grained_action.json",
 62 |         "Moments_in_Time_Raw/videos/",
 63 |         "video",
 64 |         False,
 65 |     ),
 66 |     "Unexpected Action": ("unexpected_action.json", "FunQA_test/test/", "video", False),
 67 |     "Object Existence": (
 68 |         "object_existence.json",
 69 |         "clevrer/video_validation/",
 70 |         "video",
 71 |         False,
 72 |     ),
 73 |     "Object Interaction": (
 74 |         "object_interaction.json",
 75 |         "star/Charades_v1_480/",
 76 |         "video",
 77 |         True,
 78 |     ),  # has start & end
 79 |     "Object Shuffle": ("object_shuffle.json", "perception/videos/", "video", False),
 80 |     "Moving Direction": (
 81 |         "moving_direction.json",
 82 |         "clevrer/video_validation/",
 83 |         "video",
 84 |         False,
 85 |     ),
 86 |     "Action Localization": (
 87 |         "action_localization.json",
 88 |         "sta/sta_video/",
 89 |         "video",
 90 |         True,
 91 |     ),  # has start & end
 92 |     "Scene Transition": ("scene_transition.json", "scene_qa/video/", "video", False),
 93 |     "Action Count": ("action_count.json", "perception/videos/", "video", False),
 94 |     "Moving Count": ("moving_count.json", "clevrer/video_validation/", "video", False),
 95 |     "Moving Attribute": (
 96 |         "moving_attribute.json",
 97 |         "clevrer/video_validation/",
 98 |         "video",
 99 |         False,
100 |     ),
101 |     "State Change": ("state_change.json", "perception/videos/", "video", False),
102 |     "Fine-grained Pose": ("fine_grained_pose.json", "nturgbd/", "video", False),
103 |     "Character Order": ("character_order.json", "perception/videos/", "video", False),
104 |     "Egocentric Navigation": ("egocentric_navigation.json", "vlnqa/", "video", False),
105 |     "Episodic Reasoning": (
106 |         "episodic_reasoning.json",
107 |         "tvqa/frames_fps3_hq/",
108 |         "frame",
109 |         True,
110 |     ),  # has start & end, read frame
111 |     "Counterfactual Inference": (
112 |         "counterfactual_inference.json",
113 |         "clevrer/video_validation/",
114 |         "video",
115 |         False,
116 |     ),
117 | }
118 | 
119 | class EvalDataset(torch.utils.data.IterableDataset):
120 |     """Dataset for supervised fine-tuning."""
121 | 
122 |     def __init__(
123 |         self,
124 |         data_path: str,
125 |     ) -> None:
126 |         super(EvalDataset, self).__init__()
127 | 
128 |         self.data_path = data_path
129 | 
130 |         list_data_dict = []
131 |         for task_name, task in tasks.items():
132 |             json_file = os.path.join(data_path, "json", task[0])
133 |             vis_folder = os.path.join(data_path, "video", task[1])
134 |             with open(json_file, "r") as f:
135 |                 json_data = json.load(f)
136 |             for data in json_data:
137 |                 video_path = os.path.join(vis_folder, data["video"])
138 |                 answer = data["answer"]
139 |                 question = data["question"]
140 |                 answer_idx = -1
141 |                 letters = []
142 |                 options = data["candidates"]
143 |                 options_string = ""
144 |                 for option_idx, c in enumerate(options):
145 |                     letters.append(f"{chr(ord('A') + option_idx)}")
146 |                     options_string += f"({chr(ord('A') + option_idx)}) {c}\n"
147 |                     if c == answer:
148 |                         answer_idx = option_idx
149 |                 prompt = f"Question: {question}\nOptions:\n{options_string}Answer with the option's letter from the given choices directly and only give the best option."
150 |                 list_data_dict.append(
151 |                     {
152 |                         "task_type": task_name,
153 |                         "bound": (data["start"], data["end"]) if task[3] else task[3],
154 |                         "question": question,
155 |                         "prompt": prompt,
156 |                         "answer": answer_idx,
157 |                         "answer_word": data["answer"],
158 |                         "video_name": data["video"].split(".")[0],
159 |                         "video": video_path,
160 |                         "data_type": task[2],
161 |                         "letters": ",".join(letters),
162 |                     }
163 |                 )
164 | 
165 |         self.data = list_data_dict
166 | 
167 |     def __len__(self) -> int:
168 |         return len(self.data)
169 | 
170 |     # pyre-fixme[3]: Return type must be annotated.
171 |     def __iter__(self):
172 |         return iter(self.data)
173 | 
174 |     # pyre-fixme[3]: Return type must be annotated.
175 |     # pyre-fixme[2]: Parameter must be annotated.
176 |     def __getitem__(self, i):
177 |         return self.data[i]
178 | 
179 | 
180 | def train(args) -> None:
181 |     dist.init_process_group(backend="nccl", timeout=datetime.timedelta(hours=8))
182 |     
183 |     version = args.version
184 |     model_name = args.model_name
185 |     model_path = args.model_path
186 | 
187 |     # torch.distributed.barrier()
188 |     tokenizer, model, image_processor, context_len = load_pretrained_model(
189 |         model_path,  
190 |         None,
191 |         model_name,
192 |         device_map=None,
193 |     )
194 |     model.get_model().config.drop_threshold = 0.8
195 |     model.config.use_cache = True
196 |     model.cuda()
197 |     dataset = EvalDataset(
198 |         data_path=args.data_path,
199 |     )
200 |     world_size = torch.distributed.get_world_size()
201 |     world_rank = torch.distributed.get_rank()
202 |     shard_dataset = IterableDatasetShard(
203 |         dataset,
204 |         batch_size=1,
205 |         num_processes=world_size,
206 |         process_index=world_rank,
207 |     )
208 |     torch.distributed.barrier()
209 |     output = []
210 |     final_output = [None] * world_size
211 | 
212 |     for line in tqdm(shard_dataset):
213 |         video_name = line["video_name"]
214 |         answer = line["answer"]
215 |         qs = line["prompt"]
216 |         task_type = line["task_type"]
217 |         video_path = line["video"]
218 |         bound = line["bound"]
219 |         data_type = line["data_type"]
220 |         letters = line["letters"].split(",")
221 | 
222 |         if os.path.exists(video_path):
223 |             if data_type == "video":
224 |                 vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
225 |                 max_frame = len(vr) - 1
226 |                 fps = float(vr.get_avg_fps())
227 |                 if bound:
228 |                     start, end = bound[0], bound[1]
229 |                     start_idx = max(0, round(start * fps))
230 |                     end_idx = min(round(end * fps), max_frame)
231 |                     frame_indices = np.array(
232 |                         [
233 |                             i
234 |                             for i in range(
235 |                                 start_idx,
236 |                                 end_idx,
237 |                                 round(fps / 2),
238 |                             )
239 |                         ]
240 |                     )
241 |                 else:
242 |                     frame_indices = np.array(
243 |                         [
244 |                             i
245 |                             for i in range(
246 |                                 0,
247 |                                 len(vr),
248 |                                 round(fps / 2),
249 |                             )
250 |                         ]
251 |                     )
252 |                 video = []
253 |                 for frame_index in frame_indices:
254 |                     img = vr[frame_index].asnumpy()
255 |                     video.append(img)
256 |                 video = np.stack(video)
257 |             else:
258 |                 max_frame = len(os.listdir(video_path))
259 |                 images_group = list()
260 |                 fps = 3
261 |                 if bound:
262 |                     start, end = bound[0], bound[1]
263 |                 else:
264 |                     start, end = -100000, 100000
265 |                 start_idx = max(1, round(start * fps))
266 |                 end_idx = min(round(end * fps), max_frame)
267 |                 frame_indices = [
268 |                     i
269 |                     for i in range(
270 |                         start_idx,
271 |                         end_idx,
272 |                         round(fps / 2),
273 |                     )
274 |                 ]
275 |                 for frame_index in frame_indices:
276 |                     img = Image.open(
277 |                         os.path.join(video_path, f"{frame_index:05d}.jpg")
278 |                     ).convert("RGB")
279 |                     images_group.append(np.array(img))
280 |                 video = np.stack(images_group)
281 |             
282 |             image_sizes = [video[0].shape[:2]]
283 |             video = process_images(video, image_processor, model.config)
284 |             video = [item.unsqueeze(0) for item in video]
285 |         else:
286 |             video = np.zeros((1, 1024, 1024, 3)).astype(np.uint8)
287 |             image_sizes = [(1024, 1024)]
288 |             video = process_images(video, image_processor, model.config)
289 | 
290 |         if getattr(model.config, "mm_use_im_start_end", False):
291 |             qs = (
292 |                 DEFAULT_IM_START_TOKEN
293 |                 + DEFAULT_IMAGE_TOKEN
294 |                 + DEFAULT_IM_END_TOKEN
295 |                 + "\n"
296 |                 + qs
297 |             )
298 |         else:
299 |             qs = DEFAULT_IMAGE_TOKEN + "\n" + qs
300 | 
301 |         conv = conv_templates[version].copy()
302 |         conv.append_message(conv.roles[0], qs)
303 |         conv.append_message(conv.roles[1], None)
304 |         prompt = conv.get_prompt()
305 | 
306 |         input_ids = (
307 |             tokenizer_image_token(
308 |                 prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
309 |             )
310 |             .unsqueeze(0)
311 |             .cuda()
312 |         )
313 | 
314 |         if "llama3" in version:
315 |             input_ids = input_ids[0][1:].unsqueeze(0)  # remove bos
316 | 
317 |         stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
318 |         keywords = [stop_str]
319 |         stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
320 |             
321 |         with torch.inference_mode():
322 |             output_ids = model.generate(
323 |                 input_ids,
324 |                 images=video,
325 |                 image_sizes=image_sizes,
326 |                 do_sample=False,
327 |                 temperature=0.0,
328 |                 max_new_tokens=5,  
329 |                 use_cache=True,
330 |                 stopping_criteria=[stopping_criteria],
331 |             )
332 |         if isinstance(output_ids, tuple):
333 |             output_ids = output_ids[0]
334 |         pred = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[
335 |             0
336 |         ].strip()
337 |         if pred.endswith(stop_str):
338 |             pred = pred[: -len(stop_str)]
339 |             pred = pred.strip()
340 |         pred = pred.replace("Answer", "")
341 | 
342 |         pred_answer = re.findall(
343 |             f"[\(,\ ]*[{letters[0]}-{letters[-1]}][\),\ ]*", pred
344 |         )
345 | 
346 |         pred_answer = pred_answer[0].strip()
347 |         pred_answer = pred_answer.strip("()")
348 |         if pred_answer in letters:
349 |             pred_idx = letters.index(pred_answer)
350 |             pred = letters[pred_idx]
351 |         else:
352 |             print("pred_answer: ", pred_answer, " pred: ", pred, flush=True)
353 |             pred_idx = 2
354 |             pred = letters[pred_idx]
355 | 
356 |         ans_id = uuid.uuid4()
357 |         output.append(
358 |             {
359 |                 "question": line["question"],
360 |                 "prompt": qs,
361 |                 "answer": answer,
362 |                 "pred": pred_idx,
363 |                 "task_type": task_type,
364 |                 "answer_id": str(ans_id),
365 |                 "model_id": model_name,
366 |                 "video_name": video_name,
367 |                 "metadata": {},
368 |             }
369 |         )
370 | 
371 |     dist.barrier()
372 |     dist.all_gather_object(
373 |         final_output,
374 |         output,
375 |     )
376 |     
377 |     all_output = list(chain(*final_output))
378 |     global_rank = dist.get_rank()
379 |     if global_rank == 0:
380 |         if os.path.exists("/tmp/generated_text"):
381 |             shutil.rmtree("/tmp/generated_text")
382 |         os.mkdir("/tmp/generated_text")
383 | 
384 |         with open(
385 |             os.path.join("/tmp/generated_text", "outputs.json"),
386 |             "w",
387 |         ) as f:
388 |             json.dump(all_output, f)
389 | 
390 |         task_types = tasks.keys()
391 |         task_acc = {x: [] for x in task_types}
392 |         acc = []
393 | 
394 |         for i, x in enumerate(all_output):
395 |             value = 1
396 |             if x["pred"] != x["answer"]:
397 |                 value = 0
398 |             acc.append(value)
399 |             task_acc[x["task_type"]].append(value)
400 | 
401 |         acc = sum(acc) * 100 / len(acc)
402 |         task_acc = {x: sum(task_acc[x]) * 100 / len(task_acc[x]) for x in task_acc}
403 |         print(f"Accuracy: ", acc)
404 |         print("Task ccuracy", task_acc)
405 | 
406 |         task_acc["avg"] = acc
407 | 
408 |         with open(os.path.join("/tmp/generated_text", "result.json"), "w") as f:
409 |             json.dump(task_acc, f)
410 | 
411 | if __name__ == "__main__":
412 |     parser = argparse.ArgumentParser()
413 | 
414 |     parser.add_argument('--model_path', default="./checkpoints/longvu_qwen")
415 |     parser.add_argument('--model_name', default="cambrian_qwen")
416 |     parser.add_argument('--version', default="qwen")
417 |     parser.add_argument('--local-rank', default=0)
418 |     parser.add_argument('--data_path', required=True)
419 |     args = parser.parse_args()
420 | 
421 |     if "llama3" in args.version:
422 |         args.model_name = "cambrian_llama3"
423 | 
424 |     train(args)
425 | 


--------------------------------------------------------------------------------
/eval/eval_videomme.py:
--------------------------------------------------------------------------------
  1 | # Adopted from https://github.com/lm-sys/FastChat. Below is the original copyright:
  2 | # Adopted from tatsu-lab@stanford_alpaca. Below is the original copyright:
  3 | # Make it more memory efficient by monkey patching the LLaMA model with FlashAttn.
  4 | 
  5 | # pyre-strict
  6 | 
  7 | # Need to call this before importing transformers.
  8 | 
  9 | 
 10 | import datetime
 11 | import json
 12 | import logging
 13 | import os
 14 | import re
 15 | import shutil
 16 | import uuid
 17 | from itertools import chain
 18 | import argparse
 19 | 
 20 | import sys
 21 | sys.path.append('./')
 22 | import numpy as np
 23 | import pysubs2
 24 | 
 25 | import torch
 26 | 
 27 | from longvu.builder import load_pretrained_model
 28 | from longvu.constants import (
 29 |     DEFAULT_IM_END_TOKEN,
 30 |     DEFAULT_IM_START_TOKEN,
 31 |     DEFAULT_IMAGE_TOKEN,
 32 |     IMAGE_TOKEN_INDEX,
 33 | )
 34 | from longvu.conversation import conv_templates, SeparatorStyle
 35 | from longvu.mm_datautils import (
 36 |     KeywordsStoppingCriteria,
 37 |     process_images,
 38 |     tokenizer_image_token,
 39 | )
 40 | 
 41 | from decord import cpu, VideoReader  # @manual=fbsource//third-party/pypi/decord:decord
 42 | from pyarrow import parquet as pq
 43 | from torch import distributed as dist
 44 | from tqdm import tqdm
 45 | 
 46 | from transformers.trainer_pt_utils import IterableDatasetShard
 47 | 
 48 | # pyre-fixme[3]: Return type must be annotated.
 49 | # pyre-fixme[2]: Parameter must be annotated.
 50 | def load_parquet(parquet_file):
 51 |     table = pq.read_table(parquet_file)
 52 | 
 53 |     # Convert PyArrow Table to pandas DataFrame
 54 |     df = table.to_pandas()
 55 | 
 56 |     jsons = []
 57 |     for record in df.itertuples():
 58 | 
 59 |         if len(jsons) < int(record.video_id):
 60 |             jsons.append(
 61 |                 {
 62 |                     "video_id": record.video_id,
 63 |                     "youtube_id": record.videoID,
 64 |                     "url": record.url,
 65 |                     "duration": record.duration,
 66 |                     "domain": record.domain,
 67 |                     "sub_category": record.sub_category,
 68 |                     "questions": [
 69 |                         {
 70 |                             "question_id": record.question_id,
 71 |                             "task_type": record.task_type,
 72 |                             "question": record.question,
 73 |                             "choices": list(record.options),
 74 |                             "answer": record.answer,
 75 |                         }
 76 |                     ],
 77 |                 }
 78 |             )
 79 |         else:
 80 |             jsons[-1]["questions"].append(
 81 |                 {
 82 |                     "question_id": record.question_id,
 83 |                     "task_type": record.task_type,
 84 |                     "question": record.question,
 85 |                     "choices": list(record.options),
 86 |                     "answer": record.answer,
 87 |                 }
 88 |             )
 89 | 
 90 |     return jsons
 91 | 
 92 | 
 93 | class EvalDataset(torch.utils.data.IterableDataset):
 94 |     """Dataset for supervised fine-tuning."""
 95 | 
 96 |     video_formats = [".mp4", ".avi", ".mov", ".mkv"]
 97 | 
 98 |     def __init__(
 99 |         self,
100 |         data_path: str,
101 |     ) -> None:
102 |         super(EvalDataset, self).__init__()
103 |         logging.warning("Loading data...")
104 | 
105 |         self.data_path = data_path
106 | 
107 |         data_list = load_parquet(
108 |             os.path.join(self.data_path, "test-00000-of-00001.parquet")
109 |         )
110 | 
111 |         list_data_dict = []
112 | 
113 |         for item in data_list:
114 |             video_ytid = item["url"].split("watch?v=")[-1]
115 |             video_path = os.path.join(self.data_path, "data", f"{video_ytid}.mp4")
116 |             for fmt in self.video_formats:  # Added this line
117 |                 temp_path = os.path.join(self.data_path, "data", f"{video_ytid}{fmt}")
118 |                 if os.path.exists(temp_path):
119 |                     video_path = temp_path
120 |                     break
121 | 
122 |             subtitle_path = os.path.join(
123 |                 self.data_path, "subtitle", f"{video_ytid}.srt"
124 |             )
125 | 
126 |             list_data_dict.append(
127 |                 {
128 |                     "questions": item["questions"],
129 |                     "video": video_path,
130 |                     "subtitle": subtitle_path,
131 |                     "video_name": video_ytid,
132 |                     "duration": item["duration"],
133 |                 }
134 |             )
135 | 
136 |         # pyre-fixme[4]: Attribute must be annotated.
137 |         self.data = list_data_dict
138 | 
139 |     def __len__(self) -> int:
140 |         return len(self.data)
141 | 
142 |     # pyre-fixme[3]: Return type must be annotated.
143 |     def __iter__(self):
144 |         return iter(self.data)
145 | 
146 |     # pyre-fixme[3]: Return type must be annotated.
147 |     # pyre-fixme[2]: Parameter must be annotated.
148 |     def __getitem__(self, i):
149 |         return self.data[i]
150 | 
151 | 
152 | def train(args) -> None:
153 |     dist.init_process_group(backend="nccl", timeout=datetime.timedelta(hours=8))
154 |     
155 |     version = args.version
156 |     model_name = args.model_name
157 |     model_path = args.model_path
158 | 
159 |     torch.distributed.barrier()
160 |     tokenizer, model, image_processor, context_len = load_pretrained_model(
161 |         model_path,  # pyre-fixme
162 |         None,
163 |         model_name,
164 |         device_map=None,
165 |     )
166 |     model.get_model().config.drop_threshold = 0.65
167 |     model.config.use_cache = True
168 |     model.cuda()
169 |     dataset = EvalDataset(
170 |         data_path=args.data_path,
171 |     )
172 |     world_size = torch.distributed.get_world_size()
173 |     world_rank = torch.distributed.get_rank()
174 |     shard_dataset = IterableDatasetShard(
175 |         dataset,
176 |         batch_size=1,
177 |         num_processes=world_size,
178 |         process_index=world_rank,
179 |     )
180 |     torch.distributed.barrier()
181 |     output = []
182 |     final_output = [None] * world_size
183 | 
184 |     for line in tqdm(shard_dataset):
185 |         video_name = line["video_name"]
186 |         video_path = line["video"]
187 |         subtitle_path = line["subtitle"]
188 |         questions = line["questions"]
189 | 
190 |         if os.path.exists(video_path):
191 |             vr = VideoReader(video_path, ctx=cpu(0))
192 |             fps = round(vr.get_avg_fps())
193 |             sample_fps = 0.5
194 |             frame_idx = [i for i in range(0, len(vr), round(fps / sample_fps))]
195 |             video = vr.get_batch(frame_idx).asnumpy()
196 |             
197 |             image_sizes = [video[0].shape[:2]]
198 |             video = process_images(video, image_processor, model.config)
199 |             video = [item.unsqueeze(0) for item in video]
200 |         else:
201 |             frame_idx = [1]
202 |             fps = 1
203 |             print(f"video_path {video_path} does not exist")
204 | 
205 |             video = np.zeros((1, 1024, 1024, 3)).astype(np.uint8)
206 |             image_sizes = [(1024, 1024)]
207 |             video = process_images(video, image_processor, model.config)
208 | 
209 |         if (
210 |             os.path.exists(video_path)
211 |             and os.path.exists(subtitle_path)
212 |         ):
213 |             subs = pysubs2.load(subtitle_path, encoding="utf-8")
214 |             subtitles = []
215 |             for select_id in range(0, len(frame_idx)):
216 |                 seleced_frame_id = frame_idx[select_id]
217 |                 sub_text = ""
218 |                 cur_time = pysubs2.make_time(fps=fps, frames=seleced_frame_id)
219 |                 for sub in subs:
220 |                     if sub.start < cur_time and sub.end > cur_time:
221 |                         sub_text = sub.text.replace("\\N", " ")
222 |                         break
223 |                 if sub_text.strip():
224 |                     if (
225 |                         "[Music]" not in sub_text
226 |                         and "[Applause]" not in sub_text
227 |                         and sub_text not in subtitles
228 |                     ):
229 |                         if len(subtitles) > 0:
230 |                             if sub_text not in subtitles[-1]:
231 |                                 subtitles.append(sub_text)
232 |                         else:
233 |                             subtitles.append(sub_text)
234 |             if len(tokenizer("\n".join(subtitles)).input_ids) > 6000:
235 |                 interval = len(subtitles) // 200
236 |                 indices = np.arange(0, len(subtitles), interval)
237 |                 subtitles = [subtitles[i] for i in indices]
238 |             subtitles = "\n".join(subtitles)
239 |             subtitles = f"This video's subtitles are listed below:\n{subtitles}\n"
240 |         else:
241 |             subtitles = ""
242 | 
243 |         for question in questions:
244 |             q = question["question"]
245 |             ops = question["choices"]
246 |             instruct = f"Question: {q}\n"
247 |             instruct += "Options:\n"
248 |             for op in ops:
249 |                 instruct += f"{op}\n"
250 |             instruct += (
251 |                 "Respond with only the letter (A, B, C, or D) of the correct option.\n"
252 |             )
253 |             qs = subtitles + instruct
254 | 
255 |             if getattr(model.config, "mm_use_im_start_end", False):
256 |                 qs = (
257 |                     DEFAULT_IM_START_TOKEN
258 |                     + DEFAULT_IMAGE_TOKEN
259 |                     + DEFAULT_IM_END_TOKEN
260 |                     + "\n"
261 |                     + qs
262 |                 )
263 |             else:
264 |                 qs = DEFAULT_IMAGE_TOKEN + "\n" + qs
265 | 
266 |             conv = conv_templates[version].copy()
267 |             conv.append_message(conv.roles[0], qs)
268 |             conv.append_message(conv.roles[1], None)
269 |             prompt = conv.get_prompt()
270 | 
271 |             input_ids = (
272 |                 tokenizer_image_token(
273 |                     prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
274 |                 )
275 |                 .unsqueeze(0)
276 |                 .cuda()
277 |             )
278 | 
279 |             if "llama3" in version:
280 |                 input_ids = input_ids[0][1:].unsqueeze(0)  # remove bos
281 | 
282 |             stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
283 |             keywords = [stop_str]
284 |             stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
285 | 
286 |             
287 |             with torch.inference_mode():
288 |                 output_ids = model.generate(
289 |                     input_ids,
290 |                     images=video,
291 |                     image_sizes=image_sizes,
292 |                     do_sample=False,
293 |                     temperature=0.0,
294 |                     max_new_tokens=5,  
295 |                     use_cache=True,
296 |                     stopping_criteria=[stopping_criteria],
297 |                 )
298 |             if isinstance(output_ids, tuple):
299 |                 output_ids = output_ids[0]
300 |             pred = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[
301 |                 0
302 |             ].strip()
303 |             if pred.endswith(stop_str):
304 |                 pred = pred[: -len(stop_str)]
305 |                 pred = pred.strip()
306 |             pred = pred.replace("Answer", "")
307 | 
308 |             letters = ["A", "B", "C", "D"]
309 | 
310 |             pred_answer = re.findall("[\(\ \[]*([A-D])[\)\.\ \]]*", pred)
311 | 
312 |             pred_answer = pred_answer[0].strip()
313 |             pred_answer = pred_answer.strip("()")
314 |             if pred_answer in letters:
315 |                 pred_idx = letters.index(pred_answer)
316 |                 pred = letters[pred_idx]
317 |             else:
318 |                 print("pred_answer: ", pred_answer, " pred: ", pred, flush=True)
319 |                 pred_idx = 2
320 |                 pred = letters[pred_idx]
321 | 
322 |             ans_id = uuid.uuid4()
323 |             output.append(
324 |                 {
325 |                     "question": question["question"],
326 |                     "answer": question["answer"],
327 |                     "pred": pred,
328 |                     "answer_id": str(ans_id),
329 |                     "model_id": model_name,
330 |                     "video_name": video_name,
331 |                     "metadata": {},
332 |                     "duration": line["duration"],
333 |                 }
334 |             )
335 | 
336 |     dist.barrier()
337 |     dist.all_gather_object(
338 |         final_output,
339 |         output,
340 |     )
341 |     all_output = list(chain(*final_output))
342 |     global_rank = dist.get_rank()
343 |     if global_rank == 0:
344 |         if os.path.exists("/tmp/generated_text"):
345 |             shutil.rmtree("/tmp/generated_text")
346 |         os.mkdir("/tmp/generated_text")
347 | 
348 |         with open(
349 |             os.path.join("/tmp/generated_text", "outputs.json"),
350 |             "w",
351 |         ) as f:
352 |             json.dump(all_output, f)
353 | 
354 |         correct = 0
355 |         duration_correct = {"short": 0, "medium": 0, "long": 0}
356 |         duration_all = {"short": 0, "medium": 0, "long": 0}
357 | 
358 |         for output in all_output:
359 |             duration = output["duration"]
360 |             duration_all[duration] += 1
361 |             if output["pred"] == output["answer"]:
362 |                 correct += 1
363 |                 duration_correct[duration] += 1
364 | 
365 |         result = {"averge_acc": correct / len(all_output)}
366 | 
367 |         for duration, correct_count in duration_correct.items():
368 |             result[f"{duration}_acc"] = (
369 |                 correct_count / duration_all[duration]
370 |                 if duration_all[duration] > 0
371 |                 else 0
372 |             )
373 | 
374 |         print(f"Accuracy: {result}", flush=True)
375 | 
376 | 
377 | 
378 | if __name__ == "__main__":
379 |     parser = argparse.ArgumentParser()
380 | 
381 |     parser.add_argument('--model_path', default="./checkpoints/longvu_qwen")
382 |     parser.add_argument('--model_name', default="cambrian_qwen")
383 |     parser.add_argument('--version', default="qwen")
384 |     parser.add_argument('--data_path', required=True)
385 |     args = parser.parse_args()
386 | 
387 |     if "llama3" in args.version:
388 |         args.model_name = "cambrian_llama3"
389 | 
390 |     train(args)
391 | 


--------------------------------------------------------------------------------
/examples/video1.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Vision-CAIR/LongVU/1b0be5cc28614b03a423bbaf129ae364750b5012/examples/video1.mp4


--------------------------------------------------------------------------------
/examples/video2.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Vision-CAIR/LongVU/1b0be5cc28614b03a423bbaf129ae364750b5012/examples/video2.mp4


--------------------------------------------------------------------------------
/examples/video3.mp4:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/Vision-CAIR/LongVU/1b0be5cc28614b03a423bbaf129ae364750b5012/examples/video3.mp4


--------------------------------------------------------------------------------
/longvu/__init__.py:
--------------------------------------------------------------------------------
1 | # pyre-unsafe
2 | from .language_model.cambrian_qwen import CambrianQwenModel
3 | from .language_model.cambrian_llama import CambrianLlamaModel
4 | 


--------------------------------------------------------------------------------
/longvu/apply_delta.py:
--------------------------------------------------------------------------------
 1 | # pyre-unsafe
 2 | """
 3 | Usage:
 4 | python3 -m fastchat.model.apply_delta --base ~/model_weights/llama-7b --target ~/model_weights/vicuna-7b --delta lmsys/vicuna-7b-delta
 5 | """
 6 | 
 7 | import argparse
 8 | 
 9 | import torch
10 | from tqdm import tqdm
11 | from transformers import AutoModelForCausalLM, AutoTokenizer
12 | 
13 | from . import LlavaLlamaForCausalLM
14 | 
15 | 
16 | def apply_delta(base_model_path, target_model_path, delta_path):
17 |     print("Loading base model")
18 |     base = AutoModelForCausalLM.from_pretrained(
19 |         base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
20 |     )
21 | 
22 |     print("Loading delta")
23 |     delta = LlavaLlamaForCausalLM.from_pretrained(
24 |         delta_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
25 |     )
26 |     delta_tokenizer = AutoTokenizer.from_pretrained(delta_path)
27 | 
28 |     print("Applying delta")
29 |     for name, param in tqdm(delta.state_dict().items(), desc="Applying delta"):
30 |         if name not in base.state_dict():
31 |             assert name in [
32 |                 "model.mm_projector.weight",
33 |                 "model.mm_projector.bias",
34 |             ], f"{name} not in base model"
35 |             continue
36 |         if param.data.shape == base.state_dict()[name].shape:
37 |             param.data += base.state_dict()[name]
38 |         else:
39 |             assert name in [
40 |                 "model.embed_tokens.weight",
41 |                 "lm_head.weight",
42 |             ], f"{name} dimension mismatch: {param.data.shape} vs {base.state_dict()[name].shape}"
43 |             bparam = base.state_dict()[name]
44 |             param.data[: bparam.shape[0], : bparam.shape[1]] += bparam
45 | 
46 |     print("Saving target model")
47 |     delta.save_pretrained(target_model_path)
48 |     delta_tokenizer.save_pretrained(target_model_path)
49 | 
50 | 
51 | if __name__ == "__main__":
52 |     parser = argparse.ArgumentParser()
53 |     parser.add_argument("--base-model-path", type=str, required=True)
54 |     parser.add_argument("--target-model-path", type=str, required=True)
55 |     parser.add_argument("--delta-path", type=str, required=True)
56 | 
57 |     args = parser.parse_args()
58 | 
59 |     apply_delta(args.base_model_path, args.target_model_path, args.delta_path)
60 | 


--------------------------------------------------------------------------------
/longvu/builder.py:
--------------------------------------------------------------------------------
  1 | #    Copyright 2023 Haotian Liu
  2 | #
  3 | #    Licensed under the Apache License, Version 2.0 (the "License");
  4 | #    you may not use this file except in compliance with the License.
  5 | #    You may obtain a copy of the License at
  6 | #
  7 | #        http://www.apache.org/licenses/LICENSE-2.0
  8 | #
  9 | #    Unless required by applicable law or agreed to in writing, software
 10 | #    distributed under the License is distributed on an "AS IS" BASIS,
 11 | #    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12 | #    See the License for the specific language governing permissions and
 13 | #    limitations under the License.
 14 | 
 15 | # pyre-unsafe
 16 | 
 17 | 
 18 | import os
 19 | import shutil
 20 | import warnings
 21 | 
 22 | import torch
 23 | from longvu.constants import (
 24 |     DEFAULT_IM_END_TOKEN,
 25 |     DEFAULT_IM_START_TOKEN,
 26 |     DEFAULT_IMAGE_PATCH_TOKEN,
 27 | )
 28 | 
 29 | from longvu.language_model.cambrian_llama import CambrianLlamaForCausalLM
 30 | from longvu.language_model.cambrian_qwen import CambrianQwenForCausalLM
 31 | 
 32 | from transformers import (
 33 |     AutoConfig,
 34 |     AutoModelForCausalLM,
 35 |     AutoTokenizer,
 36 |     BitsAndBytesConfig,
 37 | )
 38 | 
 39 | 
 40 | def load_pretrained_model(
 41 |     model_path,
 42 |     model_base,
 43 |     model_name,
 44 |     load_8bit=False,
 45 |     load_4bit=False,
 46 |     device_map="auto",
 47 |     device="cuda",
 48 |     use_flash_attn=False,
 49 |     model_args=None,
 50 |     **kwargs,
 51 | ):
 52 |     kwargs = {"device_map": device_map, **kwargs}
 53 | 
 54 |     if device != "cuda":
 55 |         kwargs["device_map"] = {"": device}
 56 | 
 57 |     if load_8bit:
 58 |         kwargs["load_in_8bit"] = True
 59 |     elif load_4bit:
 60 |         kwargs["load_in_4bit"] = True
 61 |         kwargs["quantization_config"] = BitsAndBytesConfig(
 62 |             load_in_4bit=True,
 63 |             bnb_4bit_compute_dtype=torch.float16,
 64 |             bnb_4bit_use_double_quant=True,
 65 |             bnb_4bit_quant_type="nf4",
 66 |         )
 67 |     else:
 68 |         kwargs["torch_dtype"] = torch.float16
 69 | 
 70 |     if use_flash_attn:
 71 |         kwargs["attn_implementation"] = "flash_attention_2"
 72 | 
 73 |     if "cambrian" in model_name.lower():
 74 |         # Load Cambrian model
 75 |         if "lora" in model_name.lower() and model_base is None:
 76 |             warnings.warn(
 77 |                 "There is `lora` in model name but no `model_base` is provided. If you are loading a LoRA model, please provide the `model_base` argument. Detailed instruction: https://github.com/haotian-liu/LLaVA#launch-a-model-worker-lora-weights-unmerged."
 78 |             )
 79 |         if "lora" in model_name.lower() and model_base is not None:
 80 |             # pyre-fixme[21]: Could not find module
 81 |             #  `core_ai.llava.language_model.cambrian_llama`.
 82 |             from core_ai.llava.language_model.cambrian_llama import CambrianConfig
 83 | 
 84 |             lora_cfg_pretrained = CambrianConfig.from_pretrained(model_path)
 85 |             tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
 86 |             print("Loading Cambrian from base model...")
 87 |             model = CambrianLlamaForCausalLM.from_pretrained(
 88 |                 model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, **kwargs
 89 |             )
 90 |             token_num, tokem_dim = model.lm_head.out_features, model.lm_head.in_features
 91 |             if model.lm_head.weight.shape[0] != token_num:
 92 |                 model.lm_head.weight = torch.nn.Parameter(
 93 |                     torch.empty(
 94 |                         token_num, tokem_dim, device=model.device, dtype=model.dtype
 95 |                     )
 96 |                 )
 97 |                 model.model.embed_tokens.weight = torch.nn.Parameter(
 98 |                     torch.empty(
 99 |                         token_num, tokem_dim, device=model.device, dtype=model.dtype
100 |                     )
101 |                 )
102 | 
103 |             print("Loading additional Cambrian weights...")
104 |             if os.path.exists(os.path.join(model_path, "non_lora_trainables.bin")):
105 |                 non_lora_trainables = torch.load(
106 |                     os.path.join(model_path, "non_lora_trainables.bin"),
107 |                     map_location="cpu",
108 |                 )
109 |             else:
110 |                 # this is probably from HF Hub
111 |                 from huggingface_hub import hf_hub_download
112 | 
113 |                 def load_from_hf(repo_id, filename, subfolder=None):
114 |                     cache_file = hf_hub_download(
115 |                         repo_id=repo_id, filename=filename, subfolder=subfolder
116 |                     )
117 |                     return torch.load(cache_file, map_location="cpu")
118 | 
119 |                 non_lora_trainables = load_from_hf(
120 |                     model_path, "non_lora_trainables.bin"
121 |                 )
122 |             non_lora_trainables = {
123 |                 (k[11:] if k.startswith("base_model.") else k): v
124 |                 for k, v in non_lora_trainables.items()
125 |             }
126 |             if any(k.startswith("model.model.") for k in non_lora_trainables):
127 |                 non_lora_trainables = {
128 |                     (k[6:] if k.startswith("model.") else k): v
129 |                     for k, v in non_lora_trainables.items()
130 |                 }
131 |             model.load_state_dict(non_lora_trainables, strict=False)
132 | 
133 |             from peft import PeftModel
134 | 
135 |             print("Loading LoRA weights...")
136 |             model = PeftModel.from_pretrained(model, model_path)
137 |             print("Merging LoRA weights...")
138 |             model = model.merge_and_unload()
139 |             print("Model is loaded...")
140 |         elif model_base is not None:
141 |             # this may be mm projector only
142 |             print(f"Loading Cambrian-1 from base model... {model_base}")
143 |             tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
144 |             cfg_pretrained = AutoConfig.from_pretrained(model_path)
145 |             model = CambrianLlamaForCausalLM.from_pretrained(
146 |                 model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs
147 |             )
148 | 
149 |             mm_projector_weights = torch.load(
150 |                 os.path.join(model_path, "mm_projector.bin"), map_location="cpu"
151 |             )
152 |             mm_projector_weights = {
153 |                 k: v.to(torch.float16) for k, v in mm_projector_weights.items()
154 |             }
155 |             model.load_state_dict(mm_projector_weights, strict=False)
156 |         else:
157 |             if "qwen" in model_name.lower():
158 |                 tokenizer = AutoTokenizer.from_pretrained(model_path)
159 |                 model = CambrianQwenForCausalLM.from_pretrained(
160 |                     model_path, low_cpu_mem_usage=True, **kwargs
161 |                 )
162 |             else:
163 |                 print(f"Loading Cambrian from {model_path}")
164 |                 tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
165 |                 model = CambrianLlamaForCausalLM.from_pretrained(
166 |                     model_path, low_cpu_mem_usage=True, **kwargs
167 |                 )
168 |     else:
169 |         # Load language model
170 |         if model_base is not None:
171 |             # PEFT model
172 |             from peft import PeftModel
173 | 
174 |             tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
175 |             model = AutoModelForCausalLM.from_pretrained(
176 |                 model_base, low_cpu_mem_usage=True, **kwargs
177 |             )
178 |             print(f"Loading LoRA weights from {model_path}")
179 |             model = PeftModel.from_pretrained(model, model_path)
180 |             print(f"Merging weights")
181 |             model = model.merge_and_unload()
182 |             print("Convert to FP16...")
183 |             model.to(torch.float16)
184 |         else:
185 |             use_fast = False
186 |             if "mpt" in model_name.lower():
187 |                 tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
188 |                 model = AutoModelForCausalLM.from_pretrained(
189 |                     model_path, low_cpu_mem_usage=True, trust_remote_code=True, **kwargs
190 |                 )
191 |             else:
192 |                 tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
193 |                 model = AutoModelForCausalLM.from_pretrained(
194 |                     model_path, low_cpu_mem_usage=True, **kwargs
195 |                 )
196 | 
197 |     image_processor = None
198 | 
199 |     if "llava" in model_name.lower():
200 |         mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False)
201 |         mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True)
202 |         if mm_use_im_patch_token:
203 |             tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
204 |         if mm_use_im_start_end:
205 |             tokenizer.add_tokens(
206 |                 [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True
207 |             )
208 |         model.resize_token_embeddings(len(tokenizer))
209 | 
210 |         vision_tower = model.get_vision_tower()
211 |         if not vision_tower.is_loaded:
212 |             try:
213 |                 vision_tower.load_model(device_map=device_map)
214 |             except ValueError:
215 |                 # ClipVisionTower doesn't support loading with device_map 'auto'
216 |                 vision_tower.load_model()
217 |                 vision_tower.to(device="cuda", dtype=torch.float16)
218 |         if device_map != "auto":
219 |             vision_tower.to(device=device_map, dtype=torch.float16)
220 |         image_processor = vision_tower.image_processor
221 |     elif "cambrian" in model_name.lower():
222 |         mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False)
223 |         mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True)
224 |         if mm_use_im_patch_token:
225 |             tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
226 |         if mm_use_im_start_end:
227 |             tokenizer.add_tokens(
228 |                 [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True
229 |             )
230 |         model.resize_token_embeddings(len(tokenizer))
231 | 
232 |         vision_tower_aux_list = model.get_vision_tower_aux_list()
233 | 
234 |         for vision_tower_aux in vision_tower_aux_list:
235 |             if not vision_tower_aux.is_loaded:
236 |                 vision_tower_aux.load_model(device_map=device_map)
237 |             vision_tower_aux.to(device=device, dtype=torch.float16)
238 | 
239 |         image_processor = [
240 |             vision_tower_aux.image_processor
241 |             for vision_tower_aux in vision_tower_aux_list
242 |         ]
243 | 
244 |     if hasattr(model.config, "max_sequence_length"):
245 |         context_len = model.config.max_sequence_length
246 |     else:
247 |         context_len = 2048
248 | 
249 |     return tokenizer, model, image_processor, context_len
250 | 


--------------------------------------------------------------------------------
/longvu/consolidate.py:
--------------------------------------------------------------------------------
 1 | # pyre-unsafe
 2 | """
 3 | Usage:
 4 | python3 -m llava.model.consolidate --src ~/model_weights/llava-7b --dst ~/model_weights/llava-7b_consolidate
 5 | """
 6 | 
 7 | import argparse
 8 | 
 9 | import torch
10 | from transformers import AutoModelForCausalLM, AutoTokenizer
11 | from longvu import *  # noqa
12 | from .utils import auto_upgrade
13 | 
14 | 
15 | def consolidate_ckpt(src_path, dst_path):
16 |     print("Loading model")
17 |     auto_upgrade(src_path)
18 |     src_model = AutoModelForCausalLM.from_pretrained(
19 |         src_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
20 |     )
21 |     src_tokenizer = AutoTokenizer.from_pretrained(src_path, use_fast=False)
22 |     src_model.save_pretrained(dst_path)
23 |     src_tokenizer.save_pretrained(dst_path)
24 | 
25 | 
26 | if __name__ == "__main__":
27 |     parser = argparse.ArgumentParser()
28 |     parser.add_argument("--src", type=str, required=True)
29 |     parser.add_argument("--dst", type=str, required=True)
30 | 
31 |     args = parser.parse_args()
32 | 
33 |     consolidate_ckpt(args.src, args.dst)
34 | 


--------------------------------------------------------------------------------
/longvu/constants.py:
--------------------------------------------------------------------------------
 1 | CONTROLLER_HEART_BEAT_EXPIRATION = 30
 2 | WORKER_HEART_BEAT_INTERVAL = 15
 3 | 
 4 | LOGDIR = "."
 5 | 
 6 | # Model Constants
 7 | IGNORE_INDEX = -100
 8 | IMAGE_TOKEN_INDEX = -200
 9 | DEFAULT_IMAGE_TOKEN = "<image>"
10 | DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
11 | DEFAULT_IM_START_TOKEN = "<im_start>"
12 | DEFAULT_IM_END_TOKEN = "<im_end>"
13 | IMAGE_PLACEHOLDER = "<image-placeholder>"
14 | 


--------------------------------------------------------------------------------
/longvu/file_io.py:
--------------------------------------------------------------------------------
1 | from iopath.common.file_io import HTTPURLHandler, PathManager as PathManagerBase
2 | 
3 | __all__ = ["PathManager"]
4 | 
5 | 
6 | PathManager = PathManagerBase()
7 | PathManager.register_handler(HTTPURLHandler())
8 | 


--------------------------------------------------------------------------------
/longvu/language_model/cambrian_qwen.py:
--------------------------------------------------------------------------------
  1 | #    Copyright 2023 Haotian Liu
  2 | #
  3 | #    Licensed under the Apache License, Version 2.0 (the "License");
  4 | #    you may not use this file except in compliance with the License.
  5 | #    You may obtain a copy of the License at
  6 | #
  7 | #        http://www.apache.org/licenses/LICENSE-2.0
  8 | #
  9 | #    Unless required by applicable law or agreed to in writing, software
 10 | #    distributed under the License is distributed on an "AS IS" BASIS,
 11 | #    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12 | #    See the License for the specific language governing permissions and
 13 | #    limitations under the License.
 14 | 
 15 | 
 16 | from typing import List, Optional, Tuple, Union
 17 | 
 18 | import torch
 19 | import torch.nn as nn
 20 | import torch.nn.functional as F
 21 | from torch.nn import CrossEntropyLoss
 22 | 
 23 | from transformers import AutoConfig, AutoModelForCausalLM
 24 | from transformers.cache_utils import Cache, DynamicCache
 25 | from transformers.generation.utils import GenerateOutput
 26 | 
 27 | from transformers.modeling_outputs import (
 28 |     BaseModelOutputWithPast,
 29 |     CausalLMOutputWithPast,
 30 | )
 31 | from transformers.utils import logging
 32 | 
 33 | from ..cambrian_arch import CambrianMetaForCausalLM, CambrianMetaModel
 34 | 
 35 | IS_XLA_AVAILABLE = False
 36 | 
 37 | from transformers import Qwen2Config, Qwen2ForCausalLM, Qwen2Model
 38 | 
 39 | logger = logging.get_logger(__name__)
 40 | 
 41 | 
 42 | class CambrianConfig(Qwen2Config):
 43 |     model_type = "cambrian_qwen"
 44 | 
 45 |     debug = "debug"
 46 | 
 47 | 
 48 | class CambrianQwenModel(CambrianMetaModel, Qwen2Model):
 49 |     config_class = CambrianConfig
 50 | 
 51 |     def __init__(self, config: Qwen2Config):
 52 |         super(CambrianQwenModel, self).__init__(config)
 53 | 
 54 |     def forward(
 55 |         self,
 56 |         # pyre-fixme[9]: input_ids has type `LongTensor`; used as `None`.
 57 |         input_ids: torch.LongTensor = None,
 58 |         attention_mask: Optional[torch.Tensor] = None,
 59 |         position_ids: Optional[torch.LongTensor] = None,
 60 |         past_key_values: Optional[List[torch.FloatTensor]] = None,
 61 |         inputs_embeds: Optional[torch.FloatTensor] = None,
 62 |         use_cache: Optional[bool] = None,
 63 |         output_attentions: Optional[bool] = None,
 64 |         output_hidden_states: Optional[bool] = None,
 65 |         return_dict: Optional[bool] = None,
 66 |         cache_position: Optional[torch.LongTensor] = None,
 67 |         vision_tower_aux_feature_list: Optional[List[torch.FloatTensor]] = None,
 68 |         vision_tower_aux_attention_masks_list: Optional[List[torch.Tensor]] = None,
 69 |         final_vision_feature_size: Optional[List[tuple]] = None,
 70 |         global_context_feature: Optional[torch.Tensor] = None,
 71 |     ) -> Union[Tuple, BaseModelOutputWithPast]:
 72 |         output_attentions = (
 73 |             output_attentions
 74 |             if output_attentions is not None
 75 |             # pyre-fixme[16]: `CambrianQwenModel` has no attribute `config`.
 76 |             else self.config.output_attentions
 77 |         )
 78 |         output_hidden_states = (
 79 |             output_hidden_states
 80 |             if output_hidden_states is not None
 81 |             else self.config.output_hidden_states
 82 |         )
 83 |         use_cache = use_cache if use_cache is not None else self.config.use_cache
 84 | 
 85 |         return_dict = (
 86 |             return_dict if return_dict is not None else self.config.use_return_dict
 87 |         )
 88 | 
 89 |         if (input_ids is None) ^ (inputs_embeds is not None):
 90 |             raise ValueError(
 91 |                 "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
 92 |             )
 93 | 
 94 |         # pyre-fixme[16]: `CambrianQwenModel` has no attribute `gradient_checkpointing`.
 95 |         # pyre-fixme[16]: `CambrianQwenModel` has no attribute `training`.
 96 |         if self.gradient_checkpointing and self.training:
 97 |             if use_cache:
 98 |                 logger.warning_once(
 99 |                     "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
100 |                 )
101 |                 use_cache = False
102 | 
103 |         use_legacy_cache = False
104 |         if use_cache and not isinstance(past_key_values, Cache):
105 |             use_legacy_cache = True
106 |             # pyre-fixme[6]: For 1st argument expected
107 |             #  `Optional[Tuple[Tuple[FloatTensor]]]` but got
108 |             #  `Optional[List[FloatTensor]]`.
109 |             past_key_values = DynamicCache.from_legacy_cache(past_key_values)
110 |             logger.warning_once(
111 |                 "We detected that you are passing `past_key_values` as a tuple and this is deprecated and will be removed in v4.43. "
112 |                 "Please use an appropriate `Cache` class (https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)"
113 |             )
114 | 
115 |         if inputs_embeds is None:
116 |             # pyre-fixme[16]: `CambrianQwenModel` has no attribute `embed_tokens`.
117 |             inputs_embeds = self.embed_tokens(input_ids)
118 | 
119 |         if cache_position is None:
120 |             past_seen_tokens = (
121 |                 # pyre-fixme[16]: Item `List` of `Union[List[torch._C.FloatTensor],
122 |                 #  DynamicCache]` has no attribute `get_seq_length`.
123 |                 past_key_values.get_seq_length() if past_key_values is not None else 0
124 |             )
125 |             cache_position = torch.arange(
126 |                 past_seen_tokens,
127 |                 past_seen_tokens + inputs_embeds.shape[1],
128 |                 device=inputs_embeds.device,
129 |             )
130 |         if position_ids is None:
131 |             position_ids = cache_position.unsqueeze(0)
132 | 
133 |         # pyre-fixme[16]: `CambrianQwenModel` has no attribute `_update_causal_mask`.
134 |         causal_mask = self._update_causal_mask(
135 |             attention_mask,
136 |             inputs_embeds,
137 |             cache_position,
138 |             past_key_values,
139 |             output_attentions,
140 |         )
141 | 
142 |         hidden_states = inputs_embeds
143 | 
144 |         # decoder layers
145 |         all_hidden_states = () if output_hidden_states else None
146 |         all_self_attns = () if output_attentions else None
147 |         next_decoder_cache = None
148 | 
149 |         # pyre-fixme[16]: `CambrianQwenModel` has no attribute `layers`.
150 |         for i, decoder_layer in enumerate(self.layers):
151 |             if output_hidden_states:
152 |                 all_hidden_states += (hidden_states,)
153 | 
154 |             if self.gradient_checkpointing and self.training:
155 |                 # pyre-fixme[16]: `CambrianQwenModel` has no attribute
156 |                 #  `_gradient_checkpointing_func`.
157 |                 layer_outputs = self._gradient_checkpointing_func(
158 |                     decoder_layer.__call__,
159 |                     hidden_states,
160 |                     causal_mask,
161 |                     position_ids,
162 |                     past_key_values,
163 |                     output_attentions,
164 |                     use_cache,
165 |                     cache_position,
166 |                 )
167 |             else:
168 |                 layer_outputs = decoder_layer(
169 |                     hidden_states,
170 |                     attention_mask=causal_mask,
171 |                     position_ids=position_ids,
172 |                     past_key_value=past_key_values,
173 |                     output_attentions=output_attentions,
174 |                     use_cache=use_cache,
175 |                     cache_position=cache_position,
176 |                 )
177 | 
178 |             hidden_states = layer_outputs[0]
179 | 
180 |             if use_cache:
181 |                 next_decoder_cache = layer_outputs[2 if output_attentions else 1]
182 | 
183 |             if output_attentions:
184 |                 all_self_attns += (layer_outputs[1],)
185 | 
186 |         # pyre-fixme[16]: `CambrianQwenModel` has no attribute `norm`.
187 |         hidden_states = self.norm(hidden_states)
188 | 
189 |         # add hidden states from the last decoder layer
190 |         if output_hidden_states:
191 |             all_hidden_states += (hidden_states,)
192 | 
193 |         next_cache = None
194 |         if use_cache:
195 |             next_cache = (
196 |                 next_decoder_cache.to_legacy_cache()
197 |                 if use_legacy_cache
198 |                 else next_decoder_cache
199 |             )
200 | 
201 |         if not return_dict:
202 |             return tuple(
203 |                 v
204 |                 for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
205 |                 if v is not None
206 |             )
207 |         return BaseModelOutputWithPast(
208 |             last_hidden_state=hidden_states,
209 |             past_key_values=next_cache,
210 |             hidden_states=all_hidden_states,
211 |             attentions=all_self_attns,
212 |         )
213 | 
214 | 
215 | class CambrianQwenForCausalLM(Qwen2ForCausalLM, CambrianMetaForCausalLM):
216 |     config_class = CambrianConfig
217 | 
218 |     def __init__(self, config):
219 |         # super(Qwen2ForCausalLM, self).__init__(config)
220 |         Qwen2ForCausalLM.__init__(self, config)
221 |         config.model_type = "cambrian_qwen"
222 |         config.rope_scaling = None
223 | 
224 |         self.model = CambrianQwenModel(config)
225 |         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
226 |         # Initialize weights and apply final processing
227 |         self.post_init()
228 | 
229 |     def get_model(self):
230 |         return self.model
231 | 
232 |     def forward(
233 |         self,
234 |         # pyre-fixme[9]: input_ids has type `LongTensor`; used as `None`.
235 |         input_ids: torch.LongTensor = None,
236 |         attention_mask: Optional[torch.Tensor] = None,
237 |         position_ids: Optional[torch.LongTensor] = None,
238 |         past_key_values: Optional[List[torch.FloatTensor]] = None,
239 |         inputs_embeds: Optional[torch.FloatTensor] = None,
240 |         labels: Optional[torch.LongTensor] = None,
241 |         use_cache: Optional[bool] = None,
242 |         output_attentions: Optional[bool] = None,
243 |         output_hidden_states: Optional[bool] = None,
244 |         images: Optional[torch.FloatTensor] = None,
245 |         image_aux_attention_masks_list: Optional[List[torch.Tensor]] = None,
246 |         image_sizes: Optional[List[List[int]]] = None,
247 |         return_dict: Optional[bool] = None,
248 |         modalities: Optional[List[str]] = ["image"],
249 |         dpo_forward: Optional[bool] = False,
250 |         cache_position=None,
251 |     ) -> Union[Tuple, CausalLMOutputWithPast]:
252 | 
253 |         input_image_features = None
254 |         highres_image_features = None
255 |         frame_split_sizes = None
256 | 
257 |         if inputs_embeds is None:
258 |             (
259 |                 input_ids,
260 |                 position_ids,
261 |                 attention_mask,
262 |                 past_key_values,
263 |                 inputs_embeds,
264 |                 labels,
265 |                 vision_tower_aux_feature_list,
266 |                 vision_tower_aux_attention_masks_list,
267 |                 final_vision_feature_size,
268 |                 global_context_feature,
269 |             ) = self.prepare_inputs_labels_for_multimodal(
270 |                 input_ids,
271 |                 position_ids,
272 |                 attention_mask,
273 |                 past_key_values,
274 |                 labels,
275 |                 images,
276 |                 image_aux_attention_masks_list,
277 |                 image_sizes,
278 |             )
279 | 
280 |         if dpo_forward:
281 |             # pyre-fixme[29]: `CambrianQwenModel` is not a function.
282 |             outputs = self.model(
283 |                 input_ids=input_ids,
284 |                 attention_mask=attention_mask,
285 |                 position_ids=position_ids,
286 |                 past_key_values=past_key_values,
287 |                 inputs_embeds=inputs_embeds,
288 |                 use_cache=use_cache,
289 |                 output_attentions=output_attentions,
290 |                 output_hidden_states=output_hidden_states,
291 |                 return_dict=return_dict,
292 |             )
293 | 
294 |             hidden_states = outputs[0]
295 |             logits = self.lm_head(hidden_states)
296 |             return logits, labels
297 | 
298 |         else:
299 |             if hasattr(self, "vision_tower_aux_feature_list"):
300 |                 # pyre-fixme[29]: `CambrianQwenModel` is not a function.
301 |                 outputs = self.model(
302 |                     input_ids=input_ids,
303 |                     attention_mask=attention_mask,
304 |                     position_ids=position_ids,
305 |                     past_key_values=past_key_values,
306 |                     inputs_embeds=inputs_embeds,
307 |                     use_cache=use_cache,
308 |                     output_attentions=output_attentions,
309 |                     output_hidden_states=output_hidden_states,
310 |                     return_dict=return_dict,
311 |                     vision_tower_aux_feature_list=(
312 |                         # pyre-fixme[61]: `vision_tower_aux_feature_list` is
313 |                         #  undefined, or not always defined.
314 |                         vision_tower_aux_feature_list
315 |                         if inputs_embeds is None
316 |                         # pyre-fixme[16]: `CambrianQwenForCausalLM` has no attribute
317 |                         #  `vision_tower_aux_feature_list`.
318 |                         else self.vision_tower_aux_feature_list
319 |                     ),
320 |                     vision_tower_aux_attention_masks_list=(
321 |                         # pyre-fixme[61]: `vision_tower_aux_attention_masks_list` is
322 |                         #  undefined, or not always defined.
323 |                         vision_tower_aux_attention_masks_list
324 |                         if inputs_embeds is None
325 |                         # pyre-fixme[16]: `CambrianQwenForCausalLM` has no attribute
326 |                         #  `vision_tower_aux_attention_masks_list`.
327 |                         else self.vision_tower_aux_attention_masks_list
328 |                     ),
329 |                     final_vision_feature_size=(
330 |                         # pyre-fixme[61]: `final_vision_feature_size` is undefined,
331 |                         #  or not always defined.
332 |                         final_vision_feature_size
333 |                         if inputs_embeds is None
334 |                         # pyre-fixme[16]: `CambrianQwenForCausalLM` has no attribute
335 |                         #  `final_vision_feature_size`.
336 |                         else self.final_vision_feature_size
337 |                     ),
338 |                     global_context_feature=(
339 |                         # pyre-fixme[61]: `global_context_feature` is undefined, or
340 |                         #  not always defined.
341 |                         global_context_feature
342 |                         if inputs_embeds is None
343 |                         # pyre-fixme[16]: `CambrianQwenForCausalLM` has no attribute
344 |                         #  `global_context_feature`.
345 |                         else self.global_context_feature
346 |                     ),
347 |                 )
348 |             else:
349 |                 # pyre-fixme[29]: `CambrianQwenModel` is not a function.
350 |                 outputs = self.model(
351 |                     input_ids=input_ids,
352 |                     attention_mask=attention_mask,
353 |                     position_ids=position_ids,
354 |                     past_key_values=past_key_values,
355 |                     inputs_embeds=inputs_embeds,
356 |                     use_cache=use_cache,
357 |                     output_attentions=output_attentions,
358 |                     output_hidden_states=output_hidden_states,
359 |                     return_dict=return_dict,
360 |                     # final_vision_feature_size=final_vision_feature_size,
361 |                 )
362 | 
363 |             hidden_states = outputs[0]
364 |             logits = self.lm_head(hidden_states)
365 |             logits = logits.float()
366 | 
367 |             loss = None
368 |             if labels is not None:
369 |                 # Shift so that tokens < n predict n
370 |                 shift_logits = logits[..., :-1, :].contiguous()
371 |                 shift_labels = labels[..., 1:].contiguous()
372 |                 # Flatten the tokens
373 |                 loss_fct = CrossEntropyLoss()
374 |                 # pyre-fixme[16]: `CambrianQwenForCausalLM` has no attribute `config`.
375 |                 shift_logits = shift_logits.view(-1, self.config.vocab_size)
376 |                 shift_labels = shift_labels.view(-1)
377 |                 # Enable model parallelism
378 |                 shift_labels = shift_labels.to(shift_logits.device)
379 |                 loss = loss_fct(shift_logits, shift_labels)
380 | 
381 |             if not return_dict:
382 |                 output = (logits,) + outputs[1:]
383 |                 return (loss,) + output if loss is not None else output
384 | 
385 |             return CausalLMOutputWithPast(
386 |                 loss=loss,
387 |                 logits=logits,
388 |                 past_key_values=outputs.past_key_values,
389 |                 hidden_states=outputs.hidden_states,
390 |                 attentions=outputs.attentions,
391 |             )
392 | 
393 |     @torch.no_grad()
394 |     def generate(
395 |         self,
396 |         inputs: Optional[torch.Tensor] = None,
397 |         images: Optional[torch.Tensor] = None,
398 |         image_sizes: Optional[torch.Tensor] = None,
399 |         **kwargs,
400 |     ) -> Union[GenerateOutput, torch.LongTensor]:
401 |         position_ids = kwargs.pop("position_ids", None)
402 |         attention_mask = kwargs.pop("attention_mask", None)
403 |         if "inputs_embeds" in kwargs:
404 |             raise NotImplementedError("`inputs_embeds` is not supported")
405 | 
406 |         if images is not None:
407 |             (
408 |                 inputs,
409 |                 position_ids,
410 |                 attention_mask,
411 |                 _,
412 |                 inputs_embeds,
413 |                 _,
414 |                 vision_tower_aux_feature_list,
415 |                 vision_tower_aux_attention_masks_list,
416 |                 final_vision_feature_size,
417 |                 global_context_feature,
418 |             ) = self.prepare_inputs_labels_for_multimodal(
419 |                 inputs,
420 |                 position_ids,
421 |                 attention_mask,
422 |                 None,
423 |                 None,
424 |                 images,
425 |                 image_sizes=image_sizes,
426 |             )
427 |             # pyre-fixme[16]: `CambrianQwenForCausalLM` has no attribute
428 |             #  `vision_tower_aux_feature_list`.
429 |             self.vision_tower_aux_feature_list = vision_tower_aux_feature_list
430 |             # pyre-fixme[16]: `CambrianQwenForCausalLM` has no attribute
431 |             #  `vision_tower_aux_attention_masks_list`.
432 |             self.vision_tower_aux_attention_masks_list = (
433 |                 vision_tower_aux_attention_masks_list
434 |             )
435 |             # pyre-fixme[16]: `CambrianQwenForCausalLM` has no attribute
436 |             #  `final_vision_feature_size`.
437 |             self.final_vision_feature_size = final_vision_feature_size
438 |             # pyre-fixme[16]: `CambrianQwenForCausalLM` has no attribute
439 |             #  `global_context_feature`.
440 |             self.global_context_feature = global_context_feature
441 |         else:
442 |             inputs_embeds = self.get_model().embed_tokens(inputs)
443 | 
444 |         # pyre-fixme[16]: `Qwen2ForCausalLM` has no attribute `generate`.
445 |         return super().generate(
446 |             position_ids=position_ids,
447 |             attention_mask=attention_mask,
448 |             inputs_embeds=inputs_embeds,
449 |             **kwargs,
450 |         )
451 | 
452 |     def prepare_inputs_for_generation(
453 |         self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs
454 |     ):
455 |         images = kwargs.pop("images", None)
456 |         image_sizes = kwargs.pop("image_sizes", None)
457 |         inputs = super().prepare_inputs_for_generation(
458 |             input_ids,
459 |             past_key_values=past_key_values,
460 |             inputs_embeds=inputs_embeds,
461 |             **kwargs,
462 |         )
463 |         if images is not None:
464 |             inputs["images"] = images
465 |         if image_sizes is not None:
466 |             inputs["image_sizes"] = image_sizes
467 |         return inputs
468 | 
469 | 
470 | AutoConfig.register("cambrian_qwen", CambrianConfig)
471 | AutoModelForCausalLM.register(CambrianConfig, CambrianQwenForCausalLM)
472 | 


--------------------------------------------------------------------------------
/longvu/make_delta.py:
--------------------------------------------------------------------------------
 1 | # pyre-unsafe
 2 | """
 3 | Usage:
 4 | python3 -m llava.model.make_delta --base ~/model_weights/llama-7b --target ~/model_weights/llava-7b --delta ~/model_weights/llava-7b-delta --hub-repo-id liuhaotian/llava-7b-delta
 5 | """
 6 | 
 7 | import argparse
 8 | 
 9 | import torch
10 | from tqdm import tqdm
11 | from transformers import AutoModelForCausalLM, AutoTokenizer
12 | 
13 | from .utils import auto_upgrade
14 | 
15 | 
16 | def make_delta(base_model_path, target_model_path, delta_path, hub_repo_id):
17 |     print("Loading base model")
18 |     base = AutoModelForCausalLM.from_pretrained(
19 |         base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
20 |     )
21 | 
22 |     print("Loading target model")
23 |     auto_upgrade(target_model_path)
24 |     target = AutoModelForCausalLM.from_pretrained(
25 |         target_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
26 |     )
27 | 
28 |     print("Calculating delta")
29 |     for name, param in tqdm(target.state_dict().items(), desc="Calculating delta"):
30 |         if name not in base.state_dict():
31 |             assert name in [
32 |                 "model.mm_projector.weight",
33 |                 "model.mm_projector.bias",
34 |             ], f"{name} not in base model"
35 |             continue
36 |         if param.data.shape == base.state_dict()[name].shape:
37 |             param.data -= base.state_dict()[name]
38 |         else:
39 |             assert name in [
40 |                 "model.embed_tokens.weight",
41 |                 "lm_head.weight",
42 |             ], f"{name} dimension mismatch: {param.data.shape} vs {base.state_dict()[name].shape}"
43 |             bparam = base.state_dict()[name]
44 |             param.data[: bparam.shape[0], : bparam.shape[1]] -= bparam
45 | 
46 |     print("Saving delta")
47 |     if hub_repo_id:
48 |         kwargs = {"push_to_hub": True, "repo_id": hub_repo_id}
49 |     else:
50 |         kwargs = {}
51 |     target.save_pretrained(delta_path, **kwargs)
52 |     target_tokenizer = AutoTokenizer.from_pretrained(target_model_path)
53 |     target_tokenizer.save_pretrained(delta_path, **kwargs)
54 | 
55 | 
56 | if __name__ == "__main__":
57 |     parser = argparse.ArgumentParser()
58 |     parser.add_argument("--base-model-path", type=str, required=True)
59 |     parser.add_argument("--target-model-path", type=str, required=True)
60 |     parser.add_argument("--delta-path", type=str, required=True)
61 |     parser.add_argument("--hub-repo-id", type=str, default=None)
62 |     args = parser.parse_args()
63 | 
64 |     make_delta(
65 |         args.base_model_path, args.target_model_path, args.delta_path, args.hub_repo_id
66 |     )
67 | 


--------------------------------------------------------------------------------
/longvu/mm_trainer.py:
--------------------------------------------------------------------------------
  1 | # pyre-strict
  2 | import os
  3 | from typing import Dict, List, Optional
  4 | 
  5 | import torch
  6 | import torch.nn as nn
  7 | 
  8 | from longvu.mm_datautils import get_mm_adapter_state_maybe_zero_3
  9 | from torch.utils.data import DataLoader, Sampler
 10 | 
 11 | from transformers import Trainer
 12 | 
 13 | from transformers.trainer import ALL_LAYERNORM_LAYERS, get_parameter_names, has_length
 14 | 
 15 | 
 16 | # pyre-fixme[3]: Return type must be annotated.
 17 | # pyre-fixme[2]: Parameter must be annotated.
 18 | def split_to_even_chunks(indices, lengths, num_chunks):
 19 |     """
 20 |     Split a list of indices into `chunks` chunks of roughly equal lengths.
 21 |     """
 22 | 
 23 |     if len(indices) % num_chunks != 0:
 24 |         return [indices[i::num_chunks] for i in range(num_chunks)]
 25 | 
 26 |     num_indices_per_chunk = len(indices) // num_chunks
 27 | 
 28 |     chunks = [[] for _ in range(num_chunks)]
 29 |     chunks_lengths = [0 for _ in range(num_chunks)]
 30 |     for index in indices:
 31 |         shortest_chunk = chunks_lengths.index(min(chunks_lengths))
 32 |         chunks[shortest_chunk].append(index)
 33 |         chunks_lengths[shortest_chunk] += lengths[index]
 34 |         if len(chunks[shortest_chunk]) == num_indices_per_chunk:
 35 |             chunks_lengths[shortest_chunk] = float("inf")
 36 | 
 37 |     return chunks
 38 | 
 39 | 
 40 | # pyre-fixme[3]: Return type must be annotated.
 41 | # pyre-fixme[2]: Parameter must be annotated.
 42 | def get_length_grouped_indices(
 43 |     lengths, batch_size, world_size, generator=None, merge=True
 44 | ):
 45 |     # We need to use torch for the random part as a distributed sampler will set the random seed for torch.
 46 |     indices = torch.randperm(len(lengths), generator=generator)
 47 |     megabatch_size = world_size * batch_size
 48 |     megabatches = [
 49 |         indices[i : i + megabatch_size].tolist()
 50 |         for i in range(0, len(lengths), megabatch_size)
 51 |     ]
 52 |     megabatches = [
 53 |         sorted(megabatch, key=lambda i: lengths[i], reverse=True)
 54 |         for megabatch in megabatches
 55 |     ]
 56 |     megabatches = [
 57 |         split_to_even_chunks(megabatch, lengths, world_size)
 58 |         for megabatch in megabatches
 59 |     ]
 60 | 
 61 |     return [i for megabatch in megabatches for batch in megabatch for i in batch]
 62 | 
 63 | 
 64 | # pyre-fixme[3]: Return type must be annotated.
 65 | # pyre-fixme[2]: Parameter must be annotated.
 66 | def get_modality_length_grouped_indices(
 67 |     lengths, batch_size, world_size, generator=None
 68 | ):
 69 |     # We need to use torch for the random part as a distributed sampler will set the random seed for torch.
 70 |     assert all(l != 0 for l in lengths), "Should not have zero length."
 71 |     if all(l > 0 for l in lengths) or all(l < 0 for l in lengths):
 72 |         # all samples are in the same modality
 73 |         return get_length_grouped_indices(
 74 |             lengths, batch_size, world_size, generator=generator
 75 |         )
 76 |     mm_indices, mm_lengths = zip(*[(i, l) for i, l in enumerate(lengths) if l > 0])
 77 |     lang_indices, lang_lengths = zip(*[(i, -l) for i, l in enumerate(lengths) if l < 0])
 78 | 
 79 |     mm_shuffle = [
 80 |         mm_indices[i]
 81 |         for i in get_length_grouped_indices(
 82 |             mm_lengths, batch_size, world_size, generator=None
 83 |         )
 84 |     ]
 85 |     lang_shuffle = [
 86 |         lang_indices[i]
 87 |         for i in get_length_grouped_indices(
 88 |             lang_lengths, batch_size, world_size, generator=None
 89 |         )
 90 |     ]
 91 |     megabatch_size = world_size * batch_size
 92 |     mm_megabatches = [
 93 |         mm_shuffle[i : i + megabatch_size]
 94 |         for i in range(0, len(mm_shuffle), megabatch_size)
 95 |     ]
 96 |     lang_megabatches = [
 97 |         lang_shuffle[i : i + megabatch_size]
 98 |         for i in range(0, len(lang_shuffle), megabatch_size)
 99 |     ]
100 | 
101 |     last_mm = mm_megabatches[-1]
102 |     last_lang = lang_megabatches[-1]
103 |     additional_batch = last_mm + last_lang
104 |     megabatches = mm_megabatches[:-1] + lang_megabatches[:-1]
105 |     megabatch_indices = torch.randperm(len(megabatches), generator=generator)
106 |     megabatches = [megabatches[i] for i in megabatch_indices]
107 | 
108 |     if len(additional_batch) > 0:
109 |         megabatches.append(sorted(additional_batch))
110 | 
111 |     return [i for megabatch in megabatches for i in megabatch]
112 | 
113 | 
114 | # pyre-fixme[3]: Return type must be annotated.
115 | # pyre-fixme[2]: Parameter must be annotated.
116 | class LengthGroupedSampler(Sampler):
117 |     r"""
118 |     Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while
119 |     keeping a bit of randomness.
120 |     """
121 | 
122 |     def __init__(
123 |         self,
124 |         batch_size: int,
125 |         world_size: int,
126 |         lengths: Optional[List[int]] = None,
127 |         generator=None,
128 |         group_by_modality: bool = False,
129 |     ):
130 |         if lengths is None:
131 |             raise ValueError("Lengths must be provided.")
132 | 
133 |         self.batch_size = batch_size
134 |         self.world_size = world_size
135 |         self.lengths = lengths
136 |         self.generator = generator
137 |         self.group_by_modality = group_by_modality
138 | 
139 |     def __len__(self):
140 |         return len(self.lengths)
141 | 
142 |     def __iter__(self):
143 |         if self.group_by_modality:
144 |             indices = get_modality_length_grouped_indices(
145 |                 self.lengths, self.batch_size, self.world_size, generator=self.generator
146 |             )
147 |         else:
148 |             indices = get_length_grouped_indices(
149 |                 self.lengths, self.batch_size, self.world_size, generator=self.generator
150 |             )
151 |         return iter(indices)
152 | 
153 | 
154 | def get_mm_adapter_state_maybe_zero_3(
155 |     named_params: Dict[str, torch.Tensor], keys_to_match: List[str]
156 | ) -> Dict[str, torch.Tensor]:
157 |     to_return = {
158 |         k: t
159 |         for k, t in named_params
160 |         if any(key_match in k for key_match in keys_to_match)
161 |     }
162 |     to_return = {
163 |         k: maybe_zero_3(v, ignore_status=True, name=k).cpu()  # pyre-ignore
164 |         for k, v in to_return.items()
165 |     }
166 |     return to_return
167 | 
168 | 
169 | def maybe_zero_3(
170 |     param: torch.Tensor, ignore_status: bool = False, name: Optional[str] = None
171 | ) -> torch.Tensor:
172 |     return param.detach().cpu().clone()
173 | 
174 | 
175 | class LLaVATrainer(Trainer):
176 |     def __init__(
177 |         self,
178 |         train_dataloader: Optional[DataLoader] = None,
179 |         # pyre-fixme[2]: Parameter must be annotated.
180 |         *args,
181 |         # pyre-fixme[2]: Parameter must be annotated.
182 |         **kwargs,
183 |     ) -> None:
184 |         self.train_dataloader = train_dataloader
185 |         super().__init__(*args, **kwargs)
186 | 
187 |     def get_train_dataloader(self) -> DataLoader:
188 |         if self.train_dataloader is not None:
189 |             print("Using sonic dataloader")
190 |             # pyre-fixme[16]: `LLaVATrainer` has no attribute `accelerator`.
191 |             return self.accelerator.prepare(self.train_dataloader)
192 |         # pyre-fixme[16]: `Trainer` has no attribute `get_train_dataloader`.
193 |         return super().get_train_dataloader()
194 | 
195 |     def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
196 |         # pyre-fixme[16]: `LLaVATrainer` has no attribute `train_dataset`.
197 |         if self.train_dataset is None or not has_length(self.train_dataset):
198 |             return None
199 | 
200 |         # pyre-fixme[16]: `LLaVATrainer` has no attribute `args`.
201 |         if self.args.group_by_modality_length:
202 |             lengths = self.train_dataset.modality_lengths
203 |             return LengthGroupedSampler(
204 |                 # self.args.train_batch_size * self.args.gradient_accumulation_steps, # TODO: seems that we should not have gradient_accumulation_steps
205 |                 self.args.train_batch_size,
206 |                 world_size=self.args.world_size,
207 |                 lengths=lengths,
208 |                 group_by_modality=True,
209 |             )
210 |         else:
211 |             # pyre-fixme[16]: `Trainer` has no attribute `_get_train_sampler`.
212 |             return super()._get_train_sampler()
213 | 
214 |     # pyre-fixme[3]: Return type must be annotated.
215 |     def create_optimizer(self):
216 |         """
217 |         Setup the optimizer.
218 | 
219 |         We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
220 |         Trainer's init through `optimizers`, or subclass and override this method in a subclass.
221 |         """
222 |         # pyre-fixme[16]: `Trainer` has no attribute `model`.
223 |         opt_model = self.model
224 |         # if self.args.unfreeze_mm_vision_tower:
225 |         #     opt_model.get_model().vision_tower_aux_list = nn.ModuleList(opt_model.get_vision_tower_aux_list())
226 |         #     self.param_to_name = map_params_to_module_names([opt_model])
227 |         # pyre-fixme[16]: `Trainer` has no attribute `optimizer`.
228 |         if self.optimizer is None:
229 |             decay_parameters = get_parameter_names(opt_model, ALL_LAYERNORM_LAYERS)
230 |             decay_parameters = [name for name in decay_parameters if "bias" not in name]
231 |             # pyre-fixme[16]: `Trainer` has no attribute `mm_projector_lr`.
232 |             assert not (self.args.mm_projector_lr and self.args.mm_vision_sampler_lr)
233 |             if self.args.mm_projector_lr is not None:
234 |                 projector_parameters = [
235 |                     name
236 |                     for name, _ in opt_model.named_parameters()
237 |                     if "mm_projector" in name
238 |                 ]
239 |                 optimizer_grouped_parameters = [
240 |                     {
241 |                         "params": [
242 |                             p
243 |                             for n, p in opt_model.named_parameters()
244 |                             if (
245 |                                 n in decay_parameters
246 |                                 and n not in projector_parameters
247 |                                 and p.requires_grad
248 |                             )
249 |                         ],
250 |                         "weight_decay": self.args.weight_decay,
251 |                     },
252 |                     {
253 |                         "params": [
254 |                             p
255 |                             for n, p in opt_model.named_parameters()
256 |                             if (
257 |                                 n not in decay_parameters
258 |                                 and n not in projector_parameters
259 |                                 and p.requires_grad
260 |                             )
261 |                         ],
262 |                         "weight_decay": 0.0,
263 |                     },
264 |                     {
265 |                         "params": [
266 |                             p
267 |                             for n, p in opt_model.named_parameters()
268 |                             if (
269 |                                 n in decay_parameters
270 |                                 and n in projector_parameters
271 |                                 and p.requires_grad
272 |                             )
273 |                         ],
274 |                         "weight_decay": self.args.weight_decay,
275 |                         "lr": self.args.mm_projector_lr,
276 |                     },
277 |                     {
278 |                         "params": [
279 |                             p
280 |                             for n, p in opt_model.named_parameters()
281 |                             if (
282 |                                 n not in decay_parameters
283 |                                 and n in projector_parameters
284 |                                 and p.requires_grad
285 |                             )
286 |                         ],
287 |                         "weight_decay": 0.0,
288 |                         "lr": self.args.mm_projector_lr,
289 |                     },
290 |                 ]
291 |             elif self.args.mm_vision_sampler_lr is not None:
292 |                 vision_sampler_parameters = [
293 |                     name
294 |                     for name, _ in opt_model.named_parameters()
295 |                     if ("vision_sampler" in name) or ("vision_query" in name)
296 |                 ]
297 |                 optimizer_grouped_parameters = [
298 |                     {
299 |                         "params": [
300 |                             p
301 |                             for n, p in opt_model.named_parameters()
302 |                             if (
303 |                                 n in decay_parameters
304 |                                 and n not in vision_sampler_parameters
305 |                                 and p.requires_grad
306 |                             )
307 |                         ],
308 |                         "weight_decay": self.args.weight_decay,
309 |                     },
310 |                     {
311 |                         "params": [
312 |                             p
313 |                             for n, p in opt_model.named_parameters()
314 |                             if (
315 |                                 n not in decay_parameters
316 |                                 and n not in vision_sampler_parameters
317 |                                 and p.requires_grad
318 |                             )
319 |                         ],
320 |                         "weight_decay": 0.0,
321 |                     },
322 |                     {
323 |                         "params": [
324 |                             p
325 |                             for n, p in opt_model.named_parameters()
326 |                             if (
327 |                                 n in decay_parameters
328 |                                 and n in vision_sampler_parameters
329 |                                 and p.requires_grad
330 |                             )
331 |                         ],
332 |                         "weight_decay": self.args.weight_decay,
333 |                         "lr": self.args.mm_vision_sampler_lr,
334 |                     },
335 |                     {
336 |                         "params": [
337 |                             p
338 |                             for n, p in opt_model.named_parameters()
339 |                             if (
340 |                                 n not in decay_parameters
341 |                                 and n in vision_sampler_parameters
342 |                                 and p.requires_grad
343 |                             )
344 |                         ],
345 |                         "weight_decay": 0.0,
346 |                         "lr": self.args.mm_vision_sampler_lr,
347 |                     },
348 |                 ]
349 |             elif (
350 |                 self.args.unfreeze_mm_vision_tower
351 |                 and self.args.mm_vision_tower_lr is not None
352 |             ):
353 |                 vision_tower_parameters = [
354 |                     name
355 |                     for name, _ in opt_model.named_parameters()
356 |                     if "vision_tower" in name
357 |                 ]
358 |                 optimizer_grouped_parameters = [
359 |                     {
360 |                         "params": [
361 |                             p
362 |                             for n, p in opt_model.named_parameters()
363 |                             if (
364 |                                 n in decay_parameters
365 |                                 and n not in vision_tower_parameters
366 |                                 and p.requires_grad
367 |                             )
368 |                         ],
369 |                         "weight_decay": self.args.weight_decay,
370 |                     },
371 |                     {
372 |                         "params": [
373 |                             p
374 |                             for n, p in opt_model.named_parameters()
375 |                             if (
376 |                                 n not in decay_parameters
377 |                                 and n not in vision_tower_parameters
378 |                                 and p.requires_grad
379 |                             )
380 |                         ],
381 |                         "weight_decay": 0.0,
382 |                     },
383 |                     {
384 |                         "params": [
385 |                             p
386 |                             for n, p in opt_model.named_parameters()
387 |                             if (
388 |                                 n in decay_parameters
389 |                                 and n in vision_tower_parameters
390 |                                 and p.requires_grad
391 |                             )
392 |                         ],
393 |                         "weight_decay": self.args.weight_decay,
394 |                         "lr": self.args.mm_vision_tower_lr,
395 |                     },
396 |                     {
397 |                         "params": [
398 |                             p
399 |                             for n, p in opt_model.named_parameters()
400 |                             if (
401 |                                 n not in decay_parameters
402 |                                 and n in vision_tower_parameters
403 |                                 and p.requires_grad
404 |                             )
405 |                         ],
406 |                         "weight_decay": 0.0,
407 |                         "lr": self.args.mm_vision_tower_lr,
408 |                     },
409 |                 ]
410 |             else:
411 |                 optimizer_grouped_parameters = [
412 |                     {
413 |                         "params": [
414 |                             p
415 |                             for n, p in opt_model.named_parameters()
416 |                             if (n in decay_parameters and p.requires_grad)
417 |                         ],
418 |                         "weight_decay": self.args.weight_decay,
419 |                     },
420 |                     {
421 |                         "params": [
422 |                             p
423 |                             for n, p in opt_model.named_parameters()
424 |                             if (n not in decay_parameters and p.requires_grad)
425 |                         ],
426 |                         "weight_decay": 0.0,
427 |                     },
428 |                 ]
429 | 
430 |             optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(
431 |                 self.args
432 |             )
433 | 
434 |             self.optimizer = optimizer_cls(
435 |                 optimizer_grouped_parameters, **optimizer_kwargs
436 |             )
437 |             if optimizer_cls.__name__ == "Adam8bit":
438 |                 import bitsandbytes
439 | 
440 |                 manager = bitsandbytes.optim.GlobalOptimManager.get_instance()
441 | 
442 |                 skipped = 0
443 |                 for module in opt_model.modules():
444 |                     if isinstance(module, nn.Embedding):
445 |                         skipped += sum(
446 |                             {
447 |                                 p.data_ptr(): p.numel() for p in module.parameters()
448 |                             }.values()
449 |                         )
450 |                         manager.register_module_override(
451 |                             module, "weight", {"optim_bits": 32}
452 |                         )
453 |         return self.optimizer
454 | 
455 |     # pyre-fixme[2]: Parameter must be annotated.
456 |     def _save_checkpoint(self, model, trial, metrics=None) -> None:
457 |         # pyre-fixme[16]: `LLaVATrainer` has no attribute `args`.
458 |         if getattr(self.args, "tune_mm_mlp_adapter", False):
459 | 
460 |             from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
461 | 
462 |             # pyre-fixme[16]: `LLaVATrainer` has no attribute `state`.
463 |             checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"
464 | 
465 |             # pyre-fixme[16]: `LLaVATrainer` has no attribute `_get_output_dir`.
466 |             run_dir = self._get_output_dir(trial=trial)
467 |             output_dir = os.path.join(run_dir, checkpoint_folder)
468 | 
469 |             # Only save Adapter
470 |             keys_to_match = ["mm_projector", "vision_resampler"]
471 |             if getattr(self.args, "use_im_start_end", False):
472 |                 keys_to_match.extend(["embed_tokens", "embed_in"])
473 | 
474 |             weight_to_save = get_mm_adapter_state_maybe_zero_3(
475 |                 # pyre-fixme[16]: `LLaVATrainer` has no attribute `model`.
476 |                 self.model.named_parameters(),
477 |                 keys_to_match,
478 |             )
479 | 
480 |             if self.args.local_rank == 0 or self.args.local_rank == -1:
481 |                 self.model.config.save_pretrained(output_dir)
482 |                 torch.save(weight_to_save, os.path.join(output_dir, "mm_projector.bin"))
483 |         else:
484 |             # pyre-fixme[16]: `Trainer` has no attribute `_save_checkpoint`.
485 |             super(LLaVATrainer, self)._save_checkpoint(model, trial, metrics)
486 | 
487 |     # pyre-fixme[2]: Parameter must be annotated.
488 |     def _save(self, output_dir: Optional[str] = None, state_dict=None) -> None:
489 |         # pyre-fixme[16]: `LLaVATrainer` has no attribute `args`.
490 |         if getattr(self.args, "tune_mm_mlp_adapter", False):
491 |             pass
492 |         else:
493 |             # pyre-fixme[16]: `Trainer` has no attribute `_save`.
494 |             super(LLaVATrainer, self)._save(output_dir, state_dict)
495 | 


--------------------------------------------------------------------------------
/longvu/mm_utils.py:
--------------------------------------------------------------------------------
  1 | import ast
  2 | import base64
  3 | import math
  4 | from io import BytesIO
  5 | 
  6 | import torch
  7 | from longvu.constants import IMAGE_TOKEN_INDEX
  8 | from PIL import Image
  9 | 
 10 | from transformers import StoppingCriteria
 11 | 
 12 | 
 13 | def select_best_resolution(original_size, possible_resolutions):
 14 |     """
 15 |     Selects the best resolution from a list of possible resolutions based on the original size.
 16 | 
 17 |     Args:
 18 |         original_size (tuple): The original size of the image in the format (width, height).
 19 |         possible_resolutions (list): A list of possible resolutions in the format [(width1, height1), (width2, height2), ...].
 20 | 
 21 |     Returns:
 22 |         tuple: The best fit resolution in the format (width, height).
 23 |     """
 24 |     original_width, original_height = original_size
 25 |     best_fit = None
 26 |     max_effective_resolution = 0
 27 |     min_wasted_resolution = float("inf")
 28 | 
 29 |     for width, height in possible_resolutions:
 30 |         scale = min(width / original_width, height / original_height)
 31 |         downscaled_width, downscaled_height = int(original_width * scale), int(
 32 |             original_height * scale
 33 |         )
 34 |         effective_resolution = min(
 35 |             downscaled_width * downscaled_height, original_width * original_height
 36 |         )
 37 |         wasted_resolution = (width * height) - effective_resolution
 38 | 
 39 |         if effective_resolution > max_effective_resolution or (
 40 |             effective_resolution == max_effective_resolution
 41 |             and wasted_resolution < min_wasted_resolution
 42 |         ):
 43 |             max_effective_resolution = effective_resolution
 44 |             min_wasted_resolution = wasted_resolution
 45 |             best_fit = (width, height)
 46 | 
 47 |     return best_fit
 48 | 
 49 | 
 50 | def resize_and_pad_image(image, target_resolution):
 51 |     """
 52 |     Resize and pad an image to a target resolution while maintaining aspect ratio.
 53 | 
 54 |     Args:
 55 |         image (PIL.Image.Image): The input image.
 56 |         target_resolution (tuple): The target resolution (width, height) of the image.
 57 | 
 58 |     Returns:
 59 |         PIL.Image.Image: The resized and padded image.
 60 |     """
 61 |     original_width, original_height = image.size
 62 |     target_width, target_height = target_resolution
 63 | 
 64 |     scale_w = target_width / original_width
 65 |     scale_h = target_height / original_height
 66 | 
 67 |     if scale_w < scale_h:
 68 |         new_width = target_width
 69 |         new_height = min(math.ceil(original_height * scale_w), target_height)
 70 |     else:
 71 |         new_height = target_height
 72 |         new_width = min(math.ceil(original_width * scale_h), target_width)
 73 | 
 74 |     # Resize the image
 75 |     resized_image = image.resize((new_width, new_height))
 76 | 
 77 |     new_image = Image.new("RGB", (target_width, target_height), (0, 0, 0))
 78 |     paste_x = (target_width - new_width) // 2
 79 |     paste_y = (target_height - new_height) // 2
 80 |     new_image.paste(resized_image, (paste_x, paste_y))
 81 | 
 82 |     return new_image
 83 | 
 84 | 
 85 | def divide_to_patches(image, patch_size):
 86 |     """
 87 |     Divides an image into patches of a specified size.
 88 | 
 89 |     Args:
 90 |         image (PIL.Image.Image): The input image.
 91 |         patch_size (int): The size of each patch.
 92 | 
 93 |     Returns:
 94 |         list: A list of PIL.Image.Image objects representing the patches.
 95 |     """
 96 |     patches = []
 97 |     width, height = image.size
 98 |     for i in range(0, height, patch_size):
 99 |         for j in range(0, width, patch_size):
100 |             box = (j, i, j + patch_size, i + patch_size)
101 |             patch = image.crop(box)
102 |             patches.append(patch)
103 | 
104 |     return patches
105 | 
106 | 
107 | def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size):
108 |     """
109 |     Calculate the shape of the image patch grid after the preprocessing for images of any resolution.
110 | 
111 |     Args:
112 |         image_size (tuple): The size of the input image in the format (width, height).
113 |         grid_pinpoints (str): A string representation of a list of possible resolutions.
114 |         patch_size (int): The size of each image patch.
115 | 
116 |     Returns:
117 |         tuple: The shape of the image patch grid in the format (width, height).
118 |     """
119 |     if type(grid_pinpoints) is list:
120 |         possible_resolutions = grid_pinpoints
121 |     else:
122 |         possible_resolutions = ast.literal_eval(grid_pinpoints)
123 |     width, height = select_best_resolution(image_size, possible_resolutions)
124 |     return width // patch_size, height // patch_size
125 | 
126 | 
127 | def process_anyres_image(image, processor, grid_pinpoints):
128 |     """
129 |     Process an image with variable resolutions.
130 | 
131 |     Args:
132 |         image (PIL.Image.Image): The input image to be processed.
133 |         processor: The image processor object.
134 |         grid_pinpoints (str): A string representation of a list of possible resolutions.
135 | 
136 |     Returns:
137 |         torch.Tensor: A tensor containing the processed image patches.
138 |     """
139 |     if type(grid_pinpoints) is list:
140 |         possible_resolutions = grid_pinpoints
141 |     else:
142 |         possible_resolutions = ast.literal_eval(grid_pinpoints)
143 |     best_resolution = select_best_resolution(image.size, possible_resolutions)
144 |     image_padded = resize_and_pad_image(image, best_resolution)
145 | 
146 |     patches = divide_to_patches(image_padded, processor.crop_size["height"])
147 | 
148 |     image_original_resize = image.resize(
149 |         (processor.size["shortest_edge"], processor.size["shortest_edge"])
150 |     )
151 | 
152 |     image_patches = [image_original_resize] + patches
153 |     image_patches = [
154 |         processor.preprocess(image_patch, return_tensors="pt")["pixel_values"][0]
155 |         for image_patch in image_patches
156 |     ]
157 |     return torch.stack(image_patches, dim=0)
158 | 
159 | 
160 | def load_image_from_base64(image):
161 |     return Image.open(BytesIO(base64.b64decode(image)))
162 | 
163 | 
164 | def expand2square(pil_img, background_color):
165 |     width, height = pil_img.size
166 |     if width == height:
167 |         return pil_img
168 |     elif width > height:
169 |         result = Image.new(pil_img.mode, (width, width), background_color)
170 |         result.paste(pil_img, (0, (width - height) // 2))
171 |         return result
172 |     else:
173 |         result = Image.new(pil_img.mode, (height, height), background_color)
174 |         result.paste(pil_img, ((height - width) // 2, 0))
175 |         return result
176 | 
177 | 
178 | # def process_images(images, image_processor, model_cfg):
179 | #     image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)
180 | #     new_images = []
181 | #     if image_aspect_ratio == 'pad':
182 | #         for image in images:
183 | #             image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
184 | #             image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
185 | #             new_images.append(image)
186 | #     elif image_aspect_ratio == "anyres":
187 | #         for image in images:
188 | #             image = process_anyres_image(image, image_processor, model_cfg.image_grid_pinpoints)
189 | #             new_images.append(image)
190 | #     else:
191 | #         return image_processor(images, return_tensors='pt')['pixel_values']
192 | #     if all(x.shape == new_images[0].shape for x in new_images):
193 | #         new_images = torch.stack(new_images, dim=0)
194 | #     return new_images
195 | 
196 | 
197 | # multiple vision towers
198 | def process_images(images, image_processor, model_cfg):
199 |     processor_aux_list = image_processor
200 |     new_images_aux_list = []
201 |     for image in images:
202 |         image_aux_list = []
203 |         for processor_aux in processor_aux_list:
204 |             image_aux = image
205 |             if hasattr(processor_aux, "image_mean"):
206 |                 try:
207 |                     target_resolution = processor_aux.crop_size["height"]
208 |                 except:
209 |                     target_resolution = processor_aux.size["height"]
210 |                 image_aux = expand2square(
211 |                     image_aux, tuple(int(x * 255) for x in processor_aux.image_mean)
212 |                 ).resize((target_resolution, target_resolution))
213 |             image_aux = processor_aux.preprocess(image_aux, return_tensors="pt")[
214 |                 "pixel_values"
215 |             ][0]
216 |             image_aux_list.append(image_aux)
217 |         new_images_aux_list.append(image_aux_list)
218 |     new_images_aux_list = [
219 |         list(batch_image_aux) for batch_image_aux in zip(*new_images_aux_list)
220 |     ]
221 |     new_images_aux_list = [
222 |         torch.stack(image_aux).half().cuda() for image_aux in new_images_aux_list
223 |     ]
224 |     return new_images_aux_list
225 | 
226 | 
227 | def tokenizer_image_token(
228 |     prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None
229 | ):
230 |     prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split("<image>")]
231 | 
232 |     def insert_separator(X, sep):
233 |         return [ele for sublist in zip(X, [sep] * len(X)) for ele in sublist][:-1]
234 | 
235 |     input_ids = []
236 |     offset = 0
237 |     if (
238 |         len(prompt_chunks) > 0
239 |         and len(prompt_chunks[0]) > 0
240 |         and prompt_chunks[0][0] == tokenizer.bos_token_id
241 |     ):
242 |         offset = 1
243 |         input_ids.append(prompt_chunks[0][0])
244 | 
245 |     for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
246 |         input_ids.extend(x[offset:])
247 | 
248 |     if return_tensors is not None:
249 |         if return_tensors == "pt":
250 |             return torch.tensor(input_ids, dtype=torch.long)
251 |         raise ValueError(f"Unsupported tensor type: {return_tensors}")
252 |     return input_ids
253 | 
254 | 
255 | def tokenizer_image_token_llama3(
256 |     prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None
257 | ):
258 |     prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split("<image>")]
259 | 
260 |     def insert_separator(X, sep):
261 |         return [ele for sublist in zip(X, [sep] * len(X)) for ele in sublist][:-1]
262 | 
263 |     input_ids = []
264 |     for x in insert_separator(prompt_chunks, [image_token_index]):
265 |         input_ids.extend(x)
266 | 
267 |     if return_tensors is not None:
268 |         if return_tensors == "pt":
269 |             return torch.tensor(input_ids, dtype=torch.long)
270 |         raise ValueError(f"Unsupported tensor type: {return_tensors}")
271 |     return input_ids
272 | 
273 | 
274 | def get_model_name_from_path(model_path):
275 |     model_path = model_path.strip("/")
276 |     model_paths = model_path.split("/")
277 |     if model_paths[-1].startswith("checkpoint-"):
278 |         return model_paths[-2] + "_" + model_paths[-1]
279 |     else:
280 |         return model_paths[-1]
281 | 
282 | 
283 | class KeywordsStoppingCriteria(StoppingCriteria):
284 |     def __init__(self, keywords, tokenizer, input_ids):
285 |         self.keywords = keywords
286 |         self.keyword_ids = []
287 |         self.max_keyword_len = 0
288 |         for keyword in keywords:
289 |             cur_keyword_ids = tokenizer(keyword).input_ids
290 |             if (
291 |                 len(cur_keyword_ids) > 1
292 |                 and cur_keyword_ids[0] == tokenizer.bos_token_id
293 |             ):
294 |                 cur_keyword_ids = cur_keyword_ids[1:]
295 |             if len(cur_keyword_ids) > self.max_keyword_len:
296 |                 self.max_keyword_len = len(cur_keyword_ids)
297 |             self.keyword_ids.append(torch.tensor(cur_keyword_ids))
298 |         self.tokenizer = tokenizer
299 |         self.start_len = input_ids.shape[1]
300 | 
301 |     def call_for_batch(
302 |         self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
303 |     ) -> bool:
304 |         offset = min(output_ids.shape[1] - self.start_len, self.max_keyword_len)
305 |         self.keyword_ids = [
306 |             keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids
307 |         ]
308 |         for keyword_id in self.keyword_ids:
309 |             truncated_output_ids = output_ids[0, -keyword_id.shape[0] :]
310 |             if torch.equal(truncated_output_ids, keyword_id):
311 |                 return True
312 |         outputs = self.tokenizer.batch_decode(
313 |             output_ids[:, -offset:], skip_special_tokens=True
314 |         )[0]
315 |         for keyword in self.keywords:
316 |             if keyword in outputs:
317 |                 return True
318 |         return False
319 | 
320 |     def __call__(
321 |         self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
322 |     ) -> bool:
323 |         outputs = []
324 |         for i in range(output_ids.shape[0]):
325 |             # pyre-fixme[6]: For 1st argument expected `LongTensor` but got `Tensor`.
326 |             outputs.append(self.call_for_batch(output_ids[i].unsqueeze(0), scores))
327 |         return all(outputs)
328 | 


--------------------------------------------------------------------------------
/longvu/multimodal_encoder/base_encoder.py:
--------------------------------------------------------------------------------
  1 | from abc import ABC, abstractmethod
  2 | 
  3 | import torch
  4 | import torch.nn as nn
  5 | 
  6 | 
  7 | class ProcessorWrapper:
  8 |     def __init__(
  9 |         self,
 10 |         transform,
 11 |         height=378,
 12 |         width=378,
 13 |         image_mean=[0.48145466, 0.4578275, 0.40821073],
 14 |     ):
 15 |         self._crop_size = {
 16 |             "height": height,
 17 |             "width": width,
 18 |         }
 19 |         self._transforms = transform
 20 |         # print(transform)
 21 |         self.image_mean = image_mean
 22 | 
 23 |     @property
 24 |     def crop_size(self):
 25 |         return self._crop_size
 26 | 
 27 |     def preprocess(self, image, return_tensors="pt"):
 28 |         # Ensure image is a PIL Image
 29 |         output = {}
 30 |         output["pixel_values"] = [self._transforms(image)]
 31 |         return output
 32 | 
 33 | 
 34 | class BaseVisionTower(nn.Module):
 35 |     def __init__(self, vision_tower_name, args, delay_load=False):
 36 |         super().__init__()
 37 | 
 38 |         self.is_loaded = False
 39 |         self.args = args
 40 | 
 41 |         self.vision_tower_name = vision_tower_name
 42 |         self.select_layer = args.mm_vision_select_layer
 43 |         self.select_feature = getattr(args, "mm_vision_select_feature", "patch")
 44 |         self.unfreeze_mm_vision_tower = getattr(args, "unfreeze_mm_vision_tower", False)
 45 |         self.delay_load = delay_load
 46 | 
 47 |     @abstractmethod
 48 |     def load_model(self, device_map=None):
 49 |         raise NotImplementedError("Subclasses must implement load_model")
 50 | 
 51 |     @abstractmethod
 52 |     def _forward(self, images):
 53 |         raise NotImplementedError("Subclasses must implement forward")
 54 | 
 55 |     def forward(self, images):
 56 |         if type(images) is list:
 57 |             image_features = [self._forward(image.unsqueeze(0)) for image in images]
 58 |         else:
 59 |             image_features = self._forward(images)
 60 | 
 61 |         return image_features
 62 | 
 63 |     @property
 64 |     def dummy_feature(self):
 65 |         return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)
 66 | 
 67 |     @property
 68 |     def dtype(self):
 69 |         # Dynamically infer the dtype from the first parameter, if not explicitly specified
 70 |         if hasattr(self.vision_tower, "dtype"):
 71 |             return self.vision_tower.dtype
 72 |         else:
 73 |             params = list(self.vision_tower.parameters())
 74 |             return (
 75 |                 params[0].dtype if len(params) > 0 else torch.float32
 76 |             )  # Default to torch.float32 if no parameters
 77 | 
 78 |     @property
 79 |     def device(self):
 80 |         # Dynamically infer the device from the first parameter, if not explicitly specified
 81 |         if hasattr(self.vision_tower, "device"):
 82 |             return self.vision_tower.device
 83 |         else:
 84 |             params = list(self.vision_tower.parameters())
 85 |             return (
 86 |                 params[0].device if len(params) > 0 else torch.device("cpu")
 87 |             )  # Default to CPU if no parameters
 88 | 
 89 |     @property
 90 |     def config(self):
 91 |         if self.is_loaded:
 92 |             return self.vision_tower.config
 93 |         else:
 94 |             return self.cfg_only
 95 | 
 96 |     @property
 97 |     def hidden_size(self):
 98 |         try:
 99 |             return self.config.hidden_size
100 |         except:
101 |             return self._hidden_size
102 | 
103 |     @property
104 |     def image_size(self):  # resolution
105 |         # return self.config.image_size
106 |         try:
107 |             return self.config.image_size
108 |         except:
109 |             return self._image_size
110 | 
111 |     @property
112 |     def patch_size(self):
113 |         # return self.config.patch_size
114 |         try:
115 |             return self.config.patch_size
116 |         except:
117 |             return self._patch_size
118 | 
119 |     @property
120 |     def num_patches_per_side(self):
121 |         if self._interp_size is not None:
122 |             return int(self._interp_size**0.5)
123 |         try:
124 |             return self.image_size // self.patch_size
125 |         except:
126 |             return self._num_patches_per_side
127 | 
128 |     @property
129 |     def num_patches(self):
130 |         if self._interp_size is not None:
131 |             return self._interp_size
132 |         try:
133 |             return self.num_patches_per_side**2
134 |         except:
135 |             return self._num_patches
136 | 


--------------------------------------------------------------------------------
/longvu/multimodal_encoder/builder.py:
--------------------------------------------------------------------------------
 1 | # pyre-unsafe
 2 | import copy
 3 | 
 4 | from .dino_encoder import DinoVisionTower
 5 | from .siglip_encoder import SiglipVisionTower
 6 | 
 7 | 
 8 | def build_vision_tower_aux_list(vision_tower_cfg, **kwargs):
 9 |     vision_tower_aux_name_list = getattr(
10 |         vision_tower_cfg,
11 |         "mm_vision_tower_aux_list",
12 |         getattr(vision_tower_cfg, "vision_tower_aux_list", None),
13 |     )
14 |     vision_tower_aux_token_len_list = getattr(
15 |         vision_tower_cfg,
16 |         "mm_vision_tower_aux_token_len_list",
17 |         getattr(vision_tower_cfg, "vision_tower_aux_token_len_list", None),
18 |     )
19 |     vision_tower_aux_list = []
20 |     for vision_tower_aux_name, vision_tower_aux_token_len in zip(
21 |         vision_tower_aux_name_list, vision_tower_aux_token_len_list
22 |     ):
23 |         config = copy.deepcopy(vision_tower_cfg)
24 |         vision_tower_aux_name += "-interp{}".format(vision_tower_aux_token_len)
25 |         if "siglip" in vision_tower_aux_name.lower():
26 |             vision_tower_aux_list.append(
27 |                 SiglipVisionTower(vision_tower_aux_name, args=config, **kwargs)
28 |             )
29 | 
30 |         # SSL-based Vision Towers
31 |         elif "dinov2" in vision_tower_aux_name.lower():
32 |             vision_tower_aux_list.append(
33 |                 DinoVisionTower(vision_tower_aux_name, args=config, **kwargs)
34 |             )
35 |         else:
36 |             raise ValueError(f"Unknown vision tower: {vision_tower_aux_name}")
37 |     return vision_tower_aux_list
38 | 


--------------------------------------------------------------------------------
/longvu/multimodal_encoder/dino_encoder.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import torch.nn.functional as F
  3 | 
  4 | from transformers import AutoImageProcessor, Dinov2Config, Dinov2Model
  5 | 
  6 | from .base_encoder import BaseVisionTower, ProcessorWrapper
  7 | 
  8 | 
  9 | class DinoVisionTower(BaseVisionTower):
 10 |     def __init__(self, vision_tower, args, delay_load=False):
 11 |         super(DinoVisionTower, self).__init__(vision_tower, args, delay_load)
 12 | 
 13 |         model_path = "facebook/dinov2-giant"
 14 |         base_model_name, res, interp = model_path, 378, 576
 15 |         self._vision_tower_name = vision_tower
 16 |         self.vision_tower_name = base_model_name
 17 |         self._image_size = res
 18 |         self._interp_size = interp
 19 |         self._patch_size = 14  # default patch size
 20 | 
 21 |         if not self.delay_load:
 22 |             self.load_model()
 23 |         else:
 24 |             self.cfg_only = Dinov2Config.from_pretrained(self.vision_tower_name)
 25 | 
 26 |     def load_model(self, device_map=None):
 27 | 
 28 |         self.vision_tower = Dinov2Model.from_pretrained(self.vision_tower_name)
 29 |         """ValueError: Dinov2Model does not support `device_map='auto'`. To implement support, the model class needs to implement the `_no_split_modules` attribute."""
 30 |         self.vision_tower._no_split_modules = ["Dinov2SwiGLUFFN"]
 31 | 
 32 |         _image_size = self.vision_tower.config.image_size
 33 |         if self._image_size is None:
 34 |             self._image_size = _image_size
 35 | 
 36 |         # increase shortest edge to prevent edge case crops
 37 |         default_shortest_ratio = 8 / 7  # 224/256
 38 |         # shortest_edge = int(default_shortest_ratio * self._image_size)
 39 |         shortest_edge = self._image_size
 40 | 
 41 |         processor = AutoImageProcessor.from_pretrained(
 42 |             self.vision_tower_name,
 43 |             crop_size=dict(height=self._image_size, width=self._image_size),
 44 |             size=dict(shortest_edge=shortest_edge),
 45 |         )
 46 |         self.image_processor = processor
 47 | 
 48 |         # Assign the output channels of the projection convolution as the hidden size
 49 |         self._hidden_size = (
 50 |             self.vision_tower.embeddings.patch_embeddings.projection.out_channels
 51 |         )
 52 |         # Assign the first value of the stride of the projection convolution as the patch size
 53 |         self._patch_size = (
 54 |             self.vision_tower.embeddings.patch_embeddings.projection.stride[0]
 55 |         )
 56 | 
 57 |         # print(self._hidden_size, self._patch_size)
 58 | 
 59 |         self.vision_tower.requires_grad_(self.unfreeze_mm_vision_tower)
 60 |         self.is_loaded = True
 61 | 
 62 |     @property
 63 |     def image_size(self):
 64 |         return self._image_size
 65 | 
 66 |     def feature_select(self, outputs):
 67 |         sequence_output = outputs[
 68 |             "last_hidden_state"
 69 |         ]  # batch_size, sequence_length, hidden_size
 70 | 
 71 |         if self.select_feature == "cls_patch":
 72 |             image_features = sequence_output
 73 |         elif self.select_feature == "patch":
 74 |             image_features = sequence_output[:, 1:]
 75 |         elif self.select_feature == "cls":
 76 |             image_features = sequence_output[:, 0]
 77 |         else:
 78 |             raise ValueError(f"Unexpected select feature: {self.select_feature}")
 79 |         return image_features
 80 | 
 81 |     def interpolate(self, image_features):
 82 |         if self._interp_size is None:
 83 |             return image_features
 84 | 
 85 |         b, num_tokens, dim = image_features.shape
 86 | 
 87 |         if num_tokens != self.num_patches:
 88 |             target_h = target_w = int(self._interp_size**0.5)
 89 |             h = w = int(num_tokens**0.5)
 90 | 
 91 |             image_features = image_features.view(b, h, w, dim)
 92 |             image_features = image_features.permute(0, 3, 1, 2).contiguous()
 93 | 
 94 |             image_features = F.interpolate(
 95 |                 image_features.to(torch.float32),
 96 |                 size=(target_h, target_w),
 97 |                 mode="bilinear",
 98 |                 align_corners=False,
 99 |             ).to(image_features.dtype)
100 | 
101 |             # Permute the dimensions back to (b, target_h, target_w, dim)
102 |             image_features = image_features.permute(0, 2, 3, 1).contiguous()
103 | 
104 |             # Flatten the spatial dimensions (target_h, target_w) into a single dimension
105 |             image_features = image_features.flatten(1, 2)
106 | 
107 |         return image_features
108 | 
109 |     def _forward(self, images):
110 |         # logger.warning(f"images shape: {images.shape}")
111 |         with torch.set_grad_enabled(self.unfreeze_mm_vision_tower):
112 |             image_forward_outs = self.vision_tower.forward(
113 |                 images.to(device=self.device, dtype=self.dtype)
114 |             )
115 |             # logger.warning(f"image_forward_outs shape: {image_forward_outs['last_hidden_state'].shape}")
116 |             image_features = self.feature_select(image_forward_outs).to(images.dtype)
117 |             # logger.warning(f"image_features shape: {image_features.shape}")
118 |             interp_features = self.interpolate(image_features)
119 |             # logger.warning(f"interp_features shape: {interp_features.shape}")
120 |             return interp_features
121 | 
122 |     @property
123 |     def num_patches_per_side(self):
124 |         return int(self.num_patches**0.5)
125 | 
126 |     @property
127 |     def num_patches(self):
128 |         if self._interp_size is None:
129 |             return (self._image_size // self._patch_size) ** 2
130 |         else:
131 |             return self._interp_size
132 | 


--------------------------------------------------------------------------------
/longvu/multimodal_encoder/drop.py:
--------------------------------------------------------------------------------
 1 | # ------------------------------------------------------------------------
 2 | # Copyright (c) 2023-present, BAAI. All Rights Reserved.
 3 | #
 4 | # Licensed under the Apache License, Version 2.0 (the "License");
 5 | # you may not use this file except in compliance with the License.
 6 | # You may obtain a copy of the License at
 7 | #
 8 | #    http://www.apache.org/licenses/LICENSE-2.0
 9 | #
10 | # Unless required by applicable law or agreed to in writing, software
11 | # distributed under the License is distributed on an "AS IS" BASIS,
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | # See the License for the specific language governing permissions and
14 | # limitations under the License.
15 | # ------------------------------------------------------------------------
16 | 
17 | # pyre-unsafe
18 | """Drop regularization layers."""
19 | 
20 | from torch import nn
21 | 
22 | 
23 | class DropPathV(nn.Module):
24 |     """Set examples to zero randomly."""
25 | 
26 |     def __init__(self, p=0.1, inplace=False):
27 |         super(DropPathV, self).__init__()
28 |         self.p = p
29 |         self.inplace = inplace
30 | 
31 |     def forward(self, input):
32 |         if not self.training or self.p <= 0:
33 |             return input
34 |         keep_p = 1 - self.p
35 |         shape = (input.shape[0],) + (1,) * (input.dim() - 1)
36 |         scale = input.new_empty(shape).bernoulli_(keep_p).div_(keep_p)
37 |         return input.mul_(scale) if self.inplace else input.mul(scale)
38 | 
39 |     def extra_repr(self):
40 |         inplace_str = ", inplace" if self.inplace else ""
41 |         return "p={}{}".format(self.p, inplace_str)
42 | 


--------------------------------------------------------------------------------
/longvu/multimodal_encoder/image.py:
--------------------------------------------------------------------------------
 1 | # ------------------------------------------------------------------------
 2 | # Copyright (c) 2023-present, BAAI. All Rights Reserved.
 3 | #
 4 | # Licensed under the Apache License, Version 2.0 (the "License");
 5 | # you may not use this file except in compliance with the License.
 6 | # You may obtain a copy of the License at
 7 | #
 8 | #    http://www.apache.org/licenses/LICENSE-2.0
 9 | #
10 | # Unless required by applicable law or agreed to in writing, software
11 | # distributed under the License is distributed on an "AS IS" BASIS,
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | # See the License for the specific language governing permissions and
14 | # limitations under the License.
15 | # ------------------------------------------------------------------------
16 | 
17 | # pyre-unsafe
18 | """Image utilities."""
19 | 
20 | import numpy as np
21 | import PIL.Image
22 | import torch
23 | 
24 | 
25 | def im_resize(img, size=None, scale=None, mode="linear"):
26 |     """Resize image by the scale or size."""
27 |     if size is None:
28 |         if not isinstance(scale, (tuple, list)):
29 |             scale = (scale, scale)
30 |         h, w = img.shape[:2]
31 |         size = int(h * scale[0] + 0.5), int(w * scale[1] + 0.5)
32 |     else:
33 |         if not isinstance(size, (tuple, list)):
34 |             size = (size, size)
35 |     resize_modes = {"linear": PIL.Image.BILINEAR}
36 |     from torchvision.transforms import ToPILImage
37 | 
38 |     to_pil = ToPILImage()
39 |     img = to_pil(img.to(torch.float32).cpu())
40 |     # img = PIL.Image.fromarray(img)
41 |     return np.array(img.resize(size[::-1], resize_modes[mode]))
42 | 
43 | 
44 | def im_rescale(img, scales, max_size=0):
45 |     """Rescale image to match the detecting scales."""
46 |     im_shape = img.shape
47 |     img_list, img_scales = [], []
48 |     size_min = np.min(im_shape[:2])
49 |     size_max = np.max(im_shape[:2])
50 |     for target_size in scales:
51 |         im_scale = float(target_size) / float(size_min)
52 |         target_size_max = max_size if max_size > 0 else target_size
53 |         if np.round(im_scale * size_max) > target_size_max:
54 |             im_scale = float(target_size_max) / float(size_max)
55 |         img_list.append(im_resize(img, scale=im_scale))
56 |         img_scales.append((im_scale, im_scale))
57 |     return img_list, img_scales
58 | 
59 | 
60 | def im_vstack(arrays, fill_value=None, dtype=None, size=None, align=None):
61 |     """Stack image arrays in sequence vertically."""
62 |     if fill_value is None:
63 |         return np.vstack(arrays)
64 |     # Compute the max stack shape.
65 |     max_shape = np.max(np.stack([arr.shape for arr in arrays]), 0)
66 |     if size is not None and min(size) > 0:
67 |         max_shape[: len(size)] = size
68 |     if align is not None and min(align) > 0:
69 |         align_size = np.ceil(max_shape[: len(align)] / align)
70 |         max_shape[: len(align)] = align_size.astype("int64") * align
71 |     # Fill output with the given value.
72 |     output_dtype = dtype or arrays[0].dtype
73 |     output_shape = [len(arrays)] + list(max_shape)
74 |     output = np.empty(output_shape, output_dtype)
75 |     output[:] = fill_value
76 |     # Copy arrays.
77 |     for i, arr in enumerate(arrays):
78 |         copy_slices = (slice(0, d) for d in arr.shape)
79 |         output[(i,) + tuple(copy_slices)] = arr
80 |     return output
81 | 


--------------------------------------------------------------------------------
/longvu/multimodal_encoder/logging.py:
--------------------------------------------------------------------------------
  1 | # ------------------------------------------------------------------------
  2 | # Copyright (c) 2023-present, BAAI. All Rights Reserved.
  3 | #
  4 | # Licensed under the Apache License, Version 2.0 (the "License");
  5 | # you may not use this file except in compliance with the License.
  6 | # You may obtain a copy of the License at
  7 | #
  8 | #    http://www.apache.org/licenses/LICENSE-2.0
  9 | #
 10 | # Unless required by applicable law or agreed to in writing, software
 11 | # distributed under the License is distributed on an "AS IS" BASIS,
 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, esither express or implied.
 13 | # See the License for the specific language governing permissions and
 14 | # limitations under the License.
 15 | # ------------------------------------------------------------------------
 16 | 
 17 | # pyre-unsafe
 18 | """Logging utilities."""
 19 | 
 20 | import inspect
 21 | import logging as _logging
 22 | import os
 23 | import sys as _sys
 24 | import threading
 25 | 
 26 | 
 27 | _logger = None
 28 | _logger_lock = threading.Lock()
 29 | 
 30 | 
 31 | def get_logger():
 32 |     global _logger
 33 |     # Use double-checked locking to avoid taking lock unnecessarily.
 34 |     if _logger:
 35 |         return _logger
 36 |     _logger_lock.acquire()
 37 |     try:
 38 |         if _logger:
 39 |             return _logger
 40 |         logger = _logging.getLogger("tokenize-anything")
 41 |         logger.setLevel("INFO")
 42 |         logger.propagate = False
 43 |         logger._is_root = True
 44 |         if True:
 45 |             # Determine whether we are in an interactive environment.
 46 |             _interactive = False
 47 |             try:
 48 |                 # This is only defined in interactive shells.
 49 |                 if _sys.ps1:
 50 |                     _interactive = True
 51 |             except AttributeError:
 52 |                 # Even now, we may be in an interactive shell with `python -i`.
 53 |                 _interactive = _sys.flags.interactive
 54 |             # If we are in an interactive environment (like Jupyter), set loglevel
 55 |             # to INFO and pipe the output to stdout.
 56 |             if _interactive:
 57 |                 logger.setLevel("INFO")
 58 |                 _logging_target = _sys.stdout
 59 |             else:
 60 |                 _logging_target = _sys.stderr
 61 |             # Add the output handler.
 62 |             _handler = _logging.StreamHandler(_logging_target)
 63 |             _handler.setFormatter(_logging.Formatter("%(levelname)s %(message)s"))
 64 |             logger.addHandler(_handler)
 65 |         _logger = logger
 66 |         return _logger
 67 |     finally:
 68 |         _logger_lock.release()
 69 | 
 70 | 
 71 | def _detailed_msg(msg):
 72 |     file, lineno = inspect.stack()[:3][2][1:3]
 73 |     return "{}:{}] {}".format(os.path.split(file)[-1], lineno, msg)
 74 | 
 75 | 
 76 | def log(level, msg, *args, **kwargs):
 77 |     get_logger().log(level, _detailed_msg(msg), *args, **kwargs)
 78 | 
 79 | 
 80 | def debug(msg, *args, **kwargs):
 81 |     if is_root():
 82 |         get_logger().debug(_detailed_msg(msg), *args, **kwargs)
 83 | 
 84 | 
 85 | def error(msg, *args, **kwargs):
 86 |     get_logger().error(_detailed_msg(msg), *args, **kwargs)
 87 |     assert 0
 88 | 
 89 | 
 90 | def fatal(msg, *args, **kwargs):
 91 |     get_logger().fatal(_detailed_msg(msg), *args, **kwargs)
 92 |     assert 0
 93 | 
 94 | 
 95 | def info(msg, *args, **kwargs):
 96 |     if is_root():
 97 |         get_logger().info(_detailed_msg(msg), *args, **kwargs)
 98 | 
 99 | 
100 | def warning(msg, *args, **kwargs):
101 |     if is_root():
102 |         get_logger().warning(_detailed_msg(msg), *args, **kwargs)
103 | 
104 | 
105 | def get_verbosity():
106 |     """Return how much logging output will be produced."""
107 |     return get_logger().getEffectiveLevel()
108 | 
109 | 
110 | def set_verbosity(v):
111 |     """Set the threshold for what messages will be logged."""
112 |     get_logger().setLevel(v)
113 | 
114 | 
115 | def set_formatter(fmt=None, datefmt=None):
116 |     """Set the formatter."""
117 |     handler = _logging.StreamHandler(_sys.stderr)
118 |     handler.setFormatter(_logging.Formatter(fmt, datefmt))
119 |     logger = get_logger()
120 |     logger.removeHandler(logger.handlers[0])
121 |     logger.addHandler(handler)
122 | 
123 | 
124 | def set_root(is_root=True):
125 |     """Set logger to the root."""
126 |     get_logger()._is_root = is_root
127 | 
128 | 
129 | def is_root():
130 |     """Return logger is the root."""
131 |     return get_logger()._is_root
132 | 


--------------------------------------------------------------------------------
/longvu/multimodal_encoder/loss.py:
--------------------------------------------------------------------------------
 1 | # ------------------------------------------------------------------------
 2 | # Copyright (c) 2023-present, BAAI. All Rights Reserved.
 3 | #
 4 | # Licensed under the Apache License, Version 2.0 (the "License");
 5 | # you may not use this file except in compliance with the License.
 6 | # You may obtain a copy of the License at
 7 | #
 8 | #    http://www.apache.org/licenses/LICENSE-2.0
 9 | #
10 | # Unless required by applicable law or agreed to in writing, software
11 | # distributed under the License is distributed on an "AS IS" BASIS,
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | # See the License for the specific language governing permissions and
14 | # limitations under the License.
15 | # ------------------------------------------------------------------------
16 | 
17 | # pyre-unsafe
18 | """Loss layers."""
19 | 
20 | from torch import nn
21 | 
22 | 
23 | def reduce_loss(loss, reduction="mean"):
24 |     """Reduce the loss."""
25 |     if reduction == "mean" or reduction == "sum":
26 |         return getattr(loss, reduction)()
27 |     if reduction == "batch_mean":
28 |         return loss.sum().mul_(1.0 / loss.size(0))
29 |     return loss
30 | 
31 | 
32 | class BinaryFocalLoss(nn.Module):
33 |     """Binary focal loss."""
34 | 
35 |     def __init__(self, alpha=0.25, reduction="none"):
36 |         super(BinaryFocalLoss, self).__init__()
37 |         self.alpha = alpha
38 |         self.reduction = reduction
39 | 
40 |     def forward(self, input, target):
41 |         alpha, p = self.alpha, input.sigmoid()
42 |         neg_alpha, neg_target = 1.0 - alpha, 1.0 - target
43 |         alpha_weight = target.mul(alpha).add_(neg_target.mul(neg_alpha))
44 |         focal_weight = (1.0 - p).mul_(target).add_(p.mul(neg_target)).square()
45 |         loss = nn.functional.binary_cross_entropy_with_logits(
46 |             input, target, reduction="none"
47 |         )
48 |         return reduce_loss(loss * focal_weight.mul_(alpha_weight), self.reduction)
49 | 
50 | 
51 | class BinaryDiceLoss(nn.Module):
52 |     """Binary dice loss."""
53 | 
54 |     def __init__(self, eps=1.0, reduction="none"):
55 |         super(BinaryDiceLoss, self).__init__()
56 |         self.eps = eps
57 |         self.reduction = reduction
58 | 
59 |     def forward(self, input, target):
60 |         input = input.sigmoid()
61 |         num = input.mul(target).sum(-1).mul_(2).add_(self.eps)
62 |         den = input.add(target).sum(-1).add_(self.eps)
63 |         return reduce_loss(1.0 - num / den, self.reduction)
64 | 
65 | 
66 | class CrossEntropyLoss(nn.Module):
67 |     """Cross entropy loss with label smoothing."""
68 | 
69 |     def __init__(self, epsilon=0, reduction="none"):
70 |         super(CrossEntropyLoss, self).__init__()
71 |         self.epsilon = epsilon
72 |         self.reduction = reduction
73 | 
74 |     def forward_dense(self, input, target):
75 |         dim, target = input.shape[-1], target.squeeze_()
76 |         x = nn.functional.log_softmax(input, dim=-1)
77 |         y = nn.functional.one_hot(target, dim).float()
78 |         x = (
79 |             x.permute([0, x.dim() - 1] + list(range(x.dim()))[1:-1])
80 |             if x.dim() > 2
81 |             else x
82 |         )
83 |         y = (
84 |             y.permute([0, y.dim() - 1] + list(range(y.dim()))[1:-1])
85 |             if y.dim() > 2
86 |             else y
87 |         )
88 |         loss = nn.functional.cross_entropy(
89 |             x, y, reduction="none", label_smoothing=self.epsilon
90 |         )
91 |         return reduce_loss(loss, self.reduction)
92 | 
93 |     def forward(self, input, target):
94 |         if self.epsilon > 0:
95 |             return self.forward_dense(input, target)
96 |         return nn.functional.cross_entropy(input, target, reduction=self.reduction)
97 | 


--------------------------------------------------------------------------------
/longvu/multimodal_encoder/registry.py:
--------------------------------------------------------------------------------
 1 | # ------------------------------------------------------------------------
 2 | # Copyright (c) 2023-present, BAAI. All Rights Reserved.
 3 | #
 4 | # Licensed under the Apache License, Version 2.0 (the "License");
 5 | # you may not use this file except in compliance with the License.
 6 | # You may obtain a copy of the License at
 7 | #
 8 | #    http://www.apache.org/licenses/LICENSE-2.0
 9 | #
10 | # Unless required by applicable law or agreed to in writing, software
11 | # distributed under the License is distributed on an "AS IS" BASIS,
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | # See the License for the specific language governing permissions and
14 | # limitations under the License.
15 | # ------------------------------------------------------------------------
16 | 
17 | # pyre-unsafe
18 | """Registry utilities."""
19 | 
20 | import collections
21 | import functools
22 | 
23 | 
24 | class Registry(object):
25 |     """Registry class."""
26 | 
27 |     def __init__(self, name):
28 |         self.name = name
29 |         self.registry = collections.OrderedDict()
30 | 
31 |     def has(self, key):
32 |         return key in self.registry
33 | 
34 |     def register(self, name, func=None, **kwargs):
35 |         def decorated(inner_function):
36 |             for key in name if isinstance(name, (tuple, list)) else [name]:
37 |                 self.registry[key] = functools.partial(inner_function, **kwargs)
38 |             return inner_function
39 | 
40 |         if func is not None:
41 |             return decorated(func)
42 |         return decorated
43 | 
44 |     def get(self, name, default=None):
45 |         if name is None:
46 |             return None
47 |         if not self.has(name):
48 |             if default is not None:
49 |                 return default
50 |             raise KeyError("`%s` is not registered in <%s>." % (name, self.name))
51 |         return self.registry[name]
52 | 
53 |     def try_get(self, name):
54 |         if self.has(name):
55 |             return self.get(name)
56 |         return None
57 | 


--------------------------------------------------------------------------------
/longvu/multimodal_encoder/siglip_encoder.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | import torch.nn.functional as F
 3 | 
 4 | from transformers import SiglipImageProcessor, SiglipVisionConfig, SiglipVisionModel
 5 | 
 6 | from .base_encoder import BaseVisionTower, ProcessorWrapper
 7 | 
 8 | 
 9 | class SiglipVisionTower(BaseVisionTower):
10 |     def __init__(self, vision_tower_name, args, delay_load=False):
11 |         super(SiglipVisionTower, self).__init__(vision_tower_name, args, delay_load)
12 |         
13 |         model_path = "google/siglip-so400m-patch14-384"
14 |         base_model_name, res, interp = model_path, 384, 576
15 |         self.vision_tower_name = base_model_name
16 |         self._image_size = res if res is not None else 512
17 |         self._interp_size = interp
18 |         if not self.delay_load:
19 |             self.load_model()
20 |         elif self.unfreeze_mm_vision_tower:
21 |             self.load_model()
22 |         else:
23 |             self._hidden_size = 1152
24 | 
25 |     def load_model(self, device_map=None):
26 |         self.vision_model = "siglip"
27 |         # clip_model, processor = create_model_from_pretrained(self.vision_tower_name)
28 |         self.vision_tower = SiglipVisionModel.from_pretrained(self.vision_tower_name)
29 | 
30 |         # self.vision_tower = clip_model.visual.trunk
31 |         self.vision_tower.output_tokens = True
32 | 
33 |         self._hidden_size = self.vision_tower.config.hidden_size
34 |         self._image_size = self.vision_tower.config.image_size
35 |         self._patch_size = self.vision_tower.config.patch_size
36 |         self.image_processor = SiglipImageProcessor.from_pretrained(
37 |             self.vision_tower_name
38 |         )
39 | 
40 |         self.vision_tower.requires_grad_(self.unfreeze_mm_vision_tower)
41 |         self.is_loaded = True
42 | 
43 |     def interpolate(self, image_features):
44 |         if self._interp_size is None:
45 |             return image_features
46 | 
47 |         b, num_tokens, dim = image_features.shape
48 | 
49 |         if num_tokens != self.num_patches:
50 |             target_h = target_w = int(self._interp_size**0.5)
51 |             h = w = int(num_tokens**0.5)
52 | 
53 |             image_features = image_features.view(b, h, w, dim)
54 |             image_features = image_features.permute(0, 3, 1, 2).contiguous()
55 | 
56 |             image_features = F.interpolate(
57 |                 image_features.to(torch.float32),
58 |                 size=(target_h, target_w),
59 |                 mode="bilinear",
60 |                 align_corners=False,
61 |             ).to(image_features.dtype)
62 | 
63 |             # Permute the dimensions back to (b, target_h, target_w, dim)
64 |             image_features = image_features.permute(0, 2, 3, 1).contiguous()
65 | 
66 |             # Flatten the spatial dimensions (target_h, target_w) into a single dimension
67 |             image_features = image_features.flatten(1, 2)
68 | 
69 |         return image_features
70 | 
71 |     def _forward(self, images, interpolate_token=576):
72 |         with torch.set_grad_enabled(self.unfreeze_mm_vision_tower):
73 |             image_features = self.vision_tower.forward(
74 |                 images.to(device=self.device, dtype=self.dtype),
75 |                 output_hidden_states=True,
76 |             ).hidden_states[-1]
77 |             interp_features = self.interpolate(image_features)
78 |             return interp_features
79 | 


--------------------------------------------------------------------------------
/longvu/multimodal_encoder/utils.py:
--------------------------------------------------------------------------------
 1 | # ------------------------------------------------------------------------
 2 | # Copyright (c) 2023-present, BAAI. All Rights Reserved.
 3 | #
 4 | # Licensed under the Apache License, Version 2.0 (the "License");
 5 | # you may not use this file except in compliance with the License.
 6 | # You may obtain a copy of the License at
 7 | #
 8 | #    http://www.apache.org/licenses/LICENSE-2.0
 9 | #
10 | # Unless required by applicable law or agreed to in writing, software
11 | # distributed under the License is distributed on an "AS IS" BASIS,
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | # See the License for the specific language governing permissions and
14 | # limitations under the License.
15 | # ------------------------------------------------------------------------
16 | 
17 | # pyre-unsafe
18 | """Layer utilities."""
19 | 
20 | import cv2
21 | import numpy as np
22 | import torch
23 | 
24 | 
25 | def init_cross_conv(blocks):
26 |     """Initialize convolutional cross attention."""
27 |     for m in blocks.modules():
28 |         if isinstance(m, torch.nn.Conv2d):
29 |             torch.nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
30 |     for blk in blocks:
31 |         torch.nn.init.constant_(blk.norm3.weight, 0)
32 | 
33 | 
34 | def set_dropout(module, dropout):
35 |     """Initialize dropout."""
36 |     for m in [m for m in module.modules() if isinstance(m, torch.nn.Dropout)]:
37 |         m.p = dropout
38 | 
39 | 
40 | def set_drop_path(blocks, drop_path):
41 |     """Initialize drop path."""
42 |     if not isinstance(blocks, torch.nn.ModuleList):
43 |         blocks = getattr(blocks, "blocks", getattr(blocks, "layers", None))
44 |     for i, blk in enumerate(blocks):
45 |         for m in [m for m in blk.modules() if type(m).__name__ == "DropPath"]:
46 |             m.p = i * drop_path / (len(blocks) - 1)
47 | 
48 | 
49 | def set_sync_batch_norm(module, ddp_group):
50 |     """Set data parallelism group for sync batch norm."""
51 |     for m in module.modules():
52 |         if isinstance(m, torch.nn.SyncBatchNorm):
53 |             m.process_group = ddp_group
54 | 
55 | 
56 | def resize_pos_embed(weight, out_len):
57 |     """Resize position embedding weights."""
58 |     out_h = out_w = int(out_len**0.5)
59 |     h = w = int(weight.shape[0] ** 0.5)
60 |     weight = weight.reshape((h, w, weight.shape[1]))
61 |     out_weight = [
62 |         cv2.resize(x, (out_w, out_h), interpolation=cv2.INTER_CUBIC)
63 |         for x in np.split(weight.astype("float32", copy=False), 4, axis=-1)
64 |     ]
65 |     out_weight = np.concatenate(out_weight, axis=-1)
66 |     return out_weight.reshape((-1, weight.shape[-1])).astype(weight.dtype, copy=False)
67 | 


--------------------------------------------------------------------------------
/longvu/multimodal_projector/builder.py:
--------------------------------------------------------------------------------
 1 | # pyre-unsafe
 2 | import re
 3 | 
 4 | import torch.nn as nn
 5 | 
 6 | 
 7 | class IdentityMap(nn.Module):
 8 |     def __init__(self):
 9 |         super().__init__()
10 | 
11 |     def forward(self, x, *args, **kwargs):
12 |         return x
13 | 
14 |     @property
15 |     def config(self):
16 |         return {"mm_projector_type": "identity"}
17 | 
18 | 
19 | class SimpleResBlock(nn.Module):
20 |     def __init__(self, channels):
21 |         super().__init__()
22 |         self.pre_norm = nn.LayerNorm(channels)
23 | 
24 |         self.proj = nn.Sequential(
25 |             nn.Linear(channels, channels), nn.GELU(), nn.Linear(channels, channels)
26 |         )
27 | 
28 |     def forward(self, x):
29 |         x = self.pre_norm(x)
30 |         return x + self.proj(x)
31 | 
32 | 
33 | def build_vision_projector(config, delay_load=False, **kwargs):
34 |     projector_type = getattr(config, "mm_projector_type", "linear")
35 |     config.mm_hidden_size = 256
36 | 
37 |     if projector_type == "linear":
38 |         return nn.Linear(config.mm_hidden_size, config.hidden_size)
39 | 
40 |     mlp_gelu_match = re.match(r"^mlp(\d+)x_gelu$", projector_type)
41 |     if mlp_gelu_match:
42 |         mlp_depth = int(mlp_gelu_match.group(1))
43 |         modules = [nn.Linear(config.mm_hidden_size, config.hidden_size)]
44 |         for _ in range(1, mlp_depth):
45 |             modules.append(nn.GELU())
46 |             modules.append(nn.Linear(config.hidden_size, config.hidden_size))
47 |         return nn.Sequential(*modules)
48 | 
49 |     if projector_type == "identity":
50 |         return IdentityMap()
51 | 
52 |     raise ValueError(f"Unknown projector type: {projector_type}")
53 | 


--------------------------------------------------------------------------------
/longvu/utils.py:
--------------------------------------------------------------------------------
 1 | # pyre-unsafe
 2 | from transformers import AutoConfig
 3 | 
 4 | 
 5 | def auto_upgrade(config):
 6 |     cfg = AutoConfig.from_pretrained(config)
 7 |     if "llava" in config and "llava" not in cfg.model_type:
 8 |         assert cfg.model_type == "llama"
 9 |         print(
10 |             "You are using newer LLaVA code base, while the checkpoint of v0 is from older code base."
11 |         )
12 |         print(
13 |             "You must upgrade the checkpoint to the new code base (this can be done automatically)."
14 |         )
15 |         confirm = input("Please confirm that you want to upgrade the checkpoint. [Y/N]")
16 |         if confirm.lower() in ["y", "yes"]:
17 |             print("Upgrading checkpoint...")
18 |             assert len(cfg.architectures) == 1
19 |             setattr(cfg.__class__, "model_type", "llava")
20 |             cfg.architectures[0] = "LlavaLlamaForCausalLM"
21 |             cfg.save_pretrained(config)
22 |             print("Checkpoint upgraded.")
23 |         else:
24 |             print("Checkpoint upgrade aborted.")
25 |             exit(1)
26 | 


--------------------------------------------------------------------------------
/longvu/vision_sampler.py:
--------------------------------------------------------------------------------
  1 | import math
  2 | 
  3 | import numpy as np
  4 | import torch
  5 | import torch.utils.checkpoint
  6 | from torch import nn
  7 | 
  8 | 
  9 | # https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20
 10 | def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
 11 |     """
 12 |     grid_size: int of the grid height and width
 13 |     return:
 14 |     pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
 15 |     """
 16 |     grid_h = np.arange(grid_size, dtype=np.float32)
 17 |     grid_w = np.arange(grid_size, dtype=np.float32)
 18 |     grid = np.meshgrid(grid_w, grid_h)  # here w goes first
 19 |     grid = np.stack(grid, axis=0)
 20 | 
 21 |     grid = grid.reshape([2, 1, grid_size, grid_size])
 22 | 
 23 |     pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
 24 |     if cls_token:
 25 |         pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
 26 |     return pos_embed
 27 | 
 28 | 
 29 | def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
 30 |     assert embed_dim % 2 == 0
 31 | 
 32 |     # use half of dimensions to encode grid_h
 33 |     emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
 34 |     emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
 35 | 
 36 |     emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
 37 |     return emb
 38 | 
 39 | 
 40 | def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
 41 |     """
 42 |     embed_dim: output dimension for each position
 43 |     pos: a list of positions to be encoded: size (M,)
 44 |     out: (M, D)
 45 |     """
 46 |     assert embed_dim % 2 == 0
 47 |     omega = np.arange(embed_dim // 2, dtype=np.float32)
 48 |     omega /= embed_dim / 2.0
 49 |     omega = 1.0 / 10000**omega  # (D/2,)
 50 | 
 51 |     pos = pos.reshape(-1)  # (M,)
 52 |     out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
 53 | 
 54 |     emb_sin = np.sin(out)  # (M, D/2)
 55 |     emb_cos = np.cos(out)  # (M, D/2)
 56 | 
 57 |     emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
 58 |     return emb
 59 | 
 60 | 
 61 | class CrossAttention(nn.Module):
 62 | 
 63 |     def __init__(self, q_dim, kv_dim, hidden_dim, num_heads, attention_bias=False):
 64 |         super().__init__()
 65 |         self.hidden_dim = hidden_dim
 66 |         self.num_heads = num_heads
 67 |         self.head_dim = self.hidden_dim // self.num_heads
 68 | 
 69 |         if (self.head_dim * self.num_heads) != self.hidden_dim:
 70 |             raise ValueError(
 71 |                 f"hidden_dim must be divisible by num_heads (got `hidden_dim`: {self.hidden_dim}"
 72 |                 f" and `num_heads`: {self.num_heads})."
 73 |             )
 74 | 
 75 |         self.q_proj = nn.Sequential(
 76 |             nn.LayerNorm(q_dim),
 77 |             nn.Linear(q_dim, self.num_heads * self.head_dim, bias=attention_bias),
 78 |         )
 79 |         self.k_proj = nn.Sequential(
 80 |             nn.LayerNorm(kv_dim),
 81 |             nn.Linear(kv_dim, self.num_heads * self.head_dim, bias=attention_bias),
 82 |         )
 83 |         self.v_proj = nn.Sequential(
 84 |             nn.LayerNorm(kv_dim),
 85 |             nn.Linear(kv_dim, self.num_heads * self.head_dim, bias=attention_bias),
 86 |         )
 87 |         self.o_proj = nn.Linear(
 88 |             self.num_heads * self.head_dim, q_dim, bias=attention_bias
 89 |         )
 90 | 
 91 |     def forward(self, vision_latents, queries, attention_mask):
 92 | 
 93 |         bsz, q_len, _ = queries.size()
 94 |         bsz, v_len, _ = vision_latents.size()
 95 | 
 96 |         query_states = self.q_proj(queries)
 97 |         key_states = self.k_proj(vision_latents)
 98 |         value_states = self.v_proj(vision_latents)
 99 | 
100 |         query_states = query_states.view(
101 |             bsz, q_len, self.num_heads, self.head_dim
102 |         ).transpose(1, 2)
103 |         key_states = key_states.view(
104 |             bsz, v_len, self.num_heads, self.head_dim
105 |         ).transpose(1, 2)
106 |         value_states = value_states.view(
107 |             bsz, v_len, self.num_heads, self.head_dim
108 |         ).transpose(1, 2)
109 | 
110 |         if attention_mask is not None:
111 |             if attention_mask.size() != (bsz, 1, q_len, v_len):
112 |                 raise ValueError(
113 |                     f"Attention mask should be of size {(bsz, 1, q_len, v_len)}, but is {attention_mask.size()}"
114 |                 )
115 | 
116 |         # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
117 |         # Reference: https://github.com/pytorch/pytorch/issues/112577.
118 |         if query_states.device.type == "cuda" and attention_mask is not None:
119 |             query_states = query_states.contiguous()
120 |             key_states = key_states.contiguous()
121 |             value_states = value_states.contiguous()
122 | 
123 |         attn_output = torch.nn.functional.scaled_dot_product_attention(
124 |             query_states,
125 |             key_states,
126 |             value_states,
127 |             attn_mask=attention_mask,
128 |         )
129 | 
130 |         attn_output = attn_output.transpose(1, 2).contiguous()
131 |         attn_output = attn_output.reshape(bsz, q_len, self.hidden_dim)
132 | 
133 |         attn_output = self.o_proj(attn_output)
134 | 
135 |         return attn_output
136 | 
137 | 
138 | class AggregationBlock(nn.Module):
139 |     def __init__(
140 |         self, attention, q_dim, kv_dim, hidden_dim, num_heads, attention_bias=False
141 |     ):
142 |         super().__init__()
143 |         self.hidden_dim = hidden_dim
144 |         self.num_heads = num_heads
145 |         self.head_dim = self.hidden_dim // self.num_heads
146 | 
147 |         if (self.head_dim * self.num_heads) != self.hidden_dim:
148 |             raise ValueError(
149 |                 f"hidden_dim must be divisible by num_heads (got `hidden_dim`: {self.hidden_dim}"
150 |                 f" and `num_heads`: {self.num_heads})."
151 |             )
152 | 
153 |         self.attention = attention
154 |         if attention:
155 |             self.attention_layer = CrossAttention(
156 |                 q_dim, kv_dim, hidden_dim, num_heads, attention_bias
157 |             )
158 |         else:
159 |             self.attention_layer = MLP(kv_dim, q_dim, q_dim)
160 | 
161 |     def forward(self, vision_latents, queries, attention_mask):
162 |         if self.attention:
163 |             queries = self.attention_layer(vision_latents, queries, attention_mask)
164 |         else:
165 |             queries = self.attention_layer(vision_latents)
166 | 
167 |         return queries
168 | 
169 | 
170 | class MultiKVCrossAttention(nn.Module):
171 | 
172 |     def __init__(self, q_dim, kv_dim_list, hidden_dim, num_heads, attention_bias=False):
173 |         super().__init__()
174 | 
175 |         self.hidden_dim = hidden_dim
176 |         self.num_heads = num_heads
177 |         self.head_dim = self.hidden_dim // self.num_heads
178 | 
179 |         if (self.head_dim * self.num_heads) != self.hidden_dim:
180 |             raise ValueError(
181 |                 f"hidden_dim must be divisible by num_heads (got `hidden_dim`: {self.hidden_dim}"
182 |                 f" and `num_heads`: {self.num_heads})."
183 |             )
184 | 
185 |         self.q_proj = nn.Sequential(
186 |             nn.LayerNorm(q_dim),
187 |             nn.Linear(q_dim, self.num_heads * self.head_dim, bias=attention_bias),
188 |         )
189 |         self.num_of_kvs = len(kv_dim_list)
190 |         for i, kv_dim in enumerate(kv_dim_list):
191 |             setattr(
192 |                 self,
193 |                 "k_proj_{}".format(i),
194 |                 nn.Sequential(
195 |                     nn.LayerNorm(kv_dim),
196 |                     nn.Linear(
197 |                         kv_dim, self.num_heads * self.head_dim, bias=attention_bias
198 |                     ),
199 |                 ),
200 |             )
201 |             setattr(
202 |                 self,
203 |                 "v_proj_{}".format(i),
204 |                 nn.Sequential(
205 |                     nn.LayerNorm(kv_dim),
206 |                     nn.Linear(
207 |                         kv_dim, self.num_heads * self.head_dim, bias=attention_bias
208 |                     ),
209 |                 ),
210 |             )
211 |         self.o_proj = nn.Linear(
212 |             self.num_heads * self.head_dim, q_dim, bias=attention_bias
213 |         )
214 | 
215 |     def forward(
216 |         self,
217 |         queries,
218 |         *vision_latents_attention_mask_list,
219 |     ):
220 | 
221 |         vision_latents_list = vision_latents_attention_mask_list[: self.num_of_kvs]
222 |         attention_mask_list = vision_latents_attention_mask_list[self.num_of_kvs :]
223 | 
224 |         bsz, q_len, _ = queries.size()
225 | 
226 |         query_states = self.q_proj(queries)
227 |         key_states = torch.cat(
228 |             [
229 |                 getattr(self, "k_proj_{}".format(i))(vision_latents_list[i])
230 |                 for i in range(self.num_of_kvs)
231 |             ],
232 |             dim=1,
233 |         )
234 |         value_states = torch.cat(
235 |             [
236 |                 getattr(self, "v_proj_{}".format(i))(vision_latents_list[i])
237 |                 for i in range(self.num_of_kvs)
238 |             ],
239 |             dim=1,
240 |         )
241 | 
242 |         v_len = key_states.shape[1]
243 | 
244 |         query_states = query_states.view(
245 |             bsz, q_len, self.num_heads, self.head_dim
246 |         ).transpose(1, 2)
247 |         key_states = key_states.view(
248 |             bsz, v_len, self.num_heads, self.head_dim
249 |         ).transpose(1, 2)
250 |         value_states = value_states.view(
251 |             bsz, v_len, self.num_heads, self.head_dim
252 |         ).transpose(1, 2)
253 | 
254 |         # if kv_weight is not None:
255 |         #     kv_weight = kv_weight.unsqueeze(1).expand(-1, self.num_heads, -1, -1)
256 | 
257 |         attention_mask = torch.cat(attention_mask_list, dim=-1)
258 | 
259 |         if attention_mask is not None:
260 |             if attention_mask.size() != (bsz, 1, q_len, v_len):
261 |                 raise ValueError(
262 |                     f"Attention mask should be of size {(bsz, 1, q_len, v_len)}, but is {attention_mask.size()}"
263 |                 )
264 | 
265 |         # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
266 |         # Reference: https://github.com/pytorch/pytorch/issues/112577.
267 |         if query_states.device.type == "cuda" and attention_mask is not None:
268 |             query_states = query_states.contiguous()
269 |             key_states = key_states.contiguous()
270 |             value_states = value_states.contiguous()
271 | 
272 |         attn_output = torch.nn.functional.scaled_dot_product_attention(
273 |             query_states,
274 |             key_states,
275 |             value_states,
276 |             attn_mask=attention_mask,
277 |         )
278 |         # attn_output = spda(
279 |         #     query_states,
280 |         #     key_states,
281 |         #     value_states,
282 |         #     attn_mask=attention_mask,
283 |         #     additional_score=kv_weight
284 |         # )
285 | 
286 |         attn_output = attn_output.transpose(1, 2).contiguous()
287 |         attn_output = attn_output.reshape(bsz, q_len, self.hidden_dim)
288 | 
289 |         attn_output = self.o_proj(attn_output)
290 | 
291 |         return attn_output
292 | 
293 | 
294 | class MLP(nn.Module):
295 |     def __init__(self, d_in, d_hidden, d_out):
296 |         super().__init__()
297 |         self.linear_1 = nn.Linear(d_in, d_hidden, bias=False)
298 |         self.act = nn.GELU()
299 |         self.linear_2 = nn.Linear(d_hidden, d_out, bias=False)
300 | 
301 |     def forward(self, x):
302 |         return self.linear_2(self.act(self.linear_1(x)))
303 | 
304 | 
305 | class VisionCrossAttentionLayer(nn.Module):
306 |     def __init__(
307 |         self,
308 |         q_dim,
309 |         context_dim,
310 |         kv_dim_list,
311 |         kv_size_list,
312 |         hidden_dim=1024,
313 |         layer_idx=0,
314 |     ):
315 |         super().__init__()
316 |         num_heads = 16
317 |         self.num_of_kvs = len(kv_dim_list)
318 | 
319 |         self.proj_context = nn.Linear(context_dim, hidden_dim, bias=False)
320 |         self.proj_in = nn.Linear(q_dim + hidden_dim, hidden_dim, bias=False)
321 |         # if self.num_of_kvs > 1:
322 |         #     self.weight_mlp = MLP(q_dim+hidden_dim, hidden_dim, self.num_of_kvs)
323 |         #     self.tower_weight = nn.Parameter(torch.zeros((self.num_of_kvs)))
324 |         self.proj_out = MLP(hidden_dim, hidden_dim, q_dim)
325 | 
326 |         self.norm = nn.LayerNorm(hidden_dim)
327 | 
328 |         self.cross_attn = MultiKVCrossAttention(
329 |             hidden_dim, kv_dim_list, hidden_dim, num_heads
330 |         )
331 |         self.kv_size_list = kv_size_list
332 |         for i, kv_size in enumerate(kv_size_list):
333 |             if kv_size > 1:
334 |                 setattr(
335 |                     self,
336 |                     "pos_embed_{}".format(i),
337 |                     nn.Parameter(torch.randn(kv_size**2, hidden_dim)),
338 |                 )
339 |                 # self.register_buffer("pos_embed_{}".format(i), torch.from_numpy(get_2d_sincos_pos_embed(hidden_dim, kv_size)).float(), persistent=False)
340 | 
341 |     def forward(
342 |         self,
343 |         queries,
344 |         context_feature,
345 |         *vision_latents_attention_mask_list,
346 |     ) -> torch.FloatTensor:
347 | 
348 |         residual = queries
349 |         # queries = self.proj_in(queries)
350 |         context_feature = self.proj_context(context_feature)
351 |         # queries = queries + context_feature
352 |         queries = torch.cat([queries, context_feature], -1)
353 | 
354 |         # if self.num_of_kvs > 1:
355 |         #     kv_weight = self.weight_mlp(queries) # B * 1 * num_tower
356 |         #     kv_weight = kv_weight + self.tower_weight.view(1, 1, -1)
357 |         #     kv_weight = kv_weight.softmax(-1)
358 |         #     kv_number_list = [size**2 for size in self.kv_size_list]
359 |         #     kv_weight = torch.repeat_interleave(kv_weight, torch.tensor(kv_number_list).to(kv_weight.device), dim=-1)
360 |         # else:
361 |         #     kv_weight = None
362 | 
363 |         queries = self.proj_in(queries)
364 | 
365 |         vision_latents_list = vision_latents_attention_mask_list[: self.num_of_kvs]
366 |         attention_mask_list = vision_latents_attention_mask_list[self.num_of_kvs :]
367 | 
368 |         attention_mask_list_reshaped = []
369 |         if attention_mask_list is not None:
370 |             for attention_mask in attention_mask_list:
371 |                 attention_mask = attention_mask.view(attention_mask.shape[0], 1, 1, -1)
372 |                 attention_mask = attention_mask.expand(-1, -1, queries.shape[1], -1)
373 |                 attention_mask_list_reshaped.append(attention_mask)
374 | 
375 |         vision_latents_pos_list = []
376 |         for i, vision_latents in enumerate(vision_latents_list):
377 |             if vision_latents.shape[1] > 1:
378 |                 vision_latents_pos_list.append(
379 |                     vision_latents
380 |                     + getattr(self, "pos_embed_{}".format(i))[None, :, :].to(
381 |                         vision_latents.dtype
382 |                     )
383 |                 )
384 |             else:
385 |                 vision_latents_pos_list.append(vision_latents)
386 | 
387 |         # Cross Attention
388 |         attention_output = self.cross_attn(
389 |             queries, *vision_latents_pos_list, *attention_mask_list_reshaped
390 |         )
391 | 
392 |         # attention_output = (attention_output * combination_weight).sum(2)
393 |         queries = queries + attention_output
394 | 
395 |         queries = self.norm(queries)
396 | 
397 |         queries = self.proj_out(queries)
398 | 
399 |         queries = queries + residual
400 | 
401 |         return queries
402 | 
403 | 
404 | class VisionAggregationLayer(nn.Module):
405 |     def __init__(
406 |         self,
407 |         q_dim,
408 |         context_dim,
409 |         kv_dim_list,
410 |         kv_size_list,
411 |         hidden_dim=1024,
412 |         layer_idx=0,
413 |     ):
414 |         super().__init__()
415 |         num_heads = 16
416 |         self.num_of_kvs = len(kv_dim_list)
417 | 
418 |         self.proj_context = nn.Linear(context_dim, hidden_dim, bias=False)
419 |         self.proj_in = nn.Linear(q_dim + hidden_dim, hidden_dim, bias=False)
420 | 
421 |         self.proj_out = MLP(hidden_dim, hidden_dim, q_dim)
422 | 
423 |         self.norm = nn.LayerNorm(hidden_dim)
424 | 
425 |         if self.num_of_kvs > 1:
426 |             self.weight_mlp = MLP(q_dim + hidden_dim, hidden_dim, self.num_of_kvs)
427 | 
428 |         for i, kv_size in enumerate(kv_size_list):
429 |             if kv_size > 1:
430 |                 setattr(
431 |                     self,
432 |                     "pos_embed_{}".format(i),
433 |                     nn.Parameter(torch.randn(kv_size**2, hidden_dim)),
434 |                 )
435 |                 setattr(
436 |                     self,
437 |                     "aggregate_{}".format(i),
438 |                     AggregationBlock(
439 |                         True, hidden_dim, kv_dim_list[i], hidden_dim, num_heads
440 |                     ),
441 |                 )
442 |             else:
443 |                 setattr(
444 |                     self,
445 |                     "aggregate_{}".format(i),
446 |                     AggregationBlock(
447 |                         False, hidden_dim, kv_dim_list[i], hidden_dim, num_heads
448 |                     ),
449 |                 )
450 | 
451 |     def forward(
452 |         self,
453 |         queries,
454 |         context_feature,
455 |         *vision_latents_attention_mask_list,
456 |     ) -> torch.FloatTensor:
457 | 
458 |         residual = queries
459 |         # queries = self.proj_in(queries)
460 |         context_feature = self.proj_context(context_feature)
461 |         # queries = queries + context_feature
462 |         queries = torch.cat([queries, context_feature], -1)
463 | 
464 |         if self.num_of_kvs > 1:
465 |             combination_weight = self.weight_mlp(queries).softmax(
466 |                 -1
467 |             )  # B * 1 * num_tower
468 |             combination_weight = combination_weight.unsqueeze(-1)
469 |         else:
470 |             combination_weight = 1
471 | 
472 |         queries = self.proj_in(queries)
473 | 
474 |         vision_latents_list = vision_latents_attention_mask_list[: self.num_of_kvs]
475 |         attention_mask_list = vision_latents_attention_mask_list[self.num_of_kvs :]
476 | 
477 |         attention_mask_list_reshaped = []
478 |         if attention_mask_list is not None:
479 |             for attention_mask in attention_mask_list:
480 |                 attention_mask = attention_mask.view(attention_mask.shape[0], 1, 1, -1)
481 |                 attention_mask = attention_mask.expand(-1, -1, queries.shape[1], -1)
482 |                 attention_mask_list_reshaped.append(attention_mask)
483 | 
484 |         vision_latents_pos_list = []
485 |         for i, vision_latents in enumerate(vision_latents_list):
486 |             if vision_latents.shape[1] > 1:
487 |                 vision_latents_pos_list.append(
488 |                     vision_latents
489 |                     + getattr(self, "pos_embed_{}".format(i))[None, :, :].to(
490 |                         vision_latents.dtype
491 |                     )
492 |                 )
493 |             else:
494 |                 vision_latents_pos_list.append(vision_latents)
495 | 
496 |         aggregated_vision_latents_list = []
497 |         for i, (vision_latents, attention_mask) in enumerate(
498 |             zip(vision_latents_pos_list, attention_mask_list_reshaped)
499 |         ):
500 |             aggregated_vision_latents_list.append(
501 |                 getattr(self, "aggregate_{}".format(i))(
502 |                     vision_latents, queries, attention_mask
503 |                 )
504 |             )
505 | 
506 |         aggregated_vision_latents = torch.stack(aggregated_vision_latents_list, 2)
507 | 
508 |         queries = queries + (aggregated_vision_latents * combination_weight).sum(2)
509 | 
510 |         queries = self.norm(queries)
511 | 
512 |         queries = self.proj_out(queries)
513 | 
514 |         queries = queries + residual
515 | 
516 |         return queries
517 | 
518 | 
519 | class VisionTokenSampler(nn.Module):
520 |     def __init__(
521 |         self,
522 |         q_dim,
523 |         context_dim,
524 |         kv_dim_list,
525 |         kv_size_list,
526 |         vision_hidden_size,
527 |         num_of_layers=1,
528 |         layer_type="joint",
529 |     ):
530 |         super().__init__()
531 |         assert layer_type in ["joint", "sep"]
532 |         if layer_type == "joint":
533 |             self.layers = nn.ModuleList(
534 |                 [
535 |                     VisionCrossAttentionLayer(
536 |                         q_dim,
537 |                         context_dim,
538 |                         kv_dim_list,
539 |                         kv_size_list,
540 |                         vision_hidden_size,
541 |                         idx,
542 |                     )
543 |                     for idx in range(num_of_layers)
544 |                 ]
545 |             )
546 |         else:
547 |             self.layers = nn.ModuleList(
548 |                 [
549 |                     VisionAggregationLayer(
550 |                         q_dim,
551 |                         context_dim,
552 |                         kv_dim_list,
553 |                         kv_size_list,
554 |                         vision_hidden_size,
555 |                         idx,
556 |                     )
557 |                     for idx in range(num_of_layers)
558 |                 ]
559 |             )
560 | 
561 |     def forward(self, queries, context_feature, *vision_latents_attention_mask_list):
562 |         for layer in self.layers:
563 |             queries = layer(
564 |                 queries, context_feature, *vision_latents_attention_mask_list
565 |             )
566 |         return queries
567 | 


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
 1 | huggingface_hub==0.25.1
 2 | torch==2.1.2
 3 | numpy==1.26.4
 4 | torchvision
 5 | transformers==4.42.4
 6 | tokenizers==0.19.1
 7 | sentencepiece==0.1.99
 8 | shortuuid
 9 | accelerate==0.34.2
10 | peft==0.4.0
11 | bitsandbytes==0.41.0
12 | pydantic<2,>=1
13 | markdown2
14 | scikit-learn==1.2.2
15 | gradio==3.35.2
16 | gradio_client==0.2.9
17 | requests
18 | httpx==0.24.0
19 | uvicorn
20 | fastapi
21 | einops==0.6.1
22 | einops-exts==0.0.4
23 | timm==0.9.16
24 | decord
25 | ninja
26 | deepspeed==0.12.2
27 | protobuf
28 | iopath
29 | 


--------------------------------------------------------------------------------
/scripts/train_image_llama3_2.sh:
--------------------------------------------------------------------------------
 1 | 
 2 | BASE_CHECKPOINT="meta-llama/Llama-3.2-3B-Instruct" # replace config.json with https://huggingface.co/Vision-CAIR/LongVU_Llama3_2_3B_img/config.json
 3 | PATH_TO_JSON=""
 4 | PATH_TO_FOLDER=""
 5 | VERSION="llama3"
 6 | 
 7 | CUDA_LAUNCH_BLOCKING=1 TORCH_DISTRIBUTED_DEBUG=DETAIL torchrun --nproc_per_node=8 --nnodes=8 \
 8 | longvu/train.py \
 9 | --output_dir "/tmp/longvu/" \
10 | --input_model_filename $BASE_CHECKPOINT \
11 | --output_model_filename "./checkpoints/cambrian_llama3_2/" \
12 | --data_path $PATH_TO_JSON \
13 | --image_folder $PATH_TO_FOLDER \
14 | --model_max_length 8192 \
15 | --fp16 False \
16 | --bf16 True \
17 | --log_on_each_node False \
18 | --logging_dir /tmp/llava/test/ \
19 | --num_train_epochs 1 \
20 | --per_device_train_batch_size 4 \
21 | --per_device_eval_batch_size 4 \
22 | --gradient_accumulation_steps 1 \
23 | --save_steps 500 \
24 | --eval_steps 500 \
25 | --logging_steps 10 \
26 | --evaluation_strategy "no" \
27 | --save_strategy "steps" \
28 | --report_to "tensorboard" \
29 | --save_total_limit 1 \
30 | --learning_rate 1e-5 \
31 | --weight_decay 0. \
32 | --warmup_ratio 0.03 \
33 | --lr_scheduler_type "cosine" \
34 | --tf32 False \
35 | --version $VERSION \
36 | --mm_vision_select_layer "-2" \
37 | --mm_use_im_start_end False \
38 | --mm_use_im_patch_token False \
39 | --image_aspect_ratio pad \
40 | --group_by_modality_length True \
41 | --dataloader_num_workers 0 \
42 | --lazy_preprocess True \
43 | --tune_mm_mlp_adapter False \
44 | --freeze_mm_mlp_adapter False \
45 | --freeze_backbone False \
46 | --fsdp "full_shard auto_wrap" \
47 | --fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
48 | --gradient_checkpointing True \
49 | --mm_projector_type sva \
50 | --image_token_len 576 \
51 | --query_num_list "[576]" \
52 | --resume True \
53 | 


--------------------------------------------------------------------------------
/scripts/train_image_qwen.sh:
--------------------------------------------------------------------------------
 1 | 
 2 | BASE_CHECKPOINT="Qwen/Qwen2-7B-Instruct" # replace config.json with https://huggingface.co/Vision-CAIR/LongVU_Qwen2_7B_img/config.json
 3 | PATH_TO_JSON=""
 4 | PATH_TO_FOLDER=""
 5 | VERSION="qwen"
 6 | 
 7 | CUDA_LAUNCH_BLOCKING=1 TORCH_DISTRIBUTED_DEBUG=DETAIL torchrun --nproc_per_node=8 --nnodes=8 \
 8 | longvu/train.py \
 9 | --output_dir "/tmp/longvu/" \
10 | --input_model_filename $BASE_CHECKPOINT \
11 | --output_model_filename "./checkpoints/cambrian_qwen/" \
12 | --data_path $PATH_TO_JSON \
13 | --image_folder $PATH_TO_FOLDER \
14 | --model_max_length 8192 \
15 | --fp16 False \
16 | --bf16 True \
17 | --log_on_each_node False \
18 | --logging_dir /tmp/llava/test/ \
19 | --num_train_epochs 1 \
20 | --per_device_train_batch_size 4 \
21 | --per_device_eval_batch_size 4 \
22 | --gradient_accumulation_steps 1 \
23 | --save_steps 500 \
24 | --eval_steps 500 \
25 | --logging_steps 10 \
26 | --evaluation_strategy "no" \
27 | --save_strategy "steps" \
28 | --report_to "tensorboard" \
29 | --save_total_limit 1 \
30 | --learning_rate 1e-5 \
31 | --weight_decay 0. \
32 | --warmup_ratio 0.03 \
33 | --lr_scheduler_type "cosine" \
34 | --tf32 False \
35 | --version $VERSION \
36 | --mm_vision_select_layer "-2" \
37 | --mm_use_im_start_end False \
38 | --mm_use_im_patch_token False \
39 | --image_aspect_ratio pad \
40 | --group_by_modality_length True \
41 | --dataloader_num_workers 0 \
42 | --lazy_preprocess True \
43 | --tune_mm_mlp_adapter False \
44 | --freeze_mm_mlp_adapter False \
45 | --freeze_backbone False \
46 | --fsdp "full_shard auto_wrap" \
47 | --fsdp_transformer_layer_cls_to_wrap 'Qwen2DecoderLayer' \
48 | --gradient_checkpointing True \
49 | --mm_projector_type sva \
50 | --image_token_len 576 \
51 | --query_num_list "[576]" \
52 | --resume True \
53 | 


--------------------------------------------------------------------------------
/scripts/train_video_llama3_2.sh:
--------------------------------------------------------------------------------
 1 | 
 2 | PREV_STAGE_CHECKPOINT=""
 3 | PATH_TO_JSON=""
 4 | PATH_TO_FOLDER=""
 5 | VERSION="llama3"
 6 | 
 7 | CUDA_LAUNCH_BLOCKING=1 TORCH_DISTRIBUTED_DEBUG=DETAIL torchrun --nproc_per_node=8 --nnodes=8 \
 8 | longvu/train.py \
 9 | --output_dir "/tmp/longvu/" \
10 | --input_model_filename $PREV_STAGE_CHECKPOINT \
11 | --output_model_filename "./checkpoints/cambrian_llama3_2/" \
12 | --data_path $PATH_TO_JSON \
13 | --image_folder $PATH_TO_FOLDER \
14 | --model_max_length 8192 \
15 | --fp16 False \
16 | --bf16 True \
17 | --log_on_each_node False \
18 | --logging_dir /tmp/llava/test/ \
19 | --num_train_epochs 1 \
20 | --per_device_train_batch_size 1 \
21 | --per_device_eval_batch_size 4 \
22 | --gradient_accumulation_steps 1 \
23 | --save_steps 500 \
24 | --eval_steps 500 \
25 | --logging_steps 10 \
26 | --evaluation_strategy "no" \
27 | --save_strategy "steps" \
28 | --report_to "tensorboard" \
29 | --save_total_limit 1 \
30 | --learning_rate 5e-6 \
31 | --weight_decay 0. \
32 | --warmup_ratio 0.03 \
33 | --lr_scheduler_type "cosine" \
34 | --tf32 False \
35 | --version $VERSION \
36 | --mm_vision_select_layer "-2" \
37 | --mm_use_im_start_end False \
38 | --mm_use_im_patch_token False \
39 | --image_aspect_ratio pad \
40 | --group_by_modality_length True \
41 | --dataloader_num_workers 0 \
42 | --lazy_preprocess True \
43 | --tune_mm_mlp_adapter False \
44 | --freeze_mm_mlp_adapter False \
45 | --freeze_backbone False \
46 | --fsdp "full_shard auto_wrap" \
47 | --fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
48 | --gradient_checkpointing True \
49 | --mm_projector_type sva \
50 | --image_token_len 144 \
51 | --query_num_list "[144]" \
52 | --resume True \
53 | --lowres_token 8 \
54 | --video_fps 1 \
55 | --highres True \
56 | --drop_threshold 0.8 \
57 | 


--------------------------------------------------------------------------------
/scripts/train_video_qwen.sh:
--------------------------------------------------------------------------------
 1 | 
 2 | PREV_STAGE_CHECKPOINT=""
 3 | PATH_TO_JSON=""
 4 | PATH_TO_FOLDER=""
 5 | VERSION="qwen"
 6 | 
 7 | CUDA_LAUNCH_BLOCKING=1 TORCH_DISTRIBUTED_DEBUG=DETAIL torchrun --nproc_per_node=8 --nnodes=8 \
 8 | longvu/train.py \
 9 | --output_dir "/tmp/longvu/" \
10 | --input_model_filename $PREV_STAGE_CHECKPOINT \
11 | --output_model_filename "./checkpoints/cambrian_qwen/" \
12 | --data_path $PATH_TO_JSON \
13 | --image_folder $PATH_TO_FOLDER \
14 | --model_max_length 8192 \
15 | --fp16 False \
16 | --bf16 True \
17 | --log_on_each_node False \
18 | --logging_dir /tmp/llava/test/ \
19 | --num_train_epochs 1 \
20 | --per_device_train_batch_size 1 \
21 | --per_device_eval_batch_size 4 \
22 | --gradient_accumulation_steps 1 \
23 | --save_steps 500 \
24 | --eval_steps 500 \
25 | --logging_steps 10 \
26 | --evaluation_strategy "no" \
27 | --save_strategy "steps" \
28 | --report_to "tensorboard" \
29 | --save_total_limit 1 \
30 | --learning_rate 5e-6 \
31 | --weight_decay 0. \
32 | --warmup_ratio 0.03 \
33 | --lr_scheduler_type "cosine" \
34 | --tf32 False \
35 | --version $VERSION \
36 | --mm_vision_select_layer "-2" \
37 | --mm_use_im_start_end False \
38 | --mm_use_im_patch_token False \
39 | --image_aspect_ratio pad \
40 | --group_by_modality_length True \
41 | --dataloader_num_workers 0 \
42 | --lazy_preprocess True \
43 | --tune_mm_mlp_adapter False \
44 | --freeze_mm_mlp_adapter False \
45 | --freeze_backbone False \
46 | --fsdp "full_shard auto_wrap" \
47 | --fsdp_transformer_layer_cls_to_wrap 'Qwen2DecoderLayer' \
48 | --gradient_checkpointing True \
49 | --mm_projector_type sva \
50 | --image_token_len 144 \
51 | --query_num_list "[144]" \
52 | --resume True \
53 | --lowres_token 8 \
54 | --video_fps 1 \
55 | --highres True \
56 | --drop_threshold 0.8 \
57 | 


--------------------------------------------------------------------------------