Update app.py

app.py

@@ -19,19 +19,19 @@ transcribe_token_id = all_special_ids[-5]
 translate_token_id = all_special_ids[-6]
 
 # --- FastSpeech2 (TTS) Setup ---
-TTS_MODEL_NAME = "ford442/fastspeech2-en-ljspeech"  # OR "facebook/fastspeech2-en-ljspeech" after PR
+TTS_MODEL_NAME = "your_username/fastspeech2-en-ljspeech"  # OR "facebook/fastspeech2-en-ljspeech" after PR
 
-# 1. Load the config (now it should exist!)
-tts_config = AutoConfig.from_pretrained(TTS_MODEL_NAME)
+# 1. Load the config.  We DO need trust_remote_code here, and we explain why below.
+tts_config = AutoConfig.from_pretrained(TTS_MODEL_NAME, trust_remote_code=True)
 
-# 2. Dynamically import the model class. This is the correct way.
-module_name = tts_config.architectures[0]  # Get model class name from config
-module = importlib.import_module(f".{tts_config._name_or_path}", package="transformers.models")
+# 2. Dynamically import the model class.  This is *still* the correct way.
+module_name = tts_config.architectures[0]
+module = importlib.import_module(f"transformers.models.{tts_config.model_type}")  # Corrected module path
 model_class = getattr(module, tts_config.architectures[0])
 
 
 # 3. Load the processor and model.
-tts_processor = AutoProcessor.from_pretrained(TTS_MODEL_NAME)
+tts_processor = AutoProcessor.from_pretrained(TTS_MODEL_NAME, trust_remote_code=True) # Keep this for now
 tts_model = model_class.from_pretrained(TTS_MODEL_NAME, config=tts_config)
 
 tts_device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -44,7 +44,7 @@ vicuna_tokenizer = AutoTokenizer.from_pretrained(VICUNA_MODEL_NAME)
 vicuna_model = AutoModelForCausalLM.from_pretrained(
     VICUNA_MODEL_NAME,
     load_in_8bit=False,
-    torch_dtype=torch.bfloat16,
+    torch_dtype=torch.float32,
     device_map="auto",
 )