models.py

import torch
from diffusers import DDIMScheduler, CosineDPMSolverMultistepScheduler
from diffusers.schedulers.scheduling_dpmsolver_sde import BrownianTreeNoiseSampler
from diffusers import AudioLDM2Pipeline, StableAudioPipeline
from transformers import RobertaTokenizer, RobertaTokenizerFast, VitsTokenizer
from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
from diffusers.models.embeddings import get_1d_rotary_pos_embed
from typing import Any, Dict, List, Optional, Tuple, Union
import gradio as gr


class PipelineWrapper(torch.nn.Module):
    def __init__(self, model_id: str,
                 device: torch.device,
                 double_precision: bool = False,
                 token: Optional[str] = None, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.model_id = model_id
        self.device = device
        self.double_precision = double_precision
        self.token = token

    def get_sigma(self, timestep: int) -> float:
        sqrt_recipm1_alphas_cumprod = torch.sqrt(1.0 / self.model.scheduler.alphas_cumprod - 1)
        return sqrt_recipm1_alphas_cumprod[timestep]

    def load_scheduler(self) -> None:
        pass

    def get_fn_STFT(self) -> torch.nn.Module:
        pass

    def get_sr(self) -> int:
        return 16000

    def vae_encode(self, x: torch.Tensor) -> torch.Tensor:
        pass

    def vae_decode(self, x: torch.Tensor) -> torch.Tensor:
        pass

    def decode_to_mel(self, x: torch.Tensor) -> torch.Tensor:
        pass

    def setup_extra_inputs(self, *args, **kwargs) -> None:
        pass

    def encode_text(self, prompts: List[str], **kwargs
                    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
        pass

    def get_variance(self, timestep: torch.Tensor, prev_timestep: torch.Tensor) -> torch.Tensor:
        pass

    def get_alpha_prod_t_prev(self, prev_timestep: torch.Tensor) -> torch.Tensor:
        pass

    def get_noise_shape(self, x0: torch.Tensor, num_steps: int) -> Tuple[int, ...]:
        variance_noise_shape = (num_steps,
                                self.model.unet.config.in_channels,
                                x0.shape[-2],
                                x0.shape[-1])
        return variance_noise_shape

    def sample_xts_from_x0(self, x0: torch.Tensor, num_inference_steps: int = 50) -> torch.Tensor:
        """
        Samples from P(x_1:T|x_0)
        """
        alpha_bar = self.model.scheduler.alphas_cumprod
        sqrt_one_minus_alpha_bar = (1-alpha_bar) ** 0.5

        variance_noise_shape = self.get_noise_shape(x0, num_inference_steps + 1)
        timesteps = self.model.scheduler.timesteps.to(self.device)
        t_to_idx = {int(v): k for k, v in enumerate(timesteps)}
        xts = torch.zeros(variance_noise_shape).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: torch.Tensor, xtm1: torch.Tensor, noise_pred: torch.Tensor,
                        t: torch.Tensor, eta: float = 0, numerical_fix: bool = True, **kwargs
                        ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        # pred of x0
        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

        # direction to xt
        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

        pred_sample_direction = (1 - alpha_prod_t_prev - eta * variance) ** (0.5) * radom_noise_pred

        mu_xt = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction

        z = (xtm1 - mu_xt) / (eta * variance ** 0.5)

        # correction to avoid error accumulation
        if numerical_fix:
            xtm1 = mu_xt + (eta * variance ** 0.5)*z

        return z, xtm1, None

    def reverse_step_with_custom_noise(self, model_output: torch.Tensor, timestep: torch.Tensor, sample: torch.Tensor,
                                       variance_noise: Optional[torch.Tensor] = None, eta: float = 0, **kwargs
                                       ) -> torch.Tensor:
        # 1. get previous step value (=t-1)
        prev_timestep = timestep - self.model.scheduler.config.num_train_timesteps // \
            self.model.scheduler.num_inference_steps
        # 2. compute alphas, betas
        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
        # 3. compute predicted original sample from predicted noise also called
        # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
        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

        # 5. compute variance: "sigma_t(η)" -> see formula (16)
        # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
        # variance = self.scheduler._get_variance(timestep, prev_timestep)
        variance = self.get_variance(timestep, prev_timestep)
        # std_dev_t = eta * variance ** (0.5)
        # Take care of asymetric reverse process (asyrp)
        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
        # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
        # pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * model_output_direction
        pred_sample_direction = (1 - alpha_prod_t_prev - eta * variance) ** (0.5) * model_output_direction
        # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
        prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction
        # 8. Add noice if eta > 0
        if eta > 0:
            if variance_noise is None:
                variance_noise = torch.randn(model_output.shape, device=self.device)
            sigma_z = eta * variance ** (0.5) * variance_noise
            prev_sample = prev_sample + sigma_z

        return prev_sample

    def unet_forward(self,
                     sample: torch.FloatTensor,
                     timestep: Union[torch.Tensor, float, int],
                     encoder_hidden_states: torch.Tensor,
                     class_labels: Optional[torch.Tensor] = None,
                     timestep_cond: Optional[torch.Tensor] = None,
                     attention_mask: Optional[torch.Tensor] = None,
                     cross_attention_kwargs: Optional[Dict[str, Any]] = None,
                     added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
                     down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
                     mid_block_additional_residual: Optional[torch.Tensor] = None,
                     encoder_attention_mask: Optional[torch.Tensor] = None,
                     replace_h_space: Optional[torch.Tensor] = None,
                     replace_skip_conns: Optional[Dict[int, torch.Tensor]] = None,
                     return_dict: bool = True,
                     zero_out_resconns: Optional[Union[int, List]] = None) -> Tuple:
        pass


class AudioLDM2Wrapper(PipelineWrapper):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        if self.double_precision:
            self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, torch_dtype=torch.float64, token=self.token
                                                           ).to(self.device)
        else:
            try:
                self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, local_files_only=True, token=self.token
                                                               ).to(self.device)
            except FileNotFoundError:
                self.model = AudioLDM2Pipeline.from_pretrained(self.model_id, local_files_only=False, token=self.token
                                                               ).to(self.device)

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

    def get_fn_STFT(self) -> torch.nn.Module:
        from audioldm.audio import TacotronSTFT
        return TacotronSTFT(
            filter_length=1024,
            hop_length=160,
            win_length=1024,
            n_mel_channels=64,
            sampling_rate=16000,
            mel_fmin=0,
            mel_fmax=8000,
        )

    def vae_encode(self, x: torch.Tensor) -> torch.Tensor:
        # self.model.vae.disable_tiling()
        if x.shape[2] % 4:
            x = torch.nn.functional.pad(x, (0, 0, 4 - (x.shape[2] % 4), 0))
        return (self.model.vae.encode(x).latent_dist.mode() * self.model.vae.config.scaling_factor).float()
        # return (self.encode_no_tiling(x).latent_dist.mode() * self.model.vae.config.scaling_factor).float()

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

    def decode_to_mel(self, x: torch.Tensor) -> torch.Tensor:
        if self.double_precision:
            tmp = self.model.mel_spectrogram_to_waveform(x[:, 0].detach().double()).detach()
        tmp = self.model.mel_spectrogram_to_waveform(x[:, 0].detach().float()).detach()
        if len(tmp.shape) == 1:
            tmp = tmp.unsqueeze(0)
        return tmp

    def encode_text(self, prompts: List[str], negative: bool = False, 
                    save_compute: bool = False, cond_length: int = 0, **kwargs
                    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        tokenizers = [self.model.tokenizer, self.model.tokenizer_2]
        text_encoders = [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):
                removed_text = tokenizer.batch_decode(
                    untruncated_ids[:, tokenizer.model_max_length - 1: -1])
                print(f"The following part of your input was truncated because {text_encoder.config.model_type} can "
                      f"only handle sequences up to {tokenizer.model_max_length} tokens: {removed_text}"
                      )

            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,
                    )
                    # append the seq-len dim: (bs, hidden_size) -> (bs, seq_len, hidden_size)
                    prompt_embeds = prompt_embeds[:, None, :]
                    # make sure that we attend to this single hidden-state
                    attention_mask = attention_mask.new_ones((len(prompts), 1))
                else:
                    prompt_embeds = text_encoder(
                        text_input_ids,
                        attention_mask=attention_mask,
                    )
                    prompt_embeds = prompt_embeds[0]

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

        # print(f'prompt[0].shape: {prompt_embeds_list[0].shape}')
        # print(f'prompt[1].shape: {prompt_embeds_list[1].shape}')
        # print(f'attn[0].shape: {attention_mask_list[0].shape}')
        # print(f'attn[1].shape: {attention_mask_list[1].shape}')

        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],
        )
        projected_prompt_embeds = projection_output.hidden_states
        projected_attention_mask = projection_output.attention_mask

        generated_prompt_embeds = self.model.generate_language_model(
            projected_prompt_embeds,
            attention_mask=projected_attention_mask,
            max_new_tokens=None,
        )

        prompt_embeds = prompt_embeds.to(dtype=self.model.text_encoder_2.dtype, device=self.device)
        attention_mask = (
            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)
        )
        generated_prompt_embeds = generated_prompt_embeds.to(dtype=self.model.language_model.dtype, device=self.device)

        return generated_prompt_embeds, prompt_embeds, attention_mask

    def get_variance(self, timestep: torch.Tensor, prev_timestep: torch.Tensor) -> torch.Tensor:
        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
        beta_prod_t_prev = 1 - alpha_prod_t_prev
        variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
        return variance

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

    def unet_forward(self,
                     sample: torch.FloatTensor,
                     timestep: Union[torch.Tensor, float, int],
                     encoder_hidden_states: torch.Tensor,
                     timestep_cond: Optional[torch.Tensor] = None,
                     class_labels: Optional[torch.Tensor] = None,
                     attention_mask: Optional[torch.Tensor] = None,
                     encoder_attention_mask: Optional[torch.Tensor] = None,
                     return_dict: bool = True,
                     cross_attention_kwargs: Optional[Dict[str, Any]] = None,
                     mid_block_additional_residual: Optional[torch.Tensor] = None,
                     replace_h_space: Optional[torch.Tensor] = None,
                     replace_skip_conns: Optional[Dict[int, torch.Tensor]] = None,
                     zero_out_resconns: Optional[Union[int, List]] = None) -> Tuple:

        # Translation
        encoder_hidden_states_1 = class_labels
        class_labels = None
        encoder_attention_mask_1 = encoder_attention_mask
        encoder_attention_mask = None

        # return self.model.unet(sample, timestep,
        #                        encoder_hidden_states=generated_prompt_embeds,
        #                        encoder_hidden_states_1=encoder_hidden_states_1,
        #                        encoder_attention_mask_1=encoder_attention_mask_1,
        #                        ), None, None

        # By default samples have to be AT least a multiple of the overall upsampling factor.
        # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).
        # However, the upsampling interpolation output size can be forced to fit any upsampling size
        # on the fly if necessary.
        default_overall_up_factor = 2 ** self.model.unet.num_upsamplers

        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
        forward_upsample_size = False
        upsample_size = None

        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
            # print("Forward upsample size to force interpolation output size.")
            forward_upsample_size = True

        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension
        # expects mask of shape:
        #   [batch, key_tokens]
        # adds singleton query_tokens dimension:
        #   [batch,                    1, key_tokens]
        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
        if attention_mask is not None:
            # assume that mask is expressed as:
            #   (1 = keep,      0 = discard)
            # convert mask into a bias that can be added to attention scores:
            #       (keep = +0,     discard = -10000.0)
            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
            attention_mask = attention_mask.unsqueeze(1)

        # convert encoder_attention_mask to a bias the same way we do for attention_mask
        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)

        # 1. time
        timesteps = timestep
        if not torch.is_tensor(timesteps):
            is_mps = sample.device.type == "mps"
            if isinstance(timestep, float):
                dtype = torch.float32 if is_mps else torch.float64
            else:
                dtype = torch.int32 if is_mps else torch.int64
            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
        elif len(timesteps.shape) == 0:
            timesteps = timesteps[None].to(sample.device)

        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
        timesteps = timesteps.expand(sample.shape[0])

        t_emb = self.model.unet.time_proj(timesteps)

        # `Timesteps` does not contain any weights and will always return f32 tensors
        # but time_embedding might actually be running in fp16. so we need to cast here.
        # there might be better ways to encapsulate this.
        t_emb = t_emb.to(dtype=sample.dtype)

        emb = self.model.unet.time_embedding(t_emb, timestep_cond)
        aug_emb = None

        if self.model.unet.class_embedding is not None:
            if class_labels is None:
                raise ValueError("class_labels should be provided when num_class_embeds > 0")

            if self.model.unet.config.class_embed_type == "timestep":
                class_labels = self.model.unet.time_proj(class_labels)

                # `Timesteps` does not contain any weights and will always return f32 tensors
                # there might be better ways to encapsulate this.
                class_labels = 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)

        # 2. pre-process
        sample = self.model.unet.conv_in(sample)

        # 3. down
        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

        # 4. mid
        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 = {}
        # 5. up
        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]
                # down_block_res_samples = [torch.zeros_like(x) for x in down_block_res_samples]

            extracted_res_conns[i] = res_samples

            # if we have not reached the final block and need to forward the
            # upsample size, we do it here
            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
                )

        # 6. post-process
        if self.model.unet.conv_norm_out:
            sample = self.model.unet.conv_norm_out(sample)
            sample = self.model.unet.conv_act(sample)
        sample = self.model.unet.conv_out(sample)

        if not return_dict:
            return (sample,)

        return UNet2DConditionOutput(sample=sample), h_space, extracted_res_conns


