bernini/training/data.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 json
import math
import random
from collections import defaultdict
from collections.abc import Iterable
from typing import Any, Callable, Dict, List, Optional

import torch
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
from diffusers.pipelines.wan.pipeline_wan import prompt_clean
from einops import rearrange
from transformers import PreTrainedTokenizerBase

from bernini.models.scheduler import FlowMatchScheduler


SYSTEM_PROMPTS = {
    "default": "You are a helpful assistant.",
    "t2i": "You are a helpful assistant specialized in text-to-image generation.",
    "t2v": "You are a helpful assistant specialized in text-to-video generation.",
    "i2i": "You are a helpful assistant specialized in image editing.",
    "r2i": "You are a helpful assistant specialized in subject-to-image generation.",
    "i2v": "You are a helpful assistant specialized in image-to-video generation.",
    "v2v": "You are a helpful assistant specialized in video editing.",
    "r2v": "You are a helpful assistant specialized in subject-to-video generation.",
    "vi2v": "You are a helpful assistant specialized in video editing on content propagation.",
    "vr2v": "You are a helpful assistant specialized in video editing with reference.",
    "ads2v": "You are a helpful assistant specialized in ads insertion.",
    "vrc2v": "You are a helpful assistant for editing. You may need to adjust the subject's action or position.",
    "mv2v": "You are a helpful assistant for editing. You might need to adjust the video's style, lighting, colors, textures, and the subject's pose or action.",
}


def _is_non_empty(value: Any) -> bool:
    return isinstance(value, Iterable) and not isinstance(value, (str, bytes)) and len(value) > 0


def _load_tensor(blob: Any) -> torch.Tensor:
    if isinstance(blob, torch.Tensor):
        return blob
    buffer = io.BytesIO(blob)
    buffer.seek(0)
    return torch.load(buffer, map_location="cpu")


def shift2boundary(shift, sigma_min=0, sigma_max=1, denoising_strength=1.0, num_steps=1000):
    sigma_start = sigma_min + (sigma_max - sigma_min) * denoising_strength
    sigmas = torch.linspace(sigma_start, sigma_min, num_steps + 1)[:-1]
    return shift * sigmas / (1 + (shift - 1) * sigmas)


def find_nearest_boundary(sigmas, sigma_value):
    return int(torch.argmin((sigmas - sigma_value).abs()).item())


def compute_density_for_timestep_sampling(
    weighting_scheme: str,
    batch_size: int,
    logit_mean: float = 0.5,
    logit_std: float = 1.0,
    mode_scale: float = 1.29,
    min_t: float = 0.0,
    max_t: float = 1.0,
):
    samples = []
    for _ in range(batch_size):
        while True:
            if weighting_scheme == "logit_normal":
                u = torch.sigmoid(torch.normal(mean=logit_mean, std=logit_std, size=(1,), device="cpu"))
            elif weighting_scheme == "mode":
                raw = torch.rand(size=(1,), device="cpu")
                u = 1 - raw - mode_scale * (torch.cos(math.pi * raw / 2) ** 2 - 1 + raw)
            else:
                u = torch.rand(size=(1,), device="cpu") * (max_t - min_t) + min_t
            if min_t <= float(u.item()) <= max_t:
                samples.append(u)
                break
    return torch.cat(samples, dim=0)


