Update app.py

app.py

@@ -77,9 +77,18 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.float16,
     attn_implementation="flash_attention_2",  # To enable flashattn, you have to install flash-attn
 ).to(device)
+assistant_model = AutoModelForCausalLM.from_pretrained(
+    "m-a-p/YuE-s2-1B-general",
+    torch_dtype=torch.float16,
+    attn_implementation="flash_attention_2",  # To enable flashattn, you have to install flash-attn
+).to(device)
+# assistant_model = torch.compile(assistant_model)
 # model = torch.compile(model)
+assistant_model.eval()
 model.eval()
 
+
+
 basic_model_config = './xcodec_mini_infer/final_ckpt/config.yaml'
 resume_path = './xcodec_mini_infer/final_ckpt/ckpt_00360000.pth'
 config_path = './xcodec_mini_infer/decoders/config.yaml'
@@ -231,6 +240,7 @@ def generate_music(
                     use_cache=True,
                     top_k=50,
                     num_beams=1,
+                    assistant_model=assistant_model
                 )
                 if output_seq[0][-1].item() != mmtokenizer.eoa:
                     tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)