vllm-omni/vllm_omni/transformers_utils/configs/ming_flash_omni.py

# SPDX-License-Identifier: Apache-2.0
# Copyright 2025 The vLLM-Omni team.
# Copyright 2024 ANT Group and the HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Configuration for Ming-flash-omni-2.0 model"""

import os
from typing import Any, ClassVar

from transformers import AutoConfig, AutoTokenizer, PretrainedConfig, PreTrainedTokenizerFast
from transformers.utils import logging

logger = logging.get_logger(__name__)


class BailingMoeV2Config(PretrainedConfig):
    model_type = "bailing_moe_v2"

    def __init__(
        self,
        vocab_size=30592,
        hidden_size=1024,
        intermediate_size=None,
        num_hidden_layers=24,
        num_attention_heads=16,
        num_key_value_heads=0,
        hidden_act="silu",
        use_qkv_bias=False,
        use_qk_norm=False,
        use_bias=True,
        rms_norm_eps=1e-05,
        norm_head=False,
        tie_word_embeddings=False,
        embedding_dropout=0.0,
        attention_dropout=0.0,
        output_dropout=0.0,
        initializer_range=0.02,
        max_position_embeddings=16384,
        rope_theta=10000.0,
        use_cache=True,
        use_sliding_window=False,
        sliding_window=81920,
        max_window_layers=28,
        rope_scaling=None,
        mrope_section=None,
        pad_token_id=126081,
        num_experts=16,
        num_shared_experts=1,
        num_experts_per_tok=2,
        n_group=8,
        topk_group=4,
        routed_scaling_factor=2.5,
        moe_intermediate_size=None,
        first_k_dense_replace=0,
        head_dim=None,
        output_router_logits=False,
        partial_rotary_factor=0.5,
        router_type="topN",
        _attn_implementation="flash_attention_2",
        use_interleaved_frame_timestamp=True,
        # Multimodal token IDs
        image_patch_token=157157,
        video_patch_token=157175,
        audio_patch_token=157168,
        image_start_token=157158,
        video_start_token=157160,
        audio_start_token=157169,
        image_end_token=157159,
        video_end_token=157161,
        audio_end_token=157170,
        # Position encoding parameters
        spatial_merge_size=2,
        tokens_per_second=2,
        **kwargs,
    ):
        self.num_hidden_layers = num_hidden_layers
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.hidden_act = hidden_act
        self.use_qkv_bias = use_qkv_bias
        self.use_bias = use_bias
        self.norm_head = norm_head
        self.rms_norm_eps = rms_norm_eps
        self.embedding_dropout = embedding_dropout
        self.attention_dropout = attention_dropout
        self.output_dropout = output_dropout
        self.initializer_range = initializer_range
        self.max_position_embeddings = max_position_embeddings
        self.rope_theta = rope_theta
        self.use_cache = use_cache
        self.use_sliding_window = use_sliding_window
        self.sliding_window = sliding_window
        self.max_window_layers = max_window_layers
        self.head_dim = head_dim or self.hidden_size // self.num_attention_heads
        self.use_qk_norm = use_qk_norm  # arg unused; QK norm is always applied

        # By default, match the value of `mrope_section`
        # to `apply_3d_rotary_pos_emb` in Ming's repo:
        # https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/modeling_bailing_moe_v2.py
        if mrope_section is None:
            mrope_section = (rope_scaling or {}).get("mrope_section", [8, 12, 12])
        # Ensure mrope_section is stored inside rope_scaling
        if rope_scaling is not None and isinstance(rope_scaling, dict):
            rope_scaling = dict(rope_scaling)
            rope_scaling.setdefault("mrope_section", mrope_section)
        self.rope_scaling = rope_scaling

        # NOTE: Expose rope_parameters["mrope_section"]
        # This refers to the pattern used for GLM-Image in vllm_omni/patch.py
        rope_type = (rope_scaling or {}).get("type", (rope_scaling or {}).get("rope_type", ""))
        if rope_type in ("video_rope", "3D", "mrope"):
            self.rope_parameters = {"mrope_section": mrope_section}
        else:
            self.rope_parameters = None

        # MoE configs
        self.num_experts = num_experts
        self.num_shared_experts = num_shared_experts
        self.num_experts_per_tok = num_experts_per_tok
        self.n_group = n_group
        self.topk_group = topk_group
        self.moe_intermediate_size = moe_intermediate_size
        self.first_k_dense_replace = first_k_dense_replace
        self.output_router_logits = output_router_logits
        self.routed_scaling_factor = routed_scaling_factor
        self.partial_rotary_factor = partial_rotary_factor
        self.router_type = router_type
        self.use_interleaved_frame_timestamp = use_interleaved_frame_timestamp
        self._attn_implementation = _attn_implementation

        # Multimodal token IDs and position encoding
        self.image_patch_token = image_patch_token
        self.video_patch_token = video_patch_token
        self.audio_patch_token = audio_patch_token
        self.image_start_token = image_start_token
        self.video_start_token = video_start_token
        self.audio_start_token = audio_start_token
        self.image_end_token = image_end_token
        self.video_end_token = video_end_token
        self.audio_end_token = audio_end_token
        self.spatial_merge_size = spatial_merge_size
        self.tokens_per_second = tokens_per_second

        super().__init__(pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs)


