import os
import sys
import time
import modal
import torch
import pathlib
import funcy
import json
import pandas as pd
import numpy as np
import tempfile
import torchaudio
import pyloudnorm as pyln

from suno_utils.audio import Audio
from suno_utils.gpt.generation import GenerationConfig
from suno_utils.gpt.engine import Engine
from suno_utils.gpt.generation_engine import make_request
from suno_utils.worker.settings import s3_client

# from suno_utils.worker.modal_base import MODAL_MOUNTS
from suno_utils.gpt import chirp_v2_5 as chirp_v2
from suno_utils.utils.text import read_jsonl
from suno_utils.utils.s3 import list_s3_dir, read_from_s3
from suno_utils.tasks.hoot import (
    encode_filepaths,
    encode,
    clean_text,
    encode_and_align,
    _get_alignable_tokens,
    _assign_word_timings,
    get_aligned_lyrics,
    ctc_align,
    load_model_list,
    get_word_timing_from_audio_and_lyrics,
    _merge_into_lines,
    decode_logits,
    load_model,
)
from suno_utils.tasks.shimmerscore import shimmerscore
from suno_utils.utils.metrics import get_cer
from suno_utils.tasks.ear import load_model


from suno_utils.diffusion.generation import (
    preload_dit_model,
    preload_tokenizer,
    TOKENIZER_FILEPATH,
    SEMANTIC_MODEL_FILEPATH,
    SEMANTIC_CLUSTERS_FILEPATH,
    _retrieve_models,
    # preload_ear_model,
)


from suno_utils.tasks.mert_25 import (
    preload_models as preload_semantic_models,
    encode as encode_semantic,
)


from suno_utils.diffusion import generation as diffusion_gen
from suno_utils.tasks.upsample_engine import UpsampleEngine, Request, Job

# from suno_utils.tasks.upsample_engine_old import UpsampleEngine, Request, Job


MOUNT_PATH = "/suno/models"

aws_secret = modal.Secret.from_name("studio-aws")
SECRETS = [
    aws_secret,
    modal.Secret.from_dict(
        {
            "SUNO_ASSETS_PATH": "/suno/models/assets",
            "XDG_CACHE_HOME": "/suno/models/",
        }
    ),
    modal.Secret.from_name("api-callback-token"),
    modal.Secret.from_name("datadog-metrics"),
]

base_image = (
    modal.Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.10")
    .apt_install(
        "curl", "ffmpeg", "sox", "unzip", "libsox-fmt-mp3", "zlib1g-dev", "git", "clang"
    )
    .run_commands(
        [
            'curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"',
            "unzip -q awscliv2.zip",
            "./aws/install",
        ]
    )
    .dockerfile_commands(
        [
            "COPY --from=datadog/serverless-init:1.2.1 /datadog-init /app/datadog-init",
            'ENTRYPOINT ["/app/datadog-init"]',
        ]
    )
    .pip_install("torch==2.5.1", "torchaudio==2.5.1")
    .pip_install(
        "flashinfer-python", index_url="https://flashinfer.ai/whl/cu124/torch2.5/"
    )
    .pip_install_private_repos(
        "github.com/suno-ai/glockenspiel.git@a5dba4e50#subdirectory=descript-audio-codec&egg=descript-audio-codec",
        git_user="mcamac",
        secrets=[modal.Secret.from_name("victor-modal-github-token")],
    )
    .pip_install_private_repos(
        "github.com/suno-ai/neon.git@1c83548#subdirectory=hoot",
        git_user="mcamac",
        secrets=[modal.Secret.from_name("victor-modal-github-token")],
    )
    .pip_install(
        "boto3",
        "transformers",
        "tokenizers",
        "encodec",
        "ctc_segmentation",
        "psutil",
        "redis",
        "pydantic",
        "nnAudio",
        "rpyc",
        "biopython>=1.81",  # TODO: don't love this depdendency, for hoot
        "pynvml",  # for torch cuda utilization
        "torchsde",
        "ninja",
        "wheel",
    )
    .pip_install_from_pyproject(
        "/home/christian/code/glockenspiel/suno_utils/pyproject.toml",
    )
    .run_commands(  # This is really slow
        "git clone https://github.com/Dao-AILab/flash-attention.git",
        "cd flash-attention/hopper && python setup.py install",
        gpu="h100",
    )
    .apt_install(
        "libogg0",
        "libopus0",
        "opus-tools",
    )
    .pip_install("transformers==4.44.0", "wandb")
)


def calculate_stereo_width(waveform):
    # Split into left and right channels
    left = waveform[0]
    right = waveform[1]

    # Compute mid/side representation
    mid = (left + right) / 2
    side = (left - right) / 2

    # Compute RMS energy of mid and side channels
    mid_energy = torch.sqrt(torch.mean(mid**2))
    side_energy = torch.sqrt(torch.mean(side**2))

    # Compute stereo width based on mid/side ratio
    # Normalize to range 0-1 using sigmoid-like function
    width_ratio = (side_energy / (mid_energy + 1e-8)).item()
    stereo_width = 2 * (1 / (1 + np.exp(-width_ratio)) - 0.5)

    return stereo_width


def calculate_loudness(waveform, sr):
    meter = pyln.Meter(sr)
    lufs_db = meter.integrated_loudness(waveform)
    return lufs_db


def calculate_loudness_factor(waveform, sr):
    meter = pyln.Meter(sr)
    normalized_waveform = waveform / np.clip(np.max(np.abs(waveform)), 1e-10, None)
    lufs_db = meter.integrated_loudness(normalized_waveform)
    return lufs_db


def calculate_possible_clipped_samples(waveform):
    # Convert to numpy if it's a torch tensor
    if isinstance(waveform, torch.Tensor):
        waveform_np = waveform.numpy()
    else:
        waveform_np = waveform

    # Count samples that are at or above the clipping threshold
    return (
        np.sum(np.abs(waveform_np) >= 1.0).item()
        if isinstance(np.sum(np.abs(waveform_np) >= 1.0), torch.Tensor)
        else np.sum(np.abs(waveform_np) >= 1.0)
    )


def third_octave_bands(sr, fmin=20.0, fmax=None):
    """
    Compute 1/3-octave band center frequencies and edges.
    """
    if fmax is None:
        fmax = sr / 2.0

    k = np.arange(-30, 30)  # wide enough range
    f_center = 1000.0 * (2.0 ** (k / 3.0))  # ISO 1/3 octave centers
    f_center = f_center[(f_center >= fmin) & (f_center <= fmax)]

    f_lower = f_center / (2 ** (1 / 6))
    f_upper = f_center * (2 ** (1 / 6))
    return f_center, f_lower, f_upper


