import torch import os import json from safetensors.torch import load_file from ovi.modules.fusion import FusionModel from ovi.modules.t5 import T5EncoderModel from ovi.modules.vae2_2 import Wan2_2_VAE from ovi.modules.mmaudio.features_utils import FeaturesUtils from ovi.distributed_comms.util import get_world_size, get_local_rank, get_global_rank def init_wan_vae_2_2(ckpt_dir, rank=0): vae_config = {} vae_config['device'] = rank vae_pth = os.path.join(ckpt_dir, "Wan2.2-TI2V-5B/Wan2.2_VAE.pth") vae_config['vae_pth'] = vae_pth vae_model = Wan2_2_VAE(**vae_config) return vae_model def init_mmaudio_vae(ckpt_dir, rank=0): vae_config = {} vae_config['mode'] = '16k' vae_config['need_vae_encoder'] = True tod_vae_ckpt = os.path.join(ckpt_dir, "MMAudio/ext_weights/v1-16.pth") bigvgan_vocoder_ckpt = os.path.join(ckpt_dir, "MMAudio/ext_weights/best_netG.pt") vae_config['tod_vae_ckpt'] = tod_vae_ckpt vae_config['bigvgan_vocoder_ckpt'] = bigvgan_vocoder_ckpt vae = FeaturesUtils(**vae_config).to(rank) # vae = FeaturesUtils(**vae_config) return vae def init_fusion_score_model_ovi( rank: int = 0, meta_init=False, av2av_edit=False, concat_edit_source_latents=True, has_video=True, has_audio=True, use_siga=False, ): # import pdb; pdb.set_trace() if (has_video): video_config = "ovi/configs/model/dit/video.json" assert os.path.exists(video_config), f"{video_config} does not exist" with open(video_config) as f: video_config = json.load(f) else: video_config = None if (has_audio): audio_config = "ovi/configs/model/dit/audio.json" assert os.path.exists(audio_config), f"{audio_config} does not exist" with open(audio_config) as f: audio_config = json.load(f) else: audio_config = None # assert os.path.exists(video_config), f"{video_config} does not exist" # assert os.path.exists(audio_config), f"{audio_config} does not exist" if meta_init: with torch.device("meta"): fusion_model = FusionModel( video_config, audio_config, av2av_edit=av2av_edit, concat_edit_source_latents=concat_edit_source_latents, use_siga=use_siga, ) else: fusion_model = FusionModel( video_config, audio_config, av2av_edit=av2av_edit, concat_edit_source_latents=concat_edit_source_latents, use_siga=use_siga, ) params_all = sum(p.numel() for p in fusion_model.parameters()) if rank == 0: print( f"Score model (Fusion) all parameters:{params_all}" ) return fusion_model, video_config, audio_config def init_text_model(ckpt_dir, rank, cpu_offload=False): wan_dir = os.path.join(ckpt_dir, "Wan2.2-TI2V-5B") text_encoder_path = os.path.join(wan_dir, "models_t5_umt5-xxl-enc-bf16.pth") text_tokenizer_path = os.path.join(wan_dir, "google/umt5-xxl") text_encoder = T5EncoderModel( text_len=512, dtype=torch.bfloat16, device=rank, checkpoint_path=text_encoder_path, tokenizer_path=text_tokenizer_path, cpu_offload=cpu_offload, shard_fn=None) return text_encoder def _maybe_adapt_patch_embedding_weight(key, value, target_shape): if not (key.endswith("patch_embedding.weight") or key.endswith("patch_embedding.0.weight")): return None, None if len(value.shape) != len(target_shape): return None, None if value.shape[0] != target_shape[0] or value.shape[2:] != target_shape[2:]: return None, None ckpt_in_dim = value.shape[1] target_in_dim = target_shape[1] if ckpt_in_dim == target_in_dim * 2: return value[:, :target_in_dim, ...].contiguous(), "cropped leading input channels" return None, None def _adapt_state_dict_for_model(model, state_dict): model_state_dict = model.state_dict() adapted_state_dict = {} adapted_messages = [] skipped_messages = [] for key, value in state_dict.items(): target_value = model_state_dict.get(key) if target_value is None or not hasattr(value, "shape"): adapted_state_dict[key] = value continue if value.shape == target_value.shape: adapted_state_dict[key] = value continue adapted_value, note = _maybe_adapt_patch_embedding_weight(key, value, target_value.shape) if adapted_value is not None: adapted_state_dict[key] = adapted_value adapted_messages.append( f"{key}: {tuple(value.shape)} -> {tuple(target_value.shape)} ({note})" ) continue skipped_messages.append( f"{key}: checkpoint {tuple(value.shape)} != model {tuple(target_value.shape)}" ) return adapted_state_dict, adapted_messages, skipped_messages def load_fusion_checkpoint(model, checkpoint_path, from_meta=False): # import pdb; pdb.set_trace() if checkpoint_path and os.path.exists(checkpoint_path): if checkpoint_path.endswith(".safetensors"): df = load_file(checkpoint_path, device="cpu") elif checkpoint_path.endswith(".pt"): try: df = torch.load(checkpoint_path, map_location="cpu", weights_only=False) df = df['module'] if 'module' in df else df except Exception as e: df = torch.load(checkpoint_path, map_location="cpu", weights_only=True) df = df['app']['model'] else: raise RuntimeError("We only support .safetensors and .pt checkpoints") df, adapted_messages, skipped_messages = _adapt_state_dict_for_model(model, df) missing, unexpected = model.load_state_dict(df, strict=False, assign=from_meta) if (get_local_rank() == 0): print("****************************************************") if adapted_messages: print("adapted keys:") for message in adapted_messages: print(message) if skipped_messages: print("skipped mismatched keys:") for message in skipped_messages: print(message) print(f"missing keys: {missing}") print(f"unexpected keys: {unexpected}") print("****************************************************") del df import gc gc.collect() print(f"Successfully loaded fusion checkpoint from {checkpoint_path}") else: raise RuntimeError("{checkpoint=} does not exists'")