from collections import defaultdict
from dataclasses import dataclass, field
from enum import Enum
from suno_utils.worker.modal_model_volume import (
    MODEL_STORE_VOLUME_PREFIX,
    GPT_PATH_DICT,
    DIFFUSION_PATH_DICT,
    CODEC_PATH_DICT,
    DIFFUSION_DEFAULT_V1_PATH,
)


class VAEVersion(str, Enum):
    """VAE version constants with metadata."""

    V_VAE_25_TUNED_2 = "v_vae_25_tuned_2"
    V_VAE_25_PEAQ_1 = "v_vae_25_peaq_1"

    @property
    def gain_adjust(self) -> float:
        """Get gain adjustment for this VAE version."""
        return {
            VAEVersion.V_VAE_25_TUNED_2: 1.28,
            VAEVersion.V_VAE_25_PEAQ_1: -1.61,
        }.get(self, 0.0)

    @property
    def sample_rate(self) -> int:
        """Get sample rate in Hz."""
        return 25  # All current versions use 25Hz

    @property
    def description(self) -> str:
        """Get human-readable description."""
        return {
            VAEVersion.V_VAE_25_TUNED_2: "25hz tuned VAE version 2. No shimmer.",
            VAEVersion.V_VAE_25_PEAQ_1: "25hz PEAQ VAE version 1 (default). Can have shimmer.",
        }.get(self, "Unknown VAE version")

    @classmethod
    def from_string(cls, value: str) -> "VAEVersion":
        """Get VAEVersion from string value."""
        for version in cls:
            if version.value == value:
                return version
        else:
            raise ValueError(f"Unknown VAE version: {value}")


def get_model_version(model: str) -> str:
    """Get the model version based on the model name.

    This is pre-defined. Based on the FE-side model name, and worker's model name.
    Note that these are the descrete code version. Only semantic and coarse (if exists).

    Version format: arch.major.minor.ft
       arch: for model architecture changes. i.e. 2 for 3b model, 3 for 7b model, etc.
       major: major model version, for changes which WILL cause model input/output structure changes.
       minor: minor model version, for other changes which will NOT cause model input/output structure changes.
       ft: fine-tune version
    """
    if "up" in model or model.startswith("diff"):
        # diff conditioned only on semantic
        # the output is cycled to generate semantic and coarse
        return "4.0.0.0"
    # semantic only + diffusion v2 -- should only have semantic, no coarse
    if (
        "sem" in model
        or "auk" in model
        or "bluejay" in model
        or "6b" in model
        or "stem" in model
        or "ahi" in model
    ):
        # this is a special case for 30b infilling
        if "infill" in model:
            return "4.0.0.0"
        return "5.0.0.0"
    # accidentally added 30b as v4
    # otherwise, v4 is using diffusion models
    elif "30b" in model or "v4" in model:
        return "4.0.0.0"
    # sem + dac 12c + 13b -- 13b is marked as 3.5.0.0
    elif "13b" in model or "v3-5" in model or "v3p5" in model:
        return "3.5.0.0"
    # sem + dac 12c + 7b -- 7b is marked as 3.0.0.0
    elif "7b" in model or "v3" in model:
        return "3.0.0.0"
    # sem + dac 8c -- v2 32b is marked as 2.0.0.0
    elif "v2" in model:
        return "2.0.0.0"
    else:
        raise ValueError(f"Unknown model {model}")


def is_vae_model(model: str) -> bool:
    """Returns if the model is using diffusion or not."""
    if (
        "v4" in model
        or "-h-" in model
        or "stem" in model
        or "seed" in model
        or "auk" in model
        or "ahi" in model
        or "bluejay" in model
    ):
        return True
    return False


def get_vae_version(model: str) -> VAEVersion:
    """Returns the diffusion codec version."""
    if (
        "sem" in model
        or "diff_v2" in model
        or "auk" in model
        or "bluejay" in model
        or "-ahi" in model
        or "stem" in model
        or "seeds" in model
        or "chirp-v4-up-u-d-2" in model  # this is a sneaky data collection model name
        or "30b_t6_infill" in model  # this is putting 30b to do auk infilling
    ):
        return VAEVersion.V_VAE_25_TUNED_2
    else:
        return VAEVersion.V_VAE_25_PEAQ_1  # 25 hz -- default for now