def third_octave_response_db(waveform: torch.Tensor, sr: int):
    """
    Compute 1/3 octave magnitude response in dB from waveform.

    Args:
        waveform (torch.Tensor): shape (n_samples,) or (1, n_samples)
        sr (int): sample rate

    Returns:
        freqs (np.ndarray): band center frequencies
        mags_db (torch.Tensor): band magnitudes in dB
    """
    if waveform.ndim > 1:
        waveform = waveform.squeeze(0)

    n = waveform.numel()
    spec = torch.fft.rfft(waveform)
    mag = torch.abs(spec) / n
    freqs = torch.fft.rfftfreq(n, d=1.0 / sr)

    # Get bands
    f_center, f_lower, f_upper = third_octave_bands(sr)
    band_mags = []
    for fl, fu in zip(f_lower, f_upper):
        idx = (freqs >= fl) & (freqs < fu)
        if idx.any():
            band_mags.append(mag[idx].mean())
        else:
            band_mags.append(torch.tensor(0.0))

    band_mags = torch.stack(band_mags)

    # Convert to dB (avoid log(0))
    mags_db = 20 * torch.log10(band_mags + 1e-12)

    return f_center, mags_db


def stereo_width(waveform: torch.Tensor):
    """
    Compute the stereo width of a waveform.
    """
    # can you implement this?
    # Assume waveform shape is (2, seq_len)
    if waveform.ndim != 2 or waveform.shape[0] != 2:
        raise ValueError(
            "waveform must have shape (2, seq_len) for stereo width calculation"
        )
    left = waveform[0]
    right = waveform[1]
    # Compute correlation coefficient between L and R
    left = left - left.mean()
    right = right - right.mean()
    numerator = (left * right).mean()
    denominator = torch.sqrt((left**2).mean() * (right**2).mean()) + 1e-12
    corr = numerator / denominator
    # Stereo width: 0 = mono, 1 = fully wide (L and R uncorrelated), -1 = fully out of phase
    width = torch.sqrt(1 - corr**2)
    return width.item()


def loudness(waveform: torch.Tensor, sample_rate: float):
    """
    Compute the loudness of a waveform.
    """
    meter = pyln.Meter(sample_rate)
    loudness = meter.integrated_loudness(waveform.permute(1, 0).numpy())
    return loudness


def analyze_audio_first_last(
    audio: torch.Tensor,
    sample_rate: int,
    segment_duration: int = 60,
):
    """
    Compute 1/3-octave response and stereo width for the first and last segments
    of a single audio tensor, and return their deltas.

    Args:
        audio: Tensor of shape (channels, samples) or (samples,)
        sample_rate: sample rate in Hz
        segment_duration: segment length in seconds for the first/last comparison

    Returns:
        {
            'center_freqs': np.ndarray,
            'third_octave_first': Tensor[dB],
            'third_octave_last': Tensor[dB],
            'delta_response': Tensor[dB],      # first - last
            'total_delta': Tensor[scalar],     # L1 magnitude of delta_response
            'stereo_width_first': float or None,
            'stereo_width_last': float or None,
            'stereo_width_delta': float or None,  # first - last
        }
    """
    if not isinstance(audio, torch.Tensor):
        raise TypeError("audio must be a torch.Tensor")
    if audio.numel() == 0:
        raise ValueError("audio is empty")

    # Determine segment length in samples (clip to available length)
    seg_len = min(audio.shape[-1], segment_duration * sample_rate)

    # Helper: mono mixdown for 1/3-octave analysis
    mono = audio.mean(dim=0) if audio.ndim > 1 else audio
    first_mono = mono[:seg_len]
    last_mono = mono[-seg_len:]

    # 1/3-octave responses (dB)
    center_freqs_first, oct_first = third_octave_response_db(first_mono, sample_rate)
    center_freqs_last, oct_last = third_octave_response_db(last_mono, sample_rate)

    # Centers should match; keep the first as canonical
    if (
        len(center_freqs_first) != len(center_freqs_last)
        or (center_freqs_first != center_freqs_last).any()
    ):
        raise RuntimeError(
            "Mismatched third-octave centers between first and last segments."
        )

    delta = oct_first - oct_last
    total_delta = delta.abs().sum()

    # Stereo width (if stereo input)
    def maybe_width(x: torch.Tensor):
        if x.ndim == 2 and x.shape[0] == 2 and x.shape[1] > 0:
            return stereo_width(x)
        return None

    first_full = audio[:, :seg_len] if audio.ndim == 2 else audio
    last_full = audio[:, -seg_len:] if audio.ndim == 2 else audio

    width_first = maybe_width(first_full)
    width_last = maybe_width(last_full)
    width_delta = None
    if (width_first is not None) and (width_last is not None):
        width_delta = float(width_first - width_last)

    # loudness
    loudness_first = loudness(first_full, sample_rate)
    loudness_last = loudness(last_full, sample_rate)
    loudness_delta = None
    if (loudness_first is not None) and (loudness_last is not None):
        loudness_delta = float(loudness_first - loudness_last)

    return {
        "center_freqs": center_freqs_first,  # np.ndarray
        "third_octave_first": oct_first,  # Tensor[dB]
        "third_octave_last": oct_last,  # Tensor[dB]
        "delta_response": delta,  # Tensor[dB]
        "total_delta": total_delta,  # Tensor[scalar]
        "stereo_width_first": width_first,  # float or None
        "stereo_width_last": width_last,  # float or None
        "stereo_width_delta": width_delta,  # float or None
        "loudness_first": loudness_first,  # float or None
        "loudness_last": loudness_last,  # float or None
        "loudness_delta": loudness_delta,  # float or None
    }


def analyze_audio(filepath):
    audio, sr = torchaudio.load(filepath)

    if sr != 48000:
        audio = torchaudio.functional.resample(audio, sr, 48000)

    # calculate loudness
    lufs_db = calculate_loudness(audio.permute(1, 0).numpy(), sr)
    lufs_db_factor = calculate_loudness_factor(audio.permute(1, 0).numpy(), sr)

    # calculate stereo width
    stereo_width = calculate_stereo_width(audio)

    # clipped samples
    clipped_samples = calculate_possible_clipped_samples(audio)

    return {
        "lufs_db": lufs_db,
        "lufs_db_factor": lufs_db_factor,
        "stereo_width": stereo_width,
        "clipped_samples": clipped_samples,
    }


def _reload_models_if_needed(hoot_filepath, hoot_tokenizer_filepath):
    model_list = load_model_list(
        checkpoint_filepath=hoot_filepath,
        tokenizer_filepath=hoot_tokenizer_filepath,
        n_gpus=None,
    )
    model = model_list[0]["model"]
    tokenizer = model_list[0]["tokenizer"]
    return model_list, model, tokenizer


