import types import torch import torch.nn as nn import torch.nn.functional as F from transformers import WhisperFeatureExtractor import whisper import torch try: torch.set_default_device("cpu") except Exception: pass import accelerate from ola.model.speech_encoder.beats.BEATs import BEATsConfig, BEATs class WhisperWrappedEncoder: @classmethod def load(cls, model_config): def replace_layer_norm(module): from whisper.model import LayerNorm for name, child in module.named_children(): if isinstance(child, LayerNorm): # Check if any parameter is a meta tensor has_meta = any(p.is_meta for p in child.parameters()) if has_meta: # For meta tensors, create new layer norm with same shape new_layer_norm = nn.LayerNorm(child.normalized_shape, eps=child.eps, elementwise_affine=child.elementwise_affine) else: old_params = child.state_dict() new_layer_norm = nn.LayerNorm(child.normalized_shape, eps=child.eps, elementwise_affine=child.elementwise_affine) new_layer_norm.load_state_dict(old_params) setattr(module, name, new_layer_norm) else: replace_layer_norm(child) # Load whisper model, handling both file paths and model names speech_encoder_path = model_config.speech_encoder # First try loading directly (works for both file paths and model names) try: encoder = whisper.load_model(name=speech_encoder_path, device='cpu').encoder except (NotImplementedError, RuntimeError) as e: if "meta tensor" in str(e): # Meta tensor issue - load model without device specification print(f"Detected meta tensor issue, using alternative loading approach...") # Load checkpoint directly to avoid device issues import os if os.path.isfile(speech_encoder_path): # Load from file checkpoint = torch.load(speech_encoder_path, map_location='cpu') # Create model from checkpoint from whisper.model import ModelDimensions, Whisper dims = ModelDimensions(**checkpoint["dims"]) model = Whisper(dims) # Load state dict without moving to device model.load_state_dict(checkpoint["model_state_dict"]) # Get encoder without device movement encoder = model.encoder else: # Try loading as model name without device import whisper.model as whisper_model # This is a fallback - may need adjustment based on actual model raise RuntimeError(f"Cannot load model {speech_encoder_path} due to meta tensor issues") else: raise e replace_layer_norm(encoder) return encoder class DualWrappedEncoder(nn.Module): def __init__(self, config): super().__init__() self.config = config self.whisper_model = self.load_whisper(config) self.beats_model = self.load_beats(config) def load_whisper(self, model_config): def replace_layer_norm(module): from whisper.model import LayerNorm for name, child in module.named_children(): if isinstance(child, LayerNorm): # Check if any parameter is a meta tensor has_meta = any(p.is_meta for p in child.parameters()) if has_meta: # For meta tensors, create new layer norm with same shape new_layer_norm = nn.LayerNorm(child.normalized_shape, eps=child.eps, elementwise_affine=child.elementwise_affine) else: old_params = child.state_dict() new_layer_norm = nn.LayerNorm(child.normalized_shape, eps=child.eps, elementwise_affine=child.elementwise_affine) new_layer_norm.load_state_dict(old_params) setattr(module, name, new_layer_norm) else: replace_layer_norm(child) # Load whisper model, handling both file paths and model names speech_encoder_path = model_config.speech_encoder # First try loading directly (works for both file paths and model names) # try: # breakpoint() import torch from whisper.model import Whisper, ModelDimensions # 1) Load checkpoint to CPU (weights are real tensors here) ckpt = torch.load("/data1/cxy/model/THUdyh/Ola-7b/large-v3.pt", map_location="cpu") dims = ModelDimensions(**ckpt["dims"]) # 2) Build the module skeleton, then MATERIALIZE tensors on CPU model = Whisper(dims) model.to_empty(device="cpu") # <-- crucial when meta is involved # 3) Load weights missing, unexpected = model.load_state_dict(ckpt["model_state_dict"], strict=True) model.eval() encoder = model.encoder print("missing:", missing) print("unexpected:", unexpected) # with accelerate.init_empty_weights(): # encoder = whisper.load_model(name=speech_encoder_path, device='cpu').encoder # state = torch.load("/data1/cxy/model/THUdyh/Ola-7b/large-v3.pt", map_location='cpu')['model_state_dict']['encoder.positional_embedding'] # breakpoint() # except (NotImplementedError, RuntimeError) as e: # if "meta tensor" in str(e): # # Meta tensor issue - load model without device specification # print(f"Detected meta tensor issue, using alternative loading approach...") # # Load checkpoint directly to avoid device issues # import os # if os.path.isfile(speech_encoder_path): # # Load from file # checkpoint = torch.load(speech_encoder_path, map_location='cpu') # # Create model from checkpoint # # breakpoint() # from whisper.model import ModelDimensions, Whisper # dims = ModelDimensions(**checkpoint["dims"]) # model = Whisper(dims) # # Load state dict without moving to device # model.load_state_dict(checkpoint["model_state_dict"]) # # Get encoder without device movement # encoder = model.encoder # else: # # Try loading as model name without device # import whisper.model as whisper_model # # This is a fallback - may need adjustment based on actual model # raise RuntimeError(f"Cannot load model {speech_encoder_path} due to meta tensor issues") # else: # raise e replace_layer_norm(encoder) return encoder def load_beats(self, model_config): beats_path = model_config.music_encoder print("Loading BEATs Model") beats_ckpt = torch.load(beats_path, map_location='cpu') beats_cfg = BEATsConfig(beats_ckpt['cfg']) beats = BEATs(beats_cfg) beats = beats.to_empty(device='cpu') # Load state dict beats.load_state_dict(beats_ckpt['model'], strict=True) # breakpoint() # 检查BEATs模型权重是否有问题 print("Checking BEATs model weights for NaN/Inf values...") nan_count = 0 inf_count = 0 for name, param in beats.named_parameters(): if torch.isnan(param).any(): print(f"ERROR - BEATs parameter {name} contains NaN values!") print(f"Debug - Parameter shape: {param.shape}") print(f"Debug - Parameter dtype: {param.dtype}") print(f"Debug - Parameter device: {param.device}") print(f"Debug - NaN count: {torch.isnan(param).sum().item()}") nan_count += 1 if torch.isinf(param).any(): print(f"ERROR - BEATs parameter {name} contains Inf values!") print(f"Debug - Parameter shape: {param.shape}") print(f"Debug - Inf count: {torch.isinf(param).sum().item()}") inf_count += 1 if nan_count > 0 or inf_count > 0: print(f"ERROR - Found NaN values in {nan_count} parameters and Inf values in {inf_count} parameters") print("This indicates the BEATs model weights are corrupted!") raise ValueError(f"BEATs model weights are corrupted: {nan_count} NaN parameters, {inf_count} Inf parameters") else: print("BEATs model weights are clean (no NaN or Inf values)") return beats def forward(self, x, raw_wav=None, audio_padding_mask=None): with torch.no_grad(): self.beats_model = self.beats_model.float() # Debug: Check input data print(f"Debug - Speech encoder input x range: {x.min().item()} to {x.max().item()}") print(f"Debug - Speech encoder input x has nan: {torch.isnan(x).any().item()}") print(f"Debug - Speech encoder input raw_wav range: {raw_wav.min().item()} to {raw_wav.max().item()}") print(f"Debug - Speech encoder input raw_wav has nan: {torch.isnan(raw_wav).any().item()}") # Check Whisper model print(f"Debug - Whisper model device: {next(self.whisper_model.parameters()).device}") print(f"Debug - Input x device: {x.device}") speech_embeds = self.whisper_model(x) print(f"Debug - Whisper output range: {speech_embeds.min().item()} to {speech_embeds.max().item()}") print(f"Debug - Whisper output has nan: {torch.isnan(speech_embeds).any().item()}") # Check BEATs model print(f"Debug - BEATs model device: {next(self.beats_model.parameters()).device}") print(f"Debug - Input raw_wav device: {raw_wav.device}") # Check if BEATs model has nan weights (should be fixed now) has_nan_weights = False for name, param in self.beats_model.named_parameters(): if torch.isnan(param).any(): print(f"WARNING - BEATs parameter {name} still has nan values after fix!") has_nan_weights = True if not has_nan_weights: print("Debug - BEATs model weights are clean (no nan)") try: # 详细检查BEATs模型输入 raw_wav_float = raw_wav.float() print(f"Debug - BEATs input raw_wav_float range: {raw_wav_float.min().item()} to {raw_wav_float.max().item()}") print(f"Debug - BEATs input raw_wav_float shape: {raw_wav_float.shape}") print(f"Debug - BEATs input raw_wav_float has nan: {torch.isnan(raw_wav_float).any().item()}") print(f"Debug - BEATs input raw_wav_float has inf: {torch.isinf(raw_wav_float).any().item()}") print(f"Debug - BEATs input raw_wav_float dtype: {raw_wav_float.dtype}") print(f"Debug - BEATs input raw_wav_float device: {raw_wav_float.device}") # 检查输入是否在BEATs期望的范围内 [-1, 1] if raw_wav_float.min().item() < -1.0 or raw_wav_float.max().item() > 1.0: print(f"WARNING - BEATs input out of expected range [-1, 1]! Clipping to valid range.") raw_wav_float = torch.clamp(raw_wav_float, -1.0, 1.0) print(f"Debug - After clipping range: {raw_wav_float.min().item()} to {raw_wav_float.max().item()}") else: print("Debug - BEATs input is within expected range [-1, 1]") if audio_padding_mask is not None: print(f"Debug - BEATs input padding_mask range: {audio_padding_mask.min().item()} to {audio_padding_mask.max().item()}") print(f"Debug - BEATs input padding_mask shape: {audio_padding_mask.shape}") print(f"Debug - BEATs input padding_mask has nan: {torch.isnan(audio_padding_mask).any().item()}") print(f"Debug - BEATs input padding_mask dtype: {audio_padding_mask.dtype}") else: print("Debug - BEATs input padding_mask is None") # 在调用BEATs之前，让我们检查模型状态 print("Debug - BEATs model training mode:", self.beats_model.training) print("Debug - BEATs model device:", next(self.beats_model.parameters()).device) # 让我们逐步调试BEATs的内部处理 print("Debug - Calling BEATs extract_features...") audio_embeds, _ = self.beats_model.extract_features(raw_wav_float, padding_mask=audio_padding_mask, feature_only=True) print(f"Debug - BEATs output range: {audio_embeds.min().item()} to {audio_embeds.max().item()}") print(f"Debug - BEATs output has nan: {torch.isnan(audio_embeds).any().item()}") print(f"Debug - BEATs output shape: {audio_embeds.shape}") print(f"Debug - BEATs output dtype: {audio_embeds.dtype}") # 检查BEATs输出是否有NaN值 if torch.isnan(audio_embeds).any(): print("ERROR - BEATs output contains NaN values!") print(f"Debug - NaN positions: {torch.isnan(audio_embeds).sum().item()} out of {audio_embeds.numel()}") print(f"Debug - NaN ratio: {torch.isnan(audio_embeds).float().mean().item():.4f}") # 不替换，直接抛出异常来找出根本原因 raise ValueError("BEATs model produced NaN values - this indicates a bug in the model or input data") except Exception as e: print(f"ERROR - BEATs model failed: {e}") print("Falling back to Whisper-only mode") # Create zero audio embeddings with the same shape as expected audio_embeds = torch.zeros(speech_embeds.shape[0], speech_embeds.shape[1], 1024, device=speech_embeds.device, dtype=speech_embeds.dtype) if audio_embeds.size(1) < speech_embeds.size(1): audio_embeds = F.pad(audio_embeds, (0, 0, 0, speech_embeds.size(1) - audio_embeds.size(1))) elif audio_embeds.size(1) > speech_embeds.size(1): speech_embeds = F.pad(speech_embeds, (0, 0, 0, audio_embeds.size(1) - speech_embeds.size(1))) speech_embeds = torch.cat((speech_embeds, audio_embeds), dim=-1) speech_embeds = speech_embeds.to(torch.bfloat16) # 最终检查是否有NaN值 if torch.isnan(speech_embeds).any(): print("ERROR - Final speech embeddings contain NaN values!") print(f"Debug - NaN positions: {torch.isnan(speech_embeds).sum().item()} out of {speech_embeds.numel()}") print(f"Debug - NaN ratio: {torch.isnan(speech_embeds).float().mean().item():.4f}") raise ValueError("Final speech embeddings contain NaN values - this indicates a bug in the speech encoder") return speech_embeds