code/llava_arch.py

"""
LLaVA Architecture with Integrated Mask Prediction for Image Editing

This module contains:
- LlavaMetaModel: Base model with vision tower, diffusion components, and mask prediction
- LlavaMetaForCausalLM: Mixin for causal LM with multimodal support
- MaskPredictor: Predicts edit regions from LLM hidden states
- BF16SafeLayerNorm: Numerically stable LayerNorm for BF16 training

Key Innovation: MaskPredictor enables mask-free inference by learning to predict
edit regions from LLM understanding, eliminating the need for external segmentation.
"""

from abc import ABC, abstractmethod
from typing import Optional, Tuple, List
import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers import FlowMatchEulerDiscreteScheduler, DPMSolverMultistepScheduler
from diffusers.models.normalization import RMSNorm

from .mobile_block import MobileConditioningProjector
from .multimodal_llava_encoder.builder import build_vision_tower
from .multimodal_llava_projector.builder import build_vision_projector
from .multimodal_projector.builder import build_down_projector
from .multimodal_decoder.builder import build_vae, build_sana
from blip3o.constants import (
    DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, 
    DEFAULT_IMAGE_PATCH_TOKEN, IGNORE_INDEX, IMAGE_TOKEN_INDEX
)


# ============================================================
# BF16-Safe LayerNorm
# ============================================================

class BF16SafeLayerNorm(nn.Module):
    """
    LayerNorm that's safe for BF16 training.
    Performs normalization in float32 for numerical stability.
    """
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.bias = nn.Parameter(torch.zeros(hidden_size))
        self.eps = eps
        self.hidden_size = hidden_size

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        input_dtype = x.dtype
        x = x.float()
        mean = x.mean(-1, keepdim=True)
        variance = (x - mean).pow(2).mean(-1, keepdim=True)
        x = (x - mean) / torch.sqrt(variance + self.eps)
        x = self.weight.float() * x + self.bias.float()
        return x.to(input_dtype)

    def reset_parameters(self):
        nn.init.ones_(self.weight)
        nn.init.zeros_(self.bias)


# ============================================================
# Mask Predictor - Enables Mask-Free Inference
# ============================================================

