update app.py

app.py

@@ -32,7 +32,7 @@ if torch.cuda.is_available():
 setup_eval_logging()
 OUTPUT_DIR = Path("./output/gradio")
 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
-NUM_SAMPLE=8
+NUM_SAMPLE=5
 snapshot_download(repo_id="google/flan-t5-large")
 a=AutoModel.from_pretrained('bert-base-uncased')
 b=AutoModel.from_pretrained('roberta-base')
@@ -193,8 +193,6 @@ def generate_audio_gradio(
         log.info("Using FlowMatching for generation.")
         generation_func = generate_fm
         sampler_arg_name = "fm"
-    all_audios=[]
-    all_scores=[]
     audios = generation_func(
         [prompt]*NUM_SAMPLE,
         negative_text=[negative_prompt]*NUM_SAMPLE,
@@ -204,17 +202,12 @@ def generate_audio_gradio(
         cfg_strength=cfg_strength,
         **{sampler_arg_name: sampler},
     )
-    for i in range(NUM_SAMPLE):
-        audio = audios.float().cpu()[i]
-        text_embed = laion_clap_model.get_text_embedding(prompt, use_tensor=True).squeeze()
-        audio_embed = laion_clap_model.get_audio_embedding_from_data(audio, use_tensor=True).squeeze()
-        score = torch.cosine_similarity(text_embed,
-                                    audio_embed,
-                                    dim=-1).mean()
-        all_audios.append(audio)
-        all_scores.append(score)
-    winner_idx = torch.argmax(torch.tensor(all_scores)).item()
-    audio=all_audios[winner_idx]
+    text_embed = laion_clap_model.get_text_embedding(prompt, use_tensor=True).squeeze()
+    audio_embed = laion_clap_model.get_audio_embedding_from_data(audios[:,0,:].float().cpu(), use_tensor=True).squeeze()
+    scores = torch.cosine_similarity(text_embed.expand(audio_embed.shape[0], -1),
+                                audio_embed,
+                                dim=-1)
+    audio=audios[torch.argmax(scores).item()].float().cpu()
     safe_prompt = (
         "".join(c for c in prompt if c.isalnum() or c in (" ", "_"))
         .rstrip()