deepanway commited on
Commit
f1069cc
1 Parent(s): 906f008

Uplaod files

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .ipynb_checkpoints/README-checkpoint.md +13 -0
  2. .ipynb_checkpoints/app-checkpoint.py +260 -0
  3. .ipynb_checkpoints/modelling_deberta_v2-checkpoint.py +1750 -0
  4. .ipynb_checkpoints/models-checkpoint.py +738 -0
  5. app.py +260 -0
  6. audioldm/__init__.py +8 -0
  7. audioldm/__main__.py +183 -0
  8. audioldm/__pycache__/__init__.cpython-311.pyc +0 -0
  9. audioldm/__pycache__/__init__.cpython-37.pyc +0 -0
  10. audioldm/__pycache__/__init__.cpython-39.pyc +0 -0
  11. audioldm/__pycache__/ldm.cpython-311.pyc +0 -0
  12. audioldm/__pycache__/ldm.cpython-37.pyc +0 -0
  13. audioldm/__pycache__/ldm.cpython-39.pyc +0 -0
  14. audioldm/__pycache__/pipeline.cpython-311.pyc +0 -0
  15. audioldm/__pycache__/pipeline.cpython-37.pyc +0 -0
  16. audioldm/__pycache__/pipeline.cpython-39.pyc +0 -0
  17. audioldm/__pycache__/utils.cpython-311.pyc +0 -0
  18. audioldm/__pycache__/utils.cpython-37.pyc +0 -0
  19. audioldm/__pycache__/utils.cpython-39.pyc +0 -0
  20. audioldm/audio/__init__.py +2 -0
  21. audioldm/audio/__pycache__/__init__.cpython-311.pyc +0 -0
  22. audioldm/audio/__pycache__/__init__.cpython-37.pyc +0 -0
  23. audioldm/audio/__pycache__/__init__.cpython-39.pyc +0 -0
  24. audioldm/audio/__pycache__/audio_processing.cpython-311.pyc +0 -0
  25. audioldm/audio/__pycache__/audio_processing.cpython-37.pyc +0 -0
  26. audioldm/audio/__pycache__/audio_processing.cpython-39.pyc +0 -0
  27. audioldm/audio/__pycache__/mix.cpython-39.pyc +0 -0
  28. audioldm/audio/__pycache__/stft.cpython-311.pyc +0 -0
  29. audioldm/audio/__pycache__/stft.cpython-37.pyc +0 -0
  30. audioldm/audio/__pycache__/stft.cpython-39.pyc +0 -0
  31. audioldm/audio/__pycache__/tools.cpython-311.pyc +0 -0
  32. audioldm/audio/__pycache__/tools.cpython-37.pyc +0 -0
  33. audioldm/audio/__pycache__/tools.cpython-39.pyc +0 -0
  34. audioldm/audio/__pycache__/torch_tools.cpython-39.pyc +0 -0
  35. audioldm/audio/audio_processing.py +100 -0
  36. audioldm/audio/stft.py +186 -0
  37. audioldm/audio/tools.py +85 -0
  38. audioldm/clap/__init__.py +0 -0
  39. audioldm/clap/__pycache__/__init__.cpython-39.pyc +0 -0
  40. audioldm/clap/__pycache__/encoders.cpython-39.pyc +0 -0
  41. audioldm/clap/encoders.py +170 -0
  42. audioldm/clap/open_clip/__init__.py +25 -0
  43. audioldm/clap/open_clip/__pycache__/__init__.cpython-39.pyc +0 -0
  44. audioldm/clap/open_clip/__pycache__/factory.cpython-39.pyc +0 -0
  45. audioldm/clap/open_clip/__pycache__/feature_fusion.cpython-39.pyc +0 -0
  46. audioldm/clap/open_clip/__pycache__/htsat.cpython-39.pyc +0 -0
  47. audioldm/clap/open_clip/__pycache__/loss.cpython-39.pyc +0 -0
  48. audioldm/clap/open_clip/__pycache__/model.cpython-39.pyc +0 -0
  49. audioldm/clap/open_clip/__pycache__/openai.cpython-39.pyc +0 -0
  50. audioldm/clap/open_clip/__pycache__/pann_model.cpython-39.pyc +0 -0
