Restructure models into lm/vlm/vam sub-packages by modality

Split the flat model modules into three modality-scoped sub-packages,
each with a config.py (configuration) and model.py (modeling code):

- models/lm/ : pure text -> MiniMindConfig + MiniMindForCausalLM
- models/vlm/ : text + vision -> VLMConfig + MiniMindVLM
- models/vam/ : text + audio -> OmniConfig + MiniMindOmni (+ TalkerModule)

models/__init__.py re-exports all symbols so existing
`from omni.models import ...` references keep working. Updated
trainers/examples/utils that imported the old flat submodule paths.
README reflects the new layout.

Co-Authored-By: opencode <noreply@opencode.ai>

README.md

@@ -14,8 +14,10 @@ Omni 是一个以 **多模态 (omni)** 为目标的 LLM 训练 / 推理框架，
   `norm.py`（`RMSNorm`）、`rope.py`（`precompute_freqs_cis` / `apply_rotary_pos_emb` / `repeat_kv`）、
   `attention.py`（`Attention`）、`mlp.py`（`FeedForward` / `MOEFeedForward`）、
   `block.py`（`MiniMindBlock`）、`model.py`（`MiniMindModel` Transformer 主干）。
-- `models/`：把 `core` 组件**拼装**成成品模型——
-  `MiniMindForCausalLM`（文本）、`MiniMindVLM`（视觉）、`MiniMindOmni`（语音/全模态）。
+- `models/`：把 `core` 组件**拼装**成成品模型，按模态能力分为三个子包（每个含 `config.py` 配置 + `model.py` 建模）：
+  - `models/lm/`：纯文本——`MiniMindConfig` + `MiniMindForCausalLM`
+  - `models/vlm/`：文本 + 视觉——`VLMConfig` + `MiniMindVLM`
+  - `models/vam/`：文本 + 语音/全模态——`OmniConfig` + `MiniMindOmni`（含 `TalkerModule`）
 - `encoders/`：外部模态编码器，按模态分目录——`vision/`（SigLIP）、`audio/`（SenseVoice）。
 - `projectors/`：把 encoder 输出**桥接**到 LLM 隐藏维度的拼接层（`MMVisionProjector`、`MMAudioProjector`）。
 - `serve/`：实时语音会话工程层（`SileroVAD`、`RealtimeSession`）。
@@ -31,10 +33,16 @@ src/omni/
 │   ├── mlp.py          #   FeedForward / MOEFeedForward
 │   ├── block.py        #   MiniMindBlock
 │   └── model.py        #   MiniMindModel（Transformer 主干）
-├── models/             # 模型拼装
-│   ├── minimind.py     # MiniMindConfig + MiniMindForCausalLM
-│   ├── vlm.py          # VLMConfig + MiniMindVLM（视觉多模态）
-│   ├── omni.py         # OmniConfig + MiniMindOmni + TalkerModule（语音/全模态）
+├── models/             # 模型拼装（按模态能力分子包）
+│   ├── lm/             #   纯文本
+│   │   ├── config.py   #     MiniMindConfig
+│   │   └── model.py    #     MiniMindForCausalLM
+│   ├── vlm/            #   文本 + 视觉
+│   │   ├── config.py   #     VLMConfig
+│   │   └── model.py    #     MiniMindVLM
+│   ├── vam/            #   文本 + 语音/全模态
+│   │   ├── config.py   #     OmniConfig
+│   │   └── model.py    #     MiniMindOmni + TalkerModule
 │   └── lora.py         # LoRA 注入 / 保存 / 合并
 ├── encoders/           # 多模态编码器（按模态分目录）
 │   ├── vision/        #   SiglipVisionEncoder

examples/convert_model.py

@@ -5,7 +5,7 @@ import torch
 import transformers
 import warnings
 from transformers import AutoTokenizer, AutoModelForCausalLM, Qwen3Config, Qwen3ForCausalLM, Qwen3MoeConfig, Qwen3MoeForCausalLM
-from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.models import MiniMindConfig, MiniMindForCausalLM
 from omni.models.lora import apply_lora, merge_lora
 
 warnings.filterwarnings('ignore', category=UserWarning)

examples/eval_llm.py

@@ -4,7 +4,7 @@ import random
 import warnings
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
-from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.models import MiniMindConfig, MiniMindForCausalLM
 from omni.models.lora import *  # noqa: F401,F403
 from omni.utils.training import setup_seed, get_model_params
 warnings.filterwarnings('ignore')

