add unit tests and functional tests

Signed-off-by: chanchzhang <chanchzhang@tencent.com>

Author: zhang-cheng09

src/megatron/bridge/models/qwen_omni/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from megatron.bridge.models.qwen_omni.modelling_qwen3_omni import Qwen3OmniMoeModel
+from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.model import Qwen3OmniMoeModel
 from megatron.bridge.models.qwen_omni.qwen3_omni_bridge import Qwen3OmniMoeBridge
 from megatron.bridge.models.qwen_omni.qwen3_omni_provider import Qwen3OmniMoeModelProvider
 

src/megatron/bridge/models/qwen_omni/context_parallel_utils.py

@@ -26,8 +26,8 @@
     Qwen3OmniMoeCode2WavConfig,
 )
 
-from megatron.bridge.models.qwen_omni.thinker_model import Qwen3OmniMoeThinkerModel
-from megatron.bridge.models.qwen_omni.transformer_config import Qwen3OmniTransformerConfig
+from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.thinker_model import Qwen3OmniMoeThinkerModel
+from megatron.bridge.models.qwen_omni.modeling_qwen3_omni.transformer_config import Qwen3OmniTransformerConfig
 
 
 class Qwen3OmniMoeModel(MegatronModule):
@@ -73,23 +73,26 @@ def set_input_tensor(self, input_tensor) -> None:
 
     def freeze(
         self,
-        freeze_language_model: bool,
-        freeze_vision_model: bool,
-        freeze_vision_projection: bool,
+        freeze_language_model: bool=False,
+        freeze_vision_model: bool=False,
+        freeze_vision_projection: bool=False,
+        freeze_audio_model: bool=False,
     ):
         """Freeze model modules.
 
         Make specific modules non-trainable by setting requires_grad to False.
 
         Args:
             freeze_language_model (bool): Freeze the language model module.
-            freeze_vision_model (bool): Freeze the vision model module (patch_embed, blocks, pos_embed).
+            freeze_vision_model (bool): Freeze the vision model module.
             freeze_vision_projection (bool): Freeze the vision projection modules (merger and deepstack_merger_list).
+            freeze_audio_model (bool): Freeze the audio model module.
         """
         return self.thinker.freeze(
             freeze_language_model,
             freeze_vision_model,
-            freeze_vision_projection
+            freeze_vision_projection,
+            freeze_audio_model,
         )
 
     def forward(
@@ -113,6 +116,7 @@ def forward(
         feature_attention_mask=None,
         audio_feature_lengths=None,
         cp_img_num: list[int] = None,
+        images_padded: list[bool] = None,
         use_audio_in_video=None,
         video_second_per_grid=None,
         **kwargs,
@@ -136,6 +140,7 @@ def forward(
             feature_attention_mask=feature_attention_mask,
             audio_feature_lengths=audio_feature_lengths,
             cp_img_num=cp_img_num,
+            images_padded=images_padded,
             use_audio_in_video=use_audio_in_video,
             video_second_per_grid=video_second_per_grid,
             **kwargs,