class GenerateWorker:
    def __init__(
        self,
        dit_model_filepath: str,
        output_path: str,
        test_type: str,
    ):
        self.output_path = output_path
        self.test_type = test_type
        start_time = time.time()
        print("Start loading models")

        # load diffusion model
        num_gpus = torch.cuda.device_count()
        cuda_device = torch.cuda.current_device()
        print(f"Found {num_gpus} GPUs. Using GPU {cuda_device}.")

        tokenizer_filepath = "s3://suno-data/georg/models/tokenizers/tokenizer_60k.json"
        semantic_model_filepath = "s3://suno-data/georg/models/semantic/mert_25.pt"
        semantic_clusters_filepath = (
            "s3://suno-data/georg/models/semantic/mert_25_2x4k.npy"
        )
        codec_filepath = CODEC_FILEPATH
        ear_model_filepath = EAR_MODEL_FILEPATH

        _ = diffusion_gen.preload_dit_model(
            dit_model_filepath=dit_model_filepath,
            use_ema_if_exists=True,
            compile=False,
            weights_precision=torch.bfloat16,
        )
        _ = preload_tokenizer(tokenizer_filepath)
        _ = preload_semantic_models(semantic_model_filepath, semantic_clusters_filepath)
        _ = preload_codec_models(codec_filepath)
        # _ = preload_ear_model(ear_model_filepath)

        # print(f"CHUNK_SIZE_SCHEDULE: {CHUNK_SIZE_SCHEDULE}")
        # self.diffusion_engine = UpsampleEngine(chunk_size_schedule=CHUNK_SIZE_SCHEDULE)
        self.diffusion_engine = UpsampleEngine(min_chunk_size=BLOCK_SIZE)
        # self.diffusion_engine = UpsampleEngine(
        #    min_chunk_size=750, vae_version="v_vae_25_tuned_2"
        # )

        # load hoot model
        print("Loading hoot model")
        hoot_filepath = "s3://suno-data/christian/checkpoints/hoot_v3.pt"
        hoot_tokenizer_filepath = (
            "s3://suno-data/christian/checkpoints/hoot_v3_tokenizer_v3.pt"
        )
        self.hoot_model_list, self.hoot_model, self.hoot_tokenizer = (
            _reload_models_if_needed(hoot_filepath, hoot_tokenizer_filepath)
        )

        # load ear model
        ear_model_filepath = "s3://suno-data/christian/checkpoints/ear/ear_v2_s3080.pt"
        self.ear_model = load_model(ear_model_filepath, compile=True)

        print(
            f"Finish loading models. Took {round(time.time() - start_time, 2)} seconds"
        )

    @staticmethod
    def download_models(dit_model_filepath, dir_path=MOUNT_PATH):
        """Download diffusion models."""
        print("Start downloading models")
        _ = diffusion_gen.get_model_if_needed(
            CODEC_FILEPATH,
            cache_dir=dir_path,
        )
        _ = diffusion_gen.get_model_if_needed(
            diffusion_gen.SEMANTIC_MODEL_FILEPATH, cache_dir=dir_path
        )
        _ = diffusion_gen.get_model_if_needed(
            diffusion_gen.SEMANTIC_CLUSTERS_FILEPATH, cache_dir=dir_path
        )
        _ = diffusion_gen.get_model_if_needed(
            diffusion_gen.CODEC_FILEPATH, cache_dir=dir_path
        )
        _ = diffusion_gen.get_model_if_needed(dit_model_filepath, cache_dir=dir_path)

        # _ = chirp_v2._get_model_if_needed(gpt_ckpt, cache_dir=dir_path)
        # _ = chirp_v2._get_model_if_needed(chirp_v2.TOKENIZER_PATH, cache_dir=dir_path)

        print("Finish downloading models")

    def generate(self, work_item):
        global models

        """Generate audio from a work item."""
        # when saving to s3 we will use a structure like this:
        # self.output_path/
        #   item_id/
        #     original_semantic.npz
        #     0_generated_audio.mp3
        #     0_generated_semantic.npz
        #     0_metadata.json
        #     1_generated_audio.mp3
        #     1_generated_semantic.npz
        #     1_metadata.json
        #     ...
        #     metadata.json

        for item in work_item:
            item_id = item["id"]

            if "s3_filepath" in item:
                s3_filepath = item["s3_filepath"]
                # read audio from s3 and then semantic encode
                audio = Audio.from_s3(s3_filepath, n_channels=2)

                if SAVE_CODES:
                    # encode the vae latents
                    vae_latents_cycled = codec_encode(
                        audio.convert(
                            sample_rate=48_000, byte_width=2, n_channels=2
                        ).normalize_volume(target_db=-16)
                    )

                    with tempfile.TemporaryDirectory() as td:
                        vae_latents_path = os.path.join(td, f"{item_id}_cycled_vae.npz")
                        np.savez(vae_latents_path, vae_latents=vae_latents_cycled)

                        s3_filepath = os.path.join(
                            self.output_path,
                            f"{item_id}",
                            f"{item_id}_vae.npz",
                        )
                        s3_client.upload_file(
                            vae_latents_path,
                            "suno-data",
                            s3_filepath,
                            ExtraArgs={"ContentType": "application/octet-stream"},
                        )

                if CYCLE_ONLY:
                    continue

                # also encode semantic codes
                codes = encode_semantic(
                    audio.convert(sample_rate=24_000, byte_width=2, n_channels=1)
                ).astype(np.int64)

                if SAVE_CODES:
                    with tempfile.TemporaryDirectory() as td:
                        codes_path = os.path.join(td, f"{item_id}_semantic.npz")
                        np.savez(codes_path, codes=codes)

                        s3_filepath = os.path.join(
                            self.output_path,
                            f"{item_id}",
                            f"{item_id}_semantic.npz",
                        )
                        s3_client.upload_file(
                            codes_path,
                            "suno-data",
                            s3_filepath,
                            ExtraArgs={"ContentType": "application/octet-stream"},
                        )
            else:
                # load the semantic codes from s3
                s3_filepath = f"s3://suno-data-uploads/studio/uploads/{item_id}.npz"

                try:
                    data = read_from_s3(s3_filepath, read_f=np.load)
                except Exception as e:
                    print(f"Error loading {s3_filepath}: {e}")
                    continue

                if "v3.0_raw" in data:
                    codes = data["v3.0_raw"]
                elif "v3.5_raw" in data:
                    codes = data["v3.5_raw"]
                elif "v4.0_raw" in data:
                    codes = data["v4.0_raw"]
                elif "v4.5_raw" in data:
                    codes = data["v4.5_raw"]
                elif "v5.0_raw" in data:
                    codes = data["v5.0_raw"]
                else:
                    raise ValueError("No codes found")

            semantic_codes = torch.from_numpy(codes[:, 0]).long()  # .cuda()
            # semantic_codes = semantic_codes[:3000]
            print(semantic_codes.shape)

            # if tags is a list, join it into a string
            if isinstance(item["tags"], list):
                tags_str = ", ".join(item["tags"][:5])
            else:
                tags_str = item["tags"]

            if tags_str is None:
                tags_str = ""

            lyrics = item.get("lyrics", None)
            if lyrics is None:
                lyrics = item.get("text", None)

            if lyrics is None:
                lyrics = ""

            print(tags_str)
            print(lyrics)

            # create a list of configs to run
            gen_cfgs = []
            if self.test_type == "steps":
                diffusion_steps = [8, 10, 12, 14, 16, 18, 20]
                diffusion_seed = 42
                diffusion_text_cfg_coef = 2.0
                noise_ctx_level = NOISE_CTX_LEVEL
                noise_ctx_pad_len = NOISE_CTX_PAD_LEN
                rho = RHO
                sigma_min = SIGMA_MIN
                sigma_max = SIGMA_MAX

                for diffusion_steps in diffusion_steps:
                    gen_cfgs.append(
                        (
                            f"steps_{diffusion_steps}",
                            diffusion_gen.DiffusionGenerationConfig(
                                steps=diffusion_steps,
                                lyrics=lyrics,
                                tags=tags_str,
                                text_cfg_coef=diffusion_text_cfg_coef,
                                codec_scale_factor=CODEC_SCALE_FACTOR,
                                scale_ctx_vector=SCALE_CTX_VECTOR,
                                noise_ctx_level=noise_ctx_level,
                                noise_ctx_pad_len=noise_ctx_pad_len,
                                # sigma_min=sigma_min,
                                sigma_max=sigma_max,
                                seed=diffusion_seed,
                                # noise_schedule=NOISE_SCHEDULE,
                                # rank_candidates=RANK_CANDIDATES,
                            ),
                        )
                    )

            elif self.test_type == "standard":
                # diffusion parameters
                if RANDOMIZE_DIFFUSION_PARAMS:
                    diffusion_seed = np.random.randint(0, 1000000)
                    diffusion_steps = np.random.choice([8, 10, 12, 14, 16, 18, 20])
                    diffusion_text_cfg_coef = np.random.choice(
                        [1.0, 1.5, 1.75, 2.0, 2.5, 3.0]
                    )
                    noise_ctx_level = np.random.choice([0.0, 0.25, 0.5, 0.75, 1.0])
                    rho = np.random.choice([0.9, 1.0, 1.1])
                    sigma_min = np.random.choice([0.05, 0.1, 0.25, 0.5])
                    sigma_max = np.random.choice([40.0, 50.0, 60.0, 70.0, 80.0, 100.0])
                    noise_ctx_pad_len = np.random.choice([0, 10, 20, 30])
                else:
                    if DIFFUSION_SEED is None:
                        diffusion_seed = 0
                        for char in item_id:
                            diffusion_seed = (diffusion_seed * 31 + ord(char)) % 1000000
                        print(f"Generated seed from '{item_id}': {diffusion_seed}")
                    else:
                        diffusion_seed = DIFFUSION_SEED
                        print(f"Using seed from config: {diffusion_seed}")
                    diffusion_steps = DIFFUSION_STEPS
                    noise_ctx_level = NOISE_CTX_LEVEL
                    noise_ctx_pad_len = NOISE_CTX_PAD_LEN
                    rho = RHO
                    sigma_min = SIGMA_MIN
                    sigma_max = SIGMA_MAX
                    sampler_type = SAMPLER_TYPE

                for n in range(1):
                    # diffusion_seed = np.random.randint(0, 1000000)
                    gen_cfgs.append(
                        (
                            "",
                            diffusion_gen.DiffusionGenerationConfig(
                                steps=diffusion_steps,
                                lyrics=lyrics,
                                tags=tags_str,
                                text_cfg_coef=TEXT_CFG_SCALE,
                                ctx_cfg_coef=CTX_CFG_SCALE,
                                codec_scale_factor=CODEC_SCALE_FACTOR,
                                scale_ctx_vector=SCALE_CTX_VECTOR,
                                noise_ctx_level=noise_ctx_level,
                                noise_ctx_pad_len=noise_ctx_pad_len,
                                semantic_skip_factor=SEMANTIC_SKIP_FACTOR,
                                rho=rho,
                                sigma_min=sigma_min,
                                sigma_max=sigma_max,
                                seed=diffusion_seed,
                                objective=OBJECTIVE,
                                sampler_type=sampler_type,
                                # rank_candidates=RANK_CANDIDATES,
                            ),
                        )
                    )
            else:
                raise ValueError(f"Invalid test type: {self.test_type}")

            # now run the requests
            for name, gen_cfg in gen_cfgs:

                request = Request(
                    id="dummy",
                    generation_config=gen_cfg,
                    tokens=semantic_codes,
                    input_tokens_finished=True,
                    # ignore_history_every=IGNORE_HISTORY_EVERY,
                )
                result = self.diffusion_engine.run_request(request)
                vae_latents = torch.concat(result.vae_latents)
                upsampled_audio = decode_stream_to_full_audio(vae_latents)

                metadata = {
                    # "original_audio": item["s3_filepath"],
                    "id": item_id,
                    "text": lyrics,
                    "tags": tags_str,
                    "diffusion": {
                        "steps": int(gen_cfg.steps),
                        "seed": int(gen_cfg.seed),
                        "text_cfg_coef": float(gen_cfg.text_cfg_coef),
                        "noise_ctx_level": float(gen_cfg.noise_ctx_level),
                        "codec_scale_factor": float(gen_cfg.codec_scale_factor),
                        "scale_ctx_vector": gen_cfg.scale_ctx_vector,
                        # "rho": float(gen_cfg.rho),
                    },
                }

                # we want to save out
                # audio file of the final audio
                # npz of the estimated semantics
                # metadata json with the original prompt and tags
                # copy some stuff
                # we want to copy the original audio and the npz of the original semantics

                with tempfile.TemporaryDirectory() as td:

                    # save audio to s3
                    upsampled_audio_path = os.path.join(td, f"{item_id}_{name}.mp3")
                    upsampled_audio.write_hq_mp3(upsampled_audio_path)

                    # run hoot evaluation (CER)
                    # check if the lyrics are not empty
                    if lyrics != "":
                        upsampled_audio_obj = Audio.from_file(
                            upsampled_audio_path, n_channels=2
                        )
                        if False:
                            aligned_lyrics = encode_and_align(
                                upsampled_audio_obj,
                                prior_texts=lyrics.lower(),
                            )
                            print(aligned_lyrics)

                        out = encode(upsampled_audio_obj, return_logits=True)
                        basic_cleaned_lyrics = lyrics.lower()
                        decoded_preds = self.hoot_tokenizer.decode_logits(
                            out, prior_text=basic_cleaned_lyrics
                        )
                        true_text_norm = clean_text(basic_cleaned_lyrics)
                        cer_val = round(get_cer(true_text_norm, decoded_preds), 3)
                        metadata["hoot_cer"] = cer_val  # add cer to metadata
                        print("hoot cer: ", metadata["hoot_cer"])
                    else:
                        metadata["hoot_cer"] = None

                    # run ear evaluation (quality score)
                    # check if the length of the audio is greater than 5 seconds
                    if upsampled_audio.duration_s > 5.0:
                        quality_score = self.ear_model.get_score(upsampled_audio_path)
                        metadata["ear_score"] = (
                            quality_score  # add quality score to metadata
                        )
                        print("ear score: ", metadata["ear_score"])
                    else:
                        metadata["ear_score"] = None

                    # run shimmer score evaluation
                    shimmer_score = shimmerscore(upsampled_audio_path)
                    metadata["shimmer_score"] = shimmer_score

                    # run the other evals
                    audio_analysis = analyze_audio(upsampled_audio_path)
                    # add all the keys in audio_analysis to metadata
                    for key, value in audio_analysis.items():
                        metadata[key] = value

                    s3_filepath = os.path.join(
                        self.output_path,
                        f"{item_id}",
                        (
                            f"{item_id}_{CHECKPOINT_NAME}_{name}.mp3"
                            if name != ""
                            else f"{item_id}_{CHECKPOINT_NAME}.mp3"
                        ),
                    )
                    print(f"Uploading to {s3_filepath}")
                    s3_client.upload_file(
                        upsampled_audio_path,
                        "suno-data",
                        s3_filepath,
                        ExtraArgs={
                            "ContentType": "audio/mpeg",
                        },
                    )

                    if SAVE_CODES:
                        print("Saving vae latents...")
                        # save the vae latents
                        vae_latents_path = os.path.join(
                            td, f"{item_id}_upsampled_vae.npz"
                        )
                        np.savez(
                            vae_latents_path, vae_latents=vae_latents.cpu().numpy()
                        )

                        s3_filepath = os.path.join(
                            self.output_path,
                            f"{item_id}",
                            f"{item_id}_upsampled_vae.npz",
                        )
                        s3_client.upload_file(
                            vae_latents_path,
                            "suno-data",
                            s3_filepath,
                            ExtraArgs={
                                "ContentType": "application/octet-stream",
                            },
                        )

                    # save the metadata
                    metadata_path = os.path.join(td, f"{item_id}_{name}_metadata.npz")
                    np.savez(metadata_path, **metadata)

                    # move the metadata to s3
                    s3_filepath = os.path.join(
                        self.output_path,
                        f"{item_id}",
                        (
                            f"{item_id}_{CHECKPOINT_NAME}_{name}__metadata.npz"
                            if name != ""
                            else f"{item_id}_{CHECKPOINT_NAME}__metadata.npz"
                        ),
                    )
                    s3_client.upload_file(
                        metadata_path,
                        "suno-data",
                        s3_filepath,
                        ExtraArgs={
                            "ContentType": "application/octet-stream",
                        },
                    )


