Spaces:
Runtime error
Runtime error
File size: 4,217 Bytes
5a9b731 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
"""
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() |