class NoiseScheduler:
    def __init__(
        self,
        shift_config: Optional[Dict[str, float]] = None,
        weighting_scheme_config: Optional[Dict[str, str]] = None,
        noise_tmin: float = 0.0,
        noise_tmax: float = 1.0,
        logit_mean: float = 0.5,
        logit_std: float = 1.0,
        mode_scale: float = 1.29,
    ):
        self.shift_config = {"default": 5.0, **(shift_config or {})}
        self.weighting_scheme_config = {
            "default": "logit_normal",
            "image": "logit_normal",
            "video": "mode",
            **(weighting_scheme_config or {}),
        }
        self.noise_tmin = noise_tmin
        self.noise_tmax = noise_tmax
        self.logit_mean = logit_mean
        self.logit_std = logit_std
        self.mode_scale = mode_scale
        self.flow_scheduler = {}
        for shift in set(self.shift_config.values()):
            sigmas = shift2boundary(shift)
            bound1 = find_nearest_boundary(sigmas, self.noise_tmin) / 1000
            bound2 = find_nearest_boundary(sigmas, self.noise_tmax) / 1000
            scheduler = FlowMatchScheduler(shift=shift, sigma_min=0.0, extra_one_step=True)
            scheduler.set_timesteps(1000, training=True, device="cpu")
            self.flow_scheduler[shift] = {
                "tmin": min(bound1, bound2),
                "tmax": max(bound1, bound2),
                "scheduler": scheduler,
            }

    def get_noise_sigma(self, task_name: str):
        if task_name in self.weighting_scheme_config:
            weighting_name = task_name
        elif "2" in task_name and task_name.rsplit("2", 1)[-1] == "i":
            weighting_name = "image"
        elif "2" in task_name and task_name.rsplit("2", 1)[-1] == "v":
            weighting_name = "video"
        else:
            weighting_name = "default"
        shift_name = task_name if task_name in self.shift_config else "default"
        cfg = self.flow_scheduler[self.shift_config[shift_name]]
        u = compute_density_for_timestep_sampling(
            self.weighting_scheme_config[weighting_name],
            batch_size=1,
            logit_mean=self.logit_mean,
            logit_std=self.logit_std,
            mode_scale=self.mode_scale,
            min_t=cfg["tmin"],
            max_t=cfg["tmax"],
        )
        scheduler = cfg["scheduler"]
        timestep_id = (u * scheduler.num_train_timesteps).long()
        timestep = scheduler.timesteps[timestep_id]
        sigma = scheduler.get_noise_sigma(timestep)
        return sigma.reshape(-1), timestep.reshape(-1)


def encode_renderer_messages(
    conversations: List[Dict[str, Any]],
    tokenizer: PreTrainedTokenizerBase,
    task_name: str,
    drop_text: bool,
    drop_video: bool,
    drop_img: bool,
) -> Dict[str, torch.Tensor]:
    image_vit_mask, video_vit_mask = [], []
    image_drop_mask, video_drop_mask = [], []
    vae_type_list = []
    texts = []
    for message in conversations:
        msg_type = message.get("type")
        if msg_type == "special_token":
            continue
        has_loss = message.get("has_loss", 1 if msg_type == "video_gen" else 0)
        if msg_type == "cot_text":
            msg_type, has_loss = "text", 0
        if msg_type == "text" and has_loss == 0:
            if not drop_text:
                texts.append(message.get("text", ""))
        elif msg_type in ("image", "image_gen"):
            image_vit_mask.append(has_loss)
            image_drop_mask.append(int(drop_img))
            if not drop_img or has_loss == 1:
                vae_type_list.append(0)
        elif msg_type in ("video", "frame_gen", "video_gen"):
            video_vit_mask.append(has_loss)
            video_drop_mask.append(int(drop_video))
            if not drop_video or has_loss == 1:
                vae_type_list.append(1)
        else:
            raise ValueError(f"Unknown message type: {msg_type}")

    prompt = " ".join(texts)
    prompt = SYSTEM_PROMPTS.get(task_name, SYSTEM_PROMPTS["default"]) + prompt_clean(prompt)
    tokenized = tokenizer(prompt, add_special_tokens=True, return_attention_mask=True, return_tensors="pt")
    input_ids = tokenized.input_ids.squeeze(0)
    attention_mask = tokenized.attention_mask.squeeze(0)
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "t5_input_lens": torch.tensor([input_ids.shape[0]], dtype=torch.long),
        "image_vit_mask": torch.tensor(image_vit_mask, dtype=torch.bool),
        "video_vit_mask": torch.tensor(video_vit_mask, dtype=torch.bool),
        "image_drop_mask": torch.tensor(image_drop_mask, dtype=torch.bool),
        "video_drop_mask": torch.tensor(video_drop_mask, dtype=torch.bool),
        "vae_type_list": torch.tensor(vae_type_list, dtype=torch.long),
    }