class MaskPredictor(nn.Module):
    """
    Predicts edit mask from LLM hidden states.
    
    This is the KEY component that enables mask-free inference.
    During training: Supervised by SAM-generated masks
    During inference: Predicts mask directly from LLM understanding
    
    Architecture:
    1. Attention pooling to focus on instruction-relevant tokens
    2. Project to spatial features
    3. Decode to mask
    """

    def __init__(self, hidden_size: int, latent_channels: int, latent_size: int = 32):
        super().__init__()
        self.latent_size = latent_size
        self.hidden_size = hidden_size

        # Attention pooling to focus on instruction-relevant tokens
        self.attention_pool = nn.Sequential(
            nn.Linear(hidden_size, hidden_size // 4),
            nn.Tanh(),
            nn.Linear(hidden_size // 4, 1),
        )

        # Layer norm for stability
        self.input_norm = BF16SafeLayerNorm(hidden_size)

        # Project pooled features to spatial representation
        intermediate_size = hidden_size // 2
        spatial_dim = latent_size * latent_size * 64

        self.hidden_proj = nn.Sequential(
            nn.Linear(hidden_size, intermediate_size),
            nn.LayerNorm(intermediate_size),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(intermediate_size, intermediate_size),
            nn.LayerNorm(intermediate_size),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(intermediate_size, spatial_dim),
        )

        # Decode to mask with sufficient capacity
        self.mask_decoder = nn.Sequential(
            nn.Conv2d(64, 256, 3, padding=1),
            nn.GroupNorm(32, 256),
            nn.GELU(),
            nn.Conv2d(256, 128, 3, padding=1),
            nn.GroupNorm(16, 128),
            nn.GELU(),
            nn.Conv2d(128, 64, 3, padding=1),
            nn.GroupNorm(8, 64),
            nn.GELU(),
            nn.Conv2d(64, 1, 1),
        )

        self._init_weights()

    def _init_weights(self):
        """Initialize weights for stable training."""
        # Initialize attention pooling
        for module in self.attention_pool:
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight, gain=0.1)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)

        # Initialize LayerNorm
        self.input_norm.reset_parameters()

        # Initialize projection layers
        for module in self.hidden_proj:
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight, gain=0.1)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.LayerNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)

        # Initialize conv layers
        for module in self.mask_decoder:
            if isinstance(module, nn.Conv2d):
                nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.GroupNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)

        # Initialize final layer with small weights for stable start
        for module in reversed(list(self.mask_decoder)):
            if isinstance(module, nn.Conv2d):
                nn.init.normal_(module.weight, mean=0.0, std=0.01)
                nn.init.zeros_(module.bias)
                break

    def forward(self, hidden_states: torch.Tensor, return_logits: bool = False) -> torch.Tensor:
        """
        Predict edit mask from LLM hidden states.

        Args:
            hidden_states: [B, seq_len, hidden_size] from LLM
            return_logits: If True, return logits instead of probabilities

        Returns:
            mask: [B, 1, H, W] predicted edit mask
        """
        batch_size = hidden_states.shape[0]
        device = hidden_states.device

        # Check for NaN/Inf in input
        if torch.isnan(hidden_states).any() or torch.isinf(hidden_states).any():
            if return_logits:
                return torch.zeros(batch_size, 1, self.latent_size, self.latent_size,
                                   device=device, dtype=torch.float32, requires_grad=True)
            return torch.full((batch_size, 1, self.latent_size, self.latent_size), 0.5,
                              device=device, dtype=torch.float32, requires_grad=True)

        # Normalize hidden states
        hidden_states = self.input_norm(hidden_states)

        # Get dtype from first layer
        target_dtype = self.attention_pool[0].weight.dtype
        hidden_states = hidden_states.to(target_dtype)

        # Attention pooling: learn which tokens are important
        attn_weights = self.attention_pool(hidden_states)
        attn_weights = F.softmax(attn_weights, dim=1)

        # Weighted sum of hidden states
        pooled = (hidden_states * attn_weights).sum(dim=1)

        # Project to spatial features
        spatial = self.hidden_proj(pooled)
        spatial = spatial.view(-1, 64, self.latent_size, self.latent_size)

        # Decode to mask logits
        mask_logits = self.mask_decoder(spatial)

        if return_logits:
            return mask_logits.float()

        return torch.sigmoid(mask_logits.float())


# ============================================================
# Diffusion Connector
# ============================================================

class DiffusionConnector(nn.Module):
    def __init__(self, input_dim=896, hidden_dim=1024, output_dim=2304, eps=1e-5):
        super().__init__()
        self.linear1 = nn.Linear(input_dim, hidden_dim)
        self.act = nn.GELU(approximate="tanh")
        self.linear2 = nn.Linear(hidden_dim, output_dim)
        self.norm = RMSNorm(output_dim, eps=eps, elementwise_affine=True)

        nn.init.xavier_uniform_(self.linear1.weight)
        nn.init.zeros_(self.linear1.bias)
        nn.init.xavier_uniform_(self.linear2.weight)
        nn.init.zeros_(self.linear2.bias)
        with torch.no_grad():
            self.norm.weight.fill_(math.sqrt(5.5))

    def forward(self, x):
        x = self.linear1(x)
        x = self.act(x)
        x = self.linear2(x)
        x = self.norm(x)
        return x


# ============================================================
# Mask Encoder - Encodes masks for diffusion conditioning
# ============================================================

class MaskEncoder(nn.Module):
    """Encodes binary mask into latent conditioning for diffusion."""
    
    def __init__(self, latent_channels: int = 32):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 64, 3, padding=1),
            nn.GroupNorm(8, 64),
            nn.SiLU(),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.GroupNorm(16, 128),
            nn.SiLU(),
            nn.Conv2d(128, latent_channels, 3, padding=1),
        )
        self._init_weights()

    def _init_weights(self):
        for module in self.encoder:
            if isinstance(module, nn.Conv2d):
                nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.GroupNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)
        # Last layer: small random weights, NOT zeros!
        nn.init.normal_(self.encoder[-1].weight, mean=0.0, std=0.01)
        nn.init.zeros_(self.encoder[-1].bias)

    def forward(self, mask: torch.Tensor) -> torch.Tensor:
        return self.encoder(mask.to(torch.bfloat16))


