deepanwayx commited on
Commit
e9dbc42
1 Parent(s): 0174916

update app

Browse files
.ipynb_checkpoints/app-checkpoint.py ADDED
@@ -0,0 +1,292 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import json
3
+ import torch
4
+ import wavio
5
+ import numpy as np
6
+ from tqdm import tqdm
7
+ from huggingface_hub import snapshot_download
8
+
9
+ from audioldm.audio.stft import TacotronSTFT
10
+ from audioldm.variational_autoencoder import AutoencoderKL
11
+
12
+ from transformers import AutoTokenizer, T5ForConditionalGeneration
13
+ from modelling_deberta_v2 import DebertaV2ForTokenClassificationRegression
14
+
15
+ import sys
16
+ sys.path.insert(0, "diffusers/src")
17
+
18
+ from diffusers import DDPMScheduler
19
+ from models import MusicAudioDiffusion
20
+
21
+ from gradio import Markdown
22
+
23
+ class MusicFeaturePredictor:
24
+ def __init__(self, path, device="cuda:0", cache_dir=None, local_files_only=False):
25
+ self.beats_tokenizer = AutoTokenizer.from_pretrained(
26
+ "microsoft/deberta-v3-large",
27
+ use_fast=False,
28
+ cache_dir=cache_dir,
29
+ local_files_only=local_files_only,
30
+ )
31
+ self.beats_model = DebertaV2ForTokenClassificationRegression.from_pretrained(
32
+ "microsoft/deberta-v3-large",
33
+ cache_dir=cache_dir,
34
+ local_files_only=local_files_only,
35
+ )
36
+ self.beats_model.eval()
37
+ self.beats_model.to(device)
38
+
39
+ beats_ckpt = f"{path}/beats/microsoft-deberta-v3-large.pt"
40
+ beats_weight = torch.load(beats_ckpt, map_location="cpu")
41
+ self.beats_model.load_state_dict(beats_weight)
42
+
43
+ self.chords_tokenizer = AutoTokenizer.from_pretrained(
44
+ "google/flan-t5-large",
45
+ cache_dir=cache_dir,
46
+ local_files_only=local_files_only,
47
+ )
48
+ self.chords_model = T5ForConditionalGeneration.from_pretrained(
49
+ "google/flan-t5-large",
50
+ cache_dir=cache_dir,
51
+ local_files_only=local_files_only,
52
+ )
53
+ self.chords_model.eval()
54
+ self.chords_model.to(device)
55
+
56
+ chords_ckpt = f"{path}/chords/flan-t5-large.bin"
57
+ chords_weight = torch.load(chords_ckpt, map_location="cpu")
58
+ self.chords_model.load_state_dict(chords_weight)
59
+
60
+ def generate_beats(self, prompt):
61
+ tokenized = self.beats_tokenizer(
62
+ prompt, max_length=512, padding=True, truncation=True, return_tensors="pt"
63
+ )
64
+ tokenized = {k: v.to(self.beats_model.device) for k, v in tokenized.items()}
65
+
66
+ with torch.no_grad():
67
+ out = self.beats_model(**tokenized)
68
+
69
+ max_beat = (
70
+ 1 + torch.argmax(out["logits"][:, 0, :], -1).detach().cpu().numpy()
71
+ ).tolist()[0]
72
+ intervals = (
73
+ out["values"][:, :, 0]
74
+ .detach()
75
+ .cpu()
76
+ .numpy()
77
+ .astype("float32")
78
+ .round(4)
79
+ .tolist()
80
+ )
81
+
82
+ intervals = np.cumsum(intervals)
83
+ predicted_beats_times = []
84
+ for t in intervals:
85
+ if t < 10:
86
+ predicted_beats_times.append(round(t, 2))
87
+ else:
88
+ break
89
+ predicted_beats_times = list(np.array(predicted_beats_times)[:50])
90
+
91
+ if len(predicted_beats_times) == 0:
92
+ predicted_beats = [[], []]
93
+ else:
94
+ beat_counts = []
95
+ for i in range(len(predicted_beats_times)):
96
+ beat_counts.append(float(1.0 + np.mod(i, max_beat)))
97
+ predicted_beats = [[predicted_beats_times, beat_counts]]
98
+
99
+ return max_beat, predicted_beats_times, predicted_beats
100
+
101
+ def generate(self, prompt):
102
+ max_beat, predicted_beats_times, predicted_beats = self.generate_beats(prompt)
103
+
104
+ chords_prompt = "Caption: {} \\n Timestamps: {} \\n Max Beat: {}".format(
105
+ prompt,
106
+ " , ".join([str(round(t, 2)) for t in predicted_beats_times]),
107
+ max_beat,
108
+ )
109
+
110
+ tokenized = self.chords_tokenizer(
111
+ chords_prompt,
112
+ max_length=512,
113
+ padding=True,
114
+ truncation=True,
115
+ return_tensors="pt",
116
+ )
117
+ tokenized = {k: v.to(self.chords_model.device) for k, v in tokenized.items()}
118
+
119
+ generated_chords = self.chords_model.generate(
120
+ input_ids=tokenized["input_ids"],
121
+ attention_mask=tokenized["attention_mask"],
122
+ min_length=8,
123
+ max_length=128,
124
+ num_beams=5,
125
+ early_stopping=True,
126
+ num_return_sequences=1,
127
+ )
128
+
129
+ generated_chords = self.chords_tokenizer.decode(
130
+ generated_chords[0],
131
+ skip_special_tokens=True,
132
+ clean_up_tokenization_spaces=True,
133
+ ).split(" n ")
134
+
135
+ predicted_chords, predicted_chords_times = [], []
136
+ for item in generated_chords:
137
+ c, ct = item.split(" at ")
138
+ predicted_chords.append(c)
139
+ predicted_chords_times.append(float(ct))
140
+
141
+ return predicted_beats, predicted_chords, predicted_chords_times
142
+
143
+
144
+ class Mustango:
145
+ def __init__(
146
+ self,
147
+ device="cuda:0",
148
+ cache_dir=None,
149
+ local_files_only=False,
150
+ ):
151
+
152
+ path = snapshot_download(repo_id="declare-lab/mustango", cache_dir=cache_dir)
153
+ pretrained_path = snapshot_download(repo_id="declare-lab/mustango-pretrained", cache_dir=cache_dir)
154
+
155
+ self.music_model = MusicFeaturePredictor(
156
+ path, device, cache_dir=cache_dir, local_files_only=local_files_only
157
+ )
158
+
159
+ vae_config = json.load(open(f"{path}/configs/vae_config.json"))
160
+ stft_config = json.load(open(f"{path}/configs/stft_config.json"))
161
+ main_config = json.load(open(f"{path}/configs/main_config.json"))
162
+
163
+ self.vae = AutoencoderKL(**vae_config).to(device)
164
+ self.stft = TacotronSTFT(**stft_config).to(device)
165
+ self.model = MusicAudioDiffusion(
166
+ main_config["text_encoder_name"],
167
+ main_config["scheduler_name"],
168
+ unet_model_config_path=f"{path}/configs/music_diffusion_model_config.json",
169
+ ).to(device)
170
+ self.model.device = device
171
+
172
+ vae_weights = torch.load(
173
+ f"{path}/vae/pytorch_model_vae.bin", map_location=device
174
+ )
175
+ stft_weights = torch.load(
176
+ f"{path}/stft/pytorch_model_stft.bin", map_location=device
177
+ )
178
+ self.main_weights = torch.load(
179
+ f"{path}/ldm/pytorch_model_ldm.bin", map_location=device
180
+ )
181
+
182
+ self.main_weights_pretrained = torch.load(
183
+ f"{pretrained_path}/ldm/pytorch_model_ldm.bin", map_location=device
184
+ )
185
+
186
+ self.vae.load_state_dict(vae_weights)
187
+ self.stft.load_state_dict(stft_weights)
188
+
189
+ self.vae.eval()
190
+ self.stft.eval()
191
+ self.model.eval()
192
+
193
+ self.scheduler = DDPMScheduler.from_pretrained(
194
+ main_config["scheduler_name"], subfolder="scheduler"
195
+ )
196
+
197
+ def generate(self, model_name, prompt, steps=100, guidance=3, samples=1, disable_progress=True):
198
+ """Genrate music for a single prompt string."""
199
+
200
+ if model_name == "declare-lab/mustango":
201
+ self.model.load_state_dict(self.main_weights)
202
+ else:
203
+ self.model.load_state_dict(self.main_weights_pretrained)
204
+
205
+ with torch.no_grad():
206
+ beats, chords, chords_times = self.music_model.generate(prompt)
207
+ latents = self.model.inference(
208
+ [prompt],
209
+ beats,
210
+ [chords],
211
+ [chords_times],
212
+ self.scheduler,
213
+ steps,
214
+ guidance,
215
+ samples,
216
+ disable_progress,
217
+ )
218
+ mel = self.vae.decode_first_stage(latents)
219
+ wave = self.vae.decode_to_waveform(mel)
220
+
221
+ return wave[0]
222
+
223
+
224
+ # Initialize Mustango
225
+ if torch.cuda.is_available():
226
+ mustango = Mustango()
227
+ else:
228
+ mustango = Mustango(device="cpu")
229
+
230
+ # output_wave = mustango.generate("This techno song features a synth lead playing the main melody.", 5, 3, disable_progress=False)
231
+
232
+ def gradio_generate(model_name, prompt, steps, guidance):
233
+ output_wave = mustango.generate(model_name, prompt, steps, guidance)
234
+ # output_filename = f"{prompt.replace(' ', '_')}_{steps}_{guidance}"[:250] + ".wav"
235
+ output_filename = "temp.wav"
236
+ wavio.write(output_filename, output_wave, rate=16000, sampwidth=2)
237
+
238
+ return output_filename
239
+
240
+
241
+ title="Mustango: Toward Controllable Text-to-Music Generation"
242
+ description_text = """
243
+ <p><a href="https://huggingface.co/spaces/declare-lab/mustango/blob/main/app.py?duplicate=true"> <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a> For faster inference without waiting in queue, you may duplicate the space and upgrade to a GPU in the settings. <br/><br/>
244
+ Generate music using Mustango by providing a text prompt.
245
+ <br/><br/> This is the demo for Mustango for controllable text to music generation: <a href="https://arxiv.org/abs/2311.08355">Read our paper.</a>
246
+ <p/>
247
+ """
248
+ #description_text = ""
249
+ # Gradio input and output components
250
+ input_text = gr.Textbox(lines=2, label="Prompt")
251
+ output_audio = gr.Audio(label="Generated Music", type="filepath")
252
+ denoising_steps = gr.Slider(minimum=100, maximum=200, value=100, step=1, label="Steps", interactive=True)
253
+ guidance_scale = gr.Slider(minimum=1, maximum=10, value=3, step=0.1, label="Guidance Scale", interactive=True)
254
+ model_name = gr.Radio(["declare-lab/mustango", "declare-lab/mustango-pretrained",], label="Choose a model type", value="declare-lab/mustango", type="value")
255
+
256
+
257
+ # CSS styling for the Duplicate button
258
+ css = '''
259
+ #duplicate-button {
260
+ margin: auto;
261
+ color: white;
262
+ background: #1565c0;
263
+ border-radius: 100vh;
264
+ }
265
+ '''
266
+
267
+ # Gradio interface
268
+ gr_interface = gr.Interface(
269
+ fn=gradio_generate,
270
+ inputs=[model_name, input_text, denoising_steps, guidance_scale],
271
+ outputs=[output_audio],
272
+ description=description_text,
273
+ allow_flagging=False,
274
+ examples=[
275
+ ["declare-lab/mustango", "This techno song features a synth lead playing the main melody. This is accompanied by programmed percussion playing a simple kick focused beat. The hi-hat is accented in an open position on the 3-and count of every bar. The synth plays the bass part with a voicing that sounds like a cello. This techno song can be played in a club. The chord sequence is Gm, A7, Eb, Bb, C, F, Gm. The beat counts to 2. The tempo of this song is 128.0 beats per minute. The key of this song is G minor.", 200, 3],
276
+ ["declare-lab/mustango", "This is a new age piece. There is a flute playing the main melody with a lot of staccato notes. The rhythmic background consists of a medium tempo electronic drum beat with percussive elements all over the spectrum. There is a playful atmosphere to the piece. This piece can be used in the soundtrack of a children's TV show or an advertisement jingle.", 200, 3],
277
+ ["declare-lab/mustango", "The song is an instrumental. The song is in medium tempo with a classical guitar playing a lilting melody in accompaniment style. The song is emotional and romantic. The song is a romantic instrumental song. The chord sequence is Gm, F6, Ebm. The time signature is 4/4. This song is in Adagio. The key of this song is G minor.", 200, 3],
278
+ ["declare-lab/mustango", "This folk song features a female voice singing the main melody. This is accompanied by a tabla playing the percussion. A guitar strums chords. For most parts of the song, only one chord is played. At the last bar, a different chord is played. This song has minimal instruments. This song has a story-telling mood. This song can be played in a village scene in an Indian movie. The chord sequence is Bbm, Ab. The beat is 3. The tempo of this song is Allegro. The key of this song is Bb minor.", 200, 3],
279
+ ["declare-lab/mustango", "This is a live performance of a classical music piece. There is an orchestra performing the piece with a violin lead playing the main melody. The atmosphere is sentimental and heart-touching. This piece could be playing in the background at a classy restaurant. The chord progression in this song is Am7, Gm, Dm, A7, Dm. The beat is 3. This song is in Largo. The key of this song is D minor.", 200, 3],
280
+ ["declare-lab/mustango", "This is a techno piece with drums and beats and a leading melody. A synth plays chords. The music kicks off with a powerful and relentless drumbeat. Over the pounding beats, a leading melody emerges. In the middle of the song, a flock of seagulls flies over the venue and make loud bird sounds. It has strong danceability and can be played in a club. The tempo is 120 bpm. The chords played by the synth are Am, Cm, Dm, Gm.", 200, 3],
281
+ ],
282
+ cache_examples=True,
283
+ )
284
+
285
+ with gr.Blocks(css=css) as demo:
286
+ title=gr.HTML(f"<h1><center>{title}</center></h1>")
287
+ dupe = gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
288
+ gr_interface.render()
289
+
290
+
291
+ # Launch Gradio app
292
+ demo.queue().launch()
.ipynb_checkpoints/models-checkpoint.py ADDED
@@ -0,0 +1,740 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import yaml
2
+ import random
3
+ import inspect
4
+ import numpy as np
5
+ from tqdm import tqdm
6
+
7
+ import torch
8
+ import torch.nn as nn
9
+ import torch.nn.functional as F
10
+
11
+ from einops import repeat
12
+ from tools.torch_tools import wav_to_fbank
13
+
14
+ from audioldm.audio.stft import TacotronSTFT
15
+ from audioldm.variational_autoencoder import AutoencoderKL
16
+ from audioldm.utils import default_audioldm_config, get_metadata
17
+
18
+ from transformers import CLIPTokenizer, AutoTokenizer, AutoProcessor
19
+ from transformers import CLIPTextModel, T5EncoderModel, AutoModel, ClapAudioModel, ClapTextModel
20
+
21
+ import sys
22
+ sys.path.insert(0, "diffusers/src")
23
+
24
+ import diffusers
25
+ from diffusers.utils import randn_tensor
26
+ from diffusers import DDPMScheduler, UNet2DConditionModel, UNet2DConditionModelMusic
27
+ from diffusers import AutoencoderKL as DiffuserAutoencoderKL
28
+ from layers.layers import chord_tokenizer, beat_tokenizer, Chord_Embedding, Beat_Embedding, Music_PositionalEncoding, Fundamental_Music_Embedding
29
+
30
+ def build_pretrained_models(name):
31
+ checkpoint = torch.load(get_metadata()[name]["path"], map_location="cpu")
32
+ scale_factor = checkpoint["state_dict"]["scale_factor"].item()
33
+
34
+ vae_state_dict = {k[18:]: v for k, v in checkpoint["state_dict"].items() if "first_stage_model." in k}
35
+
36
+ config = default_audioldm_config(name)
37
+ vae_config = config["model"]["params"]["first_stage_config"]["params"]
38
+ vae_config["scale_factor"] = scale_factor
39
+
40
+ vae = AutoencoderKL(**vae_config)
41
+ vae.load_state_dict(vae_state_dict)
42
+
43
+ fn_STFT = TacotronSTFT(
44
+ config["preprocessing"]["stft"]["filter_length"],
45
+ config["preprocessing"]["stft"]["hop_length"],
46
+ config["preprocessing"]["stft"]["win_length"],
47
+ config["preprocessing"]["mel"]["n_mel_channels"],
48
+ config["preprocessing"]["audio"]["sampling_rate"],
49
+ config["preprocessing"]["mel"]["mel_fmin"],
50
+ config["preprocessing"]["mel"]["mel_fmax"],
51
+ )
52
+
53
+ vae.eval()
54
+ fn_STFT.eval()
55
+ return vae, fn_STFT
56
+
57
+
58
+ class AudioDiffusion(nn.Module):
59
+ def __init__(
60
+ self,
61
+ text_encoder_name,
62
+ scheduler_name,
63
+ unet_model_name=None,
64
+ unet_model_config_path=None,
65
+ snr_gamma=None,
66
+ freeze_text_encoder=True,
67
+ uncondition=False,
68
+
69
+ ):
70
+ super().__init__()
71
+
72
+ assert unet_model_name is not None or unet_model_config_path is not None, "Either UNet pretrain model name or a config file path is required"
73
+
74
+ self.text_encoder_name = text_encoder_name
75
+ self.scheduler_name = scheduler_name
76
+ self.unet_model_name = unet_model_name
77
+ self.unet_model_config_path = unet_model_config_path
78
+ self.snr_gamma = snr_gamma
79
+ self.freeze_text_encoder = freeze_text_encoder
80
+ self.uncondition = uncondition
81
+
82
+ # https://huggingface.co/docs/diffusers/v0.14.0/en/api/schedulers/overview
83
+ self.noise_scheduler = DDPMScheduler.from_pretrained(self.scheduler_name, subfolder="scheduler")
84
+ self.inference_scheduler = DDPMScheduler.from_pretrained(self.scheduler_name, subfolder="scheduler")
85
+
86
+ if unet_model_config_path:
87
+ unet_config = UNet2DConditionModel.load_config(unet_model_config_path)
88
+ self.unet = UNet2DConditionModel.from_config(unet_config, subfolder="unet")
89
+ self.set_from = "random"
90
+ print("UNet initialized randomly.")
91
+ else:
92
+ self.unet = UNet2DConditionModel.from_pretrained(unet_model_name, subfolder="unet")
93
+ self.set_from = "pre-trained"
94
+ self.group_in = nn.Sequential(nn.Linear(8, 512), nn.Linear(512, 4))
95
+ self.group_out = nn.Sequential(nn.Linear(4, 512), nn.Linear(512, 8))
96
+ print("UNet initialized from stable diffusion checkpoint.")
97
+
98
+ if "stable-diffusion" in self.text_encoder_name:
99
+ self.tokenizer = CLIPTokenizer.from_pretrained(self.text_encoder_name, subfolder="tokenizer")
100
+ self.text_encoder = CLIPTextModel.from_pretrained(self.text_encoder_name, subfolder="text_encoder")
101
+ elif "t5" in self.text_encoder_name:
102
+ self.tokenizer = AutoTokenizer.from_pretrained(self.text_encoder_name)
103
+ self.text_encoder = T5EncoderModel.from_pretrained(self.text_encoder_name)
104
+ else:
105
+ self.tokenizer = AutoTokenizer.from_pretrained(self.text_encoder_name)
106
+ self.text_encoder = AutoModel.from_pretrained(self.text_encoder_name)
107
+
108
+ def compute_snr(self, timesteps):
109
+ """
110
+ Computes SNR as per https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849
111
+ """
112
+ alphas_cumprod = self.noise_scheduler.alphas_cumprod
113
+ sqrt_alphas_cumprod = alphas_cumprod**0.5
114
+ sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5
115
+
116
+ # Expand the tensors.
117
+ # Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026
118
+ sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[timesteps].float()
119
+ while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):
120
+ sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]
121
+ alpha = sqrt_alphas_cumprod.expand(timesteps.shape)
122
+
123
+ sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(device=timesteps.device)[timesteps].float()
124
+ while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape):
125
+ sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None]
126
+ sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape)
127
+
128
+ # Compute SNR.
129
+ snr = (alpha / sigma) ** 2
130
+ return snr
131
+
132
+ def encode_text(self, prompt):
133
+ device = self.text_encoder.device
134
+ batch = self.tokenizer(
135
+ prompt, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"
136
+ )
137
+ input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)
138
+
139
+ if self.freeze_text_encoder:
140
+ with torch.no_grad():
141
+ encoder_hidden_states = self.text_encoder(
142
+ input_ids=input_ids, attention_mask=attention_mask
143
+ )[0]
144
+ else:
145
+ encoder_hidden_states = self.text_encoder(
146
+ input_ids=input_ids, attention_mask=attention_mask
147
+ )[0]
148
+
149
+ boolean_encoder_mask = (attention_mask == 1).to(device)
150
+ return encoder_hidden_states, boolean_encoder_mask
151
+
152
+ def forward(self, latents, prompt, validation_mode=False):
153
+ device = self.text_encoder.device
154
+ num_train_timesteps = self.noise_scheduler.num_train_timesteps
155
+ self.noise_scheduler.set_timesteps(num_train_timesteps, device=device)
156
+
157
+ encoder_hidden_states, boolean_encoder_mask = self.encode_text(prompt)
158
+
159
+ if self.uncondition:
160
+ mask_indices = [k for k in range(len(prompt)) if random.random() < 0.1]
161
+ if len(mask_indices) > 0:
162
+ encoder_hidden_states[mask_indices] = 0
163
+
164
+ bsz = latents.shape[0]
165
+
166
+ if validation_mode:
167
+ timesteps = (self.noise_scheduler.num_train_timesteps//2) * torch.ones((bsz,), dtype=torch.int64, device=device)
168
+ else:
169
+ # Sample a random timestep for each instance
170
+ timesteps = torch.randint(0, self.noise_scheduler.num_train_timesteps, (bsz,), device=device)
171
+ # print('in if ', timesteps)
172
+ timesteps = timesteps.long()
173
+ # print('outside if ' , timesteps)
174
+ noise = torch.randn_like(latents)
175
+ noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)
176
+
177
+ # Get the target for loss depending on the prediction type
178
+ if self.noise_scheduler.config.prediction_type == "epsilon":
179
+ target = noise
180
+ elif self.noise_scheduler.config.prediction_type == "v_prediction":
181
+ target = self.noise_scheduler.get_velocity(latents, noise, timesteps)
182
+ else:
183
+ raise ValueError(f"Unknown prediction type {self.noise_scheduler.config.prediction_type}")
184
+
185
+ if self.set_from == "random":
186
+ model_pred = self.unet(
187
+ noisy_latents, timesteps, encoder_hidden_states,
188
+ encoder_attention_mask=boolean_encoder_mask
189
+ ).sample
190
+
191
+ elif self.set_from == "pre-trained":
192
+ compressed_latents = self.group_in(noisy_latents.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()
193
+ model_pred = self.unet(
194
+ compressed_latents, timesteps, encoder_hidden_states,
195
+ encoder_attention_mask=boolean_encoder_mask
196
+ ).sample
197
+ model_pred = self.group_out(model_pred.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()
198
+
199
+ if self.snr_gamma is None:
200
+ loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
201
+ else:
202
+ # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
203
+ # Adaptef from huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py
204
+ snr = self.compute_snr(timesteps)
205
+ mse_loss_weights = (
206
+ torch.stack([snr, self.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr
207
+ )
208
+ loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
209
+ loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights
210
+ loss = loss.mean()
211
+
212
+ return loss
213
+
214
+ @torch.no_grad()
215
+ def inference(self, prompt, inference_scheduler, num_steps=20, guidance_scale=3, num_samples_per_prompt=1,
216
+ disable_progress=True):
217
+ device = self.text_encoder.device
218
+ classifier_free_guidance = guidance_scale > 1.0
219
+ batch_size = len(prompt) * num_samples_per_prompt
220
+
221
+ if classifier_free_guidance:
222
+ prompt_embeds, boolean_prompt_mask = self.encode_text_classifier_free(prompt, num_samples_per_prompt)
223
+ else:
224
+ prompt_embeds, boolean_prompt_mask = self.encode_text(prompt)
225
+ prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
226
+ boolean_prompt_mask = boolean_prompt_mask.repeat_interleave(num_samples_per_prompt, 0)
227
+
228
+ inference_scheduler.set_timesteps(num_steps, device=device)
229
+ timesteps = inference_scheduler.timesteps
230
+
231
+ num_channels_latents = self.unet.in_channels
232
+ latents = self.prepare_latents(batch_size, inference_scheduler, num_channels_latents, prompt_embeds.dtype, device)
233
+
234
+ num_warmup_steps = len(timesteps) - num_steps * inference_scheduler.order
235
+ progress_bar = tqdm(range(num_steps), disable=disable_progress)
236
+
237
+ for i, t in enumerate(timesteps):
238
+ # expand the latents if we are doing classifier free guidance
239
+ latent_model_input = torch.cat([latents] * 2) if classifier_free_guidance else latents
240
+ latent_model_input = inference_scheduler.scale_model_input(latent_model_input, t)
241
+
242
+ noise_pred = self.unet(
243
+ latent_model_input, t, encoder_hidden_states=prompt_embeds,
244
+ encoder_attention_mask=boolean_prompt_mask
245
+ ).sample
246
+
247
+ # perform guidance
248
+ if classifier_free_guidance:
249
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
250
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
251
+
252
+ # compute the previous noisy sample x_t -> x_t-1
253
+ latents = inference_scheduler.step(noise_pred, t, latents).prev_sample
254
+
255
+ # call the callback, if provided
256
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % inference_scheduler.order == 0):
257
+ progress_bar.update(1)
258
+
259
+ if self.set_from == "pre-trained":
260
+ latents = self.group_out(latents.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()
261
+ return latents
262
+
263
+ def prepare_latents(self, batch_size, inference_scheduler, num_channels_latents, dtype, device):
264
+ shape = (batch_size, num_channels_latents, 256, 16)
265
+ latents = randn_tensor(shape, generator=None, device=device, dtype=dtype)
266
+ # scale the initial noise by the standard deviation required by the scheduler
267
+ latents = latents * inference_scheduler.init_noise_sigma
268
+ return latents
269
+
270
+ def encode_text_classifier_free(self, prompt, num_samples_per_prompt):
271
+ device = self.text_encoder.device
272
+ batch = self.tokenizer(
273
+ prompt, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"
274
+ )
275
+ input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)
276
+
277
+ with torch.no_grad():
278
+ prompt_embeds = self.text_encoder(
279
+ input_ids=input_ids, attention_mask=attention_mask
280
+ )[0]
281
+
282
+ prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
283
+ attention_mask = attention_mask.repeat_interleave(num_samples_per_prompt, 0)
284
+
285
+ # get unconditional embeddings for classifier free guidance
286
+ uncond_tokens = [""] * len(prompt)
287
+
288
+ max_length = prompt_embeds.shape[1]
289
+ uncond_batch = self.tokenizer(
290
+ uncond_tokens, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt",
291
+ )
292
+ uncond_input_ids = uncond_batch.input_ids.to(device)
293
+ uncond_attention_mask = uncond_batch.attention_mask.to(device)
294
+
295
+ with torch.no_grad():
296
+ negative_prompt_embeds = self.text_encoder(
297
+ input_ids=uncond_input_ids, attention_mask=uncond_attention_mask
298
+ )[0]
299
+
300
+ negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
301
+ uncond_attention_mask = uncond_attention_mask.repeat_interleave(num_samples_per_prompt, 0)
302
+
303
+ # For classifier free guidance, we need to do two forward passes.
304
+ # We concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes
305
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
306
+ prompt_mask = torch.cat([uncond_attention_mask, attention_mask])
307
+ boolean_prompt_mask = (prompt_mask == 1).to(device)
308
+
309
+ return prompt_embeds, boolean_prompt_mask
310
+
311
+
312
+ class MusicAudioDiffusion(nn.Module):
313
+ def __init__(
314
+ self,
315
+ text_encoder_name,
316
+ scheduler_name,
317
+ unet_model_name=None,
318
+ unet_model_config_path=None,
319
+ snr_gamma=None,
320
+ freeze_text_encoder=True,
321
+ uncondition=False,
322
+
323
+ d_fme = 1024, #FME
324
+ fme_type = "se",
325
+ base = 1,
326
+ if_trainable = True,
327
+ translation_bias_type = "nd",
328
+ emb_nn = True,
329
+ d_pe = 1024, #PE
330
+ if_index = True,
331
+ if_global_timing = True,
332
+ if_modulo_timing = False,
333
+ d_beat = 1024, #Beat
334
+ d_oh_beat_type = 7,
335
+ beat_len = 50,
336
+ d_chord = 1024, #Chord
337
+ d_oh_chord_type = 12,
338
+ d_oh_inv_type = 4,
339
+ chord_len = 20,
340
+
341
+ ):
342
+ super().__init__()
343
+
344
+ assert unet_model_name is not None or unet_model_config_path is not None, "Either UNet pretrain model name or a config file path is required"
345
+
346
+ self.text_encoder_name = text_encoder_name
347
+ self.scheduler_name = scheduler_name
348
+ self.unet_model_name = unet_model_name
349
+ self.unet_model_config_path = unet_model_config_path
350
+ self.snr_gamma = snr_gamma
351
+ self.freeze_text_encoder = freeze_text_encoder
352
+ self.uncondition = uncondition
353
+
354
+ # https://huggingface.co/docs/diffusers/v0.14.0/en/api/schedulers/overview
355
+ self.noise_scheduler = DDPMScheduler.from_pretrained(self.scheduler_name, subfolder="scheduler")
356
+ self.inference_scheduler = DDPMScheduler.from_pretrained(self.scheduler_name, subfolder="scheduler")
357
+
358
+ if unet_model_config_path:
359
+ unet_config = UNet2DConditionModelMusic.load_config(unet_model_config_path)
360
+ self.unet = UNet2DConditionModelMusic.from_config(unet_config, subfolder="unet")
361
+ self.set_from = "random"
362
+ print("UNet initialized randomly.")
363
+ else:
364
+ self.unet = UNet2DConditionModel.from_pretrained(unet_model_name, subfolder="unet")
365
+ self.set_from = "pre-trained"
366
+ self.group_in = nn.Sequential(nn.Linear(8, 512), nn.Linear(512, 4))
367
+ self.group_out = nn.Sequential(nn.Linear(4, 512), nn.Linear(512, 8))
368
+ print("UNet initialized from stable diffusion checkpoint.")
369
+
370
+ if "stable-diffusion" in self.text_encoder_name:
371
+ self.tokenizer = CLIPTokenizer.from_pretrained(self.text_encoder_name, subfolder="tokenizer")
372
+ self.text_encoder = CLIPTextModel.from_pretrained(self.text_encoder_name, subfolder="text_encoder")
373
+ elif "t5" in self.text_encoder_name:
374
+ self.tokenizer = AutoTokenizer.from_pretrained(self.text_encoder_name)
375
+ self.text_encoder = T5EncoderModel.from_pretrained(self.text_encoder_name)
376
+ else:
377
+ self.tokenizer = AutoTokenizer.from_pretrained(self.text_encoder_name)
378
+ self.text_encoder = AutoModel.from_pretrained(self.text_encoder_name)
379
+
380
+ self.device = self.text_encoder.device
381
+ #Music Feature Encoder
382
+ self.FME = Fundamental_Music_Embedding(d_model = d_fme, base= base, if_trainable = False, type = fme_type,emb_nn=emb_nn,translation_bias_type = translation_bias_type)
383
+ self.PE = Music_PositionalEncoding(d_model = d_pe, if_index = if_index, if_global_timing = if_global_timing, if_modulo_timing = if_modulo_timing, device = self.device)
384
+ # self.PE2 = Music_PositionalEncoding(d_model = d_pe, if_index = if_index, if_global_timing = if_global_timing, if_modulo_timing = if_modulo_timing, device = self.device)
385
+ self.beat_tokenizer = beat_tokenizer(seq_len_beat=beat_len, if_pad = True)
386
+ self.beat_embedding_layer = Beat_Embedding(self.PE, d_model = d_beat, d_oh_beat_type = d_oh_beat_type)
387
+ self.chord_embedding_layer = Chord_Embedding(self.FME, self.PE, d_model = d_chord, d_oh_type = d_oh_chord_type, d_oh_inv = d_oh_inv_type)
388
+ self.chord_tokenizer = chord_tokenizer(seq_len_chord=chord_len, if_pad = True)
389
+
390
+
391
+ def compute_snr(self, timesteps):
392
+ """
393
+ Computes SNR as per https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849
394
+ """
395
+ alphas_cumprod = self.noise_scheduler.alphas_cumprod
396
+ sqrt_alphas_cumprod = alphas_cumprod**0.5
397
+ sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5
398
+
399
+ # Expand the tensors.
400
+ # Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026
401
+ sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[timesteps].float()
402
+ while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):
403
+ sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]
404
+ alpha = sqrt_alphas_cumprod.expand(timesteps.shape)
405
+
406
+ sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(device=timesteps.device)[timesteps].float()
407
+ while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape):
408
+ sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None]
409
+ sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape)
410
+
411
+ # Compute SNR.
412
+ snr = (alpha / sigma) ** 2
413
+ return snr
414
+
415
+ def encode_text(self, prompt):
416
+ device = self.text_encoder.device
417
+ batch = self.tokenizer(
418
+ prompt, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"
419
+ )
420
+ input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device) #cuda
421
+ if self.freeze_text_encoder:
422
+ with torch.no_grad():
423
+ encoder_hidden_states = self.text_encoder(
424
+ input_ids=input_ids, attention_mask=attention_mask
425
+ )[0] #batch, len_text, dim
426
+ else:
427
+ encoder_hidden_states = self.text_encoder(
428
+ input_ids=input_ids, attention_mask=attention_mask
429
+ )[0]
430
+ boolean_encoder_mask = (attention_mask == 1).to(device) ##batch, len_text
431
+ return encoder_hidden_states, boolean_encoder_mask
432
+
433
+ def encode_beats(self, beats):
434
+ device = self.device
435
+ out_beat = []
436
+ out_beat_timing = []
437
+ out_mask = []
438
+ for beat in beats:
439
+ tokenized_beats,tokenized_beats_timing, tokenized_beat_mask = self.beat_tokenizer(beat)
440
+ out_beat.append(tokenized_beats)
441
+ out_beat_timing.append(tokenized_beats_timing)
442
+ out_mask.append(tokenized_beat_mask)
443
+ out_beat, out_beat_timing, out_mask = torch.tensor(out_beat).to(device), torch.tensor(out_beat_timing).to(device), torch.tensor(out_mask).to(device) #batch, len_beat
444
+ embedded_beat = self.beat_embedding_layer(out_beat, out_beat_timing, device)
445
+
446
+ return embedded_beat, out_mask
447
+
448
+ def encode_chords(self, chords,chords_time):
449
+ device = self.device
450
+ out_chord_root = []
451
+ out_chord_type = []
452
+ out_chord_inv = []
453
+ out_chord_timing = []
454
+ out_mask = []
455
+ for chord, chord_time in zip(chords,chords_time): #batch loop
456
+ tokenized_chord_root, tokenized_chord_type, tokenized_chord_inv, tokenized_chord_time, tokenized_chord_mask = self.chord_tokenizer(chord, chord_time)
457
+ out_chord_root.append(tokenized_chord_root)
458
+ out_chord_type.append(tokenized_chord_type)
459
+ out_chord_inv.append(tokenized_chord_inv)
460
+ out_chord_timing.append(tokenized_chord_time)
461
+ out_mask.append(tokenized_chord_mask)
462
+ #chords: (B, LEN, 4)
463
+ out_chord_root, out_chord_type, out_chord_inv, out_chord_timing, out_mask = torch.tensor(out_chord_root).to(device), torch.tensor(out_chord_type).to(device), torch.tensor(out_chord_inv).to(device), torch.tensor(out_chord_timing).to(device), torch.tensor(out_mask).to(device)
464
+ embedded_chord = self.chord_embedding_layer(out_chord_root, out_chord_type, out_chord_inv, out_chord_timing, device)
465
+ return embedded_chord, out_mask
466
+ # return out_chord_root, out_mask
467
+
468
+
469
+ def forward(self, latents, prompt, beats, chords,chords_time, validation_mode=False):
470
+ device = self.text_encoder.device
471
+ num_train_timesteps = self.noise_scheduler.num_train_timesteps
472
+ self.noise_scheduler.set_timesteps(num_train_timesteps, device=device)
473
+
474
+ encoder_hidden_states, boolean_encoder_mask = self.encode_text(prompt)
475
+
476
+ # with torch.no_grad():
477
+ encoded_beats, beat_mask = self.encode_beats(beats) #batch, len_beats, dim; batch, len_beats
478
+ encoded_chords, chord_mask = self.encode_chords(chords,chords_time)
479
+
480
+
481
+ if self.uncondition:
482
+ mask_indices = [k for k in range(len(prompt)) if random.random() < 0.1]
483
+ if len(mask_indices) > 0:
484
+ encoder_hidden_states[mask_indices] = 0
485
+ encoded_chords[mask_indices] = 0
486
+ encoded_beats[mask_indices] = 0
487
+
488
+ bsz = latents.shape[0]
489
+
490
+ if validation_mode:
491
+ timesteps = (self.noise_scheduler.num_train_timesteps//2) * torch.ones((bsz,), dtype=torch.int64, device=device)
492
+ else:
493
+ timesteps = torch.randint(0, self.noise_scheduler.num_train_timesteps, (bsz,), device=device)
494
+
495
+
496
+ timesteps = timesteps.long()
497
+
498
+ noise = torch.randn_like(latents)
499
+ noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)
500
+
501
+ # Get the target for loss depending on the prediction type
502
+ if self.noise_scheduler.config.prediction_type == "epsilon":
503
+ target = noise
504
+ elif self.noise_scheduler.config.prediction_type == "v_prediction":
505
+ target = self.noise_scheduler.get_velocity(latents, noise, timesteps)
506
+ else:
507
+ raise ValueError(f"Unknown prediction type {self.noise_scheduler.config.prediction_type}")
508
+
509
+ if self.set_from == "random":
510
+ # model_pred = torch.zeros((bsz,8,256,16)).to(device)
511
+ model_pred = self.unet(
512
+ noisy_latents, timesteps, encoder_hidden_states, encoded_beats, encoded_chords,
513
+ encoder_attention_mask=boolean_encoder_mask, beat_attention_mask = beat_mask, chord_attention_mask = chord_mask
514
+ ).sample
515
+
516
+ elif self.set_from == "pre-trained":
517
+ compressed_latents = self.group_in(noisy_latents.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()
518
+ model_pred = self.unet(
519
+ compressed_latents, timesteps, encoder_hidden_states,
520
+ encoder_attention_mask=boolean_encoder_mask
521
+ ).sample
522
+ model_pred = self.group_out(model_pred.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()
523
+
524
+ if self.snr_gamma is None:
525
+ loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
526
+ else:
527
+ # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
528
+ # Adaptef from huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py
529
+ snr = self.compute_snr(timesteps)
530
+ mse_loss_weights = (
531
+ torch.stack([snr, self.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr
532
+ )
533
+ loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
534
+ loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights
535
+ loss = loss.mean()
536
+
537
+ return loss
538
+
539
+ @torch.no_grad()
540
+ def inference(self, prompt, beats, chords,chords_time, inference_scheduler, num_steps=20, guidance_scale=3, num_samples_per_prompt=1,
541
+ disable_progress=True):
542
+ device = self.text_encoder.device
543
+ classifier_free_guidance = guidance_scale > 1.0
544
+ batch_size = len(prompt) * num_samples_per_prompt
545
+
546
+ if classifier_free_guidance:
547
+ prompt_embeds, boolean_prompt_mask = self.encode_text_classifier_free(prompt, num_samples_per_prompt)
548
+ encoded_beats, beat_mask = self.encode_beats_classifier_free(beats, num_samples_per_prompt) #batch, len_beats, dim; batch, len_beats
549
+ encoded_chords, chord_mask = self.encode_chords_classifier_free(chords, chords_time, num_samples_per_prompt)
550
+ else:
551
+ prompt_embeds, boolean_prompt_mask = self.encode_text(prompt)
552
+ prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
553
+ boolean_prompt_mask = boolean_prompt_mask.repeat_interleave(num_samples_per_prompt, 0)
554
+
555
+ encoded_beats, beat_mask = self.encode_beats(beats) #batch, len_beats, dim; batch, len_beats
556
+ encoded_beats = encoded_beats.repeat_interleave(num_samples_per_prompt, 0)
557
+ beat_mask = beat_mask.repeat_interleave(num_samples_per_prompt, 0)
558
+
559
+ encoded_chords, chord_mask = self.encode_chords(chords,chords_time)
560
+ encoded_chords = encoded_chords.repeat_interleave(num_samples_per_prompt, 0)
561
+ chord_mask = chord_mask.repeat_interleave(num_samples_per_prompt, 0)
562
+
563
+ # print(f"encoded_chords:{encoded_chords.shape}, chord_mask:{chord_mask.shape}, prompt_embeds:{prompt_embeds.shape},boolean_prompt_mask:{boolean_prompt_mask.shape} ")
564
+ inference_scheduler.set_timesteps(num_steps, device=device)
565
+ timesteps = inference_scheduler.timesteps
566
+
567
+ num_channels_latents = self.unet.in_channels
568
+ latents = self.prepare_latents(batch_size, inference_scheduler, num_channels_latents, prompt_embeds.dtype, device)
569
+
570
+ num_warmup_steps = len(timesteps) - num_steps * inference_scheduler.order
571
+ progress_bar = tqdm(range(num_steps), disable=disable_progress)
572
+
573
+ for i, t in enumerate(timesteps):
574
+ # expand the latents if we are doing classifier free guidance
575
+ latent_model_input = torch.cat([latents] * 2) if classifier_free_guidance else latents
576
+ latent_model_input = inference_scheduler.scale_model_input(latent_model_input, t)
577
+
578
+ noise_pred = self.unet(
579
+ latent_model_input, t, encoder_hidden_states=prompt_embeds,
580
+ encoder_attention_mask=boolean_prompt_mask,
581
+ beat_features = encoded_beats, beat_attention_mask = beat_mask, chord_features = encoded_chords,chord_attention_mask = chord_mask
582
+ ).sample
583
+
584
+ # perform guidance
585
+ if classifier_free_guidance: #should work for beats and chords too
586
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
587
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
588
+
589
+ # compute the previous noisy sample x_t -> x_t-1
590
+ latents = inference_scheduler.step(noise_pred, t, latents).prev_sample
591
+
592
+ # call the callback, if provided
593
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % inference_scheduler.order == 0):
594
+ progress_bar.update(1)
595
+
596
+ if self.set_from == "pre-trained":
597
+ latents = self.group_out(latents.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()
598
+ return latents
599
+
600
+ def prepare_latents(self, batch_size, inference_scheduler, num_channels_latents, dtype, device):
601
+ shape = (batch_size, num_channels_latents, 256, 16)
602
+ latents = randn_tensor(shape, generator=None, device=device, dtype=dtype)
603
+ # scale the initial noise by the standard deviation required by the scheduler
604
+ latents = latents * inference_scheduler.init_noise_sigma
605
+ return latents
606
+
607
+ def encode_text_classifier_free(self, prompt, num_samples_per_prompt):
608
+ device = self.text_encoder.device
609
+ batch = self.tokenizer(
610
+ prompt, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"
611
+ )
612
+ input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)
613
+
614
+ with torch.no_grad():
615
+ prompt_embeds = self.text_encoder(
616
+ input_ids=input_ids, attention_mask=attention_mask
617
+ )[0]
618
+
619
+ prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
620
+ attention_mask = attention_mask.repeat_interleave(num_samples_per_prompt, 0)
621
+
622
+ # get unconditional embeddings for classifier free guidance
623
+ # print(len(prompt), 'this is prompt len')
624
+ uncond_tokens = [""] * len(prompt)
625
+
626
+ max_length = prompt_embeds.shape[1]
627
+ uncond_batch = self.tokenizer(
628
+ uncond_tokens, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt",
629
+ )
630
+ uncond_input_ids = uncond_batch.input_ids.to(device)
631
+ uncond_attention_mask = uncond_batch.attention_mask.to(device)
632
+
633
+ with torch.no_grad():
634
+ negative_prompt_embeds = self.text_encoder(
635
+ input_ids=uncond_input_ids, attention_mask=uncond_attention_mask
636
+ )[0]
637
+
638
+ negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
639
+ uncond_attention_mask = uncond_attention_mask.repeat_interleave(num_samples_per_prompt, 0)
640
+
641
+ # For classifier free guidance, we need to do two forward passes.
642
+ # We concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes
643
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
644
+ prompt_mask = torch.cat([uncond_attention_mask, attention_mask])
645
+ boolean_prompt_mask = (prompt_mask == 1).to(device)
646
+
647
+ return prompt_embeds, boolean_prompt_mask
648
+
649
+
650
+ def encode_beats_classifier_free(self, beats, num_samples_per_prompt):
651
+ device = self.device
652
+ with torch.no_grad():
653
+ out_beat = []
654
+ out_beat_timing = []
655
+ out_mask = []
656
+ for beat in beats:
657
+ tokenized_beats,tokenized_beats_timing, tokenized_beat_mask = self.beat_tokenizer(beat)
658
+ out_beat.append(tokenized_beats)
659
+ out_beat_timing.append(tokenized_beats_timing)
660
+ out_mask.append(tokenized_beat_mask)
661
+ out_beat, out_beat_timing, out_mask = torch.tensor(out_beat).to(device), torch.tensor(out_beat_timing).to(device), torch.tensor(out_mask).to(device) #batch, len_beat
662
+ embedded_beat = self.beat_embedding_layer(out_beat, out_beat_timing, device)
663
+
664
+ embedded_beat = embedded_beat.repeat_interleave(num_samples_per_prompt, 0)
665
+ out_mask = out_mask.repeat_interleave(num_samples_per_prompt, 0)
666
+
667
+ uncond_beats = [[[],[]]] * len(beats)
668
+
669
+ max_length = embedded_beat.shape[1]
670
+ with torch.no_grad():
671
+ out_beat_unc = []
672
+ out_beat_timing_unc = []
673
+ out_mask_unc = []
674
+ for beat in uncond_beats:
675
+ tokenized_beats, tokenized_beats_timing, tokenized_beat_mask = self.beat_tokenizer(beat)
676
+ out_beat_unc.append(tokenized_beats)
677
+ out_beat_timing_unc.append(tokenized_beats_timing)
678
+ out_mask_unc.append(tokenized_beat_mask)
679
+ out_beat_unc, out_beat_timing_unc, out_mask_unc = torch.tensor(out_beat_unc).to(device), torch.tensor(out_beat_timing_unc).to(device), torch.tensor(out_mask_unc).to(device) #batch, len_beat
680
+ embedded_beat_unc = self.beat_embedding_layer(out_beat_unc, out_beat_timing_unc, device)
681
+
682
+ embedded_beat_unc = embedded_beat_unc.repeat_interleave(num_samples_per_prompt, 0)
683
+ out_mask_unc = out_mask_unc.repeat_interleave(num_samples_per_prompt, 0)
684
+
685
+ embedded_beat = torch.cat([embedded_beat_unc, embedded_beat])
686
+ out_mask = torch.cat([out_mask_unc, out_mask])
687
+
688
+ return embedded_beat, out_mask
689
+
690
+
691
+ def encode_chords_classifier_free(self, chords, chords_time, num_samples_per_prompt):
692
+ device = self.device
693
+ with torch.no_grad():
694
+ out_chord_root = []
695
+ out_chord_type = []
696
+ out_chord_inv = []
697
+ out_chord_timing = []
698
+ out_mask = []
699
+ for chord, chord_time in zip(chords,chords_time): #batch loop
700
+ tokenized_chord_root, tokenized_chord_type, tokenized_chord_inv, tokenized_chord_time, tokenized_chord_mask = self.chord_tokenizer(chord, chord_time)
701
+ out_chord_root.append(tokenized_chord_root)
702
+ out_chord_type.append(tokenized_chord_type)
703
+ out_chord_inv.append(tokenized_chord_inv)
704
+ out_chord_timing.append(tokenized_chord_time)
705
+ out_mask.append(tokenized_chord_mask)
706
+ out_chord_root, out_chord_type, out_chord_inv, out_chord_timing, out_mask = torch.tensor(out_chord_root).to(device), torch.tensor(out_chord_type).to(device), torch.tensor(out_chord_inv).to(device), torch.tensor(out_chord_timing).to(device), torch.tensor(out_mask).to(device)
707
+ embedded_chord = self.chord_embedding_layer(out_chord_root, out_chord_type, out_chord_inv, out_chord_timing, device)
708
+
709
+ embedded_chord = embedded_chord.repeat_interleave(num_samples_per_prompt, 0)
710
+ out_mask = out_mask.repeat_interleave(num_samples_per_prompt, 0)
711
+
712
+ chords_unc=[[]] * len(chords)
713
+ chords_time_unc=[[]] * len(chords_time)
714
+
715
+ max_length = embedded_chord.shape[1]
716
+
717
+ with torch.no_grad():
718
+ out_chord_root_unc = []
719
+ out_chord_type_unc = []
720
+ out_chord_inv_unc = []
721
+ out_chord_timing_unc = []
722
+ out_mask_unc = []
723
+ for chord, chord_time in zip(chords_unc,chords_time_unc): #batch loop
724
+ tokenized_chord_root, tokenized_chord_type, tokenized_chord_inv, tokenized_chord_time, tokenized_chord_mask = self.chord_tokenizer(chord, chord_time)
725
+ out_chord_root_unc.append(tokenized_chord_root)
726
+ out_chord_type_unc.append(tokenized_chord_type)
727
+ out_chord_inv_unc.append(tokenized_chord_inv)
728
+ out_chord_timing_unc.append(tokenized_chord_time)
729
+ out_mask_unc.append(tokenized_chord_mask)
730
+ out_chord_root_unc, out_chord_type_unc, out_chord_inv_unc, out_chord_timing_unc, out_mask_unc = torch.tensor(out_chord_root_unc).to(device), torch.tensor(out_chord_type_unc).to(device), torch.tensor(out_chord_inv_unc).to(device), torch.tensor(out_chord_timing_unc).to(device), torch.tensor(out_mask_unc).to(device)
731
+ embedded_chord_unc = self.chord_embedding_layer(out_chord_root_unc, out_chord_type_unc, out_chord_inv_unc, out_chord_timing_unc, device)
732
+
733
+
734
+ embedded_chord_unc = embedded_chord_unc.repeat_interleave(num_samples_per_prompt, 0)
735
+ out_mask_unc = out_mask_unc.repeat_interleave(num_samples_per_prompt, 0)
736
+
737
+ embedded_chord = torch.cat([embedded_chord_unc, embedded_chord])
738
+ out_mask = torch.cat([out_mask_unc, out_mask])
739
+
740
+ return embedded_chord, out_mask
.ipynb_checkpoints/requirements-checkpoint.txt ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch==2.0.1
2
+ torchaudio==2.0.2
3
+ torchvision==0.15.2
4
+ transformers==4.31.0
5
+ accelerate==0.21.0
6
+ datasets==2.1.0
7
+ einops==0.6.1
8
+ h5py==3.8.0
9
+ huggingface_hub==0.19.4
10
+ importlib_metadata==6.3.0
11
+ librosa==0.9.2
12
+ matplotlib==3.5.2
13
+ numpy==1.23.0
14
+ omegaconf==2.3.0
15
+ packaging==23.1
16
+ pandas==1.4.1
17
+ progressbar33==2.4
18
+ protobuf==3.20.*
19
+ resampy==0.4.2
20
+ safetensors==0.3.2
21
+ sentencepiece==0.1.99
22
+ scikit_image==0.19.3
23
+ scikit_learn==1.2.2
24
+ scipy==1.8.0
25
+ soundfile==0.12.1
26
+ ssr_eval==0.0.6
27
+ torchlibrosa==0.1.0
28
+ tqdm==4.63.1
29
+ wandb==0.12.14
30
+ ipython==8.12.0
31
+ gradio==4.3.0
32
+ wavio==0.0.7
app.py CHANGED
@@ -144,12 +144,13 @@ class MusicFeaturePredictor:
144
  class Mustango:
145
  def __init__(
146
  self,
147
- name="declare-lab/mustango",
148
  device="cuda:0",
149
  cache_dir=None,
150
  local_files_only=False,
151
  ):
152
- path = snapshot_download(repo_id=name, cache_dir=cache_dir)
 
 
153
 
154
  self.music_model = MusicFeaturePredictor(
155
  path, device, cache_dir=cache_dir, local_files_only=local_files_only
@@ -174,15 +175,16 @@ class Mustango:
174
  stft_weights = torch.load(
175
  f"{path}/stft/pytorch_model_stft.bin", map_location=device
176
  )
177
- main_weights = torch.load(
178
  f"{path}/ldm/pytorch_model_ldm.bin", map_location=device
179
  )
 
 
 
 
180
 
181
  self.vae.load_state_dict(vae_weights)
182
  self.stft.load_state_dict(stft_weights)
183
- self.model.load_state_dict(main_weights)
184
-
185
- print("Successfully loaded checkpoint from:", name)
186
 
187
  self.vae.eval()
188
  self.stft.eval()
@@ -192,8 +194,13 @@ class Mustango:
192
  main_config["scheduler_name"], subfolder="scheduler"
193
  )
194
 
195
- def generate(self, prompt, steps=100, guidance=3, samples=1, disable_progress=True):
196
  """Genrate music for a single prompt string."""
 
 
 
 
 
197
 
198
  with torch.no_grad():
199
  beats, chords, chords_times = self.music_model.generate(prompt)
@@ -222,8 +229,8 @@ else:
222
 
223
  # output_wave = mustango.generate("This techno song features a synth lead playing the main melody.", 5, 3, disable_progress=False)
224
 
225
- def gradio_generate(prompt, steps, guidance):
226
- output_wave = mustango.generate(prompt, steps, guidance)
227
  # output_filename = f"{prompt.replace(' ', '_')}_{steps}_{guidance}"[:250] + ".wav"
228
  output_filename = "temp.wav"
229
  wavio.write(output_filename, output_wave, rate=16000, sampwidth=2)
@@ -244,6 +251,8 @@ input_text = gr.Textbox(lines=2, label="Prompt")
244
  output_audio = gr.Audio(label="Generated Music", type="filepath")
245
  denoising_steps = gr.Slider(minimum=100, maximum=200, value=100, step=1, label="Steps", interactive=True)
246
  guidance_scale = gr.Slider(minimum=1, maximum=10, value=3, step=0.1, label="Guidance Scale", interactive=True)
 
 
247
 
248
  # CSS styling for the Duplicate button
249
  css = '''
@@ -258,17 +267,17 @@ css = '''
258
  # Gradio interface
259
  gr_interface = gr.Interface(
260
  fn=gradio_generate,
261
- inputs=[input_text, denoising_steps, guidance_scale],
262
  outputs=[output_audio],
263
  description=description_text,
264
  allow_flagging=False,
265
  examples=[
266
- ["This techno song features a synth lead playing the main melody. This is accompanied by programmed percussion playing a simple kick focused beat. The hi-hat is accented in an open position on the 3-and count of every bar. The synth plays the bass part with a voicing that sounds like a cello. This techno song can be played in a club. The chord sequence is Gm, A7, Eb, Bb, C, F, Gm. The beat counts to 2. The tempo of this song is 128.0 beats per minute. The key of this song is G minor.", 100, 3],
267
- ["This is a new age piece. There is a flute playing the main melody with a lot of staccato notes. The rhythmic background consists of a medium tempo electronic drum beat with percussive elements all over the spectrum. There is a playful atmosphere to the piece. This piece can be used in the soundtrack of a children's TV show or an advertisement jingle.", 100, 3],
268
- ["The song is an instrumental. The song is in medium tempo with a classical guitar playing a lilting melody in accompaniment style. The song is emotional and romantic. The song is a romantic instrumental song. The chord sequence is Gm, F6, Ebm. The time signature is 4/4. This song is in Adagio. The key of this song is G minor.", 100, 3],
269
- ["This folk song features a female voice singing the main melody. This is accompanied by a tabla playing the percussion. A guitar strums chords. For most parts of the song, only one chord is played. At the last bar, a different chord is played. This song has minimal instruments. This song has a story-telling mood. This song can be played in a village scene in an Indian movie. The chord sequence is Bbm, Ab. The beat is 3. The tempo of this song is Allegro. The key of this song is Bb minor.", 100, 3],
270
- ["This is a live performance of a classical music piece. There is an orchestra performing the piece with a violin lead playing the main melody. The atmosphere is sentimental and heart-touching. This piece could be playing in the background at a classy restaurant. The chord progression in this song is Am7, Gm, Dm, A7, Dm. The beat is 3. This song is in Largo. The key of this song is D minor.", 100, 3],
271
- ["This is a techno piece with drums and beats and a leading melody. A synth plays chords. The music kicks off with a powerful and relentless drumbeat. Over the pounding beats, a leading melody emerges. In the middle of the song, a flock of seagulls flies over the venue and make loud bird sounds. It has strong danceability and can be played in a club. The tempo is 120 bpm. The chords played by the synth are Am, Cm, Dm, Gm.", 100, 3],
272
  ],
273
  cache_examples=True,
274
  )
 
144
  class Mustango:
145
  def __init__(
146
  self,
 
147
  device="cuda:0",
148
  cache_dir=None,
149
  local_files_only=False,
150
  ):
151
+
152
+ path = snapshot_download(repo_id="declare-lab/mustango", cache_dir=cache_dir)
153
+ pretrained_path = snapshot_download(repo_id="declare-lab/mustango-pretrained", cache_dir=cache_dir)
154
 
155
  self.music_model = MusicFeaturePredictor(
156
  path, device, cache_dir=cache_dir, local_files_only=local_files_only
 
175
  stft_weights = torch.load(
176
  f"{path}/stft/pytorch_model_stft.bin", map_location=device
177
  )
178
+ self.main_weights = torch.load(
179
  f"{path}/ldm/pytorch_model_ldm.bin", map_location=device
180
  )
181
+
182
+ self.main_weights_pretrained = torch.load(
183
+ f"{pretrained_path}/ldm/pytorch_model_ldm.bin", map_location=device
184
+ )
185
 
186
  self.vae.load_state_dict(vae_weights)
187
  self.stft.load_state_dict(stft_weights)
 
 
 
188
 
189
  self.vae.eval()
190
  self.stft.eval()
 
194
  main_config["scheduler_name"], subfolder="scheduler"
195
  )
196
 
197
+ def generate(self, model_name, prompt, steps=100, guidance=3, samples=1, disable_progress=True):
198
  """Genrate music for a single prompt string."""
199
+
200
+ if model_name == "declare-lab/mustango":
201
+ self.model.load_state_dict(self.main_weights)
202
+ else:
203
+ self.model.load_state_dict(self.main_weights_pretrained)
204
 
205
  with torch.no_grad():
206
  beats, chords, chords_times = self.music_model.generate(prompt)
 
229
 
230
  # output_wave = mustango.generate("This techno song features a synth lead playing the main melody.", 5, 3, disable_progress=False)
231
 
232
+ def gradio_generate(model_name, prompt, steps, guidance):
233
+ output_wave = mustango.generate(model_name, prompt, steps, guidance)
234
  # output_filename = f"{prompt.replace(' ', '_')}_{steps}_{guidance}"[:250] + ".wav"
235
  output_filename = "temp.wav"
236
  wavio.write(output_filename, output_wave, rate=16000, sampwidth=2)
 
251
  output_audio = gr.Audio(label="Generated Music", type="filepath")
252
  denoising_steps = gr.Slider(minimum=100, maximum=200, value=100, step=1, label="Steps", interactive=True)
253
  guidance_scale = gr.Slider(minimum=1, maximum=10, value=3, step=0.1, label="Guidance Scale", interactive=True)
254
+ model_name = gr.Radio(["declare-lab/mustango", "declare-lab/mustango-pretrained",], label="Choose a model type", value="declare-lab/mustango", type="value")
255
+
256
 
257
  # CSS styling for the Duplicate button
258
  css = '''
 
267
  # Gradio interface
268
  gr_interface = gr.Interface(
269
  fn=gradio_generate,
270
+ inputs=[model_name, input_text, denoising_steps, guidance_scale],
271
  outputs=[output_audio],
272
  description=description_text,
273
  allow_flagging=False,
274
  examples=[
275
+ ["declare-lab/mustango", "This techno song features a synth lead playing the main melody. This is accompanied by programmed percussion playing a simple kick focused beat. The hi-hat is accented in an open position on the 3-and count of every bar. The synth plays the bass part with a voicing that sounds like a cello. This techno song can be played in a club. The chord sequence is Gm, A7, Eb, Bb, C, F, Gm. The beat counts to 2. The tempo of this song is 128.0 beats per minute. The key of this song is G minor.", 200, 3],
276
+ ["declare-lab/mustango", "This is a new age piece. There is a flute playing the main melody with a lot of staccato notes. The rhythmic background consists of a medium tempo electronic drum beat with percussive elements all over the spectrum. There is a playful atmosphere to the piece. This piece can be used in the soundtrack of a children's TV show or an advertisement jingle.", 200, 3],
277
+ ["declare-lab/mustango", "The song is an instrumental. The song is in medium tempo with a classical guitar playing a lilting melody in accompaniment style. The song is emotional and romantic. The song is a romantic instrumental song. The chord sequence is Gm, F6, Ebm. The time signature is 4/4. This song is in Adagio. The key of this song is G minor.", 200, 3],
278
+ ["declare-lab/mustango", "This folk song features a female voice singing the main melody. This is accompanied by a tabla playing the percussion. A guitar strums chords. For most parts of the song, only one chord is played. At the last bar, a different chord is played. This song has minimal instruments. This song has a story-telling mood. This song can be played in a village scene in an Indian movie. The chord sequence is Bbm, Ab. The beat is 3. The tempo of this song is Allegro. The key of this song is Bb minor.", 200, 3],
279
+ ["declare-lab/mustango", "This is a live performance of a classical music piece. There is an orchestra performing the piece with a violin lead playing the main melody. The atmosphere is sentimental and heart-touching. This piece could be playing in the background at a classy restaurant. The chord progression in this song is Am7, Gm, Dm, A7, Dm. The beat is 3. This song is in Largo. The key of this song is D minor.", 200, 3],
280
+ ["declare-lab/mustango", "This is a techno piece with drums and beats and a leading melody. A synth plays chords. The music kicks off with a powerful and relentless drumbeat. Over the pounding beats, a leading melody emerges. In the middle of the song, a flock of seagulls flies over the venue and make loud bird sounds. It has strong danceability and can be played in a club. The tempo is 120 bpm. The chords played by the synth are Am, Cm, Dm, Gm.", 200, 3],
281
  ],
282
  cache_examples=True,
283
  )