Create download_models.py

download_models.py

@@ -0,0 +1,57 @@
+"""
+Model download script for Bark TTS and Whisper ASR
+This script downloads and caches the models during Docker build
+"""
+import torch
+from transformers import AutoProcessor, BarkModel
+import whisper
+import os
+
+def download_bark_model():
+    """Download Bark TTS model"""
+    print('Downloading Bark model...')
+    try:
+        processor = AutoProcessor.from_pretrained('suno/bark-small')
+        model = BarkModel.from_pretrained('suno/bark-small')
+        print('✅ Bark model downloaded successfully')
+        return True
+    except Exception as e:
+        print(f'❌ Error downloading Bark model: {e}')
+        return False
+
+def download_whisper_model():
+    """Download Whisper ASR model"""
+    print('Downloading Whisper model...')
+    try:
+        whisper_model = whisper.load_model('base')
+        print('✅ Whisper model downloaded successfully')
+        return True
+    except Exception as e:
+        print(f'❌ Error downloading Whisper model: {e}')
+        return False
+
+def main():
+    """Main function to download all models"""
+    print('🚀 Starting model download process...')
+    print('-' * 50)
+    
+    bark_success = download_bark_model()
+    print()
+    
+    whisper_success = download_whisper_model()
+    print()
+    
+    if bark_success and whisper_success:
+        print('🎉 All models downloaded successfully!')
+    else:
+        print('⚠️  Some models failed to download. Check logs above.')
+        if not bark_success:
+            print('   - Bark model download failed')
+        if not whisper_success:
+            print('   - Whisper model download failed')
+    
+    print('-' * 50)
+    print('Model download process complete!')
+
+if __name__ == "__main__":
+    main()