# prod diffusion model
# DIT_MODEL_FILEPATH = "s3://suno-data/tony/tmp/diff/dit_v3_dpo_t10_3k_5e6_b100_t25.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/georg/tmp/2b_prefix_ft.pt"
# 2b v45
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/dpo/v1_t6_1E6_beta100_n8_bt4.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t6_2E6_beta100_n8_bt4.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t6_5E6_beta100_n16_bt4.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t6_5E6_beta100_n8_bt4_30k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t10_5E6_beta100_n8_bt4_9k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v45_2b_shared_ctx_s3784.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_infill_t6_5E6_beta100_n16_bt4_9k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_infill_t11_5E6_beta100_n16_bt4_9k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_infill_t6_5E6_beta100_n16_bt4_6k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t12_5E6_beta100_n16_bt4_9k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t13_5E6_beta100_n16_bt4_9k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t13_5E6_beta100_n16_bt4_9k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t15_5E6_beta100_n16_bt4_9k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_10k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_20k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_infill_shared_base_200k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_infill_v1_t6_5E6_beta100_n16_bt4_30k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_100k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_5e5_150k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_infill_shared_base_600k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_t3_5e5_200k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_base_ear_sft_50k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_infill_v1_t17_5E6_beta100_n16_bt4_9k_repro_main_3k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t17_5E6_beta100_n16_bt4_9k_repro_mai_ema.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_infill_v1_t17_5E6_beta100_n16_bt4_9k_repro_main_3k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_base_ear_sft_100k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t17_5E6_beta1000_n16_bt4_9k_repro_main_6k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t17_5E6_beta1000_n16_bt4_9k_repro_main_3k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t17_5E6_beta1000_n16_bt4_2k_repro_main.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t6_5E6_beta1000_n16_bt4_2k_repro_main.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_base_ear_sft_t3+t5_40k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_t4_5e5_180k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_t4_5e5_270k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t6_2E6_beta100_n16_bt4_1k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_base_ear_sft_t3_190k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_v45_infill_ear_sft_3e5_t6_100k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_base_ear_sft_t3_210k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t17_2E6_beta100_n16_bt4_1p5k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t17_2E6_beta5000_n8_bt4_1p5k_fix.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_base_ear_sft_t3_230k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t6_2E6_beta5000_n8_bt4_3k_fix.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t6_5E6_beta100_n8_bt4_3k_acc16_fix_3k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t6_5E6_beta100_n8_bt4_3k_acc16_fix_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_v45_infill_ear_sft_3e5_t6_250k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t18_1E6_beta100_n32_bt4_3k_acc4_fix_12k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_sft_t6_dpo_t18_1E6_beta100_n32_bt4_1k_acc8_fix_8k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_sft_t6_dpo_t18_1E6_beta100_n32_bt4_1k_acc8_fix_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t18_1E6_beta100_n32_bt4_1k_acc8_fix_8k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_genius_t6_sampled_10k_5E6_beta100_n12_bt4_3k_acc4_fix_6k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_sft_t18_dpo_t18_5E6_beta100_n16_bt4_1k_acc8_fix_8k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/dit_v6_dpo_t11_9k_5e6_b100.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t6_5E6_beta100_n16_bt4.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t18_1E6_beta100_n32_bt4_1k_acc8_fix_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_sft_t6_dpo_t6_5E6_beta100_n32_bt4_1k_acc4_fix_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v2_t3_1E6_beta100_n16_bt4_acc4_3k_12k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_v45_infill_ear_sft_5e5_t6_c0_2_30k.pt"

