import numpy as np
import torch
import typing as tp
import math
from torchaudio import transforms as T

from .utils import prepare_audio
from .sampling import sample, sample_k
from ..data.utils import PadCrop


def generate_diffusion_uncond(
    model,
    steps: int = 250,
    batch_size: int = 1,
    sample_size: int = 2097152,
    seed: int = -1,
    device: str = "cuda",
    init_audio: tp.Optional[tp.Tuple[int, torch.Tensor]] = None,
    init_noise_level: float = 1.0,
    return_latents=False,
    **sampler_kwargs,
) -> torch.Tensor:
    # The length of the output in audio samples
    audio_sample_size = sample_size

    # If this is latent diffusion, change sample_size instead to the downsampled latent size
    if model.pretransform is not None:
        sample_size = sample_size // model.pretransform.downsampling_ratio

    # Seed
    # The user can explicitly set the seed to deterministically generate the same output. Otherwise, use a random seed.
    seed = seed if seed != -1 else np.random.randint(0, 2**32 - 1, dtype=np.uint32)
    print(seed)
    torch.manual_seed(seed)
    # Define the initial noise immediately after setting the seed
    noise = torch.randn([batch_size, model.io_channels, sample_size], device=device)

    if init_audio is not None:
        # The user supplied some initial audio (for inpainting or variation). Let us prepare the input audio.
        in_sr, init_audio = init_audio

        io_channels = model.io_channels

        # For latent models, set the io_channels to the autoencoder's io_channels
        if model.pretransform is not None:
            io_channels = model.pretransform.io_channels

        # Prepare the initial audio for use by the model
        init_audio = prepare_audio(
            init_audio,
            in_sr=in_sr,
            target_sr=model.sample_rate,
            target_length=audio_sample_size,
            target_channels=io_channels,
            device=device,
        )

        # For latent models, encode the initial audio into latents
        if model.pretransform is not None:
            init_audio = model.pretransform.encode(init_audio)

        init_audio = init_audio.repeat(batch_size, 1, 1)
    else:
        # The user did not supply any initial audio for inpainting or variation. Generate new output from scratch.
        init_audio = None
        init_noise_level = None

    # Inpainting mask

    if init_audio is not None:
        # variations
        sampler_kwargs["sigma_max"] = init_noise_level
        mask = None
    else:
        mask = None

    # Now the generative AI part:
    # k-diffusion denoising process go!
    sampled = sample_k(
        model.model, noise, init_audio, mask, steps, **sampler_kwargs, device=device
    )

    # Denoising process done.
    # If this is latent diffusion, decode latents back into audio
    if model.pretransform is not None and not return_latents:
        sampled = model.pretransform.decode(sampled)

    # Return audio
    return sampled


