Skip to content

vllm_omni.transformers_utils.configs.auk

Config for the AuK instruction-driven audio generation model.

The upstream checkpoint ships config.yaml plus loose safetensors, so an assembled directory is produced by tools/prepare_auk_checkpoint.py: the Qwen2.5-Omni-3B config.json with this model_type, AuK's architecture, and the five sections below added at the top level.

Subclassing Qwen2_5OmniConfig is load-bearing, not stylistic. The encoder stage reuses the in-tree Qwen2.5-Omni thinker processor, whose get_hf_config resolves self.ctx.get_hf_config(Qwen2_5OmniConfig) and asserts the type, so a bare PretrainedConfig subclass would fail before the first request. Inheriting also leaves the parent's nested thinker_config / talker_config / token2wav_config handling and its get_text_config in place, which is what the encoder stage reads.

AUK_FLASH_T_GRID module-attribute 

AUK_FLASH_T_GRID: tuple[float, ...] = (
    0.0,
    0.07612049579620361,
    0.2928932309150696,
    0.6173166036605835,
    1.0,
)

DEFAULT_DIT module-attribute 

DEFAULT_DIT: dict[str, Any] = {
    "dim": 1536,
    "heads": 24,
    "dim_head": 64,
    "ff_mult": 2,
    "text_hidden_dim": 2048,
    "num_layers": 10,
    "num_single_layers": 20,
    "attn_mask_enabled": True,
}

DEFAULT_SAMPLING module-attribute 

DEFAULT_SAMPLING: dict[str, Any] = {
    "nfe": 32,
    "cfg": 2.0,
    "sway": -1.0,
}

DEFAULT_VAE module-attribute 

DEFAULT_VAE: dict[str, Any] = {
    "latent_dim": 64,
    "downsample_rate": 480,
    "target_sample_rate": 24000,
    "model_init_kwargs": {
        "upsample_rates": [5, 4, 3, 2, 2, 2],
        "upsample_kernel_sizes": [10, 8, 6, 4, 4, 4],
        "upsample_initial_channel": 1536,
        "resblock_kernel_sizes": [3, 7, 11],
        "resblock_dilation_sizes": [
            [1, 3, 5],
            [1, 3, 5],
            [1, 3, 5],
        ],
        "downsample_rates": [2, 2, 2, 3, 4, 5],
        "downsample_channels": [
            12,
            24,
            48,
            96,
            192,
            384,
            768,
        ],
        "snake_logscale": True,
        "latent_dim": 64,
        "use_vae": True,
        "causal": True,
        "flow_hidden_channels": 256,
        "act_causal": True,
    },
}

VARIANTS module-attribute 

VARIANTS = ('base', 'flash')

AuKConfig

Bases: Qwen2_5OmniConfig

Assembled AuK checkpoint config: a Qwen2.5-Omni config plus DiT and VAE.

defaults instance-attribute 

defaults = {**DEFAULT_SAMPLING, **defaults or {}}

dit instance-attribute 

dit = {**DEFAULT_DIT, **dit or {}}

downsample_rate property 

downsample_rate: int

flash_t_grid instance-attribute 

flash_t_grid = (
    list(flash_t_grid)
    if flash_t_grid is not None
    else list(AUK_FLASH_T_GRID)
)

is_flash property 

is_flash: bool

Whether this checkpoint is the distilled 4-step student.

latent_dim property 

latent_dim: int

model_type class-attribute instance-attribute 

model_type = 'auk'

target_sample_rate property 

target_sample_rate: int

vae instance-attribute 

vae = {**DEFAULT_VAE, **vae or {}}

variant instance-attribute 

variant = variant