Update app.py

app.py

@@ -63,7 +63,7 @@ from tqdm import tqdm
 from einops import rearrange
 from codecmanipulator import CodecManipulator
 from mmtokenizer import _MMSentencePieceTokenizer
-from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
 import glob
 import time
 import copy
@@ -74,16 +74,14 @@ from post_process_audio import replace_low_freq_with_energy_matched
 
 device = "cuda:0"
 
-quantization_config = BitsAndBytesConfig(load_in_8bit=True)
-
 model = AutoModelForCausalLM.from_pretrained(
     "m-a-p/YuE-s1-7B-anneal-en-cot",
-    torch_dtype="auto",
-    # attn_implementation="flash_attention_2",
+    torch_dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2",
     quantization_config=quantization_config,
     low_cpu_mem_usage=True,
-    device_map="auto"     
-)
+    # device_map="auto"     
+).to(device)
 model.eval()
 
 basic_model_config = './xcodec_mini_infer/final_ckpt/config.yaml'