mirror of
https://github.com/OpenVGLab/OmniLottie.git
synced 2026-09-17 07:36:27 +00:00
67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from transformers import Qwen2_5_VLForConditionalGeneration, AutoConfig
|
|
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLCausalLMOutputWithPast
|
|
from typing import Any, Dict, List, Optional, Tuple, Union
|
|
|
|
|
|
import transformers.models.qwen2_5_vl.modeling_qwen2_5_vl as qwen_modeling
|
|
|
|
|
|
|
|
class LottieDecoder(nn.Module):
|
|
"""
|
|
Autoregressive generative model for OmniLottie
|
|
"""
|
|
|
|
def __init__(self,
|
|
pix_len,
|
|
text_len,
|
|
model_path="Qwen/Qwen2.5-VL-3B-Instruct",
|
|
**kwargs):
|
|
super().__init__()
|
|
|
|
self.pix_len = pix_len
|
|
self.text_len = text_len
|
|
|
|
self.vocab_size = 192400
|
|
self.bos_token_id = 192398
|
|
self.eos_token_id = 192399
|
|
self.pad_token_id = 151643
|
|
|
|
print(f"Loading model from {model_path}...")
|
|
|
|
config = AutoConfig.from_pretrained(
|
|
model_path,
|
|
vocab_size=self.vocab_size,
|
|
bos_token_id=self.bos_token_id,
|
|
eos_token_id=self.eos_token_id,
|
|
pad_token_id=self.pad_token_id,
|
|
trust_remote_code=True
|
|
)
|
|
|
|
self.transformer = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
|
model_path,
|
|
config=config,
|
|
torch_dtype=torch.bfloat16,
|
|
attn_implementation="eager",
|
|
ignore_mismatched_sizes=True
|
|
)
|
|
|
|
self.transformer.resize_token_embeddings(self.vocab_size)
|
|
|
|
self.train()
|
|
|
|
def forward(self,
|
|
input_ids=None,
|
|
attention_mask=None,
|
|
pixel_values=None,
|
|
image_grid_thw=None,
|
|
pixel_values_videos = None,
|
|
video_grid_thw = None,
|
|
labels=None,
|
|
past_key_values=None,
|
|
use_cache=False,
|
|
**kwargs):
|
|
|
|
return NotImplementedError |