Spaces:
Running on Zero
Running on Zero
| import sys | |
| import torchaudio | |
| import torch | |
| import torch.nn.functional as F | |
| def load_fx_encoder_plusplus_2048(model_args, device, *args, **kwargs): | |
| from utils.feature_extractors.fx_encoder_plus_plus import load_model | |
| assert model_args is not None, "model_args must be provided for fx_encoder type" | |
| ckpt_path=model_args.ckpt_path | |
| model=load_model( | |
| model_path=ckpt_path, | |
| device=device, | |
| ) | |
| def effects_encoder_fn(x): | |
| assert x.ndim == 3, f"Input tensor x must be 2D, got {x.ndim}D" | |
| assert x.shape[1] == 2, f"Input tensor x must have 2 channels, got {x.shape[1]} channels" | |
| emb=model.fx_encoder(x) | |
| emb=emb["embedding"] # Extract the embedding from the dictionary | |
| return emb | |
| return lambda x: effects_encoder_fn(x) | |
| def add_isotropic_noise(z, sigma=0.1): | |
| """ | |
| z: [..., D] normalized embeddings (e.g., from CLAP or a regressor) | |
| sigma: scale of noise to inject | |
| Returns: z with orthogonal Gaussian noise added | |
| """ | |
| n=torch.randn_like(z) # isotropic noise | |
| z_noisy = F.normalize(z + sigma * n, dim=-1) | |
| return z_noisy | |
| def load_CLAP(model_args, device, *args, **kwargs): | |
| #original_path = sys.path.copy() | |
| from utils.laion_clap.hook import CLAP_Module | |
| model= CLAP_Module(enable_fusion=False, amodel= 'HTSAT-base') | |
| #sys.path = original_path | |
| print("checkpoint",model_args.ckpt_path) | |
| #print current sys.path | |
| print("sys.path", sys.path) | |
| model.load_ckpt(model_args.ckpt_path) | |
| model.to(device) | |
| normalize = model_args.normalize | |
| if model_args.use_adaptor: | |
| if model_args.adaptor_type == "MLP_CLAP_regressor": | |
| from networks.MLP_CLAP_regressor import MLP_CLAP_regressor | |
| adaptor=MLP_CLAP_regressor() | |
| ckpt=torch.load(model_args.adaptor_checkpoint, map_location=device, weights_only=False) | |
| adaptor.load_state_dict(ckpt["network"], strict=True) | |
| adaptor.to(device) | |
| def clap_fn(x, type=None): | |
| B, C, T = x.shape | |
| if C > 1: | |
| x= x.mean(dim=1, keepdim=True) # Convert to mono if stereo | |
| with torch.no_grad(): | |
| x=torchaudio.functional.resample(x, orig_freq=44100, new_freq=48000) | |
| x= x.squeeze(1) # Remove channel dimension for CLAP | |
| emb=model.get_audio_embedding_from_data(x,use_tensor=True) | |
| if type is not None: | |
| if type == "wet": | |
| #print("wet mode") | |
| if model_args.use_adaptor: | |
| emb= adaptor(emb) # Apply the adaptor if specified | |
| if model_args.add_noise: | |
| emb= torch.nn.functional.normalize(emb, p=2, dim=-1) # Normalize before adding noise | |
| emb = add_isotropic_noise(emb, sigma=model_args.noise_sigma) | |
| # Normalize the embeddings | |
| if normalize: | |
| emb = torch.nn.functional.normalize(emb, p=2, dim=-1) | |
| return emb | |
| return lambda x, type: clap_fn(x, type=type) | |