class Qwen3VLMoeVisionConfig(PretrainedConfig):
    """Configuration class for Qwen3 MoE Vision Transformer"""

    model_type = "qwen3_moe_vit"

    def __init__(
        self,
        depth=27,
        hidden_size=1152,
        hidden_act="gelu_pytorch_tanh",
        intermediate_size=4304,
        num_heads=16,
        in_channels=3,
        patch_size=16,
        spatial_merge_size=2,
        temporal_patch_size=2,
        out_hidden_size=3584,
        num_position_embeddings=2304,
        deepstack_visual_indexes=[8, 16, 24],
        initializer_range=0.02,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.depth = depth
        self.hidden_size = hidden_size
        self.hidden_act = hidden_act
        self.intermediate_size = intermediate_size
        self.num_heads = num_heads
        self.in_channels = in_channels
        self.patch_size = patch_size
        self.spatial_merge_size = spatial_merge_size
        self.temporal_patch_size = temporal_patch_size
        self.out_hidden_size = out_hidden_size
        self.num_position_embeddings = num_position_embeddings
        self.initializer_range = initializer_range
        self.deepstack_visual_indexes = deepstack_visual_indexes

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike, **kwargs) -> "PretrainedConfig":
        cls._set_token_in_kwargs(kwargs)

        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)

        if "vision_config" in config_dict:
            config_dict = config_dict["vision_config"]

        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
            logger.warning(
                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
            )

        return cls.from_dict(config_dict, **kwargs)


