Convert omni-o.pth to HF format + add from_pretrained support

- scripts/convert_omni_o_to_hf.py: convert .pth -> config.json + model.safetensors
- src/models/vam/model.py: from_pretrained override loads encoders after meta init
- src/models/vam/model.py: __init__ skips encoder loading on meta device
- scripts/omni_o_call.py: auto-detect HF dir and load via VAM.from_pretrained
- auto_map in config.json points to modeling_omni_o.py for trust_remote_code

scripts/convert_omni_o_to_hf.py

@@ -0,0 +1,98 @@
+import os, sys, json, argparse, shutil, torch, warnings
+warnings.filterwarnings('ignore')
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
+from models.vam import VAM, VAMConfig
+from transformers import AutoTokenizer
+
+
+def infer_config(state_dict):
+    hs = state_dict['model.embed_tokens.weight'].shape[1]
+    nl = max(int(k.split('.')[2]) for k in state_dict if k.startswith('model.layers.')) + 1
+    use_moe = any('expert' in k for k in state_dict if 'mlp' in k)
+    return dict(hidden_size=hs, num_hidden_layers=nl, use_moe=use_moe,
+                num_attention_heads=hs // 96, num_key_value_heads=hs // 192)
+
+
+def convert(torch_path, output_dir, tokenizer_dir='checkpoint/omni/native_hf',
+            sensevoice_dir='checkpoint/sensevoice', siglip_dir='checkpoint/siglip',
+            dtype=torch.float16, device='cpu'):
+    os.makedirs(output_dir, exist_ok=True)
+
+    print(f'Loading checkpoint: {torch_path}')
+    state = torch.load(torch_path, map_location=device, weights_only=True)
+    cfg = infer_config(state)
+    print(f'  hidden_size={cfg["hidden_size"]}, num_hidden_layers={cfg["num_hidden_layers"]}, use_moe={cfg["use_moe"]}')
+
+    print('Creating model...')
+    config = VAMConfig(**cfg)
+    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    model = VAM(config,
+                audio_encoder_path=os.path.join(root, sensevoice_dir),
+                vision_model_path=os.path.join(root, siglip_dir))
+
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    if missing:
+        print(f'  Missing keys (expected for encoders): {len(missing)}')
+    if unexpected:
+        print(f'  Unexpected keys: {len(unexpected)}')
+
+    del state
+    model = model.to(dtype).half()
+
+    VAM.register_for_auto_class("AutoModelForCausalLM")
+    VAMConfig.register_for_auto_class()
+
+    print(f'Saving to {output_dir}...')
+    model.save_pretrained(output_dir, safe_serialization=True)
+
+    tokenizer_path = os.path.join(root, tokenizer_dir)
+    if os.path.exists(tokenizer_path):
+        for fn in ['tokenizer.json', 'tokenizer_config.json', 'generation_config.json', 'chat_template.jinja']:
+            src = os.path.join(tokenizer_path, fn)
+            if os.path.exists(src):
+                shutil.copy2(src, os.path.join(output_dir, fn))
+        print('Tokenzier files copied')
+    else:
+        print(f'Tokenzier not found at {tokenizer_path}, skipping')
+
+    modeling_code = '''import sys, os
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src"))
+from models.vam import VAM, VAMConfig
+'''
+    with open(os.path.join(output_dir, 'modeling_omni_o.py'), 'w') as f:
+        f.write(modeling_code)
+
+    config_path = os.path.join(output_dir, 'config.json')
+    if os.path.exists(config_path):
+        with open(config_path, 'r') as f:
+            cf = json.load(f)
+        cf['auto_map'] = {
+            "AutoConfig": "modeling_omni_o.VAMConfig",
+            "AutoModelForCausalLM": "modeling_omni_o.VAM",
+        }
+        cf['model_type'] = 'omni-o'
+        with open(config_path, 'w') as f:
+            json.dump(cf, f, indent=2, ensure_ascii=False)
+
+    params = sum(p.numel() for p in model.parameters()) / 1e6
+    print(f'Done! Model params: {params:.2f}M')
+    print(f'HF model saved to: {output_dir}')
+    print(f'  VAM.from_pretrained("{output_dir}", audio_encoder_path=..., vision_model_path=...)')
+    print(f'  AutoModelForCausalLM.from_pretrained("{output_dir}", trust_remote_code=True)')
+
+
+if __name__ == '__main__':
+    p = argparse.ArgumentParser(description='Convert omni-o .pth to HuggingFace format')
+    p.add_argument('torch_path', help='Path to .pth checkpoint')
+    p.add_argument('output_dir', help='Output directory for HF model')
+    p.add_argument('--tokenizer_dir', default='checkpoint/omni/native_hf')
+    p.add_argument('--sensevoice_dir', default='checkpoint/sensevoice')
+    p.add_argument('--siglip_dir', default='checkpoint/siglip')
+    p.add_argument('--dtype', default='float16')
+    p.add_argument('--device', default='cpu')
+    args = p.parse_args()
+
+    convert(args.torch_path, args.output_dir, args.tokenizer_dir,
+            args.sensevoice_dir, args.siglip_dir,
+            getattr(torch, args.dtype), args.device)

scripts/omni_o_call.py

@@ -294,35 +294,46 @@ def init_model(args):
         from funasr import AutoModel
         M['asr'] = AutoModel(model=os.path.join(root, args.sensevoice_dir), trust_remote_code=True, device=args.device, disable_update=True)
 
-    config = VAMConfig(
-        hidden_size=args.hidden_size,
-        num_hidden_layers=args.num_hidden_layers,
-        num_attention_heads=args.hidden_size // 96,
-        num_key_value_heads=args.hidden_size // 192,
-        use_moe=args.use_moe,
-    )
     ckpt_dir = os.path.join(root, args.load_from)
-    weight = args.weight
-    if not weight.endswith('.pth'):
-        if args.use_moe and not weight.endswith('_moe'):
-            weight = f'{weight}_moe.pth'
-        else:
-            weight = f'{weight}.pth'
-    ckpt_path = os.path.join(ckpt_dir, weight)
-
-    model = VAM(config,
-                audio_encoder_path=os.path.join(root, args.sensevoice_dir),
-                vision_model_path=os.path.join(root, args.siglip_dir))
-    state = torch.load(ckpt_path, map_location='cpu', weights_only=True)
-    missing, unexpected = model.load_state_dict(state, strict=False)
-    if missing:
-        print(f'  Missing keys (expected for encoders): {len(missing)}')
-    if unexpected:
-        print(f'  Unexpected keys: {len(unexpected)}')
-    if model.audio_encoder is not None:
-        model.audio_encoder.to(args.device)
-    if model.vision_encoder is not None:
-        model.vision_encoder.to(args.device)
+    is_hf = os.path.exists(os.path.join(ckpt_dir, 'config.json')) and \
+            (os.path.exists(os.path.join(ckpt_dir, 'model.safetensors')) or
+             os.path.exists(os.path.join(ckpt_dir, 'pytorch_model.bin')))
+
+    if is_hf:
+        model = VAM.from_pretrained(
+            ckpt_dir,
+            audio_encoder_path=os.path.join(root, args.sensevoice_dir),
+            vision_model_path=os.path.join(root, args.siglip_dir),
+        )
+    else:
+        config = VAMConfig(
+            hidden_size=args.hidden_size,
+            num_hidden_layers=args.num_hidden_layers,
+            num_attention_heads=args.hidden_size // 96,
+            num_key_value_heads=args.hidden_size // 192,
+            use_moe=args.use_moe,
+        )
+        weight = args.weight
+        if not weight.endswith('.pth'):
+            if args.use_moe and not weight.endswith('_moe'):
+                weight = f'{weight}_moe.pth'
+            else:
+                weight = f'{weight}.pth'
+        ckpt_path = os.path.join(ckpt_dir, weight)
+        model = VAM(config,
+                    audio_encoder_path=os.path.join(root, args.sensevoice_dir),
+                    vision_model_path=os.path.join(root, args.siglip_dir))
+        state = torch.load(ckpt_path, map_location='cpu', weights_only=True)
+        missing, unexpected = model.load_state_dict(state, strict=False)
+        if missing:
+            print(f'  Missing keys (expected for encoders): {len(missing)}')
+        if unexpected:
+            print(f'  Unexpected keys: {len(unexpected)}')
+        if model.audio_encoder is not None:
+            model.audio_encoder.to(args.device)
+        if model.vision_encoder is not None:
+            model.vision_encoder.to(args.device)
+
     M['model'] = model.half().eval().to(args.device)
 
     tok_dir = os.path.join(root, args.tokenizer_dir)

src/models/vam/model.py

@@ -91,12 +91,34 @@ class VAM(LMForCausalLM):
         self.audio_proj = MMAudioProjector(config.audio_hidden_size, config.hidden_size)
         self.vision_proj = MMVisionProjector(config.image_hidden_size, config.hidden_size, target_tokens=config.image_token_len)
         self.audio_pad_token, self.audio_stop_token, self.audio_spk_token = config.audio_pad_token, config.audio_stop_token, config.audio_spk_token
-        audio_encoder = SenseVoiceAudioEncoder(audio_encoder_path) if audio_encoder_path else SenseVoiceAudioEncoder()
-        object.__setattr__(self, 'audio_encoder', audio_encoder)
-        object.__setattr__(self, 'audio_processor', audio_encoder.processor)
-        vision_encoder = SiglipVisionEncoder(vision_model_path) if vision_model_path else SiglipVisionEncoder()
-        object.__setattr__(self, 'vision_encoder', vision_encoder)
-        object.__setattr__(self, 'vision_processor', vision_encoder.processor)
+        meta_init = any(p.device.type == 'meta' for p in self.parameters())
+        if meta_init:
+            object.__setattr__(self, 'audio_encoder', None)
+            object.__setattr__(self, 'audio_processor', None)
+            object.__setattr__(self, 'vision_encoder', None)
+            object.__setattr__(self, 'vision_processor', None)
+        else:
+            audio_enc = SenseVoiceAudioEncoder(audio_encoder_path) if audio_encoder_path else SenseVoiceAudioEncoder()
+            object.__setattr__(self, 'audio_encoder', audio_enc)
+            object.__setattr__(self, 'audio_processor', audio_enc.processor)
+            vision_enc = SiglipVisionEncoder(vision_model_path) if vision_model_path else SiglipVisionEncoder()
+            object.__setattr__(self, 'vision_encoder', vision_enc)
+            object.__setattr__(self, 'vision_processor', vision_enc.processor)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        audio_encoder_path = kwargs.pop('audio_encoder_path', None)
+        vision_model_path = kwargs.pop('vision_model_path', None)
+        model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        if audio_encoder_path and model.audio_encoder is None:
+            enc, proc = cls.load_sensevoice(audio_encoder_path)
+            object.__setattr__(model, 'audio_encoder', enc)
+            object.__setattr__(model, 'audio_processor', proc)
+        if vision_model_path and model.vision_encoder is None:
+            enc, proc = cls.load_vision(vision_model_path)
+            object.__setattr__(model, 'vision_encoder', enc)
+            object.__setattr__(model, 'vision_processor', proc)
+        return model
 
     @staticmethod
     def load_sensevoice(path):