Update app.py

app.py

@@ -53,8 +53,18 @@ e2tts = E2TTS(
 )
 
 # Load the pre-trained model
-state_dict = torch.load(model_path, map_location=torch.device('cpu'))
-e2tts.load_state_dict(state_dict)
+checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
+if 'model_state_dict' in checkpoint:
+    state_dict = checkpoint['model_state_dict']
+elif 'ema_model_state_dict' in checkpoint:
+    state_dict = checkpoint['ema_model_state_dict']
+else:
+    state_dict = checkpoint  # Assume the checkpoint is the state dict itself
+
+# Filter out unexpected keys
+model_dict = e2tts.state_dict()
+filtered_state_dict = {k: v for k, v in state_dict.items() if k in model_dict}
+e2tts.load_state_dict(filtered_state_dict, strict=False)
 e2tts.eval()
 
 def generate_podcast_script(content, duration):