.ipynb_checkpoints/README-checkpoint.md ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: Mustango
3
+ emoji: 🐢
4
+ colorFrom: indigo
5
+ colorTo: green
6
+ sdk: gradio
7
+ sdk_version: 4.3.0
8
+ app_file: app.py
9
+ pinned: false
10
+ license: apache-2.0
11
+ ---
12
+
13
+ Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
.ipynb_checkpoints/app-checkpoint.py ADDED
@@ -0,0 +1,260 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import json
3
+ import torch
4
+ import wavio
5
+ from tqdm import tqdm
6
+ from huggingface_hub import snapshot_download
7
+
8
+ from audioldm.audio.stft import TacotronSTFT
9
+ from audioldm.variational_autoencoder import AutoencoderKL
10
+
11
+ from transformers import AutoTokenizer, T5ForConditionalGeneration
12
+ from modelling_deberta_v2 import DebertaV2ForTokenClassificationRegression
13
+
14
+ from diffusers import DDPMScheduler
15
+ from models import MusicAudioDiffusion
16
+
17
+ from gradio import Markdown
18
+
19
+ class MusicFeaturePredictor:
20
+ def __init__(self, path, device="cuda:0", cache_dir=None, local_files_only=False):
21
+ self.beats_tokenizer = AutoTokenizer.from_pretrained(
22
+ "microsoft/deberta-v3-large",
23
+ cache_dir=cache_dir,
24
+ local_files_only=local_files_only,
25
+ )
26
+ self.beats_model = DebertaV2ForTokenClassificationRegression.from_pretrained(
27
+ "microsoft/deberta-v3-large",
28
+ cache_dir=cache_dir,
29
+ local_files_only=local_files_only,
30
+ )
31
+ self.beats_model.eval()
32
+ self.beats_model.to(device)
33
+
34
+ beats_ckpt = f"{path}/beats/microsoft-deberta-v3-large.pt"
35
+ beats_weight = torch.load(beats_ckpt, map_location="cpu")
36
+ self.beats_model.load_state_dict(beats_weight)
37
+
38
+ self.chords_tokenizer = AutoTokenizer.from_pretrained(
39
+ "google/flan-t5-large",
40
+ cache_dir=cache_dir,
41
+ local_files_only=local_files_only,
42
+ )
43
+ self.chords_model = T5ForConditionalGeneration.from_pretrained(
44
+ "google/flan-t5-large",
45
+ cache_dir=cache_dir,
46
+ local_files_only=local_files_only,
47
+ )
48
+ self.chords_model.eval()
49
+ self.chords_model.to(device)
50
+
51
+ chords_ckpt = f"{path}/chords/flan-t5-large.bin"
52
+ chords_weight = torch.load(chords_ckpt, map_location="cpu")
53
+ self.chords_model.load_state_dict(chords_weight)
54
+
55
+ def generate_beats(self, prompt):
56
+ tokenized = self.beats_tokenizer(
57
+ prompt, max_length=512, padding=True, truncation=True, return_tensors="pt"
58
+ )
59
+ tokenized = {k: v.to(self.beats_model.device) for k, v in tokenized.items()}
60
+
61
+ with torch.no_grad():
62
+ out = self.beats_model(**tokenized)
63
+
64
+ max_beat = (
65
+ 1 + torch.argmax(out["logits"][:, 0, :], -1).detach().cpu().numpy()
66
+ ).tolist()[0]
67
+ intervals = (
68
+ out["values"][:, :, 0]
69
+ .detach()
70
+ .cpu()
71
+ .numpy()
72
+ .astype("float32")
73
+ .round(4)
74
+ .tolist()
75
+ )
76
+
77
+ intervals = np.cumsum(intervals)
78
+ predicted_beats_times = []
79
+ for t in intervals:
80
+ if t < 10:
81
+ predicted_beats_times.append(round(t, 2))
82
+ else:
83
+ break
84
+ predicted_beats_times = list(np.array(predicted_beats_times)[:50])
85
+
86
+ if len(predicted_beats_times) == 0:
87
+ predicted_beats = [[], []]
88
+ else:
89
+ beat_counts = []
90
+ for i in range(len(predicted_beats_times)):
91
+ beat_counts.append(float(1.0 + np.mod(i, max_beat)))
92
+ predicted_beats = [[predicted_beats_times, beat_counts]]
93
+
94
+ return max_beat, predicted_beats_times, predicted_beats
95
+
96
+ def generate(self, prompt):
97
+ max_beat, predicted_beats_times, predicted_beats = self.generate_beats(prompt)
98
+
99
+ chords_prompt = "Caption: {} \\n Timestamps: {} \\n Max Beat: {}".format(
100
+ prompt,
101
+ " , ".join([str(round(t, 2)) for t in predicted_beats_times]),
102
+ max_beat,
103
+ )
104
+
105
+ tokenized = self.chords_tokenizer(
106
+ chords_prompt,
107
+ max_length=512,
108
+ padding=True,
109
+ truncation=True,
110
+ return_tensors="pt",
111
+ )
112
+ tokenized = {k: v.to(self.chords_model.device) for k, v in tokenized.items()}
113
+
114
+ generated_chords = self.chords_model.generate(
115
+ input_ids=tokenized["input_ids"],
116
+ attention_mask=tokenized["attention_mask"],
117
+ min_length=8,
118
+ max_length=128,
119
+ num_beams=5,
120
+ early_stopping=True,
121
+ num_return_sequences=1,
122
+ )
123
+
124
+ generated_chords = self.chords_tokenizer.decode(
125
+ generated_chords[0],
126
+ skip_special_tokens=True,
127
+ clean_up_tokenization_spaces=True,
128
+ ).split(" n ")
129
+
130
+ predicted_chords, predicted_chords_times = [], []
131
+ for item in generated_chords:
132
+ c, ct = item.split(" at ")
133
+ predicted_chords.append(c)
134
+ predicted_chords_times.append(float(ct))
135
+
136
+ return predicted_beats, predicted_chords, predicted_chords_times
137
+
138
+
139
+ class Mustango:
140
+ def __init__(
141
+ self,
142
+ name="declare-lab/mustango",
143
+ device="cuda:0",
144
+ cache_dir=None,
145
+ local_files_only=False,
146
+ ):
147
+ path = snapshot_download(repo_id=name, cache_dir=cache_dir)
148
+
149
+ self.music_model = MusicFeaturePredictor(
150
+ path, device, cache_dir=cache_dir, local_files_only=local_files_only
151
+ )
152
+
153
+ vae_config = json.load(open(f"{path}/configs/vae_config.json"))
154
+ stft_config = json.load(open(f"{path}/configs/stft_config.json"))
155
+ main_config = json.load(open(f"{path}/configs/main_config.json"))
156
+
157
+ self.vae = AutoencoderKL(**vae_config).to(device)
158
+ self.stft = TacotronSTFT(**stft_config).to(device)
159
+ self.model = MusicAudioDiffusion(
160
+ main_config["text_encoder_name"],
161
+ main_config["scheduler_name"],
162
+ unet_model_config_path=f"{path}/configs/music_diffusion_model_config.json",
163
+ ).to(device)
164
+
165
+ vae_weights = torch.load(
166
+ f"{path}/vae/pytorch_model_vae.bin", map_location=device
167
+ )
168
+ stft_weights = torch.load(
169
+ f"{path}/stft/pytorch_model_stft.bin", map_location=device
170
+ )
171
+ main_weights = torch.load(
172
+ f"{path}/ldm/pytorch_model_ldm.bin", map_location=device
173
+ )
174
+
175
+ self.vae.load_state_dict(vae_weights)
176
+ self.stft.load_state_dict(stft_weights)
177
+ self.model.load_state_dict(main_weights)
178
+
179
+ print("Successfully loaded checkpoint from:", name)
180
+
181
+ self.vae.eval()
182
+ self.stft.eval()
183
+ self.model.eval()
184
+
185
+ self.scheduler = DDPMScheduler.from_pretrained(
186
+ main_config["scheduler_name"], subfolder="scheduler"
187
+ )
188
+
189
+ def generate(self, prompt, steps=100, guidance=3, samples=1, disable_progress=True):
190
+ """Genrate music for a single prompt string."""
191
+
192
+ with torch.no_grad():
193
+ beats, chords, chords_times = self.music_model.generate(prompt)
194
+ latents = self.model.inference(
195
+ [prompt],
196
+ beats,
197
+ [chords],
198
+ [chords_times],
199
+ self.scheduler,
200
+ steps,
201
+ guidance,
202
+ samples,
203
+ disable_progress,
204
+ )
205
+ mel = self.vae.decode_first_stage(latents)
206
+ wave = self.vae.decode_to_waveform(mel)
207
+
208
+ return wave[0]
209
+
210
+
211
+ # Initialize Mustango
212
+ if torch.cuda.is_available():
213
+ mustango = Mustango()
214
+ else:
215
+ mustango = Mustango(device="cpu")
216
+
217
+ def gradio_generate(prompt, steps, guidance):
218
+ output_wave = mustango.generate(prompt, steps, guidance)
219
+ # output_filename = f"{prompt.replace(' ', '_')}_{steps}_{guidance}"[:250] + ".wav"
220
+ output_filename = "temp.wav"
221
+ wavio.write(output_filename, output_wave, rate=16000, sampwidth=2)
222
+
223
+ return output_filename
224
+
225
+ # description_text = """
226
+ # <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/>
227
+ # Generate music using Mustango by providing a text prompt.
228
+ # <br/><br/> Meet Mustango, an exciting addition to the vibrant landscape of Multimodal Large Language Models \
229
+ # designed for controlled music generation. Mustango leverages Latent Diffusion Model (LDM), Flan-T5, and \
230
+ # musical features to do the magic! \
231
+ # <p/>
232
+ # """
233
+ description_text = ""
234
+ # Gradio input and output components
235
+ input_text = gr.inputs.Textbox(lines=2, label="Prompt")
236
+ output_audio = gr.outputs.Audio(label="Generated Music", type="filepath")
237
+ denoising_steps = gr.Slider(minimum=100, maximum=200, value=100, step=1, label="Steps", interactive=True)
238
+ guidance_scale = gr.Slider(minimum=1, maximum=10, value=3, step=0.1, label="Guidance Scale", interactive=True)
239
+
240
+ # Gradio interface
241
+ gr_interface = gr.Interface(
242
+ fn=gradio_generate,
243
+ inputs=[input_text, denoising_steps, guidance_scale],
244
+ outputs=[output_audio],
245
+ title="Mustango: Toward Controllable Text-to-Music Generation",
246
+ description=description_text,
247
+ allow_flagging=False,
248
+ examples=[
249
+ ["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."],
250
+ ["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."],
251
+ ["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."],
252
+ ["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."],
253
+ ["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."],
254
+ ["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."],
255
+ ],
256
+ cache_examples=False,
257
+ )
258
+
259
+ # Launch Gradio app
260
+ gr_interface.launch()
.ipynb_checkpoints/modelling_deberta_v2-checkpoint.py ADDED
@@ -0,0 +1,1750 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2020 Microsoft and the Hugging Face Inc. team.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """ PyTorch DeBERTa-v2 model."""
16
+
17
+ from collections.abc import Sequence
18
+ from typing import Optional, Tuple, Union
19
+
20
+ import torch
21
+ import torch.utils.checkpoint
22
+ from torch import nn
23
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss
24
+
25
+ from transformers.activations import ACT2FN
26
+ from transformers.modeling_outputs import (
27
+ ModelOutput,
28
+ BaseModelOutput,
29
+ MaskedLMOutput,
30
+ MultipleChoiceModelOutput,
31
+ QuestionAnsweringModelOutput,
32
+ SequenceClassifierOutput,
33
+ TokenClassifierOutput,
34
+ )
35
+ from transformers.modeling_utils import PreTrainedModel
36
+ from transformers.pytorch_utils import softmax_backward_data
37
+ from transformers.utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
38
+ from transformers.models.deberta_v2.configuration_deberta_v2 import DebertaV2Config
39
+
40
+
41
+ logger = logging.get_logger(__name__)
42
+
43
+ _CONFIG_FOR_DOC = "DebertaV2Config"
44
+ _CHECKPOINT_FOR_DOC = "microsoft/deberta-v2-xlarge"
45
+ _QA_TARGET_START_INDEX = 2
46
+ _QA_TARGET_END_INDEX = 9
47
+
48
+ DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST = [
49
+ "microsoft/deberta-v2-xlarge",
50
+ "microsoft/deberta-v2-xxlarge",
51
+ "microsoft/deberta-v2-xlarge-mnli",
52
+ "microsoft/deberta-v2-xxlarge-mnli",
53
+ ]
54
+
55
+
56
+ # Copied from transformers.models.deberta.modeling_deberta.ContextPooler
57
+ class ContextPooler(nn.Module):
58
+ def __init__(self, config):
59
+ super().__init__()
60
+ self.dense = nn.Linear(config.pooler_hidden_size, config.pooler_hidden_size)
61
+ self.dropout = StableDropout(config.pooler_dropout)
62
+ self.config = config
63
+
64
+ def forward(self, hidden_states):
65
+ # We "pool" the model by simply taking the hidden state corresponding
66
+ # to the first token.
67
+
68
+ context_token = hidden_states[:, 0]
69
+ context_token = self.dropout(context_token)
70
+ pooled_output = self.dense(context_token)
71
+ pooled_output = ACT2FN[self.config.pooler_hidden_act](pooled_output)
72
+ return pooled_output
73
+
74
+ @property
75
+ def output_dim(self):
76
+ return self.config.hidden_size
77
+
78
+
79
+ # Copied from transformers.models.deberta.modeling_deberta.XSoftmax with deberta->deberta_v2
80
+ class XSoftmax(torch.autograd.Function):
81
+ """
82
+ Masked Softmax which is optimized for saving memory
83
+
84
+ Args:
85
+ input (`torch.tensor`): The input tensor that will apply softmax.
86
+ mask (`torch.IntTensor`):
87
+ The mask matrix where 0 indicate that element will be ignored in the softmax calculation.
88
+ dim (int): The dimension that will apply softmax
89
+
90
+ Example:
91
+
92
+ ```python
93
+ >>> import torch
94
+ >>> from transformers.models.deberta_v2.modeling_deberta_v2 import XSoftmax
95
+
96
+ >>> # Make a tensor
97
+ >>> x = torch.randn([4, 20, 100])
98
+
99
+ >>> # Create a mask
100
+ >>> mask = (x > 0).int()
101
+
102
+ >>> # Specify the dimension to apply softmax
103
+ >>> dim = -1
104
+
105
+ >>> y = XSoftmax.apply(x, mask, dim)
106
+ ```"""
107
+
108
+ @staticmethod
109
+ def forward(self, input, mask, dim):
110
+ self.dim = dim
111
+ rmask = ~(mask.to(torch.bool))
112
+
113
+ output = input.masked_fill(rmask, torch.tensor(torch.finfo(input.dtype).min))
114
+ output = torch.softmax(output, self.dim)
115
+ output.masked_fill_(rmask, 0)
116
+ self.save_for_backward(output)
117
+ return output
118
+
119
+ @staticmethod
120
+ def backward(self, grad_output):
121
+ (output,) = self.saved_tensors
122
+ inputGrad = softmax_backward_data(self, grad_output, output, self.dim, output)
123
+ return inputGrad, None, None
124
+
125
+ @staticmethod
126
+ def symbolic(g, self, mask, dim):
127
+ import torch.onnx.symbolic_helper as sym_help
128
+ from torch.onnx.symbolic_opset9 import masked_fill, softmax
129
+
130
+ mask_cast_value = g.op("Cast", mask, to_i=sym_help.cast_pytorch_to_onnx["Long"])
131
+ r_mask = g.op(
132
+ "Cast",
133
+ g.op("Sub", g.op("Constant", value_t=torch.tensor(1, dtype=torch.int64)), mask_cast_value),
134
+ to_i=sym_help.cast_pytorch_to_onnx["Bool"],
135
+ )
136
+ output = masked_fill(
137
+ g, self, r_mask, g.op("Constant", value_t=torch.tensor(torch.finfo(self.type().dtype()).min))
138
+ )
139
+ output = softmax(g, output, dim)
140
+ return masked_fill(g, output, r_mask, g.op("Constant", value_t=torch.tensor(0, dtype=torch.bool)))
141
+
142
+
143
+ # Copied from transformers.models.deberta.modeling_deberta.DropoutContext
144
+ class DropoutContext(object):
145
+ def __init__(self):
146
+ self.dropout = 0
147
+ self.mask = None
148
+ self.scale = 1
149
+ self.reuse_mask = True
150
+
151
+
152
+ # Copied from transformers.models.deberta.modeling_deberta.get_mask
153
+ def get_mask(input, local_context):
154
+ if not isinstance(local_context, DropoutContext):
155
+ dropout = local_context
156
+ mask = None
157
+ else:
158
+ dropout = local_context.dropout
159
+ dropout *= local_context.scale
160
+ mask = local_context.mask if local_context.reuse_mask else None
161
+
162
+ if dropout > 0 and mask is None:
163
+ mask = (1 - torch.empty_like(input).bernoulli_(1 - dropout)).to(torch.bool)
164
+
165
+ if isinstance(local_context, DropoutContext):
166
+ if local_context.mask is None:
167
+ local_context.mask = mask
168
+
169
+ return mask, dropout
170
+
171
+
172
+ # Copied from transformers.models.deberta.modeling_deberta.XDropout
173
+ class XDropout(torch.autograd.Function):
174
+ """Optimized dropout function to save computation and memory by using mask operation instead of multiplication."""
175
+
176
+ @staticmethod
177
+ def forward(ctx, input, local_ctx):
178
+ mask, dropout = get_mask(input, local_ctx)
179
+ ctx.scale = 1.0 / (1 - dropout)
180
+ if dropout > 0:
181
+ ctx.save_for_backward(mask)
182
+ return input.masked_fill(mask, 0) * ctx.scale
183
+ else:
184
+ return input
185
+
186
+ @staticmethod
187
+ def backward(ctx, grad_output):
188
+ if ctx.scale > 1:
189
+ (mask,) = ctx.saved_tensors
190
+ return grad_output.masked_fill(mask, 0) * ctx.scale, None
191
+ else:
192
+ return grad_output, None
193
+
194
+ @staticmethod
195
+ def symbolic(g: torch._C.Graph, input: torch._C.Value, local_ctx: Union[float, DropoutContext]) -> torch._C.Value:
196
+ from torch.onnx import symbolic_opset12
197
+
198
+ dropout_p = local_ctx
199
+ if isinstance(local_ctx, DropoutContext):
200
+ dropout_p = local_ctx.dropout
201
+ # StableDropout only calls this function when training.
202
+ train = True
203
+ # TODO: We should check if the opset_version being used to export
204
+ # is > 12 here, but there's no good way to do that. As-is, if the
205
+ # opset_version < 12, export will fail with a CheckerError.
206
+ # Once https://github.com/pytorch/pytorch/issues/78391 is fixed, do something like:
207
+ # if opset_version < 12:
208
+ # return torch.onnx.symbolic_opset9.dropout(g, input, dropout_p, train)
209
+ return symbolic_opset12.dropout(g, input, dropout_p, train)
210
+
211
+
212
+ # Copied from transformers.models.deberta.modeling_deberta.StableDropout
213
+ class StableDropout(nn.Module):
214
+ """
215
+ Optimized dropout module for stabilizing the training
216
+
217
+ Args:
218
+ drop_prob (float): the dropout probabilities
219
+ """
220
+
221
+ def __init__(self, drop_prob):
222
+ super().__init__()
223
+ self.drop_prob = drop_prob
224
+ self.count = 0
225
+ self.context_stack = None
226
+
227
+ def forward(self, x):
228
+ """
229
+ Call the module
230
+
231
+ Args:
232
+ x (`torch.tensor`): The input tensor to apply dropout
233
+ """
234
+ if self.training and self.drop_prob > 0:
235
+ return XDropout.apply(x, self.get_context())
236
+ return x
237
+
238
+ def clear_context(self):
239
+ self.count = 0
240
+ self.context_stack = None
241
+
242
+ def init_context(self, reuse_mask=True, scale=1):
243
+ if self.context_stack is None:
244
+ self.context_stack = []
245
+ self.count = 0
246
+ for c in self.context_stack:
247
+ c.reuse_mask = reuse_mask
248
+ c.scale = scale
249
+
250
+ def get_context(self):
251
+ if self.context_stack is not None:
252
+ if self.count >= len(self.context_stack):
253
+ self.context_stack.append(DropoutContext())
254
+ ctx = self.context_stack[self.count]
255
+ ctx.dropout = self.drop_prob
256
+ self.count += 1
257
+ return ctx
258
+ else:
259
+ return self.drop_prob
260
+
261
+
262
+ # Copied from transformers.models.deberta.modeling_deberta.DebertaSelfOutput with DebertaLayerNorm->LayerNorm
263
+ class DebertaV2SelfOutput(nn.Module):
264
+ def __init__(self, config):
265
+ super().__init__()
266
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
267
+ self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
268
+ self.dropout = StableDropout(config.hidden_dropout_prob)
269
+
270
+ def forward(self, hidden_states, input_tensor):
271
+ hidden_states = self.dense(hidden_states)
272
+ hidden_states = self.dropout(hidden_states)
273
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
274
+ return hidden_states
275
+
276
+
277
+ # Copied from transformers.models.deberta.modeling_deberta.DebertaAttention with Deberta->DebertaV2
278
+ class DebertaV2Attention(nn.Module):
279
+ def __init__(self, config):
280
+ super().__init__()
281
+ self.self = DisentangledSelfAttention(config)
282
+ self.output = DebertaV2SelfOutput(config)
283
+ self.config = config
284
+
285
+ def forward(
286
+ self,
287
+ hidden_states,
288
+ attention_mask,
289
+ output_attentions=False,
290
+ query_states=None,
291
+ relative_pos=None,
292
+ rel_embeddings=None,
293
+ ):
294
+ self_output = self.self(
295
+ hidden_states,
296
+ attention_mask,
297
+ output_attentions,
298
+ query_states=query_states,
299
+ relative_pos=relative_pos,
300
+ rel_embeddings=rel_embeddings,
301
+ )
302
+ if output_attentions:
303
+ self_output, att_matrix = self_output
304
+ if query_states is None:
305
+ query_states = hidden_states
306
+ attention_output = self.output(self_output, query_states)
307
+
308
+ if output_attentions:
309
+ return (attention_output, att_matrix)
310
+ else:
311
+ return attention_output
312
+
313
+
314
+ # Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->DebertaV2
315
+ class DebertaV2Intermediate(nn.Module):
316
+ def __init__(self, config):
317
+ super().__init__()
318
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
319
+ if isinstance(config.hidden_act, str):
320
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
321
+ else:
322
+ self.intermediate_act_fn = config.hidden_act
323
+
324
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
325
+ hidden_states = self.dense(hidden_states)
326
+ hidden_states = self.intermediate_act_fn(hidden_states)
327
+ return hidden_states
328
+
329
+
330
+ # Copied from transformers.models.deberta.modeling_deberta.DebertaOutput with DebertaLayerNorm->LayerNorm
331
+ class DebertaV2Output(nn.Module):
332
+ def __init__(self, config):
333
+ super().__init__()
334
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
335
+ self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
336
+ self.dropout = StableDropout(config.hidden_dropout_prob)
337
+ self.config = config
338
+
339
+ def forward(self, hidden_states, input_tensor):
340
+ hidden_states = self.dense(hidden_states)
341
+ hidden_states = self.dropout(hidden_states)
342
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
343
+ return hidden_states
344
+
345
+
346
+ # Copied from transformers.models.deberta.modeling_deberta.DebertaLayer with Deberta->DebertaV2
347
+ class DebertaV2Layer(nn.Module):
348
+ def __init__(self, config):
349
+ super().__init__()
350
+ self.attention = DebertaV2Attention(config)
351
+ self.intermediate = DebertaV2Intermediate(config)
352
+ self.output = DebertaV2Output(config)
353
+
354
+ def forward(
355
+ self,
356
+ hidden_states,
357
+ attention_mask,
358
+ query_states=None,
359
+ relative_pos=None,
360
+ rel_embeddings=None,
361
+ output_attentions=False,
362
+ ):
363
+ attention_output = self.attention(
364
+ hidden_states,
365
+ attention_mask,
366
+ output_attentions=output_attentions,
367
+ query_states=query_states,
368
+ relative_pos=relative_pos,
369
+ rel_embeddings=rel_embeddings,
370
+ )
371
+ if output_attentions:
372
+ attention_output, att_matrix = attention_output
373
+ intermediate_output = self.intermediate(attention_output)
374
+ layer_output = self.output(intermediate_output, attention_output)
375
+ if output_attentions:
376
+ return (layer_output, att_matrix)
377
+ else:
378
+ return layer_output
379
+
380
+
381
+ class ConvLayer(nn.Module):
382
+ def __init__(self, config):
383
+ super().__init__()
384
+ kernel_size = getattr(config, "conv_kernel_size", 3)
385
+ groups = getattr(config, "conv_groups", 1)
386
+ self.conv_act = getattr(config, "conv_act", "tanh")
387
+ self.conv = nn.Conv1d(
388
+ config.hidden_size, config.hidden_size, kernel_size, padding=(kernel_size - 1) // 2, groups=groups
389
+ )
390
+ self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
391
+ self.dropout = StableDropout(config.hidden_dropout_prob)
392
+ self.config = config
393
+
394
+ def forward(self, hidden_states, residual_states, input_mask):
395
+ out = self.conv(hidden_states.permute(0, 2, 1).contiguous()).permute(0, 2, 1).contiguous()
396
+ rmask = (1 - input_mask).bool()
397
+ out.masked_fill_(rmask.unsqueeze(-1).expand(out.size()), 0)
398
+ out = ACT2FN[self.conv_act](self.dropout(out))
399
+
400
+ layer_norm_input = residual_states + out
401
+ output = self.LayerNorm(layer_norm_input).to(layer_norm_input)
402
+
403
+ if input_mask is None:
404
+ output_states = output
405
+ else:
406
+ if input_mask.dim() != layer_norm_input.dim():
407
+ if input_mask.dim() == 4:
408
+ input_mask = input_mask.squeeze(1).squeeze(1)
409
+ input_mask = input_mask.unsqueeze(2)
410
+
411
+ input_mask = input_mask.to(output.dtype)
412
+ output_states = output * input_mask
413
+
414
+ return output_states
415
+
416
+
417
+ class DebertaV2Encoder(nn.Module):
418
+ """Modified BertEncoder with relative position bias support"""
419
+
420
+ def __init__(self, config):
421
+ super().__init__()
422
+
423
+ self.layer = nn.ModuleList([DebertaV2Layer(config) for _ in range(config.num_hidden_layers)])
424
+ self.relative_attention = getattr(config, "relative_attention", False)
425
+
426
+ if self.relative_attention:
427
+ self.max_relative_positions = getattr(config, "max_relative_positions", -1)
428
+ if self.max_relative_positions < 1:
429
+ self.max_relative_positions = config.max_position_embeddings
430
+
431
+ self.position_buckets = getattr(config, "position_buckets", -1)
432
+ pos_ebd_size = self.max_relative_positions * 2
433
+
434
+ if self.position_buckets > 0:
435
+ pos_ebd_size = self.position_buckets * 2
436
+
437
+ self.rel_embeddings = nn.Embedding(pos_ebd_size, config.hidden_size)
438
+
439
+ self.norm_rel_ebd = [x.strip() for x in getattr(config, "norm_rel_ebd", "none").lower().split("|")]
440
+
441
+ if "layer_norm" in self.norm_rel_ebd:
442
+ self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=True)
443
+
444
+ self.conv = ConvLayer(config) if getattr(config, "conv_kernel_size", 0) > 0 else None
445
+ self.gradient_checkpointing = False
446
+
447
+ def get_rel_embedding(self):
448
+ rel_embeddings = self.rel_embeddings.weight if self.relative_attention else None
449
+ if rel_embeddings is not None and ("layer_norm" in self.norm_rel_ebd):
450
+ rel_embeddings = self.LayerNorm(rel_embeddings)
451
+ return rel_embeddings
452
+
453
+ def get_attention_mask(self, attention_mask):
454
+ if attention_mask.dim() <= 2:
455
+ extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
456
+ attention_mask = extended_attention_mask * extended_attention_mask.squeeze(-2).unsqueeze(-1)
457
+ elif attention_mask.dim() == 3:
458
+ attention_mask = attention_mask.unsqueeze(1)
459
+
460
+ return attention_mask
461
+
462
+ def get_rel_pos(self, hidden_states, query_states=None, relative_pos=None):
463
+ if self.relative_attention and relative_pos is None:
464
+ q = query_states.size(-2) if query_states is not None else hidden_states.size(-2)
465
+ relative_pos = build_relative_position(
466
+ q,
467
+ hidden_states.size(-2),
468
+ bucket_size=self.position_buckets,
469
+ max_position=self.max_relative_positions,
470
+ device=hidden_states.device,
471
+ )
472
+ return relative_pos
473
+
474
+ def forward(
475
+ self,
476
+ hidden_states,
477
+ attention_mask,
478
+ output_hidden_states=True,
479
+ output_attentions=False,
480
+ query_states=None,
481
+ relative_pos=None,
482
+ return_dict=True,
483
+ ):
484
+ if attention_mask.dim() <= 2:
485
+ input_mask = attention_mask
486
+ else:
487
+ input_mask = attention_mask.sum(-2) > 0
488
+ attention_mask = self.get_attention_mask(attention_mask)
489
+ relative_pos = self.get_rel_pos(hidden_states, query_states, relative_pos)
490
+
491
+ all_hidden_states = () if output_hidden_states else None
492
+ all_attentions = () if output_attentions else None
493
+
494
+ if isinstance(hidden_states, Sequence):
495
+ next_kv = hidden_states[0]
496
+ else:
497
+ next_kv = hidden_states
498
+ rel_embeddings = self.get_rel_embedding()
499
+ output_states = next_kv
500
+ for i, layer_module in enumerate(self.layer):
501
+ if output_hidden_states:
502
+ all_hidden_states = all_hidden_states + (output_states,)
503
+
504
+ if self.gradient_checkpointing and self.training:
505
+
506
+ def create_custom_forward(module):
507
+ def custom_forward(*inputs):
508
+ return module(*inputs, output_attentions)
509
+
510
+ return custom_forward
511
+
512
+ output_states = torch.utils.checkpoint.checkpoint(
513
+ create_custom_forward(layer_module),
514
+ next_kv,
515
+ attention_mask,
516
+ query_states,
517
+ relative_pos,
518
+ rel_embeddings,
519
+ )
520
+ else:
521
+ output_states = layer_module(
522
+ next_kv,
523
+ attention_mask,
524
+ query_states=query_states,
525
+ relative_pos=relative_pos,
526
+ rel_embeddings=rel_embeddings,
527
+ output_attentions=output_attentions,
528
+ )
529
+
530
+ if output_attentions:
531
+ output_states, att_m = output_states
532
+
533
+ if i == 0 and self.conv is not None:
534
+ output_states = self.conv(hidden_states, output_states, input_mask)
535
+
536
+ if query_states is not None:
537
+ query_states = output_states
538
+ if isinstance(hidden_states, Sequence):
539
+ next_kv = hidden_states[i + 1] if i + 1 < len(self.layer) else None
540
+ else:
541
+ next_kv = output_states
542
+
543
+ if output_attentions:
544
+ all_attentions = all_attentions + (att_m,)
545
+
546
+ if output_hidden_states:
547
+ all_hidden_states = all_hidden_states + (output_states,)
548
+
549
+ if not return_dict:
550
+ return tuple(v for v in [output_states, all_hidden_states, all_attentions] if v is not None)
551
+ return BaseModelOutput(
552
+ last_hidden_state=output_states, hidden_states=all_hidden_states, attentions=all_attentions
553
+ )
554
+
555
+
556
+ def make_log_bucket_position(relative_pos, bucket_size, max_position):
557
+ sign = torch.sign(relative_pos)
558
+ mid = bucket_size // 2
559
+ abs_pos = torch.where(
560
+ (relative_pos < mid) & (relative_pos > -mid),
561
+ torch.tensor(mid - 1).type_as(relative_pos),
562
+ torch.abs(relative_pos),
563
+ )
564
+ log_pos = (
565
+ torch.ceil(torch.log(abs_pos / mid) / torch.log(torch.tensor((max_position - 1) / mid)) * (mid - 1)) + mid
566
+ )
567
+ bucket_pos = torch.where(abs_pos <= mid, relative_pos.type_as(log_pos), log_pos * sign)
568
+ return bucket_pos
569
+
570
+
571
+ def build_relative_position(query_size, key_size, bucket_size=-1, max_position=-1, device=None):
572
+ """
573
+ Build relative position according to the query and key
574
+
575
+ We assume the absolute position of query \\(P_q\\) is range from (0, query_size) and the absolute position of key
576
+ \\(P_k\\) is range from (0, key_size), The relative positions from query to key is \\(R_{q \\rightarrow k} = P_q -
577
+ P_k\\)
578
+
579
+ Args:
580
+ query_size (int): the length of query
581
+ key_size (int): the length of key
582
+ bucket_size (int): the size of position bucket
583
+ max_position (int): the maximum allowed absolute position
584
+ device (`torch.device`): the device on which tensors will be created.
585
+
586
+ Return:
587
+ `torch.LongTensor`: A tensor with shape [1, query_size, key_size]
588
+ """
589
+
590
+ q_ids = torch.arange(0, query_size, device=device)
591
+ k_ids = torch.arange(0, key_size, device=device)
592
+ rel_pos_ids = q_ids[:, None] - k_ids[None, :]
593
+ if bucket_size > 0 and max_position > 0:
594
+ rel_pos_ids = make_log_bucket_position(rel_pos_ids, bucket_size, max_position)
595
+ rel_pos_ids = rel_pos_ids.to(torch.long)
596
+ rel_pos_ids = rel_pos_ids[:query_size, :]
597
+ rel_pos_ids = rel_pos_ids.unsqueeze(0)
598
+ return rel_pos_ids
599
+
600
+
601
+ @torch.jit.script
602
+ # Copied from transformers.models.deberta.modeling_deberta.c2p_dynamic_expand
603
+ def c2p_dynamic_expand(c2p_pos, query_layer, relative_pos):
604
+ return c2p_pos.expand([query_layer.size(0), query_layer.size(1), query_layer.size(2), relative_pos.size(-1)])
605
+
606
+
607
+ @torch.jit.script
608
+ # Copied from transformers.models.deberta.modeling_deberta.p2c_dynamic_expand
609
+ def p2c_dynamic_expand(c2p_pos, query_layer, key_layer):
610
+ return c2p_pos.expand([query_layer.size(0), query_layer.size(1), key_layer.size(-2), key_layer.size(-2)])
611
+
612
+
613
+ @torch.jit.script
614
+ # Copied from transformers.models.deberta.modeling_deberta.pos_dynamic_expand
615
+ def pos_dynamic_expand(pos_index, p2c_att, key_layer):
616
+ return pos_index.expand(p2c_att.size()[:2] + (pos_index.size(-2), key_layer.size(-2)))
617
+
618
+
619
+ class DisentangledSelfAttention(nn.Module):
620
+ """
621
+ Disentangled self-attention module
622
+
623
+ Parameters:
624
+ config (`DebertaV2Config`):
625
+ A model config class instance with the configuration to build a new model. The schema is similar to
626
+ *BertConfig*, for more details, please refer [`DebertaV2Config`]
627
+
628
+ """
629
+
630
+ def __init__(self, config):
631
+ super().__init__()
632
+ if config.hidden_size % config.num_attention_heads != 0:
633
+ raise ValueError(
634
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
635
+ f"heads ({config.num_attention_heads})"
636
+ )
637
+ self.num_attention_heads = config.num_attention_heads
638
+ _attention_head_size = config.hidden_size // config.num_attention_heads
639
+ self.attention_head_size = getattr(config, "attention_head_size", _attention_head_size)
640
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
641
+ self.query_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
642
+ self.key_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
643
+ self.value_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
644
+
645
+ self.share_att_key = getattr(config, "share_att_key", False)
646
+ self.pos_att_type = config.pos_att_type if config.pos_att_type is not None else []
647
+ self.relative_attention = getattr(config, "relative_attention", False)
648
+
649
+ if self.relative_attention:
650
+ self.position_buckets = getattr(config, "position_buckets", -1)
651
+ self.max_relative_positions = getattr(config, "max_relative_positions", -1)
652
+ if self.max_relative_positions < 1:
653
+ self.max_relative_positions = config.max_position_embeddings
654
+ self.pos_ebd_size = self.max_relative_positions
655
+ if self.position_buckets > 0:
656
+ self.pos_ebd_size = self.position_buckets
657
+
658
+ self.pos_dropout = StableDropout(config.hidden_dropout_prob)
659
+
660
+ if not self.share_att_key:
661
+ if "c2p" in self.pos_att_type:
662
+ self.pos_key_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
663
+ if "p2c" in self.pos_att_type:
664
+ self.pos_query_proj = nn.Linear(config.hidden_size, self.all_head_size)
665
+
666
+ self.dropout = StableDropout(config.attention_probs_dropout_prob)
667
+
668
+ def transpose_for_scores(self, x, attention_heads):
669
+ new_x_shape = x.size()[:-1] + (attention_heads, -1)
670
+ x = x.view(new_x_shape)
671
+ return x.permute(0, 2, 1, 3).contiguous().view(-1, x.size(1), x.size(-1))
672
+
673
+ def forward(
674
+ self,
675
+ hidden_states,
676
+ attention_mask,
677
+ output_attentions=False,
678
+ query_states=None,
679
+ relative_pos=None,
680
+ rel_embeddings=None,
681
+ ):
682
+ """
683
+ Call the module
684
+
685
+ Args:
686
+ hidden_states (`torch.FloatTensor`):
687
+ Input states to the module usually the output from previous layer, it will be the Q,K and V in
688
+ *Attention(Q,K,V)*
689
+
690
+ attention_mask (`torch.BoolTensor`):
691
+ An attention mask matrix of shape [*B*, *N*, *N*] where *B* is the batch size, *N* is the maximum
692
+ sequence length in which element [i,j] = *1* means the *i* th token in the input can attend to the *j*
693
+ th token.
694
+
695
+ output_attentions (`bool`, optional):
696
+ Whether return the attention matrix.
697
+
698
+ query_states (`torch.FloatTensor`, optional):
699
+ The *Q* state in *Attention(Q,K,V)*.
700
+
701
+ relative_pos (`torch.LongTensor`):
702
+ The relative position encoding between the tokens in the sequence. It's of shape [*B*, *N*, *N*] with
703
+ values ranging in [*-max_relative_positions*, *max_relative_positions*].
704
+
705
+ rel_embeddings (`torch.FloatTensor`):
706
+ The embedding of relative distances. It's a tensor of shape [\\(2 \\times
707
+ \\text{max_relative_positions}\\), *hidden_size*].
708
+
709
+
710
+ """
711
+ if query_states is None:
712
+ query_states = hidden_states
713
+ query_layer = self.transpose_for_scores(self.query_proj(query_states), self.num_attention_heads)
714
+ key_layer = self.transpose_for_scores(self.key_proj(hidden_states), self.num_attention_heads)
715
+ value_layer = self.transpose_for_scores(self.value_proj(hidden_states), self.num_attention_heads)
716
+
717
+ rel_att = None
718
+ # Take the dot product between "query" and "key" to get the raw attention scores.
719
+ scale_factor = 1
720
+ if "c2p" in self.pos_att_type:
721
+ scale_factor += 1
722
+ if "p2c" in self.pos_att_type:
723
+ scale_factor += 1
724
+ scale = torch.sqrt(torch.tensor(query_layer.size(-1), dtype=torch.float) * scale_factor)
725
+ attention_scores = torch.bmm(query_layer, key_layer.transpose(-1, -2)) / scale.to(dtype=query_layer.dtype)
726
+ if self.relative_attention:
727
+ rel_embeddings = self.pos_dropout(rel_embeddings)
728
+ rel_att = self.disentangled_attention_bias(
729
+ query_layer, key_layer, relative_pos, rel_embeddings, scale_factor
730
+ )
731
+
732
+ if rel_att is not None:
733
+ attention_scores = attention_scores + rel_att
734
+ attention_scores = attention_scores
735
+ attention_scores = attention_scores.view(
736
+ -1, self.num_attention_heads, attention_scores.size(-2), attention_scores.size(-1)
737
+ )
738
+
739
+ # bsz x height x length x dimension
740
+ attention_probs = XSoftmax.apply(attention_scores, attention_mask, -1)
741
+ attention_probs = self.dropout(attention_probs)
742
+ context_layer = torch.bmm(
743
+ attention_probs.view(-1, attention_probs.size(-2), attention_probs.size(-1)), value_layer
744
+ )
745
+ context_layer = (
746
+ context_layer.view(-1, self.num_attention_heads, context_layer.size(-2), context_layer.size(-1))
747
+ .permute(0, 2, 1, 3)
748
+ .contiguous()
749
+ )
750
+ new_context_layer_shape = context_layer.size()[:-2] + (-1,)
751
+ context_layer = context_layer.view(new_context_layer_shape)
752
+ if output_attentions:
753
+ return (context_layer, attention_probs)
754
+ else:
755
+ return context_layer
756
+
757
+ def disentangled_attention_bias(self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor):
758
+ if relative_pos is None:
759
+ q = query_layer.size(-2)
760
+ relative_pos = build_relative_position(
761
+ q,
762
+ key_layer.size(-2),
763
+ bucket_size=self.position_buckets,
764
+ max_position=self.max_relative_positions,
765
+ device=query_layer.device,
766
+ )
767
+ if relative_pos.dim() == 2:
768
+ relative_pos = relative_pos.unsqueeze(0).unsqueeze(0)
769
+ elif relative_pos.dim() == 3:
770
+ relative_pos = relative_pos.unsqueeze(1)
771
+ # bsz x height x query x key
772
+ elif relative_pos.dim() != 4:
773
+ raise ValueError(f"Relative position ids must be of dim 2 or 3 or 4. {relative_pos.dim()}")
774
+
775
+ att_span = self.pos_ebd_size
776
+ relative_pos = relative_pos.long().to(query_layer.device)
777
+
778
+ rel_embeddings = rel_embeddings[0 : att_span * 2, :].unsqueeze(0)
779
+ if self.share_att_key:
780
+ pos_query_layer = self.transpose_for_scores(
781
+ self.query_proj(rel_embeddings), self.num_attention_heads
782
+ ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1)
783
+ pos_key_layer = self.transpose_for_scores(self.key_proj(rel_embeddings), self.num_attention_heads).repeat(
784
+ query_layer.size(0) // self.num_attention_heads, 1, 1
785
+ )
786
+ else:
787
+ if "c2p" in self.pos_att_type:
788
+ pos_key_layer = self.transpose_for_scores(
789
+ self.pos_key_proj(rel_embeddings), self.num_attention_heads
790
+ ).repeat(
791
+ query_layer.size(0) // self.num_attention_heads, 1, 1
792
+ ) # .split(self.all_head_size, dim=-1)
793
+ if "p2c" in self.pos_att_type:
794
+ pos_query_layer = self.transpose_for_scores(
795
+ self.pos_query_proj(rel_embeddings), self.num_attention_heads
796
+ ).repeat(
797
+ query_layer.size(0) // self.num_attention_heads, 1, 1
798
+ ) # .split(self.all_head_size, dim=-1)
799
+
800
+ score = 0
801
+ # content->position
802
+ if "c2p" in self.pos_att_type:
803
+ scale = torch.sqrt(torch.tensor(pos_key_layer.size(-1), dtype=torch.float) * scale_factor)
804
+ c2p_att = torch.bmm(query_layer, pos_key_layer.transpose(-1, -2))
805
+ c2p_pos = torch.clamp(relative_pos + att_span, 0, att_span * 2 - 1)
806
+ c2p_att = torch.gather(
807
+ c2p_att,
808
+ dim=-1,
809
+ index=c2p_pos.squeeze(0).expand([query_layer.size(0), query_layer.size(1), relative_pos.size(-1)]),
810
+ )
811
+ score += c2p_att / scale.to(dtype=c2p_att.dtype)
812
+
813
+ # position->content
814
+ if "p2c" in self.pos_att_type:
815
+ scale = torch.sqrt(torch.tensor(pos_query_layer.size(-1), dtype=torch.float) * scale_factor)
816
+ if key_layer.size(-2) != query_layer.size(-2):
817
+ r_pos = build_relative_position(
818
+ key_layer.size(-2),
819
+ key_layer.size(-2),
820
+ bucket_size=self.position_buckets,
821
+ max_position=self.max_relative_positions,
822
+ device=query_layer.device,
823
+ )
824
+ r_pos = r_pos.unsqueeze(0)
825
+ else:
826
+ r_pos = relative_pos
827
+
828
+ p2c_pos = torch.clamp(-r_pos + att_span, 0, att_span * 2 - 1)
829
+ p2c_att = torch.bmm(key_layer, pos_query_layer.transpose(-1, -2))
830
+ p2c_att = torch.gather(
831
+ p2c_att,
832
+ dim=-1,
833
+ index=p2c_pos.squeeze(0).expand([query_layer.size(0), key_layer.size(-2), key_layer.size(-2)]),
834
+ ).transpose(-1, -2)
835
+ score += p2c_att / scale.to(dtype=p2c_att.dtype)
836
+
837
+ return score
838
+
839
+
840
+ # Copied from transformers.models.deberta.modeling_deberta.DebertaEmbeddings with DebertaLayerNorm->LayerNorm
841
+ class DebertaV2Embeddings(nn.Module):
842
+ """Construct the embeddings from word, position and token_type embeddings."""
843
+
844
+ def __init__(self, config):
845
+ super().__init__()
846
+ pad_token_id = getattr(config, "pad_token_id", 0)
847
+ self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
848
+ self.word_embeddings = nn.Embedding(config.vocab_size, self.embedding_size, padding_idx=pad_token_id)
849
+
850
+ self.position_biased_input = getattr(config, "position_biased_input", True)
851
+ if not self.position_biased_input:
852
+ self.position_embeddings = None
853
+ else:
854
+ self.position_embeddings = nn.Embedding(config.max_position_embeddings, self.embedding_size)
855
+
856
+ if config.type_vocab_size > 0:
857
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, self.embedding_size)
858
+
859
+ if self.embedding_size != config.hidden_size:
860
+ self.embed_proj = nn.Linear(self.embedding_size, config.hidden_size, bias=False)
861
+ self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
862
+ self.dropout = StableDropout(config.hidden_dropout_prob)
863
+ self.config = config
864
+
865
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
866
+ self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
867
+
868
+ def forward(self, input_ids=None, token_type_ids=None, position_ids=None, mask=None, inputs_embeds=None):
869
+ if input_ids is not None:
870
+ input_shape = input_ids.size()
871
+ else:
872
+ input_shape = inputs_embeds.size()[:-1]
873
+
874
+ seq_length = input_shape[1]
875
+
876
+ if position_ids is None:
877
+ position_ids = self.position_ids[:, :seq_length]
878
+
879
+ if token_type_ids is None:
880
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
881
+
882
+ if inputs_embeds is None:
883
+ inputs_embeds = self.word_embeddings(input_ids)
884
+
885
+ if self.position_embeddings is not None:
886
+ position_embeddings = self.position_embeddings(position_ids.long())
887
+ else:
888
+ position_embeddings = torch.zeros_like(inputs_embeds)
889
+
890
+ embeddings = inputs_embeds
891
+ if self.position_biased_input:
892
+ embeddings += position_embeddings
893
+ if self.config.type_vocab_size > 0:
894
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
895
+ embeddings += token_type_embeddings
896
+
897
+ if self.embedding_size != self.config.hidden_size:
898
+ embeddings = self.embed_proj(embeddings)
899
+
900
+ embeddings = self.LayerNorm(embeddings)
901
+
902
+ if mask is not None:
903
+ if mask.dim() != embeddings.dim():
904
+ if mask.dim() == 4:
905
+ mask = mask.squeeze(1).squeeze(1)
906
+ mask = mask.unsqueeze(2)
907
+ mask = mask.to(embeddings.dtype)
908
+
909
+ embeddings = embeddings * mask
910
+
911
+ embeddings = self.dropout(embeddings)
912
+ return embeddings
913
+
914
+
915
+ # Copied from transformers.models.deberta.modeling_deberta.DebertaPreTrainedModel with Deberta->DebertaV2
916
+ class DebertaV2PreTrainedModel(PreTrainedModel):
917
+ """
918
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
919
+ models.
920
+ """
921
+
922
+ config_class = DebertaV2Config
923
+ base_model_prefix = "deberta"
924
+ _keys_to_ignore_on_load_missing = ["position_ids"]
925
+ _keys_to_ignore_on_load_unexpected = ["position_embeddings"]
926
+ supports_gradient_checkpointing = True
927
+
928
+ def _init_weights(self, module):
929
+ """Initialize the weights."""
930
+ if isinstance(module, nn.Linear):
931
+ # Slightly different from the TF version which uses truncated_normal for initialization
932
+ # cf https://github.com/pytorch/pytorch/pull/5617
933
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
934
+ if module.bias is not None:
935
+ module.bias.data.zero_()
936
+ elif isinstance(module, nn.Embedding):
937
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
938
+ if module.padding_idx is not None:
939
+ module.weight.data[module.padding_idx].zero_()
940
+
941
+ def _set_gradient_checkpointing(self, module, value=False):
942
+ if isinstance(module, DebertaV2Encoder):
943
+ module.gradient_checkpointing = value
944
+
945
+
946
+ DEBERTA_START_DOCSTRING = r"""
947
+ The DeBERTa model was proposed in [DeBERTa: Decoding-enhanced BERT with Disentangled
948
+ Attention](https://arxiv.org/abs/2006.03654) by Pengcheng He, Xiaodong Liu, Jianfeng Gao, Weizhu Chen. It's build
949
+ on top of BERT/RoBERTa with two improvements, i.e. disentangled attention and enhanced mask decoder. With those two
950
+ improvements, it out perform BERT/RoBERTa on a majority of tasks with 80GB pretraining data.
951
+
952
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
953
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
954
+ and behavior.
955
+
956
+
957
+ Parameters:
958
+ config ([`DebertaV2Config`]): Model configuration class with all the parameters of the model.
959
+ Initializing with a config file does not load the weights associated with the model, only the
960
+ configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
961
+ """
962
+
963
+ DEBERTA_INPUTS_DOCSTRING = r"""
964
+ Args:
965
+ input_ids (`torch.LongTensor` of shape `({0})`):
966
+ Indices of input sequence tokens in the vocabulary.
967
+
968
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
969
+ [`PreTrainedTokenizer.__call__`] for details.
970
+
971
+ [What are input IDs?](../glossary#input-ids)
972
+ attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
973
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
974
+
975
+ - 1 for tokens that are **not masked**,
976
+ - 0 for tokens that are **masked**.
977
+
978
+ [What are attention masks?](../glossary#attention-mask)
979
+ token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*):
980
+ Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
981
+ 1]`:
982
+
983
+ - 0 corresponds to a *sentence A* token,
984
+ - 1 corresponds to a *sentence B* token.
985
+
986
+ [What are token type IDs?](../glossary#token-type-ids)
987
+ position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
988
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
989
+ config.max_position_embeddings - 1]`.
990
+
991
+ [What are position IDs?](../glossary#position-ids)
992
+ inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
993
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
994
+ is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
995
+ model's internal embedding lookup matrix.
996
+ output_attentions (`bool`, *optional*):
997
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
998
+ tensors for more detail.
999
+ output_hidden_states (`bool`, *optional*):
1000
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
1001
+ more detail.
1002
+ return_dict (`bool`, *optional*):
1003
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
1004
+ """
1005
+
1006
+
1007
+ @add_start_docstrings(
1008
+ "The bare DeBERTa Model transformer outputting raw hidden-states without any specific head on top.",
1009
+ DEBERTA_START_DOCSTRING,
1010
+ )
1011
+ # Copied from transformers.models.deberta.modeling_deberta.DebertaModel with Deberta->DebertaV2
1012
+ class DebertaV2Model(DebertaV2PreTrainedModel):
1013
+ def __init__(self, config):
1014
+ super().__init__(config)
1015
+
1016
+ self.embeddings = DebertaV2Embeddings(config)
1017
+ self.encoder = DebertaV2Encoder(config)
1018
+ self.z_steps = 0
1019
+ self.config = config
1020
+ # Initialize weights and apply final processing
1021
+ self.post_init()
1022
+
1023
+ def get_input_embeddings(self):
1024
+ return self.embeddings.word_embeddings
1025
+
1026
+ def set_input_embeddings(self, new_embeddings):
1027
+ self.embeddings.word_embeddings = new_embeddings
1028
+
1029
+ def _prune_heads(self, heads_to_prune):
1030
+ """
1031
+ Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
1032
+ class PreTrainedModel
1033
+ """
1034
+ raise NotImplementedError("The prune function is not implemented in DeBERTa model.")
1035
+
1036
+ @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
1037
+ @add_code_sample_docstrings(
1038
+ checkpoint=_CHECKPOINT_FOR_DOC,
1039
+ output_type=BaseModelOutput,
1040
+ config_class=_CONFIG_FOR_DOC,
1041
+ )
1042
+ def forward(
1043
+ self,
1044
+ input_ids: Optional[torch.Tensor] = None,
1045
+ attention_mask: Optional[torch.Tensor] = None,
1046
+ token_type_ids: Optional[torch.Tensor] = None,
1047
+ position_ids: Optional[torch.Tensor] = None,
1048
+ inputs_embeds: Optional[torch.Tensor] = None,
1049
+ output_attentions: Optional[bool] = None,
1050
+ output_hidden_states: Optional[bool] = None,
1051
+ return_dict: Optional[bool] = None,
1052
+ ) -> Union[Tuple, BaseModelOutput]:
1053
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1054
+ output_hidden_states = (
1055
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1056
+ )
1057
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1058
+
1059
+ if input_ids is not None and inputs_embeds is not None:
1060
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
1061
+ elif input_ids is not None:
1062
+ input_shape = input_ids.size()
1063
+ elif inputs_embeds is not None:
1064
+ input_shape = inputs_embeds.size()[:-1]
1065
+ else:
1066
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
1067
+
1068
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
1069
+
1070
+ if attention_mask is None:
1071
+ attention_mask = torch.ones(input_shape, device=device)
1072
+ if token_type_ids is None:
1073
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
1074
+
1075
+ embedding_output = self.embeddings(
1076
+ input_ids=input_ids,
1077
+ token_type_ids=token_type_ids,
1078
+ position_ids=position_ids,
1079
+ mask=attention_mask,
1080
+ inputs_embeds=inputs_embeds,
1081
+ )
1082
+
1083
+ encoder_outputs = self.encoder(
1084
+ embedding_output,
1085
+ attention_mask,
1086
+ output_hidden_states=True,
1087
+ output_attentions=output_attentions,
1088
+ return_dict=return_dict,
1089
+ )
1090
+ encoded_layers = encoder_outputs[1]
1091
+
1092
+ if self.z_steps > 1:
1093
+ hidden_states = encoded_layers[-2]
1094
+ layers = [self.encoder.layer[-1] for _ in range(self.z_steps)]
1095
+ query_states = encoded_layers[-1]
1096
+ rel_embeddings = self.encoder.get_rel_embedding()
1097
+ attention_mask = self.encoder.get_attention_mask(attention_mask)
1098
+ rel_pos = self.encoder.get_rel_pos(embedding_output)
1099
+ for layer in layers[1:]:
1100
+ query_states = layer(
1101
+ hidden_states,
1102
+ attention_mask,
1103
+ output_attentions=False,
1104
+ query_states=query_states,
1105
+ relative_pos=rel_pos,
1106
+ rel_embeddings=rel_embeddings,
1107
+ )
1108
+ encoded_layers.append(query_states)
1109
+
1110
+ sequence_output = encoded_layers[-1]
1111
+
1112
+ if not return_dict:
1113
+ return (sequence_output,) + encoder_outputs[(1 if output_hidden_states else 2) :]
1114
+
1115
+ return BaseModelOutput(
1116
+ last_hidden_state=sequence_output,
1117
+ hidden_states=encoder_outputs.hidden_states if output_hidden_states else None,
1118
+ attentions=encoder_outputs.attentions,
1119
+ )
1120
+
1121
+
1122
+ @add_start_docstrings("""DeBERTa Model with a `language modeling` head on top.""", DEBERTA_START_DOCSTRING)
1123
+ class DebertaV2ForMaskedLM(DebertaV2PreTrainedModel):
1124
+ _keys_to_ignore_on_load_unexpected = [r"pooler"]
1125
+ _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias", "cls.predictions.decoder.weight"]
1126
+
1127
+ def __init__(self, config):
1128
+ super().__init__(config)
1129
+
1130
+ self.deberta = DebertaV2Model(config)
1131
+ self.cls = DebertaV2OnlyMLMHead(config)
1132
+
1133
+ # Initialize weights and apply final processing
1134
+ self.post_init()
1135
+
1136
+ def get_output_embeddings(self):
1137
+ return self.cls.predictions.decoder
1138
+
1139
+ def set_output_embeddings(self, new_embeddings):
1140
+ self.cls.predictions.decoder = new_embeddings
1141
+
1142
+ @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
1143
+ @add_code_sample_docstrings(
1144
+ checkpoint=_CHECKPOINT_FOR_DOC,
1145
+ output_type=MaskedLMOutput,
1146
+ config_class=_CONFIG_FOR_DOC,
1147
+ mask="[MASK]",
1148
+ )
1149
+ # Copied from transformers.models.deberta.modeling_deberta.DebertaForMaskedLM.forward with Deberta->DebertaV2
1150
+ def forward(
1151
+ self,
1152
+ input_ids: Optional[torch.Tensor] = None,
1153
+ attention_mask: Optional[torch.Tensor] = None,
1154
+ token_type_ids: Optional[torch.Tensor] = None,
1155
+ position_ids: Optional[torch.Tensor] = None,
1156
+ inputs_embeds: Optional[torch.Tensor] = None,
1157
+ labels: Optional[torch.Tensor] = None,
1158
+ output_attentions: Optional[bool] = None,
1159
+ output_hidden_states: Optional[bool] = None,
1160
+ return_dict: Optional[bool] = None,
1161
+ ) -> Union[Tuple, MaskedLMOutput]:
1162
+ r"""
1163
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1164
+ Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
1165
+ config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
1166
+ loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
1167
+ """
1168
+
1169
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1170
+
1171
+ outputs = self.deberta(
1172
+ input_ids,
1173
+ attention_mask=attention_mask,
1174
+ token_type_ids=token_type_ids,
1175
+ position_ids=position_ids,
1176
+ inputs_embeds=inputs_embeds,
1177
+ output_attentions=output_attentions,
1178
+ output_hidden_states=output_hidden_states,
1179
+ return_dict=return_dict,
1180
+ )
1181
+
1182
+ sequence_output = outputs[0]
1183
+ prediction_scores = self.cls(sequence_output)
1184
+
1185
+ masked_lm_loss = None
1186
+ if labels is not None:
1187
+ loss_fct = CrossEntropyLoss() # -100 index = padding token
1188
+ masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
1189
+
1190
+ if not return_dict:
1191
+ output = (prediction_scores,) + outputs[1:]
1192
+ return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
1193
+
1194
+ return MaskedLMOutput(
1195
+ loss=masked_lm_loss,
1196
+ logits=prediction_scores,
1197
+ hidden_states=outputs.hidden_states,
1198
+ attentions=outputs.attentions,
1199
+ )
1200
+
1201
+
1202
+ # copied from transformers.models.bert.BertPredictionHeadTransform with bert -> deberta
1203
+ class DebertaV2PredictionHeadTransform(nn.Module):
1204
+ def __init__(self, config):
1205
+ super().__init__()
1206
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
1207
+ if isinstance(config.hidden_act, str):
1208
+ self.transform_act_fn = ACT2FN[config.hidden_act]
1209
+ else:
1210
+ self.transform_act_fn = config.hidden_act
1211
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
1212
+
1213
+ def forward(self, hidden_states):
1214
+ hidden_states = self.dense(hidden_states)
1215
+ hidden_states = self.transform_act_fn(hidden_states)
1216
+ hidden_states = self.LayerNorm(hidden_states)
1217
+ return hidden_states
1218
+
1219
+
1220
+ # copied from transformers.models.bert.BertLMPredictionHead with bert -> deberta
1221
+ class DebertaV2LMPredictionHead(nn.Module):
1222
+ def __init__(self, config):
1223
+ super().__init__()
1224
+ self.transform = DebertaV2PredictionHeadTransform(config)
1225
+
1226
+ # The output weights are the same as the input embeddings, but there is
1227
+ # an output-only bias for each token.
1228
+ self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1229
+
1230
+ self.bias = nn.Parameter(torch.zeros(config.vocab_size))
1231
+
1232
+ # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
1233
+ self.decoder.bias = self.bias
1234
+
1235
+ def forward(self, hidden_states):
1236
+ hidden_states = self.transform(hidden_states)
1237
+ hidden_states = self.decoder(hidden_states)
1238
+ return hidden_states
1239
+
1240
+
1241
+ # copied from transformers.models.bert.BertOnlyMLMHead with bert -> deberta
1242
+ class DebertaV2OnlyMLMHead(nn.Module):
1243
+ def __init__(self, config):
1244
+ super().__init__()
1245
+ self.predictions = DebertaV2LMPredictionHead(config)
1246
+
1247
+ def forward(self, sequence_output):
1248
+ prediction_scores = self.predictions(sequence_output)
1249
+ return prediction_scores
1250
+
1251
+
1252
+ @add_start_docstrings(
1253
+ """
1254
+ DeBERTa Model transformer with a sequence classification/regression head on top (a linear layer on top of the
1255
+ pooled output) e.g. for GLUE tasks.
1256
+ """,
1257
+ DEBERTA_START_DOCSTRING,
1258
+ )
1259
+ class DebertaV2ForSequenceClassification(DebertaV2PreTrainedModel):
1260
+ def __init__(self, config):
1261
+ super().__init__(config)
1262
+
1263
+ num_labels = getattr(config, "num_labels", 2)
1264
+ self.num_labels = num_labels
1265
+
1266
+ self.deberta = DebertaV2Model(config)
1267
+ self.pooler = ContextPooler(config)
1268
+ output_dim = self.pooler.output_dim
1269
+
1270
+ self.classifier = nn.Linear(output_dim, num_labels)
1271
+ drop_out = getattr(config, "cls_dropout", None)
1272
+ drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
1273
+ self.dropout = StableDropout(drop_out)
1274
+
1275
+ # Initialize weights and apply final processing
1276
+ self.post_init()
1277
+
1278
+ def get_input_embeddings(self):
1279
+ return self.deberta.get_input_embeddings()
1280
+
1281
+ def set_input_embeddings(self, new_embeddings):
1282
+ self.deberta.set_input_embeddings(new_embeddings)
1283
+
1284
+ @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
1285
+ @add_code_sample_docstrings(
1286
+ checkpoint=_CHECKPOINT_FOR_DOC,
1287
+ output_type=SequenceClassifierOutput,
1288
+ config_class=_CONFIG_FOR_DOC,
1289
+ )
1290
+ # Copied from transformers.models.deberta.modeling_deberta.DebertaForSequenceClassification.forward with Deberta->DebertaV2
1291
+ def forward(
1292
+ self,
1293
+ input_ids: Optional[torch.Tensor] = None,
1294
+ attention_mask: Optional[torch.Tensor] = None,
1295
+ token_type_ids: Optional[torch.Tensor] = None,
1296
+ position_ids: Optional[torch.Tensor] = None,
1297
+ inputs_embeds: Optional[torch.Tensor] = None,
1298
+ labels: Optional[torch.Tensor] = None,
1299
+ output_attentions: Optional[bool] = None,
1300
+ output_hidden_states: Optional[bool] = None,
1301
+ return_dict: Optional[bool] = None,
1302
+ ) -> Union[Tuple, SequenceClassifierOutput]:
1303
+ r"""
1304
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1305
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1306
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1307
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1308
+ """
1309
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1310
+
1311
+ outputs = self.deberta(
1312
+ input_ids,
1313
+ token_type_ids=token_type_ids,
1314
+ attention_mask=attention_mask,
1315
+ position_ids=position_ids,
1316
+ inputs_embeds=inputs_embeds,
1317
+ output_attentions=output_attentions,
1318
+ output_hidden_states=output_hidden_states,
1319
+ return_dict=return_dict,
1320
+ )
1321
+
1322
+ encoder_layer = outputs[0]
1323
+ pooled_output = self.pooler(encoder_layer)
1324
+ pooled_output = self.dropout(pooled_output)
1325
+ logits = self.classifier(pooled_output)
1326
+
1327
+ loss = None
1328
+ if labels is not None:
1329
+ if self.config.problem_type is None:
1330
+ if self.num_labels == 1:
1331
+ # regression task
1332
+ loss_fn = nn.MSELoss()
1333
+ logits = logits.view(-1).to(labels.dtype)
1334
+ loss = loss_fn(logits, labels.view(-1))
1335
+ elif labels.dim() == 1 or labels.size(-1) == 1:
1336
+ label_index = (labels >= 0).nonzero()
1337
+ labels = labels.long()
1338
+ if label_index.size(0) > 0:
1339
+ labeled_logits = torch.gather(
1340
+ logits, 0, label_index.expand(label_index.size(0), logits.size(1))
1341
+ )
1342
+ labels = torch.gather(labels, 0, label_index.view(-1))
1343
+ loss_fct = CrossEntropyLoss()
1344
+ loss = loss_fct(labeled_logits.view(-1, self.num_labels).float(), labels.view(-1))
1345
+ else:
1346
+ loss = torch.tensor(0).to(logits)
1347
+ else:
1348
+ log_softmax = nn.LogSoftmax(-1)
1349
+ loss = -((log_softmax(logits) * labels).sum(-1)).mean()
1350
+ elif self.config.problem_type == "regression":
1351
+ loss_fct = MSELoss()
1352
+ if self.num_labels == 1:
1353
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
1354
+ else:
1355
+ loss = loss_fct(logits, labels)
1356
+ elif self.config.problem_type == "single_label_classification":
1357
+ loss_fct = CrossEntropyLoss()
1358
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
1359
+ elif self.config.problem_type == "multi_label_classification":
1360
+ loss_fct = BCEWithLogitsLoss()
1361
+ loss = loss_fct(logits, labels)
1362
+ if not return_dict:
1363
+ output = (logits,) + outputs[1:]
1364
+ return ((loss,) + output) if loss is not None else output
1365
+
1366
+ return SequenceClassifierOutput(
1367
+ loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions
1368
+ )
1369
+
1370
+
1371
+ @add_start_docstrings(
1372
+ """
1373
+ DeBERTa Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
1374
+ Named-Entity-Recognition (NER) tasks.
1375
+ """,
1376
+ DEBERTA_START_DOCSTRING,
1377
+ )
1378
+ # Copied from transformers.models.deberta.modeling_deberta.DebertaForTokenClassification with Deberta->DebertaV2
1379
+ class DebertaV2ForTokenClassification(DebertaV2PreTrainedModel):
1380
+ _keys_to_ignore_on_load_unexpected = [r"pooler"]
1381
+
1382
+ def __init__(self, config):
1383
+ super().__init__(config)
1384
+ self.num_labels = config.num_labels
1385
+
1386
+ self.deberta = DebertaV2Model(config)
1387
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
1388
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
1389
+
1390
+ # Initialize weights and apply final processing
1391
+ self.post_init()
1392
+
1393
+ @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
1394
+ @add_code_sample_docstrings(
1395
+ checkpoint=_CHECKPOINT_FOR_DOC,
1396
+ output_type=TokenClassifierOutput,
1397
+ config_class=_CONFIG_FOR_DOC,
1398
+ )
1399
+ def forward(
1400
+ self,
1401
+ input_ids: Optional[torch.Tensor] = None,
1402
+ attention_mask: Optional[torch.Tensor] = None,
1403
+ token_type_ids: Optional[torch.Tensor] = None,
1404
+ position_ids: Optional[torch.Tensor] = None,
1405
+ inputs_embeds: Optional[torch.Tensor] = None,
1406
+ labels: Optional[torch.Tensor] = None,
1407
+ output_attentions: Optional[bool] = None,
1408
+ output_hidden_states: Optional[bool] = None,
1409
+ return_dict: Optional[bool] = None,
1410
+ ) -> Union[Tuple, TokenClassifierOutput]:
1411
+ r"""
1412
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1413
+ Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
1414
+ """
1415
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1416
+
1417
+ outputs = self.deberta(
1418
+ input_ids,
1419
+ attention_mask=attention_mask,
1420
+ token_type_ids=token_type_ids,
1421
+ position_ids=position_ids,
1422
+ inputs_embeds=inputs_embeds,
1423
+ output_attentions=output_attentions,
1424
+ output_hidden_states=output_hidden_states,
1425
+ return_dict=return_dict,
1426
+ )
1427
+
1428
+ sequence_output = outputs[0]
1429
+
1430
+ sequence_output = self.dropout(sequence_output)
1431
+ logits = self.classifier(sequence_output)
1432
+
1433
+ loss = None
1434
+ if labels is not None:
1435
+ loss_fct = CrossEntropyLoss()
1436
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
1437
+
1438
+ if not return_dict:
1439
+ output = (logits,) + outputs[1:]
1440
+ return ((loss,) + output) if loss is not None else output
1441
+
1442
+ return TokenClassifierOutput(
1443
+ loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions
1444
+ )
1445
+
1446
+ class TokenClassifierRegressionOutput(ModelOutput):
1447
+ """
1448
+ Base class for outputs of token classification models.
1449
+
1450
+ Args:
1451
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided) :
1452
+ Classification loss.
1453
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`):
1454
+ Classification scores (before SoftMax).
1455
+ hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
1456
+ Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
1457
+ one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
1458
+
1459
+ Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
1460
+ attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
1461
+ Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
1462
+ sequence_length)`.
1463
+
1464
+ Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
1465
+ heads.
1466
+ """
1467
+
1468
+ loss: Optional[torch.FloatTensor] = None
1469
+ logits: torch.FloatTensor = None
1470
+ values: torch.FloatTensor = None
1471
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
1472
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
1473
+
1474
+ class DebertaV2ForTokenClassificationRegression(DebertaV2PreTrainedModel):
1475
+ _keys_to_ignore_on_load_unexpected = [r"pooler"]
1476
+
1477
+ def __init__(self, config):
1478
+ super().__init__(config)
1479
+ self.num_labels = 4
1480
+
1481
+ self.deberta = DebertaV2Model(config)
1482
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
1483
+
1484
+ self.hidden1 = nn.Linear(config.hidden_size, config.hidden_size)
1485
+ self.classifier = nn.Linear(config.hidden_size, self.num_labels)
1486
+
1487
+ self.hidden2 = nn.Linear(config.hidden_size, config.hidden_size)
1488
+ self.regressor = nn.Linear(config.hidden_size, 1)
1489
+
1490
+ # Initialize weights and apply final processing
1491
+ self.post_init()
1492
+
1493
+ @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
1494
+ @add_code_sample_docstrings(
1495
+ checkpoint=_CHECKPOINT_FOR_DOC,
1496
+ output_type=TokenClassifierOutput,
1497
+ config_class=_CONFIG_FOR_DOC,
1498
+ )
1499
+ def forward(
1500
+ self,
1501
+ input_ids: Optional[torch.Tensor] = None,
1502
+ attention_mask: Optional[torch.Tensor] = None,
1503
+ token_type_ids: Optional[torch.Tensor] = None,
1504
+ position_ids: Optional[torch.Tensor] = None,
1505
+ inputs_embeds: Optional[torch.Tensor] = None,
1506
+ labels: Optional[torch.Tensor] = None,
1507
+ output_attentions: Optional[bool] = None,
1508
+ output_hidden_states: Optional[bool] = None,
1509
+ return_dict: Optional[bool] = None,
1510
+ ) -> Union[Tuple, TokenClassifierOutput]:
1511
+ r"""
1512
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1513
+ Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
1514
+ """
1515
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1516
+
1517
+ outputs = self.deberta(
1518
+ input_ids,
1519
+ attention_mask=attention_mask,
1520
+ token_type_ids=token_type_ids,
1521
+ position_ids=position_ids,
1522
+ inputs_embeds=inputs_embeds,
1523
+ output_attentions=output_attentions,
1524
+ output_hidden_states=output_hidden_states,
1525
+ return_dict=return_dict,
1526
+ )
1527
+
1528
+ sequence_output = outputs[0]
1529
+
1530
+ sequence_output = self.dropout(sequence_output)
1531
+
1532
+ logits = self.classifier(self.hidden1(sequence_output))
1533
+ values = self.regressor(self.hidden2(sequence_output))
1534
+
1535
+ loss = None
1536
+ if labels is not None:
1537
+ loss_fct = CrossEntropyLoss()
1538
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
1539
+
1540
+ if not return_dict:
1541
+ output = (logits,) + outputs[1:]
1542
+ return ((loss,) + output) if loss is not None else output
1543
+
1544
+ return TokenClassifierRegressionOutput(
1545
+ loss=loss, logits=logits, values=values, hidden_states=outputs.hidden_states, attentions=outputs.attentions
1546
+ )
1547
+
1548
+
1549
+ @add_start_docstrings(
1550
+ """
1551
+ DeBERTa Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
1552
+ layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
1553
+ """,
1554
+ DEBERTA_START_DOCSTRING,
1555
+ )
1556
+ class DebertaV2ForQuestionAnswering(DebertaV2PreTrainedModel):
1557
+ _keys_to_ignore_on_load_unexpected = [r"pooler"]
1558
+
1559
+ def __init__(self, config):
1560
+ super().__init__(config)
1561
+ self.num_labels = config.num_labels
1562
+
1563
+ self.deberta = DebertaV2Model(config)
1564
+ self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
1565
+
1566
+ # Initialize weights and apply final processing
1567
+ self.post_init()
1568
+
1569
+ @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
1570
+ @add_code_sample_docstrings(
1571
+ checkpoint=_CHECKPOINT_FOR_DOC,
1572
+ output_type=QuestionAnsweringModelOutput,
1573
+ config_class=_CONFIG_FOR_DOC,
1574
+ qa_target_start_index=_QA_TARGET_START_INDEX,
1575
+ qa_target_end_index=_QA_TARGET_END_INDEX,
1576
+ )
1577
+ # Copied from transformers.models.deberta.modeling_deberta.DebertaForQuestionAnswering.forward with Deberta->DebertaV2
1578
+ def forward(
1579
+ self,
1580
+ input_ids: Optional[torch.Tensor] = None,
1581
+ attention_mask: Optional[torch.Tensor] = None,
1582
+ token_type_ids: Optional[torch.Tensor] = None,
1583
+ position_ids: Optional[torch.Tensor] = None,
1584
+ inputs_embeds: Optional[torch.Tensor] = None,
1585
+ start_positions: Optional[torch.Tensor] = None,
1586
+ end_positions: Optional[torch.Tensor] = None,
1587
+ output_attentions: Optional[bool] = None,
1588
+ output_hidden_states: Optional[bool] = None,
1589
+ return_dict: Optional[bool] = None,
1590
+ ) -> Union[Tuple, QuestionAnsweringModelOutput]:
1591
+ r"""
1592
+ start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1593
+ Labels for position (index) of the start of the labelled span for computing the token classification loss.
1594
+ Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
1595
+ are not taken into account for computing the loss.
1596
+ end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1597
+ Labels for position (index) of the end of the labelled span for computing the token classification loss.
1598
+ Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
1599
+ are not taken into account for computing the loss.
1600
+ """
1601
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1602
+
1603
+ outputs = self.deberta(
1604
+ input_ids,
1605
+ attention_mask=attention_mask,
1606
+ token_type_ids=token_type_ids,
1607
+ position_ids=position_ids,
1608
+ inputs_embeds=inputs_embeds,
1609
+ output_attentions=output_attentions,
1610
+ output_hidden_states=output_hidden_states,
1611
+ return_dict=return_dict,
1612
+ )
1613
+
1614
+ sequence_output = outputs[0]
1615
+
1616
+ logits = self.qa_outputs(sequence_output)
1617
+ start_logits, end_logits = logits.split(1, dim=-1)
1618
+ start_logits = start_logits.squeeze(-1).contiguous()
1619
+ end_logits = end_logits.squeeze(-1).contiguous()
1620
+
1621
+ total_loss = None
1622
+ if start_positions is not None and end_positions is not None:
1623
+ # If we are on multi-GPU, split add a dimension
1624
+ if len(start_positions.size()) > 1:
1625
+ start_positions = start_positions.squeeze(-1)
1626
+ if len(end_positions.size()) > 1:
1627
+ end_positions = end_positions.squeeze(-1)
1628
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
1629
+ ignored_index = start_logits.size(1)
1630
+ start_positions = start_positions.clamp(0, ignored_index)
1631
+ end_positions = end_positions.clamp(0, ignored_index)
1632
+
1633
+ loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
1634
+ start_loss = loss_fct(start_logits, start_positions)
1635
+ end_loss = loss_fct(end_logits, end_positions)
1636
+ total_loss = (start_loss + end_loss) / 2
1637
+
1638
+ if not return_dict:
1639
+ output = (start_logits, end_logits) + outputs[1:]
1640
+ return ((total_loss,) + output) if total_loss is not None else output
1641
+
1642
+ return QuestionAnsweringModelOutput(
1643
+ loss=total_loss,
1644
+ start_logits=start_logits,
1645
+ end_logits=end_logits,
1646
+ hidden_states=outputs.hidden_states,
1647
+ attentions=outputs.attentions,
1648
+ )
1649
+
1650
+
1651
+ @add_start_docstrings(
1652
+ """
1653
+ DeBERTa Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
1654
+ softmax) e.g. for RocStories/SWAG tasks.
1655
+ """,
1656
+ DEBERTA_START_DOCSTRING,
1657
+ )
1658
+ class DebertaV2ForMultipleChoice(DebertaV2PreTrainedModel):
1659
+ def __init__(self, config):
1660
+ super().__init__(config)
1661
+
1662
+ num_labels = getattr(config, "num_labels", 2)
1663
+ self.num_labels = num_labels
1664
+
1665
+ self.deberta = DebertaV2Model(config)
1666
+ self.pooler = ContextPooler(config)
1667
+ output_dim = self.pooler.output_dim
1668
+
1669
+ self.classifier = nn.Linear(output_dim, 1)
1670
+ drop_out = getattr(config, "cls_dropout", None)
1671
+ drop_out = self.config.hidden_dropout_prob if drop_out is None else drop_out
1672
+ self.dropout = StableDropout(drop_out)
1673
+
1674
+ self.init_weights()
1675
+
1676
+ def get_input_embeddings(self):
1677
+ return self.deberta.get_input_embeddings()
1678
+
1679
+ def set_input_embeddings(self, new_embeddings):
1680
+ self.deberta.set_input_embeddings(new_embeddings)
1681
+
1682
+ @add_start_docstrings_to_model_forward(DEBERTA_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
1683
+ @add_code_sample_docstrings(
1684
+ checkpoint=_CHECKPOINT_FOR_DOC,
1685
+ output_type=MultipleChoiceModelOutput,
1686
+ config_class=_CONFIG_FOR_DOC,
1687
+ )
1688
+ def forward(
1689
+ self,
1690
+ input_ids: Optional[torch.Tensor] = None,
1691
+ attention_mask: Optional[torch.Tensor] = None,
1692
+ token_type_ids: Optional[torch.Tensor] = None,
1693
+ position_ids: Optional[torch.Tensor] = None,
1694
+ inputs_embeds: Optional[torch.Tensor] = None,
1695
+ labels: Optional[torch.Tensor] = None,
1696
+ output_attentions: Optional[bool] = None,
1697
+ output_hidden_states: Optional[bool] = None,
1698
+ return_dict: Optional[bool] = None,
1699
+ ) -> Union[Tuple, MultipleChoiceModelOutput]:
1700
+ r"""
1701
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1702
+ Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
1703
+ num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
1704
+ `input_ids` above)
1705
+ """
1706
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1707
+ num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
1708
+
1709
+ flat_input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
1710
+ flat_position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
1711
+ flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
1712
+ flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
1713
+ flat_inputs_embeds = (
1714
+ inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
1715
+ if inputs_embeds is not None
1716
+ else None
1717
+ )
1718
+
1719
+ outputs = self.deberta(
1720
+ flat_input_ids,
1721
+ position_ids=flat_position_ids,
1722
+ token_type_ids=flat_token_type_ids,
1723
+ attention_mask=flat_attention_mask,
1724
+ inputs_embeds=flat_inputs_embeds,
1725
+ output_attentions=output_attentions,
1726
+ output_hidden_states=output_hidden_states,
1727
+ return_dict=return_dict,
1728
+ )
1729
+
1730
+ encoder_layer = outputs[0]
1731
+ pooled_output = self.pooler(encoder_layer)
1732
+ pooled_output = self.dropout(pooled_output)
1733
+ logits = self.classifier(pooled_output)
1734
+ reshaped_logits = logits.view(-1, num_choices)
1735
+
1736
+ loss = None
1737
+ if labels is not None:
1738
+ loss_fct = CrossEntropyLoss()
1739
+ loss = loss_fct(reshaped_logits, labels)
1740
+
1741
+ if not return_dict:
1742
+ output = (reshaped_logits,) + outputs[1:]
1743
+ return ((loss,) + output) if loss is not None else output
1744
+
1745
+ return MultipleChoiceModelOutput(
1746
+ loss=loss,
1747
+ logits=reshaped_logits,
1748
+ hidden_states=outputs.hidden_states,
1749
+ attentions=outputs.attentions,
1750
+ )
.ipynb_checkpoints/models-checkpoint.py ADDED
@@ -0,0 +1,738 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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.beat_embedding_layer.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).cuda(), torch.tensor(out_beat_timing).cuda(), torch.tensor(out_mask).cuda() #batch, len_beat
444
+ embedded_beat = self.beat_embedding_layer(out_beat, out_beat_timing)
445
+
446
+ return embedded_beat, out_mask
447
+
448
+ def encode_chords(self, chords,chords_time):
449
+ out_chord_root = []
450
+ out_chord_type = []
451
+ out_chord_inv = []
452
+ out_chord_timing = []
453
+ out_mask = []
454
+ for chord, chord_time in zip(chords,chords_time): #batch loop
455
+ tokenized_chord_root, tokenized_chord_type, tokenized_chord_inv, tokenized_chord_time, tokenized_chord_mask = self.chord_tokenizer(chord, chord_time)
456
+ out_chord_root.append(tokenized_chord_root)
457
+ out_chord_type.append(tokenized_chord_type)
458
+ out_chord_inv.append(tokenized_chord_inv)
459
+ out_chord_timing.append(tokenized_chord_time)
460
+ out_mask.append(tokenized_chord_mask)
461
+ #chords: (B, LEN, 4)
462
+ out_chord_root, out_chord_type, out_chord_inv, out_chord_timing, out_mask = torch.tensor(out_chord_root).cuda(), torch.tensor(out_chord_type).cuda(), torch.tensor(out_chord_inv).cuda(), torch.tensor(out_chord_timing).cuda(), torch.tensor(out_mask).cuda()
463
+ embedded_chord = self.chord_embedding_layer(out_chord_root, out_chord_type, out_chord_inv, out_chord_timing)
464
+ return embedded_chord, out_mask
465
+ # return out_chord_root, out_mask
466
+
467
+
468
+ def forward(self, latents, prompt, beats, chords,chords_time, validation_mode=False):
469
+ device = self.text_encoder.device
470
+ num_train_timesteps = self.noise_scheduler.num_train_timesteps
471
+ self.noise_scheduler.set_timesteps(num_train_timesteps, device=device)
472
+
473
+ encoder_hidden_states, boolean_encoder_mask = self.encode_text(prompt)
474
+
475
+ # with torch.no_grad():
476
+ encoded_beats, beat_mask = self.encode_beats(beats) #batch, len_beats, dim; batch, len_beats
477
+ encoded_chords, chord_mask = self.encode_chords(chords,chords_time)
478
+
479
+
480
+ if self.uncondition:
481
+ mask_indices = [k for k in range(len(prompt)) if random.random() < 0.1]
482
+ if len(mask_indices) > 0:
483
+ encoder_hidden_states[mask_indices] = 0
484
+ encoded_chords[mask_indices] = 0
485
+ encoded_beats[mask_indices] = 0
486
+
487
+ bsz = latents.shape[0]
488
+
489
+ if validation_mode:
490
+ timesteps = (self.noise_scheduler.num_train_timesteps//2) * torch.ones((bsz,), dtype=torch.int64, device=device)
491
+ else:
492
+ timesteps = torch.randint(0, self.noise_scheduler.num_train_timesteps, (bsz,), device=device)
493
+
494
+
495
+ timesteps = timesteps.long()
496
+
497
+ noise = torch.randn_like(latents)
498
+ noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)
499
+
500
+ # Get the target for loss depending on the prediction type
501
+ if self.noise_scheduler.config.prediction_type == "epsilon":
502
+ target = noise
503
+ elif self.noise_scheduler.config.prediction_type == "v_prediction":
504
+ target = self.noise_scheduler.get_velocity(latents, noise, timesteps)
505
+ else:
506
+ raise ValueError(f"Unknown prediction type {self.noise_scheduler.config.prediction_type}")
507
+
508
+ if self.set_from == "random":
509
+ # model_pred = torch.zeros((bsz,8,256,16)).to(device)
510
+ model_pred = self.unet(
511
+ noisy_latents, timesteps, encoder_hidden_states, encoded_beats, encoded_chords,
512
+ encoder_attention_mask=boolean_encoder_mask, beat_attention_mask = beat_mask, chord_attention_mask = chord_mask
513
+ ).sample
514
+
515
+ elif self.set_from == "pre-trained":
516
+ compressed_latents = self.group_in(noisy_latents.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()
517
+ model_pred = self.unet(
518
+ compressed_latents, timesteps, encoder_hidden_states,
519
+ encoder_attention_mask=boolean_encoder_mask
520
+ ).sample
521
+ model_pred = self.group_out(model_pred.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()
522
+
523
+ if self.snr_gamma is None:
524
+ loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
525
+ else:
526
+ # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
527
+ # Adaptef from huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py
528
+ snr = self.compute_snr(timesteps)
529
+ mse_loss_weights = (
530
+ torch.stack([snr, self.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr
531
+ )
532
+ loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
533
+ loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights
534
+ loss = loss.mean()
535
+
536
+ return loss
537
+
538
+ @torch.no_grad()
539
+ def inference(self, prompt, beats, chords,chords_time, inference_scheduler, num_steps=20, guidance_scale=3, num_samples_per_prompt=1,
540
+ disable_progress=True):
541
+ device = self.text_encoder.device
542
+ classifier_free_guidance = guidance_scale > 1.0
543
+ batch_size = len(prompt) * num_samples_per_prompt
544
+
545
+ if classifier_free_guidance:
546
+ prompt_embeds, boolean_prompt_mask = self.encode_text_classifier_free(prompt, num_samples_per_prompt)
547
+ encoded_beats, beat_mask = self.encode_beats_classifier_free(beats, num_samples_per_prompt) #batch, len_beats, dim; batch, len_beats
548
+ encoded_chords, chord_mask = self.encode_chords_classifier_free(chords, chords_time, num_samples_per_prompt)
549
+ else:
550
+ prompt_embeds, boolean_prompt_mask = self.encode_text(prompt)
551
+ prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
552
+ boolean_prompt_mask = boolean_prompt_mask.repeat_interleave(num_samples_per_prompt, 0)
553
+
554
+ encoded_beats, beat_mask = self.encode_beats(beats) #batch, len_beats, dim; batch, len_beats
555
+ encoded_beats = encoded_beats.repeat_interleave(num_samples_per_prompt, 0)
556
+ beat_mask = beat_mask.repeat_interleave(num_samples_per_prompt, 0)
557
+
558
+ encoded_chords, chord_mask = self.encode_chords(chords,chords_time)
559
+ encoded_chords = encoded_chords.repeat_interleave(num_samples_per_prompt, 0)
560
+ chord_mask = chord_mask.repeat_interleave(num_samples_per_prompt, 0)
561
+
562
+ # print(f"encoded_chords:{encoded_chords.shape}, chord_mask:{chord_mask.shape}, prompt_embeds:{prompt_embeds.shape},boolean_prompt_mask:{boolean_prompt_mask.shape} ")
563
+ inference_scheduler.set_timesteps(num_steps, device=device)
564
+ timesteps = inference_scheduler.timesteps
565
+
566
+ num_channels_latents = self.unet.in_channels
567
+ latents = self.prepare_latents(batch_size, inference_scheduler, num_channels_latents, prompt_embeds.dtype, device)
568
+
569
+ num_warmup_steps = len(timesteps) - num_steps * inference_scheduler.order
570
+ progress_bar = tqdm(range(num_steps), disable=disable_progress)
571
+
572
+ for i, t in enumerate(timesteps):
573
+ # expand the latents if we are doing classifier free guidance
574
+ latent_model_input = torch.cat([latents] * 2) if classifier_free_guidance else latents
575
+ latent_model_input = inference_scheduler.scale_model_input(latent_model_input, t)
576
+
577
+ noise_pred = self.unet(
578
+ latent_model_input, t, encoder_hidden_states=prompt_embeds,
579
+ encoder_attention_mask=boolean_prompt_mask,
580
+ beat_features = encoded_beats, beat_attention_mask = beat_mask, chord_features = encoded_chords,chord_attention_mask = chord_mask
581
+ ).sample
582
+
583
+ # perform guidance
584
+ if classifier_free_guidance: #should work for beats and chords too
585
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
586
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
587
+
588
+ # compute the previous noisy sample x_t -> x_t-1
589
+ latents = inference_scheduler.step(noise_pred, t, latents).prev_sample
590
+
591
+ # call the callback, if provided
592
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % inference_scheduler.order == 0):
593
+ progress_bar.update(1)
594
+
595
+ if self.set_from == "pre-trained":
596
+ latents = self.group_out(latents.permute(0, 2, 3, 1).contiguous()).permute(0, 3, 1, 2).contiguous()
597
+ return latents
598
+
599
+ def prepare_latents(self, batch_size, inference_scheduler, num_channels_latents, dtype, device):
600
+ shape = (batch_size, num_channels_latents, 256, 16)
601
+ latents = randn_tensor(shape, generator=None, device=device, dtype=dtype)
602
+ # scale the initial noise by the standard deviation required by the scheduler
603
+ latents = latents * inference_scheduler.init_noise_sigma
604
+ return latents
605
+
606
+ def encode_text_classifier_free(self, prompt, num_samples_per_prompt):
607
+ device = self.text_encoder.device
608
+ batch = self.tokenizer(
609
+ prompt, max_length=self.tokenizer.model_max_length, padding=True, truncation=True, return_tensors="pt"
610
+ )
611
+ input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device)
612
+
613
+ with torch.no_grad():
614
+ prompt_embeds = self.text_encoder(
615
+ input_ids=input_ids, attention_mask=attention_mask
616
+ )[0]
617
+
618
+ prompt_embeds = prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
619
+ attention_mask = attention_mask.repeat_interleave(num_samples_per_prompt, 0)
620
+
621
+ # get unconditional embeddings for classifier free guidance
622
+ # print(len(prompt), 'this is prompt len')
623
+ uncond_tokens = [""] * len(prompt)
624
+
625
+ max_length = prompt_embeds.shape[1]
626
+ uncond_batch = self.tokenizer(
627
+ uncond_tokens, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt",
628
+ )
629
+ uncond_input_ids = uncond_batch.input_ids.to(device)
630
+ uncond_attention_mask = uncond_batch.attention_mask.to(device)
631
+
632
+ with torch.no_grad():
633
+ negative_prompt_embeds = self.text_encoder(
634
+ input_ids=uncond_input_ids, attention_mask=uncond_attention_mask
635
+ )[0]
636
+
637
+ negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_samples_per_prompt, 0)
638
+ uncond_attention_mask = uncond_attention_mask.repeat_interleave(num_samples_per_prompt, 0)
639
+
640
+ # For classifier free guidance, we need to do two forward passes.
641
+ # We concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes
642
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
643
+ prompt_mask = torch.cat([uncond_attention_mask, attention_mask])
644
+ boolean_prompt_mask = (prompt_mask == 1).to(device)
645
+
646
+ return prompt_embeds, boolean_prompt_mask
647
+
648
+
649
+ def encode_beats_classifier_free(self, beats, num_samples_per_prompt):
650
+ with torch.no_grad():
651
+ out_beat = []
652
+ out_beat_timing = []
653
+ out_mask = []
654
+ for beat in beats:
655
+ tokenized_beats,tokenized_beats_timing, tokenized_beat_mask = self.beat_tokenizer(beat)
656
+ out_beat.append(tokenized_beats)
657
+ out_beat_timing.append(tokenized_beats_timing)
658
+ out_mask.append(tokenized_beat_mask)
659
+ out_beat, out_beat_timing, out_mask = torch.tensor(out_beat).cuda(), torch.tensor(out_beat_timing).cuda(), torch.tensor(out_mask).cuda() #batch, len_beat
660
+ embedded_beat = self.beat_embedding_layer(out_beat, out_beat_timing)
661
+
662
+ embedded_beat = embedded_beat.repeat_interleave(num_samples_per_prompt, 0)
663
+ out_mask = out_mask.repeat_interleave(num_samples_per_prompt, 0)
664
+
665
+ uncond_beats = [[[],[]]] * len(beats)
666
+
667
+ max_length = embedded_beat.shape[1]
668
+ with torch.no_grad():
669
+ out_beat_unc = []
670
+ out_beat_timing_unc = []
671
+ out_mask_unc = []
672
+ for beat in uncond_beats:
673
+ tokenized_beats, tokenized_beats_timing, tokenized_beat_mask = self.beat_tokenizer(beat)
674
+ out_beat_unc.append(tokenized_beats)
675
+ out_beat_timing_unc.append(tokenized_beats_timing)
676
+ out_mask_unc.append(tokenized_beat_mask)
677
+ out_beat_unc, out_beat_timing_unc, out_mask_unc = torch.tensor(out_beat_unc).cuda(), torch.tensor(out_beat_timing_unc).cuda(), torch.tensor(out_mask_unc).cuda() #batch, len_beat
678
+ embedded_beat_unc = self.beat_embedding_layer(out_beat_unc, out_beat_timing_unc)
679
+
680
+ embedded_beat_unc = embedded_beat_unc.repeat_interleave(num_samples_per_prompt, 0)
681
+ out_mask_unc = out_mask_unc.repeat_interleave(num_samples_per_prompt, 0)
682
+
683
+ embedded_beat = torch.cat([embedded_beat_unc, embedded_beat])
684
+ out_mask = torch.cat([out_mask_unc, out_mask])
685
+
686
+ return embedded_beat, out_mask
687
+
688
+
689
+ def encode_chords_classifier_free(self, chords, chords_time, num_samples_per_prompt):
690
+
691
+ with torch.no_grad():
692
+ out_chord_root = []
693
+ out_chord_type = []
694
+ out_chord_inv = []
695
+ out_chord_timing = []
696
+ out_mask = []
697
+ for chord, chord_time in zip(chords,chords_time): #batch loop
698
+ tokenized_chord_root, tokenized_chord_type, tokenized_chord_inv, tokenized_chord_time, tokenized_chord_mask = self.chord_tokenizer(chord, chord_time)
699
+ out_chord_root.append(tokenized_chord_root)
700
+ out_chord_type.append(tokenized_chord_type)
701
+ out_chord_inv.append(tokenized_chord_inv)
702
+ out_chord_timing.append(tokenized_chord_time)
703
+ out_mask.append(tokenized_chord_mask)
704
+ out_chord_root, out_chord_type, out_chord_inv, out_chord_timing, out_mask = torch.tensor(out_chord_root).cuda(), torch.tensor(out_chord_type).cuda(), torch.tensor(out_chord_inv).cuda(), torch.tensor(out_chord_timing).cuda(), torch.tensor(out_mask).cuda()
705
+ embedded_chord = self.chord_embedding_layer(out_chord_root, out_chord_type, out_chord_inv, out_chord_timing)
706
+
707
+ embedded_chord = embedded_chord.repeat_interleave(num_samples_per_prompt, 0)
708
+ out_mask = out_mask.repeat_interleave(num_samples_per_prompt, 0)
709
+
710
+ chords_unc=[[]] * len(chords)
711
+ chords_time_unc=[[]] * len(chords_time)
712
+
713
+ max_length = embedded_chord.shape[1]
714
+
715
+ with torch.no_grad():
716
+ out_chord_root_unc = []
717
+ out_chord_type_unc = []
718
+ out_chord_inv_unc = []
719
+ out_chord_timing_unc = []
720
+ out_mask_unc = []
721
+ for chord, chord_time in zip(chords_unc,chords_time_unc): #batch loop
722
+ tokenized_chord_root, tokenized_chord_type, tokenized_chord_inv, tokenized_chord_time, tokenized_chord_mask = self.chord_tokenizer(chord, chord_time)
723
+ out_chord_root_unc.append(tokenized_chord_root)
724
+ out_chord_type_unc.append(tokenized_chord_type)
725
+ out_chord_inv_unc.append(tokenized_chord_inv)
726
+ out_chord_timing_unc.append(tokenized_chord_time)
727
+ out_mask_unc.append(tokenized_chord_mask)
728
+ 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).cuda(), torch.tensor(out_chord_type_unc).cuda(), torch.tensor(out_chord_inv_unc).cuda(), torch.tensor(out_chord_timing_unc).cuda(), torch.tensor(out_mask_unc).cuda()
729
+ embedded_chord_unc = self.chord_embedding_layer(out_chord_root_unc, out_chord_type_unc, out_chord_inv_unc, out_chord_timing_unc)
730
+
731
+
732
+ embedded_chord_unc = embedded_chord_unc.repeat_interleave(num_samples_per_prompt, 0)
733
+ out_mask_unc = out_mask_unc.repeat_interleave(num_samples_per_prompt, 0)
734
+
735
+ embedded_chord = torch.cat([embedded_chord_unc, embedded_chord])
736
+ out_mask = torch.cat([out_mask_unc, out_mask])
737
+
738
+ return embedded_chord, out_mask
app.py ADDED
@@ -0,0 +1,260 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import json
3
+ import torch
4
+ import wavio
5
+ from tqdm import tqdm
6
+ from huggingface_hub import snapshot_download
7
+
8
+ from audioldm.audio.stft import TacotronSTFT
9
+ from audioldm.variational_autoencoder import AutoencoderKL
10
+
11
+ from transformers import AutoTokenizer, T5ForConditionalGeneration
12
+ from modelling_deberta_v2 import DebertaV2ForTokenClassificationRegression
13
+
14
+ from diffusers import DDPMScheduler
15
+ from models import MusicAudioDiffusion
16
+
17
+ from gradio import Markdown
18
+
19
+ class MusicFeaturePredictor:
20
+ def __init__(self, path, device="cuda:0", cache_dir=None, local_files_only=False):
21
+ self.beats_tokenizer = AutoTokenizer.from_pretrained(
22
+ "microsoft/deberta-v3-large",
23
+ cache_dir=cache_dir,
24
+ local_files_only=local_files_only,
25
+ )
26
+ self.beats_model = DebertaV2ForTokenClassificationRegression.from_pretrained(
27
+ "microsoft/deberta-v3-large",
28
+ cache_dir=cache_dir,
29
+ local_files_only=local_files_only,
30
+ )
31
+ self.beats_model.eval()
32
+ self.beats_model.to(device)
33
+
34
+ beats_ckpt = f"{path}/beats/microsoft-deberta-v3-large.pt"
35
+ beats_weight = torch.load(beats_ckpt, map_location="cpu")
36
+ self.beats_model.load_state_dict(beats_weight)
37
+
38
+ self.chords_tokenizer = AutoTokenizer.from_pretrained(
39
+ "google/flan-t5-large",
40
+ cache_dir=cache_dir,
41
+ local_files_only=local_files_only,
42
+ )
43
+ self.chords_model = T5ForConditionalGeneration.from_pretrained(
44
+ "google/flan-t5-large",
45
+ cache_dir=cache_dir,
46
+ local_files_only=local_files_only,
47
+ )
48
+ self.chords_model.eval()
49
+ self.chords_model.to(device)
50
+
51
+ chords_ckpt = f"{path}/chords/flan-t5-large.bin"
52
+ chords_weight = torch.load(chords_ckpt, map_location="cpu")
53
+ self.chords_model.load_state_dict(chords_weight)
54
+
55
+ def generate_beats(self, prompt):
56
+ tokenized = self.beats_tokenizer(
57
+ prompt, max_length=512, padding=True, truncation=True, return_tensors="pt"
58
+ )
59
+ tokenized = {k: v.to(self.beats_model.device) for k, v in tokenized.items()}
60
+
61
+ with torch.no_grad():
62
+ out = self.beats_model(**tokenized)
63
+
64
+ max_beat = (
65
+ 1 + torch.argmax(out["logits"][:, 0, :], -1).detach().cpu().numpy()
66
+ ).tolist()[0]
67
+ intervals = (
68
+ out["values"][:, :, 0]
69
+ .detach()
70
+ .cpu()
71
+ .numpy()
72
+ .astype("float32")
73
+ .round(4)
74
+ .tolist()
75
+ )
76
+
77
+ intervals = np.cumsum(intervals)
78
+ predicted_beats_times = []
79
+ for t in intervals:
80
+ if t < 10:
81
+ predicted_beats_times.append(round(t, 2))
82
+ else:
83
+ break
84
+ predicted_beats_times = list(np.array(predicted_beats_times)[:50])
85
+
86
+ if len(predicted_beats_times) == 0:
87
+ predicted_beats = [[], []]
88
+ else:
89
+ beat_counts = []
90
+ for i in range(len(predicted_beats_times)):
91
+ beat_counts.append(float(1.0 + np.mod(i, max_beat)))
92
+ predicted_beats = [[predicted_beats_times, beat_counts]]
93
+
94
+ return max_beat, predicted_beats_times, predicted_beats
95
+
96
+ def generate(self, prompt):
97
+ max_beat, predicted_beats_times, predicted_beats = self.generate_beats(prompt)
98
+
99
+ chords_prompt = "Caption: {} \\n Timestamps: {} \\n Max Beat: {}".format(
100
+ prompt,
101
+ " , ".join([str(round(t, 2)) for t in predicted_beats_times]),
102
+ max_beat,
103
+ )
104
+
105
+ tokenized = self.chords_tokenizer(
106
+ chords_prompt,
107
+ max_length=512,
108
+ padding=True,
109
+ truncation=True,
110
+ return_tensors="pt",
111
+ )
112
+ tokenized = {k: v.to(self.chords_model.device) for k, v in tokenized.items()}
113
+
114
+ generated_chords = self.chords_model.generate(
115
+ input_ids=tokenized["input_ids"],
116
+ attention_mask=tokenized["attention_mask"],
117
+ min_length=8,
118
+ max_length=128,
119
+ num_beams=5,
120
+ early_stopping=True,
121
+ num_return_sequences=1,
122
+ )
123
+
124
+ generated_chords = self.chords_tokenizer.decode(
125
+ generated_chords[0],
126
+ skip_special_tokens=True,
127
+ clean_up_tokenization_spaces=True,
128
+ ).split(" n ")
129
+
130
+ predicted_chords, predicted_chords_times = [], []
131
+ for item in generated_chords:
132
+ c, ct = item.split(" at ")
133
+ predicted_chords.append(c)
134
+ predicted_chords_times.append(float(ct))
135
+
136
+ return predicted_beats, predicted_chords, predicted_chords_times
137
+
138
+
139
+ class Mustango:
140
+ def __init__(
141
+ self,
142
+ name="declare-lab/mustango",
143
+ device="cuda:0",
144
+ cache_dir=None,
145
+ local_files_only=False,
146
+ ):
147
+ path = snapshot_download(repo_id=name, cache_dir=cache_dir)
148
+
149
+ self.music_model = MusicFeaturePredictor(
150
+ path, device, cache_dir=cache_dir, local_files_only=local_files_only
151
+ )
152
+
153
+ vae_config = json.load(open(f"{path}/configs/vae_config.json"))
154
+ stft_config = json.load(open(f"{path}/configs/stft_config.json"))
155
+ main_config = json.load(open(f"{path}/configs/main_config.json"))
156
+
157
+ self.vae = AutoencoderKL(**vae_config).to(device)
158
+ self.stft = TacotronSTFT(**stft_config).to(device)
159
+ self.model = MusicAudioDiffusion(
160
+ main_config["text_encoder_name"],
161
+ main_config["scheduler_name"],
162
+ unet_model_config_path=f"{path}/configs/music_diffusion_model_config.json",
163
+ ).to(device)
164
+
165
+ vae_weights = torch.load(
166
+ f"{path}/vae/pytorch_model_vae.bin", map_location=device
167
+ )
168
+ stft_weights = torch.load(
169
+ f"{path}/stft/pytorch_model_stft.bin", map_location=device
170
+ )
171
+ main_weights = torch.load(
172
+ f"{path}/ldm/pytorch_model_ldm.bin", map_location=device
173
+ )
174
+
175
+ self.vae.load_state_dict(vae_weights)
176
+ self.stft.load_state_dict(stft_weights)
177
+ self.model.load_state_dict(main_weights)
178
+
179
+ print("Successfully loaded checkpoint from:", name)
180
+
181
+ self.vae.eval()
182
+ self.stft.eval()
183
+ self.model.eval()
184
+
185
+ self.scheduler = DDPMScheduler.from_pretrained(
186
+ main_config["scheduler_name"], subfolder="scheduler"
187
+ )
188
+
189
+ def generate(self, prompt, steps=100, guidance=3, samples=1, disable_progress=True):
190
+ """Genrate music for a single prompt string."""
191
+
192
+ with torch.no_grad():
193
+ beats, chords, chords_times = self.music_model.generate(prompt)
194
+ latents = self.model.inference(
195
+ [prompt],
196
+ beats,
197
+ [chords],
198
+ [chords_times],
199
+ self.scheduler,
200
+ steps,
201
+ guidance,
202
+ samples,
203
+ disable_progress,
204
+ )
205
+ mel = self.vae.decode_first_stage(latents)
206
+ wave = self.vae.decode_to_waveform(mel)
207
+
208
+ return wave[0]
209
+
210
+
211
+ # Initialize Mustango
212
+ if torch.cuda.is_available():
213
+ mustango = Mustango()
214
+ else:
215
+ mustango = Mustango(device="cpu")
216
+
217
+ def gradio_generate(prompt, steps, guidance):
218
+ output_wave = mustango.generate(prompt, steps, guidance)
219
+ # output_filename = f"{prompt.replace(' ', '_')}_{steps}_{guidance}"[:250] + ".wav"
220
+ output_filename = "temp.wav"
221
+ wavio.write(output_filename, output_wave, rate=16000, sampwidth=2)
222
+
223
+ return output_filename
224
+
225
+ # description_text = """
226
+ # <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/>
227
+ # Generate music using Mustango by providing a text prompt.
228
+ # <br/><br/> Meet Mustango, an exciting addition to the vibrant landscape of Multimodal Large Language Models \
229
+ # designed for controlled music generation. Mustango leverages Latent Diffusion Model (LDM), Flan-T5, and \
230
+ # musical features to do the magic! \
231
+ # <p/>
232
+ # """
233
+ description_text = ""
234
+ # Gradio input and output components
235
+ input_text = gr.inputs.Textbox(lines=2, label="Prompt")
236
+ output_audio = gr.outputs.Audio(label="Generated Music", type="filepath")
237
+ denoising_steps = gr.Slider(minimum=100, maximum=200, value=100, step=1, label="Steps", interactive=True)
238
+ guidance_scale = gr.Slider(minimum=1, maximum=10, value=3, step=0.1, label="Guidance Scale", interactive=True)
239
+
240
+ # Gradio interface
241
+ gr_interface = gr.Interface(
242
+ fn=gradio_generate,
243
+ inputs=[input_text, denoising_steps, guidance_scale],
244
+ outputs=[output_audio],
245
+ title="Mustango: Toward Controllable Text-to-Music Generation",
246
+ description=description_text,
247
+ allow_flagging=False,
248
+ examples=[
249
+ ["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."],
250
+ ["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."],
251
+ ["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."],
252
+ ["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."],
253
+ ["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."],
254
+ ["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."],
255
+ ],
256
+ cache_examples=False,
257
+ )
258
+
259
+ # Launch Gradio app
260
+ gr_interface.launch()
audioldm/__init__.py ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ from .ldm import LatentDiffusion
2
+ from .utils import seed_everything, save_wave, get_time, get_duration
3
+ from .pipeline import *
4
+
5
+
6
+
7
+
8
+
audioldm/__main__.py ADDED
@@ -0,0 +1,183 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/python3
2
+ import os
3
+ from audioldm import text_to_audio, style_transfer, build_model, save_wave, get_time, round_up_duration, get_duration
4
+ import argparse
5
+
6
+ CACHE_DIR = os.getenv(
7
+ "AUDIOLDM_CACHE_DIR",
8
+ os.path.join(os.path.expanduser("~"), ".cache/audioldm"))
9
+
10
+ parser = argparse.ArgumentParser()
11
+
12
+ parser.add_argument(
13
+ "--mode",
14
+ type=str,
15
+ required=False,
16
+ default="generation",
17
+ help="generation: text-to-audio generation; transfer: style transfer",
18
+ choices=["generation", "transfer"]
19
+ )
20
+
21
+ parser.add_argument(
22
+ "-t",
23
+ "--text",
24
+ type=str,
25
+ required=False,
26
+ default="",
27
+ help="Text prompt to the model for audio generation",
28
+ )
29
+
30
+ parser.add_argument(
31
+ "-f",
32
+ "--file_path",
33
+ type=str,
34
+ required=False,
35
+ default=None,
36
+ help="(--mode transfer): Original audio file for style transfer; Or (--mode generation): the guidance audio file for generating simialr audio",
37
+ )
38
+
39
+ parser.add_argument(
40
+ "--transfer_strength",
41
+ type=float,
42
+ required=False,
43
+ default=0.5,
44
+ help="A value between 0 and 1. 0 means original audio without transfer, 1 means completely transfer to the audio indicated by text",
45
+ )
46
+
47
+ parser.add_argument(
48
+ "-s",
49
+ "--save_path",
50
+ type=str,
51
+ required=False,
52
+ help="The path to save model output",
53
+ default="./output",
54
+ )
55
+
56
+ parser.add_argument(
57
+ "--model_name",
58
+ type=str,
59
+ required=False,
60
+ help="The checkpoint you gonna use",
61
+ default="audioldm-s-full",
62
+ choices=["audioldm-s-full", "audioldm-l-full", "audioldm-s-full-v2"]
63
+ )
64
+
65
+ parser.add_argument(
66
+ "-ckpt",
67
+ "--ckpt_path",
68
+ type=str,
69
+ required=False,
70
+ help="The path to the pretrained .ckpt model",
71
+ default=None,
72
+ )
73
+
74
+ parser.add_argument(
75
+ "-b",
76
+ "--batchsize",
77
+ type=int,
78
+ required=False,
79
+ default=1,
80
+ help="Generate how many samples at the same time",
81
+ )
82
+
83
+ parser.add_argument(
84
+ "--ddim_steps",
85
+ type=int,
86
+ required=False,
87
+ default=200,
88
+ help="The sampling step for DDIM",
89
+ )
90
+
91
+ parser.add_argument(
92
+ "-gs",
93
+ "--guidance_scale",
94
+ type=float,
95
+ required=False,
96
+ default=2.5,
97
+ help="Guidance scale (Large => better quality and relavancy to text; Small => better diversity)",
98
+ )
99
+
100
+ parser.add_argument(
101
+ "-dur",
102
+ "--duration",
103
+ type=float,
104
+ required=False,
105
+ default=10.0,
106
+ help="The duration of the samples",
107
+ )
108
+
109
+ parser.add_argument(
110
+ "-n",
111
+ "--n_candidate_gen_per_text",
112
+ type=int,
113
+ required=False,
114
+ default=3,
115
+ help="Automatic quality control. This number control the number of candidates (e.g., generate three audios and choose the best to show you). A Larger value usually lead to better quality with heavier computation",
116
+ )
117
+
118
+ parser.add_argument(
119
+ "--seed",
120
+ type=int,
121
+ required=False,
122
+ default=42,
123
+ help="Change this value (any integer number) will lead to a different generation result.",
124
+ )
125
+
126
+ args = parser.parse_args()
127
+
128
+ if(args.ckpt_path is not None):
129
+ print("Warning: ckpt_path has no effect after version 0.0.20.")
130
+
131
+ assert args.duration % 2.5 == 0, "Duration must be a multiple of 2.5"
132
+
133
+ mode = args.mode
134
+ if(mode == "generation" and args.file_path is not None):
135
+ mode = "generation_audio_to_audio"
136
+ if(len(args.text) > 0):
137
+ print("Warning: You have specified the --file_path. --text will be ignored")
138
+ args.text = ""
139
+
140
+ save_path = os.path.join(args.save_path, mode)
141
+
142
+ if(args.file_path is not None):
143
+ save_path = os.path.join(save_path, os.path.basename(args.file_path.split(".")[0]))
144
+
145
+ text = args.text
146
+ random_seed = args.seed
147
+ duration = args.duration
148
+ guidance_scale = args.guidance_scale
149
+ n_candidate_gen_per_text = args.n_candidate_gen_per_text
150
+
151
+ os.makedirs(save_path, exist_ok=True)
152
+ audioldm = build_model(model_name=args.model_name)
153
+
154
+ if(args.mode == "generation"):
155
+ waveform = text_to_audio(
156
+ audioldm,
157
+ text,
158
+ args.file_path,
159
+ random_seed,
160
+ duration=duration,
161
+ guidance_scale=guidance_scale,
162
+ ddim_steps=args.ddim_steps,
163
+ n_candidate_gen_per_text=n_candidate_gen_per_text,
164
+ batchsize=args.batchsize,
165
+ )
166
+
167
+ elif(args.mode == "transfer"):
168
+ assert args.file_path is not None
169
+ assert os.path.exists(args.file_path), "The original audio file \'%s\' for style transfer does not exist." % args.file_path
170
+ waveform = style_transfer(
171
+ audioldm,
172
+ text,
173
+ args.file_path,
174
+ args.transfer_strength,
175
+ random_seed,
176
+ duration=duration,
177
+ guidance_scale=guidance_scale,
178
+ ddim_steps=args.ddim_steps,
179
+ batchsize=args.batchsize,
180
+ )
181
+ waveform = waveform[:,None,:]
182
+
183
+ save_wave(waveform, save_path, name="%s_%s" % (get_time(), text))
audioldm/__pycache__/__init__.cpython-311.pyc ADDED
Binary file (387 Bytes). View file
 