# sft model
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/diff_v2_2b_2mil_ft_infill_20250421_v2.pt"

# current prod model
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v45_2b_step_2mil_ft_8k_infill_apr21_t1_18_cs0.pt"

# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t18+syn_5E6_beta100_n16_bt2_3k_acc_4_6k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t18+syn_5E6_beta100_n16_bt2_1k_acc_8_8k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_v45_infill_ear_sft_1e5_t7_70k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t18+syn_5E6_beta100_n16_bt2_1k_acc4_sft.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_syn_corr_5E6_beta100_n16_bt2_3k_acc4_sft_12k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_v45_infill_ear_sft_1e5_t7_380k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_1E6_beta100_n32_bt4_acc4_3k_12k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_1E6_beta100_n32_bt4_acc4_3k_3k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_5E6_beta100_n16_bt2_acc8_1k_sft_8k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_5E6_beta100_n16_bt2_acc8_1k_sft_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_syn_5E6_beta100_n16_bt2_acc8_1k_sft_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_syn_5E6_beta100_n16_bt2_acc8_1k_sft_0_1_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_5E6_beta100_n32_bt4_acc4_3k_1k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_1E6_beta100_n16_bt2_acc8_1k_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_1E6_beta100_n32_bt4_acc4_3k_1k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t2_1E6_beta100_n16_bt4_acc4_1k_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t2_1E6_beta100_n16_bt4_acc4_1k_1k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_ear_1E6_beta100_n4_bt2_acc8_1k_8k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_ear_neg_0_5_1E6_beta100_n4_bt2_acc8_1k_8k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_ear_1E6_beta100_n8_bt2_acc8_2k_16k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_apr21_t1_18_cs0_t1_ear_1E6_beta100_n4_bt2_acc8_4k_32k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_ear_1E6_beta100_n4_bt2_acc8_4k_32k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t6_1E6_beta100_n16_bt4_acc4_1k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t7_1E6_beta100_n16_bt4_acc4_1k_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_og_d3_t7_1E6_beta100_n16_bt4_acc4_1k_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_og_d3_t9_1E6_beta100_n16_bt4_acc4_1k_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_ear_t1_1E6_beta100_n4_bt2_acc8_4k_32k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t9_2E6_beta100_n16_bt4_acc4_3k_12k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t19_1E6_beta100_n16_bt4_acc4_3k_12k.pt"  #
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t10_1E6_beta100_n16_bt4_acc4_3k_12k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_resume_750k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t12_1E6_beta100_n16_bt4_acc4_3k_12k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t13_1E6_beta100_n16_bt4_acc4_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t15_1E6_beta100_n16_bt4_acc4_3k_12k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t16_1E6_beta100_n16_bt4_acc4_3k_10k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t17_1E6_beta100_n16_bt4_acc4_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t17_1E6_beta100_n16_bt4_acc4_3k_12k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t19_1E6_beta100_n8_bt4_acc8_3k_24k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t18_1E6_beta100_n8_bt4_acc8_3k_24k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t20_1E6_beta100_n16_bt16_acc4_6k_24k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t21_1E6_beta100_n16_bt16_acc4_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_resume_1_25m.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t10_1E6_beta100_n16_bt16_acc8_3k_24k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t19_5E6_beta100_n16_bt4_acc4_3k_12k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d4_2_5E6_beta100_n16_bt2_acc4_4k_16k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_4b_1m.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_resume_1_5m.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_4b_1_25m.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d4_3_5E6_beta100_n16_bt2_acc4_4k_16k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d4_t6_freeze_5E6_beta100_n16_bt2_acc4_4k_8k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_v45_infill_5e5_sft_log_snr_20k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_resume_1_75m.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_4b_resume_2m.pt"
# DIT_MODEL_FILEPATH = (
#    "s3://suno-data/christian/checkpoints/diffusion/4n_25hz_v45_distill_adv_100k.pt"
# )
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v45_2b_step_2mil_ft_8k_infill_apr21_d3_v10.pt"
# DIT_MODEL_FILEPATH = (
#    "s3://suno-data/christian/checkpoints/diffusion/4n_25hz_v45_distill_adv2.pt"
# )
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_infill_shared_flow_5e5_sft_50k.pt"
# DIT_MODEL_FILEPATH = (
#    "s3://suno-data/christian/checkpoints/diffusion/8n_2b_adv_bs2_N5_5e6_001_50k.pt"
# )
# DIT_MODEL_FILEPATH = (
#    "s3://suno-data/christian/checkpoints/diffusion/4n_2b_adv_bs2_N5_5e6_acc4_t1.pt"
# )
# /app/suno/modal/models/tony/tmp/diff/v45_2b_step_2mil_ft_8k_infill_apr21_t1_18_cs0.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_bootstrap_t2_1E6_beta100_n16_bt2_acc2_4k_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_bootstrap_t2_1E6_beta100_n16_bt2_acc2_4k_8k.pt"
DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d4_t39_1E6_beta100_n16_bt2_acc4_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_r1_1E6_beta100_n4_bt2_acc2_4k_8k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_infill_shared_flow_revert_resume_1_4m.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/4n_2b_flow_distill_bs1_N5_c5e5_g5e7_acc2_ema_alt_dmd_cfg2_400k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diff_v2_2b_2mil_ft_v0.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_resume_1_75m.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_2b_flow_sft_bs4_t8_5e6_300k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_infill_shared_flow_60s_180k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_bootstrap_t4_1E6_beta100_n8_bt2_acc2_4k_test_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_bootstrap_t4_1E6_beta100_n8_bt2_acc2_4k_test_9k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/4n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg2_cosine_210k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/4n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_2_cosine_nu1_0p1_160k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/4n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg2_cosine_400k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/4n_25hz_2b_flow_1e6_syn_sft_t1_30k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/4n_25hz_2b_flow_5e5_sft_t8_250k.pt"
#
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/4n_25hz_2b_flow_1e6_syn_sft_t2_7k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/4n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg2_cosine_700k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/4n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_2_cosine_nu1_0p1_500k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/4n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_1_cosine_nu1_0p1_residual_5k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_2_cosine_nu1_0p1_residual_6k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_2_cosine_nu1_0p1_residual_13k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_1p5_cosine_nu1_0p1_residual_10k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_1p5_cosine_nu1_0p1_residual_50k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_1p5_cosine_nu1_0p1_residual_70k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_1p25_cosine_nu1_0p1_residual_20k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_1E6_beta100_n8_bt2_acc2_4k_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_1p25_residual_aug_80k.pt"
# DIT_MODEL_FILEPATH = (
#    "s3://suno-data/christian/checkpoints/diffusion/4n_25hz_2b_flow_5e5_sft_t8_500k.pt"
# )
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_1E6_beta100_n8_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_1p5_cosine_nu1_0p1_residual_170k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_1p25_cosine_nu1_0p1_residual_120k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e7_alt_dmd_cfg_1p25_residual_sft_50k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e7_alt_dmd_cfg_1p25_residual_sft_100k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e7_alt_dmd_cfg_2_residual_sft_rm_50k.pt"
# DIT_MODEL_FILEPATH = (
#    "s3://suno-data/christian/checkpoints/diffusion/4n_25hz_2b_flow_5e5_sft_t8_725k.pt"
# )
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e7_alt_dmd_cfg_1p25_residual_sft_120k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_2b_flow_distill_bs1_N5_c5e5_g1e6_alt_dmd_cfg_2_residual_sft_65k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_2b_flow_distill_bs1_N5_c5e5_g1e6_alt_dmd_cfg_2_residual_sft_10k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_ctx_t0_1E6_beta100_n4_bt2_acc2_4k_4k.pt"
#
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_2b_flow_distill_bs1_N5_c5e5_g1e6_alt_dmd_cfg_2_residual_sft_40k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_ctx_t0_1E6_beta100_n4_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_2b_flow_distill_bs1_N5_c5e5_g1e6_alt_dmd_cfg_2_residual_sft_125k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_t1_1E6_beta100_n4_bt2_acc2_8k_6k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_t0_1E6_beta100_n4_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_2b_flow_distill_bs1_N5_c5e5_g1e6_alt_dmd_cfg_2_residual_sft_220k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_v1_t18_1E6_beta100_n4_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_v2_d4_v39_1E6_beta100_n4_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s1039_v1_t18_1E6_beta100_n4_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s1039_v1_t18_1E6_beta100_n15_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s1039_lm_t0_1E6_beta100_n4_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/4n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_residual_sft_45k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s1039_lm_t1_0_6_1E6_beta100_n4_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s1039_lm_t1_0_6_5E7_beta100_n4_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s1039_lm_t1_0_6_0822_5E7_beta100_n8_bt2_acc2_4k_8k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s1039_lm_t1_0_6_0822_5E7_beta100_n8_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_lm_t1_0_6_0822_5E7_beta100_n8_bt2_acc2_4k_last.pt"
#
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_lm_t1_0_6_0822_1_sft_0p1_5E7_beta100_n8_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_lm_t1_0_6_0823_0_0p1_5E7_beta100_n8_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_lm_t1_0_6_0824_0_1E6_beta100_n8_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_lm_t1_0_7_cut_0824_0_1E6_beta100_n8_bt2_acc4_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_lm_t1_0_7_cut_0824_0_1E6_beta100_n8_bt2_acc4_2k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_lm_t1_0_7_cut_history_4x_1E6_beta100_n8_bt2_acc4_2k_last.pt"
#
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_lm_t1_0_7_cut_history_4x_1E6_beta100_n8_bt2_acc4_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_carp_t1_v1_1E6_beta100_n4_bt2_acc4_3k_last.pt"
# DIT_MODEL_FILEPATH = "3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_lm_t1_0p6_cut_history_4x_1E6_beta100_n4_bt2_acc4_2k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3395_rd2_t1_merge_1E6_beta100_n8_bt2_acc4_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_rd2_lm_t2_0_7_cut_history_4x_1E6_beta100_n8_bt2_acc4_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_rd2_merge_t1_1E6_beta100_n8_bt2_acc4_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_rd2_hoot_t0_1E6_beta100_n8_bt2_acc4_3k_12k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_rd2_merge_t1_1E6_beta100_n8_bt2_acc4_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_rd2_hoot_t0_1E6_beta100_n8_bt2_acc4_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_rd2_hoot_t1_1E6_beta100_n16_bt2_acc1_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_rd2_hoot_t1_1E6_beta50_n16_bt2_acc1_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_rd2_hoot_t1_1E6_beta25_n16_bt2_acc1_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_rd2_hoot_t2_1E6_beta100_n16_bt2_acc1_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_rd2_hoot_t2_1E6_beta100_n16_bt2_acc1_6k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_residual_sft_t8_200k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s1039_rd2_lm_t2_1E6_beta100_n16_bt2_acc1_6k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3242_rd1_v1_t18_1E6_beta100_n16_bt2_acc1_6k_last.pt"
# DIT_MODEL_FILEPATH = (
#    "s3://suno-data/christian/checkpoints/diffusion/4n_25hz_2b_flow_5e5_sft_t8_500k.pt"
# )
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_v1_t18_1E6_beta100_n16_bt2_acc2_4k_4k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_v1_t18_1E6_beta100_n16_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t2_merge_1E6_beta100_n16_bt2_acc2_4k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_v1_t18_1E6_beta100_n16_bt2_acc4_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_v1_t18_5E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d0_1E6_beta100_n16_bt2_9k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d1_1E6_beta100_n16_bt2_9k_3k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d1_1E6_beta100_n16_bt2_9k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d3_1E6_beta100_n8_bt2_9k_3k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d4_1E6_beta100_n14_bt2_9k_2k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d4_5E7_beta100_n14_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d5_5E7_beta100_n16_bt2_6k_3k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d5_5E7_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d6_5E7_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d6_5E7_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d8_5E7_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d9_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d9_1E6_beta100_n16_bt2_9k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d11_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d10_1E6_beta100_n16_bt2_9k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d11_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d11_1E6_beta100_n16_bt2_1p5k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d12_1E6_beta100_n16_bt2_9k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d13_1E6_beta100_n14_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d14_1E6_beta100_n14_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d15_1E6_beta100_n16_bt2_9k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d15_1E6_beta50_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d16_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d18_1E6_beta100_n12_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d18_1E6_beta200_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d19_1E6_noise1_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d20_1E6_noise1_beta100_n16_bt2_1p5k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d21_1E6_noise1_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d21_5E7_noise1_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d21_5E7_noise1_beta100_n8_bt2_1k_500.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d22_1E6_noise1_beta100_n8_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d22_1E6_noise1_beta100_n8_bt2_1k_500.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_d23_1E6_noise1_beta100_n8_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_base_rd2_d24_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_base_rd2_d24_1E6_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_base_rd1_d25_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d25_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_base_rd1_d26_1E6_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_base_rd1_d27_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d28_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d30_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d30_1E6_beta100_n16_bt2_9k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d31_1E6_beta100_n8_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d32_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_1p5_residual_sft_t8_rm_s4862_10k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d33_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d34_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d34_5E7_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d34_5E7_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_diffv2_v1_t18_1E6_beta100_n16_bt2_sft0p01_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_d35_1E6_beta200_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_pair_t0_1E6_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_pair_t1_1E6_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_pair_t2_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_pair_t2_1E6_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_pair_t2_1E6_beta100_n16_bt2_noise_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_pair_t3_1E6_beta100_n16_bt2_noise_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_pair_t4_1E6_beta100_n16_bt2_noise_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_pair_t5_1E6_beta100_n16_bt2_noise_1k_last.pt"
# IT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_pair_t5_1E6_beta100_n16_bt2_noise_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_pair_t5_1E6_beta200_n16_bt2_noise_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_pair_t5_1E6_beta100_n16_bt2_noise_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_pair_t9_1E6_beta100_n8_bt2_noise_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_pair_t9_1E6_beta200_n16_bt2_noise_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_pair_t10_1E6_beta100_n16_bt2_noise_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_t3_pair_t10_1E6_beta100_n16_bt2_noise_500_last.pt"
# DIT_MODEL_FILEPATH = (
#    "s3://suno-data/christian/checkpoints/diffusion/v45_2b_step_2_600_000.pt"
# )
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_resume_1_75m.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t14_1E6_sft_1E4_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t14_1E6_sft_only_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t14_1E6_sft_1E4_beta100_n16_bt2_9k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_t14_1E6_sft_1E4_beta100_n16_bt2_3k_last.pt"
# $DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_t14_1E6_sft_1E4_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_t15_1E6_sft_1E4_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t14_1E6_sft_1E3_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd2_t15_1E6_sft_1E2_beta100_n16_bt2_acc4_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t15_1E6_sft_1E2_beta100_n16_bt2_acc4_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g2e6_alt_dmd_cfg_1p5_residual_s3431_100k.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t16_1E6_sft_1E2_beta100_n16_bt2_acc4_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t17_1E6_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t18_1E6_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t19_1E6_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t19_1E6_beta100_n16_bt2_3k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t20_1E6_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t21_1E6_beta100_n16_bt2_1k_last.pt"
# DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_bass_d5_t10_1E6_beta100_n16_bt2_acc4_noise1_3k_off_d4_v40_last.pt"
DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_sft_t8_rd1_t22_1E6_beta100_n16_bt2_3k_last.pt"

