bernini/data/bernini_process.py

# Copyright (c) 2026 Bytedance Ltd. and/or its affiliate
#
# 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.
import io
import copy
import json
import random
from typing import TYPE_CHECKING, Any, Callable, Dict, List

import torch
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
from einops import rearrange

if TYPE_CHECKING:
    from transformers import ProcessorMixin
    from veomni.data.chat_template import ChatTemplate


def filt_out_source_vae(
    inputs,
    vae_latents,
    vae_latent_shape,
    target_mask,
    drop_vision,
    vae_latent_key,
    vae_shape_key,
    vae_mask_key,
):
    vae_latents_filted, vae_latents_filted_shape, vae_mask = [], [], []
    if drop_vision:
        for (vae_emb, shape, is_target) in zip(vae_latents, vae_latent_shape, target_mask):
            if is_target:
                vae_mask.extend([is_target])
                vae_latents_filted.append(vae_emb)
                vae_latents_filted_shape.append(shape)
    else:
        vae_latents_filted = vae_latents
        vae_latents_filted_shape = vae_latent_shape
        vae_mask = target_mask

    inputs[vae_latent_key] = vae_latents_filted
    inputs[vae_mask_key] = torch.tensor(vae_mask)
    inputs[vae_shape_key] = vae_latents_filted_shape


def get_drop_condition(
    text_dropout_rate: float,
    img_dropout_rate: float,
    video_dropout_rate: float,
):
    drop_text, drop_video, drop_img = 0, 0, 0
    text_drop, img_drop, video_drop = random.random(), random.random(), random.random()
    if text_drop < text_dropout_rate:
        drop_text = 1
    if img_drop < img_dropout_rate:
        drop_img = 1
    if video_drop < video_dropout_rate:
        drop_video = 1
    return drop_text, drop_video, drop_img


def rearrange_vae_feature(vae_emb):
    pt, ph, pw = 1, 2, 2
    patched_vae_emb = rearrange(
        vae_emb, 'c (t pt) (h ph) (w pw) -> (t h w) c pt ph pw', pt=pt, ph=ph, pw=pw
    )
    return patched_vae_emb


def packing_vae(
    vae_rope_func: "Callable",
    vae_type_list: "List[int]",
    image_inputs: "Dict[str, Any]",
    video_inputs: "Dict[str, Any]",
    noise_sigma: float,
    max_vae_frames: int = None,
    interpolate_src_id: bool = True,
    max_trained_src_id: int = 5,
):
    image_vae_masks = list(image_inputs.pop('image_vae_mask', []))
    video_vae_masks = list(video_inputs.pop('video_vae_mask', []))

    image_vae_list = iter(image_inputs.pop('image_vae_latents', []))
    image_vae_shape_list = iter(image_inputs.pop('image_vae_shape', []))
    image_vae_mask_list = iter(image_vae_masks)

    video_vae_list = iter(video_inputs.pop('video_vae_latents', []))
    video_vae_shape_list = iter(video_inputs.pop('video_vae_shape', []))
    video_vae_mask_list = iter(video_vae_masks)

    # Source ids for the conditioning segments (the target keeps source_id 0).
    # Training assigns position-based integer ids; when more conditioning
    # segments are given than the model saw in training (`max_trained_src_id`),
    # evenly spread their ids across the trained range [1, max_trained_src_id]
    # so the rotary phases stay inside the trained manifold instead of
    # extrapolating to unseen integer ids.
    num_src = sum(1 for m in image_vae_masks + video_vae_masks if not m)
    src_sids = None
    if interpolate_src_id and num_src > max_trained_src_id:
        src_sids = torch.linspace(1.0, float(max_trained_src_id), num_src).tolist()
    src_ptr = 0  # cursor into src_sids

    input_vae_latents = []
    input_vae_shape = []
    input_vae_rope = []
    vae_latents_mask = []
    target_velocity = []
    target_lens = []

    for idx, vae_type in enumerate(vae_type_list):
        if vae_type == 0:
            vae_emb = next(image_vae_list)
            vae_shape = next(image_vae_shape_list)
            vae_mask = next(image_vae_mask_list)
        else:
            vae_emb = next(video_vae_list)
            vae_shape = next(video_vae_shape_list)
            vae_mask = next(video_vae_mask_list)

        if max_vae_frames is not None and vae_emb.shape[1] > max_vae_frames:
            vae_emb = vae_emb[:, :max_vae_frames, :, :]
            vae_shape[0] = max_vae_frames

        if not vae_mask:
            src_sid = src_sids[src_ptr] if src_sids is not None else float(idx + 1)
            src_ptr += 1

        if vae_rope_func is None:
            vae_rope = None
        elif vae_mask:
            vae_rope = vae_rope_func(vae_emb.unsqueeze(0), source_id=0).squeeze(0)
        else:
            vae_rope = vae_rope_func(vae_emb.unsqueeze(0), source_id=src_sid).squeeze(0)
        input_vae_rope.append(vae_rope)
        input_vae_shape.append(vae_shape)
        vae_emb = rearrange_vae_feature(vae_emb)

        vae_latents_mask.extend([vae_mask] * vae_emb.shape[0])
        if vae_mask:
            noise_sigma = noise_sigma.to(device=vae_emb.device, dtype=vae_emb.dtype)
            noise = torch.randn_like(vae_emb, dtype=torch.float32)
            input_noise_latent = (1 - noise_sigma) * vae_emb + noise_sigma * noise
            target_velocity.append(noise - vae_emb.float())
            target_lens.append(vae_emb.shape[0])
            input_vae_latents.append(input_noise_latent)
        else:
            input_vae_latents.append(vae_emb)
    input_vae_shape = torch.tensor(input_vae_shape)
    vae_latents_mask = torch.tensor(vae_latents_mask)
    packed_vae_latents = {
        "input_vae_shape": input_vae_shape,
        "vae_latents_mask": vae_latents_mask,
    }
    if vae_rope_func is not None and len(input_vae_rope) > 0:
        input_vae_rope = torch.cat(input_vae_rope, dim=1)
        input_vae_rope = input_vae_rope.permute(1, 0, 2)
        packed_vae_latents['input_vae_rope'] = input_vae_rope

    if len(input_vae_latents) > 0:
        input_vae_latents = torch.cat(input_vae_latents, dim=0)
        vae_seqlen = input_vae_latents.shape[0]
        target_lens = torch.tensor(target_lens)
        packed_vae_latents['input_vae_latents'] = input_vae_latents
    else:
        vae_seqlen = 0
        target_lens = torch.tensor([])

    packed_vae_latents['vae_seqlen'] = torch.tensor([vae_seqlen])
    packed_vae_latents['target_lens'] = target_lens

    if len(target_velocity) > 0:
        target_velocity = torch.cat(target_velocity, dim=0)
        packed_vae_latents['target_velocity'] = target_velocity

    return packed_vae_latents