audioldm/__pycache__/__init__.cpython-37.pyc ADDED
Binary file (299 Bytes). View file
 
audioldm/__pycache__/__init__.cpython-39.pyc ADDED
Binary file (311 Bytes). View file
 
audioldm/__pycache__/ldm.cpython-311.pyc ADDED
Binary file (31.4 kB). View file
 
audioldm/__pycache__/ldm.cpython-37.pyc ADDED
Binary file (16 kB). View file
 
audioldm/__pycache__/ldm.cpython-39.pyc ADDED
Binary file (16 kB). View file
 
audioldm/__pycache__/pipeline.cpython-311.pyc ADDED
Binary file (13.1 kB). View file
 
audioldm/__pycache__/pipeline.cpython-37.pyc ADDED
Binary file (6.41 kB). View file
 
audioldm/__pycache__/pipeline.cpython-39.pyc ADDED
Binary file (6.53 kB). View file
 
audioldm/__pycache__/utils.cpython-311.pyc ADDED
Binary file (13.6 kB). View file
 
audioldm/__pycache__/utils.cpython-37.pyc ADDED
Binary file (7.66 kB). View file
 
audioldm/__pycache__/utils.cpython-39.pyc ADDED
Binary file (7.88 kB). View file
 
audioldm/audio/__init__.py ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ from .tools import wav_to_fbank, read_wav_file
2
+ from .stft import TacotronSTFT
audioldm/audio/__pycache__/__init__.cpython-311.pyc ADDED
Binary file (303 Bytes). View file
 
