Update modal_tts.py

modal_tts.py

@@ -1,43 +1,73 @@
 import io
-
 import modal
+from pydantic import BaseModel
+
+# Request model for JSON body
+class TTSRequest(BaseModel):
+    prompt: str
+    use_music: bool = True
 
-image = modal.Image.debian_slim(python_version="3.12").pip_install(
-    "chatterbox-tts==0.1.1", "fastapi[standard]"
+# Shared image for all Modal functions
+image = (
+    modal.Image.debian_slim(python_version="3.10")
+    .pip_install("chatterbox-tts==0.1.1", "fastapi[standard]", "pydub", "ffmpeg")
+    .apt_install("ffmpeg")  # Required by pydub
 )
-app = modal.App("chatterbox-api-example", image=image)
 
-with image.imports():
-    import torchaudio as ta
-    from chatterbox.tts import ChatterboxTTS
-    from fastapi.responses import StreamingResponse
+# Attach Volume
+volume = modal.Volume.from_name("background-music")
 
+# Modal App
+app = modal.App("NewsShots_TTS_", image=image)
 
-@app.cls(gpu="a10g", scaledown_window=60 * 5, enable_memory_snapshot=True)
-@modal.concurrent(max_inputs=10)
-class Chatterbox:
+# TTS Class
+@app.cls(gpu="a10g", scaledown_window=60 * 10, volumes={"/music": volume})
+class ChatterboxWithMusic:
     @modal.enter()
     def load(self):
-        self.model = ChatterboxTTS.from_pretrained(device="cuda")
-
-    @modal.fastapi_endpoint(docs=True, method="POST")
-    def generate(self, prompt: str):
-        # Generate audio waveform from the input text
-        wav = self.model.generate(prompt)
+        from chatterbox.tts import ChatterboxTTS
+        from pydub import AudioSegment
+        self.tts_model = ChatterboxTTS.from_pretrained(device="cuda")
+        self.AudioSegment = AudioSegment
 
-        # Create an in-memory buffer to store the WAV file
+    @modal.fastapi_endpoint(method="POST")
+    def generate(self, request: TTSRequest):
+        import torchaudio
+        from fastapi.responses import StreamingResponse
+        
+        # Extract data from request body
+        prompt = request.prompt
+        use_music = request.use_music
+        
+        # Generate speech from prompt
+        wav_tensor = self.tts_model.generate(prompt)
         buffer = io.BytesIO()
-
-        # Save the generated audio to the buffer in WAV format
-        # Uses the model's sample rate and WAV format
-        ta.save(buffer, wav, self.model.sr, format="wav")
-
-        # Reset buffer position to the beginning for reading
+        torchaudio.save(buffer, wav_tensor, self.tts_model.sr, format="wav")
         buffer.seek(0)
-
-        # Return the audio as a streaming response with appropriate MIME type.
-        # This allows for browsers to playback audio directly.
-        return StreamingResponse(
-            io.BytesIO(buffer.read()),
-            media_type="audio/wav",
-        )
+        
+        # Convert to AudioSegment
+        tts_audio = self.AudioSegment.from_file(buffer, format="wav")
+        
+        # Try to load background music
+        if use_music:
+            try:
+                with open("/music/music/download.mp3", "rb") as f:
+                    music_bytes = f.read()
+                background = self.AudioSegment.from_file(io.BytesIO(music_bytes))
+                background = background - 15
+                if len(background) < len(tts_audio):
+                    background *= (len(tts_audio) // len(background) + 1)
+                background = background[:len(tts_audio)]
+                final_audio = tts_audio.overlay(background)
+            except FileNotFoundError:
+                final_audio = tts_audio
+        else:
+            final_audio = tts_audio
+        
+        # Export mixed audio to buffer
+        final_buffer = io.BytesIO()
+        final_audio.export(final_buffer, format="mp3")
+        final_buffer.seek(0)
+        
+        # Stream as response
+        return StreamingResponse(final_buffer, media_type="audio/mpeg")