class StableAudWrapper(PipelineWrapper):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        try:
            self.model = StableAudioPipeline.from_pretrained(self.model_id, token=self.token, local_files_only=True
                                                             ).to(self.device)
        except (FileNotFoundError, OSError):
            self.model = StableAudioPipeline.from_pretrained(self.model_id, token=self.token, local_files_only=False
                                                             ).to(self.device)
        self.model.transformer.eval()
        self.model.vae.eval()

        if self.double_precision:
            self.model = self.model.to(torch.float64)

    def load_scheduler(self) -> None:
        self.model.scheduler = CosineDPMSolverMultistepScheduler.from_pretrained(
            self.model_id, subfolder="scheduler", token=self.token)

    def encode_text(self, prompts: List[str], negative: bool = False, **kwargs) -> Tuple[torch.Tensor, None, torch.Tensor]:
        text_inputs = self.model.tokenizer(
            prompts,
            padding="max_length",
            max_length=self.model.tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt",
        )

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

        self.model.text_encoder.eval()
        with torch.no_grad():
            prompt_embeds = self.model.text_encoder(text_input_ids, attention_mask=attention_mask)[0]

        if negative and attention_mask is not None:  # set the masked tokens to the null embed
            prompt_embeds = torch.where(attention_mask.to(torch.bool).unsqueeze(2), prompt_embeds, 0.0)

        prompt_embeds = self.model.projection_model(text_hidden_states=prompt_embeds).text_hidden_states

        if attention_mask is None:
            raise gr.Error("Shouldn't reach here. Please raise an issue if you do.")
            """prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
            if attention_mask is not None and negative_attention_mask is None:
                negative_attention_mask = torch.ones_like(attention_mask)
            elif attention_mask is None and negative_attention_mask is not None:
                attention_mask = torch.ones_like(negative_attention_mask)"""

        if prompts == [""]:  # empty
            return torch.zeros_like(prompt_embeds, device=prompt_embeds.device), None, None

        prompt_embeds = prompt_embeds * attention_mask.unsqueeze(-1).to(prompt_embeds.dtype)
        prompt_embeds = prompt_embeds * attention_mask.unsqueeze(-1).to(prompt_embeds.dtype)
        return prompt_embeds, None, attention_mask

    def get_fn_STFT(self) -> torch.nn.Module:
        from audioldm.audio import TacotronSTFT
        return TacotronSTFT(
            filter_length=1024,
            hop_length=160,
            win_length=1024,
            n_mel_channels=64,
            sampling_rate=44100,
            mel_fmin=0,
            mel_fmax=22050,
        )

    def vae_encode(self, x: torch.Tensor) -> torch.Tensor:
        x = x.unsqueeze(0)

        audio_vae_length = int(self.model.transformer.config.sample_size * self.model.vae.hop_length)
        audio_shape = (1, self.model.vae.config.audio_channels, audio_vae_length)

        # check num_channels
        if x.shape[1] == 1 and self.model.vae.config.audio_channels == 2:
            x = x.repeat(1, 2, 1)

        audio_length = x.shape[-1]
        audio = x.new_zeros(audio_shape)
        audio[:, :, : min(audio_length, audio_vae_length)] = x[:, :, :audio_vae_length]

        encoded_audio = self.model.vae.encode(audio.to(self.device)).latent_dist
        encoded_audio = encoded_audio.sample()
        return encoded_audio

    def vae_decode(self, x: torch.Tensor) -> torch.Tensor:
        torch.cuda.empty_cache()
        # return self.model.vae.decode(1 / self.model.vae.config.scaling_factor * x).sample
        aud = self.model.vae.decode(x).sample
        return aud[:, :, self.waveform_start:self.waveform_end]

    def setup_extra_inputs(self, x: torch.Tensor, init_timestep: torch.Tensor,
                           extra_info: Optional[Any] = None,
                           audio_start_in_s: float = 0, audio_end_in_s: Optional[float] = None,
                           save_compute: bool = False) -> None:
        max_audio_length_in_s = self.model.transformer.config.sample_size * self.model.vae.hop_length / \
            self.model.vae.config.sampling_rate
        if audio_end_in_s is None:
            audio_end_in_s = max_audio_length_in_s

        if audio_end_in_s - audio_start_in_s > max_audio_length_in_s:
            raise ValueError(
                f"The total audio length requested ({audio_end_in_s-audio_start_in_s}s) is longer "
                f"than the model maximum possible length ({max_audio_length_in_s}). "
                f"Make sure that 'audio_end_in_s-audio_start_in_s<={max_audio_length_in_s}'."
            )

        self.waveform_start = int(audio_start_in_s * self.model.vae.config.sampling_rate)
        self.waveform_end = int(audio_end_in_s * self.model.vae.config.sampling_rate)

        self.seconds_start_hidden_states, self.seconds_end_hidden_states = self.model.encode_duration(
            audio_start_in_s, audio_end_in_s, self.device, False, 1)

        if save_compute:
            self.seconds_start_hidden_states = torch.cat([self.seconds_start_hidden_states, self.seconds_start_hidden_states], dim=0)
            self.seconds_end_hidden_states = torch.cat([self.seconds_end_hidden_states, self.seconds_end_hidden_states], dim=0)

        self.audio_duration_embeds = torch.cat([self.seconds_start_hidden_states,
                                                self.seconds_end_hidden_states], dim=2)

        # 7. Prepare rotary positional embedding
        self.rotary_embedding = get_1d_rotary_pos_embed(
            self.model.rotary_embed_dim,
            x.shape[2] + self.audio_duration_embeds.shape[1],
            use_real=True,
            repeat_interleave_real=False,
        )

        self.model.scheduler._init_step_index(init_timestep)

        # fix lower_order_nums for the reverse step - Option 1: only start from first order
        # self.model.scheduler.lower_order_nums = 0
        # self.model.scheduler.model_outputs = [None] * self.model.scheduler.config.solver_order
        # fix lower_order_nums for the reverse step - Option 2: start from the correct order with history
        t_to_idx = {float(v): k for k, v in enumerate(self.model.scheduler.timesteps)}
        idx = len(self.model.scheduler.timesteps) - t_to_idx[float(init_timestep)] - 1
        self.model.scheduler.model_outputs = [None, extra_info[idx] if extra_info is not None else None]
        self.model.scheduler.lower_order_nums = min(self.model.scheduler.step_index,
                                                    self.model.scheduler.config.solver_order)

        # if rand check:
        #     x *= self.model.scheduler.init_noise_sigma
        # return x

    def sample_xts_from_x0(self, x0: torch.Tensor, num_inference_steps: int = 50) -> torch.Tensor:
        """
        Samples from P(x_1:T|x_0)
        """

        sigmas = self.model.scheduler.sigmas
        shapes = self.get_noise_shape(x0, num_inference_steps + 1)
        xts = torch.zeros(shapes).to(x0.device)
        xts[0] = x0

        timesteps = self.model.scheduler.timesteps.to(self.device)
        t_to_idx = {float(v): k for k, v in enumerate(timesteps)}
        for t in reversed(timesteps):
            # idx = t_to_idx[int(t)]
            idx = num_inference_steps - t_to_idx[float(t)]
            n = torch.randn_like(x0)
            xts[idx] = x0 + n * sigmas[t_to_idx[float(t)]]
        return xts

    def get_zs_from_xts(self, xt: torch.Tensor, xtm1: torch.Tensor, data_pred: torch.Tensor,
                        t: torch.Tensor, numerical_fix: bool = True, first_order: bool = False, **kwargs
                        ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        # pred of x0
        sigmas = self.model.scheduler.sigmas
        timesteps = self.model.scheduler.timesteps
        solver_order = self.model.scheduler.config.solver_order

        if self.model.scheduler.step_index is None:
            self.model.scheduler._init_step_index(t)
        curr_step_index = self.model.scheduler.step_index

        # Improve numerical stability for small number of steps
        lower_order_final = (curr_step_index == len(timesteps) - 1) and (
            self.model.scheduler.config.euler_at_final
            or (self.model.scheduler.config.lower_order_final and len(timesteps) < 15)
            or self.model.scheduler.config.final_sigmas_type == "zero")
        lower_order_second = ((curr_step_index == len(timesteps) - 2) and
                              self.model.scheduler.config.lower_order_final and len(timesteps) < 15)

        data_pred = self.model.scheduler.convert_model_output(data_pred, sample=xt)
        for i in range(solver_order - 1):
            self.model.scheduler.model_outputs[i] = self.model.scheduler.model_outputs[i + 1]
        self.model.scheduler.model_outputs[-1] = data_pred

        # instead of brownian noise, here we calculate the noise ourselves
        if (curr_step_index == len(timesteps) - 1) and self.model.scheduler.config.final_sigmas_type == "zero":
            z = torch.zeros_like(xt)
        elif first_order or solver_order == 1 or self.model.scheduler.lower_order_nums < 1 or lower_order_final:
            sigma_t, sigma_s = sigmas[curr_step_index + 1], sigmas[curr_step_index]
            h = torch.log(sigma_s) - torch.log(sigma_t)
            z = (xtm1 - (sigma_t / sigma_s * torch.exp(-h)) * xt - (1 - torch.exp(-2.0 * h)) * data_pred) \
                / (sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)))
        elif solver_order == 2 or self.model.scheduler.lower_order_nums < 2 or lower_order_second:
            sigma_t = sigmas[curr_step_index + 1]
            sigma_s0 = sigmas[curr_step_index]
            sigma_s1 = sigmas[curr_step_index - 1]
            m0, m1 = self.model.scheduler.model_outputs[-1], self.model.scheduler.model_outputs[-2]
            h, h_0 = torch.log(sigma_s0) - torch.log(sigma_t), torch.log(sigma_s1) - torch.log(sigma_s0)
            r0 = h_0 / h
            D0, D1 = m0, (1.0 / r0) * (m0 - m1)

            # sde-dpmsolver++
            z = (xtm1 - (sigma_t / sigma_s0 * torch.exp(-h)) * xt
                 - (1 - torch.exp(-2.0 * h)) * D0
                 - 0.5 * (1 - torch.exp(-2.0 * h)) * D1) \
                / (sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)))

        # correction to avoid error accumulation
        if numerical_fix:
            if first_order or solver_order == 1 or self.model.scheduler.lower_order_nums < 1 or lower_order_final:
                xtm1 = self.model.scheduler.dpm_solver_first_order_update(data_pred, sample=xt, noise=z)
            elif solver_order == 2 or self.model.scheduler.lower_order_nums < 2 or lower_order_second:
                xtm1 = self.model.scheduler.multistep_dpm_solver_second_order_update(
                    self.model.scheduler.model_outputs, sample=xt, noise=z)
            # If not perfect recon - maybe TODO fix self.model.scheduler.model_outputs as well?

        if self.model.scheduler.lower_order_nums < solver_order:
            self.model.scheduler.lower_order_nums += 1
        # upon completion increase step index by one
        self.model.scheduler._step_index += 1

        return z, xtm1, self.model.scheduler.model_outputs[-2]

    def get_sr(self) -> int:
        return self.model.vae.config.sampling_rate

    def get_noise_shape(self, x0: torch.Tensor, num_steps: int) -> Tuple[int, int, int]:
        variance_noise_shape = (num_steps,
                                self.model.transformer.config.in_channels,
                                int(self.model.transformer.config.sample_size))
        return variance_noise_shape

    def reverse_step_with_custom_noise(self, model_output: torch.Tensor, timestep: torch.Tensor, sample: torch.Tensor,
                                       variance_noise: Optional[torch.Tensor] = None,
                                       first_order: bool = False, **kwargs
                                       ) -> torch.Tensor:
        if self.model.scheduler.step_index is None:
            self.model.scheduler._init_step_index(timestep)

        # Improve numerical stability for small number of steps
        lower_order_final = (self.model.scheduler.step_index == len(self.model.scheduler.timesteps) - 1) and (
            self.model.scheduler.config.euler_at_final
            or (self.model.scheduler.config.lower_order_final and len(self.model.scheduler.timesteps) < 15)
            or self.model.scheduler.config.final_sigmas_type == "zero"
        )
        lower_order_second = (
            (self.model.scheduler.step_index == len(self.model.scheduler.timesteps) - 2) and
            self.model.scheduler.config.lower_order_final and len(self.model.scheduler.timesteps) < 15
        )

        model_output = self.model.scheduler.convert_model_output(model_output, sample=sample)
        for i in range(self.model.scheduler.config.solver_order - 1):
            self.model.scheduler.model_outputs[i] = self.model.scheduler.model_outputs[i + 1]
        self.model.scheduler.model_outputs[-1] = model_output

        if variance_noise is None:
            if self.model.scheduler.noise_sampler is None:
                self.model.scheduler.noise_sampler = BrownianTreeNoiseSampler(
                    model_output, sigma_min=self.model.scheduler.config.sigma_min,
                    sigma_max=self.model.scheduler.config.sigma_max, seed=None)
            variance_noise = self.model.scheduler.noise_sampler(
                self.model.scheduler.sigmas[self.model.scheduler.step_index],
                self.model.scheduler.sigmas[self.model.scheduler.step_index + 1]).to(model_output.device)

        if first_order or self.model.scheduler.config.solver_order == 1 or \
                self.model.scheduler.lower_order_nums < 1 or lower_order_final:
            prev_sample = self.model.scheduler.dpm_solver_first_order_update(
                model_output, sample=sample, noise=variance_noise)
        elif self.model.scheduler.config.solver_order == 2 or \
                self.model.scheduler.lower_order_nums < 2 or lower_order_second:
            prev_sample = self.model.scheduler.multistep_dpm_solver_second_order_update(
                self.model.scheduler.model_outputs, sample=sample, noise=variance_noise)

        if self.model.scheduler.lower_order_nums < self.model.scheduler.config.solver_order:
            self.model.scheduler.lower_order_nums += 1

        # upon completion increase step index by one
        self.model.scheduler._step_index += 1

        return prev_sample

    def unet_forward(self,
                     sample: torch.FloatTensor,
                     timestep: Union[torch.Tensor, float, int],
                     encoder_hidden_states: torch.Tensor,
                     encoder_attention_mask: Optional[torch.Tensor] = None,
                     return_dict: bool = True,
                     **kwargs) -> Tuple:

        # Create text_audio_duration_embeds and audio_duration_embeds
        embeds = torch.cat([encoder_hidden_states, self.seconds_start_hidden_states, self.seconds_end_hidden_states],
                           dim=1)
        if encoder_attention_mask is None:
            # handle the batched case
            if embeds.shape[0] > 1:
                embeds[0] = torch.zeros_like(embeds[0], device=embeds.device)
            else:
                embeds = torch.zeros_like(embeds, device=embeds.device)

        noise_pred = self.model.transformer(sample,
                                            timestep.unsqueeze(0),
                                            encoder_hidden_states=embeds,
                                            global_hidden_states=self.audio_duration_embeds,
                                            rotary_embedding=self.rotary_embedding)

        if not return_dict:
            return (noise_pred.sample,)

        return noise_pred, None, None


def load_model(model_id: str, device: torch.device,
               double_precision: bool = False, token: Optional[str] = None) -> PipelineWrapper:
    if 'audioldm2' in model_id:
        ldm_stable = AudioLDM2Wrapper(model_id=model_id, device=device, double_precision=double_precision, token=token)
    elif 'stable-audio' in model_id:
        ldm_stable = StableAudWrapper(model_id=model_id, device=device, double_precision=double_precision, token=token)
    ldm_stable.load_scheduler()
    torch.cuda.empty_cache()
    return ldm_stable