audioldm/audio/__pycache__/__init__.cpython-37.pyc ADDED
Binary file (237 Bytes). View file
 
audioldm/audio/__pycache__/__init__.cpython-39.pyc ADDED
Binary file (249 Bytes). View file
 
audioldm/audio/__pycache__/audio_processing.cpython-311.pyc ADDED
Binary file (4.22 kB). View file
 
audioldm/audio/__pycache__/audio_processing.cpython-37.pyc ADDED
Binary file (2.75 kB). View file
 
audioldm/audio/__pycache__/audio_processing.cpython-39.pyc ADDED
Binary file (2.77 kB). View file
 
audioldm/audio/__pycache__/mix.cpython-39.pyc ADDED
Binary file (1.7 kB). View file
 
audioldm/audio/__pycache__/stft.cpython-311.pyc ADDED
Binary file (10 kB). View file
 
audioldm/audio/__pycache__/stft.cpython-37.pyc ADDED
Binary file (4.93 kB). View file
 
audioldm/audio/__pycache__/stft.cpython-39.pyc ADDED
Binary file (4.97 kB). View file
 
audioldm/audio/__pycache__/tools.cpython-311.pyc ADDED
Binary file (4.45 kB). View file
 
audioldm/audio/__pycache__/tools.cpython-37.pyc ADDED
Binary file (2.17 kB). View file
 
