# Copyright (c) 2023-2024 DeepSeek. # # Permission is hereby granted, free of charge, to any person obtaining a copy of # this software and associated documentation files (the "Software"), to deal in # the Software without restriction, including without limitation the rights to # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of # the Software, and to permit persons to whom the Software is furnished to do so, # subject to the following conditions: # # The above copyright notice and this permission notice shall be included in all # copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. from math import e import torch from attrdict import AttrDict from einops import rearrange from transformers import ( AutoConfig, AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM, PreTrainedModel, ) from transformers.modeling_outputs import CausalLMOutputWithPast from torch.nn import CrossEntropyLoss from transformers.configuration_utils import PretrainedConfig from janus.models.clip_encoder import CLIPVisionTower from janus.models.projector import MlpProjector class vision_head(torch.nn.Module): def __init__(self, params): super().__init__() self.output_mlp_projector = torch.nn.Linear( params.n_embed, params.image_token_embed ) self.vision_activation = torch.nn.GELU() self.vision_head = torch.nn.Linear( params.image_token_embed, params.image_token_size ) def forward(self, x): x = self.output_mlp_projector(x) x = self.vision_activation(x) x = self.vision_head(x) return x def model_name_to_cls(cls_name): if "MlpProjector" in cls_name: cls = MlpProjector elif "CLIPVisionTower" in cls_name: cls = CLIPVisionTower elif "VQ" in cls_name: from janus.models.vq_model import VQ_models cls = VQ_models[cls_name] elif "vision_head" in cls_name: cls = vision_head else: raise ValueError(f"class_name {cls_name} is invalid.") return cls class VisionConfig(PretrainedConfig): model_type = "vision" cls: str = "" params: AttrDict = {} def __init__(self, **kwargs): super().__init__(**kwargs) self.cls = kwargs.get("cls", "") if not isinstance(self.cls, str): self.cls = self.cls.__name__ self.params = AttrDict(kwargs.get("params", {})) class AlignerConfig(PretrainedConfig): model_type = "aligner" cls: str = "" params: AttrDict = {} def __init__(self, **kwargs): super().__init__(**kwargs) self.cls = kwargs.get("cls", "") if not isinstance(self.cls, str): self.cls = self.cls.__name__ self.params = AttrDict(kwargs.get("params", {})) class GenVisionConfig(PretrainedConfig): model_type = "gen_vision" cls: str = "" params: AttrDict = {} def __init__(self, **kwargs): super().__init__(**kwargs) self.cls = kwargs.get("cls", "") if not isinstance(self.cls, str): self.cls = self.cls.__name__ self.params = AttrDict(kwargs.get("params", {})) class GenAlignerConfig(PretrainedConfig): model_type = "gen_aligner" cls: str = "" params: AttrDict = {} def __init__(self, **kwargs): super().__init__(**kwargs) self.cls = kwargs.get("cls", "") if not isinstance(self.cls, str): self.cls = self.cls.__name__ self.params = AttrDict(kwargs.get("params", {})) class GenHeadConfig(PretrainedConfig): model_type = "gen_head" cls: str = "" params: AttrDict = {} def __init__(self, **kwargs): super().__init__(**kwargs) self.cls = kwargs.get("cls", "") if not isinstance(self.cls, str): self.cls = self.cls.__name__ self.params = AttrDict(kwargs.get("params", {})) from dataclasses import dataclass @dataclass class VLChatProcessorOutput(): sft_format: str input_ids: torch.Tensor pixel_values: torch.Tensor num_image_tokens: torch.IntTensor def __len__(self): return len(self.input_ids) class MultiModalityConfig(PretrainedConfig): model_type = "multi_modality" vision_config: VisionConfig aligner_config: AlignerConfig gen_vision_config: GenVisionConfig gen_aligner_config: GenAlignerConfig gen_head_config: GenHeadConfig language_config: LlamaConfig def __init__(self, **kwargs): super().__init__(**kwargs) vision_config = kwargs.get("vision_config", {}) self.vision_config = VisionConfig(**vision_config) aligner_config = kwargs.get("aligner_config", {}) self.aligner_config = AlignerConfig(**aligner_config) gen_vision_config = kwargs.get("gen_vision_config", {}) self.gen_vision_config = GenVisionConfig(**gen_vision_config) gen_aligner_config = kwargs.get("gen_aligner_config", {}) self.gen_aligner_config = GenAlignerConfig(**gen_aligner_config) gen_head_config = kwargs.get("gen_head_config", {}) self.gen_head_config = GenHeadConfig(**gen_head_config) language_config = kwargs.get("language_config", {}) if isinstance(language_config, LlamaConfig): self.language_config = language_config else: self.language_config = LlamaConfig(**language_config) class MultiModalityPreTrainedModel(PreTrainedModel): config_class = MultiModalityConfig base_model_prefix = "multi_modality" _no_split_modules = [] _skip_keys_device_placement = "past_key_values" class MultiModalityCausalLM(MultiModalityPreTrainedModel): def __init__(self, config: MultiModalityConfig): super().__init__(config) vision_config = config.vision_config vision_cls = model_name_to_cls(vision_config.cls) self.vision_model = vision_cls(**vision_config.params) aligner_config = config.aligner_config aligner_cls = model_name_to_cls(aligner_config.cls) self.aligner = aligner_cls(aligner_config.params) gen_vision_config = config.gen_vision_config gen_vision_cls = model_name_to_cls(gen_vision_config.cls) self.gen_vision_model = gen_vision_cls() gen_aligner_config = config.gen_aligner_config gen_aligner_cls = model_name_to_cls(gen_aligner_config.cls) self.gen_aligner = gen_aligner_cls(gen_aligner_config.params) gen_head_config = config.gen_head_config gen_head_cls = model_name_to_cls(gen_head_config.cls) self.gen_head = gen_head_cls(gen_head_config.params) self.gen_embed = torch.nn.Embedding( gen_vision_config.params.image_token_size, gen_vision_config.params.n_embed ) language_config = config.language_config self.language_model = LlamaForCausalLM(language_config) def prepare_inputs_embeds( self, input_ids: torch.LongTensor, pixel_values: torch.FloatTensor, images_seq_mask: torch.LongTensor=None, images_emb_mask: torch.LongTensor=None, **kwargs, ): """ Args: input_ids (torch.LongTensor): [b, T] pixel_values (torch.FloatTensor): [b, n_images, 3, h, w] images_seq_mask (torch.BoolTensor): [b, T] images_emb_mask (torch.BoolTensor): [b, n_images, n_image_tokens] assert torch.sum(images_seq_mask) == torch.sum(images_emb_mask) Returns: input_embeds (torch.Tensor): [b, T, D] """ # bs, n = pixel_values.shape[0:2] # images = rearrange(pixel_values, "b n c h w -> (b n) c h w") # # [b x n, T2, D] # images_embeds = self.aligner(self.vision_model(images)) # # # [b x n, T2, D] -> [b, n x T2, D] # images_embeds = rearrange(images_embeds, "(b n) t d -> b (n t) d", b=bs, n=n) # # [b, n, T2] -> [b, n x T2] # # images_emb_mask = rearrange(images_emb_mask, "b n t -> b (n t)") # # # [b, T, D] # # input_ids[input_ids < 0] = 0 # ignore the image embeddings # inputs_embeds = self.language_model.get_input_embeddings()(input_ids) # # # replace with the image embeddings # # inputs_embeds[images_seq_mask] = images_embeds[images_emb_mask] # # return inputs_embeds, images_embeds bs, n = pixel_values.shape[0:2] print('px.shape', pixel_values.shape) images = rearrange(pixel_values, "b n c h w -> (b n) c h w") # [b x n, T2, D] images_embeds = self.aligner(self.vision_model(images)) # [b x n, T2, D] -> [b, n x T2, D] images_embeds = rearrange(images_embeds, "(b n) t d -> b (n t) d", b=bs, n=n) # [b, n, T2] -> [b, n x T2] images_emb_mask = rearrange(images_emb_mask, "b n t -> b (n t)") # [b, T, D] input_ids[input_ids < 0] = 0 # ignore the image embeddings inputs_embeds = self.language_model.get_input_embeddings()(input_ids) # replace with the image embeddings print('input_ids' ,input_ids.shape) print('images_seq_mask ',images_seq_mask.shape) print('inputs_embeds ',inputs_embeds.shape) print('images_embeds ',images_embeds.shape) print('images_emb_mask ',images_emb_mask.shape) inputs_embeds[images_seq_mask] = images_embeds[images_emb_mask] return inputs_embeds def prepare_gen_img_embeds(self, image_ids: torch.LongTensor): return self.gen_aligner(self.gen_embed(image_ids)) def forward(self,vl_chat_processor, input_ids, labels=None, task="understanding", return_dict=True, pixel_values=None, images_seq_mask=None, images_emb_mask=None, **kwargs): if task == "understanding": inputs_embeds = self.prepare_inputs_embeds(input_ids, pixel_values, images_seq_mask, images_emb_mask) return self.language_model.forward( inputs_embeds=inputs_embeds, labels=labels, **kwargs ) elif task == "generation": print('LLLLLLLLLLL ',pixel_values) print(kwargs) image_token_num_per_image = 576 cfg_weight = 5 temperature = 1 tokens = torch.zeros((2*input_ids.size(0), input_ids.size(1)), dtype=torch.int).cuda() for i in range(2): tokens[i*input_ids.size(0):(i+1)*input_ids.size(0), :] = input_ids if i % 2 != 0: tokens[i*input_ids.size(0):(i+1)*input_ids.size(0), 1:-1] = 100015 # pad_id inputs_embeds = self.language_model.get_input_embeddings()(tokens) generated_tokens = torch.zeros((2*input_ids.size(0), image_token_num_per_image), dtype=torch.int).cuda() outputs = self.language_model.model(inputs_embeds=inputs_embeds, use_cache=True, past_key_values=None, labels=labels) hidden_states = outputs.last_hidden_state logits = self.gen_head(hidden_states) logits_cond = logits[0::2, :] logits_uncond = logits[1::2, :] all_logits = logits_uncond + cfg_weight * (logits_cond - logits_uncond) loss_fct = CrossEntropyLoss() shift_logits = all_logits[..., :-1, :].contiguous() shift_logits = shift_logits.view(-1, self.config.gen_head_config.params.image_token_size) if labels is not None: shift_labels = labels[..., 1:].contiguous() shift_labels = shift_labels.view(-1) shift_labels = shift_labels.to(shift_logits.device) loss = loss_fct(shift_logits, shift_labels) else: loss = None if not return_dict: output = (logits,) + outputs[1:] return ((loss,) + output) if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) elif task == "generation_direct": outputs = self.language_model.model(input_ids=input_ids, **kwargs) hidden_states = outputs[0] # possibly outputs[0] logits = self.gen_head(hidden_states) loss = None logits = logits.float() # Shift so that tokens < n predict n shift_logits = logits[..., :-1, :].contiguous() shift_logits = shift_logits.view(-1, self.config.gen_head_config.params.image_token_size) if labels is not None: shift_labels = labels[..., 1:].contiguous() # Flatten the tokens loss_fct = CrossEntropyLoss() shift_labels = shift_labels.view(-1) # Enable model parallelism shift_labels = shift_labels.to(shift_logits.device) loss = loss_fct(shift_logits, shift_labels) else: loss = None if not return_dict: output = (logits,) + outputs[1:] return ((loss,) + output) if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) elif task == "image_editing": # image_token_num_per_image = 576 # img_size = 384 # patch_size = 16 # cfg_weight = kwargs.get('cfg_weight', 5) # cfg_weight2 = kwargs.get('cfg_weight2', 5) # temperature = kwargs.get('temperature', 1.0) # parallel_size = kwargs.get('parallel_size', input_ids.size(0)) # # # 构造tokens: 每个输入生成3个版本 (cond_full, cond_part, uncond) # tokens = torch.zeros((3 * input_ids.size(0), input_ids.size(1)), dtype=torch.int).cuda() # pre_data = [] # img_len = len(kwargs['source_image']) # # # 处理输入图像 # import PIL.Image # images = [PIL.Image.open(image_path).convert("RGB") for image_path in kwargs['source_image']] # encoder_pixel_values = vl_chat_processor.image_processor(images, return_tensors="pt")['pixel_values'] # # # 为每个样本构造3种条件的tokens # for i in range(3 * input_ids.size(0)): # tokens[i, :] = input_ids[i // 3, :] # if i % 3 == 2: # uncond版本，用pad_id替换中间tokens # tokens[i, 1:-1] = 100015 # pad_id # # # 添加数据到pre_data # pre_data.append(VLChatProcessorOutput( # sft_format=kwargs['sft_format'][i // 3], # pixel_values=encoder_pixel_values[i // 3, :], # input_ids=tokens[i - 2], # num_image_tokens=[vl_chat_processor.num_image_tokens] * 1 # )) # pre_data.append(VLChatProcessorOutput( # sft_format=kwargs['sft_format'][i // 3], # pixel_values=encoder_pixel_values[i // 3, :], # input_ids=tokens[i - 1], # num_image_tokens=[vl_chat_processor.num_image_tokens] * 1 # )) # pre_data.append(VLChatProcessorOutput( # sft_format=kwargs['sft_format'][i // 3], # pixel_values=None, # input_ids=tokens[i], # num_image_tokens=[] # )) # # # 批处理输入数据 # prepare_inputs = vl_chat_processor.batchify(pre_data) # # # 准备输入embeddings # inputs_embeds = self.prepare_inputs_embeds( # input_ids=tokens.cuda(), # pixel_values=prepare_inputs['pixel_values'].to(torch.bfloat16).cuda(), # images_emb_mask=prepare_inputs['images_emb_mask'].cuda(), # images_seq_mask=prepare_inputs['images_seq_mask'].cuda() # ) # # # 处理输入图像的编码 # input_image_pixel_values = vl_chat_processor.image_processor(images, return_tensors="pt")[ # 'pixel_values'].to(torch.bfloat16).cuda() # quant_input, emb_loss_input, info_input = self.gen_vision_model.encode(input_image_pixel_values) # image_tokens_input = info_input[2].detach().reshape(input_image_pixel_values.shape[0], -1) # image_embeds_input = self.prepare_gen_img_embeds(image_tokens_input) # # # 将输入图像embeddings插入到正确位置 # ppp = (tokens == 100580).nonzero() # 找到图像token位置 # for ii, ind in enumerate(ppp): # if ii % 4 == 0: # offset = ind[1] + 2 # inputs_embeds[ind[0], offset: offset + image_embeds_input.shape[1], :] = image_embeds_input[ # (ii // 2) % img_len] # # # **训练模式：只计算loss，不生成图像** # labels = None # if labels is not None: # outputs = self.language_model.model( # inputs_embeds=inputs_embeds, # use_cache=True, # past_key_values=None # ) # hidden_states = outputs.last_hidden_state # logits = self.gen_head(hidden_states) # # # 分离三种条件的logits # logit_cond_full = logits[0::3, :] # logit_cond_part = logits[1::3, :] # logit_uncond = logits[2::3, :] # # # 计算组合logits # logit_cond = (logit_cond_full + cfg_weight2 * logit_cond_part) / (1 + cfg_weight2) # all_logits = logit_uncond + cfg_weight * (logit_cond - logit_uncond) # # # 计算loss # loss_fct = CrossEntropyLoss() # shift_logits = all_logits[..., :-1, :].contiguous() # shift_logits = shift_logits.view(-1, self.config.gen_head_config.params.image_token_size) # shift_labels = labels[..., 1:].contiguous().view(-1).to(shift_logits.device) # loss = loss_fct(shift_logits, shift_labels) # # if not return_dict: # output = (all_logits,) + outputs[1:] # return ((loss,) + output) if loss is not None else output # # return CausalLMOutputWithPast( # loss=loss, # logits=all_logits, # past_key_values=outputs.past_key_values, # hidden_states=outputs.hidden_states, # attentions=outputs.attentions, # ) # # # **推理模式：自回归生成图像** # else: # import numpy as np # with torch.inference_mode(): # generated_tokens = torch.zeros((parallel_size, image_token_num_per_image), dtype=torch.int).cuda() # outputs = None # # # 自回归生成循环 # for i in range(image_token_num_per_image): # # 前向传播 # outputs = self.language_model.model( # inputs_embeds=inputs_embeds, # use_cache=True, # past_key_values=outputs.past_key_values if i != 0 else None # ) # hidden_states = outputs.last_hidden_state # # # 获取最后一个token的logits # logits = self.gen_head(hidden_states[:, -1, :]) # # # 分离三种条件的logits # logit_cond_full = logits[0::3, :] # logit_cond_part = logits[1::3, :] # logit_uncond = logits[2::3, :] # # # 计算组合logits # logit_cond = (logit_cond_full + cfg_weight2 * logit_cond_part) / (1 + cfg_weight2) # combined_logits = logit_uncond + cfg_weight * (logit_cond - logit_uncond) # # # 采样下一个token # probs = torch.softmax(combined_logits / temperature, dim=-1) # next_token = torch.multinomial(probs, num_samples=1) # generated_tokens[:, i] = next_token.squeeze(dim=-1) # # # 为下一步准备输入embeddings # if i < image_token_num_per_image - 1: # # 扩展next_token到3个副本 # next_token_expanded = torch.cat([ # next_token.unsqueeze(dim=1), # next_token.unsqueeze(dim=1), # next_token.unsqueeze(dim=1) # ], dim=1).view(-1) # # # 获取下一个token的embeddings # img_embeds = self.prepare_gen_img_embeds(next_token_expanded) # inputs_embeds = img_embeds.unsqueeze(dim=1) # # # 解码生成的tokens为图像 # dec = self.gen_vision_model.decode_code( # generated_tokens.to(dtype=torch.int), # shape=[parallel_size, 8, img_size // patch_size, img_size // patch_size] # ) # dec = dec.to(torch.float32).cpu().numpy().transpose(0, 2, 3, 1) # dec = np.clip((dec + 1) / 2 * 255, 0, 255) # # # 构造输出图像数组，确保形状正确 # visual_img = np.zeros((parallel_size, img_size, img_size, 3), dtype=np.uint8) # visual_img[:, :, :] = dec # # # 可选：保存调试图像（仅在推理模式下） # # if kwargs.get('save_debug_image', False): # # # 确保保存的是单张图像，形状为(384, 384, 3) # debug_img = visual_img[0] if visual_img.shape[0] > 0 else visual_img # PIL.Image.fromarray(debug_img).save('/home/ps/Bxh/align-anything/debug_output.png') image_token_num_per_image = 576 img_size = 384 patch_size = 16 cfg_weight = 5 temperature = 1 tokens = torch.zeros((3 * input_ids.size(0), input_ids.size(1)), dtype=torch.int).cuda() pre_data = [] img_len = len(kwargs['source_image']) # print(kwargs['source_image'].size(0)) print(kwargs['source_image']) print(len(kwargs['source_image'][0])) import PIL.Image images = [PIL.Image.open(image_path).convert("RGB") for image_path in kwargs['source_image']] # images = [PIL.Image.open(image_path).convert("RGB") for image_path in kwargs['source_image']] print('len_images : ',len(images)) encoder_pixel_values = vl_chat_processor.image_processor(images, return_tensors="pt")['pixel_values'] print(encoder_pixel_values.shape) print(encoder_pixel_values[0].shape) print(encoder_pixel_values[0][0][0][:2]) # print((encoder_pixel_values[0]!= encoder_pixel_values[1]).sum()) # print((encoder_pixel_values[0] != encoder_pixel_values[2]).sum()) # print((encoder_pixel_values[0] != encoder_pixel_values[3]).sum()) for i in range(3 * input_ids.size(0)): print(input_ids.shape) print(input_ids.size(0)) tokens[i * input_ids.size(0):(i + 1) * input_ids.size(0),:] = input_ids[i // 3,:] if i % 3 == 2: tokens[i * input_ids.size(0):(i + 1) * input_ids.size(0), 1:-1] = 100002 print(encoder_pixel_values[i//3,:].shape) print(len(kwargs['sft_format'][i//3])) print(tokens[i].shape) pre_data.append(VLChatProcessorOutput(sft_format=kwargs['sft_format'][i//3], pixel_values=encoder_pixel_values[i//3,:], input_ids=tokens[i - 2], num_image_tokens=[vl_chat_processor.num_image_tokens] * 1)) pre_data.append(VLChatProcessorOutput(sft_format=kwargs['sft_format'][i//3], pixel_values=encoder_pixel_values[i//3,:], input_ids=tokens[i - 1], num_image_tokens=[vl_chat_processor.num_image_tokens] * 1)) pre_data.append(VLChatProcessorOutput(sft_format=kwargs['sft_format'][i//3], pixel_values=None, input_ids=tokens[i], num_image_tokens=[])) # print(tokens.shape) # _, src_image = self.prepare_inputs_embeds(tokens[0], kwargs['source_image']) ppp = (tokens == 100580).nonzero() # print(tokens[0][583],tokens[0][584],tokens[0][576],tokens[0][577]) # print(input_ids.size(0)) # print(tokens[0][2], tokens[0][3]) # print(tokens[0][1161], tokens[0][1162]) # print(ppp) # print(src_image.shape) # img_len = src_image.shape[0] # # inputs_embeds_2 = self.language_model.get_input_embeddings()(tokens[1]) # # inputs_embeds_3 = self.language_model.get_input_embeddings()(tokens[2]) # inputs_embeds = self.language_model.get_input_embeddings()(tokens) # print(inputs_embeds.shape) prepare_inputs = vl_chat_processor.batchify(pre_data) # print('prepare_inputs pixel_values', prepare_inputs['pixel_values'].shape) # print('prepare_inputs images_emb_mask', prepare_inputs['images_emb_mask'].shape) # print('prepare_inputs images_seq_mask', prepare_inputs['images_seq_mask'].shape) inputs_embeds = self.prepare_inputs_embeds( input_ids=tokens.cuda(), pixel_values=prepare_inputs['pixel_values'].to(torch.bfloat16).cuda(), images_emb_mask=prepare_inputs['images_emb_mask'].cuda(), images_seq_mask=prepare_inputs['images_seq_mask'].cuda() ) input_image_pixel_values = vl_chat_processor.image_processor(images, return_tensors="pt")['pixel_values'].to(torch.bfloat16).cuda() quant_input, emb_loss_input, info_input = self.gen_vision_model.encode(input_image_pixel_values) image_tokens_input = info_input[2].detach().reshape(input_image_pixel_values.shape[0], -1) image_embeds_input = self.prepare_gen_img_embeds(image_tokens_input) # print('image_embeds_input', image_embeds_input.shape) # print('inputs_embeds', inputs_embeds.shape) for ii, ind in enumerate(ppp): # print('nmsl: ', ii, ind) if ii % 4 == 0: offset = ind[1] + 2 inputs_embeds[ind[0], offset: offset + image_embeds_input.shape[1], :] = image_embeds_input[ii // 4] generated_tokens = torch.zeros((3 * input_ids.size(0), image_token_num_per_image), dtype=torch.int).cuda() outputs = self.language_model.model(inputs_embeds=inputs_embeds, use_cache=True, past_key_values=None, labels=labels) hidden_states = outputs.last_hidden_state print('HHHH',hidden_states.shape) # torch.save(inputs_embeds, '/data/bxh_data/unify_model/hidden_states.pt') logits = self.gen_head(hidden_states) print('logits.shape', logits.shape) # [3, 1760, 16384]) print(labels.shape) # [3, 1760] # logits_cond = logits[0::2, :] # logits_uncond = logits[1::2, :] logit_cond_full = logits[0::3, :] logit_cond_part = logits[1::3, :] logit_uncond = logits[2::3, :] cfg_weight2 = 5 logit_cond = (logit_cond_full + cfg_weight2 * (logit_cond_part)) / (1 + cfg_weight2) all_logits = logit_uncond + cfg_weight * (logit_cond - logit_uncond) # all_logits = logits_uncond + cfg_weight * (logits_cond - logits_uncond) loss_fct = CrossEntropyLoss() shift_logits = all_logits[..., :-1, :].contiguous() shift_logits = shift_logits.view(-1, self.config.gen_head_config.params.image_token_size) if labels is not None: shift_labels = labels[..., 1:].contiguous() shift_labels = shift_labels.view(-1) shift_labels = shift_labels.to(shift_logits.device) print(shift_logits.shape, shift_labels.shape) loss = loss_fct(shift_logits, shift_labels) else: loss = None if not return_dict: output = (logits,) + outputs[1:] return ((loss,) + output) if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) AutoConfig.register("vision", VisionConfig) AutoConfig.register("aligner", AlignerConfig) AutoConfig.register("gen_vision", GenVisionConfig) AutoConfig.register("gen_aligner", GenAlignerConfig) AutoConfig.register("gen_head", GenHeadConfig) AutoConfig.register("multi_modality", MultiModalityConfig) AutoModelForCausalLM.register(MultiModalityConfig, MultiModalityCausalLM)