# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # coding: utf-8 import random from typing import List, Tuple, Optional, Dict from einops import rearrange import torch import torch.nn.functional as F from torch import nn from torch.nn.attention.flex_attention import create_block_mask from transformers.configuration_utils import PretrainedConfig from transformers.modeling_utils import PreTrainedModel from data.data_utils import ( create_sparse_mask, get_flattened_position_ids_extrapolate, get_flattened_position_ids_interpolate, get_flattened_position_ids_interpolate_video, get_flattened_position_ids_extrapolate_video, ) from .qwen2_navit import NaiveCache, Qwen2ForCausalLM from .modeling_utils import MLPconnector, TimestepEmbedder, PositionEmbedding3D from config.config_factory import TrainingArguments from common.utils.misc import AutoEncoderParams from common.utils.distributed import get_global_rank from common.utils.logging import get_logger from modeling.vit.qwen2_5_vl_vit import Qwen2_5_VisionTransformerPretrainedModel from modeling.qwen2 import Qwen2Tokenizer from common.val.utils import map_splits_to_samples, make_packed_vit_token_embed, uncond_split_pro from data.common import shift_position_ids from copy import deepcopy class LanceConfig(PretrainedConfig): def __init__( self, visual_gen=True, visual_und=True, llm_config=None, vit_config=None, vae_config: AutoEncoderParams = None, latent_patch_size=(1, 2, 2), # pt ph pw max_latent_size=32, vit_max_num_patch_per_side=70, connector_act="gelu_pytorch_tanh", interpolate_pos=False, timestep_shift=1.0, **kwargs, ): super().__init__(**kwargs) self.visual_gen = visual_gen self.visual_und = visual_und self.llm_config = llm_config self.vit_config = vit_config self.vae_config = vae_config self.latent_patch_size = latent_patch_size self.max_num_frames = kwargs.get("max_num_frames", 25) self.max_latent_size = max_latent_size self.vit_max_num_patch_per_side = vit_max_num_patch_per_side self.connector_act = connector_act self.interpolate_pos = interpolate_pos self.timestep_shift = timestep_shift class Lance(PreTrainedModel): config_class = LanceConfig base_model_prefix = "lance" def __init__( self, language_model: Qwen2ForCausalLM, vit_model: Qwen2_5_VisionTransformerPretrainedModel, vit_type: str = "qwen2_5_vl", config: LanceConfig = None, **kwargs ): super().__init__(config) self.language_model: Qwen2ForCausalLM = language_model self.hidden_size = config.llm_config.hidden_size self.use_moe = "Mo" in config.llm_config.layer_module self.num_heads = config.llm_config.num_attention_heads self.logger = get_logger() self.log_rank0 = self.logger.info if get_global_rank() == 0 else lambda x: None if config.visual_gen: self.latent_patch_size = config.latent_patch_size self.timestep_shift = config.timestep_shift self.latent_downsample_spatial = config.vae_config.downsample_spatial * config.latent_patch_size[-1] self.latent_downsample_temporal = config.vae_config.downsample_temporal self.max_num_latent_frames = config.max_num_frames // self.latent_downsample_temporal + 1 self.latent_channel = config.vae_config.z_channels self.max_latent_size = config.max_latent_size self.patch_latent_dim = self.latent_patch_size[0] * self.latent_patch_size[1] * self.latent_patch_size[2] * self.latent_channel self.time_embedder = TimestepEmbedder(self.hidden_size) self.vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size) # vision input self.llm2vae = nn.Linear(self.hidden_size, self.patch_latent_dim) # vision ouput self.latent_pos_embed = PositionEmbedding3D(self.max_num_latent_frames, self.max_latent_size, self.hidden_size) safety = 1024 self.pos_shift = self.max_latent_size * self.max_latent_size * self.max_num_latent_frames + safety if config.visual_und: self.vit_model: Qwen2_5_VisionTransformerPretrainedModel = vit_model self.vit_patch_size = config.vit_config.patch_size self.vit_max_num_patch_per_side = config.vit_max_num_patch_per_side self.vit_type = vit_type if self.vit_type == "qwen2_5_vl": self.vit_hidden_size: int = config.vit_config.out_hidden_size self.connector: MLPconnector = MLPconnector(self.vit_hidden_size, self.hidden_size, config.connector_act) elif self.vit_type == "qwen_2_5_vl_original": pass else: raise ValueError(f"vit_model_type {self.vit_type} not supported") self.vit_model.eval() if config.interpolate_pos: self.get_flattened_position_ids = get_flattened_position_ids_interpolate else: self.get_flattened_position_ids = get_flattened_position_ids_extrapolate self.config = config self.training_args: TrainingArguments = kwargs.get("training_args") def update_tokenizer(self, tokenizer): self.tokenizer: Qwen2Tokenizer = tokenizer self.vocab_size_efficient = len(tokenizer) def process_attention_mask(self, current_attn_modes, current_split_lens, current_seq_len, device, BLOCK_SIZE=128): current_attn_modes_ = ["full" if mode_ in ["full_noise", "full_noise_target"] else mode_ for mode_ in current_attn_modes] sparse_mask = create_sparse_mask(current_seq_len, current_split_lens, current_attn_modes_, device) current_seq_len_sum = sum(current_seq_len) attention_mask = create_block_mask( sparse_mask, B=1, H=self.num_heads, Q_LEN=current_seq_len_sum, KV_LEN=current_seq_len_sum, device=device, BLOCK_SIZE=BLOCK_SIZE, _compile=False ) return attention_mask def forward( self, sequence_length: int, packed_text_ids: torch.LongTensor, packed_text_indexes: torch.LongTensor, sample_lens: List[int], sample_type: List[str], sample_N_target: List[int], packed_position_ids: torch.LongTensor, nested_attention_masks: List[torch.Tensor] = None, split_lens: List[int] = None, attn_modes: List[str] = None, ce_loss_indexes: Optional[torch.BoolTensor] = None, packed_label_ids: Optional[torch.LongTensor] = None, packed_vit_tokens: Optional[torch.Tensor] = None, packed_vit_token_indexes: Optional[torch.LongTensor] = None, packed_vit_position_ids: Optional[torch.LongTensor] = None, vit_token_seqlens: Optional[torch.IntTensor] = None, vit_video_grid_thw: Optional[torch.IntTensor] = None, vae_video_grid_thw: Optional[torch.IntTensor] = None, video_grid_thw: Optional[torch.IntTensor] = None, # for visual generation padded_latent: Optional[torch.Tensor] = None, patchified_vae_latent_shapes: Optional[List[Tuple[int, int]]] = None, packed_latent_position_ids: Optional[torch.LongTensor] = None, packed_vae_token_indexes: Optional[torch.LongTensor] = None, packed_timesteps: Optional[torch.LongTensor] = None, mse_loss_indexes: Optional[torch.BoolTensor] = None, vit_data_mode: Optional[List[str]] = None, # Indicates whether each VIT split is online or offline. sample_task: Optional[torch.LongTensor] = None, sample_modality: Optional[torch.LongTensor] = None, BLOCK_SIZE: int = 128, ) -> torch.Tensor: """ Args: sequence_length: length of sequence. packed_text_ids: 1-D int tensor, packed text token ids. packed_text_indexes: 1-D int tensor, packed text token indexes in sequence. sample_lens: A list of N ints, length of each sample in packed_sequence. nested_attention_masks: A list of N 2-D float tensor, where 0.0 means attention and -inf means ignore. packed_position_ids: packed 1-D positions, an image has only one global position shared by all latent tokens. packed_vit_tokens: packed patchified image tokens for vit model. packed_vit_position_ids: 1-D int tensor, the position of each token for vit model. packed_vit_token_indexes: 1-D int tensor, packed vit token indexes in sequence. vit_token_seqlens: 1-D int tensor, the length of each image tokens for vit model. packed_label_ids: 1-D int tensor, packed label token ids. ce_loss_indexes: 1-D bool tensor, where to compute ce loss. padded_latent: padded latent from VAE encoder. patchified_vae_latent_shapes: A list of (h, w) tuples, patchfied latent shapes of each image. packed_latent_position_ids: 1-D int tensor, the position of each token for latent. packed_vae_token_indexes: 1-D int tensor, padded image token indexes in sequence. packed_timesteps: 1-D float tensor, flow timesteps. 0 indicates use clean image. mse_loss_indexes: 1-D bool tensor, where to compute mse loss. """ N_vit_split = attn_modes.count("full") device = packed_text_ids.device apply_qwen_2_5_vl_pos_emb = getattr(self.training_args, "apply_qwen_2_5_vl_pos_emb", False) sample_splits = map_splits_to_samples(sample_lens, split_lens) if apply_qwen_2_5_vl_pos_emb: # TODO : packed_position_ids = [] sample_lens_tensor = torch.tensor(sample_lens, device=device, dtype=torch.long) cu_sample_lens = torch.cat([torch.zeros(1, device=device, dtype=torch.long), sample_lens_tensor.cumsum(0)[:-1]]) for i_sample in range(len(sample_lens) - 1): text_ids = packed_text_ids[cu_sample_lens[i_sample] : cu_sample_lens[i_sample + 1]] left, right = sample_splits[i_sample][0], sample_splits[i_sample][-1] + 1 grid_thw_rope = video_grid_thw[i_sample] i_sample_task = sample_task[cu_sample_lens[i_sample] : cu_sample_lens[i_sample + 1]] i_sample_modality = sample_modality[cu_sample_lens[i_sample] : cu_sample_lens[i_sample + 1]] current_packed_position_ids, rope_deltas = self.language_model.get_rope_index( input_ids=text_ids.unsqueeze(0), image_grid_thw=grid_thw_rope, video_grid_thw=grid_thw_rope, second_per_grid_ts=[1.0]*len(grid_thw_rope), attention_mask=torch.ones([1, len(text_ids)], dtype=torch.long, device=device), ) current_packed_position_ids = shift_position_ids(current_packed_position_ids, pos_shift = 1000, attn_modes = attn_modes[left:right], split_lens = split_lens[left:right], shift_attn_mode=['full_noise',"full"], pro_type = 10, i_sample_task=i_sample_task, i_sample_modality=i_sample_modality) packed_position_ids.append(current_packed_position_ids) packed_position_ids = torch.cat(packed_position_ids, dim=-1) packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids) packed_sequence = packed_text_embedding.new_zeros(size=(sequence_length, self.hidden_size)) packed_sequence[packed_text_indexes] = packed_text_embedding[packed_text_indexes] if nested_attention_masks is None: attn_modes_ = ["full" if mode=="full_noise" else mode for mode in attn_modes] sparse_mask = create_sparse_mask(sample_lens, split_lens, attn_modes_, packed_text_embedding.device) seqlen = sum(sample_lens) attention_mask = create_block_mask(sparse_mask, B=1, H=self.num_heads, Q_LEN=seqlen, KV_LEN=seqlen, device=packed_text_embedding.device, BLOCK_SIZE=BLOCK_SIZE, _compile=True) else: attention_mask = nested_attention_masks if N_vit_split > 0: if self.vit_type in ("qwen2_5_vl", "qwen_2_5_vl_original"): with torch.no_grad(): packed_vit_token_embed = make_packed_vit_token_embed(packed_vit_tokens, vit_data_mode, vit_video_grid_thw, self.vit_model) if self.vit_type == "qwen2_5_vl": packed_vit_token_embed = self.connector(packed_vit_token_embed) packed_sequence[packed_vit_token_indexes] = packed_vit_token_embed # flow matching loss if self.config.visual_gen: pt, ph, pw = self.latent_patch_size packed_latent = [] for latent, (t, h, w) in zip(padded_latent, patchified_vae_latent_shapes): patches = rearrange(latent, "(t pt) (h ph) (w pw) c -> (t h w) (pt ph pw c)", t=t, pt=pt, h=h, ph=ph, w=w, pw=pw) packed_latent.append(patches) packed_latent_clean = torch.cat(packed_latent, dim=0) noise = torch.randn_like(packed_latent_clean) if getattr(self.training_args, "incre_time_pro", 0) <=0: packed_timesteps = torch.sigmoid(packed_timesteps) packed_timesteps = self.timestep_shift * packed_timesteps / (1 + (self.timestep_shift - 1) * packed_timesteps) packed_latent = (1 - packed_timesteps[:, None]) * packed_latent_clean + packed_timesteps[:, None] * noise packed_timestep_embeds = self.time_embedder(packed_timesteps) latent_token_pos_emb = self.latent_pos_embed(packed_latent_position_ids) packed_latent = self.vae2llm(packed_latent) + packed_timestep_embeds + latent_token_pos_emb packed_sequence[packed_vae_token_indexes] = packed_latent.to(packed_sequence.dtype) extra_inputs = {} if self.use_moe: packed_und_token_indexes = packed_text_indexes if packed_vit_token_indexes is not None: packed_und_token_indexes = torch.cat([packed_text_indexes, packed_vit_token_indexes], dim=0) extra_inputs.update( packed_und_token_indexes=packed_und_token_indexes, packed_gen_token_indexes=packed_vae_token_indexes, ) last_hidden_state = self.language_model( packed_sequence=packed_sequence, sample_lens=sample_lens, attention_mask=attention_mask, packed_position_ids=packed_position_ids, **extra_inputs, ) mse, frame_mse, total_mse_tokens = None, None, None if self.config.visual_gen: packed_mse_preds = self.llm2vae(last_hidden_state[mse_loss_indexes]) total_mse_tokens = packed_mse_preds.shape[0] target = noise - packed_latent_clean has_mse = packed_timesteps > 0 mse = (packed_mse_preds - target[has_mse]) ** 2 ce = None if ce_loss_indexes is not None: V_eff = self.vocab_size_efficient ignore_index = -100 h = last_hidden_state[ce_loss_indexes] logits = self.language_model.lm_head(h)[..., :V_eff] targets = packed_label_ids.to(dtype=torch.long) invalid = (targets >= V_eff) | (targets < 0) targets = torch.where(invalid, torch.full_like(targets, ignore_index), targets) ce = F.cross_entropy(logits, targets, reduction="none", ignore_index=ignore_index) return dict(mse=mse, ce=ce, frame_mse=frame_mse, total_mse_tokens=total_mse_tokens) @torch.no_grad() def validation_gen( self, val_packed_text_ids: torch.LongTensor, val_packed_text_indexes: torch.LongTensor, val_packed_vit_tokens: torch.LongTensor, val_packed_vit_token_indexes: torch.LongTensor, val_sample_lens: List[int], val_packed_position_ids: torch.LongTensor, val_split_lens: List[int] = None, val_attn_modes: List[str] = None, val_sample_N_target: List[int] = None, vit_video_grid_thw: Optional[torch.IntTensor] = None, vae_video_grid_thw: Optional[torch.IntTensor] = None, video_grid_thw: Optional[torch.IntTensor] = None, val_mse_loss_indexes: Optional[torch.BoolTensor] = None, val_packed_vae_token_indexes: Optional[torch.LongTensor] = None, val_padded_latent: Optional[torch.Tensor] = None, sample_task: Optional[torch.LongTensor] = None, sample_modality: Optional[torch.LongTensor] = None, video_sizes: List[Tuple[int, int, int]] = [[1, 256, 256]], val_padded_videos: torch.Tensor = None, timestep_shift: float = 4.0, num_timesteps: int = 24, cfg_interval: Optional[Tuple[float, float]] = [0, 1], cfg_renorm_min: float = 0.0, cfg_renorm_type: str = "global", cfg_text_scale: float = 1.0, cfg_vit_scale: float = 1.0, # HACK device=None, dtype=None, new_token_ids=None, BLOCK_SIZE: int = 128, apply_chat_template: bool = False, apply_qwen_2_5_vl_pos_emb: bool = False, image_token_id: int = 151655, caption: Optional[List[str]] = None, index: str = "", **kwargs, ): start_id = new_token_ids["start_of_image"] end_id = new_token_ids["end_of_image"] pt, ph, pw = self.latent_patch_size index_dtype = val_packed_text_ids.dtype cu_sample_lens = torch.nn.functional.pad(torch.cumsum(torch.tensor(val_sample_lens, device=device), dim=0), (1, 0)) sample_splits = map_splits_to_samples(val_sample_lens, val_split_lens) if val_packed_vit_tokens is not None and vit_video_grid_thw is not None: vit_sample_len = vit_video_grid_thw[:, 0] * vit_video_grid_thw[:, 1] * vit_video_grid_thw[:, 2] cu_vit_sample_lens = torch.cat([torch.zeros(1, device=vit_video_grid_thw.device, dtype=vit_sample_len.dtype), vit_sample_len.cumsum(0)]) self.vit_model = self.vit_model.to(device=device, dtype=dtype) val_packed_vit_tokens = torch.cat(val_packed_vit_tokens, dim=0) x_t_all = [] max_samples = kwargs.get("max_samples", 16) num_samples = len(val_sample_lens) max_samples = min(num_samples, max_samples) gen_idx = 0 curr_vae_split_idx, curr_vit_split_idx = 0, 0 padded_videos = [] for i_sample in range(num_samples): left, right = sample_splits[i_sample][0], sample_splits[i_sample][-1] + 1 # --- for interleave --- current_split_lens = val_split_lens[left:right] current_attn_modes = val_attn_modes[left:right] N_noise_element = current_attn_modes.count("noise") + current_attn_modes.count("full_noise") + current_attn_modes.count("full_noise_target") N_vit_split = current_attn_modes.count("full") if right > len(val_attn_modes): break if N_noise_element<=0: curr_vit_split_idx += N_vit_split continue if gen_idx >= max_samples: break # 1. Get the slice information of the current sample within the entire batch sample_start_idx = cu_sample_lens[i_sample] sample_end_idx = cu_sample_lens[i_sample + 1] current_seq_len = val_sample_lens[i_sample] current_pos_ids = val_packed_position_ids[sample_start_idx:sample_end_idx] i_sample_task = sample_task[sample_start_idx:sample_end_idx] i_sample_modality = sample_modality[sample_start_idx:sample_end_idx] vae_mask = (val_packed_vae_token_indexes >= sample_start_idx) & (val_packed_vae_token_indexes < sample_end_idx) current_vae_token_indexes_local = val_packed_vae_token_indexes[vae_mask] - sample_start_idx # --- VAE MSE token part: indices of the positions in x_t that need to be updated --- vae_mse_mask = (val_mse_loss_indexes >= sample_start_idx) & (val_mse_loss_indexes < sample_end_idx) current_vae_mse_indexes_local = val_mse_loss_indexes[vae_mse_mask] - sample_start_idx # Indices of x_t positions that need updates. current_vae_mse_indexes_local_in_vae = ( current_vae_mse_indexes_local - current_vae_mse_indexes_local[0] + torch.where(current_vae_token_indexes_local == current_vae_mse_indexes_local[0])[0] ) num_vid_tokens_list, vid_shape_list, vae_position_ids, curr_padded_latent = [], [], [], [] # 2. Generate vit uncond features (optional) cfg_vit_pro = False if cfg_vit_scale > 1.0 and "full" in current_attn_modes: vit_uncond_sequence, vit_uncond_attn_modes, vit_uncond_split_lens, vit_uncond_vae_index, _, vit_uncond_packed_gen_token_indexes, vit_uncond_packed_und_token_indexes, vit_uncond_text_ids, vit_uncond_seq_len, vit_uncond_pad = uncond_split_pro(self.language_model, current_attn_modes, current_split_lens, vae_video_grid_thw, vit_video_grid_thw, curr_vae_split_idx, curr_vit_split_idx, device, dtype, start_id, image_token_id, end_id, BLOCK_SIZE, is_text_uncond = True, is_vit_uncond = True) cfg_vit_pro = True for i_target in range(N_noise_element): T, H, W = video_sizes[curr_vae_split_idx] t = (T - 1) // self.latent_downsample_temporal + 1 h = H // self.latent_downsample_spatial w = W // self.latent_downsample_spatial vid_shape_list.append([t, h, w]) num_vid_tokens_list.append(t * h * w) # prepare packed_vae_position_ids vae_position_ids.append( get_flattened_position_ids_extrapolate_video(t, h, w, max_latent_size=self.max_latent_size) ) if len(current_vae_mse_indexes_local) != len(current_vae_token_indexes_local): padded_latent_ = val_padded_latent[curr_vae_split_idx] # (T,H,W,C) patches = rearrange(padded_latent_, "(t pt) (h ph) (w pw) c -> (t h w) (pt ph pw c)", t=t, pt=pt, h=h, ph=ph, w=w, pw=pw) curr_padded_latent.append(patches) if val_padded_videos is not None: padded_videos.append(val_padded_videos[curr_vae_split_idx]) curr_vae_split_idx += 1 num_vid_tokens = sum(num_vid_tokens_list) vae_position_ids = torch.cat(vae_position_ids, 0) if curr_padded_latent != []: curr_padded_latent = torch.cat(curr_padded_latent, dim=0).to(dtype) # 2. Reconstruct the input sequence and attention mask for the current sample current_sequence = torch.zeros((current_seq_len, self.hidden_size), device=device, dtype=dtype) # --- Text part --- text_mask = (val_packed_text_indexes >= sample_start_idx) & (val_packed_text_indexes < sample_end_idx) current_text_indexes_local = val_packed_text_indexes[text_mask] - sample_start_idx current_text_ids = val_packed_text_ids[sample_start_idx:sample_end_idx] current_text_embedding = self.language_model.model.embed_tokens(current_text_ids).to(dtype=dtype) current_sequence[current_text_indexes_local] = current_text_embedding[current_text_indexes_local] if cfg_text_scale > 1.0: if cfg_vit_pro: vit_uncond_attn_modes_, vit_uncond_split_lens_ = vit_uncond_attn_modes, vit_uncond_split_lens vit_uncond_attn_mask = self.process_attention_mask(vit_uncond_attn_modes_, vit_uncond_split_lens_, [vit_uncond_seq_len, vit_uncond_pad], device = device, BLOCK_SIZE = BLOCK_SIZE) # --- VIT part: support ti2i --- if N_vit_split != 0: vit_sample_start_idx = cu_vit_sample_lens[curr_vit_split_idx] vit_sample_end_idx = cu_vit_sample_lens[curr_vit_split_idx + N_vit_split] current_val_packed_vit_tokens = val_packed_vit_tokens[vit_sample_start_idx:vit_sample_end_idx].to(dtype) current_val_vit_video_grid_thw = vit_video_grid_thw[curr_vit_split_idx : curr_vit_split_idx + N_vit_split] curr_vit_split_idx += N_vit_split if self.vit_type in ["qwen2_5_vl", "qwen_2_5_vl_original"]: packed_vit_token_embed = self.vit_model(hidden_states=current_val_packed_vit_tokens, grid_thw=current_val_vit_video_grid_thw) if self.vit_type in ["qwen2_5_vl"]: packed_vit_token_embed = self.connector(packed_vit_token_embed).to(dtype) else: raise NotImplementedError(f"{self.vit_type} is not supported") vit_mask = (val_packed_vit_token_indexes >= sample_start_idx) & (val_packed_vit_token_indexes < sample_end_idx) current_vit_indexes_local = val_packed_vit_token_indexes[vit_mask] - sample_start_idx current_sequence[current_vit_indexes_local] = packed_vit_token_embed current_seq_len_pad = (current_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE * BLOCK_SIZE current_pad = current_seq_len_pad - current_seq_len if current_pad > 0: current_split_lens = current_split_lens + [current_pad] current_attn_modes = current_attn_modes + ["causal"] current_split_lens_, current_attn_modes_ = current_split_lens, current_attn_modes attention_mask = self.process_attention_mask(current_attn_modes_, current_split_lens_, [current_seq_len, current_pad], device = device, BLOCK_SIZE = BLOCK_SIZE) validation_noise_seed = kwargs.get("validation_noise_seed", -1) if validation_noise_seed > 0: generator = torch.Generator(device=device).manual_seed(validation_noise_seed + get_global_rank() * max_samples + i_sample) else: generator = None x_t = torch.randn(num_vid_tokens, self.patch_latent_dim, generator=generator, device=device, dtype=dtype) if curr_padded_latent != []: curr_padded_latent[current_vae_mse_indexes_local_in_vae] = x_t[current_vae_mse_indexes_local_in_vae] x_t = curr_padded_latent timesteps = torch.linspace(1, 0, num_timesteps + 1, device=x_t.device) timesteps = timestep_shift * timesteps / (1 + (timestep_shift - 1) * timesteps) dts = timesteps[:-1] - timesteps[1:] timesteps = timesteps[:-1] if apply_qwen_2_5_vl_pos_emb: grid_thw_rope = video_grid_thw[i_sample] current_pos_ids, _ = self.language_model.get_rope_index( input_ids=current_text_ids.unsqueeze(0), image_grid_thw=grid_thw_rope, video_grid_thw=grid_thw_rope, second_per_grid_ts=[1.0]*len(grid_thw_rope), attention_mask=torch.ones([1, len(current_text_ids)], dtype=torch.long, device=device), ) current_pos_ids = shift_position_ids( current_pos_ids, pos_shift=1000, attn_modes=current_attn_modes, split_lens=current_split_lens, shift_attn_mode=["full_noise", "full"], pro_type=10, i_sample_task=i_sample_task, i_sample_modality=i_sample_modality, ) if cfg_text_scale > 1.0: uncond_mask = i_sample_modality!=0 _, uncond_pos_ids, uncond_attn_mask, _, _, uncond_extra_inputs, uncond_seq_len = self.uncond_split_pro_new( uncond_mask, current_text_ids, current_attn_modes, current_split_lens, device, dtype, BLOCK_SIZE, grid_thw_rope, apply_qwen_2_5_vl_pos_emb, i_sample_task=i_sample_task, i_sample_modality=i_sample_modality, ) for _ in range(1): timestep = torch.zeros(x_t.shape[0], device=x_t.device) for i, timestep_ in enumerate(timesteps): timestep[current_vae_mse_indexes_local_in_vae] = torch.tensor([timestep_] * current_vae_mse_indexes_local_in_vae.shape[0], device=x_t.device) if timestep_ > cfg_interval[0] and timestep_ <= cfg_interval[1]: cfg_text_scale_ = cfg_text_scale cfg_vit_scale_ = cfg_vit_scale else: cfg_text_scale_ = 1.0 cfg_vit_scale_ = 1.0 # --- vae encoder --- timestep_embed = self.time_embedder(timestep) latent_pos_embed = self.latent_pos_embed(vae_position_ids) vae_embed = self.vae2llm(x_t) + timestep_embed + latent_pos_embed vae_embed = vae_embed.to(current_sequence.dtype) current_sequence[current_vae_token_indexes_local] = vae_embed extra_inputs = {} if self.use_moe: if N_vit_split != 0: packed_und_token_indexes = torch.cat([current_text_indexes_local, current_vit_indexes_local], dim=0) else: packed_und_token_indexes = current_text_indexes_local extra_inputs.update( packed_und_token_indexes=packed_und_token_indexes.to(dtype=index_dtype), packed_gen_token_indexes=current_vae_token_indexes_local.to(dtype=index_dtype), ) self.language_model.to(current_sequence.dtype) cond_hidden_state = self.language_model( packed_sequence=current_sequence[:current_seq_len], sample_lens=[current_seq_len], attention_mask=attention_mask, packed_position_ids=current_pos_ids.to(dtype=index_dtype), mode_forward="validation", **extra_inputs, ) v_t = self.llm2vae(cond_hidden_state[current_vae_mse_indexes_local]) # cfg text forward if cfg_text_scale_ > 1.0: uncond_sequence = current_sequence[uncond_mask] cfg_text_v_t = self.uncond_forward(uncond_sequence, uncond_pos_ids, uncond_seq_len, uncond_attn_mask, uncond_extra_inputs, current_vae_mse_indexes_local, current_seq_len) if cfg_vit_pro: if i_sample_task is not None: i_sample_task_text_uncond = i_sample_task[i_sample_modality!=0] i_sample_modality_text_uncond = i_sample_modality[i_sample_modality!=0] else: i_sample_task_text_uncond, i_sample_modality_text_uncond = None, None if i_sample_task is not None: i_sample_task_text_vit_uncond = i_sample_task_text_uncond[i_sample_modality_text_uncond!=4] i_sample_modality_text_vit_uncond = i_sample_modality_text_uncond[i_sample_modality_text_uncond!=4] else: i_sample_task_text_vit_uncond, i_sample_modality_text_vit_uncond = None, None cfg_text_vit_v_t = self.uncond_forward(vae_embed, vit_uncond_sequence, vit_uncond_text_ids, vit_uncond_seq_len, vit_uncond_packed_und_token_indexes, vit_uncond_packed_gen_token_indexes, vit_uncond_attn_mask, vit_uncond_vae_index, grid_thw_rope, current_vae_mse_indexes_local, current_seq_len, apply_qwen_2_5_vl_pos_emb, device,i_sample_task_text_vit_uncond,i_sample_modality_text_vit_uncond) v_t_ = cfg_text_vit_v_t + cfg_text_scale_ * (v_t - cfg_text_v_t) + cfg_vit_scale_ * (cfg_text_v_t - cfg_text_vit_v_t) else: v_t_ = cfg_text_v_t + cfg_text_scale_ * (v_t - cfg_text_v_t) if cfg_renorm_type == "global": norm_v_t = torch.norm(v_t) norm_v_t_ = torch.norm(v_t_) scale = (norm_v_t / (norm_v_t_ + 1e-8)).clamp(min=cfg_renorm_min, max=1.0) elif cfg_renorm_type == "channel": norm_v_t = torch.norm(v_t, dim=-1, keepdim=True) norm_v_t_ = torch.norm(v_t_, dim=-1, keepdim=True) scale = (norm_v_t / (norm_v_t_ + 1e-8)).clamp(min=cfg_renorm_min, max=1.0) elif cfg_renorm_type.lower() in ("", "none", "null"): scale = 1 else: raise NotImplementedError(f"{cfg_renorm_type} is not suppoprted") v_t = v_t_ * scale x_t[current_vae_mse_indexes_local_in_vae] = x_t[current_vae_mse_indexes_local_in_vae] - v_t.to(x_t.device) * dts[i] curr_seq_target, patch = 0, [] for i_target in range(N_noise_element): pt, ph, pw = self.latent_patch_size t, h, w = vid_shape_list[i_target] len_target = t * h * w x_t_ = rearrange(x_t[curr_seq_target : curr_seq_target + len_target], "(t h w) (pt ph pw c) -> (t pt) (h ph) (w pw) c", t=t, h=h, w=w, pt=pt, ph=ph, pw=pw) patch.append(x_t_) curr_seq_target += len_target x_t_all.append(patch) gen_idx += 1 if caption != None: return x_t_all, [caption], padded_videos, index return x_t_all def uncond_split_pro_new( self, uncond_mask, current_text_ids, current_attn_modes, current_split_lens, device, dtype, BLOCK_SIZE, grid_thw_rope=None, apply_qwen_2_5_vl_pos_emb=False, i_sample_task=None, i_sample_modality=None, uncond_pos_ids=None, ): start = 0 uncond_split_lens, uncond_attn_modes, uncond_packed_gen_token_indexes = [], [], [] for i_visual, attn_mode_ in enumerate(current_attn_modes): split_len_ = current_split_lens[i_visual] end = start + split_len_ split_in_uncond = int(uncond_mask[start:end].sum()) start += split_len_ if split_in_uncond == 0: continue else: if attn_mode_ in ["noise", "full_noise"]: start_gen, end_gen = sum(uncond_split_lens) + 1, sum(uncond_split_lens) + 1 + split_len_ - 2 uncond_packed_gen_token_indexes.extend(range(start_gen, end_gen)) uncond_split_lens.append(split_in_uncond) uncond_attn_modes.append(attn_mode_) uncond_seq_len = sum(uncond_split_lens) uncond_seq_len_pad = (uncond_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE * BLOCK_SIZE uncond_pad = uncond_seq_len_pad - uncond_seq_len if uncond_pad > 0: uncond_split_lens.append(uncond_pad) uncond_attn_modes.append("causal") uncond_packed_gen_token_indexes = torch.tensor(uncond_packed_gen_token_indexes, dtype=torch.long, device=device) all_indexes = torch.arange(0, uncond_seq_len).to(device) und_token_mask = ~torch.isin(all_indexes, uncond_packed_gen_token_indexes) uncond_packed_und_token_indexes = all_indexes[und_token_mask] uncond_extra_inputs = {} if self.use_moe: uncond_extra_inputs.update( packed_und_token_indexes=uncond_packed_und_token_indexes, packed_gen_token_indexes=uncond_packed_gen_token_indexes, ) # Build the unconditional attention mask. uncond_attn_mask = self.process_attention_mask(uncond_attn_modes, uncond_split_lens, [uncond_seq_len, uncond_pad], device=device, BLOCK_SIZE=BLOCK_SIZE) # Extract text ids for the unconditional sequence. uncond_text_ids = current_text_ids[uncond_mask] uncond_sample_task = i_sample_task[uncond_mask] if i_sample_task is not None else None uncond_sample_modality = i_sample_modality[uncond_mask] if i_sample_modality is not None else None if apply_qwen_2_5_vl_pos_emb: uncond_pos_ids, uncond_rope_deltas = self.language_model.get_rope_index( input_ids=uncond_text_ids.unsqueeze(0), image_grid_thw=grid_thw_rope, video_grid_thw=grid_thw_rope, second_per_grid_ts=[1.0] * len(grid_thw_rope), attention_mask=torch.ones([1, len(uncond_text_ids)], dtype=torch.long, device=device), ) uncond_pos_ids = shift_position_ids( uncond_pos_ids, pos_shift=1000, attn_modes=uncond_attn_modes, split_lens=uncond_split_lens, shift_attn_mode=["full_noise", "full"], pro_type=10, i_sample_task=uncond_sample_task, i_sample_modality=uncond_sample_modality, ) else: uncond_pos_ids = torch.tensor(uncond_pos_ids, dtype=torch.long, device=device)[:uncond_seq_len] return ( uncond_text_ids, uncond_pos_ids, uncond_attn_mask, uncond_attn_modes, uncond_split_lens, uncond_extra_inputs, uncond_seq_len, ) def uncond_forward( self, uncond_sequence, uncond_pos_ids, uncond_seq_len, uncond_attn_mask, uncond_extra_inputs, current_vae_mse_indexes_local, current_seq_len, ): uncond_hidden_state = self.language_model( packed_sequence=uncond_sequence[:uncond_seq_len], sample_lens=[uncond_seq_len], attention_mask=uncond_attn_mask, packed_position_ids=uncond_pos_ids, mode_forward="validation", # NOTE **uncond_extra_inputs, ) uncond_current_vae_mse_indexes_local = current_vae_mse_indexes_local - (current_seq_len - uncond_seq_len) cfg_text_v_t = self.llm2vae(uncond_hidden_state[uncond_current_vae_mse_indexes_local]) return cfg_text_v_t @torch.no_grad() def validation_video_to_text( self, val_packed_text_ids: torch.LongTensor, val_packed_text_indexes: torch.LongTensor, val_packed_position_ids: torch.LongTensor, val_ce_loss_indexes: torch.LongTensor, val_sample_N_target: List[int], val_split_lens: List[int], val_attn_modes: List[str], val_sample_lens: List[int], val_sample_type: List[str], val_packed_vit_tokens: Optional[torch.Tensor] = None, val_vit_video_grid_thw: Optional[torch.IntTensor] = None, max_samples: int = 1, max_length: int = 256, device: torch.device = None, dtype: torch.dtype = None, new_token_ids: Dict[str, int] = None, pad_token_id: int = None, vocab_size: int = None, do_sample: bool = False, temperature: float = 1.0, caption: any = "", tokenizer: any = None, apply_chat_template: bool = False, apply_qwen_2_5_vl_pos_emb: bool = False, image_token_id: int = 151655, BLOCK_SIZE: int = 128, visualize_generation_progress: bool = False, index: str = "", ): # Special tokens. start_id = new_token_ids["start_of_image"] end_id = new_token_ids["end_of_image"] bos_id = new_token_ids["bos_token_id"] eos_id = new_token_ids["eos_token_id"] # Per-sample lengths. cu_sample_lens = torch.nn.functional.pad(torch.cumsum(torch.tensor(val_sample_lens, device=device), dim=0), (1, 0)) sample_splits = map_splits_to_samples(val_sample_lens, val_split_lens) # Length of each VIT token sequence in each sample. vit_sample_len = val_vit_video_grid_thw[:, 0] * val_vit_video_grid_thw[:, 1] * val_vit_video_grid_thw[:, 2] # shape: (N,) , N = 1 * 16 * 16, cu_vit_sample_lens = torch.cat([torch.zeros(1, device=val_vit_video_grid_thw.device, dtype=vit_sample_len.dtype), vit_sample_len.cumsum(0)]) if val_packed_vit_tokens is not None: val_packed_vit_tokens = torch.cat(val_packed_vit_tokens, dim=0) max_samples = min(len(val_sample_lens), max_samples) cnt_samples = 0 generated_sequence_all = [] L = len(val_sample_lens) curr_vit_split_idx = 0 for i_sample in range(L): left, right = sample_splits[i_sample][0], sample_splits[i_sample][-1] + 1 # --- for interleave --- current_split_lens = val_split_lens[left:right] current_attn_modes = val_attn_modes[left:right] N_target = val_sample_N_target[i_sample] N_vit_split = current_attn_modes.count("full") if val_sample_type[i_sample] != "und": curr_vit_split_idx += N_vit_split continue cnt_samples += 1 if cnt_samples > max_samples: break assert N_target == 1 # Get slice information for the current video VIT sample in the batch. vit_sample_start_idx = cu_vit_sample_lens[curr_vit_split_idx] vit_sample_end_idx = cu_vit_sample_lens[curr_vit_split_idx + N_vit_split] current_val_packed_vit_tokens = val_packed_vit_tokens[vit_sample_start_idx:vit_sample_end_idx] current_val_vit_video_grid_thw = val_vit_video_grid_thw[curr_vit_split_idx : curr_vit_split_idx + N_vit_split] curr_vit_split_idx += N_vit_split if N_vit_split > 0 : if self.vit_type in ["qwen2_5_vl", "qwen_2_5_vl_original"]: packed_vit_token_embed = self.vit_model(hidden_states=current_val_packed_vit_tokens, grid_thw=current_val_vit_video_grid_thw) if self.vit_type in ["qwen2_5_vl"]: packed_vit_token_embed = self.connector(packed_vit_token_embed).to(dtype) else: raise NotImplementedError(f"{self.vit_type} is not supported") sample_start_idx = cu_sample_lens[i_sample] sample_end_idx = cu_sample_lens[i_sample + 1] current_pos_ids = val_packed_position_ids[sample_start_idx:sample_end_idx] text_mask_ce = (val_ce_loss_indexes >= sample_start_idx) & (val_ce_loss_indexes < sample_end_idx) current_ce_loss_indexes_local = val_ce_loss_indexes[text_mask_ce] - sample_start_idx if text_mask_ce.any(): current_text_ids = val_packed_text_ids[sample_start_idx:sample_end_idx][: current_ce_loss_indexes_local[0] + 1] else: current_text_ids = val_packed_text_ids[sample_start_idx:sample_end_idx] num_text_ids = current_text_ids.shape[0] num_last_split = num_text_ids - sum(current_split_lens[:-N_target]) current_split_lens = current_split_lens[:-N_target] if num_last_split > 1: current_split_lens.extend([num_last_split - 1]) max_seq_len = (max_length + num_text_ids + BLOCK_SIZE - 1) // BLOCK_SIZE * BLOCK_SIZE num_pad = max_seq_len - num_text_ids current_text_ids = torch.cat( [current_text_ids, torch.full((num_pad,), pad_token_id, dtype=torch.long, device=device)], dim=0 ) packed_text_embedding = self.language_model.model.embed_tokens(current_text_ids).to(dtype) if N_vit_split > 0 : mask = current_text_ids == image_token_id mask_unsqueezed = mask.unsqueeze(-1) mask_expanded = mask_unsqueezed.expand_as(packed_text_embedding) image_mask = mask_expanded.to(packed_text_embedding.device) curr_packed_sequence = packed_text_embedding.masked_scatter(image_mask, packed_vit_token_embed) else: curr_packed_sequence = packed_text_embedding step = num_text_ids - 1 generated_sequence = [] if apply_qwen_2_5_vl_pos_emb: current_packed_position_ids, rope_deltas = self.language_model.get_rope_index( input_ids=current_text_ids.unsqueeze(0), image_grid_thw=current_val_vit_video_grid_thw, video_grid_thw=current_val_vit_video_grid_thw, second_per_grid_ts=[1.0], attention_mask=torch.ones([1, max_seq_len], dtype=torch.long, device=device), # Full-one attention mask. ) else: current_pos_ids = current_pos_ids[:num_text_ids] pos_pad_start = int(current_pos_ids[-1] + 1) current_pad = torch.arange(pos_pad_start, pos_pad_start + num_pad, device=device) current_packed_position_ids = torch.cat([current_pos_ids, current_pad], dim=0) current_sample_lens = [max_seq_len] seqlen = sum(current_sample_lens) current_attn_modes_ = current_attn_modes[: len(current_split_lens)] + ["causal", "causal"] current_attn_modes_ = ["full" if mode_=="full_noise" else mode_ for mode_ in current_attn_modes_] while step < (max_seq_len - 1): current_text_len = (step + 1) - (num_text_ids - 1) current_split_lens_ = current_split_lens + [current_text_len, num_pad + 1 - current_text_len] sparse_mask = create_sparse_mask(current_sample_lens, current_split_lens_, current_attn_modes_, device) attention_mask = create_block_mask(sparse_mask, B=1, H=self.num_heads, Q_LEN=seqlen, KV_LEN=seqlen, device=device, BLOCK_SIZE=BLOCK_SIZE, _compile=False) extra_inputs = {"mode": "und"} if self.use_moe: packed_und_token_indexes = torch.arange(0, max_seq_len, device=device) extra_inputs.update( packed_und_token_indexes=packed_und_token_indexes, packed_gen_token_indexes=None, ) last_hidden_state = self.language_model( packed_sequence=curr_packed_sequence.to(dtype=dtype), sample_lens=current_sample_lens, attention_mask=attention_mask, packed_position_ids=current_packed_position_ids, mode_forward="validation", **extra_inputs, ) pred_logits = self.language_model.lm_head(last_hidden_state[step : step + 1, :]) pred_logits[:, vocab_size:] = float("-inf") if do_sample: probs = nn.functional.softmax(pred_logits / temperature, dim=-1) curr_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) else: curr_tokens = torch.argmax(pred_logits, dim=-1) generated_sequence.append(curr_tokens) if visualize_generation_progress: print(f"curr_tokens: {curr_tokens}", curr_tokens.item(), ", eos_id:", eos_id) if curr_tokens.item() == eos_id: break curr_packed_sequence[step + 1] = self.language_model.model.embed_tokens(curr_tokens) step += 1 generated_sequence = torch.stack([i.to(device) for i in generated_sequence], dim=0) generated_sequence_all.append(generated_sequence) return generated_sequence_all, caption, index @torch.no_grad() def validation_und_KVcache( self, val_packed_text_ids: torch.LongTensor, val_packed_text_indexes: torch.LongTensor, val_packed_position_ids: torch.LongTensor, val_ce_loss_indexes: torch.LongTensor, val_sample_N_target: List[int], val_split_lens: List[int], val_attn_modes: List[str], val_sample_lens: List[int], val_sample_type: List[str], val_packed_vit_tokens: Optional[torch.Tensor] = None, val_vit_video_grid_thw: Optional[torch.IntTensor] = None, max_samples: int = 1, max_length: int = 256, device: torch.device = None, dtype: torch.dtype = None, new_token_ids: Dict[str, int] = None, pad_token_id: int = None, vocab_size: int = None, do_sample: bool = False, temperature: float = 1.0, caption: any = "", tokenizer: any = None, apply_chat_template: bool = False, apply_qwen_2_5_vl_pos_emb: bool = False, image_token_id: int = 151655, BLOCK_SIZE: int = 128, visualize_generation_progress: bool = False, index: str = "", ): eos_id = new_token_ids["eos_token_id"] cu_sample_lens = torch.nn.functional.pad(torch.cumsum(torch.tensor(val_sample_lens, device=device), dim=0), (1, 0)) sample_splits = map_splits_to_samples(val_sample_lens, val_split_lens) vit_sample_len = val_vit_video_grid_thw[:, 0] * val_vit_video_grid_thw[:, 1] * val_vit_video_grid_thw[:, 2] cu_vit_sample_lens = torch.cat([torch.zeros(1, device=val_vit_video_grid_thw.device, dtype=vit_sample_len.dtype), vit_sample_len.cumsum(0)]) if val_packed_vit_tokens is not None: self.vit_model = self.vit_model.to(device=device, dtype=dtype) val_packed_vit_tokens = torch.cat(val_packed_vit_tokens, dim=0) max_samples = min(len(val_sample_lens), max_samples) cnt_samples = 0 generated_sequence_all = [] curr_vit_split_idx = 0 def _slice_position_ids(position_ids, start, end): if position_ids.dim() == 3: return position_ids[:, :, start:end] return position_ids[start:end] def _update_und_context(gen_context, sequence, position_ids, start, end, is_causal): query_len = end - start if query_len <= 0: return gen_context query_index = int(gen_context["kv_lens"][0].item()) output = self.language_model.forward_inference( packed_query_sequence=sequence[start:end], query_lens=torch.tensor([query_len], dtype=torch.int32, device=device), packed_query_position_ids=_slice_position_ids(position_ids, start, end), packed_query_indexes=torch.arange(query_index, query_index + query_len, dtype=torch.long, device=device), past_key_values=gen_context["past_key_values"], key_values_lens=gen_context["kv_lens"], packed_key_value_indexes=torch.arange(0, query_index, dtype=torch.long, device=device), update_past_key_values=True, is_causal=is_causal, mode="und", ) gen_context["past_key_values"] = output.past_key_values gen_context["kv_lens"] += query_len return gen_context self.language_model.eval() self.eval() for i_sample in range(len(val_sample_lens)): left, right = sample_splits[i_sample][0], sample_splits[i_sample][-1] + 1 current_split_lens = val_split_lens[left:right] current_attn_modes = val_attn_modes[left:right] N_target = val_sample_N_target[i_sample] N_vit_split = current_attn_modes.count("full") if val_sample_type[i_sample] != "und": curr_vit_split_idx += N_vit_split continue cnt_samples += 1 if cnt_samples > max_samples: break assert N_target == 1 vit_sample_start_idx = int(cu_vit_sample_lens[curr_vit_split_idx].item()) vit_sample_end_idx = int(cu_vit_sample_lens[curr_vit_split_idx + N_vit_split].item()) current_val_packed_vit_tokens = val_packed_vit_tokens[vit_sample_start_idx:vit_sample_end_idx].to(device=device, dtype=dtype) current_val_vit_video_grid_thw = val_vit_video_grid_thw[curr_vit_split_idx: curr_vit_split_idx + N_vit_split] curr_vit_split_idx += N_vit_split packed_vit_token_embed = None if N_vit_split > 0: if self.vit_type in ["qwen2_5_vl", "qwen_2_5_vl_original"]: packed_vit_token_embed = self.vit_model(hidden_states=current_val_packed_vit_tokens, grid_thw=current_val_vit_video_grid_thw) if self.vit_type in ["qwen2_5_vl"]: packed_vit_token_embed = self.connector(packed_vit_token_embed).to(dtype) else: raise NotImplementedError(f"{self.vit_type} is not supported") sample_start_idx = int(cu_sample_lens[i_sample].item()) sample_end_idx = int(cu_sample_lens[i_sample + 1].item()) current_pos_ids = val_packed_position_ids[sample_start_idx:sample_end_idx] text_mask_ce = (val_ce_loss_indexes >= sample_start_idx) & (val_ce_loss_indexes < sample_end_idx) current_ce_loss_indexes_local = val_ce_loss_indexes[text_mask_ce] - sample_start_idx if text_mask_ce.any().item(): current_text_ids = val_packed_text_ids[sample_start_idx:sample_end_idx][: current_ce_loss_indexes_local[0] + 1] else: current_text_ids = val_packed_text_ids[sample_start_idx:sample_end_idx] num_text_ids = current_text_ids.shape[0] context_len = num_text_ids - 1 num_last_split = num_text_ids - sum(current_split_lens[:-N_target]) current_split_lens = current_split_lens[:-N_target] if num_last_split > 1: current_split_lens.extend([num_last_split - 1]) current_attn_modes = current_attn_modes[: len(current_split_lens)] packed_sequence = self.language_model.model.embed_tokens(current_text_ids).to(dtype) if N_vit_split > 0: image_mask = (current_text_ids == image_token_id).unsqueeze(-1).expand_as(packed_sequence) packed_sequence = packed_sequence.masked_scatter(image_mask.to(packed_sequence.device), packed_vit_token_embed) pos_len = num_text_ids + max_length if apply_qwen_2_5_vl_pos_emb: pos_text_ids = torch.cat( [current_text_ids, torch.full((max_length,), pad_token_id, dtype=torch.long, device=device)], dim=0 ) current_packed_position_ids, _ = self.language_model.get_rope_index( input_ids=pos_text_ids.unsqueeze(0), image_grid_thw=current_val_vit_video_grid_thw, video_grid_thw=current_val_vit_video_grid_thw, second_per_grid_ts=[1.0] * max(N_vit_split, 1), attention_mask=torch.ones([1, pos_len], dtype=torch.long, device=device), ) else: current_pos_ids = current_pos_ids[:num_text_ids] pos_pad_start = int(current_pos_ids[-1] + 1) current_pad = torch.arange(pos_pad_start, pos_pad_start + max_length, device=device) current_packed_position_ids = torch.cat([current_pos_ids, current_pad], dim=0) gen_context = self.init_gen_context(device=device, dtype=torch.int32) current_start = 0 for attn_mode, split_len in zip(current_attn_modes, current_split_lens): current_end = min(current_start + split_len, context_len) if current_end <= current_start: continue is_causal = attn_mode not in ["full", "full_noise", "full_noise_target"] gen_context = _update_und_context(gen_context, packed_sequence, current_packed_position_ids, current_start, current_end, is_causal) current_start = current_end if current_start >= context_len: break if current_start < context_len: gen_context = _update_und_context(gen_context, packed_sequence, current_packed_position_ids, current_start, context_len, True) curr_tokens = current_text_ids[context_len:context_len + 1] generated_sequence = [] for step in range(max_length): packed_text_embedding = self.language_model.model.embed_tokens(curr_tokens).to(dtype) query_index = int(gen_context["kv_lens"][0].item()) output = self.language_model.forward_inference( packed_query_sequence=packed_text_embedding, query_lens=torch.ones(1, dtype=torch.int32, device=device), packed_query_position_ids=_slice_position_ids(current_packed_position_ids, context_len + step, context_len + step + 1), packed_query_indexes=torch.arange(query_index, query_index + 1, dtype=torch.long, device=device), past_key_values=gen_context["past_key_values"], key_values_lens=gen_context["kv_lens"], packed_key_value_indexes=torch.arange(0, query_index, dtype=torch.long, device=device), update_past_key_values=True, is_causal=True, mode="und", ) gen_context["past_key_values"] = output.past_key_values gen_context["kv_lens"] += 1 pred_logits = self.language_model.lm_head(output.packed_query_sequence) pred_logits[:, vocab_size:] = float("-inf") if do_sample: probs = nn.functional.softmax(pred_logits / temperature, dim=-1) curr_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) else: curr_tokens = torch.argmax(pred_logits, dim=-1) generated_sequence.append(curr_tokens) if visualize_generation_progress: print(f"curr_tokens: {curr_tokens}", curr_tokens.item(), ", eos_id:", eos_id) if curr_tokens.item() == eos_id: break generated_sequence = torch.stack([i.to(device) for i in generated_sequence], dim=0) generated_sequence_all.append(generated_sequence) return generated_sequence_all, caption, index def prepare_vit_images_validation(self, curr_kvlens, curr_rope, vit_tokens, new_token_ids, device): packed_vit_token_indexes = list() vit_token_seqlens, packed_vit_tokens, packed_vit_position_ids = list(), list(), list() packed_text_ids, packed_text_indexes = list(), list() packed_seqlens, packed_position_ids, packed_indexes = list(), list(), list() packed_key_value_indexes = list() _curr = curr = 0 newlens, new_rope = list(), list() for vit_token, curr_kvlen, curr_position_id in zip(vit_tokens, curr_kvlens, curr_rope): packed_key_value_indexes.extend(range(curr, curr + curr_kvlen)) curr += curr_kvlen packed_text_ids.append(new_token_ids["start_of_image"]) packed_text_indexes.append(_curr) packed_indexes.append(curr) curr += 1 _curr += 1 packed_vit_tokens.append(vit_token) num_img_tokens = len(vit_tokens[0]) // 4 vit_token_seqlens.append(num_img_tokens) packed_vit_token_indexes.extend(range(_curr, _curr + num_img_tokens)) packed_indexes.extend(range(curr, curr + num_img_tokens)) curr += num_img_tokens _curr += num_img_tokens packed_text_ids.append(new_token_ids['end_of_image']) packed_text_indexes.append(_curr) packed_indexes.append(curr) curr += 1 _curr += 1 packed_position_ids.extend([curr_position_id] * (num_img_tokens + 2)) packed_seqlens.append(num_img_tokens + 2) newlens.append(curr_kvlen + num_img_tokens + 2) new_rope.append(curr_position_id + 1) generation_input = { "packed_text_ids": torch.tensor(packed_text_ids, dtype=torch.long, device=device), "packed_text_indexes": torch.tensor(packed_text_indexes, dtype=torch.long, device=device), "vit_token_seqlens": torch.tensor(vit_token_seqlens, dtype=torch.int, device=device), "packed_vit_tokens": torch.cat(packed_vit_tokens, dim=0).to(device), "packed_vit_token_indexes": torch.tensor(packed_vit_token_indexes, dtype=torch.long, device=device), "packed_position_ids": torch.tensor(packed_position_ids, dtype=torch.long, device=device), "packed_seqlens": torch.tensor(packed_seqlens, dtype=torch.int, device=device), "packed_indexes": torch.tensor(packed_indexes, dtype=torch.long, device=device), "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long, device=device), "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int, device=device), } return generation_input, newlens, new_rope @torch.no_grad() def forward_cache_update_vit_validation( self, past_key_values: NaiveCache, vit_vae_video_grid_thw: torch.IntTensor, packed_text_ids: torch.LongTensor, packed_text_indexes: torch.LongTensor, packed_vit_tokens: torch.Tensor, packed_vit_token_indexes: torch.LongTensor, vit_token_seqlens: torch.IntTensor, packed_position_ids: torch.LongTensor, packed_seqlens: torch.IntTensor, packed_indexes: torch.LongTensor, packed_key_value_indexes: torch.LongTensor, key_values_lens: torch.IntTensor, device: torch.device = None, dtype: torch.dtype = None, ): packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids).to(dtype) packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size), dtype = dtype) packed_sequence[packed_text_indexes] = packed_text_embedding if self.vit_type in ["qwen2_5_vl", "qwen_2_5_vl_original"]: packed_vit_token_embed = self.vit_model( hidden_states=packed_vit_tokens, grid_thw=vit_vae_video_grid_thw, ) if self.vit_type in ["qwen2_5_vl"]: packed_vit_token_embed = self.connector(packed_vit_token_embed).to(dtype) packed_sequence[packed_vit_token_indexes] = packed_vit_token_embed else: raise NotImplementedError(f"{self.vit_type} is not supported") extra_inputs = {} if self.use_moe: extra_inputs = {"mode": "und"} output = self.language_model.forward_inference( packed_query_sequence=packed_sequence, query_lens=packed_seqlens, packed_query_position_ids=packed_position_ids, packed_query_indexes=packed_indexes, past_key_values=past_key_values, packed_key_value_indexes=packed_key_value_indexes, key_values_lens=key_values_lens, update_past_key_values=True, is_causal=False, **extra_inputs, ) past_key_values = output.past_key_values return past_key_values def prepare_start_tokens(self, curr_kvlens, curr_rope, new_token_ids, device): packed_start_tokens, packed_key_value_indexes = list(), list() packed_query_position_ids = list() curr = 0 for curr_kvlen, curr_position_id in zip(curr_kvlens, curr_rope): packed_key_value_indexes.extend(range(curr, curr + curr_kvlen)) packed_start_tokens.append(new_token_ids["bos_token_id"]) packed_query_position_ids.append(curr_position_id) curr += curr_kvlen generation_input = { "packed_start_tokens": torch.tensor(packed_start_tokens, dtype=torch.long).to(device), "packed_query_position_ids": torch.tensor(packed_query_position_ids, dtype=torch.long).to(device), "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int).to(device), "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long).to(device), } return generation_input @torch.no_grad() def generate_text( self, past_key_values: NaiveCache, packed_key_value_indexes: torch.LongTensor, key_values_lens: torch.IntTensor, packed_start_tokens: torch.LongTensor, packed_query_position_ids: torch.LongTensor, max_length: int, do_sample: bool = False, temperature: float = 1.0, end_token_id: int = None, vocab_size: int = None, ): step = 0 generated_sequence = [] curr_tokens = packed_start_tokens while step < max_length: generated_sequence.append(curr_tokens) packed_text_embedding = self.language_model.model.embed_tokens(curr_tokens) query_lens = torch.ones_like(curr_tokens) packed_query_indexes = torch.cumsum(key_values_lens, dim=0) + torch.arange(0, len(key_values_lens), device=key_values_lens.device, dtype=key_values_lens.dtype) uppacked = list(packed_key_value_indexes.split(key_values_lens.tolist(), dim=0)) for i in range(len(uppacked)): uppacked[i] += i packed_key_value_indexes = torch.cat(uppacked, dim=0) extra_inputs = {} if self.use_moe: extra_inputs = {"mode": "und"} output = self.language_model.forward_inference( packed_query_sequence=packed_text_embedding, query_lens=query_lens, packed_query_position_ids=packed_query_position_ids, packed_query_indexes=packed_query_indexes, past_key_values=past_key_values, key_values_lens=key_values_lens, packed_key_value_indexes=packed_key_value_indexes, update_past_key_values=True, is_causal=True, **extra_inputs, ) past_key_values = output.past_key_values packed_query_sequence = output.packed_query_sequence pred_logits = self.language_model.lm_head(packed_query_sequence) pred_logits[:, vocab_size:] = float('-inf') # ++ if do_sample: probs = nn.functional.softmax(pred_logits / temperature, dim=-1) curr_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) else: curr_tokens = torch.argmax(pred_logits, dim=-1) uppacked = list(packed_key_value_indexes.split(key_values_lens.tolist(), dim=0)) for i in range(len(uppacked)): uppacked[i] = torch.cat([uppacked[i], torch.tensor([uppacked[i][-1] + 1], device=uppacked[i].device)], dim=0) packed_key_value_indexes = torch.cat(uppacked, dim=0) key_values_lens = key_values_lens + 1 packed_query_position_ids = packed_query_position_ids + 1 step += 1 if end_token_id is not None and curr_tokens[0].item() == end_token_id: generated_sequence.append(curr_tokens) break output_device = generated_sequence[0].device return torch.stack([i.to(output_device) for i in generated_sequence], dim=0) def init_gen_context(self, device: torch.device, dtype: torch.dtype): gen_context = { 'kv_lens': torch.tensor([0], device=device, dtype=dtype), 'past_key_values': NaiveCache(self.config.llm_config.num_hidden_layers), } return gen_context @torch.no_grad() def validation_gen_KVcache( self, val_packed_text_ids: torch.LongTensor, val_packed_text_indexes: torch.LongTensor, val_packed_vit_tokens: torch.LongTensor, val_packed_vit_token_indexes: torch.LongTensor, val_sample_lens: List[int], val_packed_position_ids: torch.LongTensor, val_split_lens: List[int] = None, val_attn_modes: List[str] = None, val_sample_N_target: List[int] = None, vit_video_grid_thw: Optional[torch.IntTensor] = None, # NOTE: used only for TI2I. vae_video_grid_thw: Optional[torch.IntTensor] = None, video_grid_thw: Optional[torch.IntTensor] = None, val_mse_loss_indexes: Optional[torch.BoolTensor] = None, val_packed_vae_token_indexes: Optional[torch.LongTensor] = None, val_padded_latent: Optional[torch.Tensor] = None, sample_task: Optional[torch.LongTensor] = None, sample_modality: Optional[torch.LongTensor] = None, video_sizes: List[Tuple[int, int, int]] = [[1, 256, 256]], val_padded_videos: torch.Tensor = None, timestep_shift: float = 4.0, num_timesteps: int = 24, cfg_interval: Optional[Tuple[float, float]] = [0, 1], cfg_renorm_min: float = 0.0, cfg_renorm_type: str = "global", cfg_text_scale: float = 1.0, cfg_vit_scale: float = 1.0, device=None, dtype=None, new_token_ids=None, BLOCK_SIZE: int = 128, apply_chat_template: bool = False, apply_qwen_2_5_vl_pos_emb: bool = False, image_token_id: int = 151655, caption: Optional[List[str]] = None, index: str = "", **kwargs, ): cfg_vision_scale = cfg_vit_scale pt, ph, pw = self.latent_patch_size index_dtype = val_packed_text_ids.dtype cu_sample_lens = torch.nn.functional.pad(torch.cumsum(torch.tensor(val_sample_lens, device=device), dim=0), (1, 0)) sample_splits = map_splits_to_samples(val_sample_lens, val_split_lens) if val_packed_vit_tokens is not None and vit_video_grid_thw is not None: vit_sample_len = vit_video_grid_thw[:, 0] * vit_video_grid_thw[:, 1] * vit_video_grid_thw[:, 2] # shape: (N,) , N = 1 * 16 * 16, cu_vit_sample_lens = torch.cat([torch.zeros(1, device=vit_video_grid_thw.device, dtype=vit_sample_len.dtype), vit_sample_len.cumsum(0)]) self.vit_model = self.vit_model.to(device=device, dtype=dtype) val_packed_vit_tokens = torch.cat(val_packed_vit_tokens, dim=0) x_t_all = [] max_samples = kwargs.get("max_samples", 16) L = max(len(val_sample_lens) - 1, 1) max_samples = min(L, max_samples) gen_idx = 0 curr_vae_split_idx, curr_vit_split_idx = 0, 0 padded_videos = [] for i_sample in range(L): # fix: need -1. left, right = sample_splits[i_sample][0], sample_splits[i_sample][-1] + 1 current_split_lens = val_split_lens[left:right] current_attn_modes = val_attn_modes[left:right] N_target = val_sample_N_target[i_sample] N_noise_element = current_attn_modes.count("noise") + current_attn_modes.count("full_noise") + current_attn_modes.count("full_noise_target") N_vit_split = current_attn_modes.count("full") if right > len(val_attn_modes): break if N_noise_element<=0: curr_vit_split_idx += N_vit_split continue if gen_idx >= max_samples: break # 1. Get slice information for the current sample in the batch. sample_start_idx = cu_sample_lens[i_sample] sample_end_idx = cu_sample_lens[i_sample + 1] current_seq_len = val_sample_lens[i_sample] current_pos_ids = val_packed_position_ids[sample_start_idx:sample_end_idx] i_sample_task = sample_task[sample_start_idx:sample_end_idx] i_sample_modality = sample_modality[sample_start_idx:sample_end_idx] # --- Visual feature embeddings --- vae_mask = (val_packed_vae_token_indexes >= sample_start_idx) & (val_packed_vae_token_indexes < sample_end_idx) current_vae_token_indexes_local = val_packed_vae_token_indexes[vae_mask] - sample_start_idx # --- VAE MSE token part: indices of the positions in x_t that need to be updated --- vae_mse_mask = (val_mse_loss_indexes >= sample_start_idx) & (val_mse_loss_indexes < sample_end_idx) current_vae_mse_indexes_local = val_mse_loss_indexes[vae_mse_mask] - sample_start_idx # Indices of x_t positions that need updates. current_vae_mse_indexes_local_in_vae = ( current_vae_mse_indexes_local - current_vae_mse_indexes_local[0] + torch.where(current_vae_token_indexes_local == current_vae_mse_indexes_local[0])[0] ) num_vid_tokens_list, vid_shape_list, vae_position_ids, curr_padded_latent = [], [], [], [] # 2. Generate VIT unconditional features (optional). cfg_vision_pro = False if cfg_vision_scale > 1.0 and "full" in current_attn_modes: cfg_vision_pro = True vision_uncond_mask = i_sample_modality <= 1 _, vision_uncond_pos_ids, _ = self.uncond_split_pro_kvcache(vision_uncond_mask, current_text_ids, device, dtype, apply_qwen_2_5_vl_pos_emb, grid_thw_rope = grid_thw_rope[-N_target:], current_attn_modes=current_attn_modes, current_split_lens=current_split_lens, i_sample_task=i_sample_task, i_sample_modality=i_sample_modality ) # NOTE: grid_thw_rope excludes VIT/VAE condition entries. for i_target in range(N_noise_element): T, H, W = video_sizes[curr_vae_split_idx] t = (T - 1) // self.latent_downsample_temporal + 1 h = H // self.latent_downsample_spatial w = W // self.latent_downsample_spatial vid_shape_list.append([t, h, w]) num_vid_tokens_list.append(t * h * w) # Prepare packed_vae_position_ids vae_position_ids.append( get_flattened_position_ids_extrapolate_video(t, h, w, max_latent_size=self.max_latent_size) # Patch size is 1 in latent space. # NOT USED during extrapolation. ) if len(current_vae_mse_indexes_local) != len(current_vae_token_indexes_local): padded_latent_ = val_padded_latent[curr_vae_split_idx] patches = rearrange(padded_latent_, "(t pt) (h ph) (w pw) c -> (t h w) (pt ph pw c)", t=t, pt=pt, h=h, ph=ph, w=w, pw=pw) curr_padded_latent.append(patches) if val_padded_videos is not None: padded_videos.append(val_padded_videos[curr_vae_split_idx]) curr_vae_split_idx += 1 num_vid_tokens = sum(num_vid_tokens_list) vae_position_ids = torch.cat(vae_position_ids, 0) if curr_padded_latent != []: curr_padded_latent = torch.cat(curr_padded_latent, dim=0).to(dtype) # 2. Rebuild the input sequence and attention mask for the current sample. current_sequence = torch.zeros((current_seq_len, self.hidden_size), device=device, dtype=dtype) # --- Text part --- text_mask = (val_packed_text_indexes >= sample_start_idx) & (val_packed_text_indexes < sample_end_idx) current_text_indexes_local = val_packed_text_indexes[text_mask] - sample_start_idx current_text_ids = val_packed_text_ids[sample_start_idx:sample_end_idx] current_text_embedding = self.language_model.model.embed_tokens(current_text_ids).to(dtype=dtype) current_sequence[current_text_indexes_local] = current_text_embedding[current_text_indexes_local] # --- VIT part: supports TI2I --- if N_vit_split != 0: vit_sample_start_idx = cu_vit_sample_lens[curr_vit_split_idx] vit_sample_end_idx = cu_vit_sample_lens[curr_vit_split_idx + N_vit_split] current_val_packed_vit_tokens = val_packed_vit_tokens[vit_sample_start_idx:vit_sample_end_idx].to(dtype) current_val_vit_video_grid_thw = vit_video_grid_thw[curr_vit_split_idx : curr_vit_split_idx + N_vit_split] curr_vit_split_idx += N_vit_split if self.vit_type in ["qwen2_5_vl", "qwen_2_5_vl_original"]: packed_vit_token_embed = self.vit_model(hidden_states=current_val_packed_vit_tokens, grid_thw=current_val_vit_video_grid_thw) if self.vit_type in ["qwen2_5_vl"]: packed_vit_token_embed = self.connector(packed_vit_token_embed).to(dtype) else: raise NotImplementedError(f"{self.vit_type} is not supported") vit_mask = (val_packed_vit_token_indexes >= sample_start_idx) & (val_packed_vit_token_indexes < sample_end_idx) current_vit_indexes_local = val_packed_vit_token_indexes[vit_mask] - sample_start_idx current_sequence[current_vit_indexes_local] = packed_vit_token_embed # --- Keep input, mask, and length aligned with training by padding to a multiple of BLOCK_SIZE --- current_seq_len_pad = (current_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE * BLOCK_SIZE current_pad = current_seq_len_pad - current_seq_len if current_pad > 0: current_split_lens = current_split_lens + [current_pad] current_attn_modes = current_attn_modes + ["causal"] validation_noise_seed = kwargs.get("validation_noise_seed", -1) if validation_noise_seed > 0: generator = torch.Generator(device=device).manual_seed(validation_noise_seed + get_global_rank() * max_samples + i_sample) else: generator = None x_t = torch.randn(num_vid_tokens, self.patch_latent_dim, generator=generator, device=device, dtype=dtype) if curr_padded_latent != []: curr_padded_latent[current_vae_mse_indexes_local_in_vae] = x_t[current_vae_mse_indexes_local_in_vae] x_t = curr_padded_latent timesteps = torch.linspace(1, 0, num_timesteps + 1, device=x_t.device) timesteps = timestep_shift * timesteps / (1 + (timestep_shift - 1) * timesteps) dts = timesteps[:-1] - timesteps[1:] timesteps = timesteps[:-1] if apply_qwen_2_5_vl_pos_emb: grid_thw_rope = video_grid_thw[i_sample] current_pos_ids, _ = self.language_model.get_rope_index( input_ids=current_text_ids.unsqueeze(0), image_grid_thw=grid_thw_rope, video_grid_thw=grid_thw_rope, second_per_grid_ts=[1.0]*len(grid_thw_rope), attention_mask=torch.ones([1, len(current_text_ids)], dtype=torch.long, device=device), ) current_pos_ids = shift_position_ids(current_pos_ids, pos_shift = 1000, attn_modes = current_attn_modes, split_lens = current_split_lens, shift_attn_mode=['full_noise',"full"], pro_type = 10, i_sample_task=i_sample_task, i_sample_modality=i_sample_modality) if cfg_text_scale > 1.0: uncond_mask = i_sample_modality!=0 _, uncond_pos_ids, _ = self.uncond_split_pro_kvcache(uncond_mask, current_text_ids, device, dtype, apply_qwen_2_5_vl_pos_emb, grid_thw_rope = grid_thw_rope, current_attn_modes=current_attn_modes, current_split_lens=current_split_lens, i_sample_task=i_sample_task, i_sample_modality=i_sample_modality) extra_inputs = {} if self.use_moe: if N_vit_split != 0: packed_und_token_indexes = torch.cat([current_text_indexes_local, current_vit_indexes_local], dim=0) else: packed_und_token_indexes = current_text_indexes_local extra_inputs.update( packed_und_token_indexes=packed_und_token_indexes.to(dtype=index_dtype), packed_gen_token_indexes=current_vae_token_indexes_local.to(dtype=index_dtype), ) timestep = torch.zeros(x_t.shape[0], device=x_t.device) timestep[current_vae_mse_indexes_local_in_vae] = torch.tensor([1.] * current_vae_mse_indexes_local_in_vae.shape[0], device=x_t.device) # --- Store visual feature encodings (VAE condition) --- timestep_embed = self.time_embedder(timestep) latent_pos_embed = self.latent_pos_embed(vae_position_ids) vae_embed = self.vae2llm(x_t) + timestep_embed + latent_pos_embed vae_embed = vae_embed.to(current_sequence.dtype) current_sequence[current_vae_token_indexes_local] = vae_embed # For kv cache gen_context = self.init_gen_context(device=device, dtype=torch.int32) # gen_context initializes kv_lens, ropes, and past_key_values. cfg_text_context = deepcopy(gen_context) cfg_vision_context = deepcopy(gen_context ) current_cond_start, current_cond_end = 0, 0 self.language_model.eval() self.eval() for i_attn_mode_, current_cond_len in zip(current_attn_modes, current_split_lens): current_cond_end += current_cond_len if i_attn_mode_ == "noise": vae_in_packed_sequence_index = torch.arange(current_cond_start, current_cond_end, dtype=torch.long, device=device) packed_seqlens_vae = current_cond_len target_packed_vae_token_indexes = torch.arange(1, current_cond_len-1, dtype=torch.long, device=device) target_packed_text_indexes = torch.tensor([0, current_cond_len-1], dtype=torch.long, device=device) break if i_attn_mode_ == 'causal': is_causal = True else: is_causal = False gen_context = self.update_gen_context(current_sequence, current_pos_ids, gen_context, extra_inputs, current_cond_start, current_cond_end, current_cond_len, device, dtype, is_causal = is_causal) if cfg_text_scale > 1.0 and i_sample_modality[current_cond_start] != 0: cfg_text_context = self.update_gen_context(current_sequence, current_pos_ids, cfg_text_context, extra_inputs, current_cond_start, current_cond_end, current_cond_len, device, dtype, is_causal = is_causal) if cfg_vision_scale > 1.0 and i_sample_modality[current_cond_start] > 1: cfg_vision_context = self.update_gen_context(current_sequence, current_pos_ids, cfg_vision_context, extra_inputs, current_cond_start, current_cond_end, current_cond_len, device, dtype, is_causal = is_causal) current_cond_start = current_cond_end for _ in range(1): timestep = torch.zeros(x_t.shape[0], device=x_t.device) for i, timestep_ in enumerate(timesteps): timestep[current_vae_mse_indexes_local_in_vae] = torch.tensor([timestep_] * current_vae_mse_indexes_local_in_vae.shape[0], device=x_t.device) if timestep_ > cfg_interval[0] and timestep_ <= cfg_interval[1]: cfg_text_scale_ = cfg_text_scale cfg_vision_scale_ = cfg_vision_scale else: cfg_text_scale_ = 1.0 cfg_vision_scale_ = 1.0 # --- Visual feature encoding --- timestep_embed = self.time_embedder(timestep) latent_pos_embed = self.latent_pos_embed(vae_position_ids) vae_embed = self.vae2llm(x_t) + timestep_embed + latent_pos_embed vae_embed = vae_embed.to(current_sequence.dtype) current_sequence[current_vae_token_indexes_local] = vae_embed packed_sequence_vae = current_sequence[vae_in_packed_sequence_index] extra_inputs_vae = {} if self.use_moe: extra_inputs_vae = {"mode": "gen", "packed_vae_token_indexes": target_packed_vae_token_indexes, "packed_text_indexes": target_packed_text_indexes} v_t_output = self.language_model.forward_inference( packed_query_sequence=packed_sequence_vae, query_lens=torch.tensor([packed_seqlens_vae],dtype=torch.int32, device=device), packed_query_position_ids=current_pos_ids[:, :, current_cond_start:current_cond_end], packed_query_indexes=vae_in_packed_sequence_index, past_key_values=gen_context['past_key_values'], key_values_lens=gen_context['kv_lens'], packed_key_value_indexes=torch.arange(0,gen_context['kv_lens'][0], dtype=torch.int64, device=device), update_past_key_values=False, is_causal=False, **extra_inputs_vae, ) v_t = self.llm2vae(v_t_output.packed_query_sequence) v_t = v_t[target_packed_vae_token_indexes] # --- Apply CFG --- if cfg_text_scale_ > 1.0: cfg_text_output = self.language_model.forward_inference( packed_query_sequence=packed_sequence_vae, query_lens=torch.tensor([packed_seqlens_vae],dtype=torch.int32, device=device), packed_query_position_ids=uncond_pos_ids[:,:,cfg_text_context['kv_lens'][0]:cfg_text_context['kv_lens'][0]+packed_seqlens_vae], packed_query_indexes=vae_in_packed_sequence_index - sum(i_sample_modality==0), past_key_values=cfg_text_context['past_key_values'], key_values_lens=cfg_text_context['kv_lens'], packed_key_value_indexes=torch.arange(0,cfg_text_context['kv_lens'][0], dtype=torch.int64, device=device), update_past_key_values=False, is_causal=False, **extra_inputs_vae, ) cfg_text_v_t = self.llm2vae(cfg_text_output.packed_query_sequence) cfg_text_v_t = cfg_text_v_t[target_packed_vae_token_indexes] if cfg_vision_pro: cfg_vision_output = self.language_model.forward_inference( packed_query_sequence=packed_sequence_vae, query_lens=torch.tensor([packed_seqlens_vae],dtype=torch.int32, device=device), packed_query_position_ids=vision_uncond_pos_ids[:,:,cfg_vision_context['kv_lens'][0]:cfg_vision_context['kv_lens'][0]+packed_seqlens_vae], packed_query_indexes=vae_in_packed_sequence_index - sum(i_sample_modality==4), past_key_values=cfg_vision_context['past_key_values'], key_values_lens=cfg_vision_context['kv_lens'], packed_key_value_indexes=torch.arange(0,cfg_vision_context['kv_lens'][0], dtype=torch.int64, device=device), update_past_key_values=False, is_causal=False, **extra_inputs_vae, ) cfg_text_vision_v_t = self.llm2vae(cfg_vision_output.packed_query_sequence) cfg_text_vision_v_t = cfg_text_vision_v_t[target_packed_vae_token_indexes] v_t_ = cfg_text_vision_v_t + cfg_text_scale_ * (v_t - cfg_text_v_t) + cfg_vision_scale_ * (cfg_text_v_t - cfg_text_vision_v_t) else: v_t_ = cfg_text_v_t + cfg_text_scale_ * (v_t - cfg_text_v_t) if cfg_renorm_type == "global": norm_v_t = torch.norm(v_t) norm_v_t_ = torch.norm(v_t_) scale = (norm_v_t / (norm_v_t_ + 1e-8)).clamp(min=cfg_renorm_min, max=1.0) elif cfg_renorm_type == "channel": norm_v_t = torch.norm(v_t, dim=-1, keepdim=True) norm_v_t_ = torch.norm(v_t_, dim=-1, keepdim=True) scale = (norm_v_t / (norm_v_t_ + 1e-8)).clamp(min=cfg_renorm_min, max=1.0) elif cfg_renorm_type.lower() in ("", "none", "null"): scale = 1 else: raise NotImplementedError(f"{cfg_renorm_type} is not suppoprted") v_t = v_t_ * scale x_t[current_vae_mse_indexes_local_in_vae] = x_t[current_vae_mse_indexes_local_in_vae] - v_t.to(x_t.device) * dts[i] # velocity pointing from data to noise # ---- Reshape each sample independently to [T,H,W,C], avoiding use of the last sample's t/h/w for the whole batch ---- curr_seq_target, patch = 0, [] for i_target in range(N_noise_element): pt, ph, pw = self.latent_patch_size t, h, w = vid_shape_list[i_target] len_target = t * h * w x_t_ = rearrange(x_t[curr_seq_target : curr_seq_target + len_target], "(t h w) (pt ph pw c) -> (t pt) (h ph) (w pw) c", t=t, h=h, w=w, pt=pt, ph=ph, pw=pw) patch.append(x_t_) curr_seq_target += len_target x_t_all.append(patch) gen_idx += 1 if caption != None: return x_t_all, [caption], padded_videos, index return x_t_all def get_uncond_attn_modes_split_lens(self, current_attn_modes, current_split_lens, uncond_mask): # Filter unconditional sample parts according to uncond_mask. curr = 0 uncond_attn_modes, uncond_split_lens = [], [] for i, split_len in enumerate(current_split_lens): mask_slice = uncond_mask[curr:curr+split_len] if mask_slice.all(): uncond_attn_modes.append(current_attn_modes[i]) uncond_split_lens.append(split_len) curr += split_len return uncond_attn_modes, uncond_split_lens def uncond_split_pro_kvcache( self, uncond_mask, current_text_ids, device, dtype, apply_qwen_2_5_vl_pos_emb=False, uncond_pos_ids=None, grid_thw_rope=None, current_attn_modes=None, current_split_lens=None, i_sample_task=None, i_sample_modality=None, ): # Extract text ids for the unconditional sequence. uncond_text_ids = current_text_ids[uncond_mask] uncond_seq_len = len(uncond_text_ids) if apply_qwen_2_5_vl_pos_emb: uncond_pos_ids, uncond_rope_deltas = self.language_model.get_rope_index( input_ids=uncond_text_ids.unsqueeze(0), image_grid_thw=grid_thw_rope, video_grid_thw=grid_thw_rope, second_per_grid_ts=[1.0] * len(grid_thw_rope), attention_mask=torch.ones([1, len(uncond_text_ids)], dtype=torch.long, device=device), ) uncond_attn_modes, uncond_split_lens = self.get_uncond_attn_modes_split_lens( current_attn_modes, current_split_lens, uncond_mask) i_sample_task = i_sample_task[uncond_mask] i_sample_modality = i_sample_modality[uncond_mask] uncond_pos_ids = shift_position_ids(uncond_pos_ids, pos_shift = 1000, attn_modes = uncond_attn_modes, split_lens = uncond_split_lens, shift_attn_mode=['full_noise',"full"], pro_type = 10, i_sample_task=i_sample_task, i_sample_modality=i_sample_modality) else: uncond_pos_ids = torch.tensor(uncond_pos_ids, dtype=torch.long, device=device)[:uncond_seq_len] return ( uncond_text_ids, uncond_pos_ids, uncond_seq_len, ) def update_gen_context(self, current_sequence, current_pos_ids, gen_context, extra_inputs, current_cond_start, current_cond_end, current_cond_len, device, dtype, is_causal = True): extra_inputs_cond = {} extra_inputs_gen_mask = (extra_inputs["packed_gen_token_indexes"] >= current_cond_start) & (extra_inputs["packed_gen_token_indexes"] < current_cond_end) extra_inputs_cond["packed_vae_token_indexes"] = extra_inputs["packed_gen_token_indexes"][extra_inputs_gen_mask] - gen_context['kv_lens'] extra_inputs_und_mask = (extra_inputs["packed_und_token_indexes"] >= current_cond_start) & (extra_inputs["packed_und_token_indexes"] < current_cond_end) extra_inputs_cond["packed_text_indexes"] = extra_inputs["packed_und_token_indexes"][extra_inputs_und_mask] - gen_context['kv_lens'] if extra_inputs_cond["packed_vae_token_indexes"].shape[0] > 0 : mode_ = "gen" else: mode_ = "und" output = self.language_model.forward_inference( packed_query_sequence=current_sequence[current_cond_start:current_cond_end], query_lens=torch.tensor([current_cond_len],dtype=torch.int32, device=device), packed_query_position_ids=current_pos_ids[:, :, current_cond_start:current_cond_end], packed_query_indexes=torch.arange(gen_context['kv_lens'][0],gen_context['kv_lens'][0] + current_cond_len, dtype=torch.long, device=device), # Positions for the current new input. past_key_values=gen_context['past_key_values'], packed_key_value_indexes=torch.arange(0,gen_context['kv_lens'][0], dtype=torch.int64, device=device), # Positions for the past KV cache. key_values_lens=gen_context['kv_lens'], update_past_key_values=True, is_causal=is_causal, mode = mode_, **extra_inputs_cond ) gen_context['past_key_values'] = output.past_key_values gen_context['kv_lens'] += current_cond_len return gen_context