def get_vae_version_gain_adjust(vae_version: str | VAEVersion | None) -> float:
    if isinstance(vae_version, VAEVersion):
        return vae_version.gain_adjust
    elif isinstance(vae_version, str):
        # Try to convert string to VAEVersion for backward compatibility
        vae_enum = VAEVersion.from_string(vae_version)
        if vae_enum:
            return vae_enum.gain_adjust
        # Fallback for string comparison (for backward compatibility)
        if vae_version == VAEVersion.V_VAE_25_TUNED_2.value:
            return 1.28
        elif vae_version == VAEVersion.V_VAE_25_PEAQ_1.value:
            return -1.61
    return 0.0


@dataclass
class ModalEngineConfig:
    duration: int = 120
    model: str = "7b_special"
    compile: bool = True
    gpu: str = "A100"
    model_concurrency: int = 1  # maximum number of clips to process in parallel
    upsample_concurrency: int = 48
    decoder_concurrency: int = 60
    max_sequences: int = 6  # maximum number of streams that we reserve
    max_length_s: int = 120
    max_history_duration_s: int = 60
    world_size: int = 1
    batch_increment: int = 16
    max_batch_size: int = 128
    min_batch_size: int = None

    keep_warm_gpt: dict = field(default_factory=lambda: defaultdict(int, {"dev": 0, "prod": 1}))
    keep_warm_decoder: dict = field(default_factory=lambda: defaultdict(int, {"dev": 0, "prod": 1}))
    keep_warm_upsample: dict = field(default_factory=lambda: defaultdict(int, {"dev": 0, "prod": 1}))
    keep_warm_buffer: dict = field(default_factory=lambda: defaultdict(int, {"dev": 0, "prod": 1}))

    def __post_init__(self):
        if self.world_size > 1:
            assert self.compile, "TP without compile is slow"

    @property
    def model_version(self):
        """
        Model version, used for logging and compatibility checks.
        """
        return get_model_version(self.model)

    @property
    def vae_version(self):
        """VAE version, used for logging and compatibility checks."""
        return get_vae_version(self.model)

    @property
    def gain_adjust(self):
        return get_vae_version_gain_adjust(self.vae_version)

    @property
    def model_type(self):
        if "v2" in self.model:
            return "v2"
        elif "7b" in self.model:
            return "7b"
        elif "13b" in self.model:
            return "13b"
        elif "30b" in self.model:
            return "30b"
        elif "sem" in self.model:
            return "sem"
        else:
            raise ValueError(f"Unknown model {self.model}")

    @property
    def gpt_ckpt_path(self) -> str:
        path_dict = {
            model_name: f"{MODEL_STORE_VOLUME_PREFIX}{path}"
            for model_name, path in GPT_PATH_DICT.items()
        }
        return path_dict[self.model]

    @property
    def diffusion_ckpt_path(self) -> str:
        path_dict = {
            model_name: f"{MODEL_STORE_VOLUME_PREFIX}{path}"
            for model_name, path in DIFFUSION_PATH_DICT.items()
        }
        # :meltingface:
        return path_dict.get(self.model, f"{MODEL_STORE_VOLUME_PREFIX}{DIFFUSION_DEFAULT_V1_PATH}")

    @property
    def codec_ckpt_path(self) -> str:
        """Codec checkpoint path correlated with diffusion versions."""
        return f"{MODEL_STORE_VOLUME_PREFIX}{CODEC_PATH_DICT[self.vae_version.value]}"