# ============================================================
# Spatial Reference Encoder
# ============================================================

class SpatialRefEncoder(nn.Module):
    """Encodes reference image latents for spatial conditioning."""
    
    def __init__(self, latent_channels: int = 32):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(latent_channels, 64, 3, padding=1),
            nn.GroupNorm(8, 64),
            nn.SiLU(),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.GroupNorm(16, 128),
            nn.SiLU(),
            nn.Conv2d(128, latent_channels, 3, padding=1),
        )
        self._init_weights()

    def _init_weights(self):
        for module in self.encoder:
            if isinstance(module, nn.Conv2d):
                nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu')
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.GroupNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)
        # Last layer: small random weights
        nn.init.normal_(self.encoder[-1].weight, mean=0.0, std=0.01)
        nn.init.zeros_(self.encoder[-1].bias)

    def forward(self, latents: torch.Tensor) -> torch.Tensor:
        return self.encoder(latents)


# ============================================================
# LlavaMetaModel - Base Model with All Components
# ============================================================

class LlavaMetaModel:
    """
    Base model containing:
    - Vision tower for image understanding
    - DiT for diffusion generation
    - VAE for latent encoding/decoding
    - MaskPredictor for edit region prediction
    - MaskEncoder for mask conditioning
    - Conditioning weights (mask_weight, spatial_weight)
    """

    def __init__(self, config):
        super(LlavaMetaModel, self).__init__(config)

        # Vision components
        if hasattr(config, "mm_vision_tower"):
            self.vision_tower = build_vision_tower(config, delay_load=True)
            self.mm_projector = build_vision_projector(config)

        # Diffusion components
        if hasattr(config, "diffusion_name_or_path"):
            self.dit = build_sana(config)
            self.vae = build_vae(config)
            
            # Diffusion connector
            self.diffusion_connector = MobileConditioningProjector(
                input_dim=896, 
                hidden_dim=512, 
                output_dim=2304, 
                num_layers=config.vlm_num_layers
            )

            # Noise scheduler
            if getattr(config, 'is_train', False):
                print("Using FlowMatchEulerDiscreteScheduler for training")
                self.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
                    config.diffusion_name_or_path, subfolder="scheduler"
                )
            else:
                print("Using DPMSolverMultistepScheduler for inference")
                self.noise_scheduler = DPMSolverMultistepScheduler.from_pretrained(
                    config.diffusion_name_or_path, subfolder="scheduler"
                )

        # Get latent config
        latent_channels = getattr(config, 'latent_channels', 32)
        latent_size = getattr(config, 'latent_size', 32)

        # ============================================================
        # Mask Prediction Components (for image editing)
        # ============================================================
        
        # Mask predictor: predicts edit region from LLM hidden states
        if getattr(config, 'use_mask_predictor', True):
            self.mask_predictor = MaskPredictor(
                hidden_size=config.hidden_size,
                latent_channels=latent_channels,
                latent_size=latent_size
            )
        else:
            self.mask_predictor = None

        # Mask encoder: encodes mask for diffusion conditioning
        if getattr(config, 'use_mask_conditioning', True):
            self.mask_encoder = MaskEncoder(latent_channels=latent_channels)
            # CRITICAL: This is inside self (LlavaMetaModel), so it gets saved!
            self.mask_weight = nn.Parameter(torch.tensor(1.0))
        else:
            self.mask_encoder = None
            self.mask_weight = None

        # Spatial reference encoder
        if getattr(config, 'use_spatial_conditioning', False):
            self.spatial_ref_encoder = SpatialRefEncoder(latent_channels=latent_channels)
            self.spatial_weight = nn.Parameter(torch.tensor(0.5))
        else:
            self.spatial_ref_encoder = None
            self.spatial_weight = None

        # Operation embedding for edit type
        if getattr(config, 'use_operation_embedding', False):
            num_operations = getattr(config, 'num_operation_types', 10)
            self.operation_embedding = nn.Embedding(num_operations, latent_channels)
        else:
            self.operation_embedding = None

    def get_vision_tower(self):
        vision_tower = getattr(self, 'vision_tower', None)
        if type(vision_tower) is list:
            vision_tower = vision_tower[0]
        return vision_tower
    
    def get_sana(self):
        dit = getattr(self, 'dit', None)
        if type(dit) is list:
            dit = dit[0]
        if dit is not None:
            dit.to(self.device)
        return dit

    def get_sana_vae(self):
        vae = getattr(self, 'vae', None)
        if type(vae) is list:
            vae = vae[0]
        if vae is not None:
            vae.to(self.device)
        return vae

    def reinitialize_mask_components(self):
        """
        Reinitialize mask-related components.
        Call after loading pretrained weights if these components weren't in the original model.
        """
        print("Reinitializing mask components...")
        
        if self.mask_predictor is not None:
            self.mask_predictor._init_weights()
            print("  ✓ mask_predictor reinitialized")

        if self.mask_encoder is not None:
            self.mask_encoder._init_weights()
            print("  ✓ mask_encoder reinitialized")

        if self.spatial_ref_encoder is not None:
            self.spatial_ref_encoder._init_weights()
            print("  ✓ spatial_ref_encoder reinitialized")

        if self.mask_weight is not None:
            nn.init.ones_(self.mask_weight)
            print("  ✓ mask_weight set to 1.0")

        if self.spatial_weight is not None:
            nn.init.constant_(self.spatial_weight, 0.5)
            print("  ✓ spatial_weight set to 0.5")

        #if self.operation_embedding is not None:
        #    nn.init.normal_(self.operation_embedding.weight, mean=0.0, std=0.02)
        #    print("  ✓ operation_embedding reinitialized")

        print("Reinitialization complete!")

    def initialize_vision_modules(self, model_args, fsdp=None):
        """Initialize vision and diffusion modules."""
        mm_vision_select_layer = model_args.mm_vision_select_layer
        mm_vision_select_feature = model_args.mm_vision_select_feature
        mm_patch_merge_type = model_args.mm_patch_merge_type

        # Initialize DiT
        if self.get_sana() is None:
            dit = build_sana(model_args)
            if hasattr(model_args, "is_train"):
                if model_args.is_train:
                    print("FLOW MATCHING !!")
                    self.noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_args.diffusion_name_or_path, subfolder="scheduler")
                else:
                    print("DPM SOLVER !!")
                    self.noise_scheduler = DPMSolverMultistepScheduler.from_pretrained(model_args.diffusion_name_or_path, subfolder="scheduler")

            if fsdp is not None and len(fsdp) > 0:
                self.dit = [dit]
            else:
                self.dit = dit
        else:
            if fsdp is not None and len(fsdp) > 0:
                dit = self.dit[0]
            else:
                dit = self.dit
        for p in dit.parameters():
            p.requires_grad = False
                
        if self.get_sana_vae() is None:
            vae = build_vae(model_args)

            if fsdp is not None and len(fsdp) > 0:
                self.vae = [vae]
            else:
                self.vae = vae
        else:
            if fsdp is not None and len(fsdp) > 0:
                vae = self.vae[0]
            else:
                vae = self.vae
        for p in vae.parameters():
            p.requires_grad = False
    

        if self.get_vision_tower() is None:
            print("=" * 20, "Building vision tower", "=" * 20)
            vision_tower = build_vision_tower(model_args)
            

            if fsdp is not None and len(fsdp) > 0:
                self.vision_tower = [vision_tower]
            else:
                self.vision_tower = vision_tower
        else:
            if fsdp is not None and len(fsdp) > 0:
                vision_tower = self.vision_tower[0]
            else:
                vision_tower = self.vision_tower
            vision_tower.load_model()
        
        
        if getattr(self, 'diffusion_connector', None) is None:
            #self.diffusion_connector = DiffusionConnector(input_dim=self.config.hidden_size,hidden_dim=1024,output_dim=2304)
            self.diffusion_connector = MobileConditioningProjector(input_dim=896, hidden_dim=512, output_dim=2304, num_layers=model_args.vlm_num_layers)


            '''
            norm = RMSNorm(2304, eps=1e-5, elementwise_affine=True)
            with torch.no_grad():
                norm.weight.fill_(math.sqrt(5.5))
            self.diffusion_connector = nn.Sequential(
                nn.Linear(self.config.hidden_size, 1024),
                nn.GELU(approximate="tanh"),
                nn.Linear(1024, 2304),
                norm,
            )
            '''
        else:
            for p in self.diffusion_connector.parameters():
                p.requires_grad = True
        
        # freeze all parameters in dit except for caption_projection
        for name, param in self.dit.named_parameters():
            if "caption" in name:
                param.requires_grad = True
            else:
                param.requires_grad = False
        

        for p in dit.parameters():
            p.requires_grad = True
        for p in vision_tower.parameters():
            p.requires_grad = False
           # vision_tower().eval()

        self.config.use_mm_proj = True
        self.config.mm_projector_type = getattr(model_args, 'mm_projector_type', 'linear')
        self.config.mm_vision_select_layer = mm_vision_select_layer
        self.config.mm_vision_select_feature = mm_vision_select_feature
        self.config.mm_patch_merge_type = mm_patch_merge_type
        self.config.diffusion_name_or_path = model_args.diffusion_name_or_path
        self.config.is_train = False #model_args.is_train

        if getattr(self, 'down_projector', None) is None:
            self.down_projector = build_down_projector(self.config)
        else:
            # In case it is frozen by LoRA
            for p in self.down_projector.parameters():
                p.requires_grad = True

        




