Update app.py

app.py

@@ -137,13 +137,15 @@ def predict(text, speaker):
 
     ### ### ### 
     example = dataset['test'][11]
-    # speaker_embedding = torch.tensor(example["speaker_embeddings"]).unsqueeze(0)
-    speaker_embedding = torch.tensor(example["speaker_embeddings"]).unsqueeze(0).unsqueeze(0).to(device)
-    speaker_embedding = speaker_embedding.expand(-1, inputs["input_ids"].size(1), -1)
-    spectrogram = model.generate_speech(inputs["input_ids"].to(device), speaker_embedding)
+    speaker_embedding = torch.tensor(example["speaker_embeddings"]).unsqueeze(0).to(device)
+    
+    # Ensure the speaker_embedding has the correct dimensions
+    if speaker_embedding.dim() == 2:
+        speaker_embedding = speaker_embedding.unsqueeze(1).expand(-1, inputs["input_ids"].size(1), -1)
+    elif speaker_embedding.dim() == 3:
+        speaker_embedding = speaker_embedding.expand(-1, inputs["input_ids"].size(1), -1)
     
-    # speaker_embedding = torch.tensor(speaker_embedding).unsqueeze(0)
-    # spectrogram = model.generate_speech(inputs["input_ids"], speaker_embedding)
+    spectrogram = model.generate_speech(inputs["input_ids"].to(device), speaker_embedding)
     with torch.no_grad():
         speech = vocoder(spectrogram)
     # speech = model.generate_speech(input_ids, speaker_embedding, vocoder=vocoder)