CHECKPOINT_NAME = DIT_MODEL_FILEPATH.split("/")[-1].replace(".pt", "")

# OUTPUT_STR = "farm-speaker-desk"
# OUTPUT_STR = "pencil-helmet-window"
# OUTPUT_STR = "walking-flowers"
# OUTPUT_STR = "reference-audio"
# OUTPUT_STR = "auk-test"
# OUTPUT_STR = "auk-clips-up-u-2"
OUTPUT_STR = "bluejay-lang-balanced-200"
# OUTPUT_STR = "v2-infill-data-v1"
# OUTPUT_STR = "v2-diff-step-test"
# OUTPUT_STR = "real-audio-test"
TEST_TYPE = "standard"  # standard, steps
SAVE_CODES = False
CYCLE_ONLY = False
GPU_TYPE = "H100"  # H100, A10G
OBJECTIVE = "rectified_flow" if "flow" in DIT_MODEL_FILEPATH else "v"
SAMPLER_TYPE = "pingpong" if "distill" in DIT_MODEL_FILEPATH else "dpmpp"

# extra params
RHO = 1.0
SIGMA_MIN = 0.5
SIGMA_MAX = 50.0
BLOCK_SIZE = 25 * 30  # 30s chunks
DIFFUSION_SEED = None  # 42 is default
RANDOMIZE_DIFFUSION_PARAMS = False
USE_ALTERNATE_CFG = False
TEXT_CFG_SCALE = 2.0 if "distill" not in DIT_MODEL_FILEPATH else 1.0
CTX_CFG_SCALE = 1.0
CTX_STEP_PERCENT = 1.0
RANK_CANDIDATES = 1

