#!/usr/bin/env python3 # coding=utf-8 # Copyright (c) Ant Group. All rights reserved. from typing import List, Optional, Union import torch import torch.nn as nn import torch.nn.functional as F from PIL import Image from diffusers.models.normalization import RMSNorm from transformers import PreTrainedModel from transformers.utils import logging from configuration_bailingmm2 import BailingMM2Config from modeling_bailing_moe_v2 import BailingMoeV2ForCausalLM from bailingmm_utils import process_ratio, find_first_index_of_consecutive_ones, merge_consecutive_ones from inference_profile import load_checkpoint_capabilities, resolve_model_directory import os from copy import deepcopy # vision encoder from qwen2_5_vit import Qwen2_5_VisionTransformer logger = logging.get_logger(__name__) _CONFIG_FOR_DOC = "BailingMM2Config" class BailingMM2NativeForConditionalGeneration(PreTrainedModel): config_class = BailingMM2Config base_model_prefix = "model" supports_gradient_checkpointing = True _skip_keys_device_placement = "past_key_values" _supports_flash_attn_2 = True def __init__( self, config: BailingMM2Config, empty_load=False, ): super().__init__(config) self.config: BailingMM2Config = config self.vision = None self.llm_dytpe = torch.bfloat16 if empty_load: self.model = None return if self.config.vision_config: self.vision = Qwen2_5_VisionTransformer(self.config.vision_config) self.model = BailingMoeV2ForCausalLM(self.config.llm_config) mlp_modules_img = [nn.Linear(self.vision.image_emb_dim, self.model.config.hidden_size)] for _ in range(1, self.config.mlp_depth): mlp_modules_img.append(nn.GELU()) mlp_modules_img.append(nn.Linear(self.model.config.hidden_size, self.model.config.hidden_size)) self.linear_proj = nn.Sequential(*mlp_modules_img) self.post_init() def extract_image_feature(self, pixel_values, grid_thw): with torch.cuda.amp.autocast(dtype=torch.bfloat16): image_embeds = self.vision(pixel_values, grid_thw=grid_thw) image_embeds = self.linear_proj(image_embeds) image_embeds = F.normalize(image_embeds, dim=-1) return image_embeds @torch.no_grad() def generate( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, pixel_values: Optional[torch.FloatTensor] = None, pixel_values_videos: Optional[torch.FloatTensor] = None, audio_feats: Optional[torch.FloatTensor] = None, audio_feats_lengths: Optional[torch.LongTensor] = None, audio_placeholder_loc_lens: Optional[torch.LongTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.Tensor]] = None, num_logits_to_keep: Optional[int] = 0, image_gen: Optional[bool] = False, image_gen_pixel_values_reference: Optional[torch.FloatTensor] = None, image_gen_negative_input_ids: Optional[torch.LongTensor] = None, image_gen_negative_attention_mask: Optional[torch.Tensor] = None, image_gen_steps: Optional[int] = None, image_gen_seed: Optional[int] = None, image_gen_cfg: Optional[float] = None, image_gen_image_cfg: Optional[float] = 1.0, image_gen_cfg_mode: Optional[int] = 1, image_gen_height: Optional[int] = None, image_gen_width: Optional[int] = None, image_gen_llm_hidden_states: Optional[torch.LongTensor] = None, image_gen_negative_llm_hidden_states: Optional[torch.LongTensor] = None, image_gen_text: Optional[list] = None, image_gen_highres = 512, image_gen_only_extract_hidden_states = False, image_gen_condition_embeds=None, image_gen_negative_condition_embeds=None, image_gen_condition_embeds_2=None, image_gen_negative_condition_embeds_2=None, image_gen_return_batch=False, image_gen_task=None, num_frames_per_prompt=1, **generate_kwargs, ): if audio_feats is not None or audio_feats_lengths is not None: raise ValueError("audio input is not supported by Ming Image inference") if image_gen_image_cfg not in (None, 1.0): raise ValueError( "image_gen_image_cfg is not supported by this inference path; " "guidance is controlled by image_gen_cfg only" ) image_embeds, video_embeds, audio_embeds, audio_embeds_lengths = None, None, None, None if image_gen: if not hasattr(self, "inference_profile"): raise RuntimeError( "image modules were loaded without a checkpoint inference profile" ) self.inference_profile.validate_task( image_gen_task, has_reference_image=image_gen_pixel_values_reference is not None, num_layers=num_frames_per_prompt, ) sampling = self.inference_profile.resolve_sampling_parameters( steps=image_gen_steps, cfg=image_gen_cfg, ) image_gen_steps = sampling.steps image_gen_cfg = sampling.cfg if image_gen_pixel_values_reference is not None: input_channels = image_gen_pixel_values_reference.shape[1] expected_channels = self.inference_profile.vae_input_channels if input_channels % expected_channels != 0: raise ValueError( "reference image channels do not match the checkpoint " f"VAE contract: input={input_channels}, expected a " f"multiple of {expected_channels}" ) condition_embeds, negative_condition_embeds = None, None condition_embeds_2, negative_condition_embeds_2 = None, None if (image_gen_condition_embeds is not None) or (image_gen_condition_embeds_2 is not None): if image_gen_condition_embeds is not None: condition_embeds = image_gen_condition_embeds negative_condition_embeds = condition_embeds * 0.0 if image_gen_negative_condition_embeds is None else image_gen_negative_condition_embeds if image_gen_condition_embeds_2 is not None: condition_embeds_2 = image_gen_condition_embeds_2 negative_condition_embeds_2 = condition_embeds_2 * 0.0 if image_gen_negative_condition_embeds_2 is None else image_gen_negative_condition_embeds_2 else: if image_gen_llm_hidden_states is None: assert self.model is not None assert self.vision is not None if pixel_values is not None: image_embeds = self.extract_image_feature(pixel_values, grid_thw=image_grid_thw) assert self.loaded_image_gen_modules is True, "please add `load_image_gen=True` in from_pretrained() method" assert position_ids is None condition_embeds, condition_embeds_2 = self.get_condition_embeds_for_image_gen( input_ids=input_ids, attention_mask=attention_mask, image_embeds=image_embeds, position_ids=position_ids, use_cache=use_cache, image_grid_thw=image_grid_thw, llm_hidden_states=image_gen_llm_hidden_states, ) if condition_embeds is not None: negative_condition_embeds = condition_embeds * 0.0 if condition_embeds_2 is not None: negative_condition_embeds_2 = condition_embeds_2 * 0.0 # negative prompt feature is deprecated # negative_condition_embeds = self.get_learnable_token_embeds_for_image_gen( # input_ids=image_gen_negative_input_ids, # attention_mask=image_gen_negative_attention_mask, # image_embeds=image_embeds, # position_ids=position_ids, # use_cache=use_cache, # image_grid_thw=image_grid_thw, # llm_hidden_states=image_gen_negative_llm_hidden_states, # ) if ((image_gen_negative_input_ids is not None) or (image_gen_negative_llm_hidden_states is not None)) else condition_embeds * 0.0 if image_gen_only_extract_hidden_states: return condition_embeds, negative_condition_embeds, condition_embeds_2, negative_condition_embeds_2 assert (condition_embeds is not None) or (condition_embeds_2 is not None) if (condition_embeds is not None) and (condition_embeds_2 is not None): assert condition_embeds.shape[0] == condition_embeds_2.shape[0] bsz = condition_embeds.shape[0] if condition_embeds is not None else condition_embeds_2.shape[0] if image_gen_height is None or image_gen_width is None: if isinstance(image_gen_highres, int): image_gen_height, image_gen_width = [image_gen_highres] * bsz, [image_gen_highres] * bsz elif image_gen_highres is True: image_gen_height, image_gen_width = [1024] * bsz, [1024] * bsz else: image_gen_height, image_gen_width = [512] * bsz, [512] * bsz elif isinstance(image_gen_height, torch.Tensor) or isinstance(image_gen_width, torch.Tensor): assert isinstance(image_gen_height, torch.Tensor), image_gen_height assert isinstance(image_gen_width, torch.Tensor), image_gen_width image_gen_height = image_gen_height.cpu().tolist() image_gen_width = image_gen_width.cpu().tolist() assert len(image_gen_height) == bsz assert len(image_gen_width) == bsz elif isinstance(image_gen_height, int) or isinstance(image_gen_width, int): assert isinstance(image_gen_height, int), image_gen_height assert isinstance(image_gen_width, int), image_gen_width image_gen_height = [image_gen_height] * bsz image_gen_width = [image_gen_width] * bsz else: assert isinstance(image_gen_height, list), image_gen_height assert isinstance(image_gen_width, list), image_gen_width assert len(image_gen_height) == bsz assert len(image_gen_width) == bsz image_gen_height_diffusion_list = [] image_gen_width_diffusion_list = [] image_gen_output_resize_height = [] image_gen_output_resize_width = [] for height, width in zip(image_gen_height, image_gen_width): closest_size, resize_size = process_ratio(ori_h=height, ori_w=width, highres=image_gen_highres) height, width = closest_size image_gen_height_diffusion_list.append(height) image_gen_width_diffusion_list.append(width) height, width = resize_size image_gen_output_resize_height.append(height) image_gen_output_resize_width.append(width) image_gen_height = image_gen_height_diffusion_list[0] assert all([i == image_gen_height for i in image_gen_height_diffusion_list]) image_gen_width = image_gen_width_diffusion_list[0] assert all([i == image_gen_width for i in image_gen_width_diffusion_list]) if image_gen_pixel_values_reference is not None: assert (image_gen_height, image_gen_width) == (image_gen_pixel_values_reference.shape[-2], image_gen_pixel_values_reference.shape[-1]) if image_gen_seed is None or image_gen_seed < 0: from datetime import datetime image_gen_seed = datetime.now().microsecond % 1000 sample_kwargs = { "steps": image_gen_steps, "seed": image_gen_seed, "cfg": image_gen_cfg, "height": image_gen_height, "width": image_gen_width, "cfg_mode": image_gen_cfg_mode, "ref_x": image_gen_pixel_values_reference, "encoder_hidden_states": condition_embeds, "directvlm_hidden_states": condition_embeds_2, "num_frames_per_prompt": num_frames_per_prompt, } image = self.diffusion_loss.sample( **sample_kwargs, ) if image_gen_task == "layer-decompose": output_size = ( image_gen_output_resize_width[0], image_gen_output_resize_height[0], ) image = [item.resize(output_size, Image.LANCZOS) for item in image] else: image = [ item.resize((width, height), Image.LANCZOS) for item, width, height in zip( image, image_gen_output_resize_width, image_gen_output_resize_height, ) ] if ( image_gen_task != "layer-decompose" and not image_gen_return_batch and len(image) == 1 ): image = image[0] return image if pixel_values is not None: image_embeds = self.extract_image_feature(pixel_values, grid_thw=image_grid_thw) if pixel_values_videos is not None: video_embeds = self.extract_image_feature(pixel_values_videos, grid_thw=video_grid_thw) with torch.cuda.amp.autocast(dtype=torch.bfloat16): outputs = self.model.generate( input_ids=input_ids, query_embeds_image=image_embeds, query_embeds_video=video_embeds, query_embeds_audio=audio_embeds, query_embeds_audio_lengths=audio_embeds_lengths, placeholder_audio_loc_lens=audio_placeholder_loc_lens, image_grid_thw=image_grid_thw, image_grid_thw_video=video_grid_thw, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, num_logits_to_keep=num_logits_to_keep, **generate_kwargs, ) return outputs def load_image_gen_modules(self, inference_model_path, torch_dtype=torch.float32, load_image_gen_diffusion=True, load_image_gen_others=True, device=None): inference_model_path = str(resolve_model_directory(inference_model_path)) if os.path.exists(os.path.join(inference_model_path, "byt5")): raise ValueError( "Ming Image inference does not support a byt5 component; " "the public checkpoint layout has no byt5/ directory." ) self.inference_profile = load_checkpoint_capabilities(inference_model_path) if device is not None: device = torch.device(device) elif self.model is not None: device = self.model.device else: device = torch.device(torch.cuda.current_device()) logger.info(f"load_image_gen_modules device={device}") from transformers import AutoModelForCausalLM from safetensors.torch import load_file temp_state_dict = load_file( os.path.join(inference_model_path, "mlp", "model.safetensors") ) with open(os.path.join(inference_model_path, 'mlp', 'config.json'), 'r') as f: import json metax_config = json.load(f) diffusion_c_input_dim = metax_config.get("diffusion_c_input_dim", 2048) self.img_gen_scales = metax_config.get("img_gen_scales", [4, 8, 16]) self.connector_norm = metax_config.get("connector_norm", True) self.use_vlm_directvlm_condition = metax_config.get( "use_vlm_directvlm_condition", False ) self.use_learnable_token_condition = metax_config.get( "use_learnable_token_condition", True ) self.selected_hidden_states_layers = metax_config.get( "selected_hidden_states_layers" ) self.diffusion_inner_dim = metax_config.get("diffusion_inner_dim") if load_image_gen_others: self.connector = None self.query_tokens_dict = nn.ParameterDict() # cumulative token index across the scales self.scale_indices = [] current_idx = 0 for scale in self.img_gen_scales: num_tokens = scale * scale scale_name = f"{scale}x{scale}" #weights = temp_state_dict[f"query_tokens_dict.{scale_name}"] self.query_tokens_dict[scale_name] = nn.Parameter( torch.nn.functional.normalize(torch.randn(num_tokens, self.config.llm_config.hidden_size), dim=-1) ) current_idx += scale * scale self.scale_indices.append(current_idx) self.query_tokens_dict.to(torch_dtype).to(device) if self.use_learnable_token_condition: modified_state_dict_query_tokens = { f"{scale}x{scale}": temp_state_dict[f"query_tokens_dict.{scale}x{scale}"] for scale in self.img_gen_scales } self.query_tokens_dict.load_state_dict(modified_state_dict_query_tokens, strict=True) # self.norm_query_embeds = True # load connector self.connector = AutoModelForCausalLM.from_pretrained(inference_model_path, subfolder='connector', torch_dtype=torch_dtype) for layer in self.connector.model.layers: layer.self_attn.is_causal = False self.connector.to(device) self.proj_in = nn.Linear(self.config.llm_config.hidden_size, self.connector.config.hidden_size) self.proj_out = nn.Linear(self.connector.config.hidden_size, diffusion_c_input_dim) modified_state_dict_in = { 'weight': temp_state_dict['proj_in.weight'], 'bias': temp_state_dict['proj_in.bias'] } self.proj_in.load_state_dict(modified_state_dict_in, strict=True) modified_state_dict_out = { 'weight': temp_state_dict['proj_out.weight'], 'bias': temp_state_dict['proj_out.bias'] } self.proj_out.load_state_dict(modified_state_dict_out, strict=True) self.proj_in.to(device=device, dtype=torch_dtype) self.proj_out.to(device=device, dtype=torch_dtype) self.proj_directvlm = None if self.use_vlm_directvlm_condition: directvlm_dim = self.model.config.hidden_size if self.selected_hidden_states_layers is not None: directvlm_dim = directvlm_dim * len(self.selected_hidden_states_layers) self.proj_directvlm = nn.Sequential(RMSNorm(directvlm_dim, eps=1e-5), nn.Linear(directvlm_dim, self.diffusion_inner_dim, bias=True)) modified_state_dict_directvlm = { '0.weight': temp_state_dict["proj_directvlm.0.weight"], '1.weight': temp_state_dict["proj_directvlm.1.weight"], '1.bias': temp_state_dict["proj_directvlm.1.bias"], } self.proj_directvlm.load_state_dict(modified_state_dict_directvlm, strict=True) self.proj_directvlm.to(device=device, dtype=torch_dtype) if load_image_gen_diffusion: diffusion_mlp_state_dict = { key[len("mlp.") :] : temp_state_dict[key] for key in temp_state_dict if key.startswith("mlp.") } from diffusion.generator import ImageGenerator self.diffusion_loss = ImageGenerator( model_path=inference_model_path, scheduler_path=inference_model_path, vision_dim=diffusion_c_input_dim, mlp_state_dict=diffusion_mlp_state_dict, torch_dtype=torch_dtype, device=device, use_identity_mlp=metax_config.get("use_identity_mlp", False), text_encoder_norm=metax_config.get("text_encoder_norm", False), inference_profile=self.inference_profile, ) self.diffusion_loss.to(device) self.loaded_image_gen_modules = True @classmethod def _from_int8_checkpoint(cls, vlm_directory, device, **kwargs): """Load an mllm/ component written by quant/quantize_stream.py (weight-only int8). The model is built with its parameters on the meta device, the Linear modules listed in int8_manifest.json become Int8Linear shells, and every stored tensor is loaded straight onto `device`, so BF16 weights for the quantized modules never exist in memory. """ from accelerate import init_empty_weights from quant.load_int8 import load_int8_mllm_ device_map = kwargs.pop("device_map", None) if device_map is not None: # infer.py's default "balanced" plan on a single-GPU box maps every module to GPU 0; # that is honoured. Splitting the int8 model across devices is not supported. targets = set(device_map.values()) if isinstance(device_map, dict) else {device_map} if len(targets) != 1 or not isinstance(next(iter(targets)), int): raise ValueError( "the int8 mllm checkpoint loads onto a single GPU; device_map targets " f"{sorted(map(str, targets))} (use --device-map none)" ) device = torch.device("cuda", next(iter(targets))) supported = ("torch_dtype", "dtype", "attn_implementation") unsupported = sorted(key for key in kwargs if key not in supported) if unsupported: raise ValueError( f"the int8 mllm checkpoint loads onto a single device; unsupported arguments: {unsupported}" ) device = torch.device(device) if device is not None else torch.device("cpu") if device.type == "cuda" and device.index is None: device = torch.device("cuda", torch.cuda.current_device()) config = BailingMM2Config.from_pretrained(vlm_directory) with init_empty_weights(): model = cls._from_config(config, **kwargs) report = load_int8_mllm_(model, vlm_directory, device) # Buffers built in __init__ (rotary inv_freq) are not stored in the checkpoint; they follow # the weights. Int8Linear keeps its fp32 scales through this and any later dtype cast. model.to(device) logger.info(f"int8 mllm loaded from {vlm_directory}: {report}") model.tie_weights() model.eval() return model @classmethod def from_pretrained( cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs, ): load_image_gen = False if "load_image_gen" in kwargs: load_image_gen = kwargs["load_image_gen"] del kwargs["load_image_gen"] load_image_gen_diffusion = True if "load_image_gen_diffusion" in kwargs: load_image_gen_diffusion = kwargs["load_image_gen_diffusion"] del kwargs["load_image_gen_diffusion"] load_image_gen_others = True if "load_image_gen_others" in kwargs: load_image_gen_others = kwargs["load_image_gen_others"] del kwargs["load_image_gen_others"] load_vlm = True if "load_vlm" in kwargs: load_vlm = kwargs["load_vlm"] del kwargs["load_vlm"] image_gen_device = kwargs.pop("image_gen_device", None) vlm_directory = pretrained_model_name_or_path if load_image_gen: pretrained_model_name_or_path = str( resolve_model_directory( pretrained_model_name_or_path, revision=kwargs.get("revision"), cache_dir=kwargs.get("cache_dir"), local_files_only=kwargs.get("local_files_only", False), token=kwargs.get("token"), ) ) # The package root keeps connector/mlp/transformer/vae/scheduler; # the MLLM itself (config, weights, tokenizer data) lives in mllm/. vlm_directory = os.path.join(pretrained_model_name_or_path, "mllm") if not os.path.isdir(vlm_directory): raise FileNotFoundError( "checkpoint package is missing the mllm/ component: " f"{vlm_directory}. Migrate the package to the component " "layout before loading." ) if load_vlm and os.path.exists(os.path.join(vlm_directory, "int8_manifest.json")): model = cls._from_int8_checkpoint(vlm_directory, image_gen_device, **kwargs) elif load_vlm: model = super().from_pretrained( vlm_directory, *model_args, **kwargs, ) else: model = cls( BailingMM2Config.from_dict(BailingMM2Config.get_config_dict(vlm_directory)[0]), empty_load=True, ) if load_image_gen: model.load_image_gen_modules( pretrained_model_name_or_path, torch_dtype=kwargs["torch_dtype"] if "torch_dtype" in kwargs else torch.float32, load_image_gen_diffusion=load_image_gen_diffusion, load_image_gen_others=load_image_gen_others, device=image_gen_device, ) return model def append_input_ids_with_multiscale_learnable_tokens( self, text_ids, attention_mask, scales, start_token_id, end_token_id, patch_token_id, ): default_scaled_tokens = [] default_scaled_attn_masks = [] default_gen_masks = [] for scale in scales: default_scaled_tokens.append(start_token_id) default_scaled_tokens.extend([patch_token_id for _ in range(scale * scale)]) default_scaled_tokens.append(end_token_id) default_scaled_attn_masks.extend([1 for _ in range(scale * scale + 2)]) default_gen_masks.append(0) default_gen_masks.extend([1 for _ in range(scale * scale)]) default_gen_masks.append(0) text_ids_list = text_ids.cpu().tolist() attention_mask_list = attention_mask.cpu().tolist() new_text_ids_list = [] new_attention_mask_list = [] gen_mask_list = [] new_labels_list = [] for text_ids_one_batch, attention_mask_one_batch in zip( text_ids_list, attention_mask_list ): assert len(text_ids_one_batch) == len(attention_mask_one_batch) padding_start = 0 for idx, value in enumerate(attention_mask_one_batch): if value == 0: break padding_start += 1 new_text_ids_list.append(text_ids_one_batch[:padding_start] + deepcopy(default_scaled_tokens) + text_ids_one_batch[padding_start:]) new_labels_list.append([ -100 for _ in range(padding_start)] + [1 for _ in range(len(default_scaled_tokens))] + [-100 for _ in range(len(text_ids_one_batch[padding_start:]))] ) new_attention_mask_list.append(attention_mask_one_batch[:padding_start] + deepcopy(default_scaled_attn_masks) + attention_mask_one_batch[padding_start:]) gen_mask_list.append( [0 for _ in range(len(attention_mask_one_batch[:padding_start]))] + \ deepcopy(default_gen_masks) + \ [0 for _ in range(len(attention_mask_one_batch[padding_start:]))] ) text_ids_append_lq = torch.tensor(new_text_ids_list, dtype=text_ids.dtype).to(text_ids.device) attention_mask_append_lq = torch.tensor(new_attention_mask_list, dtype=attention_mask.dtype).to(attention_mask.device) gen_mask = torch.tensor(gen_mask_list, dtype=attention_mask.dtype).to(attention_mask.device) labels = torch.tensor(new_labels_list, dtype=text_ids.dtype).to(text_ids.device) assert attention_mask_append_lq.shape == text_ids_append_lq.shape assert labels.shape == text_ids_append_lq.shape assert gen_mask.shape == text_ids_append_lq.shape return text_ids_append_lq, labels, attention_mask_append_lq, gen_mask def appand_learnable_tokens( self, text_ids, gen_mask, image_embeds, image_grid_thw, patch_token_id, ): query_tokens_embeds = torch.cat( [self.query_tokens_dict[f"{scale}x{scale}"] for scale in self.img_gen_scales], dim=0, ) if image_embeds is not None: query_tokens_embeds = query_tokens_embeds.to(image_embeds.dtype).to(image_embeds.device) assert text_ids.shape == gen_mask.shape text_ids_aslist = text_ids.cpu().view(-1).tolist() gen_mask_aslist = gen_mask.cpu().view(-1).tolist() is_patch_list = [1 if i == patch_token_id else 0 for i in text_ids_aslist] idxes_start_of_patch = find_first_index_of_consecutive_ones(is_patch_list) isgen_indicators = merge_consecutive_ones([1 if gen_mask_aslist[i] else 0 for i in idxes_start_of_patch], len(self.img_gen_scales)) if any([i == 0 for i in isgen_indicators]): assert image_grid_thw is not None assert image_grid_thw.ndim == 2 assert image_embeds is not None assert image_embeds.ndim == 2 new_image_grid_thw = [] new_image_embeds = [] cum_image_token = 0 cnt_input_image = 0 for is_gen in isgen_indicators: if is_gen: for scale in self.img_gen_scales: new_image_grid_thw.append([1, 2, scale * scale * 2]) new_image_embeds.append(query_tokens_embeds) else: thw = image_grid_thw[cnt_input_image].tolist() assert thw[0] == 1 assert thw[1] % 2 == 0 # h assert thw[2] % 2 == 0 # w n_image_token = (thw[1] // 2) * (thw[2] // 2) image_embed_one = image_embeds[cum_image_token : cum_image_token + n_image_token, :] new_image_embeds.append(image_embed_one) new_image_grid_thw.append(thw) cnt_input_image += 1 cum_image_token += n_image_token if image_grid_thw is not None: assert cnt_input_image == image_grid_thw.shape[0] assert cum_image_token == image_embeds.shape[0] else: assert cnt_input_image == 0 assert cum_image_token == 0 new_image_grid_thw = torch.tensor(new_image_grid_thw, dtype=text_ids.dtype).to(text_ids.device) new_image_embeds = torch.cat(new_image_embeds, dim=0).to(text_ids.device) total_patch_token = 0 for bid in range(new_image_grid_thw.shape[0]): thw = new_image_grid_thw[bid].tolist() assert thw[0] == 1 assert thw[1] % 2 == 0 assert thw[2] % 2 == 0 patch_h = thw[1] // 2 patch_w = thw[2] // 2 n_patch_token = patch_h * patch_w total_patch_token += n_patch_token # if torch.distributed.get_rank() == 0: # embed() # torch.distributed.barrier() assert total_patch_token == new_image_embeds.shape[0], f"{total_patch_token}, vs. {new_image_embeds.shape}" return new_image_grid_thw, new_image_embeds def get_condition_embeds_for_image_gen( self, input_ids, attention_mask, image_embeds, position_ids, use_cache, image_grid_thw, llm_hidden_states, ): input_ids, labels, attention_mask, gen_mask = self.append_input_ids_with_multiscale_learnable_tokens( input_ids, attention_mask, self.img_gen_scales, self.config.llm_config.image_patch_token + 1, self.config.llm_config.image_patch_token + 2, self.config.llm_config.image_patch_token, ) if llm_hidden_states is None: image_grid_thw, image_embeds = self.appand_learnable_tokens( input_ids, gen_mask, image_embeds, image_grid_thw, self.config.llm_config.image_patch_token, ) with torch.cuda.amp.autocast(dtype=torch.bfloat16): if image_embeds is None or input_ids.size(1) == 1: words_embeddings = self.model.get_input_embeddings()(input_ids.clip(0, self.model.get_input_embeddings().weight.shape[0] - 1)) image_mask = None audio_mask = None else: words_embeddings, image_mask, audio_mask = self.model.model.prompt_wrap_navit( input_ids=input_ids.clip(0, self.model.get_input_embeddings().weight.shape[0] - 1), config=self.model.model.config, query_embeds_image=image_embeds, ) assert input_ids.size(1) == words_embeddings.size(1), "{} vs {}".format( input_ids.size, words_embeddings.size, ) # if torch.distributed.get_rank() == 3: # embed() # torch.distributed.barrier() outputs = self.model.forward( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=None, inputs_embeds=words_embeddings, image_grid_thw=image_grid_thw, use_cache=False, image_mask=image_mask, audio_mask=None, output_hidden_states=True, ) hidden_states = outputs.hidden_states[-1] else: hidden_states = llm_hidden_states directvlm_hidden_states = None if self.use_vlm_directvlm_condition: # use hidden states use_input_mask = torch.lt(labels, 0).int().to(attention_mask.dtype) * attention_mask assert use_input_mask.ndim == 2 directvlm_max_valid_ind = use_input_mask.cumsum(-1).argmax(-1).max().item() + 1 #directvlm_max_valid_ind = min(directvlm_max_valid_ind, self.max_vlm_directvlm_length) use_input_mask = use_input_mask[:, :directvlm_max_valid_ind] if self.selected_hidden_states_layers is not None: directvlm_hidden_states = torch.cat([ outputs.hidden_states[layer_i].to(labels.device)[:, :directvlm_max_valid_ind, :] * use_input_mask.unsqueeze(-1) for layer_i in self.selected_hidden_states_layers ], dim=-1) else: directvlm_hidden_states = outputs.hidden_states[-1].to(labels.device)[:, :directvlm_max_valid_ind, :] * use_input_mask.unsqueeze(-1) directvlm_hidden_states = directvlm_hidden_states.detach() directvlm_hidden_states = self.proj_directvlm(directvlm_hidden_states) scale_embeds = None if self.use_learnable_token_condition: with torch.cuda.amp.autocast(dtype=torch.bfloat16): gen_mask = gen_mask.unsqueeze(-1).expand(gen_mask.shape[0], gen_mask.shape[1], hidden_states.shape[-1]).to(hidden_states.device).bool() hidden_states_gen = torch.masked_select(hidden_states, gen_mask).view(hidden_states.shape[0], -1, hidden_states.shape[-1]) # split hidden_states into per-scale representations scale_start_idxes = [0] + self.scale_indices[:-1] scale_end_idxes = self.scale_indices assert scale_end_idxes[-1] == hidden_states_gen.shape[1] scale, scale_start_idx, scale_end_idx = [ i for i in zip(self.img_gen_scales, scale_start_idxes, scale_end_idxes) ][-1] scale_hidden = hidden_states_gen[:, scale_start_idx : scale_end_idx, :] scale_embeds = self.proj_in(scale_hidden) seq_shape = scale_embeds.shape with torch.cuda.amp.autocast(dtype=torch.bfloat16): scale_embeds = self.connector( inputs_embeds=scale_embeds, attention_mask=torch.ones(seq_shape[0],1,seq_shape[1],seq_shape[1]).to(scale_embeds.device), output_hidden_states=True ).hidden_states[-1] scale_embeds = self.proj_out(scale_embeds) # normalize if self.connector_norm: scale_embeds = torch.nn.functional.normalize(scale_embeds, dim=-1) return scale_embeds, directvlm_hidden_states __all__ = [ "BailingMM2NativeForConditionalGeneration" ]