def _filter_out_target_clip(inputs, mask, embeds_key, grid_key):
    if embeds_key not in inputs or grid_key not in inputs:
        return
    embeds, grids = [], []
    for emb, grid, drop in zip(inputs[embeds_key], inputs[grid_key], mask):
        if not bool(drop):
            embeds.append(emb)
            grids.append(grid)
    if embeds:
        inputs[embeds_key] = torch.cat(embeds, dim=0)
        inputs[grid_key] = torch.stack(grids)
    else:
        inputs.pop(embeds_key)
        inputs.pop(grid_key)


def _filter_source_vae(inputs, vae_latents, vit_mask, drop_vision, latent_key, mask_key):
    kept_latents, kept_mask = [], []
    for vae_emb, is_target in zip(vae_latents, vit_mask):
        if not drop_vision or bool(is_target):
            kept_latents.append(vae_emb)
            kept_mask.append(bool(is_target))
    inputs[latent_key] = kept_latents
    inputs[mask_key] = torch.tensor(kept_mask, dtype=torch.bool)


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


def pack_vae_latents(
    vae_rope_func: Callable,
    vae_type_list: torch.Tensor,
    image_inputs: Dict[str, Any],
    video_inputs: Dict[str, Any],
    noise_sigma: torch.Tensor,
    max_vae_frames: Optional[int] = None,
):
    image_vae_list = iter(image_inputs.pop("image_vae_latents", []))
    image_vae_mask_list = iter(image_inputs.pop("image_vae_mask", []))
    video_vae_list = iter(video_inputs.pop("video_vae_latents", []))
    video_vae_mask_list = iter(video_inputs.pop("video_vae_mask", []))
    input_vae_latents, input_vae_rope, vae_latents_mask = [], [], []
    target_velocity, target_lens = [], []
    for idx, vae_type in enumerate(vae_type_list.tolist()):
        if vae_type == 0:
            vae_emb = next(image_vae_list)
            vae_mask = bool(next(image_vae_mask_list))
        else:
            vae_emb = next(video_vae_list)
            vae_mask = bool(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]

        source_id = 0 if vae_mask else idx + 1
        vae_rope = vae_rope_func(vae_emb.unsqueeze(0), source_id=source_id).squeeze(0)
        input_vae_rope.append(vae_rope)
        packed = _rearrange_vae_feature(vae_emb)
        vae_latents_mask.extend([vae_mask] * packed.shape[0])
        if vae_mask:
            noise = torch.randn_like(packed, dtype=torch.float32)
            input_vae_latents.append((1 - noise_sigma) * packed + noise_sigma * noise)
            target_velocity.append(noise - packed.float())
            target_lens.append(packed.shape[0])
        else:
            input_vae_latents.append(packed)

    input_vae_latents = torch.cat(input_vae_latents, dim=0)
    target_velocity = torch.cat(target_velocity, dim=0)
    input_vae_rope = torch.cat(input_vae_rope, dim=1)
    return {
        "input_vae_latents": input_vae_latents,
        "input_vae_rope": input_vae_rope.permute(1, 0, 2),
        "vae_latents_mask": torch.tensor(vae_latents_mask, dtype=torch.bool),
        "vae_seqlen": torch.tensor([input_vae_latents.shape[0]], dtype=torch.long),
        "target_velocity": target_velocity,
        "target_lens": torch.tensor(target_lens, dtype=torch.long),
    }


