main/library/audioldm2/models.py

import os
import sys
import torch
import librosa

import numpy as np
import torch.nn.functional as F

from scipy.signal import get_window
from librosa.util import pad_center
from diffusers import DDIMScheduler, AudioLDM2Pipeline
from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
from transformers import RobertaTokenizer, RobertaTokenizerFast, VitsTokenizer

sys.path.append(os.getcwd())

from main.configs.config import Config
from main.library.utils import check_audioldm2

config = Config()

class Pipeline(torch.nn.Module):
    def __init__(self, model_id, device, double_precision = False, token = None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model_id = model_id
        self.device = device
        self.double_precision = double_precision
        self.token = token

    def load_scheduler(self):
        pass

    def get_melspectrogram(self):
        pass

    def vae_encode(self, x):
        pass

    def vae_decode(self, x):
        pass

    def decode_to_mel(self, x):
        pass

    def setup_extra_inputs(self, *args, **kwargs):
        pass

    def encode_text(self, prompts, **kwargs):
        pass

    def get_variance(self, timestep, prev_timestep):
        pass

    def get_alpha_prod_t_prev(self, prev_timestep):
        pass

    def get_noise_shape(self, x0, num_steps):
        return (num_steps, self.model.unet.config.in_channels, x0.shape[-2], x0.shape[-1])

    def sample_xts_from_x0(self, x0, num_inference_steps = 50):
        alpha_bar = self.model.scheduler.alphas_cumprod
        sqrt_one_minus_alpha_bar = (1 - alpha_bar) ** 0.5
        timesteps = self.model.scheduler.timesteps.to(self.device)
        t_to_idx = {int(v): k for k, v in enumerate(timesteps)}
        xts = torch.zeros(self.get_noise_shape(x0, num_inference_steps + 1)).to(x0.device)
        xts[0] = x0

        for t in reversed(timesteps):
            idx = num_inference_steps - t_to_idx[int(t)]
            xts[idx] = x0 * (alpha_bar[t] ** 0.5) + torch.randn_like(x0) * sqrt_one_minus_alpha_bar[t]

        return xts

    def get_zs_from_xts(self, xt, xtm1, noise_pred, t, eta = 0, numerical_fix = True, **kwargs):
        alpha_bar = self.model.scheduler.alphas_cumprod

        if self.model.scheduler.config.prediction_type == 'epsilon': pred_original_sample = (xt - (1 - alpha_bar[t]) ** 0.5 * noise_pred) / alpha_bar[t] ** 0.5
        elif self.model.scheduler.config.prediction_type == 'v_prediction': pred_original_sample = (alpha_bar[t] ** 0.5) * xt - ((1 - alpha_bar[t]) ** 0.5) * noise_pred

        prev_timestep = t - self.model.scheduler.config.num_train_timesteps // self.model.scheduler.num_inference_steps
        alpha_prod_t_prev = self.get_alpha_prod_t_prev(prev_timestep)
        variance = self.get_variance(t, prev_timestep)

        if self.model.scheduler.config.prediction_type == 'epsilon': radom_noise_pred = noise_pred
        elif self.model.scheduler.config.prediction_type == 'v_prediction': radom_noise_pred = (alpha_bar[t] ** 0.5) * noise_pred + ((1 - alpha_bar[t]) ** 0.5) * xt

        mu_xt = alpha_prod_t_prev ** (0.5) * pred_original_sample + ((1 - alpha_prod_t_prev - eta * variance) ** (0.5) * radom_noise_pred)
        z = (xtm1 - mu_xt) / (eta * variance ** 0.5)

        if numerical_fix: xtm1 = mu_xt + (eta * variance ** 0.5)*z
        return z, xtm1, None

    def reverse_step_with_custom_noise(self, model_output, timestep, sample, variance_noise = None, eta = 0, **kwargs):
        prev_timestep = timestep - self.model.scheduler.config.num_train_timesteps // self.model.scheduler.num_inference_steps
        alpha_prod_t = self.model.scheduler.alphas_cumprod[timestep]
        alpha_prod_t_prev = self.get_alpha_prod_t_prev(prev_timestep)
        beta_prod_t = 1 - alpha_prod_t

        if self.model.scheduler.config.prediction_type == 'epsilon': pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
        elif self.model.scheduler.config.prediction_type == 'v_prediction': pred_original_sample = (alpha_prod_t ** 0.5) * sample - (beta_prod_t ** 0.5) * model_output

        variance = self.get_variance(timestep, prev_timestep)

        if self.model.scheduler.config.prediction_type == 'epsilon': model_output_direction = model_output
        elif self.model.scheduler.config.prediction_type == 'v_prediction': model_output_direction = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample

        prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + ((1 - alpha_prod_t_prev - eta * variance) ** (0.5) * model_output_direction)

        if eta > 0:
            if variance_noise is None: variance_noise = torch.randn(model_output.shape, device=self.device)
            prev_sample = prev_sample + (eta * variance ** (0.5) * variance_noise)

        return prev_sample

    def unet_forward(self, sample, timestep, encoder_hidden_states, class_labels = None, timestep_cond = None, attention_mask = None, cross_attention_kwargs = None, added_cond_kwargs = None, down_block_additional_residuals = None, mid_block_additional_residual = None, encoder_attention_mask = None, replace_h_space = None, replace_skip_conns = None, return_dict = True, zero_out_resconns = None):
        pass

class STFT(torch.nn.Module):
    def __init__(self, fft_size, hop_size, window_size, window_type="hann"):
        super().__init__()
        self.fft_size = fft_size
        self.hop_size = hop_size
        self.window_size = window_size
        self.window_type = window_type
        
        scale = fft_size / hop_size
        fourier_basis = np.fft.fft(np.eye(fft_size))

        cutoff = fft_size // 2 + 1
        fourier_basis = np.vstack([np.real(fourier_basis[:cutoff, :]), np.imag(fourier_basis[:cutoff, :])])
        
        self.forward_basis = torch.FloatTensor(fourier_basis[:, None, :])
        self.inverse_basis = torch.FloatTensor(np.linalg.pinv(scale * fourier_basis).T[:, None, :])
        
        if window_type:
            assert fft_size >= window_size

            fft_window = torch.from_numpy(pad_center(get_window(window_type, window_size, fftbins=True), size=fft_size)).float()
            self.forward_basis *= fft_window
            self.inverse_basis *= fft_window
        
        if not hasattr(self, "forward_basis"): self.register_buffer("forward_basis", self.forward_basis)
        if not hasattr(self, "inverse_basis"): self.register_buffer("inverse_basis", self.inverse_basis)
    
    def transform(self, signal):
        batch_size, num_samples = signal.shape
        transformed_signal = F.conv1d(F.pad(signal.view(batch_size, 1, num_samples).unsqueeze(1), (self.fft_size // 2, self.fft_size // 2, 0, 0), mode="reflect").squeeze(1), self.forward_basis, stride=self.hop_size, padding=0).cpu()
        
        cutoff = self.fft_size // 2 + 1
        real_part, imag_part = transformed_signal[:, :cutoff, :], transformed_signal[:, cutoff:, :]

        return torch.sqrt(real_part ** 2 + imag_part ** 2), torch.atan2(imag_part, real_part)

class MelSpectrogramProcessor(torch.nn.Module):
    def __init__(self, fft_size, hop_size, window_size, num_mel_bins, sample_rate, fmin, fmax):
        super().__init__()
        self.num_mel_bins = num_mel_bins
        self.sample_rate = sample_rate
        self.stft_processor = STFT(fft_size, hop_size, window_size)
        self.register_buffer("mel_filter", torch.from_numpy(librosa.filters.mel(sr=sample_rate, n_fft=fft_size, n_mels=num_mel_bins, fmin=fmin, fmax=fmax)).float())
    
    def compute_mel_spectrogram(self, waveform, normalization_fn=torch.log):
        assert torch.min(waveform) >= -1
        assert torch.max(waveform) <= 1
        
        magnitudes, _ = self.stft_processor.transform(waveform)
        return normalization_fn(torch.clamp(torch.matmul(self.mel_filter, magnitudes), min=1e-5))

class AudioLDM2(Pipeline):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, local_files_only=True, torch_dtype=torch.float16 if config.is_half else torch.float32).to(self.device)

    def load_scheduler(self):
        self.model.scheduler = DDIMScheduler.from_pretrained(self.model_id, subfolder="scheduler")

    def get_melspectrogram(self):
        return MelSpectrogramProcessor(fft_size=1024, hop_size=160, window_size=1024, num_mel_bins=64, sample_rate=16000, fmin=0, fmax=8000)

    def vae_encode(self, x):
        if x.shape[2] % 4: x = F.pad(x, (0, 0, 4 - (x.shape[2] % 4), 0))
        output = (self.model.vae.encode(x.half() if config.is_half else x.float()).latent_dist.mode() * self.model.vae.config.scaling_factor)
        return output.half() if config.is_half else output.float()

    def vae_decode(self, x):
        return self.model.vae.decode(1 / self.model.vae.config.scaling_factor * x).sample

    def decode_to_mel(self, x):
        tmp = self.model.mel_spectrogram_to_waveform(x[:, 0].detach().to(torch.float16 if config.is_half else torch.float32)).detach()

        if len(tmp.shape) == 1: tmp = tmp.unsqueeze(0)
        return tmp

    def encode_text(self, prompts, negative = False, save_compute = False, cond_length = 0, **kwargs):
        tokenizers, text_encoders = [self.model.tokenizer, self.model.tokenizer_2], [self.model.text_encoder, self.model.text_encoder_2]
        prompt_embeds_list, attention_mask_list = [], []

        for tokenizer, text_encoder in zip(tokenizers, text_encoders):
            text_inputs = tokenizer(prompts, padding="max_length" if (save_compute and negative) or isinstance(tokenizer, (RobertaTokenizer, RobertaTokenizerFast)) else True, max_length=tokenizer.model_max_length if (not save_compute) or ((not negative) or isinstance(tokenizer, (RobertaTokenizer, RobertaTokenizerFast, VitsTokenizer))) else cond_length, truncation=True, return_tensors="pt")
            text_input_ids = text_inputs.input_ids

            attention_mask = text_inputs.attention_mask
            untruncated_ids = tokenizer(prompts, padding="longest", return_tensors="pt").input_ids

            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1: -1])

            text_input_ids = text_input_ids.to(self.device)
            attention_mask = attention_mask.to(self.device)

            with torch.no_grad():
                if text_encoder.config.model_type == "clap":
                    prompt_embeds = text_encoder.get_text_features(text_input_ids, attention_mask=attention_mask)
                    prompt_embeds = prompt_embeds[:, None, :]
                    attention_mask = attention_mask.new_ones((len(prompts), 1))
                else: prompt_embeds = text_encoder(text_input_ids, attention_mask=attention_mask)[0]

            prompt_embeds_list.append(prompt_embeds)
            attention_mask_list.append(attention_mask)

        projection_output = self.model.projection_model(hidden_states=prompt_embeds_list[0], hidden_states_1=prompt_embeds_list[1], attention_mask=attention_mask_list[0], attention_mask_1=attention_mask_list[1])
        generated_prompt_embeds = self.model.generate_language_model(projection_output.hidden_states, attention_mask=projection_output.attention_mask, max_new_tokens=None)
        prompt_embeds = prompt_embeds.to(dtype=self.model.text_encoder_2.dtype, device=self.device)
        return generated_prompt_embeds.to(dtype=self.model.language_model.dtype, device=self.device), prompt_embeds, (attention_mask.to(device=self.device) if attention_mask is not None else torch.ones(prompt_embeds.shape[:2], dtype=torch.long, device=self.device))

    def get_variance(self, timestep, prev_timestep):
        alpha_prod_t = self.model.scheduler.alphas_cumprod[timestep]
        alpha_prod_t_prev = self.get_alpha_prod_t_prev(prev_timestep)
        return ((1 - alpha_prod_t_prev) / (1 - alpha_prod_t)) * (1 - alpha_prod_t / alpha_prod_t_prev)

    def get_alpha_prod_t_prev(self, prev_timestep):
        return self.model.scheduler.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.model.scheduler.final_alpha_cumprod

    def unet_forward(self, sample, timestep, encoder_hidden_states, timestep_cond = None, class_labels = None, attention_mask = None, encoder_attention_mask = None, return_dict = True, cross_attention_kwargs = None, mid_block_additional_residual = None, replace_h_space = None, replace_skip_conns = None, zero_out_resconns = None):
        encoder_hidden_states_1 = class_labels
        class_labels = None
        encoder_attention_mask_1 = encoder_attention_mask
        encoder_attention_mask = None
        default_overall_up_factor = 2 ** self.model.unet.num_upsamplers
        forward_upsample_size = False
        upsample_size = None

        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): forward_upsample_size = True

        if attention_mask is not None:
            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
            attention_mask = attention_mask.unsqueeze(1)

        if encoder_attention_mask is not None:
            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)

        if encoder_attention_mask_1 is not None:
            encoder_attention_mask_1 = (1 - encoder_attention_mask_1.to(sample.dtype)) * -10000.0
            encoder_attention_mask_1 = encoder_attention_mask_1.unsqueeze(1)

        timesteps = timestep
        if not torch.is_tensor(timesteps):
            is_mps = sample.device.type == "mps"

            dtype = (torch.float16 if is_mps else torch.float32) if isinstance(timestep, float) else (torch.int16 if is_mps else torch.int32)

            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
        elif len(timesteps.shape) == 0: timesteps = timesteps[None].to(sample.device)

        emb = self.model.unet.time_embedding(self.model.unet.time_proj(timesteps.expand(sample.shape[0])).to(dtype=sample.dtype), timestep_cond)
        aug_emb = None

        if self.model.unet.class_embedding is not None:
            if class_labels is None: raise ValueError

            if self.model.unet.config.class_embed_type == "timestep": class_labels = self.model.unet.time_proj(class_labels).to(dtype=sample.dtype)
            class_emb = self.model.unet.class_embedding(class_labels).to(dtype=sample.dtype)

            if self.model.unet.config.class_embeddings_concat: emb = torch.cat([emb, class_emb], dim=-1)
            else: emb = emb + class_emb

        emb = emb + aug_emb if aug_emb is not None else emb
        if self.model.unet.time_embed_act is not None: emb = self.model.unet.time_embed_act(emb)

        sample = self.model.unet.conv_in(sample)
        down_block_res_samples = (sample,)

        for downsample_block in self.model.unet.down_blocks:
            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention: sample, res_samples = downsample_block(hidden_states=sample, temb=emb, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask, cross_attention_kwargs=cross_attention_kwargs, encoder_attention_mask=encoder_attention_mask, encoder_hidden_states_1=encoder_hidden_states_1, encoder_attention_mask_1=encoder_attention_mask_1)
            else: sample, res_samples = downsample_block(hidden_states=sample, temb=emb)

            down_block_res_samples += res_samples

        if self.model.unet.mid_block is not None: sample = self.model.unet.mid_block(sample, emb, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask, cross_attention_kwargs=cross_attention_kwargs, encoder_attention_mask=encoder_attention_mask, encoder_hidden_states_1=encoder_hidden_states_1, encoder_attention_mask_1=encoder_attention_mask_1)

        if replace_h_space is None: h_space = sample.clone()
        else:
            h_space = replace_h_space
            sample = replace_h_space.clone()

        if mid_block_additional_residual is not None: sample = sample + mid_block_additional_residual
        extracted_res_conns = {}

        for i, upsample_block in enumerate(self.model.unet.up_blocks):
            is_final_block = i == len(self.model.unet.up_blocks) - 1
            res_samples = down_block_res_samples[-len(upsample_block.resnets):]
            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]

            if replace_skip_conns is not None and replace_skip_conns.get(i): res_samples = replace_skip_conns.get(i)

            if zero_out_resconns is not None:
                if (type(zero_out_resconns) is int and i >= (zero_out_resconns - 1)) or type(zero_out_resconns) is list and i in zero_out_resconns: res_samples = [torch.zeros_like(x) for x in res_samples]

            extracted_res_conns[i] = res_samples
            if not is_final_block and forward_upsample_size: upsample_size = down_block_res_samples[-1].shape[2:]

            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: sample = upsample_block(hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, upsample_size=upsample_size, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, encoder_hidden_states_1=encoder_hidden_states_1, encoder_attention_mask_1=encoder_attention_mask_1)
            else: sample = upsample_block(hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size)

        if self.model.unet.conv_norm_out: sample = self.model.unet.conv_act(self.model.unet.conv_norm_out(sample))
        sample = self.model.unet.conv_out(sample)

        if not return_dict: return (sample,)
        return UNet2DConditionOutput(sample=sample), h_space, extracted_res_conns

def load_model(model, device):
    check_audioldm2(model)

    ldm_stable = AudioLDM2(model_id=os.path.join("assets", "models", "audioldm2", model), device=device, double_precision=False)
    ldm_stable.load_scheduler()

    if torch.cuda.is_available(): torch.cuda.empty_cache()
    elif torch.backends.mps.is_available(): torch.mps.empty_cache()

    return ldm_stable