Spaces:
Sleeping
Sleeping
| # This script pre-downloads models for the Acoustic Intelligence app | |
| import os | |
| import sys | |
| import logging | |
| import torch | |
| import traceback | |
| # Configure logging | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', | |
| handlers=[logging.StreamHandler(sys.stdout)] | |
| ) | |
| logger = logging.getLogger("model_cache_builder") | |
| # Check if we're in a HF Space | |
| IS_SPACE = os.environ.get("SPACE_ID") is not None | |
| logger.info(f"Running in HF Space: {IS_SPACE}") | |
| # Set device | |
| if torch.cuda.is_available(): | |
| DEVICE = "cuda:0" | |
| logger.info(f"Using CUDA device: {torch.cuda.get_device_name(0)}") | |
| elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): | |
| DEVICE = "mps" | |
| logger.info("Using MPS (Metal Performance Shaders) device") | |
| else: | |
| DEVICE = "cpu" | |
| logger.info("Using CPU device") | |
| def download_nltk_data(): | |
| """Download NLTK data for text processing.""" | |
| logger.info("Downloading NLTK data...") | |
| try: | |
| import nltk | |
| nltk.download('punkt', quiet=True) | |
| nltk.download('stopwords', quiet=True) | |
| logger.info("Successfully downloaded NLTK data") | |
| except Exception as e: | |
| logger.error(f"Error downloading NLTK data: {e}") | |
| # Create fallback directories | |
| os.makedirs(os.path.expanduser("~/nltk_data/tokenizers"), exist_ok=True) | |
| os.makedirs(os.path.expanduser("~/nltk_data/corpora"), exist_ok=True) | |
| def download_tts_models(): | |
| """Download TTS models.""" | |
| logger.info("Downloading TTS models...") | |
| try: | |
| from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan | |
| processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts") | |
| model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts") | |
| vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan") | |
| logger.info("Successfully downloaded TTS models") | |
| except Exception as e: | |
| logger.error(f"Error downloading TTS models: {e}") | |
| def download_asr_model(): | |
| """Download ASR model.""" | |
| logger.info("Downloading FunASR model...") | |
| try: | |
| from funasr import AutoModel | |
| # Just initialize without loading to cache the model files | |
| # The real loading happens at runtime | |
| logger.info("Successfully initialized FunASR") | |
| except Exception as e: | |
| logger.error(f"Error initializing FunASR: {e}") | |
| def download_emotion_model(): | |
| """Download emotion recognition model.""" | |
| logger.info("Downloading emotion recognition model...") | |
| try: | |
| from transformers import Wav2Vec2Processor, Wav2Vec2ForSequenceClassification | |
| model_name = "Dpngtm/wav2vec2-emotion-recognition" | |
| processor = Wav2Vec2Processor.from_pretrained(model_name) | |
| model = Wav2Vec2ForSequenceClassification.from_pretrained(model_name) | |
| logger.info("Successfully downloaded emotion recognition model") | |
| except Exception as e: | |
| logger.error(f"Error downloading emotion recognition model: {e}") | |
| def download_speaker_diarization(): | |
| """Initialize pyannote for speaker diarization.""" | |
| logger.info("Initializing pyannote.audio...") | |
| try: | |
| from pyannote.audio import Pipeline | |
| # Don't authenticate here, just download the model weights | |
| # The actual authentication happens at runtime with user token | |
| logger.info("Successfully initialized pyannote.audio") | |
| except Exception as e: | |
| logger.error(f"Error initializing pyannote.audio: {e}") | |
| def main(): | |
| """Main function to download all necessary models.""" | |
| logger.info("Starting model download...") | |
| # Download all models | |
| download_nltk_data() | |
| download_tts_models() | |
| download_asr_model() | |
| download_emotion_model() | |
| download_speaker_diarization() | |
| logger.info("Model download complete!") | |
| if __name__ == "__main__": | |
| main() | |