def upsample_diffusion(
    model,
    audio: torch.Tensor,
    model_type: str,
    steps: int = 250,
    codes: torch.Tensor = None,
    cfg_scale: float = 6,
    batch_size: int = 1,
    start_sample: int = 0,
    seconds_total: float = 10.0,
    sample_size: int = 2097152,
    sample_rate: int = 48000,
    seed: int = -1,
    device: str = "cuda",
    return_latents: bool = False,
    **sampler_kwargs,
):
    # If this is latent diffusion, change sample_size instead to the downsampled latent size
    if model.pretransform is not None:
        sample_size = sample_size // model.pretransform.downsampling_ratio

    audio = audio.to(device)

    # Seed
    # The user can explicitly set the seed to deterministically generate the same output. Otherwise, use a random seed.
    seed = seed if seed != -1 else np.random.randint(0, 2**32 - 1)
    print(seed)
    torch.manual_seed(seed)
    # Define the initial noise immediately after setting the seed
    noise = torch.randn([batch_size, model.io_channels, sample_size], device=device)

    # Conditioning

    # create conditioning dict
    seconds_start = start_sample * sample_rate
    # seconds_total = audio.shape[-1] / sample_rate
    conditioning = [
        {
            "seconds_start": seconds_start,
            "seconds_total": seconds_total,
        }
    ]

    if model_type == "semantic":
        conditioning[0]["semantic_musicfm"] = (audio, codes)
    elif model_type == "codec":
        conditioning[0]["codec_dac_2c_25x12"] = {
            "audio": audio,
            "codes": codes.squeeze(0),
        }
    elif model_type == "mert":
        conditioning[0]["mert"] = {
            "audio": audio,
            "codes": codes.squeeze(0),
        }
    else:
        raise ValueError(f"Invalid model_type: {model_type}")

    conditioning_tensors = None
    assert (
        conditioning is not None or conditioning_tensors is not None
    ), "Must provide either conditioning or conditioning_tensors"
    if conditioning_tensors is None:
        conditioning_tensors = model.conditioner(conditioning, device)
    conditioning_tensors = model.get_conditioning_inputs(conditioning_tensors)

    # stuff we don't use
    init_audio = None
    mask = None
    negative_conditioning_tensors = None

    # Now the generative AI part:
    # k-diffusion denoising process go!
    sampled = sample_k(
        model.model,
        noise,
        init_audio,
        mask,
        steps,
        **sampler_kwargs,
        **conditioning_tensors,
        # **negative_conditioning_tensors,
        cfg_scale=cfg_scale,
        batch_cfg=True,
        rescale_cfg=True,
        device=device,
    )

    # Denoising process done.
    # If this is latent diffusion, decode latents back into audio
    if model.pretransform is not None and not return_latents:
        # cast sampled latents to pretransform dtype
        sampled = sampled.to(next(model.pretransform.parameters()).dtype)
        sampled = model.pretransform.decode(sampled)

    # Return audio
    return sampled


def denoising_diffusion_from_codes(
    model,
    codec_codes: torch.Tensor,
    steps: int = 250,
    cfg_scale: float = 1.0,
    batch_size: int = 1,
    start_sample: int = 0,
    seconds_total: float = 10.0,
    sample_size: int = 1000,
    sample_rate: int = 48000,
    seed: int = -1,
    init_noise: torch.Tensor = None,
    device: str = "cuda",
    return_latents: bool = False,
    **sampler_kwargs,
):
    # The user can explicitly set the seed to deterministically generate the same output. Otherwise, use a random seed.
    seed = seed if seed != -1 else np.random.randint(0, 2**32 - 1)
    torch.manual_seed(seed)
    # Define the initial noise immediately after setting the seed

    if init_noise is None:
        noise = torch.randn([batch_size, model.io_channels, sample_size], device=device)
    else:
        noise = init_noise

    print("noise", noise.shape)

    # conditioning_tensors is a dict with the tensors (embeddings)
    codec_codes = model.conditioner.conditioners["codec_codes"]([codec_codes], device)
    # print(semantic_embeds.shape, codec_embeds.shape)

    conditioning_tensors = {
        "codec_codes": codec_codes,
    }

    conditioning_tensors = model.get_conditioning_inputs(conditioning_tensors)

    # stuff we don't use
    init_audio = None
    mask = None
    negative_conditioning_tensors = None

    # Now the generative AI part:
    # k-diffusion denoising process go!
    sampled = sample_k(
        model.model,
        noise,
        init_audio,
        mask,
        steps,
        **sampler_kwargs,
        **conditioning_tensors,
        # **negative_conditioning_tensors,
        cfg_scale=cfg_scale,
        batch_cfg=True,
        rescale_cfg=True,
        device=device,
    )

    # Denoising process done.
    # If this is latent diffusion, decode latents back into audio

    # Return audio
    return sampled