# IGNORE_HISTORY_EVERY = 2
EAR_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/ear/ear_v3_s8558.pt"

# use this to add a suffix to the checkpoint name
EXTRA_NAME_SUFFIX = "_step10" if "distill" not in DIT_MODEL_FILEPATH else "_step2"

CHECKPOINT_NAME = CHECKPOINT_NAME + EXTRA_NAME_SUFFIX

if "dit_v6_dpo_t11_9k_5e6_b100" in DIT_MODEL_FILEPATH:
    CODEC_SCALE_FACTOR = 2.5
    SCALE_CTX_VECTOR = False
    NOISE_CTX_LEVEL = 0.0
    NOISE_CTX_PAD_LEN = 0
    DIFFUSION_STEPS = 10
    SEMANTIC_SKIP_FACTOR = 1

    CODEC_FILEPATH = "s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth"

    from suno_utils.tasks.dac_vae_100hz_peaq import (
        preload_models as preload_codec_models,
        decode as codec_decode,
        encode as codec_encode,
        decode_stream_to_full_audio,
    )

else:
    CODEC_SCALE_FACTOR = 0.4
    SCALE_CTX_VECTOR = True
    NOISE_CTX_LEVEL = 0.75
    NOISE_CTX_PAD_LEN = 0
    DIFFUSION_STEPS = 10 if "distill" not in DIT_MODEL_FILEPATH else 2
    SEMANTIC_SKIP_FACTOR = 1
    NOISE_SCHEDULE = "polyexponential"
    # NOISE_SCHEDULE = "cosine"
    # CHUNK_SIZE_SCHEDULE = [10 * 25, 20 * 25, 30 * 25]  # 10s, 20s, 30s chunks
    # CHUNK_SIZE_SCHEDULE = [30 * 25, 30 * 25, 30 * 25]  # 30s chunks

    CODEC_FILEPATH = "s3://suno-data/minz/models/dac_vae_tuned_25hz.pth"

    from suno_utils.tasks.dac_vae_fixed_25hz import (
        preload_models as preload_codec_models,
        decode as codec_decode,
        encode as codec_encode,
        decode_stream_to_full_audio,
    )


