Update app.py

app.py

@@ -5,11 +5,12 @@ import random
 import argparse
 import os
 import torch
+import torch.nn.functional as F
 import librosa
 from tqdm import tqdm
 from diffusers import DDIMScheduler
 from solospeech.model.solospeech.conditioners import SoloSpeech_TSE
-from solospeech.model.solospeech.conditioners import SoloSpeech_TSR
+# from solospeech.model.solospeech.conditioners import SoloSpeech_TSR
 from solospeech.scripts.solospeech.utils import save_audio
 import shutil
 from solospeech.vae_modules.autoencoder_wrapper import Autoencoder
@@ -61,6 +62,7 @@ parser = argparse.ArgumentParser()
 # pre-trained model path
 parser.add_argument('--eta', type=int, default=0)
 parser.add_argument("--num_infer_steps", type=int, default=200)
+parser.add_argument("--num_candidates", type=int, default=4)
 parser.add_argument('--sample-rate', type=int, default=16000)
 # random seed
 parser.add_argument('--random-seed', type=int, default=42, help="Fixed seed")
@@ -71,11 +73,11 @@ local_dir = snapshot_download(
     repo_id="OpenSound/SoloSpeech-models"
 )
 args.tse_config = os.path.join(local_dir, "config_extractor.yaml")
-args.tsr_config = os.path.join(local_dir, "config_tsr.yaml")
+# args.tsr_config = os.path.join(local_dir, "config_tsr.yaml")
 args.vae_config = os.path.join(local_dir, "config_compressor.json")
 args.autoencoder_path = os.path.join(local_dir, "compressor.ckpt")
 args.tse_ckpt = os.path.join(local_dir, "extractor.pt")
-args.tsr_ckpt = os.path.join(local_dir, "tsr.pt")
+# args.tsr_ckpt = os.path.join(local_dir, "tsr.pt")
 args.geco_ckpt = os.path.join(local_dir, "corrector.ckpt")
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
@@ -84,8 +86,8 @@ print(f"Device: {device}")
 print("Loading models...")
 with open(args.tse_config, 'r') as fp:
     args.tse_config = yaml.safe_load(fp)
-with open(args.tsr_config, 'r') as fp:
-    args.tsr_config = yaml.safe_load(fp)
+# with open(args.tsr_config, 'r') as fp:
+#     args.tsr_config = yaml.safe_load(fp)
 args.v_prediction = args.tse_config["ddim"]["v_prediction"]
 # load compressor
 autoencoder = Autoencoder(args.autoencoder_path, args.vae_config, 'stft_vae', quantization_first=True)
@@ -98,12 +100,12 @@ tse_model = SoloSpeech_TSE(
 ).to(device)
 tse_model.load_state_dict(torch.load(args.tse_ckpt)['model'])
 tse_model.eval()