def upsample_diffusion_from_semantic(
    model,
    semantic_codes: torch.Tensor,
    steps: int = 250,
    cfg_scale: float = 1.0,
    batch_size: int = 1,
    start_sample: int = 0,
    seconds_total: float = 10.0,
    sample_size: int = 1000,
    sample_rate: int = 48000,
    seed: int = -1,
    init_noise: torch.Tensor = None,
    device: str = "cuda",
    return_latents: bool = False,
    **sampler_kwargs,
):
    # The user can explicitly set the seed to deterministically generate the same output. Otherwise, use a random seed.
    seed = seed if seed != -1 else np.random.randint(0, 2**32 - 1)
    torch.manual_seed(seed)
    # Define the initial noise immediately after setting the seed

    if init_noise is None:
        noise = torch.randn([batch_size, model.io_channels, sample_size], device=device)
    else:
        noise = init_noise

    print("noise", noise.shape)

    # conditioning_tensors is a dict with the tensors (embeddings)
    semantic_codes = model.conditioner.conditioners["semantic_codes"](
        [semantic_codes], device
    )
    # print(semantic_embeds.shape, codec_embeds.shape)

    conditioning_tensors = {
        "semantic_codes": semantic_codes,
    }

    conditioning_tensors = model.get_conditioning_inputs(conditioning_tensors)

    # stuff we don't use
    init_audio = None
    mask = None
    negative_conditioning_tensors = None

    # Now the generative AI part:
    # k-diffusion denoising process go!
    sampled = sample_k(
        model.model,
        noise,
        init_audio,
        mask,
        steps,
        **sampler_kwargs,
        **conditioning_tensors,
        # **negative_conditioning_tensors,
        cfg_scale=cfg_scale,
        batch_cfg=True,
        rescale_cfg=True,
        device=device,
    )

    # Denoising process done.
    # If this is latent diffusion, decode latents back into audio

    # Return audio
    return sampled


def upsample_diffusion_from_discrete(
    model,
    discrete_codes: torch.Tensor,
    steps: int = 250,
    cfg_scale: float = 1.0,
    batch_size: int = 1,
    start_sample: int = 0,
    seconds_total: float = 10.0,
    sample_size: int = 1000,
    sample_rate: int = 48000,
    seed: int = -1,
    init_noise: torch.Tensor = None,
    device: str = "cuda",
    return_latents: bool = False,
    **sampler_kwargs,
):
    # The user can explicitly set the seed to deterministically generate the same output. Otherwise, use a random seed.
    seed = seed if seed != -1 else np.random.randint(0, 2**32 - 1)
    torch.manual_seed(seed)
    # Define the initial noise immediately after setting the seed

    if init_noise is None:
        noise = torch.randn([batch_size, model.io_channels, sample_size], device=device)
    else:
        noise = init_noise

    print("noise", noise.shape)

    # conditioning_tensors is a dict with the tensors (embeddings)
    discrete_codes = model.conditioner.conditioners["discrete_codes"](
        [discrete_codes], device
    )
    # print(semantic_embeds.shape, codec_embeds.shape)

    conditioning_tensors = {
        "discrete_codes": discrete_codes,
    }

    conditioning_tensors = model.get_conditioning_inputs(conditioning_tensors)

    # stuff we don't use
    init_audio = None
    mask = None
    negative_conditioning_tensors = None

    # Now the generative AI part:
    # k-diffusion denoising process go!
    sampled = sample_k(
        model.model,
        noise,
        init_audio,
        mask,
        steps,
        **sampler_kwargs,
        **conditioning_tensors,
        # **negative_conditioning_tensors,
        cfg_scale=cfg_scale,
        batch_cfg=True,
        rescale_cfg=True,
        device=device,
    )

    # Denoising process done.
    # If this is latent diffusion, decode latents back into audio

    # Return audio
    return sampled