class WhisperEncoderConfig(PretrainedConfig):
    """Configuration class for Whisper audio encoder"""

    model_type = "whisper_encoder"

    def __init__(
        self,
        whisper_encoder_config: dict[str, Any] | None = None,
        ds_kernel_size=3,
        ds_stride=2,
        norm_query_embeds=True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.whisper_encoder_config = whisper_encoder_config or {}
        self.ds_kernel_size = ds_kernel_size
        self.ds_stride = ds_stride
        self.norm_query_embeds = norm_query_embeds


class BailingMM2Config(PretrainedConfig):
    model_type = "bailingmm_moe_v2_lite"
    is_composition = True
    sub_configs: ClassVar = {"llm_config": AutoConfig}

    def __init__(
        self,
        mlp_depth=1,
        llm_config: BailingMoeV2Config | None = None,
        vision_config: Qwen3VLMoeVisionConfig | None = None,
        audio_config: WhisperEncoderConfig | None = None,
        **kwargs,
    ):
        self.audio_config = WhisperEncoderConfig(**audio_config) if isinstance(audio_config, dict) else audio_config
        self.vision_config = (
            Qwen3VLMoeVisionConfig(**vision_config) if isinstance(vision_config, dict) else vision_config
        )
        self.llm_config = BailingMoeV2Config(**llm_config) if isinstance(llm_config, dict) else llm_config
        self.mlp_depth = mlp_depth
        super().__init__(**kwargs)

    def get_text_config(self, decoder: bool = False) -> PretrainedConfig:  # noqa: ARG002
        return self.llm_config


class MingImageGenConfig(PretrainedConfig):
    """Configuration for Ming-flash-omni-2.0 image generation stage.

    Mirrors the layout of the HF checkpoint at
    https://huggingface.co/inclusionAI/Ming-flash-omni-2.0 where image-gen
    components live in sibling subfolders (``connector/``, ``transformer/``,
    ``vae/``, ``scheduler/``, ``mlp/``).
    """

    model_type = "ming_flash_omni_imagegen"

    def __init__(
        self,
        connector_subfolder: str = "connector",
        transformer_subfolder: str = "transformer",
        vae_subfolder: str = "vae",
        scheduler_subfolder: str = "scheduler",
        mlp_subfolder: str = "mlp",
        diffusion_c_input_dim: int = 2560,
        thinker_hidden_size: int = 4096,
        img_gen_scales: list[int] | None = None,
        text_encoder_norm: bool = True,
        # Default sampling parameters — match Ming's defaults in
        # ``modeling_bailingmm2.py::generate`` (``image_gen_steps=30``,
        # ``image_gen_cfg=2.0``).
        num_inference_steps: int = 30,
        guidance_scale: float = 2.0,
        default_height: int = 1024,
        default_width: int = 1024,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.connector_subfolder = connector_subfolder
        self.transformer_subfolder = transformer_subfolder
        self.vae_subfolder = vae_subfolder
        self.scheduler_subfolder = scheduler_subfolder
        self.mlp_subfolder = mlp_subfolder
        self.diffusion_c_input_dim = diffusion_c_input_dim
        self.thinker_hidden_size = thinker_hidden_size
        self.img_gen_scales = img_gen_scales if img_gen_scales is not None else [16]
        self.text_encoder_norm = text_encoder_norm
        self.num_inference_steps = num_inference_steps
        self.guidance_scale = guidance_scale
        self.default_height = default_height
        self.default_width = default_width

    @property
    def num_query_tokens(self) -> int:
        """Total number of learnable query tokens appended to the thinker input.

        For ``img_gen_scales=[16]`` this yields 256 tokens (a single 16x16 grid).
        """
        return sum(s * s for s in self.img_gen_scales)


class MingFlashOmniTalkerConfig(PretrainedConfig):
    """Configuration class for Ming-flash-omni-2.0 talker (TTS) stage.

    The talker uses a Qwen2 LLM backbone with CFM (Conditional Flow Matching)
    via a DiT diffusion transformer, plus an Aggregator that maps generated
    audio latents back to the LLM embedding space for autoregressive generation.
    """

    model_type = "ming_flash_omni_talker"

    def __init__(
        self,
        llm_config: dict[str, Any] | None = None,
        flowmodel: dict[str, Any] | None = None,
        aggregator: dict[str, Any] | None = None,
        steps: int = 10,
        patch_size: int = 4,
        history_patch_size: int = 32,
        latent_dim: int = 64,
        cfg_strength: float = 2.0,
        audio_vae_path: str | None = None,
        campplus_model: str | None = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.llm_config = llm_config
        self.flowmodel = flowmodel or {}
        self.aggregator = aggregator or {}
        self.steps = steps
        self.patch_size = patch_size
        self.history_patch_size = history_patch_size
        self.latent_dim = latent_dim
        self.cfg_strength = cfg_strength
        self.audio_vae_path = audio_vae_path
        self.campplus_model = campplus_model

    def get_text_config(self, decoder: bool = False) -> PretrainedConfig:  # noqa: ARG002
        if isinstance(self.llm_config, dict):
            return PretrainedConfig.from_dict(self.llm_config)
        return self.llm_config


class MingFlashOmniConfig(PretrainedConfig):
    """Configuration class for unified Ming-flash-omni-2.0 model"""

    model_type = "ming_flash_omni"
    is_composition = True
    sub_configs: ClassVar = {
        "thinker_config": BailingMM2Config,
        "image_gen_config": MingImageGenConfig,
        "talker_config": MingFlashOmniTalkerConfig,
    }

    def __init__(
        self,
        thinker_config: BailingMM2Config | dict[str, Any] | None = None,
        image_gen_config: MingImageGenConfig | dict[str, Any] | None = None,
        talker_config: MingFlashOmniTalkerConfig | dict[str, Any] | None = None,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if isinstance(thinker_config, dict):
            self.thinker_config = BailingMM2Config(**thinker_config)
        else:
            self.thinker_config = thinker_config or BailingMM2Config()

        # Image generation config
        if isinstance(image_gen_config, dict):
            self.image_gen_config = MingImageGenConfig(**image_gen_config)
        elif image_gen_config is None:
            self.image_gen_config = MingImageGenConfig()
        else:
            self.image_gen_config = image_gen_config

        # Talker config
        if isinstance(talker_config, dict):
            self.talker_config = MingFlashOmniTalkerConfig(**talker_config)
        else:
            self.talker_config = talker_config

    def get_text_config(self, decoder: bool = False) -> PretrainedConfig:  # noqa: ARG002
        return self.thinker_config.get_text_config()


# Register model_type -> config class for AutoConfig
AutoConfig.register(BailingMoeV2Config.model_type, BailingMoeV2Config)
AutoConfig.register(BailingMM2Config.model_type, BailingMM2Config)
AutoConfig.register(MingImageGenConfig.model_type, MingImageGenConfig)
AutoConfig.register(MingFlashOmniTalkerConfig.model_type, MingFlashOmniTalkerConfig)
AutoConfig.register(MingFlashOmniConfig.model_type, MingFlashOmniConfig)

# Register tokenizer mapping for composition configs so that
# AutoTokenizer.from_pretrained can resolve the tokenizer class
AutoTokenizer.register(BailingMM2Config, fast_tokenizer_class=PreTrainedTokenizerFast)
AutoTokenizer.register(MingFlashOmniTalkerConfig, fast_tokenizer_class=PreTrainedTokenizerFast)
AutoTokenizer.register(MingFlashOmniConfig, fast_tokenizer_class=PreTrainedTokenizerFast)