Update pico_model.py

pico_model.py

@@ -8,40 +8,6 @@ import torch.nn.functional as F
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers import DDPMScheduler, UNet2DConditionModel
 
-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)
@@ -243,7 +209,7 @@ class ClapText_Onset_2_Audio_Diffusion(nn.Module):
 from sklearn.metrics.pairwise import cosine_similarity
 import laion_clap
 from laion_clap.clap_module.factory import load_state_dict as clap_load_state_dict
-
+from llm_preprocess import get_event
 class PicoDiffusion(ClapText_Onset_2_Audio_Diffusion):
     def __init__(self,      
         scheduler_name,
@@ -260,31 +226,12 @@ 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=self.device)
         del diffusion_ckpt["class_emb.weight"]
         ckpt.update(diffusion_ckpt)
         self.load_state_dict(ckpt)
 
-        self.event_list = [
-            "burping_belching",             # 0
-            "car_horn_honking",             #
-            "cat_meowing",                  #    
-            "cow_mooing",                   #
-            "dog_barking",                  #  
-            "door_knocking",                #
-            "door_slamming",                #
-            "explosion",                    #  
-            "gunshot",                      # 8
-            "sheep_goat_bleating",          #
-            "sneeze",                       #
-            "spraying",                     # 
-            "thump_thud",                   #   
-            "train_horn",                   #
-            "tapping_clicking_clanking",    #
-            "woman_laughing",               #         
-            "duck_quacking",                # 16   
-            "whistling",                    #    
-        ]
+        self.event_list = get_event()
         self.events_emb = self.freeze_text_encoder.get_text_embedding(self.event_list, use_tensor=False)
 
         
@@ -300,10 +247,11 @@ class PicoDiffusion(ClapText_Onset_2_Audio_Diffusion):
         for event_timestamp in timestampCaption.split(' and '):
             # event_timestamp : event1__onset1-offset1_onset2-offset2
             (event, instance) = event_timestamp.split(' at ')
-            events.append(event)
+            
             # instance : onset1-offset1_onset2-offset2
             event_emb = self.freeze_text_encoder.get_text_embedding([event, ""], use_tensor=False)[0]
             event_id = np.argmax(cosine_similarity(event_emb.reshape(1, -1), self.events_emb))
+            events.append(self.event_list[event_id])
             for start_end in instance.split('_'):
                 (start, end) = start_end.split('-')         
                 start, end = int(float(start)*250/10), int(float(end)*250/10)