Qwen3-Omni Preprocessing and Reformatting

Author: eitanporat

src/MaxText/configs/base.yml

@@ -865,7 +865,10 @@ dtype_mm: "float32"  # Data type for multimodal model's vision encoder
 remat_policy_for_vit: "minimal"  # Remat policy for multimodal model's vision encoder. Check `remat_policy` for options.
 image_size_for_vit: 896 # Default for Gemma3, and should be overwritten by model's config
 image_path: "" # Local image path used for decoding, can be multiple paths separated by comma, exp "/path/image1.jpg,/path/image2.jpg"
+audio_path: "" # Local audio path used for decoding, can be multiple paths separated by comma, exp "/path/audio1.wav,/path/audio2.wav"
+video_path: "" # Local video path used for decoding, can be multiple paths separated by comma, exp "/path/video1.mp4,/path/video2.mp4"
 image_placeholder: "<|image|>"
+audio_placeholder: "<|audio|>"
 posemb_type_for_vit: "learn"
 # max_num_images_per_example only applies for training when your image column is a list of images.
 # -1 means no limit, and will pad to the max possible number of images determined by sequence length.

src/MaxText/configs/models/qwen3-omni-30b-a3b.yml

@@ -38,3 +38,47 @@ rope_max_timescale: 10_000_000
 
 # General Model Settings
 enable_dropout: False
+
+# Vision Encoder Configuration
+# Based on https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py
+image_size_for_vit: 768
+hidden_size_for_vit: 1152
+intermediate_size_for_vit: 4304
+num_attention_heads_for_vit: 16
+num_hidden_layers_for_vit: 27
+num_channels_for_vit: 3
+patch_size_for_vit: 16
+temporal_patch_size_for_vit: 2
+spatial_merge_size_for_vit: 2
+out_hidden_size_for_vit: 2048
+num_position_embeddings_for_vit: 2304
+deepstack_visual_indexes_for_vit: [8, 16, 24]
+
+use_multimodal: true
+use_audio: true
+# Audio Encoder Configuration (need to set use_audio=true to enable)
+# Based on https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen3_omni_moe/configuration_qwen3_omni_moe.py
+d_model_for_audio: 1280
+encoder_layers_for_audio: 32
+encoder_attention_heads_for_audio: 20
+encoder_ffn_dim_for_audio: 5120
+max_source_positions_for_audio: 1500
+num_mel_bins_for_audio: 128
+downsample_hidden_size_for_audio: 480
+output_dim_for_audio: 2048
+attention_dropout_for_audio: 0.0
+n_window_for_audio: 50
+n_window_infer_for_audio: 400
+conv_chunksize_for_audio: 500
+num_conv_layers_for_audio: 3
+max_timescale_for_audio: 10000.0
+max_sample_len_for_audio: 10000
+
+freeze_audio_encoder_params: false
+freeze_vision_encoder_params: false
+# MRoPE Settings (Multi-dimensional RoPE for multimodal)
+use_mrope: true
+mrope_section: [24, 20, 20]
+
+
+image_placeholder: "<|image|>"