-# load tsr model
-tsr_model = SoloSpeech_TSR(
-    args.tsr_config['diffwrap']['UDiT']
-).to(device)
-tsr_model.load_state_dict(torch.load(args.tsr_ckpt)['model'])
-tsr_model.eval()
+# # load tsr model
+# tsr_model = SoloSpeech_TSR(
+#     args.tsr_config['diffwrap']['UDiT']
+# ).to(device)
+# tsr_model.load_state_dict(torch.load(args.tsr_ckpt)['model'])
+# tsr_model.eval()
 # load corrector
 geco_model = ScoreModel.load_from_checkpoint(
     args.geco_ckpt,
@@ -113,7 +115,7 @@ geco_model.eval(no_ema=False)
 geco_model.cuda()
 # load sid model
 ecapatdnn_model = Encoder.from_hparams(source="yangwang825/ecapa-tdnn-vox2")
-cosine_sim = torch.nn.CosineSimilarity(dim=-1)
+# cosine_sim = torch.nn.CosineSimilarity(dim=-1)
 # load diffusion tools
 noise_scheduler = DDIMScheduler(**args.tse_config["ddim"]['diffusers'])
 # these steps reset dtype of noise_scheduler params
@@ -128,7 +130,7 @@ _ = noise_scheduler.add_noise(latents, noise, timesteps)
 
 
 @spaces.GPU
-def sample_diffusion(tse_model, tsr_model, autoencoder, std, scheduler, device,
+def sample_diffusion(tse_model, autoencoder, std, scheduler, device,
                      mixture=None, reference=None, lengths=None, reference_lengths=None, 
                      ddim_steps=50, eta=0, seed=2025
                      ):
@@ -136,7 +138,7 @@ def sample_diffusion(tse_model, tsr_model, autoencoder, std, scheduler, device,
         generator = torch.Generator(device=device).manual_seed(seed)
         scheduler.set_timesteps(ddim_steps)
         tse_pred = torch.randn(mixture.shape, generator=generator, device=device)
-        tsr_pred = torch.randn(mixture.shape, generator=generator, device=device)
+        # tsr_pred = torch.randn(mixture.shape, generator=generator, device=device)
 
         for t in scheduler.timesteps:
             tse_pred = scheduler.scale_model_input(tse_pred, t)
@@ -151,22 +153,22 @@ def sample_diffusion(tse_model, tsr_model, autoencoder, std, scheduler, device,
             tse_pred = scheduler.step(model_output=model_output, timestep=t, sample=tse_pred,
                                     eta=eta, generator=generator).prev_sample
         
-        for t in scheduler.timesteps:
-            tsr_pred = scheduler.scale_model_input(tsr_pred, t)
-            model_output, _ = tsr_model(
-                x=tsr_pred, 
-                timesteps=t, 
-                mixture=mixture, 
-                reference=tse_pred, 
-                x_len=lengths, 
-                )
-            tsr_pred = scheduler.step(model_output=model_output, timestep=t, sample=tsr_pred,
-                                    eta=eta, generator=generator).prev_sample
+        # for t in scheduler.timesteps:
+        #     tsr_pred = scheduler.scale_model_input(tsr_pred, t)
+        #     model_output, _ = tsr_model(
+        #         x=tsr_pred, 
+        #         timesteps=t, 
+        #         mixture=mixture, 
+        #         reference=tse_pred, 
+        #         x_len=lengths, 
+        #         )
+        #     tsr_pred = scheduler.step(model_output=model_output, timestep=t, sample=tsr_pred,
+        #                             eta=eta, generator=generator).prev_sample
 
         tse_pred = autoencoder(embedding=tse_pred.transpose(2,1), std=std).squeeze(1)
-        tsr_pred = autoencoder(embedding=tsr_pred.transpose(2,1), std=std).squeeze(1)
+        # tsr_pred = autoencoder(embedding=tsr_pred.transpose(2,1), std=std).squeeze(1)
 
-        return tse_pred, tsr_pred
+        return tse_pred
 
 @spaces.GPU
 def tse(test_wav, enroll_wav):
@@ -179,19 +181,21 @@ def tse(test_wav, enroll_wav):
     with torch.no_grad():
         # compressor
         reference, _ = autoencoder(audio=reference.unsqueeze(1))
-        reference_lengths = torch.LongTensor([reference.shape[-1]]).to(device)
+        reference_lengths = torch.LongTensor([reference.shape[-1]] * args.num_candidates).to(device)
         mixture_input = torch.tensor(mixture).unsqueeze(0).to(device)
         mixture_wav = mixture_input
         mixture_input, std = autoencoder(audio=mixture_input.unsqueeze(1))
-        lengths = torch.LongTensor([mixture_input.shape[-1]]).to(device)   
+        lengths = torch.LongTensor([mixture_input.shape[-1]] * args.num_candidates).to(device)   
         # extractor
-        tse_pred, tsr_pred = sample_diffusion(tse_model, tsr_model, autoencoder, std, noise_scheduler, device, mixture_input.transpose(2,1), reference.transpose(2,1), lengths, reference_lengths, ddim_steps=args.num_infer_steps, eta=args.eta, seed=args.random_seed)
-        ecapatdnn_embedding1 = ecapatdnn_model.encode_batch(tse_pred.squeeze()).squeeze()
-        ecapatdnn_embedding2 = ecapatdnn_model.encode_batch(tsr_pred.squeeze()).squeeze()
-        ecapatdnn_embedding3 = ecapatdnn_model.encode_batch(torch.tensor(reference_wav)).squeeze()
-        sim1 = cosine_sim(ecapatdnn_embedding1, ecapatdnn_embedding3).item()
-        sim2 = cosine_sim(ecapatdnn_embedding2, ecapatdnn_embedding3).item()
-        pred = tse_pred if sim1 > sim2 else tsr_pred
+        mixture_input = mixture_input.repeat(args.num_candidates, 1, 1)
+        reference = reference.repeat(args.num_candidates, 1, 1)
+        tse_pred = sample_diffusion(tse_model, autoencoder, std, noise_scheduler, device, mixture_input.transpose(2,1), reference.transpose(2,1), lengths, reference_lengths, ddim_steps=args.num_infer_steps, eta=args.eta, seed=args.random_seed)
+        tse_pred = sample_diffusion(tse_model, autoencoder, std, noise_scheduler, device, mixture_input.transpose(2,1), reference.transpose(2,1), lengths, reference_lengths, ddim_steps=args.num_infer_steps, eta=args.eta, seed=args.random_seed)
+        ecapatdnn_embedding_pred = ecapatdnn_model.encode_batch(tse_pred).squeeze()
+        ecapatdnn_embedding_ref = ecapatdnn_model.encode_batch(torch.tensor(reference_wav)).squeeze()
+        cos_sims = F.cosine_similarity(ecapatdnn_embedding_pred, ecapatdnn_embedding_ref.unsqueeze(0), dim=1)
+        _, max_idx = torch.max(cos_sims, dim=0)
+        pred = tse_pred[max_idx].unsqueeze(0)
         # corrector
         min_leng = min(pred.shape[-1], mixture_wav.shape[-1])
         x = pred[...,:min_leng]