examples/eval_toolcall.py

@@ -9,7 +9,7 @@ import torch
 from datetime import datetime
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 from openai import OpenAI
-from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.models import MiniMindConfig, MiniMindForCausalLM
 from omni.utils.training import setup_seed, get_model_params
 warnings.filterwarnings('ignore')
 

examples/serve_openai_api.py

@@ -14,7 +14,7 @@ from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Field
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
-from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.models import MiniMindConfig, MiniMindForCausalLM
 from omni.models.lora import apply_lora, load_lora
 
 warnings.filterwarnings('ignore')

src/omni/models/__init__.py

@@ -1,3 +1,16 @@
+from omni.models.lm.config import MiniMindConfig
+from omni.models.lm.model import MiniMindForCausalLM
+from omni.models.vlm.config import VLMConfig
+from omni.models.vlm.model import MiniMindVLM
+from omni.models.vam.config import OmniConfig
+from omni.models.vam.model import MiniMindOmni, TalkerModule
+from omni.models.lora import (
+    LoRA,
+    apply_lora,
+    load_lora,
+    save_lora,
+    merge_lora,
+)
 from omni.core import (
     RMSNorm,
     Attention,
@@ -8,46 +21,26 @@ from omni.core import (
     precompute_freqs_cis,
     apply_rotary_pos_emb,
 )
-from omni.models.minimind import (
-    MiniMindConfig,
-    MiniMindForCausalLM,
-)
-from omni.models.vlm import (
-    VLMConfig,
-    MiniMindVLM,
-)
-from omni.models.omni import (
-    OmniConfig,
-    MiniMindOmni,
-    TalkerModule,
-)
-from omni.models.lora import (
-    LoRA,
-    apply_lora,
-    load_lora,
-    save_lora,
-    merge_lora,
-)
 
 __all__ = [
     "MiniMindConfig",
-    "MiniMindModel",
     "MiniMindForCausalLM",
     "VLMConfig",
     "MiniMindVLM",
     "OmniConfig",
     "MiniMindOmni",
     "TalkerModule",
+    "LoRA",
+    "apply_lora",
+    "load_lora",
+    "save_lora",
+    "merge_lora",
     "RMSNorm",
     "Attention",
     "FeedForward",
     "MOEFeedForward",
     "MiniMindBlock",
+    "MiniMindModel",
     "precompute_freqs_cis",
     "apply_rotary_pos_emb",
-    "LoRA",
-    "apply_lora",
-    "load_lora",
-    "save_lora",
-    "merge_lora",
 ]

src/omni/models/lm/__init__.py

@@ -0,0 +1,4 @@
+from omni.models.lm.config import MiniMindConfig
+from omni.models.lm.model import MiniMindForCausalLM
+
+__all__ = ["MiniMindConfig", "MiniMindForCausalLM"]

src/omni/models/lm/config.py

@@ -0,0 +1,41 @@
+import math
+from transformers import PretrainedConfig
+
+
+class MiniMindConfig(PretrainedConfig):
+    model_type = "minimind"
+
+    def __init__(self, hidden_size=768, num_hidden_layers=8, use_moe=False, **kwargs):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.use_moe = use_moe
+        self.dropout = kwargs.get("dropout", 0.0)
+        self.vocab_size = kwargs.get("vocab_size", 6400)
+        self.bos_token_id = kwargs.get("bos_token_id", 1)
+        self.eos_token_id = kwargs.get("eos_token_id", 2)
+        self.flash_attn = kwargs.get("flash_attn", True)
+        self.num_attention_heads = kwargs.get("num_attention_heads", 8)
+        self.num_key_value_heads = kwargs.get("num_key_value_heads", 4)
+        self.head_dim = kwargs.get("head_dim", self.hidden_size // self.num_attention_heads)
+        self.hidden_act = kwargs.get("hidden_act", 'silu')
+        self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64)
+        self.max_position_embeddings = kwargs.get("max_position_embeddings", 32768)
+        self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-6)
+        self.rope_theta = kwargs.get("rope_theta", 1e6)
+        self.tie_word_embeddings = kwargs.get("tie_word_embeddings", True)
+        self.inference_rope_scaling = kwargs.get("inference_rope_scaling", False)
+        self.rope_scaling = {
+            "beta_fast": 32,
+            "beta_slow": 1,
+            "factor": 16,
+            "original_max_position_embeddings": 2048,
+            "attention_factor": 1.0,
+            "type": "yarn"
+        } if self.inference_rope_scaling else None
+        # MoE specific configs (ignored if use_moe = False)
+        self.num_experts = kwargs.get("num_experts", 4)
+        self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 1)
+        self.moe_intermediate_size = kwargs.get("moe_intermediate_size", self.intermediate_size)
+        self.norm_topk_prob = kwargs.get("norm_topk_prob", True)
+        self.router_aux_loss_coef = kwargs.get("router_aux_loss_coef", 5e-4)

