File size: 3,360 Bytes
7b592f7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | import random
from pathlib import Path
import numpy as np
import torch
from transformers import AutoConfig
from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
from src.audiointeraction.generate.base import AUDIO_TOKENS_PER_CHUNK # noqa: F401 (re-export for callers)
from src.audiointeraction.model import GPT, Config
from src.audiointeraction.utils import load_checkpoint
def set_seed(seed: int = 1337) -> None:
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
import json
from pathlib import Path
from safetensors.torch import load_file
def load_model(fabric, model_config_dir, checkpoint_dir):
"""Load a GPT from a local sharded safetensors directory.
`checkpoint_dir` must contain:
model.safetensors.index.json
model-00001-of-0000N.safetensors
...
"""
config = Config.from_file(Path(model_config_dir) / "model_config.yaml")
with fabric.init_module(empty_init=(fabric.world_size > 1)):
model = GPT(config)
model = fabric.setup(model)
checkpoint_dir = Path(checkpoint_dir)
index_path = checkpoint_dir / "model.safetensors.index.json"
if not index_path.is_file():
raise FileNotFoundError(
f"No model.safetensors.index.json under {checkpoint_dir}. "
f"Expected a sharded safetensors directory."
)
with open(index_path) as f:
index = json.load(f)
shard_files = sorted(set(index["weight_map"].values()))
state_dict = {}
for shard in shard_files:
state_dict.update(load_file(str(checkpoint_dir / shard), device="cpu"))
missing, unexpected = model.load_state_dict(state_dict, strict=True)
if missing or unexpected:
print(f"[load_model] missing={missing[:3]}… unexpected={unexpected[:3]}…")
return model
def load_audio_encoder(qwen_omni_ckpt, audio_tower_ckpt, device):
print(qwen_omni_ckpt)
cfg = AutoConfig.from_pretrained(qwen_omni_ckpt)
# Omni 的 config 是嵌套的：thinker_config.audio_config 才是 audio encoder 的配置
audio_cfg = cfg.thinker_config.audio_config
encoder = Qwen2_5OmniAudioEncoder._from_config(audio_cfg)
state_dict = torch.load(audio_tower_ckpt, map_location=device)
encoder.load_state_dict(state_dict)
encoder.to(device).requires_grad_(False).eval()
return encoder
def resolve_checkpoint_paths(checkpoint_dir: str):
"""Map a single checkpoint root → (model_config_dir, trained_checkpoint,
qwen_omni_ckpt, audio_tower_ckpt). The release layout is:
<checkpoint_dir>/
model_config.yaml + tokenizer.json + ... ← model_config_dir = root
audiointeraction_LM.pt
audiointeraction_ChunkwisedEncoder.pth
qwen_2_5_omni_config/
"""
ckpt = Path(checkpoint_dir)
return (
str(ckpt),
str(ckpt),
str(ckpt / "qwen25OmniConfig"),
str(ckpt / "audiointeraction_ChunkwisedEncoder.pth"),
)
def get_best_device():
if torch.cuda.is_available():
return torch.device("cuda")
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")
|