Hybrid.py

import torch
import torch.nn.functional as F
from jukebox.vqvae.vqvae import VQVAE  

from audio_diffusion.models import DiffusionAttnUnet1D
from audio_diffusion.utils import ema_update
from viz.viz import audio_spectrogram_image


# VQ-VAE encoding
def encode_with_vqvae(vqvae, audio_input):
    encoded_audio, *_ = vqvae.encode(audio_input)
    return encoded_audio

# DDIM denoising
def denoise_with_ddim(ddim_model, encoded_audio, timesteps):
    denoised_audio = sample(ddim_model, encoded_audio, steps=timesteps, eta=0)
    return denoised_audio

# VQ-VAE decoding
def decode_with_vqvae(vqvae, denoised_audio):
    decoded_audio = vqvae.decode(denoised_audio)
    return decoded_audio

# Main function to integrate VQ-VAE and DDIM
def generate_audio_with_vqvae_ddim(vqvae, ddim_model, raw_audio_input, timesteps):
    encoded_audio = encode_with_vqvae(vqvae, raw_audio_input)
    denoised_audio = denoise_with_ddim(ddim_model, encoded_audio, timesteps)
    final_audio_output = decode_with_vqvae(vqvae, denoised_audio)
    return final_audio_output

if __name__ == "__main__":
  
    vqvae_model = VQVAE(...)  # Initialize with the appropriate parameters
    vqvae_model.load_state_dict(torch.load("path_to_vqvae_model_weights.pt"))
    vqvae_model.eval()


    ddim_model = DiffusionAttnUnet1D(...)  # Initialize with the appropriate parameters
    ddim_model.load_state_dict(torch.load("path_to_ddim_model_weights.pt"))
    ddim_model.eval()

    # Provide a sample raw audio input and specify timesteps for DDIM
    raw_audio_sample = ...  # Load or generate a sample raw audio input
    timesteps = ...

    # Generate audio using the integrated model
    generated_audio = generate_audio_with_vqvae_ddim(vqvae_model, ddim_model, raw_audio_sample, timesteps)