def process_renderer_sample(
    sample: Dict[str, Any],
    tokenizer: PreTrainedTokenizerBase,
    vae_rope_func: Callable,
    vae_latent_mean: torch.Tensor,
    vae_latent_std: torch.Tensor,
    noise_scheduler: NoiseScheduler,
    text_dropout_rate: float = 0.0,
    img_dropout_rate: float = 0.0,
    video_dropout_rate: float = 0.0,
    max_vae_frames: Optional[int] = None,
    source_name: str = "",
    **kwargs,
):
    source_name = source_name or sample.get("source_name", "") or ""
    task_name = source_name.split("$")[0].lower() or "default"
    noise_sigma, noise_timestep = noise_scheduler.get_noise_sigma(task_name)
    drop_text = random.random() < text_dropout_rate
    drop_img = random.random() < img_dropout_rate
    drop_video = random.random() < video_dropout_rate
    tokenized = encode_renderer_messages(
        json.loads(sample["inputs"]), tokenizer, task_name, drop_text, drop_video, drop_img
    )

    image_inputs, video_inputs = {}, {}
    if _is_non_empty(sample.get("image_embeds", [])):
        image_inputs = {"image_embeds": [], "image_grid_thw": []}
        for emb, thw in zip(sample["image_embeds"], sample["image_grid_thw"]):
            image_inputs["image_embeds"].append(_load_tensor(emb))
            image_inputs["image_grid_thw"].append(torch.as_tensor(thw))
        image_inputs["image_grid_thw"] = torch.stack(image_inputs["image_grid_thw"])
    if _is_non_empty(sample.get("video_embeds", [])):
        video_inputs = {"video_embeds": [], "video_grid_thw": []}
        for emb, thw in zip(sample["video_embeds"], sample["video_grid_thw"]):
            video_inputs["video_embeds"].append(_load_tensor(emb))
            video_inputs["video_grid_thw"].append(torch.as_tensor(thw))
        video_inputs["video_grid_thw"] = torch.stack(video_inputs["video_grid_thw"])

    image_drop_mask = tokenized.pop("image_drop_mask")
    video_drop_mask = tokenized.pop("video_drop_mask")
    image_vit_mask = tokenized.pop("image_vit_mask")
    video_vit_mask = tokenized.pop("video_vit_mask")
    if "image_embeds" in image_inputs:
        _filter_out_target_clip(image_inputs, image_vit_mask | image_drop_mask, "image_embeds", "image_grid_thw")
    if "video_embeds" in video_inputs:
        _filter_out_target_clip(video_inputs, video_vit_mask | video_drop_mask, "video_embeds", "video_grid_thw")

    if _is_non_empty(sample.get("image_vae_latents", [])):
        image_latents = []
        for vae_blob in sample["image_vae_latents"]:
            vae = DiagonalGaussianDistribution(_load_tensor(vae_blob)).sample().squeeze(0)
            image_latents.append((vae - vae_latent_mean) / vae_latent_std)
        _filter_source_vae(image_inputs, image_latents, image_vit_mask, drop_img, "image_vae_latents", "image_vae_mask")
    if _is_non_empty(sample.get("video_vae_latents", [])):
        video_latents = []
        for vae_blob in sample["video_vae_latents"]:
            vae = DiagonalGaussianDistribution(_load_tensor(vae_blob)).sample().squeeze(0)
            video_latents.append((vae - vae_latent_mean) / vae_latent_std)
        _filter_source_vae(video_inputs, video_latents, video_vit_mask, drop_video, "video_vae_latents", "video_vae_mask")

    vae_type_list = tokenized.pop("vae_type_list")
    packed = pack_vae_latents(
        vae_rope_func,
        vae_type_list,
        image_inputs,
        video_inputs,
        noise_sigma,
        max_vae_frames=max_vae_frames,
    )
    tokenized.update(packed)
    tokenized["timesteps"] = noise_timestep
    tokenized["vlm_seqlen"] = torch.tensor([tokenized["attention_mask"].sum()], dtype=torch.long)
    tokenized["num_tokens"] = torch.tensor([tokenized["vae_seqlen"][0] + tokenized["vlm_seqlen"][0]], dtype=torch.long)
    return [tokenized]