Add

app.py

@@ -171,7 +171,7 @@ def numpy_to_mp3(audio_array, sampling_rate):
     # Normalize audio_array if it's floating-point
     if np.issubdtype(audio_array.dtype, np.floating):
         max_val = np.max(np.abs(audio_array))
-        audio_array = (audio_array / max_val) * 32767  # Normalize to 16-bit range
+        audio_array = (audio_array / max_val) * 32767 # Normalize to 16-bit range
         audio_array = audio_array.astype(np.int16)
 
     # Create an audio segment from the numpy array
@@ -196,9 +196,10 @@ sampling_rate = model.audio_encoder.config.sampling_rate
 frame_rate = model.audio_encoder.config.frame_rate
 
 import random
+import datetime
 
 @spaces.GPU
-def generate_base(subject, setting, ):
+def generate_base(subject, setting):
 
     messages = [{"role": "sytem", "content": ("You are an award-winning children's bedtime story author lauded for your inventive stories."
                                               "You want to write a bed time story for your child. They will give you the subject and setting "
@@ -209,34 +210,63 @@ def generate_base(subject, setting, ):
     response = client.chat_completion(messages, max_tokens=2048, seed=random.randint(1, 5000))
     gr.Info("Story Generated", duration=3)
     story = response.choices[0].message.content
-    yield story, None
 
     model_input = story.replace("\n", " ").strip()
-    model_input = nltk.sent_tokenize(model_input)
+    model_input_tokens = nltk.sent_tokenize(model_input)
 
     play_steps_in_s = 4.0
     play_steps = int(frame_rate * play_steps_in_s)
 
+    gr.Info("Generating Audio")
     description = "Jenny speaks at an average pace with a calm delivery in a very confined sounding environment with clear audio quality."
-    description = [description for _ in range(len(model_input))]
-    description_tokens = tokenizer(description, return_tensors="pt").input_ids.to(device)
+    story_tokens = tokenizer(model_input_tokens, return_tensors="pt", padding=True).input_ids.to(device)
+    description_tokens = tokenizer([description for _ in range(len(model_input_tokens))], return_tensors="pt").input_ids.to(device)
+    speech_output = model.generate(input_ids=description_tokens, prompt_input_ids=story_tokens)
+    speech_output = [output.cpu().numpy() for output in speech_output]
+    gr.Info("Generated Audio")
+    return None, None, {"audio": speech_output, "text": model_input}
+    
+def stream_audio(state):
+    speech_output = state["audio"]
+    sentences = state["text"]
+
+    gr.Info("Reading Story")
+
+    story = ""
+    for sentence, new_audio in zip(sentences, speech_output):
+#         print(f"i, j, time: {i}, {j} {datetime.datetime.now()}")
+        print(f"Sample of length: {round(new_audio.shape[0] / sampling_rate, 2)} seconds")
+        story += f"{sentence}\n"
+        yield story,  numpy_to_mp3(new_audio, sampling_rate=sampling_rate)
 
+    # BATCH_SIZE = 4
     # for i in range(0, len(model_input), BATCH_SIZE):
     #     inputs = model_input[i:min(i + BATCH_SIZE, len(model_input))]
+    #     story_tokens = tokenizer(inputs, return_tensors="pt", padding=True).input_ids.to(device)
+    #     description_tokens = tokenizer([description for _ in range(len(inputs))], return_tensors="pt").input_ids.to(device)
+    #     speech_output = model.generate(input_ids=description_tokens, prompt_input_ids=story_tokens)
+        
+    #     speech_output = [output.cpu().numpy() for output in speech_output]
+    #     for j, new_audio in enumerate(speech_output):
+    #         if i + j == 0:
+    #             gr.Info("Reading story", duration=3)
+    #         print(f"i, j, time: {i}, {j} {datetime.datetime.now()}")
+    #         print(f"Sample of length: {round(new_audio.shape[0] / sampling_rate, 2)} seconds")
+    #         yield story,  numpy_to_mp3(new_audio, sampling_rate=sampling_rate)
         
     #     if len(inputs) != 0:
     # input_ids = tokenizer(description, return_tensors="pt").input_ids.to(device)
-    story = tokenizer(model_input, return_tensors="pt", padding=True).input_ids.to(device)
+    # story = tokenizer(model_input, return_tensors="pt", padding=True).input_ids.to(device)
 
-    speech_output = model.generate(input_ids=description_tokens, prompt_input_ids=story)
+    # speech_output = model.generate(input_ids=description_tokens, prompt_input_ids=story)
         
-    speech_output = [output.cpu().numpy() for output in speech_output]
+    # speech_output = [output.cpu().numpy() for output in speech_output]
 
-    for i, new_audio in enumerate(speech_output):
-        if i == 0:
-            gr.Info("Reading story", duration=3)
-        print(f"Sample of length: {round(new_audio.shape[0] / sampling_rate, 2)} seconds")
-        yield story, numpy_to_mp3(new_audio, sampling_rate=sampling_rate)
+    # for i, new_audio in enumerate(speech_output):
+    #     if i == 0:
+    #         gr.Info("Reading story", duration=3)
+    #     print(f"Sample of length: {round(new_audio.shape[0] / sampling_rate, 2)} seconds")
+    #     yield story, numpy_to_mp3(new_audio, sampling_rate=sampling_rate)
     
     # print(f"{i}-th part generated")
     # pieces += [*speech_output, silence.copy()]
@@ -286,7 +316,8 @@ with gr.Blocks() as block:
 
     inputs = [subject, setting]
     outputs = [story, audio_out]
-    run_button.click(fn=generate_base, inputs=inputs, outputs=outputs)
+    state = gr.State()
+    run_button.click(fn=generate_base, inputs=inputs, outputs=[story, audio_out, state]).success(stream_audio, inputs=state, outputs=outputs)
 
 block.queue()
 block.launch(share=True)