def unpad_image(tensor, original_size):
    """
    Unpads a PyTorch tensor of a padded and resized image.

    Args:
    tensor (torch.Tensor): The image tensor, assumed to be in CxHxW format.
    original_size (tuple): The original size of PIL image (width, height).

    Returns:
    torch.Tensor: The unpadded image tensor.
    """
    original_width, original_height = original_size
    current_height, current_width = tensor.shape[1:]

    original_aspect_ratio = original_width / original_height
    current_aspect_ratio = current_width / current_height

    if original_aspect_ratio > current_aspect_ratio:
        scale_factor = current_width / original_width
        new_height = int(original_height * scale_factor)
        padding = (current_height - new_height) // 2
        unpadded_tensor = tensor[:, padding:current_height - padding, :]
    else:
        scale_factor = current_height / original_height
        new_width = int(original_width * scale_factor)
        padding = (current_width - new_width) // 2
        unpadded_tensor = tensor[:, :, padding:current_width - padding]

    return unpadded_tensor


class LlavaMetaForCausalLM(ABC):

    @abstractmethod
    def get_model(self):
        pass

    def get_vision_tower(self):
        return self.get_model().get_vision_tower()
    
    def visual(self, pixel_values: torch.Tensor) -> torch.Tensor:
        image_features = self.get_model().get_vision_tower()(pixel_values)
        image_features = self.get_model().mm_projector(image_features)
        return image_features


    def get_mm_projector(self):
        return self.get_model().mm_projector


    def get_sigmas(self, timesteps, device, n_dim=4, dtype=torch.float32):
        sigmas = self.get_model().noise_scheduler.sigmas.to(device=device, dtype=dtype)
        schedule_timesteps = self.get_model().noise_scheduler.timesteps.to(device=device)
        timesteps = timesteps.to(device)
        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]

        sigma = sigmas[step_indices].flatten()
        while len(sigma.shape) < n_dim:
            sigma = sigma.unsqueeze(-1)
        return sigma

    def mask_drop(self, latents, drop_prob=0.1):
        if drop_prob <= 0:
            return latents
        mask = torch.bernoulli(torch.zeros(latents.shape[0], device=latents.device, dtype=latents.dtype) + drop_prob)
        while len(mask.shape) < len(latents.shape):
            mask = mask.unsqueeze(-1)
        mask = 1 - mask  # need to flip 0 <-> 1
        return latents * mask

    # ============================================================
    # Convenience Properties for Mask Components
    # ============================================================
    
    @property
    def mask_predictor(self):
        return getattr(self.get_model(), 'mask_predictor', None)

    @property
    def mask_encoder(self):
        return getattr(self.get_model(), 'mask_encoder', None)

    @property
    def mask_weight(self):
        return getattr(self.get_model(), 'mask_weight', None)

    @property
    def spatial_weight(self):
        return getattr(self.get_model(), 'spatial_weight', None)

    @property
    def spatial_ref_encoder(self):
        return getattr(self.get_model(), 'spatial_ref_encoder', None)

    @property
    def operation_embedding(self):
        return getattr(self.get_model(), 'operation_embedding', None)

    # ============================================================
    # Multimodal Input Preparation
    # ============================================================

    def prepare_inputs_labels_for_multimodal(
        self, input_ids, position_ids, attention_mask, past_key_values, labels,
        gen_images=None, und_images=None
    ):
        if (gen_images is None and und_images is None) or input_ids.shape[1] == 1 or self.get_vision_tower() is None:
            return input_ids, position_ids, attention_mask, past_key_values, None, labels, None, None, None
        if gen_images is not None:
            vae = self.get_model().get_sana_vae()
            vae_device = vae.device
            prompt_image_embeds = vae.encode(gen_images.to(vae_device)).latent if gen_images is not None else None
            prompt_image_embeds = prompt_image_embeds * vae.config.scaling_factor if prompt_image_embeds is not None else None
            target_image_embeds = torch.clone(prompt_image_embeds).detach()
        else:
            target_image_embeds = None
            

        images = und_images
        if type(images) is list or images.ndim == 5:
            if type(images) is list:
                images = [x.unsqueeze(0) if x.ndim == 3 else x for x in images]
            concat_images = torch.cat([image for image in images], dim=0)
            image_features = self.visual(concat_images)
            split_sizes = [image.shape[0] for image in images]
            image_features = torch.split(image_features, split_sizes, dim=0)
            image_features = [x.flatten(0, 1) for x in image_features]
        else:
            image_features = self.visual(images) # [B, image_tokens, hidden_size]


        # Let's just add dummy tensors if they do not exist,
        # it is a headache to deal with None all the time.
        # But it is not ideal, and if you have a better idea,
        # please open an issue / submit a PR, thanks.
        _labels = labels
        _position_ids = position_ids
        _attention_mask = attention_mask
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
        else:
            attention_mask = attention_mask.bool()
        if position_ids is None:
            position_ids = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=input_ids.device)
        if labels is None:
            labels = torch.full_like(input_ids, IGNORE_INDEX)

        # remove the padding using attention_mask -- FIXME
        input_ids = [cur_input_ids[cur_attention_mask] for cur_input_ids, cur_attention_mask in zip(input_ids, attention_mask)]
        labels = [cur_labels[cur_attention_mask] for cur_labels, cur_attention_mask in zip(labels, attention_mask)]

        new_input_embeds = []
        new_labels = []
        new_input_ids = []
        cur_image_idx = 0
        for batch_idx, cur_input_ids in enumerate(input_ids):
            num_images = (cur_input_ids == IMAGE_TOKEN_INDEX).sum()
            if num_images == 0:
                cur_image_features = image_features[cur_image_idx]
                cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids)
                cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0]], dim=0)
                new_input_embeds.append(cur_input_embeds)
                new_labels.append(labels[batch_idx])
                cur_image_idx += 1
                continue
            image_token_indices = [-1] + torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0].tolist() + [cur_input_ids.shape[0]]
            cur_input_ids_noim = []
            cur_labels = labels[batch_idx]
            cur_labels_noim = []
            for i in range(len(image_token_indices) - 1):
                cur_input_ids_noim.append(cur_input_ids[image_token_indices[i]+1:image_token_indices[i+1]])
                cur_labels_noim.append(cur_labels[image_token_indices[i]+1:image_token_indices[i+1]])
            split_sizes = [x.shape[0] for x in cur_labels_noim]
            cur_input_embeds = self.get_model().embed_tokens(torch.cat(cur_input_ids_noim))
            cur_input_embeds_no_im = torch.split(cur_input_embeds, split_sizes, dim=0)
            cur_new_input_embeds = []
            cur_new_labels = []
            cur_new_input_ids = []

            for i in range(num_images + 1):
                cur_new_input_embeds.append(cur_input_embeds_no_im[i])
                cur_new_labels.append(cur_labels_noim[i])
                cur_new_input_ids.append(cur_input_ids_noim[i])
                if i < num_images:
                    if cur_image_idx < image_features.shape[0]:
                        cur_image_features = image_features[cur_image_idx]
                    else:
                        cur_image_features = image_features[-1]
                    cur_image_idx += 1
                    cur_new_input_embeds.append(cur_image_features)
                    cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=cur_labels.device, dtype=cur_labels.dtype))
                    cur_new_input_ids.append(torch.full((cur_image_features.shape[0],), IMAGE_TOKEN_INDEX, device=cur_labels.device, dtype=cur_labels.dtype))

            cur_new_input_embeds = [x.to(self.device) for x in cur_new_input_embeds]

            cur_new_input_embeds = torch.cat(cur_new_input_embeds, dim=0)
            cur_new_labels = torch.cat(cur_new_labels, dim=0)
            cur_new_input_ids = torch.cat(cur_new_input_ids, dim=0)

            new_input_embeds.append(cur_new_input_embeds)
            new_labels.append(cur_new_labels)
            new_input_ids.append(cur_new_input_ids)

        # Combine them
        max_len = max(x.shape[0] for x in new_input_embeds)
        batch_size = len(new_input_embeds)

        new_input_embeds_padded = []
        new_labels_padded = torch.full((batch_size, max_len), IGNORE_INDEX, dtype=new_labels[0].dtype, device=new_labels[0].device)
        attention_mask = torch.zeros((batch_size, max_len), dtype=attention_mask.dtype, device=attention_mask.device)
        position_ids = torch.zeros((batch_size, max_len), dtype=position_ids.dtype, device=position_ids.device)
        new_input_ids_padded = torch.full((batch_size, max_len), -300, dtype=new_input_ids[0].dtype, device=new_input_ids[0].device) if len(new_input_ids) > 0 else None
        

        for i, (cur_new_embed, cur_new_labels, cur_new_input_ids) in enumerate(zip(new_input_embeds, new_labels, new_input_ids)):
            cur_len = cur_new_embed.shape[0]
            new_input_embeds_padded.append(torch.cat((
                cur_new_embed,
                torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)
            ), dim=0))
            if cur_len > 0:
                new_labels_padded[i, :cur_len] = cur_new_labels
                attention_mask[i, :cur_len] = True
                position_ids[i, :cur_len] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)
                new_input_ids_padded[i, :cur_len] = cur_new_input_ids

        new_input_embeds = torch.stack(new_input_embeds_padded, dim=0)

        if _labels is None:
            new_labels = None
        else:
            new_labels = new_labels_padded

        if _attention_mask is None:
            attention_mask = None
        else:
            attention_mask = attention_mask.to(dtype=_attention_mask.dtype)

        if _position_ids is None:
            position_ids = None

        return None, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels, target_image_embeds


    def initialize_vision_tokenizer(self, model_args, tokenizer):
        if model_args.mm_use_im_patch_token:
            tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
            self.resize_token_embeddings(len(tokenizer))

        if model_args.mm_use_im_start_end:
            num_new_tokens = tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
            self.resize_token_embeddings(len(tokenizer))

            if num_new_tokens > 0:
                input_embeddings = self.get_input_embeddings().weight.data
                output_embeddings = self.get_output_embeddings().weight.data

                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
                    dim=0, keepdim=True)
                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
                    dim=0, keepdim=True)

                input_embeddings[-num_new_tokens:] = input_embeddings_avg
                output_embeddings[-num_new_tokens:] = output_embeddings_avg

            if model_args.tune_mm_mlp_adapter:
                for p in self.get_input_embeddings().parameters():
                    p.requires_grad = True
                for p in self.get_output_embeddings().parameters():
                    p.requires_grad = False

            if model_args.pretrain_mm_mlp_adapter:
                mm_projector_weights = torch.load(model_args.pretrain_mm_mlp_adapter, map_location='cpu')
                embed_tokens_weight = mm_projector_weights['model.embed_tokens.weight']
                assert num_new_tokens == 2
                if input_embeddings.shape == embed_tokens_weight.shape:
                    input_embeddings[-num_new_tokens:] = embed_tokens_weight[-num_new_tokens:]
                elif embed_tokens_weight.shape[0] == num_new_tokens:
                    input_embeddings[-num_new_tokens:] = embed_tokens_weight
                else:
                    raise ValueError(f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}.")
        elif model_args.mm_use_im_patch_token:
            if model_args.tune_mm_mlp_adapter:
                for p in self.get_input_embeddings().parameters():
                    p.requires_grad = False
                for p in self.get_output_embeddings().parameters():
                    p.requires_grad = False