def download_model_wrapper_d():
    # this print is necessary to have modal rerun this when MODEL changes
    # Modal tracks referenced global variables
    # Change the name of the function to force a rerun
    print("Downloading model for history encoder")
    GenerateWorker.download_models(DIT_MODEL_FILEPATH)


image = base_image.run_function(download_model_wrapper_d, secrets=SECRETS)
APP_NAME = f"batch-eval-diffusion"
app = modal.App(APP_NAME, image=image, secrets=SECRETS)

N_MAX_REPLICAS = 16


@app.cls(
    # gpu=modal.gpu.H100(count=1),
    gpu=GPU_TYPE,
    cpu=4,
    secrets=SECRETS,
    timeout=2 * 60 * 60,
    container_idle_timeout=240,
    # mounts=MODAL_MOUNTS,
    memory=15000,
    concurrency_limit=N_MAX_REPLICAS,
)
class GenerateStub:
    def __init__(self, dit_ckpt: str, output_path: str, test_type: str):
        import torch

        num_gpus = torch.cuda.device_count()
        print(f"Found {num_gpus} GPUs.")
        self.worker = GenerateWorker(
            dit_ckpt,
            output_path,
            test_type,
        )

    @modal.method()
    def generate(self, work_item: list[dict]):
        return self.worker.generate(work_item)


@app.local_entrypoint()
def main():
    # checkpoints in s3://suno-data/christian/checkpoints
    # prompts in s3://suno-data/christian/prompts
    # outputs in s3://suno-data/christian/outputs

    # load the positive prompts
    # work_items = read_from_s3(
    #    "s3://suno-data/christian/sft/pos_interesting_clips_up_u_1_20241201_full.jsonl",
    #    read_f=read_jsonl,
    # )

    # work_items = read_jsonl(
    #    "/home/christian/code/christian/metadata/sft/auk_clips_up_u_1_20241201_pos.jsonl",
    # )

    work_items = read_jsonl(
        "/home/christian/code/christian/metadata/sft/interesting_clips_bluejay_t1_20250811_public_only_lang_balanced_200.jsonl",
    )

    # work_items = read_jsonl(
    #    "/home/christian/audio/reference-audio.jsonl",
    #  )

    # work_items = read_jsonl(
    #    "/home/christian/code/christian/metadata/discogs_subset_sampled_metas.jsonl"
    #

    # work_items = read_jsonl(
    #    "/home/christian/code/christian/metadata/ear/genius_t6_sampled_10k.jsonl"
    # )

    # just do the first 128
    work_items = work_items  # [:128]
    print(f"Total work items: {len(work_items)}")

    # first check for existing ids in the output path
    # existing_ids = list_s3_dir(
    #    f"s3://suno-data/christian/outputs/{OUTPUT_STR}/",
    # )

    # existing_ids = [os.path.dirname(f[0]).split("/")[-1] for f in existing_ids]
    # existing_ids = list(set(existing_ids))
    # print(f"Total existing ids: {len(existing_ids)}")

    ## now remove these from the work items
    # work_items = [f for f in work_items if f["id"] not in existing_ids]
    print(f"Total work items remaining: {len(work_items)}")

    chunksize = 8  # num of prompts per worker

    worker = GenerateStub(
        DIT_MODEL_FILEPATH, f"christian/outputs/{OUTPUT_STR}", test_type=TEST_TYPE
    )

    work_items = list(funcy.chunks(chunksize, work_items))
    print(f"Chunksize: {chunksize}, total chunks: {len(work_items)}")

    # print("Testing inference...")
    # t0 = time.time()
    # for work_item in work_items[:1]:
    #    _ = worker.generate.remote(work_item)
    # print(f"{int(round(time.time()-t0))}s for test")

    print("Running batch inference...")
    t0 = time.time()
    _ = list(worker.generate.map(work_items))
    print(round((time.time() - t0) / 60 / 60), "h total for batch generation")
