1use crate::model_config;
6use super::traits::*;
7use anyhow::Result;
8use serde::{Serialize, Deserialize};
9
10model_config!(MusicGenConfig {
11 vocab_size: usize = 2048,
12 hidden_size: usize = 1024,
13 intermediate_size: usize = 4096,
14 num_hidden_layers: usize = 24,
15 num_attention_heads: usize = 16,
16 num_key_value_heads: usize = 16,
17 max_position_embeddings: usize = 2048,
18 num_codebooks: usize = 4,
19 audio_channels: usize = 1,
20 sampling_rate: usize = 32000,
21 frame_rate: usize = 50,
22 layer_norm_eps: f32 = 1e-5,
23 hidden_dropout: f32 = 0.0,
24 pad_token_id: i64 = 2048,
25 bos_token_id: i64 = 2048,
26 eos_token_id: i64 = 2048,
27});
28
29impl MusicGenConfig {
30 pub fn from_gguf_config(gguf: &crate::weight_loader_core::GGUFModelConfig) -> Self {
31 Self {
32 hidden_size: gguf.hidden_size,
33 num_hidden_layers: gguf.num_hidden_layers,
34 num_attention_heads: gguf.num_attention_heads,
35 num_key_value_heads: gguf.num_key_value_heads,
36 ..Default::default()
37 }
38 }
39}
40
41pub struct MusicGenModelV2 {
42 config: MusicGenConfig,
43 device: Device,
44 embed_tokens: Vec<Tensor>, layers: Vec<MusicGenDecoderLayer>,
46 norm: Tensor,
47 lm_heads: Vec<Tensor>, }
49
50pub struct MusicGenDecoderLayer {
51 self_attn_q: Tensor,
52 self_attn_k: Tensor,
53 self_attn_v: Tensor,
54 self_attn_o: Tensor,
55 gate_proj: Tensor,
56 up_proj: Tensor,
57 down_proj: Tensor,
58 input_layernorm: Tensor,
59 post_attention_layernorm: Tensor,
60 num_heads: usize,
61 num_kv_heads: usize,
62 head_dim: usize,
63}
64
65impl Model for MusicGenModelV2 {
66 type Config = MusicGenConfig;
67
68 fn new(config: MusicGenConfig) -> Result<Self> {
69 let device = Device::CPU;
70
71 let mut embed_tokens = Vec::with_capacity(config.num_codebooks);
72 let mut lm_heads = Vec::with_capacity(config.num_codebooks);
73
74 for _ in 0..config.num_codebooks {
75 embed_tokens.push(ops_fn::zeros(&[config.vocab_size, config.hidden_size], DataType::Float32, &device)?);
76 lm_heads.push(ops_fn::zeros(&[config.hidden_size, config.vocab_size], DataType::Float32, &device)?);
77 }
78
79 let norm = ops_fn::zeros(&[config.hidden_size], DataType::Float32, &device)?;
80
81 let mut layers = Vec::with_capacity(config.num_hidden_layers);
82 for _ in 0..config.num_hidden_layers {
83 layers.push(MusicGenDecoderLayer::new(&config, &device)?);
84 }
85
86 Ok(Self { config, device, embed_tokens, layers, norm, lm_heads })
87 }
88
89 fn from_weights(config: MusicGenConfig, weights: ModelWeights) -> Result<Self> {
90 let mut model = Self::new(config)?;
91
92 for i in 0..model.config.num_codebooks {
93 if let Some(w) = weights.get(&format!("model.decoder.embed_tokens.{}.weight", i)) {
94 model.embed_tokens[i] = w.clone();
95 }
96 if let Some(w) = weights.get(&format!("lm_heads.{}.weight", i)) {
97 model.lm_heads[i] = ops_fn::transpose(w)?;
98 }
99 }
100
101 if let Some(w) = weights.get("model.decoder.final_layer_norm.weight") {
102 model.norm = w.clone();
103 }
104
105 Ok(model)
106 }
107
108 fn forward(&self, inputs: &ModelInputs) -> Result<ModelOutputs> {
109 match inputs {
110 ModelInputs::Text { input_ids, .. } => {
111 let mut hidden = ops_fn::embedding(input_ids, &self.embed_tokens[0])?;
113
114 for layer in &self.layers {
115 hidden = layer.forward(&hidden)?;
116 }
117
118 hidden = ops_fn::layer_norm(&hidden, &self.norm, None, self.config.layer_norm_eps)?;
119
120 let logits = ops_fn::matmul(&hidden, &self.lm_heads[0])?;
122
123 Ok(ModelOutputs::Logits { logits, hidden_states: None })
124 }
125 _ => Err(anyhow::anyhow!("MusicGen requires token input")),
126 }
127 }
128
129 fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String> {
130 use crate::tokenizer::Tokenizer;
131 use rand::Rng;
132
133 let tokenizer = Tokenizer::new();
134 let mut tokens: Vec<u32> = tokenizer.encode(prompt);
135
136 for _ in 0..config.max_new_tokens {
137 let input_ids = Tensor::from_i64_slice(&tokens.iter().map(|&t| t as i64).collect::<Vec<_>>(), &[1, tokens.len()], &self.device)?;
138 let outputs = self.forward(&ModelInputs::text(input_ids))?;
139
140 let logits = match outputs { ModelOutputs::Logits { logits, .. } => logits, _ => return Err(anyhow::anyhow!("Expected logits")) };
141 let logits_vec: Vec<f32> = logits.to_candle()?.flatten_all()?.to_vec1()?;
142 let start = (tokens.len() - 1) * self.config.vocab_size;
143
144 let next_token = if config.do_sample && config.temperature > 0.0 {
145 let scaled: Vec<f32> = logits_vec[start..start + self.config.vocab_size].iter()
146 .map(|&x| x / config.temperature).collect();
147 let max_val = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
148 let exp_sum: f32 = scaled.iter().map(|&x| (x - max_val).exp()).sum();
149 let probs: Vec<f32> = scaled.iter().map(|&x| (x - max_val).exp() / exp_sum).collect();
150
151 let mut rng = rand::thread_rng();
152 let r: f32 = rng.gen();
153 let mut cum = 0.0;
154 let mut s = 0u32;
155 for (i, &p) in probs.iter().enumerate() { cum += p; if r <= cum { s = i as u32; break; } }
156 s
157 } else {
158 logits_vec[start..start + self.config.vocab_size].iter()
159 .enumerate().max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap()).map(|(i, _)| i as u32).unwrap_or(0)
160 };
161
162 if next_token == config.eos_token_id { break; }
163 tokens.push(next_token);
164 }
165
166 Ok(tokenizer.decode(&tokens))
167 }
168
169 fn config(&self) -> &Self::Config { &self.config }
170
171 fn memory_requirements(&self) -> MemoryRequirements {
172 let p = (self.config.vocab_size * self.config.hidden_size * self.config.num_codebooks +
173 self.config.hidden_size * self.config.hidden_size * 4 * self.config.num_hidden_layers) * 4;
174 MemoryRequirements { gpu_memory: p, cpu_memory: p / 4, kv_cache_memory: p / 8, peak_memory: p * 2 }
175 }
176
177 fn to_device(&mut self, device: &Device) -> Result<()> {
178 self.device = device.clone();
179 self.norm = self.norm.to_device(device)?;
180 for e in &mut self.embed_tokens { *e = e.to_device(device)?; }
181 for h in &mut self.lm_heads { *h = h.to_device(device)?; }
182 Ok(())
183 }
184}
185
186impl MusicGenDecoderLayer {
187 fn new(config: &MusicGenConfig, device: &Device) -> Result<Self> {
188 let head_dim = config.hidden_size / config.num_attention_heads;
189
190 Ok(Self {
191 self_attn_q: ops_fn::zeros(&[config.hidden_size, config.num_attention_heads * head_dim], DataType::Float32, device)?,
192 self_attn_k: ops_fn::zeros(&[config.hidden_size, config.num_key_value_heads * head_dim], DataType::Float32, device)?,
193 self_attn_v: ops_fn::zeros(&[config.hidden_size, config.num_key_value_heads * head_dim], DataType::Float32, device)?,
194 self_attn_o: ops_fn::zeros(&[config.num_attention_heads * head_dim, config.hidden_size], DataType::Float32, device)?,
195 gate_proj: ops_fn::zeros(&[config.hidden_size, config.intermediate_size], DataType::Float32, device)?,
196 up_proj: ops_fn::zeros(&[config.hidden_size, config.intermediate_size], DataType::Float32, device)?,
197 down_proj: ops_fn::zeros(&[config.intermediate_size, config.hidden_size], DataType::Float32, device)?,
198 input_layernorm: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
199 post_attention_layernorm: ops_fn::zeros(&[config.hidden_size], DataType::Float32, device)?,
200 num_heads: config.num_attention_heads,
201 num_kv_heads: config.num_key_value_heads,
202 head_dim,
203 })
204 }
205
206 fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
207 let shape = hidden_states.shape();
208 let (batch_size, seq_len, _) = (shape[0], shape[1], shape[2]);
209
210 let residual = hidden_states.clone();
211 let hidden = ops_fn::layer_norm(hidden_states, &self.input_layernorm, None, 1e-5)?;
212
213 let q = ops_fn::matmul(&hidden, &self.self_attn_q)?.to_candle()?
214 .reshape(&[batch_size, seq_len, self.num_heads, self.head_dim])?.transpose(1, 2)?;
215 let k = ops_fn::matmul(&hidden, &self.self_attn_k)?.to_candle()?
216 .reshape(&[batch_size, seq_len, self.num_kv_heads, self.head_dim])?.transpose(1, 2)?;
217 let v = ops_fn::matmul(&hidden, &self.self_attn_v)?.to_candle()?
218 .reshape(&[batch_size, seq_len, self.num_kv_heads, self.head_dim])?.transpose(1, 2)?;
219
220 let num_groups = self.num_heads / self.num_kv_heads;
221 let (k, v) = if num_groups > 1 {
222 (k.unsqueeze(2)?.broadcast_as(&[batch_size, self.num_kv_heads, num_groups, seq_len, self.head_dim])?.reshape(&[batch_size, self.num_heads, seq_len, self.head_dim])?,
223 v.unsqueeze(2)?.broadcast_as(&[batch_size, self.num_kv_heads, num_groups, seq_len, self.head_dim])?.reshape(&[batch_size, self.num_heads, seq_len, self.head_dim])?)
224 } else { (k, v) };
225
226 let scale = (self.head_dim as f32).powf(-0.5);
227 let scores = (q.contiguous()?.matmul(&k.transpose(2, 3)?.contiguous()?)? * (scale as f64))?;
228
229 let device = scores.device();
230 let mask = { let mut m = vec![0.0f32; seq_len * seq_len]; for i in 0..seq_len { for j in (i+1)..seq_len { m[i*seq_len+j] = f32::NEG_INFINITY; } } candle_core::Tensor::from_vec(m, &[1,1,seq_len,seq_len], device)? };
231 let attn = candle_nn::ops::softmax_last_dim(&scores.broadcast_add(&mask)?)?.matmul(&v.contiguous()?)?
232 .transpose(1, 2)?.reshape(&[batch_size, seq_len, self.num_heads * self.head_dim])?;
233
234 let hidden = ops_fn::add(&residual, &ops_fn::matmul(&Tensor::from_candle(attn), &self.self_attn_o)?)?;
235
236 let residual = hidden.clone();
237 let hidden = ops_fn::layer_norm(&hidden, &self.post_attention_layernorm, None, 1e-5)?;
238 let gate = ops_fn::silu(&ops_fn::matmul(&hidden, &self.gate_proj)?)?;
239 let up = ops_fn::matmul(&hidden, &self.up_proj)?;
240 ops_fn::add(&residual, &ops_fn::matmul(&ops_fn::mul(&gate, &up)?, &self.down_proj)?)
241 }
242}
243
244#[cfg(test)]
245mod tests {
246 use super::*;
247
248 #[test]
249 fn test_musicgen_config() {
250 let config = MusicGenConfig::default();
251 assert_eq!(config.hidden_size, 1024);
252 assert_eq!(config.num_codebooks, 4);
253 }
254
255 #[test]
256 fn test_musicgen_forward() {
257 let config = MusicGenConfig { vocab_size: 100, hidden_size: 32, intermediate_size: 128, num_hidden_layers: 1, num_attention_heads: 2, num_key_value_heads: 2, num_codebooks: 2, ..Default::default() };
258 let model = MusicGenModelV2::new(config).unwrap();
259 let outputs = model.forward(&ModelInputs::text(ops_fn::zeros(&[1, 4], DataType::Int64, &Device::CPU).unwrap())).unwrap();
260 match outputs { ModelOutputs::Logits { logits, .. } => assert_eq!(logits.shape(), &[1, 4, 100]), _ => panic!() }
261 }
262}