Update requirements for TANGO

README.md

@@ -63,6 +63,8 @@ More generated samples are shown [here](https://github.com/declare-lab/tango/blo
 
 ## Prerequisites
 
+Our code is built on pytorch version 1.13.1+cu117. We mention `torch==1.13.1` in the requirements file but you might need to install a specific cuda version of torch depending on your GPU device type.
+
 Install `requirements.txt`. You will also need to install the `diffusers` package from the directory provided in this repo:
 
 ```bash

requirements.txt

@@ -10,7 +10,7 @@ huggingface_hub==0.13.3
 importlib_metadata==6.3.0
 librosa==0.9.2
 matplotlib==3.5.2
-numpy==1.22.0
+numpy==1.23.0
 omegaconf==2.3.0
 packaging==23.1
 pandas==1.4.1

tango.py

@@ -19,9 +19,9 @@ def __init__(self, name="declare-lab/tango", device="cuda:0"):
         self.stft = TacotronSTFT(**stft_config).to(device)
         self.model = AudioDiffusion(**main_config).to(device)
 
-        vae_weights = torch.load("{}/pytorch_model_vae.bin".format(path))
-        stft_weights = torch.load("{}/pytorch_model_stft.bin".format(path))
-        main_weights = torch.load("{}/pytorch_model_main.bin".format(path))
+        vae_weights = torch.load("{}/pytorch_model_vae.bin".format(path), map_location=device)
+        stft_weights = torch.load("{}/pytorch_model_stft.bin".format(path), map_location=device)
+        main_weights = torch.load("{}/pytorch_model_main.bin".format(path), map_location=device)
 
         self.vae.load_state_dict(vae_weights)
         self.stft.load_state_dict(stft_weights)

train.py

@@ -307,7 +307,8 @@ def main():
 
     if args.hf_model:
         hf_model_path = snapshot_download(repo_id=args.hf_model)
-        model.load_state_dict(torch.load("{}/pytorch_model_main.bin".format(hf_model_path)))
+        model.load_state_dict(torch.load("{}/pytorch_model_main.bin".format(hf_model_path), map_location="cpu"))
+        accelerator.print("Successfully loaded checkpoint from:", args.hf_model)
 
     if args.prefix:
         prefix = args.prefix