MODEL_CONFIG_DICT = {
    "v2": ModalEngineConfig(
        model="v2",  # v2, v2_ft13 for prod for now
        duration=80,
        max_history_duration_s=20,
        max_sequences=20,
        model_concurrency=16,
        max_batch_size=24,
        batch_increment=8,
        compile=True,
        gpu="H100",  # A100 same setting
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 1}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 1}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
    ),
    "7b": ModalEngineConfig(
        model="7b",  # 7b model
        duration=120,
        max_history_duration_s=60,
        max_sequences=40,
        model_concurrency=12,
        compile=True,
        gpu="H100",  # A100 40 same setting
    ),
    "7b_dpo": ModalEngineConfig(
        model="7b_dpo",  # 7b_dpo model
        duration=120,
        max_history_duration_s=60,
        max_sequences=40,
        model_concurrency=8,
        compile=True,
        gpu="H100",  # A100 40 same setting
    ),
    "7b_ipo": ModalEngineConfig(
        model="7b_ipo",  # 7b_ipo model
        duration=120,
        max_history_duration_s=60,
        max_sequences=72,
        model_concurrency=64,
        compile=True,
        gpu="H100",  # A100 40 same setting
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 1}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 1}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
    ),
    "13b_ft_1": ModalEngineConfig(
        model="13b_ft_1",  # 13b_ft dpo model, 2 heads
        duration=240,
        max_history_duration_s=120,
        max_sequences=80,  #
        model_concurrency=60,  #
        min_batch_size=60 + 24,
        max_batch_size=60 + 24,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 1}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 1}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
    ),
    "13b_special_8": ModalEngineConfig(
        model="13b_special_8",  # 13b_special model
        duration=240,
        max_history_duration_s=120,
        max_sequences=36,  # max 40
        max_batch_size=72 + 28,
        model_concurrency=80,
        min_batch_size=72 + 28,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 20}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 10}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
    ),
    "13b_tech": ModalEngineConfig(
        model="13b_tech",  # 13b_special model
        duration=240,
        max_history_duration_s=120,
        max_sequences=36,  # max 40
        max_batch_size=56 + 24,
        model_concurrency=64,
        upsample_concurrency=56,
        min_batch_size=56 + 24,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 3}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 3}),
        keep_warm_upsample=defaultdict(int, {"dev": 0, "prod": 3}),
    ),
    # this is the A100 80GB config
    # "13b_special_8": ModalEngineConfig(
    #     model="13b_special_8",  # 13b_special model
    #     duration=240,
    #     max_history_duration_s=120,
    #     max_sequences=36,
    #     model_concurrency=32,
    #     compile=False,
    #     gpu="A100-80GB",  # A100 - 36/28/2, A10G
    #     world_size=1,
    #     keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 1}),
    #     keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 1}),
    # ),
    "13b_short": ModalEngineConfig(
        model="13b_short",  # 13b_special model off 13b data
        duration=30,
        max_history_duration_s=10,
        max_sequences=40,
        max_batch_size=40 + 16,
        model_concurrency=40,
        min_batch_size=40,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 1}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 1}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
    ),
    "13b_special_31": ModalEngineConfig(
        model="13b_special_31",  # 13b_special model
        duration=240,
        max_history_duration_s=120,
        max_sequences=36,  # max 40
        max_batch_size=48 + 24,
        model_concurrency=48,
        min_batch_size=48 + 24,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 1}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 1}),
    ),
    # main prod worker, high keep warm
    "13b_special_32": ModalEngineConfig(
        model="13b_special_32",  # 13b_special model
        duration=240,
        max_history_duration_s=120,
        max_sequences=36,  # max 40
        max_batch_size=108,
        model_concurrency=88,
        upsample_concurrency=80,
        min_batch_size=108,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 20}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 20}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 20}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
    ),
    "13b_special_33": ModalEngineConfig(
        model="13b_special_33",  # 13b_special model
        duration=240,
        max_history_duration_s=120,
        max_sequences=36,  # max 40
        max_batch_size=56 + 24,
        model_concurrency=56,
        min_batch_size=56 + 24,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 5}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 5}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 5}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
    ),
    "13b_special_32_fast": ModalEngineConfig(
        model="13b_special_32_fast",  # 13b_special model
        duration=240,
        max_history_duration_s=120,
        max_sequences=36,  # max 40
        max_batch_size=56 + 24,
        model_concurrency=56,
        min_batch_size=56 + 24,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 1}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 1}),
        keep_warm_upsample=defaultdict(int, {"dev": 0, "prod": 1}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
    ),
    "13b_special_test": ModalEngineConfig(
        model="13b_special_test",  # 13b_special model
        duration=240,
        max_history_duration_s=120,
        max_sequences=36,  # max 40
        max_batch_size=56 + 24,
        model_concurrency=56,
        min_batch_size=56 + 24,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_upsample=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
    ),
    "13b_special_test_2": ModalEngineConfig(
        model="13b_special_test_2",  # 13b_special model
        duration=240,
        max_history_duration_s=120,
        max_sequences=36,  # max 40
        max_batch_size=56 + 24,
        model_concurrency=56,
        min_batch_size=56 + 24,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_upsample=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
    ),
    "13b_special_test_3": ModalEngineConfig(
        model="13b_special_test_3",  # 13b_special model
        duration=240,
        max_history_duration_s=120,
        max_sequences=36,  # max 40
        max_batch_size=56 + 24,
        model_concurrency=56,
        min_batch_size=56 + 24,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_upsample=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
    ),
    "13b_special_test_diff": ModalEngineConfig(
        model="13b_special_test_diff",  # 13b_special model
        duration=240,
        max_history_duration_s=120,
        max_sequences=36,  # max 40
        max_batch_size=56 + 24,
        model_concurrency=56,
        min_batch_size=56 + 24,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_upsample=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
    ),
    "13b_special_test_tech_1": ModalEngineConfig(
        model="13b_special_test_tech_1",  # 13b_special model
        duration=240,
        max_history_duration_s=120,
        max_sequences=36,  # max 40
        max_batch_size=56 + 24,
        model_concurrency=56,
        min_batch_size=56 + 24,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_upsample=defaultdict(int, {"dev": 0, "prod": 5}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
    ),
    "13b_upload_4": ModalEngineConfig(
        model="13b_upload_4",  # 13b_special model off 13b data
        duration=240,
        max_history_duration_s=60,
        max_sequences=36,  # max 40
        max_batch_size=56 + 24,
        model_concurrency=56,
        min_batch_size=56 + 24,
        compile=True,
        gpu="H100",  # A100 - 36/28/2, A10G
        world_size=1,
        batch_increment=16,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 2}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 2}),
    ),
    "30b": ModalEngineConfig(
        model="30b",
        duration=240,
        max_history_duration_s=120,
        max_sequences=60,  # can do 88 but that's pushing...
        model_concurrency=40,  # 4 stream 25, 2 stream 50, 1 stream 90
        compile=True,
        gpu="H200",  # A100 32/16/4, A10G
        world_size=1,
        batch_increment=32,
        upsample_concurrency=60,
        min_batch_size=64,
        max_batch_size=64,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 2}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 2}),
    ),
    # # main prod worker, high keep warm
    "30b_t6": ModalEngineConfig(
        model="30b_t6",
        duration=240,
        max_history_duration_s=120,
        max_sequences=60,  # can do 88 but that's pushing...
        model_concurrency=40,  # 4 stream 25, 2 stream 50, 1 stream 90
        compile=True,
        gpu="H200",  # A100 32/16/4, A10G
        world_size=1,
        batch_increment=32,
        upsample_concurrency=60,
        min_batch_size=64,
        max_batch_size=64,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 3}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 3}),
        keep_warm_upsample=defaultdict(int, {"dev": 0, "prod": 3}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
    ),
    "30b_t6_test": ModalEngineConfig(
        model="30b_t6_test",
        duration=240,
        max_history_duration_s=120,
        max_sequences=80,  # can do 88 but that's pushing...
        model_concurrency=40,  # 4 stream 25, 2 stream 50, 1 stream 90
        compile=True,
        gpu="H100",  # A100 32/16/4, A10G
        world_size=2,
        batch_increment=32,
        upsample_concurrency=72,
        min_batch_size=60,
        max_batch_size=60 + 32,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 2}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 2}),
        keep_warm_upsample=defaultdict(int, {"dev": 0, "prod": 2}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
    ),
    "30b_t6_tech": ModalEngineConfig(
        model="30b_t6_tech",
        duration=240,
        max_history_duration_s=120,
        max_sequences=60,  # can do 88 but that's pushing...
        model_concurrency=40,  # 4 stream 25, 2 stream 50, 1 stream 90
        compile=True,
        gpu="H200",  # A100 32/16/4, A10G
        world_size=1,
        batch_increment=32,
        upsample_concurrency=60,
        min_batch_size=64,
        max_batch_size=64,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 3}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 3}),
    ),
    "30b_t6_eval": ModalEngineConfig(
        model="30b_t6_eval",
        duration=240,
        max_history_duration_s=120,
        max_sequences=60,  # can do 88 but that's pushing...
        model_concurrency=40,  # 4 stream 25, 2 stream 50, 1 stream 90
        compile=True,
        gpu="H200",  # A100 32/16/4, A10G
        world_size=1,
        batch_increment=32,
        upsample_concurrency=60,
        min_batch_size=64,
        max_batch_size=64,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 0}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 0}),
    ),
    "30b_test_diff": ModalEngineConfig(
        model="30b_test_diff",
        duration=240,
        max_history_duration_s=120,
        max_sequences=60,  # can do 88 but that's pushing...
        model_concurrency=40,  # 4 stream 25, 2 stream 50, 1 stream 90
        compile=True,
        gpu="H200",  # A100 32/16/4, A10G
        world_size=1,
        batch_increment=32,
        upsample_concurrency=60,
        min_batch_size=64,
        max_batch_size=64,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 3}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 3}),
    ),
    "30b_t6_infill": ModalEngineConfig(
        model="30b_t6_infill",
        duration=240,
        max_history_duration_s=120,
        max_sequences=60,  # can do 88 but that's pushing...
        model_concurrency=40,  # 4 stream 25, 2 stream 50, 1 stream 90
        compile=True,
        gpu="H200",  # A100 32/16/4, A10G
        world_size=1,
        batch_increment=32,
        upsample_concurrency=60,
        min_batch_size=64,
        max_batch_size=64,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 1}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 1}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 1}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
    ),
    "30b_classical": ModalEngineConfig(
        model="30b_classical",
        duration=240,
        max_history_duration_s=120,
        max_sequences=60,  # can do 88 but that's pushing...
        model_concurrency=40,  # 4 stream 25, 2 stream 50, 1 stream 90
        compile=True,
        gpu="H200",  # A100 32/16/4, A10G
        world_size=1,
        batch_increment=32,
        upsample_concurrency=60,
        min_batch_size=64,
        max_batch_size=64,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 0}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 0}),
    ),
    "30b_dance": ModalEngineConfig(
        model="30b_dance",
        duration=240,
        max_history_duration_s=120,
        max_sequences=60,  # can do 88 but that's pushing...
        model_concurrency=40,  # 4 stream 25, 2 stream 50, 1 stream 90
        compile=True,
        gpu="H200",  # A100 32/16/4, A10G
        world_size=1,
        batch_increment=32,
        upsample_concurrency=60,
        min_batch_size=64,
        max_batch_size=64,
        keep_warm_gpt=defaultdict(int, {"dev": 0, "prod": 0}),
        keep_warm_decoder=defaultdict(int, {"dev": 0, "prod": 0}),
    ),
    "6b_sem": ModalEngineConfig(
        model="6b_sem",
        duration=480,
        max_history_duration_s=240,
        max_sequences=72,
        max_batch_size=120,
        model_concurrency=42,
        min_batch_size=120,
        batch_increment=32,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 3}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 3}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 3}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
        decoder_concurrency=32,
        upsample_concurrency=32,
    ),
    "6b_sem_bluejay": ModalEngineConfig(
        model="6b_sem_bluejay",
        duration=480,
        max_history_duration_s=240,
        max_sequences=72,  # 18 for SFT
        max_batch_size=120,
        model_concurrency=42,
        min_batch_size=120,
        batch_increment=32,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 3}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 3}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 3}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 1}),
        decoder_concurrency=32,
        upsample_concurrency=32,
    ),
    "6b_sem_bluejay_test": ModalEngineConfig(  # DPO
        model="6b_sem_bluejay_test",
        duration=480,
        max_history_duration_s=240,
        max_sequences=72,
        max_batch_size=120,
        model_concurrency=42,
        min_batch_size=120,
        batch_increment=32,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
        decoder_concurrency=32,
        upsample_concurrency=32,
    ),
    "6b_sem_bluejay_test_2": ModalEngineConfig(
        model="6b_sem_bluejay_test_2",
        duration=480,
        max_history_duration_s=240,
        max_sequences=72,
        max_batch_size=120,
        model_concurrency=42,
        min_batch_size=120,
        batch_increment=32,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
        decoder_concurrency=32,
        upsample_concurrency=32,
    ),
    "6b_sem_cover": ModalEngineConfig(
        model="6b_sem_cover",
        duration=480,
        max_history_duration_s=240,
        max_sequences=24,
        max_batch_size=24 + 12,
        model_concurrency=24,
        min_batch_size=24,
        batch_increment=12,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 0}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 0}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 0}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
        decoder_concurrency=32,
        upsample_concurrency=32,
    ),
    "6b_sem_task": ModalEngineConfig(
        model="6b_sem_task",
        duration=480,
        max_history_duration_s=240,
        max_sequences=72,
        max_batch_size=120,
        model_concurrency=42,
        min_batch_size=120,
        batch_increment=32,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 3}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 3}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 3}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
        decoder_concurrency=32,
        upsample_concurrency=32,
    ),
    "6b_sem_eval": ModalEngineConfig(
        model="6b_sem_eval",
        duration=480,
        max_history_duration_s=240,
        max_sequences=72,
        max_batch_size=120,
        model_concurrency=42,
        min_batch_size=120,
        batch_increment=32,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 0}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 0}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 0}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
        decoder_concurrency=32,
        upsample_concurrency=32,
    ),
    "6b_sem_test": ModalEngineConfig(
        model="6b_sem_test",
        duration=480,
        max_history_duration_s=240,
        max_sequences=72,
        max_batch_size=120,
        model_concurrency=42,
        min_batch_size=120,
        batch_increment=32,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
        decoder_concurrency=32,
        upsample_concurrency=32,
    ),
    "6b_sem_test_2": ModalEngineConfig(
        model="6b_sem_test_2",
        duration=480,
        max_history_duration_s=240,
        max_sequences=72,
        max_batch_size=120,
        model_concurrency=42,
        min_batch_size=120,
        batch_increment=32,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
        decoder_concurrency=32,
        upsample_concurrency=32,
    ),
    "6b_sem_t1": ModalEngineConfig(
        model="6b_sem_t1",
        duration=480,
        max_history_duration_s=240,
        max_sequences=72,
        max_batch_size=120,
        model_concurrency=36,
        min_batch_size=120,
        batch_increment=32,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 40}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 40}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 40}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
        decoder_concurrency=40,
        upsample_concurrency=40,
    ),
    "6b_sem_tech": ModalEngineConfig(
        model="6b_sem_tech",
        duration=480,
        max_history_duration_s=240,
        max_sequences=72,
        max_batch_size=120,
        model_concurrency=36,
        min_batch_size=120,
        batch_increment=32,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 4}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
        decoder_concurrency=40,
        upsample_concurrency=40,
    ),
    "6b_sem_t2": ModalEngineConfig(  # DPO
        model="6b_sem_t2",
        duration=480,
        max_history_duration_s=240,
        max_sequences=72,
        max_batch_size=120,
        model_concurrency=36,
        min_batch_size=120,
        batch_increment=32,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 1}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 1}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 1}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
        decoder_concurrency=40,
        upsample_concurrency=40,
    ),
    "3b_sem_orig": ModalEngineConfig(
        model="3b_sem_orig",
        duration=480,
        max_history_duration_s=240,
        max_sequences=72,
        max_batch_size=120,
        model_concurrency=42,
        min_batch_size=120,
        batch_increment=32,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 0}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 0}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 0}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
        decoder_concurrency=40,
        upsample_concurrency=40,
    ),
    "3b_sem_test": ModalEngineConfig(
        model="3b_sem_test",
        duration=480,
        max_history_duration_s=240,
        max_sequences=72,
        max_batch_size=120,
        model_concurrency=42,
        min_batch_size=120,
        batch_increment=32,
        compile=True,
        gpu="H100",
        world_size=1,
        keep_warm_gpt=defaultdict(int, {"dev": 1, "prod": 0}),
        keep_warm_decoder=defaultdict(int, {"dev": 1, "prod": 0}),
        keep_warm_upsample=defaultdict(int, {"dev": 1, "prod": 0}),
        keep_warm_buffer=defaultdict(int, {"dev": 0, "prod": 0}),
        decoder_concurrency=40,
        upsample_concurrency=40,
    ),
    "diff_v1": ModalEngineConfig(
        model="diff_v1",
        duration=240,
    ),
    "diff_v1_test": ModalEngineConfig(
        model="diff_v1_test",
        duration=240,
    ),
    "diff_v1_test_2": ModalEngineConfig(
        model="diff_v1_test_2",
        duration=240,
    ),
    "diff_v1_test_3": ModalEngineConfig(
        model="diff_v1_test_3",
        duration=240,
    ),
    "diff_v1_test_4": ModalEngineConfig(
        model="diff_v1_test_4",
        duration=240,
    ),
    "diff_v2": ModalEngineConfig(
        model="diff_v2",
        duration=240,
    ),
    "diff_v2_test": ModalEngineConfig(
        model="diff_v2_test",
        duration=240,
    ),
    "diff_v2_data": ModalEngineConfig(
        model="diff_v2_data",
        duration=240,
    ),
    "stems_v0": ModalEngineConfig(
        model="stems_v0",
        duration=240,
    ),
    "stems_v1": ModalEngineConfig(
        model="stems_v1",
        duration=240,
    ),
    "stems_v1_8_output": ModalEngineConfig(
        model="stems_v1_8_output",
        duration=240,
    ),
    "stems_v1_12_output": ModalEngineConfig(
        model="stems_v1_12_output",
        duration=240,
    ),
    "diff_seeds_v0": ModalEngineConfig(
        model="diff_seeds_v0",
        duration=30,
    ),
}