def upsample_diffusion_from_uncond(
    model,
    steps: int = 250,
    cfg_scale: float = 1.0,
    batch_size: int = 1,
    start_sample: int = 0,
    seconds_total: float = 10.0,
    sample_size: int = 1000,
    sample_rate: int = 48000,
    seed: int = -1,
    init_noise: torch.Tensor = None,
    init_audio: torch.Tensor = None,
    mask: torch.Tensor = None,
    device: str = "cuda",
    return_latents: bool = False,
    **sampler_kwargs,
):
    # The user can explicitly set the seed to deterministically generate the same output. Otherwise, use a random seed.
    seed = seed if seed != -1 else np.random.randint(0, 2**32 - 1)
    torch.manual_seed(seed)
    # Define the initial noise immediately after setting the seed

    if init_noise is None:
        noise = torch.randn([batch_size, model.io_channels, sample_size], device=device)
    else:
        noise = init_noise

    print("noise", noise.shape)

    # stuff we don't use
    negative_conditioning_tensors = None
    conditioning_tensors = None

    conditioning = [
        {
            "seconds_start": 0.0,
            "seconds_total": seconds_total,
        }
    ]

    if conditioning_tensors is None:
        conditioning_tensors = model.conditioner(conditioning, device)
    conditioning_tensors = model.get_conditioning_inputs(conditioning_tensors)

    # Now the generative AI part:
    # k-diffusion denoising process go!
    sampled = sample_k(
        model.model,
        noise,
        init_audio,
        mask,
        steps,
        **sampler_kwargs,
        **conditioning_tensors,
        # **negative_conditioning_tensors,
        cfg_scale=cfg_scale,
        batch_cfg=True,
        rescale_cfg=True,
        device=device,
    )

    # Denoising process done.
    # If this is latent diffusion, decode latents back into audio

    # Return audio
    return sampled


def upsample_diffusion_from_lyrics(
    model,
    lyrics: str,
    steps: int = 250,
    cfg_scale: float = 1.0,
    batch_size: int = 1,
    start_sample: int = 0,
    seconds_total: float = 10.0,
    sample_size: int = 1000,
    sample_rate: int = 48000,
    seed: int = -1,
    init_noise: torch.Tensor = None,
    device: str = "cuda",
    return_latents: bool = False,
    **sampler_kwargs,
):
    # The user can explicitly set the seed to deterministically generate the same output. Otherwise, use a random seed.
    seed = seed if seed != -1 else np.random.randint(0, 2**32 - 1)
    torch.manual_seed(seed)
    # Define the initial noise immediately after setting the seed

    if init_noise is None:
        noise = torch.randn([batch_size, model.io_channels, sample_size], device=device)
    else:
        noise = init_noise

    print("noise", noise.shape)

    # stuff we don't use
    init_audio = None
    mask = None
    negative_conditioning_tensors = None

    # conditioning_tensors is a dict with the tensors (embeddings)
    lyrics_tensor = model.conditioner.conditioners["lyrics"]([lyrics], device)

    conditioning_tensors = {
        "lyrics": lyrics_tensor,
    }

    conditioning_tensors = model.get_conditioning_inputs(conditioning_tensors)

    # Now the generative AI part:
    # k-diffusion denoising process go!
    sampled = sample_k(
        model.model,
        noise,
        init_audio,
        mask,
        steps,
        **sampler_kwargs,
        **conditioning_tensors,
        # **negative_conditioning_tensors,
        cfg_scale=cfg_scale,
        batch_cfg=True,
        rescale_cfg=True,
        device=device,
    )

    # Denoising process done.
    # If this is latent diffusion, decode latents back into audio

    # Return audio
    return sampled


