KingNish commited on
Commit
d9f5774
·
verified ·
1 Parent(s): 56ebb10

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -72,7 +72,7 @@ device = "cuda:0"
72
  # Load model and tokenizer outside the generation function (load once)
73
  print("Loading model...")
74
  model = AutoModelForCausalLM.from_pretrained(
75
- "m-a-p/YuE-s1-7B-anneal-en-cot", # "m-a-p/YuE-s1-7B-anneal-en-icl",
76
  torch_dtype=torch.float16,
77
  attn_implementation="flash_attention_2",
78
  ).to(device)
@@ -220,7 +220,7 @@ def generate_music(
220
  logits_processor=LogitsProcessorList([BlockTokenRangeProcessor(0, 32002), BlockTokenRangeProcessor(32016, 32016)]),
221
  guidance_scale=guidance_scale,
222
  use_cache=True,
223
- num_beams=1
224
  )
225
  if output_seq[0][-1].item() != mmtokenizer.eoa:
226
  tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)
 
72
  # Load model and tokenizer outside the generation function (load once)
73
  print("Loading model...")
74
  model = AutoModelForCausalLM.from_pretrained(
75
+ "m-a-p/YuE-s1-7B-anneal-en-icl", # "m-a-p/YuE-s1-7B-anneal-en-cot",
76
  torch_dtype=torch.float16,
77
  attn_implementation="flash_attention_2",
78
  ).to(device)
 
220
  logits_processor=LogitsProcessorList([BlockTokenRangeProcessor(0, 32002), BlockTokenRangeProcessor(32016, 32016)]),
221
  guidance_scale=guidance_scale,
222
  use_cache=True,
223
+ num_beams=3
224
  )
225
  if output_seq[0][-1].item() != mmtokenizer.eoa:
226
  tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)