# 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. """Bernini Renderer model: a UMT5 text encoder + a Wan2.2 dual-expert diffusion decoder.""" import torch import contextlib from transformers import UMT5EncoderModel from transformers.configuration_utils import PretrainedConfig from dataclasses import dataclass from typing import Any, Optional import torch.nn as nn from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.modeling_utils import PreTrainedModel from .wan_diffusion import GEN_Wanx22 # Captured at import time, before VeOmni's build_foundation_model wraps model # construction in init_empty_weights() (which monkeypatches register_parameter # to push params to the meta device). Submodules below load real pretrained # weights via transformers v5 from_pretrained, which is incompatible with that # patch; restore the original around those loads. FSDP2 reshards afterwards. _ORIG_REGISTER_PARAMETER = nn.Module.register_parameter @contextlib.contextmanager def load_pretrained_submodules(): patched = nn.Module.register_parameter nn.Module.register_parameter = _ORIG_REGISTER_PARAMETER try: yield finally: nn.Module.register_parameter = patched class BerniniRendererConfig(PretrainedConfig): model_type = "bernini_renderer" def __init__( self, wan22_base: str = None, diff_dec_config_path: str = None, skip_transformer_1: bool = False, skip_transformer_2: bool = False, switch_dit_boundary: float = 0.875, max_sequence_length: int = 512, shift: float = 3.0, boundary_ratio: float = 0.3, use_unipc: bool = True, scratch: bool = False, ema_decay: float = None, use_src_id_rotary_emb: bool = True, interpolate_src_id: bool = True, max_trained_src_id: int = 5, **kwargs, ): super().__init__(**kwargs) self.wan22_base = wan22_base or diff_dec_config_path self.diff_dec_config_path = diff_dec_config_path or wan22_base self.skip_transformer_1 = skip_transformer_1 self.skip_transformer_2 = skip_transformer_2 self.switch_dit_boundary = switch_dit_boundary self.max_sequence_length = max_sequence_length self.shift = shift self.boundary_ratio = boundary_ratio self.use_unipc = use_unipc self.scratch = scratch self.ema_decay = ema_decay self.use_src_id_rotary_emb = use_src_id_rotary_emb # When the number of reference sources 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 noisy target keeps source_id 0. self.interpolate_src_id = interpolate_src_id self.max_trained_src_id = max_trained_src_id self.architectures = ["BerniniRendererModel"] @dataclass class BerniniRendererOutputWithPast(CausalLMOutputWithPast): diff_loss: Optional[torch.Tensor] = None class EMAModule(nn.Module): def __init__(self, model: nn.Module, ema_decay: float = 0.9999, scratch: bool = False): import copy super().__init__() self.ema_decay = ema_decay self.scratch = scratch for name, module in model.named_children(): self.add_module(name, copy.deepcopy(module)) self.requires_grad_(False) self.eval() @torch.no_grad() def update(self, model: nn.Module): model_params = dict(model.named_parameters()) ema_params = dict(self.named_parameters()) for name, param in model_params.items(): name = name.removeprefix("_fsdp_wrapped_module.") if name not in ema_params: continue if self.scratch: ema_params[name].data.copy_(param.detach().data) else: ema_params[name].data.mul_(self.ema_decay).add_(param.detach().data, alpha=1 - self.ema_decay) self.scratch = False self.eval() class BerniniRendererModel(PreTrainedModel): config_class = BerniniRendererConfig supports_gradient_checkpointing = True _supports_flash_attn = True _no_split_modules = ["UMT5Block", "WanTransformerBlock"] def __init__(self, config: BerniniRendererConfig): super().__init__(config) self.max_sequence_length = config.max_sequence_length # Match the parent model dtype so FSDP2 sees uniform parameter dtypes. model_dtype = getattr(config, "dtype", None) or torch.bfloat16 with load_pretrained_submodules(): self.t5_text_encoder = UMT5EncoderModel.from_pretrained( config.wan22_base, subfolder="text_encoder", torch_dtype=model_dtype ) self.diff_dec = GEN_Wanx22(config) self.t5_text_encoder.requires_grad_(False) if config.ema_decay is not None: self.ema = EMAModule(self, ema_decay=config.ema_decay, scratch=config.scratch) self.post_init() def _set_gradient_checkpointing(self, enable=True, gradient_checkpointing_func=None): for module in self.modules(): if hasattr(module, "gradient_checkpointing"): module.gradient_checkpointing = enable if gradient_checkpointing_func is not None: module._gradient_checkpointing_func = gradient_checkpointing_func def init_weights(self): """Re-materialize pretrained weights instead of HF random init. ``post_init`` (during ``__init__``) and VeOmni's meta + FSDP2 flow both call this. In ``__init__`` the submodules already hold real pretrained weights, so this is a no-op. Under FSDP2 the framework calls ``to_empty()`` (turning params into uninitialized sharded DTensors) then ``init_weights()``; HF's default would random-init and clobber the Wan2.2/UMT5 pretrained weights (and the DTensor RNG broadcast it issues hangs the run). Instead, reload the pretrained submodules and copy them into the (possibly sharded) params. To avoid every rank materializing the (large fp32) pretrained submodules, only rank0 builds the reference state dict; the other ranks receive each tensor through ``distribute_tensor``'s rank0 scatter (``src_data_rank=0``). Whether a param has a source is decided purely from its name so every rank issues the same collectives. """ import torch.distributed as dist from torch.distributed.tensor import DTensor, distribute_tensor params = dict(self.named_parameters()) if not any(p.is_meta or isinstance(p, DTensor) for p in params.values()): return is_dist = dist.is_available() and dist.is_initialized() is_rank0 = (not is_dist) or dist.get_rank() == 0 reference = self._pretrained_reference_state_dict() if is_rank0 else {} def _source_key(name: str): # EMA params mirror the main weights (matches the deepcopy init in # EMAModule); strip the ``ema.`` prefix to find their source tensor. base = name.removeprefix("ema.") if base.startswith("t5_text_encoder.") or base.startswith("diff_dec."): return base return None for name, param in params.items(): key = _source_key(name) if key is None: # Deterministic across ranks: no collective issued for these. continue if isinstance(param, DTensor): if is_rank0: source = reference[key].to(dtype=param.dtype) else: # Placeholder for shape/dtype only; data is ignored because # distribute_tensor scatters from src_data_rank=0. source = torch.empty(param.shape, dtype=param.dtype) sharded = distribute_tensor( source, param.device_mesh, param.placements, src_data_rank=0, ) param.data.copy_(sharded) else: # Unsharded meta path (single process): load locally on rank0. param.data.copy_(reference[key].to(device=param.device, dtype=param.dtype)) def _pretrained_reference_state_dict(self): """Build a CPU state dict of the pretrained UMT5 + Wan2.2 submodules.""" with load_pretrained_submodules(): t5 = UMT5EncoderModel.from_pretrained( self.config.wan22_base, subfolder="text_encoder", torch_dtype=torch.float32 ) diff_dec = GEN_Wanx22(self.config) state_dict = {f"t5_text_encoder.{k}": v for k, v in t5.state_dict().items()} state_dict.update({f"diff_dec.{k}": v for k, v in diff_dec.state_dict().items()}) return state_dict def get_ignore_modules_in_mixed_precision(self): from diffusers.models.embeddings import TimestepEmbedding from diffusers.models.normalization import FP32LayerNorm return (TimestepEmbedding, FP32LayerNorm) def encode_prompt(self, input_ids, attention_mask): """Encode token ids into padded T5 embeddings `[B, max_len, hidden]`.""" 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.max_sequence_length)] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( [ torch.cat([u, u.new_zeros(self.max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds ], dim=0, ) return prompt_embeds def get_t5_text_embeddings(self, input_ids, attention_mask, input_lens): input_ids = input_ids.squeeze(0) attention_mask = attention_mask.squeeze(0) input_lens = input_lens.squeeze(0) # SequenceParallelCollator pads 1-D packed metadata (including # ``t5_input_lens``) to the Ulysses size with zeros. Drop those # padding entries and trim matching padded token tails before splitting. input_lens = input_lens[input_lens > 0] valid_token_len = int(input_lens.sum().item()) input_ids = input_ids[:valid_token_len] attention_mask = attention_mask[:valid_token_len] input_ids_list = torch.split(input_ids, input_lens.tolist()) attention_mask_list = torch.split(attention_mask, input_lens.tolist()) padded_input_ids = [] padded_attention_mask = [] seq_lens = torch.clamp(input_lens, max=self.max_sequence_length) for ids, mask in zip(input_ids_list, attention_mask_list): seq_len = ids.size(0) if seq_len < self.max_sequence_length: pad_len = self.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.max_sequence_length] mask = mask[: self.max_sequence_length] padded_input_ids.append(ids) padded_attention_mask.append(mask) with torch.no_grad(): prompt_embeds = self.t5_text_encoder( torch.stack(padded_input_ids), torch.stack(padded_attention_mask) ).last_hidden_state prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( [torch.cat([u, u.new_zeros(self.max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0, ) return [self.max_sequence_length] * len(input_lens), prompt_embeds.view(1, -1, prompt_embeds.size(-1)) def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, t5_input_lens: Optional[torch.Tensor] = None, input_vae_latents: Optional[torch.Tensor] = None, input_vae_rope: Optional[torch.Tensor] = None, vae_latents_mask: Optional[torch.Tensor] = None, vae_seqlen: Optional[torch.Tensor] = None, timesteps: Optional[torch.Tensor] = None, target_velocity: Optional[Any] = None, **kwargs, ): if hasattr(target_velocity, "tensor"): target_velocity = target_velocity.tensor batch_text_seqlen, batch_text_embs = self.get_t5_text_embeddings(input_ids, attention_mask, t5_input_lens) valid_sample_count = len(batch_text_seqlen) if vae_seqlen is not None: vae_seqlen = vae_seqlen.squeeze(0) vae_seqlen = vae_seqlen[vae_seqlen > 0].unsqueeze(0) if timesteps is not None: timesteps = timesteps.squeeze(0)[:valid_sample_count].unsqueeze(0) diff_loss = self.diff_dec( input_vae_latents=input_vae_latents, input_vae_rope=input_vae_rope, vae_latents_mask=vae_latents_mask, vae_seqlen=vae_seqlen, text_embs=batch_text_embs, batch_text_seqlen=batch_text_seqlen, timesteps=timesteps, target_velocity=target_velocity, ) return BerniniRendererOutputWithPast(diff_loss=diff_loss) @torch.no_grad() def sample( self, input_ids, attention_mask, uncond_input_ids, uncond_attention_mask, image_vae_latents=None, multi_video_vae_latents=None, multi_image_vae_latents=None, num_frames: int = 1, width: int = 832, height: int = 480, num_inference_steps: int = 50, guidance_mode: str = "rv2v", omega_vid: float = 3.0, omega_img: float = 3.0, omega_txt: float = 4.0, omega_scale: float = 0.75, flow_shift: float = 5.0, seed: int = 42, device="cuda", eta: float = 0.5, norm_threshold=(50.0, 50.0), momentum: float = -0.5, ): self.t5_text_encoder = self.t5_text_encoder.to(device) prompt_embeds = self.encode_prompt(input_ids, attention_mask) uncond_prompt_embeds = ( self.encode_prompt(uncond_input_ids, uncond_attention_mask) if uncond_input_ids is not None else None ) self.t5_text_encoder = self.t5_text_encoder.to("cpu") torch.cuda.empty_cache() return self.diff_dec.sample( prompt_embeds=prompt_embeds, uncond_prompt_embeds=uncond_prompt_embeds, image_vae_latents=image_vae_latents, multi_video_vae_latents=multi_video_vae_latents, multi_image_vae_latents=multi_image_vae_latents, num_frames=num_frames, width=width, height=height, num_inference_steps=num_inference_steps, guidance_mode=guidance_mode, omega_vid=omega_vid, omega_img=omega_img, omega_txt=omega_txt, omega_scale=omega_scale, flow_shift=flow_shift, seed=seed, device=device, eta=eta, norm_threshold=norm_threshold, momentum=momentum, )