modified: app.py

app.py

@@ -4,8 +4,33 @@ import os
 import shutil
 import tempfile
 import spaces
-from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
 import torch
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer'))
+sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer', 'descriptaudiocodec'))
+import argparse
+import numpy as np
+import json
+from omegaconf import OmegaConf
+import torchaudio
+from torchaudio.transforms import Resample
+import soundfile as sf
+
+import uuid
+from tqdm import tqdm
+from einops import rearrange
+from codecmanipulator import CodecManipulator
+from mmtokenizer import _MMSentencePieceTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
+import glob
+import time
+import copy
+from collections import Counter
+from models.soundstream_hubert_new import SoundStream
+from vocoder import build_codec_model, process_audio
+from post_process_audio import replace_low_freq_with_energy_matched
+import re
 
 is_shared_ui = True if "innova-ai/YuE-music-generator-demo" in os.environ['SPACE_ID'] else False
 
@@ -116,33 +141,15 @@ model = AutoModelForCausalLM.from_pretrained(
 model.to(device)
 model.eval()
 
-import os
-import sys
-sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer'))
-sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xcodec_mini_infer', 'descriptaudiocodec'))
-import argparse
-import torch
-import numpy as np
-import json
-from omegaconf import OmegaConf
-import torchaudio
-from torchaudio.transforms import Resample
-import soundfile as sf
+mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
 
-import uuid
-from tqdm import tqdm
-from einops import rearrange
-from codecmanipulator import CodecManipulator
-from mmtokenizer import _MMSentencePieceTokenizer
-from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
-import glob
-import time
-import copy
-from collections import Counter
-from models.soundstream_hubert_new import SoundStream
-from vocoder import build_codec_model, process_audio
-from post_process_audio import replace_low_freq_with_energy_matched
-import re
+codectool = CodecManipulator("xcodec", 0, 1)
+model_config = OmegaConf.load(basic_model_config)
+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.eval()
 
 def generate_music(
     stage1_model="m-a-p/YuE-s1-7B-anneal-en-cot",
@@ -174,22 +181,6 @@ def generate_music(
     stage1_output_dir = os.path.join(output_dir, f"stage1")
     os.makedirs(stage1_output_dir, exist_ok=True)
 
-    # load tokenizer and model
-    device = torch.device(f"cuda:{cuda_idx}" if torch.cuda.is_available() else "cpu")
-
-    # Now you can use `device` to move your tensors or models to the GPU (if available)
-    print(f"Using device: {device}")
-
-    mmtokenizer = _MMSentencePieceTokenizer("./mm_tokenizer_v0.2_hf/tokenizer.model")
-
-    codectool = CodecManipulator("xcodec", 0, 1)
-    model_config = OmegaConf.load(basic_model_config)
-    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.eval()
-
     class BlockTokenRangeProcessor(LogitsProcessor):
         def __init__(self, start_id, end_id):
             self.blocked_token_ids = list(range(start_id, end_id))
@@ -216,13 +207,7 @@ def generate_music(
 
     # Call the function and print the result
     stage1_output_set = []
-    # Tips:
-    # genre tags support instrumental，genre，mood，vocal timbr and vocal gender
-    # # all kinds of tags are needed
-    # with open(genre_txt) as f:
-    #     genres = f.read().strip()
-    # with open(lyrics_txt) as f:
-    #     lyrics = split_lyrics(f.read())
+
     genres = genre_txt.strip()
     lyrics = split_lyrics(lyrics_txt+"\n")
     # intruction