def bernini_process_sample(
    sample: Dict[str, Any],
    processor: "ProcessorMixin",
    chat_template: "ChatTemplate",
    position_id_func: "Callable",
    vae_rope_func: Callable,
    vae_latent_mean: torch.Tensor,
    vae_latent_std: torch.Tensor,
    text_dropout_rate: float,
    img_dropout_rate: float,
    video_dropout_rate: float,
    max_vae_frames: int = None,  
    noise_sigma=torch.tensor(0),
    noise_timestep=torch.tensor(0), 
    noise_sigma_low=None,
    noise_timestep_low=None,
    vit_mask_ratio: float = 1.0, # mask all tokens when inference
    interpolate_src_id: bool = True,
    max_trained_src_id: int = 5,
    **kwargs,
):
    """
    Processes multimodal example with qwen2_5_vl's pre-processor.
    """
    
    task_name = kwargs.get("source_name", "").split("$")[0].lower()
    drop_text, drop_video, drop_img = get_drop_condition(
        text_dropout_rate,
        img_dropout_rate,
        video_dropout_rate
    )

    conversations = json.loads(sample["inputs"])

    token_num_inputs = {}
    if "image_embeds" in sample and sample['image_embeds'] is not None and len(sample['image_embeds']) > 0:
        raw_image_embeds = []
        raw_image_grid_thw = []
        for vit_emb, thw in zip(sample['image_embeds'], sample['image_grid_thw']):
            buffer = io.BytesIO(vit_emb)
            buffer.seek(0)
            vit_emb = torch.load(buffer)
            raw_image_embeds.append(vit_emb)
            raw_image_grid_thw.append(thw)

        merge_length = processor.image_processor.merge_size**2
        token_num_inputs["image"] = (
            torch.tensor(raw_image_grid_thw).prod(dim=-1) // merge_length
        )

    if "video_embeds" in sample and sample['video_embeds'] is not None and len(sample['video_embeds']) > 0:
        raw_video_embeds = []
        raw_video_grid_thw = []
        for vit_emb, thw in zip(sample['video_embeds'], sample['video_grid_thw']):
            buffer = io.BytesIO(vit_emb)
            buffer.seek(0)
            vit_emb = torch.load(buffer)
            raw_video_embeds.append(vit_emb)
            raw_video_grid_thw.append(thw)

        merge_length = processor.image_processor.merge_size**2
        token_num_inputs["video"] = (
            torch.tensor(raw_video_grid_thw).prod(dim=-1) // merge_length
        )

    tokenized_example = chat_template.encode_messages(
        conversations,
        token_num_inputs,
        task_name,
        drop_text=drop_text,
        drop_video=drop_video,
        drop_img=drop_img,
        vit_mask_ratio=vit_mask_ratio,
        **kwargs,
    )

    for k, v in tokenized_example.items():
        if isinstance(v, str):
            continue
        if torch.is_tensor(v):
            tokenized_example[k] = v
        else:
            tokenized_example[k] = torch.as_tensor(v)
    
    # Packing vit embeds
    vit_type_list, vit_img_and_vid_id_list = tokenized_example.pop(
        'vit_type_list'), tokenized_example.pop('vit_img_and_vid_id_list')
    visual_embeds = []
    image_grid_thw, video_grid_thw = [], []
    for vit_type, vit_id in zip(vit_type_list, vit_img_and_vid_id_list):
        if vit_type == 0: # image
            image_grid_thw.append(raw_image_grid_thw[vit_id])
            visual_embeds.append(raw_image_embeds[vit_id])
        elif vit_type == 1: # video
            video_grid_thw.append(raw_video_grid_thw[vit_id])
            visual_embeds.append(raw_video_embeds[vit_id])
    
    if len(image_grid_thw) > 0:
        image_grid_thw = torch.tensor(image_grid_thw)
    if len(video_grid_thw) > 0:
        video_grid_thw = torch.tensor(video_grid_thw)
    if len(visual_embeds) > 0:
        visual_embeds = torch.cat(visual_embeds, dim=0)
        tokenized_example['visual_embeds'] = visual_embeds
    else:
        visual_embeds = torch.randn(0, 3584)

    
    input_ids = tokenized_example["input_ids"]
    tokenized_example["position_ids"] = position_id_func(
        input_ids=input_ids.unsqueeze(0),
        image_grid_thw=image_grid_thw if len(image_grid_thw) > 0 else None,
        video_grid_thw=video_grid_thw if len(video_grid_thw) > 0 else None,
        attention_mask=tokenized_example["attention_mask"].unsqueeze(0),
    )[0].squeeze(1).clone()  # (dim, l)
    tokenized_example["mllm_seqlen"] = tokenized_example["attention_mask"].sum().reshape(1)

    image_inputs, video_inputs = {}, {}
    # Packing vae latents
    image_target_mask, video_target_mask = tokenized_example.pop(
        'image_target_mask'), tokenized_example.pop('video_target_mask')

    if sample.get('image_vae_latents', None) is not None and len(sample['image_vae_latents']) > 0:
        image_vae_latents = []
        image_vae_shape = []
        for vae_emb, _ in zip(sample['image_vae_latents'], image_target_mask):
            buffer = io.BytesIO(vae_emb)
            buffer.seek(0)
            vae_emb = torch.load(buffer)
            _, _, t, h, w = vae_emb.shape
            image_vae_shape.append([t, h, w])
            vae_emb = DiagonalGaussianDistribution(vae_emb).mode()
            vae_emb = vae_emb.squeeze(0)
            vae_emb = (vae_emb - vae_latent_mean) / vae_latent_std
            image_vae_latents.append(vae_emb)
        filt_out_source_vae(
            image_inputs,
            image_vae_latents,
            image_vae_shape,
            image_target_mask,
            drop_img,
            'image_vae_latents',
            'image_vae_shape',
            'image_vae_mask'
        )

    if "video_vae_latents" in sample and len(sample['video_vae_latents']) > 0:
        video_vae_latents = []
        video_vae_shape = []
        for vae_emb, _ in zip(sample['video_vae_latents'], video_target_mask):
            buffer = io.BytesIO(vae_emb)
            buffer.seek(0)
            vae_emb = torch.load(buffer)
            _, _, t, h, w = vae_emb.shape
            video_vae_shape.append([t, h, w])
            vae_emb = DiagonalGaussianDistribution(vae_emb).mode()
            vae_emb = vae_emb.squeeze(0)
            vae_emb = (vae_emb - vae_latent_mean) / vae_latent_std
            video_vae_latents.append(vae_emb)
        filt_out_source_vae(
            video_inputs,
            video_vae_latents,
            video_vae_shape,
            video_target_mask,
            drop_video,
            'video_vae_latents',
            'video_vae_shape',
            'video_vae_mask'
        )

    vae_type_list = tokenized_example.pop('vae_type_list')

    if noise_sigma_low is not None:
        image_inputs_copy = copy.deepcopy(image_inputs)
        video_inputs_copy = copy.deepcopy(video_inputs)
        packed_vae_latents_low = packing_vae(
            vae_rope_func,
            vae_type_list,
            image_inputs_copy,
            video_inputs_copy,
            noise_sigma_low,
            max_vae_frames,
            interpolate_src_id=interpolate_src_id,
            max_trained_src_id=max_trained_src_id,
        )
        for k, v in packed_vae_latents_low.items():
            tokenized_example[k + '_low'] = v

    packed_vae_latents = packing_vae(
        vae_rope_func,
        vae_type_list,
        image_inputs,
        video_inputs,
        noise_sigma,
        max_vae_frames,
        interpolate_src_id=interpolate_src_id,
        max_trained_src_id=max_trained_src_id,
    )
    tokenized_example.update(packed_vae_latents)
    tokenized_example['timesteps'] = torch.tensor([noise_timestep])
    if noise_timestep_low is not None:
        tokenized_example['timesteps_low'] = torch.tensor([noise_timestep_low])

    tokenized_example["num_tokens"] = torch.tensor(
        [tokenized_example["vae_seqlen"][0] + tokenized_example["mllm_seqlen"][0]]
    )

    tokenized_example['task_name'] = task_name
    return [tokenized_example]