def upsample_diffusion_from_semantic_and_text(
    model,
    semantic_codes: torch.Tensor,
    tags: str = "",
    lyrics: str = "",
    steps: int = 250,
    cfg_scale: float = 1.0,
    batch_size: int = 1,
    start_sample: int = 0,
    seconds_total: float = 30.0,
    sample_size: int = 1000,
    sample_rate: int = 48000,
    seed: int = -1,
    init_noise: torch.Tensor = None,
    device: str = "cuda",
    return_latents: bool = False,
    init_audio: torch.Tensor = None,
    mask: torch.Tensor = None,
    compile: bool = True,
    latent_context: torch.Tensor = None,
    **sampler_kwargs,
):
    # The user can explicitly set the seed to deterministically generate the same output. Otherwise, use a random seed.
    seed = seed if seed != -1 else np.random.randint(0, 2**32 - 1)
    torch.manual_seed(seed)
    # Define the initial noise immediately after setting the seed

    if init_noise is None:
        noise = torch.randn([batch_size, model.io_channels, sample_size], device=device)
    else:
        noise = init_noise

    # stuff we don't use
    # mask = None
    negative_conditioning_tensors = None

    # conditioning_tensors is a dict with the tensors (embeddings)
    tags_and_lyrics_tensor = model.conditioner.conditioners["tags_and_lyrics"](
        [(tags, lyrics)], device
    )
    semantic_codes_tensor = model.conditioner.conditioners["semantic_codes"](
        [semantic_codes], device
    )

    conditioning_tensors = {
        "tags_and_lyrics": tags_and_lyrics_tensor,
        "semantic_codes": semantic_codes_tensor,
    }

    if latent_context is not None:
        latent_context_tensor = model.conditioner.conditioners["latent_context"](
            [latent_context], device
        )
        conditioning_tensors["latent_context"] = latent_context_tensor

    # create unconditional conditioning
    if cfg_scale != 1.0:
        empty_tags_and_lyrics_tensor = model.conditioner.conditioners[
            "tags_and_lyrics"
        ]([("", "")], device)
        empty_semantic_codes_tensor = model.conditioner.conditioners["semantic_codes"](
            [torch.ones_like(semantic_codes) * 4000], device
        )
        # use the empty tags but keep semantic conditioning the same
        conditioning_tensors["empty_tags_and_lyrics"] = empty_tags_and_lyrics_tensor
        conditioning_tensors["empty_semantic_codes"] = semantic_codes_tensor

        if latent_context is not None:  # supply latent context if it exists
            # empty_latent_context = model.vae_pad_embed.view(-1, 1).repeat(1, 3000)
            # empty_latent_context_tensor = model.conditioner.conditioners[
            #    "latent_context"
            # ]([empty_latent_context], device)
            conditioning_tensors["empty_latent_context"] = latent_context_tensor

        conditioning_tensors = model.get_conditioning_inputs(
            conditioning_tensors, negative=False, empty=True
        )
    else:
        conditioning_tensors = model.get_conditioning_inputs(conditioning_tensors)

    # Now the generative AI part:
    # k-diffusion denoising process go!
    sampled = sample_k(
        torch.compile(model.model, disable=not compile),
        noise,
        init_audio,
        mask,
        steps,
        **sampler_kwargs,
        **conditioning_tensors,
        # **negative_conditioning_tensors,
        # latent_context=latent_context,
        cfg_scale=cfg_scale,
        batch_cfg=True,
        rescale_cfg=True,
        device=device,
        # sigma_min=0.001,
        # sigma_max=1,
    )

    # Denoising process done.
    # If this is latent diffusion, decode latents back into audio
    # rescale the latents by scale factor, but model might not have scale_factor
    # if hasattr(model, "scale_factor"):
    #    print("rescaling by", model.scale_factor)
    #    sampled /= model.scale_factor

    # Return audio
    return sampled