src/omni/models/{minimind.py → lm/model.py}

@@ -1,50 +1,11 @@
-import math
 import torch
 import torch.nn.functional as F
 from torch import nn
-from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
+from transformers import PreTrainedModel, GenerationMixin
 from transformers.modeling_outputs import MoeCausalLMOutputWithPast
 
 from omni.core import MiniMindModel
-
-
-class MiniMindConfig(PretrainedConfig):
-    model_type = "minimind"
-
-    def __init__(self, hidden_size=768, num_hidden_layers=8, use_moe=False, **kwargs):
-        super().__init__(**kwargs)
-        self.hidden_size = hidden_size
-        self.num_hidden_layers = num_hidden_layers
-        self.use_moe = use_moe
-        self.dropout = kwargs.get("dropout", 0.0)
-        self.vocab_size = kwargs.get("vocab_size", 6400)
-        self.bos_token_id = kwargs.get("bos_token_id", 1)
-        self.eos_token_id = kwargs.get("eos_token_id", 2)
-        self.flash_attn = kwargs.get("flash_attn", True)
-        self.num_attention_heads = kwargs.get("num_attention_heads", 8)
-        self.num_key_value_heads = kwargs.get("num_key_value_heads", 4)
-        self.head_dim = kwargs.get("head_dim", self.hidden_size // self.num_attention_heads)
-        self.hidden_act = kwargs.get("hidden_act", 'silu')
-        self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64)
-        self.max_position_embeddings = kwargs.get("max_position_embeddings", 32768)
-        self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-6)
-        self.rope_theta = kwargs.get("rope_theta", 1e6)
-        self.tie_word_embeddings = kwargs.get("tie_word_embeddings", True)
-        self.inference_rope_scaling = kwargs.get("inference_rope_scaling", False)
-        self.rope_scaling = {
-            "beta_fast": 32,
-            "beta_slow": 1,
-            "factor": 16,
-            "original_max_position_embeddings": 2048,
-            "attention_factor": 1.0,
-            "type": "yarn"
-        } if self.inference_rope_scaling else None
-        # MoE specific configs (ignored if use_moe = False)
-        self.num_experts = kwargs.get("num_experts", 4)
-        self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 1)
-        self.moe_intermediate_size = kwargs.get("moe_intermediate_size", self.intermediate_size)
-        self.norm_topk_prob = kwargs.get("norm_topk_prob", True)
-        self.router_aux_loss_coef = kwargs.get("router_aux_loss_coef", 5e-4)
+from omni.models.lm.config import MiniMindConfig
 
 
 class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):

src/omni/models/vam/__init__.py

@@ -0,0 +1,4 @@
+from omni.models.vam.config import OmniConfig
+from omni.models.vam.model import MiniMindOmni, TalkerModule
+
+__all__ = ["OmniConfig", "MiniMindOmni", "TalkerModule"]

src/omni/models/vam/config.py

