Add variable audio lengths

app.py

@@ -95,8 +95,10 @@ def generate_base(story):
     speech_output = model.generate(input_ids=description_tokens.input_ids,
                                    prompt_input_ids=story_tokens.input_ids,
                                    attention_mask=description_tokens.attention_mask,
-                                   prompt_attention_mask=story_tokens.attention_mask)
-    speech_output = [output.cpu().numpy() for output in speech_output]
+                                   prompt_attention_mask=story_tokens.attention_mask,
+                                  return_dict_in_generate=True,
+                                  )
+    speech_output = [output.cpu().numpy()[:output_length] for (output, output_length) in zip(speech_output.sequences, speech_output.audios_length)]
     return None, None, speech_output