# 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 copy import torch import torch.nn as nn import torch.distributed as dist import os from transformers import UMT5EncoderModel, Qwen2_5_VLConfig from transformers.configuration_utils import PretrainedConfig from transformers.modeling_utils import PreTrainedModel from transformers.utils import logging from .diffloss_fm import DiffLoss_FM from .wan_diffusion import GEN_Wanx22 from .modeling_qwen2_5_vl import Qwen2_5_VLForConditionalGeneration logger = logging.get_logger(__name__) def _join_subfolder(base_subfolder, leaf): if base_subfolder: return f"{base_subfolder}/{leaf}" return leaf class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x): dtype = x.dtype x = x.float() x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) return (x * self.weight).to(dtype) class MLPConnector(nn.Module): """Local connector implementation used by the Bernini checkpoint. Keep the same parameter layout as the original training connector so the released checkpoint loads without the external training package. """ def __init__( self, in_dim, num_layers_for_gen=1, out_dim_for_gen=4096, enable_gen_branch=True, gen_head_type="mlp", num_layers_for_vit=1, out_dim_for_vit=3584, enable_vit_branch=True, ): super().__init__() self.enable_gen_branch = enable_gen_branch self.enable_vit_branch = enable_vit_branch if enable_gen_branch: self.proj_gen = nn.Sequential( nn.Linear(in_dim, out_dim_for_gen), nn.GELU(), RMSNorm(out_dim_for_gen), nn.Linear(out_dim_for_gen, out_dim_for_gen), ) if enable_vit_branch: self.pred_vit = nn.Sequential( nn.Linear(in_dim, out_dim_for_vit), nn.GELU(), nn.Linear(out_dim_for_vit, out_dim_for_vit), RMSNorm(out_dim_for_vit), nn.Linear(out_dim_for_vit, out_dim_for_vit), ) @staticmethod def _run_projection(proj, x): param = next(proj.parameters(), None) if param is not None and (x.device != param.device or x.dtype != param.dtype): x = x.to(device=param.device, dtype=param.dtype) return proj(x) def for_gen(self, x): return self._run_projection(self.proj_gen, x) def for_vit(self, x): return self._run_projection(self.pred_vit, x) def with_skip_config(config, skip_transformer_1=False, skip_transformer_2=False): """Create a copy of config with skip_transformer_* flags set.""" config_copy = copy.deepcopy(config) config_copy.skip_transformer_1 = skip_transformer_1 config_copy.skip_transformer_2 = skip_transformer_2 return config_copy def _join_if_present(base, *parts): if base is None: return None return os.path.join(base, *parts) class BerniniConfig(PretrainedConfig): model_type = "bernini" def __init__( self, base_dir=None, mllm_config_path=None, mllm_subfolder=None, processor_config_path=None, processor_subfolder=None, mllm_attn_implementation="sdpa", diff_dec_config_path=None, transformer_config_path=None, transformer_2_config_path=None, scheduler_config_path=None, bernini_ckpt_subfolder=None, scratch_mllm=False, scratch=False, noise_tmin=0.0, # [0, 0.875] for WAN2.2 high noise; [0.875, 1.0] for WAN2.2 low noise noise_tmax=1.0, # [0, 1.0] for WAN2.1 flow_shift=5, use_unipc=False, target_fps=16, switch_dit_boundary=0.875, shift=3.0, cotrain=False, # setting for clip fmmar num_mask_token=256, clip_diff_cfg=None, connector_cfg=None, mask_ratio_infer_cfg=None, feature_type_from_stage_one=None, additional_special_tokens=[], tie_word_embeddings=False, ema_decay=None, partial_pretrain_model=None, use_src_id_rotary_emb=False, interpolate_src_id=True, max_trained_src_id=5, max_sequence_length=512, # t5 embedding t5_text_encoder_path=None, t5_text_encoder_subfolder=None, t5_tokenizer_path=None, t5_tokenizer_subfolder=None, t5_max_sequence_length=512, t5_combine_type="kl_loss", vae_model_path=None, vae_subfolder=None, vae_config_path=None, wovae_task_list=['und_img', 'und_txt', 'und_vid'], **kwargs, ): super().__init__(**kwargs) self.base_dir = base_dir self.mllm_config_path = mllm_config_path if mllm_config_path is not None else base_dir self.mllm_subfolder = mllm_subfolder self.mllm_attn_implementation = mllm_attn_implementation self.processor_config_path = ( processor_config_path if processor_config_path is not None else self.mllm_config_path ) self.processor_subfolder = processor_subfolder self.diff_dec_config_path = diff_dec_config_path if diff_dec_config_path is not None else base_dir self.transformer_config_path = ( transformer_config_path if transformer_config_path is not None else _join_if_present(base_dir, "transformer_config.json") ) self.transformer_2_config_path = ( transformer_2_config_path if transformer_2_config_path is not None else _join_if_present(base_dir, "transformer_2_config.json") ) self.scheduler_config_path = scheduler_config_path or ( os.path.join(base_dir, "scheduler") if base_dir is not None else None ) self.bernini_ckpt_subfolder = bernini_ckpt_subfolder self.vae_model_path = vae_model_path if vae_model_path is not None else base_dir self.vae_subfolder = vae_subfolder self.vae_config_path = vae_config_path or ( os.path.join(self.vae_model_path, _join_subfolder(self.vae_subfolder or "vae", "config.json")) if self.vae_model_path is not None else None ) self.scratch = scratch self.scratch_mllm = scratch_mllm self.ema_decay = ema_decay self.noise_tmin = noise_tmin self.noise_tmax = noise_tmax self.flow_shift = flow_shift self.use_unipc = use_unipc self.target_fps = target_fps self.switch_dit_boundary = switch_dit_boundary self.shift = shift self.cotrain = cotrain self.use_src_id_rotary_emb = use_src_id_rotary_emb # When the number of conditioning segments exceeds `max_trained_src_id` # (the largest source_id seen in training), evenly map their ids into # the trained range [1, max_trained_src_id] instead of extrapolating to # unseen integer ids. The target segment keeps source_id 0. self.interpolate_src_id = interpolate_src_id self.max_trained_src_id = max_trained_src_id self.max_sequence_length = max_sequence_length self.wovae_task_list = wovae_task_list self.num_mask_token = num_mask_token self.clip_diff_cfg = clip_diff_cfg self.connector_cfg = connector_cfg self.mask_ratio_infer_cfg = mask_ratio_infer_cfg self.feature_type_from_stage_one = feature_type_from_stage_one self.additional_special_tokens = additional_special_tokens self.tie_word_embeddings = tie_word_embeddings self.partial_pretrain_model = partial_pretrain_model self.t5_text_encoder_path = t5_text_encoder_path if t5_text_encoder_path is not None else base_dir self.t5_text_encoder_subfolder = t5_text_encoder_subfolder self.t5_tokenizer_path = t5_tokenizer_path if t5_tokenizer_path is not None else base_dir self.t5_tokenizer_subfolder = t5_tokenizer_subfolder self.t5_max_sequence_length = t5_max_sequence_length self.t5_combine_type = t5_combine_type self.architectures = ["BerniniModel"] class BerniniModel(PreTrainedModel): config_class = BerniniConfig def __init__(self, config): super().__init__(config) self.mllm = None self.diff_dec = None self.vit_decoder = None self.base_dir = getattr(config, "base_dir", None) self.mllm_config_path = config.mllm_config_path self.mllm_subfolder = getattr(config, "mllm_subfolder", None) self.diff_dec_config_path = config.diff_dec_config_path self.processor_config_path = config.processor_config_path self.feature_type_from_stage_one = config.feature_type_from_stage_one self.num_mask_token = config.num_mask_token self.use_t5_encoder = config.t5_text_encoder_path is not None # =============== Init MLLM =============== self.mllm_attn_implementation = config.mllm_attn_implementation logger.info( f"MLLM attention implement: config.mllm_attn_implementation={config.mllm_attn_implementation}" ) if self.config.mllm_config_path is not None: mllm_config = Qwen2_5_VLConfig.from_pretrained( self.config.mllm_config_path, subfolder=self.config.mllm_subfolder, ) if self.config.scratch_mllm: self.mllm = Qwen2_5_VLForConditionalGeneration._from_config( mllm_config, attn_implementation=config.mllm_attn_implementation, torch_dtype=torch.bfloat16, ) else: self.mllm = Qwen2_5_VLForConditionalGeneration.from_pretrained( self.config.mllm_config_path, subfolder=self.config.mllm_subfolder, attn_implementation=config.mllm_attn_implementation, ) self.mask_tokens = nn.Parameter(torch.randn(1, self.num_mask_token, self.mllm.config.hidden_size) * 0.01) # =============== Init Diff Dec =============== if self.config.diff_dec_config_path: if getattr(self.config, "cotrain", False): self.diff_dec = GEN_Wanx22(with_skip_config(config, skip_transformer_1=False, skip_transformer_2=True)) self.diff_dec_low = GEN_Wanx22(with_skip_config(config, skip_transformer_1=True, skip_transformer_2=False)) else: self.diff_dec = GEN_Wanx22(config) self.diff_dec_low = None # =============== Init Connector =============== if config.connector_cfg.get('enable_gen_branch', True): assert self.diff_dec is not None if config.connector_cfg.get('enable_vit_branch', True): assert self.mllm is not None self.connector = MLPConnector( in_dim=self.mllm.config.hidden_size if self.mllm is not None else 3584, # setting for diffusion generator num_layers_for_gen=config.connector_cfg.get('num_layers_for_gen', 1), out_dim_for_gen=config.connector_cfg.get('out_dim_for_gen', 4096), enable_gen_branch=config.connector_cfg.get('enable_gen_branch', True), gen_head_type=config.connector_cfg.get('gen_head_type', 'mlp'), # setting for predict vit embed num_layers_for_vit=config.connector_cfg.get('num_layers_for_vit', 1), out_dim_for_vit=config.connector_cfg.get('out_dim_for_vit', 3584), enable_vit_branch=config.connector_cfg.get('enable_vit_branch', True), ) # =============== vit decoder =============== self.vit_decoder = DiffLoss_FM( z_channels=config.clip_diff_cfg.get('z_channels', 3584), target_channels=config.clip_diff_cfg.get('target_channels', 3584), depth=config.clip_diff_cfg.get('depth', 16), width=config.clip_diff_cfg.get('width', 1536), diff_net=config.clip_diff_cfg.get("diff_net", "SimpleMLPAdaLN"), scheduler_type=config.clip_diff_cfg.get("scheduler_type", "FlowMatchScheduler"), shift=config.clip_diff_cfg.get("shift", 3.0), num_inference_steps=config.clip_diff_cfg.get("num_inference_steps", 100), extra_one_step=config.clip_diff_cfg.get("extra_one_step", True), diffusion_batch_mul=config.clip_diff_cfg.get("diffusion_batch_mul", 1), grad_checkpointing=True, ) # =============== t5 embedding =============== if self.use_t5_encoder: logger.info(f"Initializing UMT5 encoder from {config.t5_text_encoder_path}") def load_t5_text_encoder(): return UMT5EncoderModel.from_pretrained( config.t5_text_encoder_path, subfolder=config.t5_text_encoder_subfolder, torch_dtype=torch.bfloat16, ) # Stagger loading across ranks to reduce peak memory if dist.is_initialized(): rank = dist.get_rank() world_size = dist.get_world_size() for r in range(world_size): if r == rank: self.t5_text_encoder = load_t5_text_encoder() dist.barrier() else: self.t5_text_encoder = load_t5_text_encoder() self.t5_max_sequence_length = getattr(config, 't5_max_sequence_length', 512) self.t5_text_encoder.eval() for param in self.t5_text_encoder.parameters(): param.requires_grad = False def get_t5_text_embeddings(self, input_ids, attention_mask, input_lens, pad_text_embeds=True): """ Args: input_ids: tensor with shape (1, n) where n = sum of all sequence lengths attention_mask: tensor with shape (1, n) input_lens: tensor with shape (1, b) where b = batch_size Returns: batch_text_seqlen: list of t5_max_sequence_length repeated b times batch_text_embs: tensor with shape (1, b * t5_max_sequence_length, hidden_dim) """ # Remove batch dim input_ids = input_ids.squeeze(0) # (n,) attention_mask = attention_mask.squeeze(0) # (n,) input_lens = input_lens.squeeze(0) # (b,) batch_size = input_lens.size(0) # Split concatenated sequence into individual samples by lengths input_ids_list = torch.split(input_ids, input_lens.tolist()) attention_mask_list = torch.split(attention_mask, input_lens.tolist()) # Pad each sample to t5_max_sequence_length padded_input_ids = [] padded_attention_mask = [] for ids, mask in zip(input_ids_list, attention_mask_list): seq_len = ids.size(0) if seq_len < self.t5_max_sequence_length: pad_len = self.t5_max_sequence_length - seq_len ids = torch.cat([ids, ids.new_zeros(pad_len)]) mask = torch.cat([mask, mask.new_zeros(pad_len)]) else: ids = ids[:self.t5_max_sequence_length] mask = mask[:self.t5_max_sequence_length] padded_input_ids.append(ids) padded_attention_mask.append(mask) encoder_device = next(self.t5_text_encoder.parameters()).device # Stack to batch: (batch_size, t5_max_sequence_length) input_ids_batch = torch.stack(padded_input_ids, dim=0) attention_mask_batch = torch.stack(padded_attention_mask, dim=0) input_ids_batch = input_ids_batch.to(encoder_device) attention_mask_batch = attention_mask_batch.to(encoder_device) # Get actual sequence lengths (clamped to t5_max_sequence_length) seq_lens = torch.clamp(input_lens, max=self.t5_max_sequence_length) # Get embeddings with torch.no_grad(): prompt_embeds = self.t5_text_encoder( input_ids_batch, attention_mask_batch ).last_hidden_state # (batch_size, t5_max_sequence_length, hidden_dim) # Zero out padding positions if pad_text_embeds: prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack([ torch.cat([u, u.new_zeros(self.t5_max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds ], dim=0) # (batch_size, t5_max_sequence_length, hidden_dim) # Build return values batch_text_seqlen = [self.t5_max_sequence_length] * batch_size batch_text_embs = prompt_embeds.view(1, -1, prompt_embeds.size(-1)) return batch_text_embs, batch_text_seqlen else: prompt_embeds = [u[:v].unsqueeze(0) for u, v in zip(prompt_embeds, seq_lens)] seq_lens = seq_lens.to(dtype=torch.long).cpu().tolist() batch_text_embs = torch.cat(prompt_embeds, dim=1) return batch_text_embs, seq_lens def get_t5_text_embeddings_sample(self, input_ids, attention_mask): encoder_device = next(self.t5_text_encoder.parameters()).device input_ids = input_ids.to(encoder_device) attention_mask = attention_mask.to(encoder_device) seq_lens = attention_mask.gt(0).sum(dim=1).long() with torch.no_grad(): prompt_embeds = self.t5_text_encoder( input_ids, attention_mask).last_hidden_state prompt_embeds = [u[:min(v, self.t5_max_sequence_length)] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack(prompt_embeds, dim=0) return prompt_embeds def get_ignore_modules_in_mixed_precision(self): from diffusers.models.embeddings import TimestepEmbedding from diffusers.models.normalization import FP32LayerNorm return (TimestepEmbedding, FP32LayerNorm) def post_process_input_embeds( self, input_embeds, visual_output_mask, tgt_vit_mask, inference=False ): target_vit_embed_mask = visual_output_mask.squeeze(0) target_vit_embeds = input_embeds[:, target_vit_embed_mask, :] target_vit_embeds_gt = target_vit_embeds.clone() mask_token = self.mask_tokens[:, :1] if inference: # mask all tokens mask_rate = 1 all_vit_token_num = sum(target_vit_embed_mask).detach().cpu().numpy() target_vit_embeds[:, :, :] = mask_token.expand(1, all_vit_token_num, -1) input_embeds[:, target_vit_embed_mask, :] = target_vit_embeds diff_loss_mask = torch.ones(all_vit_token_num).to(target_vit_embeds.device) elif tgt_vit_mask is not None: diff_loss_mask = tgt_vit_mask.squeeze(0).bool() token_num = int(diff_loss_mask.sum().detach().cpu().item()) target_vit_embeds[:, diff_loss_mask, :] = mask_token.expand(1, token_num, -1) input_embeds[:, target_vit_embed_mask, :] = target_vit_embeds else: # tgt_vit_mask is None all_vit_token_num = sum(target_vit_embed_mask).detach().cpu().numpy() diff_loss_mask = torch.zeros(all_vit_token_num).to(target_vit_embeds.device) return dict( input_embeds=input_embeds, diff_loss_mask=diff_loss_mask, target_vit_embeds=target_vit_embeds_gt ) def feat_from_planner_to_renderer( self, hidden_states, tgt_vit_mask, visual_output_mask, inference=False ): pred_vit_embed_mask = visual_output_mask.squeeze(0) pred_vit_embeds = hidden_states[:, pred_vit_embed_mask, :].clone() # For calculate vit decoder loss txt_and_vit_token_mask = visual_output_mask.squeeze(0).logical_not() if not inference: all_idx = torch.nonzero(pred_vit_embed_mask, as_tuple=False).squeeze(-1) # shape [N] cur_clip_mask = tgt_vit_mask.bool().logical_not() valid_clip_idx = all_idx[cur_clip_mask] pred_vit_embed_mask = torch.zeros(hidden_states.shape[1], dtype=torch.bool, device=hidden_states.device) pred_vit_embed_mask[valid_clip_idx] = True cond_embed_mask = (txt_and_vit_token_mask | pred_vit_embed_mask) diff_mllm_context_txt_mask = txt_and_vit_token_mask[cond_embed_mask] diff_mllm_context_vit_mask = pred_vit_embed_mask[cond_embed_mask] connector_param = next(self.connector.parameters()) if connector_param.device != hidden_states.device or connector_param.dtype != hidden_states.dtype: self.connector.to(device=hidden_states.device, dtype=hidden_states.dtype) diff_mllm_contexts = hidden_states[:, cond_embed_mask, :] diff_mllm_contexts = self.connector.for_gen(diff_mllm_contexts) mllm_context_seqlens = [] pred_vit_embed_seqlens = [] pred_vit_embed_mask = visual_output_mask.squeeze(0) mllm_context_seqlens.append(int(cond_embed_mask.sum().item())) pred_vit_embed_seqlens.append(int(pred_vit_embed_mask.sum().item())) return dict( diff_mllm_contexts=diff_mllm_contexts, mllm_context_seqlens=mllm_context_seqlens, pred_vit_embeds=pred_vit_embeds, pred_vit_embed_seqlens=pred_vit_embed_seqlens, diff_mllm_context_txt_mask=diff_mllm_context_txt_mask, diff_mllm_context_vit_mask=diff_mllm_context_vit_mask, ) def format_mllm_inputs_embeds( self, input_ids, visual_embeds, visual_input_mask, visual_output_mask, ): inputs_embeds = self.mllm.get_input_embeddings()(input_ids).to(dtype=torch.bfloat16) if visual_embeds is not None and len(visual_embeds) > 0: visual_mask = visual_input_mask | visual_output_mask n_visual_tokens = visual_mask.sum().long().item() n_visual_features = visual_embeds.shape[0] if n_visual_tokens != n_visual_features: raise ValueError( f"Image features and image tokens do not match: tokens: {n_visual_tokens}, features {n_visual_features}" ) visual_mask = ( visual_mask.unsqueeze(-1) .expand_as(inputs_embeds) .to(inputs_embeds.device) ) visual_embeds = visual_embeds.to( inputs_embeds.device, inputs_embeds.dtype) inputs_embeds = inputs_embeds.masked_scatter( visual_mask, visual_embeds) return inputs_embeds