Spaces:
Runtime error
Runtime error
| """ | |
| test_conditioning.py | |
| Desc: In order to make sure that we can condition on phoneme encodings, need to ensure that | |
| the model can take in 2D Embeddings | |
| """ | |
| import torch | |
| from audio_diffusion_pytorch import DiffusionModel, UNetV0, VDiffusion, VSampler | |
| import sys | |
| sys.path.append('.') | |
| # Example build model function | |
| def create_model(): | |
| return DiffusionModel( | |
| net_t=UNetV0, # The model type used for diffusion | |
| in_channels=2, # U-Net: number of input/output (audio) channels | |
| channels=[8, 32, 64, 128, 256, 512, 512, 1024, 1024], # U-Net: channels at each layer | |
| factors=[1, 4, 4, 4, 2, 2, 2, 2, 2], # U-Net: downsampling and upsampling factors at each layer | |
| items=[1, 2, 2, 2, 2, 2, 2, 4, 4], # U-Net: number of repeating items at each layer | |
| attentions=[0, 0, 0, 0, 0, 1, 1, 1, 1], # U-Net: attention enabled/disabled at each layer | |
| attention_heads=8, # U-Net: number of attention heads per attention block | |
| attention_features=64, # U-Net: number of attention features per attention block, | |
| diffusion_t=VDiffusion, # The diffusion method used | |
| sampler_t=VSampler, # The diffusion sampler used | |
| ) | |
| def create_cond_model(): | |
| return DiffusionModel( | |
| net_t=UNetV0, # The model type used for diffusion | |
| in_channels=2, # U-Net: number of input/output (audio) channels | |
| channels=[8, 32, 64, 128, 256, 512, 512, 1024, 1024], # U-Net: channels at each layer | |
| factors=[1, 4, 4, 4, 2, 2, 2, 2, 2], # U-Net: downsampling and upsampling factors at each layer | |
| items=[1, 2, 2, 2, 2, 2, 2, 4, 4], # U-Net: number of repeating items at each layer | |
| attentions=[0, 0, 0, 0, 0, 1, 1, 1, 1], # U-Net: attention enabled/disabled at each layer | |
| attention_heads=8, # U-Net: number of attention heads per attention block | |
| attention_features=64, # U-Net: number of attention features per attention block, | |
| diffusion_t=VDiffusion, # The diffusion method used | |
| sampler_t=VSampler, # The diffusion sampler used | |
| embedding_features=768, # U-Net: embedding features | |
| cross_attentions=[0, 0, 0, 1, 1, 1, 1, 1, 1], # U-Net: cross-attention enabled/disabled at each layer | |
| ) | |
| def create_mel_model(): | |
| return DiffusionModel( | |
| net_t=UNetV0, # The model type used for diffusion | |
| in_channels=64, # U-Net: number of input/output (audio) channels | |
| channels=[8, 32, 64, 128, 256, 512, 512, 1024, 1024], # U-Net: channels at each layer | |
| factors=[1, 4, 4, 4, 2, 2, 2, 2, 2], # U-Net: downsampling and upsampling factors at each layer | |
| items=[1, 2, 2, 2, 2, 2, 2, 4, 4], # U-Net: number of repeating items at each layer | |
| attentions=[0, 0, 0, 0, 0, 1, 1, 1, 1], # U-Net: attention enabled/disabled at each layer | |
| attention_heads=8, # U-Net: number of attention heads per attention block | |
| attention_features=64, # U-Net: number of attention features per attention block, | |
| diffusion_t=VDiffusion, # The diffusion method used | |
| sampler_t=VSampler, # The diffusion sampler used | |
| embedding_features=768, # U-Net: embedding features | |
| cross_attentions=[0, 0, 0, 1, 1, 1, 1, 1, 1], # U-Net: cross-attention enabled/disabled at each layer | |
| ) | |
| # A file function for conditioning 2D Conditioning with random values | |
| def test_2D_condition(): | |
| model = create_model() | |
| cond_model = create_cond_model() | |
| # Pretend that input is waveform | |
| audio_wave = torch.randn(1, 2, 2**18) # [batch, in_channels, length] | |
| embedding = torch.randn(1, 1, 768) # [batch, num_embeddings, embedding_features] | |
| y_nocond = model(audio_wave) | |
| y = cond_model(audio_wave, embedding=embedding) | |
| # Lets Try multiple Embeddings | |
| embedding = torch.randn(1, 20, 768) # [batch, num_embeddings, embedding_features] | |
| y = cond_model(audio_wave, embedding=embedding) | |
| del model | |
| del cond_model | |
| mel_model = create_mel_model() | |
| # Pretend that input is mel-spec | |
| audio_wave = torch.randn(1, 64, 2**18) # [batch, in_channels, length] | |
| embedding = torch.randn(1, 64, 768) # [batch, num_embeddings, embedding_features]; should be [batch, length, phone_embed_dim] | |
| y = mel_model(audio_wave, embedding = embedding) | |
| if __name__ == "__main__": | |
| # Note: The below function works as is | |
| test_2D_condition() |