@@ -0,0 +1,24 @@
+from omni.models.lm.config import MiniMindConfig
+
+
+class OmniConfig(MiniMindConfig):
+    model_type = "minimind-o"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.num_talker_hidden_layers = kwargs.get("num_talker_hidden_layers", 4)
+        self.talker_hidden_size = kwargs.get("talker_hidden_size", 768)
+        self.audio_ids = kwargs.get("audio_ids", [16])
+        self.audio_special_token = kwargs.get("audio_special_token", "<|audio_pad|>")
+        self.audio_hidden_size = kwargs.get("audio_hidden_size", 512)
+        self.audio_vocab_size = kwargs.get("audio_vocab_size", 2112)
+        self.audio_pad_token = kwargs.get("audio_pad_token", 2049)
+        self.audio_stop_token = kwargs.get("audio_stop_token", 2050)
+        self.audio_spk_token = kwargs.get("audio_spk_token", 2051)
+        self.spk_emb_size = kwargs.get("spk_emb_size", 192)
+        self.think_end_ids = kwargs.get("think_end_ids", [26, 234, 234])
+        self.image_ids = kwargs.get("image_ids", [12])
+        self.image_special_token = kwargs.get("image_special_token", "<|image_pad|>")
+        self.image_hidden_size = kwargs.get("image_hidden_size", 768)
+        self.image_token_len = kwargs.get("image_token_len", 64)
+        self.bridge_layer = kwargs.get("bridge_layer", self.num_hidden_layers // 2 - 1)

src/omni/models/{omni.py → vam/model.py}

@@ -10,36 +10,15 @@ from torch.nn import functional as F
 from transformers.modeling_outputs import MoeCausalLMOutputWithPast
 from transformers import SiglipVisionModel, SiglipImageProcessor, logging as hf_logging
 
-from omni.core import RMSNorm, precompute_freqs_cis, apply_rotary_pos_emb, MiniMindBlock, MOEFeedForward
-from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
-from omni.encoders.audio import SenseVoiceAudioEncoder
+from omni.core import RMSNorm, precompute_freqs_cis, MiniMindBlock, MOEFeedForward
+from omni.models.lm.config import MiniMindConfig
+from omni.models.lm.model import MiniMindForCausalLM
+from omni.models.vam.config import OmniConfig
+from omni.encoders.audio import SenseVoiceAudioEncoder, SenseVoiceAudioProcessor
 from omni.encoders.vision import SiglipVisionEncoder
 from omni.projectors import MMVisionProjector, MMAudioProjector
 
 
-class OmniConfig(MiniMindConfig):
-    model_type = "minimind-o"
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.num_talker_hidden_layers = kwargs.get("num_talker_hidden_layers", 4)
-        self.talker_hidden_size = kwargs.get("talker_hidden_size", 768)
-        self.audio_ids = kwargs.get("audio_ids", [16])
-        self.audio_special_token = kwargs.get("audio_special_token", "<|audio_pad|>")
-        self.audio_hidden_size = kwargs.get("audio_hidden_size", 512)
-        self.audio_vocab_size = kwargs.get("audio_vocab_size", 2112)
-        self.audio_pad_token = kwargs.get("audio_pad_token", 2049)
-        self.audio_stop_token = kwargs.get("audio_stop_token", 2050)
-        self.audio_spk_token = kwargs.get("audio_spk_token", 2051)
-        self.spk_emb_size = kwargs.get("spk_emb_size", 192)
-        self.think_end_ids = kwargs.get("think_end_ids", [26, 234, 234])
-        self.image_ids = kwargs.get("image_ids", [12])
-        self.image_special_token = kwargs.get("image_special_token", "<|image_pad|>")
-        self.image_hidden_size = kwargs.get("image_hidden_size", 768)
-        self.image_token_len = kwargs.get("image_token_len", 64)
-        self.bridge_layer = kwargs.get("bridge_layer", self.num_hidden_layers // 2 - 1)
-
-
 class TalkerHead(nn.Module):
     def __init__(self, in_features, out_features, num_layers=8, rank=256):
         super().__init__()

src/omni/models/vlm/__init__.py

@@ -0,0 +1,4 @@
+from omni.models.vlm.config import VLMConfig
+from omni.models.vlm.model import MiniMindVLM
+
+__all__ = ["VLMConfig", "MiniMindVLM"]

src/omni/models/vlm/config.py

@@ -0,0 +1,12 @@
+from omni.models.lm.config import MiniMindConfig
+
+
+class VLMConfig(MiniMindConfig):
+    model_type = "minimind-v"
+
+    def __init__(self, image_special_token='<|image_pad|>', image_ids=[12], **kwargs):
+        self.image_special_token = image_special_token
+        self.image_ids = image_ids
+        self.image_hidden_size = kwargs.get("image_hidden_size", 768)
+        self.image_token_len = kwargs.get("image_token_len", 64)
+        super().__init__(**kwargs)

src/omni/models/{vlm.py → vlm/model.py}

@@ -1,4 +1,3 @@
-import os
 import torch
 import torch.nn.functional as F
 import warnings
@@ -7,24 +6,14 @@ from torch import nn
 from transformers.modeling_outputs import MoeCausalLMOutputWithPast
 
 from omni.core import precompute_freqs_cis, MOEFeedForward
-from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.models.lm.model import MiniMindForCausalLM
+from omni.models.vlm.config import VLMConfig
 from omni.encoders.vision import SiglipVisionEncoder
 from omni.projectors import MMVisionProjector
 
 warnings.filterwarnings('ignore')
 
 
-class VLMConfig(MiniMindConfig):
-    model_type = "minimind-v"
-
-    def __init__(self, image_special_token='<|image_pad|>', image_ids=[12], **kwargs):
-        self.image_special_token = image_special_token
-        self.image_ids = image_ids
-        self.image_hidden_size = kwargs.get("image_hidden_size", 768)
-        self.image_token_len = kwargs.get("image_token_len", 64)
-        super().__init__(**kwargs)
-
-
 class MiniMindVLM(MiniMindForCausalLM):
     config_class = VLMConfig
 

src/omni/trainers/agent.py

@@ -19,7 +19,7 @@ from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from transformers import AutoTokenizer
-from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.models import MiniMindConfig, MiniMindForCausalLM
 from omni.datasets.lm_dataset import AgentRLDataset
 from omni.utils.training import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
 from omni.trainers.rollout_engine import create_rollout_engine, compute_per_token_logps

src/omni/trainers/distillation.py

@@ -12,7 +12,7 @@ from contextlib import nullcontext
 from torch import optim
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
-from omni.models.minimind import MiniMindConfig
+from omni.models import MiniMindConfig
 from omni.datasets.lm_dataset import SFTDataset
 from omni.utils.training import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
 

src/omni/trainers/dpo.py

@@ -12,7 +12,7 @@ from contextlib import nullcontext
 from torch import optim
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
-from omni.models.minimind import MiniMindConfig
+from omni.models import MiniMindConfig
 from omni.datasets.lm_dataset import DPODataset
 from omni.utils.training import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
 

src/omni/trainers/full_sft.py

@@ -11,7 +11,7 @@ from contextlib import nullcontext
 from torch import optim, nn
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
-from omni.models.minimind import MiniMindConfig
+from omni.models import MiniMindConfig
 from omni.datasets.lm_dataset import SFTDataset
 from omni.utils.training import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
 

src/omni/trainers/grpo.py

@@ -17,7 +17,7 @@ from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from transformers import AutoModel
-from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.models import MiniMindConfig, MiniMindForCausalLM
 from omni.datasets.lm_dataset import RLAIFDataset
 from omni.utils.training import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
 from omni.trainers.rollout_engine import create_rollout_engine

src/omni/trainers/lora.py

@@ -11,7 +11,7 @@ from contextlib import nullcontext
 from torch import optim, nn
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
-from omni.models.minimind import MiniMindConfig
+from omni.models import MiniMindConfig
 from omni.datasets.lm_dataset import SFTDataset
 from omni.models.lora import save_lora, apply_lora
 from omni.utils.training import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler

src/omni/trainers/ppo.py

@@ -16,7 +16,7 @@ from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
 from torch.nn.utils import clip_grad_norm_
 from torch.optim.lr_scheduler import CosineAnnealingLR
-from omni.models.minimind import MiniMindConfig, MiniMindForCausalLM
+from omni.models import MiniMindConfig, MiniMindForCausalLM
 from omni.datasets.lm_dataset import RLAIFDataset
 from omni.utils.training import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel
 from omni.trainers.rollout_engine import create_rollout_engine

src/omni/trainers/pretrain.py

@@ -11,7 +11,7 @@ from contextlib import nullcontext
 from torch import optim, nn
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
-from omni.models.minimind import MiniMindConfig
+from omni.models import MiniMindConfig
 from omni.datasets.lm_dataset import PretrainDataset
 from omni.utils.training import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler
 

src/omni/utils/multimodal.py

@@ -5,7 +5,7 @@ import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel
 
 from omni.utils.training import Logger, is_main_process
-from omni.models import MiniMindOmni
+from omni.models import MiniMindOmni, MiniMindVLM
 
 
 def get_vlm_model_params(model, config, ignore_patterns=('vision_encoder',)):
@@ -28,7 +28,6 @@ def get_vlm_model_params(model, config, ignore_patterns=('vision_encoder',)):
 
 def init_vlm_model(vlm_config, from_weight='pretrain_vlm', tokenizer_path='../model', vision_model_path='../model/siglip2-base-p32-256-ve', save_dir='../out', device='cuda', freeze_llm=0):
     from transformers import AutoTokenizer
-    from omni.models.vlm import MiniMindVLM
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
     model = MiniMindVLM(vlm_config, vision_model_path=vision_model_path)
 

src/omni/utils/training.py

@@ -8,7 +8,7 @@ import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import Sampler
 from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
-from omni.models.minimind import MiniMindForCausalLM
+from omni.models import MiniMindForCausalLM
 
 
 def get_model_params(model, config):