def upsample_diffusion_from_semantic_and_text_with_phonemes(
    model,
    semantic_codes: torch.Tensor,
    tags: str = "",
    lyrics: str = "",
    phonemes: str = "",
    steps: int = 250,
    cfg_scale: float = 1.0,
    batch_size: int = 1,
    start_sample: int = 0,
    seconds_total: float = 30.0,
    sample_size: int = 1000,
    sample_rate: int = 48000,
    seed: int = -1,
    init_noise: torch.Tensor = None,
    device: str = "cuda",
    return_latents: bool = False,
    init_audio: torch.Tensor = None,
    mask: torch.Tensor = None,
    compile: bool = True,
    **sampler_kwargs,
):
    # The user can explicitly set the seed to deterministically generate the same output. Otherwise, use a random seed.
    seed = seed if seed != -1 else np.random.randint(0, 2**32 - 1)
    torch.manual_seed(seed)
    # Define the initial noise immediately after setting the seed

    if init_noise is None:
        noise = torch.randn([batch_size, model.io_channels, sample_size], device=device)
    else:
        noise = init_noise

    print("noise", noise.shape)

    # stuff we don't use
    # mask = None
    negative_conditioning_tensors = None

    print(model.conditioner.conditioners.keys())

    # conditioning_tensors is a dict with the tensors (embeddings)
    tags_and_lyrics_tensor = model.conditioner.conditioners["tags_and_lyrics"](
        [(tags, lyrics)], device
    )
    phonemes_tensor = model.conditioner.conditioners["phonemes"](
        [phonemes], device
    )  # empty phonemes
    semantic_codes_tensor = model.conditioner.conditioners["semantic_codes"](
        [semantic_codes], device
    )

    conditioning_tensors = {
        "tags_and_lyrics": tags_and_lyrics_tensor,
        "semantic_codes": semantic_codes_tensor,
        "phonemes": phonemes_tensor,
    }

    # create unconditional conditioning
    if cfg_scale != 1.0:
        empty_tags_and_lyrics_tensor = model.conditioner.conditioners[
            "tags_and_lyrics"
        ]([("", "")], device)
        empty_phonemes_tensor = model.conditioner.conditioners["phonemes"](
            [""], device
        )  # empty phonemes
        empty_semantic_codes_tensor = model.conditioner.conditioners["semantic_codes"](
            [torch.ones_like(semantic_codes) * 4000], device
        )
        # use the empty tags but keep semantic conditioning the same
        conditioning_tensors["empty_phonemes"] = empty_phonemes_tensor
        conditioning_tensors["empty_tags_and_lyrics"] = empty_tags_and_lyrics_tensor
        conditioning_tensors["empty_semantic_codes"] = semantic_codes_tensor
        print("no cfg on semantic")

        conditioning_tensors = model.get_conditioning_inputs(
            conditioning_tensors, negative=False, empty=True
        )
    else:
        conditioning_tensors = model.get_conditioning_inputs(conditioning_tensors)

    print(conditioning_tensors.keys())

    # Now the generative AI part:
    # k-diffusion denoising process go!
    sampled = sample_k(
        torch.compile(model.model, disable=not compile),
        noise,
        init_audio,
        mask,
        steps,
        **sampler_kwargs,
        **conditioning_tensors,
        # **negative_conditioning_tensors,
        cfg_scale=cfg_scale,
        batch_cfg=True,
        rescale_cfg=True,
        device=device,
    )

    # Denoising process done.
    # If this is latent diffusion, decode latents back into audio
    # rescale the latents by scale factor, but model might not have scale_factor
    # if hasattr(model, "scale_factor"):
    #    print("rescaling by", model.scale_factor)
    #    sampled /= model.scale_factor

    # Return audio
    return sampled


def upsample_diffusion_from_codes(
    model,
    semantic_codes: torch.Tensor = None,
    codec_codes: torch.Tensor = None,
    keep_n_codebooks: int = 12,
    steps: int = 250,
    cfg_scale: float = 1.0,
    batch_size: int = 1,
    start_sample: int = 0,
    seconds_total: float = 10.0,
    sample_size: int = 250,
    sample_rate: int = 48000,
    seed: int = -1,
    init_noise: torch.Tensor = None,
    device: str = "cuda",
    return_latents: bool = False,
    **sampler_kwargs,
):
    # audio = audio.to(device)

    # Seed
    # The user can explicitly set the seed to deterministically generate the same output. Otherwise, use a random seed.
    seed = seed if seed != -1 else np.random.randint(0, 2**32 - 1)
    torch.manual_seed(seed)
    # Define the initial noise immediately after setting the seed

    if init_noise is None:
        noise = torch.randn([batch_size, model.io_channels, sample_size], device=device)
    else:
        noise = init_noise

    print("noise", noise.shape)

    # conditioning_tensors is a dict with the tensors (embeddings)
    semantic_embeds = model.conditioner.conditioners["semantic_codes"](
        [semantic_codes], device
    )
    codec_embeds = model.conditioner.conditioners["codec_codes"](
        [codec_codes], device, keep_n_codebooks
    )
    # print(semantic_embeds.shape, codec_embeds.shape)

    conditioning_tensors = {
        "semantic_codes": semantic_embeds,
        "codec_codes": codec_embeds,
    }

    conditioning_tensors = model.get_conditioning_inputs(conditioning_tensors)

    # stuff we don't use
    init_audio = None
    mask = None
    negative_conditioning_tensors = None

    # Now the generative AI part:
    # k-diffusion denoising process go!
    sampled = sample_k(
        model.model,
        noise,
        init_audio,
        mask,
        steps,
        **sampler_kwargs,
        **conditioning_tensors,
        # **negative_conditioning_tensors,
        cfg_scale=cfg_scale,
        batch_cfg=True,
        rescale_cfg=True,
        device=device,
    )

    # Denoising process done.
    # If this is latent diffusion, decode latents back into audio

    # Return audio
    return sampled


