1import random
2import re
3import sys
4from argparse import ArgumentParser
5from pathlib import Path
6from warnings import simplefilter
7
8sys.path.append("xcodec_mini_infer")
9simplefilter("ignore")
10
11import torch
12import torchaudio
13import yaml
14from exllamav2 import (
15 ExLlamaV2,
16 ExLlamaV2Cache,
17 ExLlamaV2Config,
18 ExLlamaV2Tokenizer,
19 Timer,
20)
21from exllamav2.generator import (
22 ExLlamaV2DynamicGenerator,
23 ExLlamaV2DynamicJob,
24 ExLlamaV2Sampler,
25)
26from rich import print
27
28from xcodec_mini_infer.models.soundstream_hubert_new import SoundStream
29
30parser = ArgumentParser()
31parser.add_argument("-m", "--model", required=True)
32parser.add_argument("-g", "--genre", required=True)
33parser.add_argument("-l", "--lyrics", required=True)
34parser.add_argument("-s", "--seed", type=int, default=None)
35parser.add_argument("-d", "--debug", action="store_true")
36parser.add_argument("--repetition_penalty", type=float, default=1.2)
37parser.add_argument("--temperature", type=float, default=1.0)
38parser.add_argument("--top_p", type=float, default=0.93)
39args = parser.parse_args()
40
41with Timer() as timer:
42 config = ExLlamaV2Config(args.model)
43 model = ExLlamaV2(config, lazy_load=True)
44 cache = ExLlamaV2Cache(model, lazy=True)
45 model.load_autosplit(cache)
46
47 tokenizer = ExLlamaV2Tokenizer(config, lazy_init=True)
48 generator = ExLlamaV2DynamicGenerator(model, cache, tokenizer)
49 generator.warmup()
50
51print(f"Loaded model in {timer.interval:.2f} seconds.")
52
53genre = Path(args.genre)
54genre = genre.read_text(encoding="utf-8") if genre.is_file() else args.genre
55genre = genre.strip()
56
57lyrics = Path(args.lyrics)
58lyrics = lyrics.read_text(encoding="utf-8") if lyrics.is_file() else args.lyrics
59lyrics = lyrics.strip()
60
61lyrics = re.findall(r"\[(\w+)\](.*?)\n(?=\[|\Z)", lyrics, re.DOTALL)
62lyrics = [f"[{l[0]}]\n{l[1].strip()}\n\n" for l in lyrics]
63lyrics_joined = "\n".join(lyrics)
64
65gen_settings = ExLlamaV2Sampler.Settings()
66gen_settings.allow_tokens(tokenizer, [32002] + list(range(45334, 46358)))
67gen_settings.temperature = args.temperature
68gen_settings.token_repetition_penalty = args.repetition_penalty
69gen_settings.top_p = args.top_p
70
71seed = args.seed if args.seed else random.randint(0, 2**64 - 1)
72stop_conditions = ["<EOA>"]
73
74output_joined = ""
75output = []
76
77with Timer() as timer:
78 for segment in lyrics:
79 current = []
80
81 input = (
82 "Generate music from the given lyrics segment by segment.\n"
83 f"[Genre] {genre}\n"
84 f"{lyrics_joined}{output_joined}[start_of_segment]{segment}<SOA><xcodec>"
85 )
86
87 input_ids = tokenizer.encode(input, encode_special_tokens=True)
88 input_len = input_ids.shape[-1]
89 max_new_tokens = config.max_seq_len - input_len
90
91 print(
92 f"Using {input_len} tokens of {config.max_seq_len} tokens "
93 f"with {max_new_tokens} tokens left."
94 )
95
96 job = ExLlamaV2DynamicJob(
97 input_ids=input_ids,
98 max_new_tokens=max_new_tokens,
99 gen_settings=gen_settings,
100 seed=seed,
101 stop_conditions=stop_conditions,
102 decode_special_tokens=True,
103 )
104
105 generator.enqueue(job)
106
107 with Timer() as inner:
108 while generator.num_remaining_jobs():
109 for result in generator.iterate():
110 if result.get("stage") == "streaming":
111 text = result.get("text")
112
113 if text:
114 current.append(text)
115 output.append(text)
116
117 if args.debug:
118 print(text, end="", flush=True)
119
120 if result.get("eos") and current:
121 current_joined = "".join(current)
122 output_joined += (
123 f"[start_of_segment]{segment}<SOA><xcodec>"
124 f"{current_joined}<EOA>[end_of_segment]"
125 )
126
127 if args.debug:
128 print()
129
130 print(f"Generated {len(current)} tokens in {inner.interval:.2f} seconds.")
131
132print(f"Finished in {timer.interval:.2f} seconds with seed {seed}.")
133
134with Timer() as timer:
135 codec_config = Path("xcodec_mini_infer/final_ckpt/config.yaml")
136 codec_config = yaml.safe_load(codec_config.read_bytes())
137 codec = SoundStream(**codec_config["generator"]["config"])
138 state_dict = torch.load("xcodec_mini_infer/final_ckpt/ckpt_00360000.pth")
139 codec.load_state_dict(state_dict["codec_model"])
140 codec = codec.eval().cuda()
141
142print(f"Loaded codec in {timer.interval:.2f} seconds.")
143
144with Timer() as timer, torch.inference_mode():
145 pattern = re.compile(r"<xcodec/0/(\d+)>")
146 output_ids = [int(o[10:-1]) for o in output if re.match(pattern, o)]
147
148 vocal = output_ids[::2]
149 vocal = torch.tensor([[vocal]]).cuda()
150 vocal = vocal.permute(1, 0, 2)
151 vocal = codec.decode(vocal)
152 vocal = vocal.squeeze(0).cpu()
153 torchaudio.save("vocal.wav", vocal, 16000)
154
155 inst = output_ids[1::2]
156 inst = torch.tensor([[inst]]).cuda()
157 inst = inst.permute(1, 0, 2)
158 inst = codec.decode(inst)
159 inst = inst.squeeze(0).cpu()
160 torchaudio.save("inst.wav", inst, 16000)
161
162print(f"Decoded audio in {timer.interval:.2f} seconds.")
