from functools import partial import logging import re from typing import Optional, Tuple, Union, List from einops import rearrange from timm.layers import LayerNorm, LayerNorm2d from timm.layers.pos_embed import resample_abs_pos_embed from timm.models.regnet import RegStage import torch from torch import nn import torch.nn.functional as F import torch.utils.checkpoint from transformers import LlamaForCausalLM from transformers.modeling_outputs import BaseModelOutput from transformers.modeling_utils import PreTrainedModel from transformers.models.auto import AutoModelForCausalLM from transformers.models.qwen2_vl.configuration_qwen2_vl import ( Qwen2VLVisionConfig, ) from transformers.models.qwen2_vl.modeling_qwen2_vl import ( PatchEmbed, Qwen2VLPreTrainedModel, Qwen2VisionTransformerPretrainedModel, Qwen2VLVisionBlock, VisionRotaryEmbedding ) from configuration import KananaVVisualProjectorConfig, KananaVConfig logger = logging.getLogger("kanana-1.5-v") def build_pos_embeds( config: KananaVVisualProjectorConfig, num_input_tokens: int, vision_hidden_size: int ): # pos emb if config.pos_emb: # ✨ 수정: num_input_tokens가 음수일 때 기본값 사용 if num_input_tokens <= 0: num_input_tokens = config.pos_emb_size if hasattr(config, 'pos_emb_size') else 576 pos_emb = torch.nn.Parameter(torch.zeros(1, num_input_tokens, vision_hidden_size)) nn.init.trunc_normal_(pos_emb, mean=0.0, std=0.02) else: pos_emb = None return pos_emb def build_eos_tokens(config: KananaVVisualProjectorConfig, output_hidden_size: int): # think tokens num_eos_tokens = config.num_eos_tokens if num_eos_tokens: eos_tokens = torch.nn.Parameter(torch.randn(1, num_eos_tokens, output_hidden_size)) nn.init.trunc_normal_(eos_tokens, mean=0.0, std=config.initializer_range) else: eos_tokens = None return eos_tokens def build_prenorm(config: KananaVVisualProjectorConfig): if getattr(config, "prenorm", False): prenorm = LayerNorm(config.encoder_hidden_size) else: prenorm = None return prenorm def build_mlp(depth: int, hidden_size: int, output_hidden_size: int): layers = [nn.Linear(hidden_size, output_hidden_size)] for _ in range(1, depth): layers.append(nn.SiLU()) layers.append(nn.Linear(output_hidden_size, output_hidden_size)) return nn.Sequential(*layers) class PatchMerge(nn.Module): def __init__(self, merge_size): super().__init__() self.merge_size = merge_size def forward(self, x, channel_last=False): if channel_last: x = rearrange(x, "B H W D -> B D H W") _, D, H, W = x.shape # 홀수 차원을 처리하기 위해 패딩 추가 pad_h = (self.merge_size - H % self.merge_size) % self.merge_size pad_w = (self.merge_size - W % self.merge_size) % self.merge_size if pad_h > 0 or pad_w > 0: print(f"🔍 PatchMerge - 패딩 추가: H={H}->{H+pad_h}, W={W}->{W+pad_w}") x = torch.nn.functional.pad(x, (0, pad_w, 0, pad_h), mode='replicate') H, W = H + pad_h, W + pad_w merged_x = rearrange( x, "B D (H h2) (W w2) -> B (D h2 w2) H W", h2=self.merge_size, w2=self.merge_size ) return merged_x class DynamicCAbstractor(nn.Module): """Dynamic C-Abstractor based on RegBlock""" def __init__(self, config: KananaVVisualProjectorConfig, num_input_tokens: int): super().__init__() self.config = config # ✨ 수정: num_input_tokens가 음수일 때 기본값 설정 if num_input_tokens <= 0: num_input_tokens = config.pos_emb_size if hasattr(config, 'pos_emb_size') else 576 self.num_input_tokens = num_input_tokens # ✨ 추가: 누락된 속성들 설정 self.merge_size = getattr(config, 'merge_size', 2) self.pos_emb_size = getattr(config, 'pos_emb_size', 576) # ✨ 최적화: 모든 레이어를 bfloat16으로 초기화 self.pos_emb = build_pos_embeds(config, num_input_tokens, config.encoder_hidden_size) if self.pos_emb is not None: self.pos_emb.data = self.pos_emb.data.to(torch.bfloat16) self.eos_tokens = build_eos_tokens(config, config.output_hidden_size) if self.eos_tokens is not None: self.eos_tokens.data = self.eos_tokens.data.to(torch.bfloat16) self.prenorm = build_prenorm(config) if self.prenorm is not None: self.prenorm = self.prenorm.to(torch.bfloat16) # ✨ 수정: build_net에서 self.net과 self.readout 설정 self.build_net() # ✨ 최적화: net 레이어들을 bfloat16으로 변환 if hasattr(self, 'net'): if isinstance(self.net, nn.ModuleList): for layer in self.net: layer = layer.to(torch.bfloat16) for module in layer.modules(): if hasattr(module, 'weight'): module.weight.data = module.weight.data.to(torch.bfloat16) if hasattr(module, 'bias') and module.bias is not None: module.bias.data = module.bias.data.to(torch.bfloat16) else: # self.net이 단일 모듈인 경우 self.net = self.net.to(torch.bfloat16) for module in self.net.modules(): if hasattr(module, 'weight'): module.weight.data = module.weight.data.to(torch.bfloat16) if hasattr(module, 'bias') and module.bias is not None: module.bias.data = module.bias.data.to(torch.bfloat16) # ✨ 최적화: readout 레이어를 bfloat16으로 변환 if hasattr(self, 'readout'): self.readout = self.readout.to(torch.bfloat16) for module in self.readout.modules(): if hasattr(module, 'weight'): module.weight.data = module.weight.data.to(torch.bfloat16) if hasattr(module, 'bias') and module.bias is not None: module.bias.data = module.bias.data.to(torch.bfloat16) def build_net(self): encoder_hidden_size = self.config.encoder_hidden_size hidden_size = self.config.hidden_size output_hidden_size = self.config.output_hidden_size depth = self.config.depth mlp_depth = self.config.mlp_depth RegBlock = partial( RegStage, stride=1, dilation=1, act_layer=nn.SiLU, norm_layer=LayerNorm2d, ) s1 = RegBlock( depth, encoder_hidden_size, hidden_size, ) sampler = PatchMerge(merge_size=self.merge_size) s2 = RegBlock( depth, self.merge_size**2 * hidden_size, hidden_size, ) if depth: self.net = nn.ModuleList([s1, sampler, s2]) self.readout = build_mlp(mlp_depth, hidden_size, output_hidden_size) else: self.net = sampler self.readout = build_mlp(mlp_depth, encoder_hidden_size, output_hidden_size) def forward(self, flattened_visual_embeds, grid_thw, **unused_kwargs): n_token_loc = torch.prod(grid_thw, dim=1) split_visual_embeds = torch.split(flattened_visual_embeds, n_token_loc.tolist()) flattened_visual_embeds = [] for _visual_embeds, _grid_thw in zip(split_visual_embeds, grid_thw): T, H, W = _grid_thw assert T == 1, "T must be 1. Video is not supported yet." reshaped_visual_embeds = rearrange( _visual_embeds, "(t h w) d -> 1 t h w d", t=T, h=H, w=W ) # remove temporal dim reshaped_visual_embeds = reshaped_visual_embeds[:, 0] if self.prenorm is not None: reshaped_visual_embeds = self.prenorm(reshaped_visual_embeds) if self.pos_emb is not None: # interpolate pos emb and add to visual embeds print(f"🔍 abstractor - pos_emb 형태: {self.pos_emb.shape}") print(f"🔍 abstractor - reshaped_visual_embeds 형태: {reshaped_visual_embeds.shape}") _local_pos_emb = resample_abs_pos_embed( posemb=self.pos_emb, old_size=tuple([int(self.pos_emb_size**0.5)] * 2), new_size=(H, W), num_prefix_tokens=0, ) _local_pos_emb = rearrange( _local_pos_emb, "1 (h w) d -> 1 h w d", h=H, w=W, ) print(f"🔍 abstractor - _local_pos_emb 형태: {_local_pos_emb.shape}") # 차원이 맞지 않는 경우 처리 if reshaped_visual_embeds.shape[-1] != _local_pos_emb.shape[-1]: print(f"🔍 abstractor - 차원 불일치 감지, pos_emb 건너뛰기") # pos_emb를 건너뛰고 visual_embeds만 사용 else: reshaped_visual_embeds = reshaped_visual_embeds + _local_pos_emb reshaped_visual_embeds = self._forward( reshaped_visual_embeds, input_size=(H, W), ) flattened_visual_embeds.append(reshaped_visual_embeds) reshaped_visual_embeds = torch.cat(flattened_visual_embeds, dim=0) output = BaseModelOutput(last_hidden_state=reshaped_visual_embeds) return output def _forward(self, x, input_size): h, w = input_size x = rearrange(x, "1 h w d -> 1 d h w", h=h, w=w) # 입력 채널 수가 맞지 않는 경우 처리 # RegStage의 첫 번째 블록에서 채널 수 확인 try: if hasattr(self.net[0], 'conv'): expected_channels = self.net[0].conv.in_channels elif hasattr(self.net[0], 'blocks') and len(self.net[0].blocks) > 0: expected_channels = self.net[0].blocks[0].conv1.in_channels else: # 기본값 사용 expected_channels = 1280 except: expected_channels = 1280 actual_channels = x.shape[1] if actual_channels != expected_channels: # 선형 변환으로 채널 수 조정 if not hasattr(self, 'channel_adapter'): # channel_adapter를 bfloat16으로 생성 self.channel_adapter = nn.Linear(actual_channels, expected_channels, dtype=torch.bfloat16).to(x.device) x = x.permute(0, 2, 3, 1) # (1, d, h, w) -> (1, h, w, d) # 입력을 bfloat16으로 변환 (한 번만) if x.dtype != torch.bfloat16: x = x.to(torch.bfloat16) x = self.channel_adapter(x) # 채널 수 조정 x = x.permute(0, 3, 1, 2) # (1, h, w, d) -> (1, d, h, w) # ✨ 최적화: 이미 bfloat16으로 초기화된 레이어들 사용 x = self.net[0](x) x = self.net[1](x) x = self.net[2](x) x = rearrange(x, "1 d h w -> (h w) d") # ✨ 최적화: 이미 bfloat16으로 초기화된 readout 사용 x = self.readout(x) return x class CustomQwen2VLVE(Qwen2VisionTransformerPretrainedModel): config_class = Qwen2VLVisionConfig _no_split_modules = ["Qwen2VLVisionBlock"] def __init__(self, config) -> None: Qwen2VLPreTrainedModel.__init__(self, config) self.spatial_merge_size = config.spatial_merge_size self.gradient_checkpointing = False self.patch_embed = PatchEmbed( patch_size=config.patch_size, temporal_patch_size=config.temporal_patch_size, in_channels=config.in_channels, embed_dim=config.embed_dim, ) head_dim = config.embed_dim // config.num_heads self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2) self.blocks = nn.ModuleList( [Qwen2VLVisionBlock(config, config._attn_implementation) for _ in range(config.depth)] ) def forward( self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, BaseModelOutput]: assert return_dict, "Only return_dict=True is supported." encoder_states = () if output_hidden_states else None hidden_states = self.patch_embed(pixel_values) rotary_pos_emb = self.rot_pos_emb(grid_thw) emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) position_embeddings = emb.cos(), emb.sin() cu_seqlens = torch.repeat_interleave( grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] ).cumsum(dim=0, dtype=torch.int32) cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) for blk in self.blocks: if output_hidden_states: encoder_states = encoder_states + (hidden_states,) if self.gradient_checkpointing and self.training: layer_outputs = torch.utils.checkpoint.checkpoint( blk.__call__, hidden_states=hidden_states, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, use_reentrant=False, ) else: layer_outputs = blk( hidden_states=hidden_states, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, ) hidden_states = layer_outputs if output_hidden_states: encoder_states = encoder_states + (hidden_states,) if not return_dict: return tuple(v for v in [hidden_states, encoder_states] if v is not None) return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states) def get_num_tokens(self): return -1 class KananaVPreTrainedModel(PreTrainedModel): """ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. """ config_class = KananaVConfig base_model_prefix = "kanana-1.5-v" supports_gradient_checkpointing = True _skip_keys_device_placement = "past_key_values" _supports_flash_attn_2 = True _supports_sdpa = True _supports_cache_class = True _supports_static_cache = False _keys_to_ignore_on_load_missing = [ r"position_ids", r"language_model.encoder.embed_tokens.weight", r"language_model.decoder.embed_tokens.weight", r"language_model.lm_head.weight", ] _no_split_modules = [ "CustomQwen2VLVE", "DynamicCAbstractor", "LlamaForCausalLM", "Parameter", ] def _init_weights(self, module): """Initialize the weights""" if ( isinstance(module, nn.Conv2d) or isinstance(module, nn.Embedding) or isinstance(module, nn.Linear) ): module.weight.data.normal_(mean=0.0, std=0.02) if hasattr(module, "bias") and module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.LayerNorm): module.bias.data.zero_() module.weight.data.fill_(1.0) elif isinstance(module, nn.Parameter): raise ValueError() class KananaVForConditionalGeneration(KananaVPreTrainedModel): config_class = KananaVConfig def __init__(self, config: KananaVConfig): super().__init__(config) logger.info("Build vision model ...") self.vision_model = CustomQwen2VLVE._from_config(config.vision_config) logger.info("Build projector ...") self.abstractor = DynamicCAbstractor(config.projector_config, num_input_tokens=self.vision_model.get_num_tokens()) logger.info("Build language model ...") self.language_model = LlamaForCausalLM._from_config(config=config.text_config) self.post_init() def forward_vision(self, pixel_values: Union[torch.Tensor, List[torch.Tensor]], image_metas: Optional[dict] = None): # ✨ 핵심 수정: pixel_values가 리스트일 경우와 텐서일 경우를 모두 처리 if isinstance(pixel_values, list): # 다중 이미지: 각 이미지를 처리하여 결과를 합침 visual_features_list = [] for i, pv in enumerate(pixel_values): single_image_metas = {k: v[i] for k, v in image_metas.items()} # grid_thw 처리 수정 vision_grid_thw = single_image_metas["vision_grid_thw"] if isinstance(vision_grid_thw, (list, tuple)): # 튜플을 리스트로 변환하여 텐서 생성 grid_thw = torch.tensor([list(vision_grid_thw)]).to(pv.device) else: grid_thw = torch.tensor([vision_grid_thw]).to(pv.device) # ✨ 최적화: 불필요한 dtype 변환 제거 v_outputs = self.vision_model( pixel_values=pv.unsqueeze(0), grid_thw=grid_thw, return_dict=True, output_hidden_states=True ) layer_index = self.config.projector_config.feature_layer_index visual_features_list.append(self._get_visual_feature_at(v_outputs.hidden_states, layer_index)) return visual_features_list # 리스트 형태로 반환 else: # 단일 이미지 - 이미 분리된 특징 텐서 # grid_thw가 리스트인 경우 첫 번째 요소 사용 grid_thw = image_metas["vision_grid_thw"] if isinstance(grid_thw, list): grid_thw = grid_thw[0] # grid_thw를 텐서로 변환 if not isinstance(grid_thw, torch.Tensor): if isinstance(grid_thw, (list, tuple)): # 튜플을 리스트로 변환하여 텐서 생성 grid_thw = torch.tensor([list(grid_thw)]) else: grid_thw = torch.tensor([grid_thw]) # 디바이스 정보 추가 if hasattr(pixel_values, 'device'): grid_thw = grid_thw.to(pixel_values.device) # pixel_values가 2D 특징 텐서인 경우 vision_model을 통해 처리 if len(pixel_values.shape) == 2: # 2D 특징 텐서를 vision_model이 처리할 수 있는 형태로 변환 # 다중 이미지와 동일한 방식으로 처리하되, 올바른 형태로 변환 # pixel_values를 (1, 3, H, W) 형태로 재구성 # 다중 이미지에서 사용하는 방식과 동일하게 처리 if len(pixel_values.shape) == 2: # 2D 텐서를 vision_model이 처리할 수 있는 형태로 변환 # 다중 이미지에서는 이미 올바른 형태로 전달되므로 동일하게 처리 # ✨ 최적화: 불필요한 dtype 변환 제거 v_outputs = self.vision_model( pixel_values=pixel_values, grid_thw=grid_thw, return_dict=True, output_hidden_states=True ) layer_index = self.config.projector_config.feature_layer_index return self._get_visual_feature_at(v_outputs.hidden_states, layer_index) else: return pixel_values # ✨ 최적화: 불필요한 dtype 변환 제거 v_outputs = self.vision_model( pixel_values=pixel_values, grid_thw=grid_thw, return_dict=True, output_hidden_states=True ) layer_index = self.config.projector_config.feature_layer_index return self._get_visual_feature_at(v_outputs.hidden_states, layer_index) def forward_projector(self, visual_features: Union[torch.Tensor, List[torch.Tensor]], image_metas: Optional[dict] = None): print(f"🔍 forward_projector - visual_features 형태: {visual_features.shape if hasattr(visual_features, 'shape') else type(visual_features)}") # ✨ 핵심 수정: visual_features가 리스트일 경우 처리 if isinstance(visual_features, list): print(f"🔍 forward_projector - 리스트 형태 처리") visual_embeds_list = [] for i, vf in enumerate(visual_features): single_image_metas = {k: v[i] for k, v in image_metas.items()} # grid_thw 처리 수정 vision_grid_thw = single_image_metas["vision_grid_thw"] if isinstance(vision_grid_thw, (list, tuple)): # 튜플을 리스트로 변환하여 텐서 생성 grid_thw = torch.tensor([list(vision_grid_thw)]).to(vf.device) else: grid_thw = torch.tensor([vision_grid_thw]).to(vf.device) print(f"🔍 forward_projector - 이미지 {i} 처리 중") print(f"🔍 forward_projector - 이미지 {i} 특징 형태: {vf.shape}") print(f"🔍 forward_projector - 이미지 {i} grid_thw: {grid_thw}") visual_embeds = self.abstractor(vf, grid_thw=grid_thw)["last_hidden_state"] print(f"🔍 forward_projector - 이미지 {i} visual_embeds 형태: {visual_embeds.shape}") visual_embeds_list.append(visual_embeds) return torch.cat(visual_embeds_list, dim=0) # 최종적으로 하나의 텐서로 합쳐서 반환 else: # 단일 이미지 print(f"🔍 forward_projector - 단일 텐서 처리") # visual_features가 이미 처리된 특징 텐서인 경우 if len(visual_features.shape) == 2: print(f"🔍 forward_projector - 이미 처리된 특징 텐서 감지") print(f"🔍 forward_projector - 특징 텐서 형태: {visual_features.shape}") # grid_thw가 리스트인 경우 첫 번째 요소 사용 grid_thw = image_metas["vision_grid_thw"] if isinstance(grid_thw, list): grid_thw = grid_thw[0] # grid_thw를 텐서로 변환 if not isinstance(grid_thw, torch.Tensor): if isinstance(grid_thw, (list, tuple)): # 튜플을 리스트로 변환하여 텐서 생성 grid_thw = torch.tensor([list(grid_thw)]) else: grid_thw = torch.tensor([grid_thw]) # 디바이스 정보 추가 if hasattr(visual_features, 'device'): grid_thw = grid_thw.to(visual_features.device) print(f"🔍 forward_projector - grid_thw: {grid_thw}") print(f"🔍 forward_projector - grid_thw 계산된 토큰 수: {torch.prod(grid_thw, dim=1)}") print(f"🔍 forward_projector - 실제 특징 텐서 토큰 수: {visual_features.shape[0]}") # grid_thw가 실제 토큰 수와 맞지 않는 경우 수정 actual_tokens = visual_features.shape[0] if torch.prod(grid_thw, dim=1).item() != actual_tokens: print(f"🔍 forward_projector - grid_thw 수정 필요") # 실제 토큰 수에 맞는 grid_thw 계산 # 이미지의 실제 비율을 고려하여 계산 T = 1 # 이미지 메타데이터에서 실제 크기 정보 사용 if 'hw_orig_resolution' in image_metas: orig_h, orig_w = image_metas['hw_orig_resolution'] if isinstance(orig_h, list): orig_h = orig_h[0] if isinstance(orig_h[0], (int, float)) else orig_h[0][0] if isinstance(orig_w, list): orig_w = orig_w[0] if isinstance(orig_w[0], (int, float)) else orig_w[0][0] # 실제 비율을 유지하면서 토큰 수에 맞게 조정 aspect_ratio = orig_w / orig_h H = int((actual_tokens / aspect_ratio) ** 0.5) W = int(actual_tokens / H) # 토큰 수가 맞지 않으면 조정 while H * W != actual_tokens and H > 1 and W > 1: if H * W > actual_tokens: H -= 1 W = int(actual_tokens / H) else: W += 1 H = int(actual_tokens / W) else: # 기본값 사용 H = int(actual_tokens ** 0.5) W = actual_tokens // H if actual_tokens % H != 0: W += 1 corrected_grid_thw = torch.tensor([[T, H, W]]) print(f"🔍 forward_projector - 수정된 grid_thw: {corrected_grid_thw}") print(f"🔍 forward_projector - 수정된 토큰 수: {torch.prod(corrected_grid_thw, dim=1)}") # 토큰 수가 맞지 않는 경우 패딩 또는 자르기 expected_tokens = torch.prod(corrected_grid_thw, dim=1).item() if expected_tokens > actual_tokens: # 패딩 padding_size = expected_tokens - actual_tokens padding = torch.zeros(padding_size, visual_features.shape[1], device=visual_features.device) visual_features = torch.cat([visual_features, padding], dim=0) print(f"🔍 forward_projector - 패딩 추가: {padding_size}개 토큰") elif expected_tokens < actual_tokens: # 자르기 visual_features = visual_features[:expected_tokens] print(f"🔍 forward_projector - 토큰 자르기: {expected_tokens}개로") grid_thw = corrected_grid_thw # 특징 텐서를 abstractor에 직접 전달 visual_embeds = self.abstractor(visual_features, grid_thw=grid_thw)["last_hidden_state"] print(f"🔍 forward_projector - abstractor 출력 형태: {visual_embeds.shape}") return visual_embeds else: # 일반적인 vision model 출력 grid_thw = image_metas["vision_grid_thw"] return self.abstractor(visual_features, grid_thw=grid_thw)["last_hidden_state"] def forward_and_project_vision(self, pixel_values, image_metas: Optional[dict] = None): visual_features = self.forward_vision(pixel_values, image_metas=image_metas) visual_embeds = self.forward_projector(visual_features, image_metas=image_metas) return visual_embeds def _get_visual_feature_at(self, v_output, layer_index): if isinstance(layer_index, list): visual_features = torch.stack(v_output, dim=1)[:, layer_index] # [B, n_scales, L, dim] else: visual_features = v_output[layer_index] # [B, L, dim] return visual_features def embed_text_tokens(self, input_ids): """Embed input_ids into text_embeds, ignoring media tokens (negative values).""" input_ids = input_ids.clone() input_ids[input_ids < 0] = 0 text_embeds = self.language_model.get_input_embeddings()(input_ids) if hasattr(self.language_model, "transformer") and hasattr( self.language_model.transformer, "word_embeddings_layernorm" ): text_embeds = self.language_model.transformer.word_embeddings_layernorm(text_embeds) return text_embeds def prepare_mm_inputs( self, input_ids: torch.FloatTensor, pixel_values: Optional[list[torch.FloatTensor]] = None, image_metas: Optional[dict] = None, attention_mask: Optional[torch.LongTensor] = None, ): """Prepare multimodal inputs from input_ids and pixel_values.""" if pixel_values is not None: # pixel_values가 리스트인 경우 각각을 변환 if isinstance(pixel_values, list): pixel_values = [pv.to(self._get_input_dtype()) for pv in pixel_values] else: pixel_values = pixel_values.to(self._get_input_dtype()) if attention_mask is None: attention_mask = input_ids.new_ones(*input_ids.shape) # Get Text Embeddings text_embeds = self.embed_text_tokens(input_ids) flattened_text_embeds = rearrange(text_embeds, "b l d -> (b l) d") flattened_input_ids = rearrange(input_ids, "b l -> (b l)") # Get Visual Embeddings if pixel_values is not None: print(f"🔍 prepare_mm_inputs - pixel_values 타입: {type(pixel_values)}") if hasattr(pixel_values, 'shape'): print(f"🔍 prepare_mm_inputs - pixel_values 형태: {pixel_values.shape}") if isinstance(pixel_values, list): print(f"🔍 prepare_mm_inputs - pixel_values 길이: {len(pixel_values)}") # 다중 이미지 처리: 각 이미지를 개별적으로 처리 if isinstance(pixel_values, list) and len(pixel_values) > 1: print(f"🔍 prepare_mm_inputs - 다중 이미지 처리 시작") visual_embeds_list = [] for i, single_pixel_values in enumerate(pixel_values): print(f"🔍 prepare_mm_inputs - 이미지 {i} 처리 중") # 각 이미지에 대한 개별 image_metas 생성 single_image_metas = {} for key, value_list in image_metas.items(): if isinstance(value_list, list): single_image_metas[key] = value_list[i] else: single_image_metas[key] = value_list # 개별 이미지 처리 single_visual_embeds = self.forward_and_project_vision( single_pixel_values.unsqueeze(0), single_image_metas ) visual_embeds_list.append(single_visual_embeds) # 모든 이미지의 visual embeds를 연결 flattened_visual_embeds = torch.cat(visual_embeds_list, dim=0) print(f"🔍 prepare_mm_inputs - 다중 이미지 처리 완료, 연결된 embeds 크기: {flattened_visual_embeds.shape}") else: # 단일 이미지 처리 (기존 방식) print(f"🔍 prepare_mm_inputs - 단일 이미지 처리") # pixel_values가 이미 처리된 특징 텐서인 경우 (다중 이미지 결합) if hasattr(pixel_values, 'shape') and len(pixel_values.shape) == 2: print(f"🔍 prepare_mm_inputs - 처리된 특징 텐서 감지, 다중 이미지로 분리 시도") # image_metas에서 이미지 개수 확인 num_images = 0 if isinstance(image_metas, dict) and "image_token_thw" in image_metas: num_images = len(image_metas["image_token_thw"]) print(f"🔍 prepare_mm_inputs - 감지된 이미지 개수: {num_images}") if num_images > 1: print(f"🔍 prepare_mm_inputs - {num_images}개 이미지로 분리 처리") visual_embeds_list = [] # 각 이미지의 실제 토큰 수 계산 current_idx = 0 for i in range(num_images): print(f"🔍 prepare_mm_inputs - 이미지 {i} 처리 중") # 각 이미지에 대한 개별 image_metas 생성 single_image_metas = {} for key, value_list in image_metas.items(): if isinstance(value_list, list): single_image_metas[key] = value_list[i] else: single_image_metas[key] = value_list # image_token_thw에서 실제 토큰 수 계산 if "image_token_thw" in single_image_metas: token_thw = single_image_metas["image_token_thw"] if isinstance(token_thw, (list, tuple)): tokens_per_image = token_thw[0] * token_thw[1] * token_thw[2] else: tokens_per_image = token_thw[0] * token_thw[1] * token_thw[2] print(f"🔍 prepare_mm_inputs - 이미지 {i} 실제 토큰 수: {tokens_per_image}") else: # 기본값 사용 tokens_per_image = pixel_values.shape[0] // num_images print(f"🔍 prepare_mm_inputs - 이미지 {i} 기본 토큰 수: {tokens_per_image}") # pixel_values에서 해당 이미지 부분 추출 start_idx = current_idx end_idx = current_idx + tokens_per_image single_pixel_values = pixel_values[start_idx:end_idx] print(f"🔍 prepare_mm_inputs - 이미지 {i} 특징 형태: {single_pixel_values.shape}") # 개별 이미지 처리 single_visual_embeds = self.forward_and_project_vision( single_pixel_values, single_image_metas ) visual_embeds_list.append(single_visual_embeds) current_idx += tokens_per_image # 모든 이미지의 visual embeds를 연결 flattened_visual_embeds = torch.cat(visual_embeds_list, dim=0) print(f"🔍 prepare_mm_inputs - 다중 이미지 처리 완료, 연결된 embeds 크기: {flattened_visual_embeds.shape}") else: # 단일 이미지 처리 print(f"🔍 prepare_mm_inputs - 단일 이미지로 처리") flattened_visual_embeds = self.forward_and_project_vision( pixel_values, image_metas ) # pixel_values가 배치 형태인 경우 개별 이미지로 분리 elif hasattr(pixel_values, 'shape') and len(pixel_values.shape) == 4 and pixel_values.shape[0] > 1: print(f"🔍 prepare_mm_inputs - 배치 형태 감지, 개별 이미지로 분리") visual_embeds_list = [] for i in range(pixel_values.shape[0]): print(f"🔍 prepare_mm_inputs - 배치 이미지 {i} 처리 중") # 각 이미지에 대한 개별 image_metas 생성 single_image_metas = {} for key, value_list in image_metas.items(): if isinstance(value_list, list): single_image_metas[key] = value_list[i] else: single_image_metas[key] = value_list # 개별 이미지 처리 if isinstance(pixel_values, list): single_pixel_values = pixel_values[i:i+1] else: # pixel_values가 텐서인 경우 single_pixel_values = pixel_values[i:i+1] single_visual_embeds = self.forward_and_project_vision( single_pixel_values, single_image_metas ) visual_embeds_list.append(single_visual_embeds) # 모든 이미지의 visual embeds를 연결 flattened_visual_embeds = torch.cat(visual_embeds_list, dim=0) print(f"🔍 prepare_mm_inputs - 다중 이미지 처리 완료, 연결된 embeds 크기: {flattened_visual_embeds.shape}") # 각 이미지의 embeds 크기 출력 for i, embeds in enumerate(visual_embeds_list): print(f"🔍 prepare_mm_inputs - 이미지 {i} embeds 크기: {embeds.shape}") else: # 단일 이미지 처리 # image_metas가 다중 이미지 정보를 포함하는 경우 첫 번째 이미지 정보만 사용 if isinstance(image_metas, dict): single_image_metas = {} for key, value_list in image_metas.items(): if isinstance(value_list, list): single_image_metas[key] = value_list[0] # 첫 번째 이미지 정보 사용 else: single_image_metas[key] = value_list print(f"🔍 prepare_mm_inputs - 단일 이미지 처리, 첫 번째 이미지 정보 사용") else: single_image_metas = image_metas # 단일 이미지 처리 시 pixel_values가 리스트인지 확인 if isinstance(pixel_values, list): single_pixel_values = pixel_values[0] # 첫 번째 이미지만 사용 else: single_pixel_values = pixel_values flattened_visual_embeds = self.forward_and_project_vision( single_pixel_values, single_image_metas ) # dtype 일치를 위해 visual_embeds를 text_embeds와 같은 dtype으로 변환 flattened_visual_embeds = flattened_visual_embeds.to(flattened_text_embeds.dtype) # visual embeds와 -1 토큰 개수 확인 및 조정 num_visual_tokens = flattened_visual_embeds.shape[0] num_neg_one_tokens = (flattened_input_ids == -1).sum().item() print(f"🔍 prepare_mm_inputs - visual embeds 개수: {num_visual_tokens}") print(f"🔍 prepare_mm_inputs - -1 토큰 개수: {num_neg_one_tokens}") if num_visual_tokens != num_neg_one_tokens: print(f"🔍 prepare_mm_inputs - 토큰 개수 불일치, 조정 필요") if num_visual_tokens > num_neg_one_tokens: # visual embeds가 많으면 자르기 flattened_visual_embeds = flattened_visual_embeds[:num_neg_one_tokens] print(f"🔍 prepare_mm_inputs - visual embeds 자르기: {num_visual_tokens} -> {num_neg_one_tokens}") else: # visual embeds가 적으면 반복해서 사용 repeat_times = num_neg_one_tokens // num_visual_tokens remainder = num_neg_one_tokens % num_visual_tokens if repeat_times > 0: # visual embeds를 반복 repeated_embeds = flattened_visual_embeds.repeat(repeat_times, 1) if remainder > 0: # 나머지 부분 추가 remainder_embeds = flattened_visual_embeds[:remainder] repeated_embeds = torch.cat([repeated_embeds, remainder_embeds], dim=0) flattened_visual_embeds = repeated_embeds else: # visual embeds가 너무 적으면 첫 번째 토큰을 반복 first_token = flattened_visual_embeds[0:1].repeat(num_neg_one_tokens, 1) flattened_visual_embeds = first_token print(f"🔍 prepare_mm_inputs - visual embeds 반복: {num_visual_tokens} -> {num_neg_one_tokens}") flattened_text_embeds[flattened_input_ids == -1] = flattened_visual_embeds input_embeds = rearrange( flattened_text_embeds, "(b l) d -> b l d", b=input_ids.shape[0] ) return_inputs = { "inputs_embeds": input_embeds, "attention_mask": attention_mask, } return return_inputs def forward( self, pixel_values: list[torch.FloatTensor], image_metas: dict[list], input_ids: torch.FloatTensor, seq_length: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.LongTensor] = None, labels: Optional[torch.LongTensor] = None, return_dict: Optional[bool] = None, ): return_dict = return_dict if return_dict is not None else self.config.use_return_dict inputs = self.prepare_mm_inputs( input_ids=input_ids, pixel_values=pixel_values, image_metas=image_metas, attention_mask=attention_mask, ) outputs = self.language_model( **inputs, labels=labels, position_ids=None, return_dict=return_dict, output_attentions=self.config.output_attentions, ) return outputs @torch.no_grad() def generate( self, pixel_values: torch.FloatTensor = None, image_metas: dict[list] = None, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.LongTensor] = None, seq_length: Optional[torch.LongTensor] = None, **generate_kwargs, ) -> torch.LongTensor: """ Overrides `generate` function to be able to use the model as a conditional generator. Args: pixel_values (`torch.FloatTensor` of shape (batch_size, num_channels, height, width)): Input images to be processed. input_ids (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*): The sequence used as a prompt for the generation. attention_mask (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*): Mask to avoid performing attention on padding token indices Returns: captions (list): A list of strings of length batch_size * num_captions. """ if input_ids is None: return self.language_model.generate(attention_mask=attention_mask, **generate_kwargs) if pixel_values is None: return self.language_model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs) if ( image_metas is not None and image_metas.get("vision_grid_thw") is not None and isinstance(image_metas.get("vision_grid_thw"), torch.Tensor) ): image_metas["vision_grid_thw"] = image_metas["vision_grid_thw"].to(input_ids.device) inputs = self.prepare_mm_inputs( input_ids=input_ids, pixel_values=pixel_values, image_metas=image_metas, attention_mask=attention_mask, ) outputs = self.language_model.generate( **inputs, **generate_kwargs, ) return outputs def _get_input_dtype(self): dtype = next(self.vision_model.parameters()).dtype return dtype