def generate_diffusion_cond(
    model,
    steps: int = 250,
    cfg_scale=6,
    conditioning: dict = None,
    conditioning_tensors: tp.Optional[dict] = None,
    negative_conditioning: dict = None,
    negative_conditioning_tensors: tp.Optional[dict] = None,
    batch_size: int = 1,
    sample_size: int = 2097152,
    sample_rate: int = 48000,
    seed: int = -1,
    device: str = "cuda",
    init_audio: tp.Optional[tp.Tuple[int, torch.Tensor]] = None,
    init_noise_level: float = 1.0,
    mask_args: dict = None,
    return_latents=False,
    **sampler_kwargs,
) -> torch.Tensor:
    """
    Generate audio from a prompt using a diffusion model.

    Args:
        model: The diffusion model to use for generation.
        steps: The number of diffusion steps to use.
        cfg_scale: Classifier-free guidance scale
        conditioning: A dictionary of conditioning parameters to use for generation.
        conditioning_tensors: A dictionary of precomputed conditioning tensors to use for generation.
        batch_size: The batch size to use for generation.
        sample_size: The length of the audio to generate, in samples.
        sample_rate: The sample rate of the audio to generate (Deprecated, now pulled from the model directly)
        seed: The random seed to use for generation, or -1 to use a random seed.
        device: The device to use for generation.
        init_audio: A tuple of (sample_rate, audio) to use as the initial audio for generation.
        init_noise_level: The noise level to use when generating from an initial audio sample.
        return_latents: Whether to return the latents used for generation instead of the decoded audio.
        **sampler_kwargs: Additional keyword arguments to pass to the sampler.
    """

    # The length of the output in audio samples
    audio_sample_size = sample_size

    # If this is latent diffusion, change sample_size instead to the downsampled latent size
    if model.pretransform is not None:
        sample_size = sample_size // model.pretransform.downsampling_ratio

    # Seed
    # The user can explicitly set the seed to deterministically generate the same output. Otherwise, use a random seed.
    seed = seed if seed != -1 else np.random.randint(0, 2**32 - 1)
    print(seed)
    torch.manual_seed(seed)
    # Define the initial noise immediately after setting the seed
    noise = torch.randn([batch_size, model.io_channels, sample_size], device=device)

    # Conditioning
    assert (
        conditioning is not None or conditioning_tensors is not None
    ), "Must provide either conditioning or conditioning_tensors"
    if conditioning_tensors is None:
        conditioning_tensors = model.conditioner(conditioning, device)
    conditioning_tensors = model.get_conditioning_inputs(conditioning_tensors)

    if negative_conditioning is not None or negative_conditioning_tensors is not None:
        if negative_conditioning_tensors is None:
            negative_conditioning_tensors = model.conditioner(
                negative_conditioning, device
            )

        negative_conditioning_tensors = model.get_conditioning_inputs(
            negative_conditioning_tensors, negative=True
        )
    else:
        negative_conditioning_tensors = {}

    if init_audio is not None:
        # The user supplied some initial audio (for inpainting or variation). Let us prepare the input audio.
        in_sr, init_audio = init_audio

        io_channels = model.io_channels

        # For latent models, set the io_channels to the autoencoder's io_channels
        if model.pretransform is not None:
            io_channels = model.pretransform.io_channels

        # Prepare the initial audio for use by the model
        init_audio = prepare_audio(
            init_audio,
            in_sr=in_sr,
            target_sr=model.sample_rate,
            target_length=audio_sample_size,
            target_channels=io_channels,
            device=device,
        )

        # For latent models, encode the initial audio into latents
        if model.pretransform is not None:
            init_audio = model.pretransform.encode(init_audio)

        init_audio = init_audio.repeat(batch_size, 1, 1)
    else:
        # The user did not supply any initial audio for inpainting or variation. Generate new output from scratch.
        init_audio = None
        init_noise_level = None
        mask_args = None

    # Inpainting mask
    if init_audio is not None and mask_args is not None:
        # Cut and paste init_audio according to cropfrom, pastefrom, pasteto
        # This is helpful for forward and reverse outpainting
        cropfrom = math.floor(mask_args["cropfrom"] / 100.0 * sample_size)
        pastefrom = math.floor(mask_args["pastefrom"] / 100.0 * sample_size)
        pasteto = math.ceil(mask_args["pasteto"] / 100.0 * sample_size)
        assert pastefrom < pasteto, "Paste From should be less than Paste To"
        croplen = pasteto - pastefrom
        if cropfrom + croplen > sample_size:
            croplen = sample_size - cropfrom
        cropto = cropfrom + croplen
        pasteto = pastefrom + croplen
        cutpaste = init_audio.new_zeros(init_audio.shape)
        cutpaste[:, :, pastefrom:pasteto] = init_audio[:, :, cropfrom:cropto]
        # print(cropfrom, cropto, pastefrom, pasteto)
        init_audio = cutpaste
        # Build a soft mask (list of floats 0 to 1, the size of the latent) from the given args
        mask = build_mask(sample_size, mask_args)
        mask = mask.to(device)
    elif init_audio is not None and mask_args is None:
        # variations
        sampler_kwargs["sigma_max"] = init_noise_level
        mask = None
    else:
        mask = None

    # Now the generative AI part:
    # k-diffusion denoising process go!
    sampled = sample_k(
        model.model,
        noise,
        init_audio,
        mask,
        steps,
        **sampler_kwargs,
        **conditioning_tensors,
        **negative_conditioning_tensors,
        cfg_scale=cfg_scale,
        batch_cfg=True,
        rescale_cfg=True,
        device=device,
    )

    # v-diffusion:
    # sampled = sample(model.model, noise, steps, 0, **conditioning_tensors, embedding_scale=cfg_scale)

    # Denoising process done.
    # If this is latent diffusion, decode latents back into audio
    if model.pretransform is not None and not return_latents:
        # cast sampled latents to pretransform dtype
        sampled = sampled.to(next(model.pretransform.parameters()).dtype)
        sampled = model.pretransform.decode(sampled)

    # Return audio
    return sampled


