Update pico_model.py

pico_model.py

@@ -12,36 +12,6 @@ from audioldm.audio.stft import TacotronSTFT
 from audioldm.variational_autoencoder.autoencoder import AutoencoderKL
 from audioldm.utils import default_audioldm_config, get_metadata
 
-
-
-def build_pretrained_models(name):
-    checkpoint = torch.load(get_metadata()[name]["path"], map_location="cpu")
-    scale_factor = checkpoint["state_dict"]["scale_factor"].item()
-
-    vae_state_dict = {k[18:]: v for k, v in checkpoint["state_dict"].items() if "first_stage_model." in k}
-
-    config = default_audioldm_config(name)
-    vae_config = config["model"]["params"]["first_stage_config"]["params"]
-    vae_config["scale_factor"] = scale_factor
-
-    vae = AutoencoderKL(**vae_config)
-    vae.load_state_dict(vae_state_dict)
-
-    fn_STFT = TacotronSTFT(
-        config["preprocessing"]["stft"]["filter_length"],
-        config["preprocessing"]["stft"]["hop_length"],
-        config["preprocessing"]["stft"]["win_length"],
-        config["preprocessing"]["mel"]["n_mel_channels"],
-        config["preprocessing"]["audio"]["sampling_rate"],
-        config["preprocessing"]["mel"]["mel_fmin"],
-        config["preprocessing"]["mel"]["mel_fmax"],
-    )
-
-    vae.eval()
-    fn_STFT.eval()
-
-    return vae, fn_STFT
-
 def _init_layer(layer):
     """Initialize a Linear or Convolutional layer. """
     nn.init.xavier_uniform_(layer.weight)
@@ -260,7 +230,7 @@ class PicoDiffusion(ClapText_Onset_2_Audio_Diffusion):
         ckpt = clap_load_state_dict(freeze_text_encoder_ckpt, skip_params=True)
         del_parameter_key = ["text_branch.embeddings.position_ids"]
         ckpt = {f"freeze_text_encoder.model.{k}":v for k, v in ckpt.items() if k not in del_parameter_key}
-        diffusion_ckpt = torch.load(diffusion_pt)
+        diffusion_ckpt = torch.load(diffusion_pt, map_location=torch.device(self.device))
         del diffusion_ckpt["class_emb.weight"]
         ckpt.update(diffusion_ckpt)
         self.load_state_dict(ckpt)