working inference example

lucidrains · Feb 2, 2023 · c40aee3 · c40aee3
1 parent 8a66cce
commit c40aee3
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 5 deletions.
diff --git a/README.md b/README.md
@@ -79,6 +79,17 @@ wavs = torch.randn(2, 1024)
 conds = quantizer(wavs = wavs, namespace = 'semantic') # (2, 8, 1024) - 8 is number of quantizers
 ```
 
+After much training, you will pass your finetuned or trained-from-scratch `AudioLM` and `MuLaN` wrapped in `MuLaNEmbedQuantizer` to the `MusicLM`
+
+```python
+musiclm = MusicLM(
+    audio_lm = audio_lm,
+    mulan_embed_quantizer = mulan_embed_quantizer
+)
+
+music = musiclm(['the crystalline sounds of the piano in a ballroom']) # torch.Tensor
+```
+
 ## Todo
 
 - [x] mulan seems to be using decoupled contrastive learning, offer that as an option

diff --git a/musiclm_pytorch/musiclm_pytorch.py b/musiclm_pytorch/musiclm_pytorch.py
@@ -549,7 +549,7 @@ def forward(
         self.eval()
 
         texts = tokenizer.tokenize(raw_texts)
-        cond_tokens = self.mulan_embed_quantizer(texts = texts)
 
-        wavs = self.audio_lm.generate(cond_tokens = cond_tokens, **audio_lm_kwargs)
-        return wavs
+        text_embeds = self.mulan_embed_quantizer(texts = texts)
+
+        return self.audio_lm(text_embeds = text_embeds, **audio_lm_kwargs)
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'musiclm-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.3',
+  version = '0.0.4',
   license='MIT',
   description = 'MusicLM - AudioLM + Audio CLIP to text to music synthesis',
   author = 'Phil Wang',
@@ -19,7 +19,7 @@
     'contrastive learning'
   ],
   install_requires=[
-    'audiolm-pytorch',
+    'audiolm-pytorch>=0.8.1',
     'beartype',
     'einops>=0.4',
     'vector-quantize-pytorch>=1.0.0',