# builds a softmask given the parameters
# returns array of values 0 to 1, size sample_size, where 0 means noise / fresh generation, 1 means keep the input audio,
# and anything between is a mixture of old/new
# ideally 0.5 is half/half mixture but i haven't figured this out yet
def build_mask(sample_size, mask_args):
    maskstart = math.floor(mask_args["maskstart"] / 100.0 * sample_size)
    maskend = math.ceil(mask_args["maskend"] / 100.0 * sample_size)
    softnessL = round(mask_args["softnessL"] / 100.0 * sample_size)
    softnessR = round(mask_args["softnessR"] / 100.0 * sample_size)
    marination = mask_args["marination"]
    # use hann windows for softening the transition (i don't know if this is correct)
    hannL = torch.hann_window(softnessL * 2, periodic=False)[:softnessL]
    hannR = torch.hann_window(softnessR * 2, periodic=False)[softnessR:]
    # build the mask.
    mask = torch.zeros((sample_size))
    mask[maskstart:maskend] = 1
    mask[maskstart : maskstart + softnessL] = hannL
    mask[maskend - softnessR : maskend] = hannR
    # marination finishes the inpainting early in the denoising schedule, and lets audio get changed in the final rounds
    if marination > 0:
        mask = mask * (1 - marination)
    # print(mask)
    return mask