audioldm/audio/__pycache__/tools.cpython-39.pyc ADDED
Binary file (2.18 kB). View file
 
audioldm/audio/__pycache__/torch_tools.cpython-39.pyc ADDED
Binary file (3.79 kB). View file
 
audioldm/audio/audio_processing.py ADDED
@@ -0,0 +1,100 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ import librosa.util as librosa_util
4
+ from scipy.signal import get_window
5
+
6
+
7
+ def window_sumsquare(
8
+ window,
9
+ n_frames,
10
+ hop_length,
11
+ win_length,
12
+ n_fft,
13
+ dtype=np.float32,
14
+ norm=None,
15
+ ):
16
+ """
17
+ # from librosa 0.6
18
+ Compute the sum-square envelope of a window function at a given hop length.
19
+
20
+ This is used to estimate modulation effects induced by windowing
21
+ observations in short-time fourier transforms.
22
+
23
+ Parameters
24
+ ----------
25
+ window : string, tuple, number, callable, or list-like
26
+ Window specification, as in `get_window`
27
+
28
+ n_frames : int > 0
29
+ The number of analysis frames
30
+
31
+ hop_length : int > 0
32
+ The number of samples to advance between frames
33
+
34
+ win_length : [optional]
35
+ The length of the window function. By default, this matches `n_fft`.
36
+
37
+ n_fft : int > 0
38
+ The length of each analysis frame.
39
+
40
+ dtype : np.dtype
41
+ The data type of the output
42
+
43
+ Returns
44
+ -------
45
+ wss : np.ndarray, shape=`(n_fft + hop_length * (n_frames - 1))`
46
+ The sum-squared envelope of the window function
47
+ """
48
+ if win_length is None:
49
+ win_length = n_fft
50
+
51
+ n = n_fft + hop_length * (n_frames - 1)
52
+ x = np.zeros(n, dtype=dtype)
53
+
54
+ # Compute the squared window at the desired length
55
+ win_sq = get_window(window, win_length, fftbins=True)
56
+ win_sq = librosa_util.normalize(win_sq, norm=norm) ** 2
57
+ win_sq = librosa_util.pad_center(win_sq, n_fft)
58
+
59
+ # Fill the envelope
60
+ for i in range(n_frames):
61
+ sample = i * hop_length
62
+ x[sample : min(n, sample + n_fft)] += win_sq[: max(0, min(n_fft, n - sample))]
63
+ return x
64
+
65
+
66
+ def griffin_lim(magnitudes, stft_fn, n_iters=30):
67
+ """
68
+ PARAMS
69
+ ------
70
+ magnitudes: spectrogram magnitudes
71
+ stft_fn: STFT class with transform (STFT) and inverse (ISTFT) methods
72
+ """
73
+
74
+ angles = np.angle(np.exp(2j * np.pi * np.random.rand(*magnitudes.size())))
75
+ angles = angles.astype(np.float32)
76
+ angles = torch.autograd.Variable(torch.from_numpy(angles))
77
+ signal = stft_fn.inverse(magnitudes, angles).squeeze(1)
78
+
79
+ for i in range(n_iters):
80
+ _, angles = stft_fn.transform(signal)
81
+ signal = stft_fn.inverse(magnitudes, angles).squeeze(1)
82
+ return signal
83
+
84
+
85
+ def dynamic_range_compression(x, normalize_fun=torch.log, C=1, clip_val=1e-5):
86
+ """
87
+ PARAMS
88
+ ------
89
+ C: compression factor
90
+ """
91
+ return normalize_fun(torch.clamp(x, min=clip_val) * C)
92
+
93
+
94
+ def dynamic_range_decompression(x, C=1):
95
+ """
96
+ PARAMS
97
+ ------
98
+ C: compression factor used to compress
99
+ """
100
+ return torch.exp(x) / C
audioldm/audio/stft.py ADDED
@@ -0,0 +1,186 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn.functional as F
3
+ import numpy as np
4
+ from scipy.signal import get_window
5
+ from librosa.util import pad_center, tiny
6
+ from librosa.filters import mel as librosa_mel_fn
7
+
8
+ from audioldm.audio.audio_processing import (
9
+ dynamic_range_compression,
10
+ dynamic_range_decompression,
11
+ window_sumsquare,
12
+ )
13
+
14
+
15
+ class STFT(torch.nn.Module):
16
+ """adapted from Prem Seetharaman's https://github.com/pseeth/pytorch-stft"""
17
+
18
+ def __init__(self, filter_length, hop_length, win_length, window="hann"):
19
+ super(STFT, self).__init__()
20
+ self.filter_length = filter_length
21
+ self.hop_length = hop_length
22
+ self.win_length = win_length
23
+ self.window = window
24
+ self.forward_transform = None
25
+ scale = self.filter_length / self.hop_length
26
+ fourier_basis = np.fft.fft(np.eye(self.filter_length))
27
+
28
+ cutoff = int((self.filter_length / 2 + 1))
29
+ fourier_basis = np.vstack(
30
+ [np.real(fourier_basis[:cutoff, :]), np.imag(fourier_basis[:cutoff, :])]
31
+ )
32
+
33
+ forward_basis = torch.FloatTensor(fourier_basis[:, None, :])
34
+ inverse_basis = torch.FloatTensor(
35
+ np.linalg.pinv(scale * fourier_basis).T[:, None, :]
36
+ )
37
+
38
+ if window is not None:
39
+ assert filter_length >= win_length
40
+ # get window and zero center pad it to filter_length
41
+ fft_window = get_window(window, win_length, fftbins=True)
42
+ fft_window = pad_center(fft_window, filter_length)
43
+ fft_window = torch.from_numpy(fft_window).float()
44
+
45
+ # window the bases
46
+ forward_basis *= fft_window
47
+ inverse_basis *= fft_window
48
+
49
+ self.register_buffer("forward_basis", forward_basis.float())
50
+ self.register_buffer("inverse_basis", inverse_basis.float())
51
+
52
+ def transform(self, input_data):
53
+ device = self.forward_basis.device
54
+ input_data = input_data.to(device)
55
+
56
+ num_batches = input_data.size(0)
57
+ num_samples = input_data.size(1)
58
+
59
+ self.num_samples = num_samples
60
+
61
+ # similar to librosa, reflect-pad the input
62
+ input_data = input_data.view(num_batches, 1, num_samples)
63
+ input_data = F.pad(
64
+ input_data.unsqueeze(1),
65
+ (int(self.filter_length / 2), int(self.filter_length / 2), 0, 0),
66
+ mode="reflect",
67
+ )
68
+ input_data = input_data.squeeze(1)
69
+
70
+ forward_transform = F.conv1d(
71
+ input_data,
72
+ torch.autograd.Variable(self.forward_basis, requires_grad=False),
73
+ stride=self.hop_length,
74
+ padding=0,
75
+ )#.cpu()
76
+
77
+ cutoff = int((self.filter_length / 2) + 1)
78
+ real_part = forward_transform[:, :cutoff, :]
79
+ imag_part = forward_transform[:, cutoff:, :]
80
+
81
+ magnitude = torch.sqrt(real_part**2 + imag_part**2)
82
+ phase = torch.autograd.Variable(torch.atan2(imag_part.data, real_part.data))
83
+
84
+ return magnitude, phase
85
+
86
+ def inverse(self, magnitude, phase):
87
+ device = self.forward_basis.device
88
+ magnitude, phase = magnitude.to(device), phase.to(device)
89
+
90
+ recombine_magnitude_phase = torch.cat(
91
+ [magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1
92
+ )
93
+
94
+ inverse_transform = F.conv_transpose1d(
95
+ recombine_magnitude_phase,
96
+ torch.autograd.Variable(self.inverse_basis, requires_grad=False),
97
+ stride=self.hop_length,
98
+ padding=0,
99
+ )
100
+
101
+ if self.window is not None:
102
+ window_sum = window_sumsquare(
103
+ self.window,
104
+ magnitude.size(-1),
105
+ hop_length=self.hop_length,
106
+ win_length=self.win_length,
107
+ n_fft=self.filter_length,
108
+ dtype=np.float32,
109
+ )
110
+ # remove modulation effects
111
+ approx_nonzero_indices = torch.from_numpy(
112
+ np.where(window_sum > tiny(window_sum))[0]
113
+ )
114
+ window_sum = torch.autograd.Variable(
115
+ torch.from_numpy(window_sum), requires_grad=False
116
+ )
117
+ window_sum = window_sum
118
+ inverse_transform[:, :, approx_nonzero_indices] /= window_sum[
119
+ approx_nonzero_indices
120
+ ]
121
+
122
+ # scale by hop ratio
123
+ inverse_transform *= float(self.filter_length) / self.hop_length
124
+
125
+ inverse_transform = inverse_transform[:, :, int(self.filter_length / 2) :]
126
+ inverse_transform = inverse_transform[:, :, : -int(self.filter_length / 2) :]
127
+
128
+ return inverse_transform
129
+
130
+ def forward(self, input_data):
131
+ self.magnitude, self.phase = self.transform(input_data)
132
+ reconstruction = self.inverse(self.magnitude, self.phase)
133
+ return reconstruction
134
+
135
+
136
+ class TacotronSTFT(torch.nn.Module):
137
+ def __init__(
138
+ self,
139
+ filter_length,
140
+ hop_length,
141
+ win_length,
142
+ n_mel_channels,
143
+ sampling_rate,
144
+ mel_fmin,
145
+ mel_fmax,
146
+ ):
147
+ super(TacotronSTFT, self).__init__()
148
+ self.n_mel_channels = n_mel_channels
149
+ self.sampling_rate = sampling_rate
150
+ self.stft_fn = STFT(filter_length, hop_length, win_length)
151
+ mel_basis = librosa_mel_fn(
152
+ sampling_rate, filter_length, n_mel_channels, mel_fmin, mel_fmax
153
+ )
154
+ mel_basis = torch.from_numpy(mel_basis).float()
155
+ self.register_buffer("mel_basis", mel_basis)
156
+
157
+ def spectral_normalize(self, magnitudes, normalize_fun):
158
+ output = dynamic_range_compression(magnitudes, normalize_fun)
159
+ return output
160
+
161
+ def spectral_de_normalize(self, magnitudes):
162
+ output = dynamic_range_decompression(magnitudes)
163
+ return output
164
+
165
+ def mel_spectrogram(self, y, normalize_fun=torch.log):
166
+ """Computes mel-spectrograms from a batch of waves
167
+ PARAMS
168
+ ------
169
+ y: Variable(torch.FloatTensor) with shape (B, T) in range [-1, 1]
170
+
171
+ RETURNS
172
+ -------
173
+ mel_output: torch.FloatTensor of shape (B, n_mel_channels, T)
174
+ """
175
+ assert torch.min(y.data) >= -1, torch.min(y.data)
176
+ assert torch.max(y.data) <= 1, torch.max(y.data)
177
+
178
+ magnitudes, phases = self.stft_fn.transform(y)
179
+ magnitudes = magnitudes.data
180
+ mel_output = torch.matmul(self.mel_basis, magnitudes)
181
+ mel_output = self.spectral_normalize(mel_output, normalize_fun)
182
+ energy = torch.norm(magnitudes, dim=1)
183
+
184
+ log_magnitudes = self.spectral_normalize(magnitudes, normalize_fun)
185
+
186
+ return mel_output, log_magnitudes, energy
audioldm/audio/tools.py ADDED
@@ -0,0 +1,85 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ import torchaudio
4
+
5
+
6
+ def get_mel_from_wav(audio, _stft):
7
+ audio = torch.clip(torch.FloatTensor(audio).unsqueeze(0), -1, 1)
8
+ audio = torch.autograd.Variable(audio, requires_grad=False)
9
+ melspec, log_magnitudes_stft, energy = _stft.mel_spectrogram(audio)
10
+ melspec = torch.squeeze(melspec, 0).numpy().astype(np.float32)
11
+ log_magnitudes_stft = (
12
+ torch.squeeze(log_magnitudes_stft, 0).numpy().astype(np.float32)
13
+ )
14
+ energy = torch.squeeze(energy, 0).numpy().astype(np.float32)
15
+ return melspec, log_magnitudes_stft, energy
16
+
17
+
18
+ def _pad_spec(fbank, target_length=1024):
19
+ n_frames = fbank.shape[0]
20
+ p = target_length - n_frames
21
+ # cut and pad
22
+ if p > 0:
23
+ m = torch.nn.ZeroPad2d((0, 0, 0, p))
24
+ fbank = m(fbank)
25
+ elif p < 0:
26
+ fbank = fbank[0:target_length, :]
27
+
28
+ if fbank.size(-1) % 2 != 0:
29
+ fbank = fbank[..., :-1]
30
+
31
+ return fbank
32
+
33
+
34
+ def pad_wav(waveform, segment_length):
35
+ waveform_length = waveform.shape[-1]
36
+ assert waveform_length > 100, "Waveform is too short, %s" % waveform_length
37
+ if segment_length is None or waveform_length == segment_length:
38
+ return waveform
39
+ elif waveform_length > segment_length:
40
+ return waveform[:segment_length]
41
+ elif waveform_length < segment_length:
42
+ temp_wav = np.zeros((1, segment_length))
43
+ temp_wav[:, :waveform_length] = waveform
44
+ return temp_wav
45
+
46
+ def normalize_wav(waveform):
47
+ waveform = waveform - np.mean(waveform)
48
+ waveform = waveform / (np.max(np.abs(waveform)) + 1e-8)
49
+ return waveform * 0.5
50
+
51
+
52
+ def read_wav_file(filename, segment_length):
53
+ # waveform, sr = librosa.load(filename, sr=None, mono=True) # 4 times slower
54
+ waveform, sr = torchaudio.load(filename) # Faster!!!
55
+ waveform = torchaudio.functional.resample(waveform, orig_freq=sr, new_freq=16000)
56
+ waveform = waveform.numpy()[0, ...]
57
+ waveform = normalize_wav(waveform)
58
+ waveform = waveform[None, ...]
59
+ waveform = pad_wav(waveform, segment_length)
60
+
61
+ waveform = waveform / np.max(np.abs(waveform))
62
+ waveform = 0.5 * waveform
63
+
64
+ return waveform
65
+
66
+
67
+ def wav_to_fbank(filename, target_length=1024, fn_STFT=None):
68
+ assert fn_STFT is not None
69
+
70
+ # mixup
71
+ waveform = read_wav_file(filename, target_length * 160) # hop size is 160
72
+
73
+ waveform = waveform[0, ...]
74
+ waveform = torch.FloatTensor(waveform)
75
+
76
+ fbank, log_magnitudes_stft, energy = get_mel_from_wav(waveform, fn_STFT)
77
+
78
+ fbank = torch.FloatTensor(fbank.T)
79
+ log_magnitudes_stft = torch.FloatTensor(log_magnitudes_stft.T)
80
+
81
+ fbank, log_magnitudes_stft = _pad_spec(fbank, target_length), _pad_spec(
82
+ log_magnitudes_stft, target_length
83
+ )
84
+
85
+ return fbank, log_magnitudes_stft, waveform
audioldm/clap/__init__.py ADDED
File without changes
audioldm/clap/__pycache__/__init__.cpython-39.pyc ADDED
Binary file (152 Bytes). View file
 
audioldm/clap/__pycache__/encoders.cpython-39.pyc ADDED
Binary file (5.1 kB). View file
 
audioldm/clap/encoders.py ADDED
@@ -0,0 +1,170 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ from audioldm.clap.open_clip import create_model
4
+ from audioldm.clap.training.data import get_audio_features
5
+ import torchaudio
6
+ from transformers import RobertaTokenizer
7
+ import torch.nn.functional as F
8
+
9
+
10
+ class CLAPAudioEmbeddingClassifierFreev2(nn.Module):
11
+ def __init__(
12
+ self,
13
+ pretrained_path="",
14
+ key="class",
15
+ sampling_rate=16000,
16
+ embed_mode="audio",
17
+ amodel = "HTSAT-tiny",
18
+ unconditional_prob=0.1,
19
+ random_mute=False,
20
+ max_random_mute_portion=0.5,
21
+ training_mode=True,
22
+ ):
23
+ super().__init__()
24
+
25
+ self.key = key
26
+ self.device = "cpu"
27
+ self.precision = "fp32"
28
+ self.amodel = amodel # or 'PANN-14'
29
+ self.tmodel = "roberta" # the best text encoder in our training
30
+ self.enable_fusion = False # False if you do not want to use the fusion model
31
+ self.fusion_type = "aff_2d"
32
+ self.pretrained = pretrained_path
33
+ self.embed_mode = embed_mode
34
+ self.embed_mode_orig = embed_mode
35
+ self.sampling_rate = sampling_rate
36
+ self.unconditional_prob = unconditional_prob
37
+ self.random_mute = random_mute
38
+ self.tokenize = RobertaTokenizer.from_pretrained("roberta-base")
39
+ self.max_random_mute_portion = max_random_mute_portion
40
+ self.training_mode = training_mode
41
+ self.model, self.model_cfg = create_model(
42
+ self.amodel,
43
+ self.tmodel,
44
+ self.pretrained,
45
+ precision=self.precision,
46
+ device=self.device,
47
+ enable_fusion=self.enable_fusion,
48
+ fusion_type=self.fusion_type,
49
+ )
50
+ for p in self.model.parameters():
51
+ p.requires_grad = False
52
+
53
+ self.model.eval()
54
+
55
+ def get_unconditional_condition(self, batchsize):
56
+ self.unconditional_token = self.model.get_text_embedding(
57
+ self.tokenizer(["", ""])
58
+ )[0:1]
59
+ return torch.cat([self.unconditional_token.unsqueeze(0)] * batchsize, dim=0)
60
+
61
+ def batch_to_list(self, batch):
62
+ ret = []
63
+ for i in range(batch.size(0)):
64
+ ret.append(batch[i])
65
+ return ret
66
+
67
+ def make_decision(self, probability):
68
+ if float(torch.rand(1)) < probability:
69
+ return True
70
+ else:
71
+ return False
72
+
73
+ def random_uniform(self, start, end):
74
+ val = torch.rand(1).item()
75
+ return start + (end - start) * val
76
+
77
+ def _random_mute(self, waveform):
78
+ # waveform: [bs, t-steps]
79
+ t_steps = waveform.size(-1)
80
+ for i in range(waveform.size(0)):
81
+ mute_size = int(
82
+ self.random_uniform(0, end=int(t_steps * self.max_random_mute_portion))
83
+ )
84
+ mute_start = int(self.random_uniform(0, t_steps - mute_size))
85
+ waveform[i, mute_start : mute_start + mute_size] = 0
86
+ return waveform
87
+
88
+ def cos_similarity(self, waveform, text):
89
+ # waveform: [bs, t_steps]
90
+ with torch.no_grad():
91
+ self.embed_mode = "audio"
92
+ audio_emb = self(waveform.cuda())
93
+ self.embed_mode = "text"
94
+ text_emb = self(text)
95
+ similarity = F.cosine_similarity(audio_emb, text_emb, dim=2), audio_emb, text_emb
96
+ return similarity.squeeze()
97
+
98
+ def forward(self, batch, key=None):
99
+ # If you want this conditioner to be unconditional, set self.unconditional_prob = 1.0
100
+ # If you want this conditioner to be fully conditional, set self.unconditional_prob = 0.0
101
+ if self.model.training == True and not self.training_mode:
102
+ print(
103
+ "The pretrained CLAP model should always be in eval mode. Reloading model just in case you change the parameters."
104
+ )
105
+ self.model, self.model_cfg = create_model(
106
+ self.amodel,
107
+ self.tmodel,
108
+ self.pretrained,
109
+ precision=self.precision,
110
+ device="cuda",
111
+ enable_fusion=self.enable_fusion,
112
+ fusion_type=self.fusion_type,
113
+ )
114
+ for p in self.model.parameters():
115
+ p.requires_grad = False
116
+ self.model.eval()
117
+
118
+ # the 'fusion' truncate mode can be changed to 'rand_trunc' if run in unfusion mode
119
+ if self.embed_mode == "audio":
120
+ with torch.no_grad():
121
+ audio_dict_list = []
122
+ assert (
123
+ self.sampling_rate == 16000
124
+ ), "We only support 16000 sampling rate"
125
+ if self.random_mute:
126
+ batch = self._random_mute(batch)
127
+ # batch: [bs, 1, t-samples]
128
+ batch = torchaudio.functional.resample(
129
+ batch, orig_freq=self.sampling_rate, new_freq=48000
130
+ )
131
+ for waveform in self.batch_to_list(batch):
132
+ audio_dict = {}
133
+ audio_dict = get_audio_features(
134
+ audio_dict,
135
+ waveform,
136
+ 480000,
137
+ data_truncating="fusion",
138
+ data_filling="repeatpad",
139
+ audio_cfg=self.model_cfg["audio_cfg"],
140
+ )
141
+ audio_dict_list.append(audio_dict)
142
+ # [bs, 512]
143
+ embed = self.model.get_audio_embedding(audio_dict_list)
144
+ elif self.embed_mode == "text":
145
+ with torch.no_grad():
146
+ # the 'fusion' truncate mode can be changed to 'rand_trunc' if run in unfusion mode
147
+ text_data = self.tokenizer(batch)
148
+ embed = self.model.get_text_embedding(text_data)
149
+
150
+ embed = embed.unsqueeze(1)
151
+ self.unconditional_token = self.model.get_text_embedding(
152
+ self.tokenizer(["", ""])
153
+ )[0:1]
154
+
155
+ for i in range(embed.size(0)):
156
+ if self.make_decision(self.unconditional_prob):
157
+ embed[i] = self.unconditional_token
158
+
159
+ # [bs, 1, 512]
160
+ return embed.detach()
161
+
162
+ def tokenizer(self, text):
163
+ result = self.tokenize(
164
+ text,
165
+ padding="max_length",
166
+ truncation=True,
167
+ max_length=512,
168
+ return_tensors="pt",
169
+ )
170
+ return {k: v.squeeze(0) for k, v in result.items()}
audioldm/clap/open_clip/__init__.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from .factory import (
2
+ list_models,
3
+ create_model,
4
+ create_model_and_transforms,
5
+ add_model_config,
6
+ )
7
+ from .loss import ClipLoss, gather_features, LPLoss, lp_gather_features, LPMetrics
8
+ from .model import (
9
+ CLAP,
10
+ CLAPTextCfg,
11
+ CLAPVisionCfg,
12
+ CLAPAudioCfp,
13
+ convert_weights_to_fp16,
14
+ trace_model,
15
+ )
16
+ from .openai import load_openai_model, list_openai_models
17
+ from .pretrained import (
18
+ list_pretrained,
19
+ list_pretrained_tag_models,
20
+ list_pretrained_model_tags,
21
+ get_pretrained_url,
22
+ download_pretrained,
23
+ )
24
+ from .tokenizer import SimpleTokenizer, tokenize
25
+ from .transform import image_transform
audioldm/clap/open_clip/__pycache__/__init__.cpython-39.pyc ADDED
Binary file (958 Bytes). View file
 
audioldm/clap/open_clip/__pycache__/factory.cpython-39.pyc ADDED
Binary file (6.67 kB). View file
 
audioldm/clap/open_clip/__pycache__/feature_fusion.cpython-39.pyc ADDED
Binary file (4.17 kB). View file
 
audioldm/clap/open_clip/__pycache__/htsat.cpython-39.pyc ADDED
Binary file (30.8 kB). View file
 
audioldm/clap/open_clip/__pycache__/loss.cpython-39.pyc ADDED
Binary file (8.06 kB). View file
 
audioldm/clap/open_clip/__pycache__/model.cpython-39.pyc ADDED
Binary file (23.8 kB). View file
 
audioldm/clap/open_clip/__pycache__/openai.cpython-39.pyc ADDED
Binary file (4.55 kB). View file
 
audioldm/clap/open_clip/__pycache__/pann_model.cpython-39.pyc ADDED
Binary file (13.3 kB). View file