Update app.py

app.py

@@ -76,8 +76,8 @@ model = AutoModelForCausalLM.from_pretrained(
     "m-a-p/YuE-s1-7B-anneal-en-cot",
     torch_dtype=torch.float16,
     attn_implementation="flash_attention_2",  # To enable flashattn, you have to install flash-attn
-)
-model.to(device)
+).to(device)
+model = torch.compile(model)
 model.eval()
 
 basic_model_config = './xcodec_mini_infer/final_ckpt/config.yaml'
@@ -90,15 +90,19 @@ mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model"
 
 codectool = CodecManipulator("xcodec", 0, 1)
 model_config = OmegaConf.load(basic_model_config)
+# Load codec model
 codec_model = eval(model_config.generator.name)(**model_config.generator.config).to(device)
 parameter_dict = torch.load(resume_path, map_location='cpu')
 codec_model.load_state_dict(parameter_dict['codec_model'])
-codec_model.to(device)
+codec_model = torch.compile(codec_model)
 codec_model.eval()
 
+# Preload and compile vocoders
 vocal_decoder, inst_decoder = build_codec_model(config_path, vocal_decoder_path, inst_decoder_path)
 vocal_decoder.to(device)
 inst_decoder.to(device)
+vocal_decoder = torch.compile(vocal_decoder)
+inst_decoder = torch.compile(inst_decoder)
 vocal_decoder.eval()
 inst_decoder.eval()