import torch, librosa, laion_clap, functools, warnings
from huggingface_hub import hf_hub_download  # Added for robust HF downloading
warnings.filterwarnings('ignore')

original_load = torch.load
torch.load = functools.partial(original_load, weights_only=False)
original_load_state_dict = torch.nn.Module.load_state_dict

def tolerant_load_state_dict(self, state_dict, strict=True, assign=False):
    return original_load_state_dict(self, state_dict, strict=False, assign=assign)
torch.nn.Module.load_state_dict = tolerant_load_state_dict

# Initialize the CLAP module with fusion disabled as requested
model = laion_clap.CLAP_Module(enable_fusion=False)

# Securely download the weights through the official Hugging Face Hub API
# This bypasses the broken wget/urllib external download logic entirely
local_checkpoint = hf_hub_download(
    repo_id="lukewys/laion_clap", 
    filename="630k-audioset-best.pt"
)

# Load the locally cached file path directly into the model
model.load_ckpt(local_checkpoint) 

torch.load = original_load
torch.nn.Module.load_state_dict = original_load_state_dict

def get_clap_embedding(path):
    audio_data, _ = librosa.load(path, sr=48000)
    audio_data = audio_data.reshape(1, -1)
    with torch.no_grad():
        audio_embed = model.get_audio_embedding_from_data(x=audio_data)
    return audio_embed.flatten()