sheikyerbouti/audioEditing

1import numpy as np2import torch3from typing import Optional, List, Tuple, NamedTuple, Union4from models import PipelineWrapper5 6 7class PromptEmbeddings(NamedTuple):8    embedding_hidden_states: torch.Tensor9    embedding_class_lables: torch.Tensor10    boolean_prompt_mask: torch.Tensor11 12 13def load_audio(audio_path: Union[str, np.array], fn_STFT, left: int = 0, right: int = 0, device: Optional[torch.device] = None14               ) -> torch.tensor:15    if type(audio_path) is str:16        import audioldm17        import audioldm.audio18 19        duration = min(audioldm.utils.get_duration(audio_path), 30)20 21        mel, _, _ = audioldm.audio.wav_to_fbank(audio_path, target_length=int(duration * 102.4), fn_STFT=fn_STFT)22        mel = mel.unsqueeze(0)23    else:24        mel = audio_path25 26    c, h, w = mel.shape27    left = min(left, w-1)28    right = min(right, w - left - 1)29    mel = mel[:, :, left:w-right]30    mel = mel.unsqueeze(0).to(device)31 32    return mel33 34 35def get_height_of_spectrogram(length: int, ldm_stable: PipelineWrapper) -> int:36    vocoder_upsample_factor = np.prod(ldm_stable.model.vocoder.config.upsample_rates) / \37        ldm_stable.model.vocoder.config.sampling_rate38 39    if length is None:40        length = ldm_stable.model.unet.config.sample_size * ldm_stable.model.vae_scale_factor * \41            vocoder_upsample_factor42 43    height = int(length / vocoder_upsample_factor)44 45    # original_waveform_length = int(length * ldm_stable.model.vocoder.config.sampling_rate)46    if height % ldm_stable.model.vae_scale_factor != 0:47        height = int(np.ceil(height / ldm_stable.model.vae_scale_factor)) * ldm_stable.model.vae_scale_factor48        print(49            f"Audio length in seconds {length} is increased to {height * vocoder_upsample_factor} "50            f"so that it can be handled by the model. It will be cut to {length} after the "51            f"denoising process."52        )53 54    return height55 56 57def get_text_embeddings(target_prompt: List[str], target_neg_prompt: List[str], ldm_stable: PipelineWrapper58                        ) -> Tuple[torch.Tensor, PromptEmbeddings, PromptEmbeddings]:59    text_embeddings_hidden_states, text_embeddings_class_labels, text_embeddings_boolean_prompt_mask = \60        ldm_stable.encode_text(target_prompt)61    uncond_embedding_hidden_states, uncond_embedding_class_lables, uncond_boolean_prompt_mask = \62        ldm_stable.encode_text(target_neg_prompt)63 64    text_emb = PromptEmbeddings(embedding_hidden_states=text_embeddings_hidden_states,65                                boolean_prompt_mask=text_embeddings_boolean_prompt_mask,66                                embedding_class_lables=text_embeddings_class_labels)67    uncond_emb = PromptEmbeddings(embedding_hidden_states=uncond_embedding_hidden_states,68                                  boolean_prompt_mask=uncond_boolean_prompt_mask,69                                  embedding_class_lables=uncond_embedding_class_lables)70 71    return text_embeddings_class_labels, text_emb, uncond_emb72