Update app.py

app.py

@@ -260,55 +260,135 @@ def inference_by_model(text: str, audio_file: str, save_path: str) -> str:
 
 #_______________generate audios and folder structure_______________________
 
-async def generate_story_audios(story: StoryCreationDTO, base_output: str):
+# async def generate_story_audios(story: StoryCreationDTO, base_output: str):
+#     """
+#     Generates audio files and folders for the entire story
+#     """
+#     story_dir = Path(base_output) / story.storyId
+#     story_dir.mkdir(parents=True, exist_ok=True)
+
+#     for chapter in story.chapters:
+#         chapter_dir = story_dir / chapter.chapterId
+#         chapter_dir.mkdir(exist_ok=True)
+
+#         # --- Chapter title audio ---
+#         prosody_file_title = await download_file_from_url(chapter.title.prosodyReference)
+#         title_save_path = chapter_dir / "title.wav"
+        
+#         tagged_text_title = generate_tagged_text(
+#             chapter.title.sentence, 
+#             chapter.title.emotion,   
+#             chapter.title.intensity    
+#         )
+        
+        # title_generated_audio_path = inference_by_model(
+        #     text=tagged_text_title,
+        #     audio_file=prosody_file_title,
+        #     save_path=title_save_path
+        # )
+        # # os.remove(prosody_file_title)
+
+        # for scene in chapter.scenes:
+        #     await download_scene_files(scene)
+        #     scene_dir = chapter_dir / scene.sceneId
+        #     scene_dir.mkdir(exist_ok=True)
+
+        #     # --- Sentences audio ---
+        #     for sentence in scene.sentences:
+        #         # Download the prosody reference audio from Supabase
+        #         prosody_file = download_cache[sentence.prosodyReference]
+        #         sentence_save_path = scene_dir / f"{sentence.sentenceId}.wav"
+        #         tagged_text = generate_tagged_text(
+        #             sentence.sentence, 
+        #             sentence.emotion,   
+        #             sentence.intensity    
+        #         )
+        #         sentence_generated_audio_path = inference_by_model(
+        #             text=tagged_text,
+        #             audio_file=prosody_file,
+        #             save_path=sentence_save_path
+        #         )
+        #         # os.remove(prosody_file)
+
+import asyncio
+from pathlib import Path
+
+async def generate_story_audios_async(story: StoryCreationDTO, base_output: str, max_concurrent_gpu: int = 1):
     """
-    Generates audio files and folders for the entire story
+    Generates audio files for the story while overlapping GPU inference and disk writes.
+    max_concurrent_gpu: semaphore to limit simultaneous GPU usage (1 if GPU is the bottleneck)
     """
     story_dir = Path(base_output) / story.storyId
     story_dir.mkdir(parents=True, exist_ok=True)
 
+    # Semaphore ensures we don't overload GPU
+    gpu_semaphore = asyncio.Semaphore(max_concurrent_gpu)
+
+    async def process_sentence(chapter_dir: Path, scene: SceneDto, sentence: SentenceDto):
+        async with gpu_semaphore:
+            # Get prosody file from cache
+            prosody_file = download_cache[sentence.prosodyReference]
+
+            sentence_save_path = chapter_dir / scene.sceneId / f"{sentence.sentenceId}.wav"
+            Path(sentence_save_path).parent.mkdir(parents=True, exist_ok=True)
+
+            tagged_text = generate_tagged_text(
+                sentence.sentence,
+                sentence.emotion,
+                sentence.intensity
+            )
+
+            # Run GPU inference in a thread pool to avoid blocking event loop
+            loop = asyncio.get_event_loop()
+            generated_path = await loop.run_in_executor(
+                None,
+                inference_by_model,
+                tagged_text,
+                prosody_file,
+                str(sentence_save_path)
+            )
+
+            return generated_path
+
+    # Prepare tasks for chapters
+    chapter_tasks = []
     for chapter in story.chapters:
         chapter_dir = story_dir / chapter.chapterId
         chapter_dir.mkdir(exist_ok=True)
 
-        # --- Chapter title audio ---
-        prosody_file_title = await download_file_from_url(chapter.title.prosodyReference)
+        # --- Chapter title ---
+        title_prosody = download_cache[chapter.title.prosodyReference]
         title_save_path = chapter_dir / "title.wav"
-        
-        tagged_text_title = generate_tagged_text(
-            chapter.title.sentence, 
-            chapter.title.emotion,   
-            chapter.title.intensity    
+        tagged_title = generate_tagged_text(
+            chapter.title.sentence,
+            chapter.title.emotion,
+            chapter.title.intensity
         )
-        
-        title_generated_audio_path = inference_by_model(
-            text=tagged_text_title,
-            audio_file=prosody_file_title,
-            save_path=title_save_path
+
+        # Run title generation immediately
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(
+            None,
+            inference_by_model,
+            tagged_title,
+            title_prosody,
+            str(title_save_path)
         )
-        # os.remove(prosody_file_title)
 
+        # --- Scenes ---
+        scene_tasks = []
         for scene in chapter.scenes:
+            # Download scene files (prosody, sfx, bg) concurrently
             await download_scene_files(scene)
-            scene_dir = chapter_dir / scene.sceneId
-            scene_dir.mkdir(exist_ok=True)
 
-            # --- Sentences audio ---
             for sentence in scene.sentences:
-                # Download the prosody reference audio from Supabase
-                prosody_file = download_cache[sentence.prosodyReference]
-                sentence_save_path = scene_dir / f"{sentence.sentenceId}.wav"
-                tagged_text = generate_tagged_text(
-                    sentence.sentence, 
-                    sentence.emotion,   
-                    sentence.intensity    
-                )
-                sentence_generated_audio_path = inference_by_model(
-                    text=tagged_text,
-                    audio_file=prosody_file,
-                    save_path=sentence_save_path
-                )
-                # os.remove(prosody_file)
+                scene_tasks.append(process_sentence(chapter_dir, scene, sentence))
+
+        # Run all sentences in this chapter concurrently but limited by GPU semaphore
+        chapter_tasks.append(asyncio.gather(*scene_tasks))
+
+    # Wait for all chapters to complete
+    await asyncio.gather(*chapter_tasks)
 
 #_______________ Concatenating the generated audios to make the final story (post-processing)_______________________