36 lines
1.0 KiB
Python
36 lines
1.0 KiB
Python
import torch
|
|
|
|
from .configuration import MossMusicConfig, MossMusicEncoderConfig
|
|
from .modeling import MossMusicModel
|
|
from .processing import MelConfig, MossMusicProcessor
|
|
|
|
HF_REPO = "OpenMOSS-Team/MOSS-Music-8B-Instruct"
|
|
|
|
|
|
def load_moss_music(
|
|
name_or_path: str = HF_REPO,
|
|
dtype: torch.dtype = torch.bfloat16,
|
|
device="cpu",
|
|
enable_time_marker: bool = True,
|
|
):
|
|
"""Model + processor from a hub id or local checkpoint dir. Time markers
|
|
(elapsed seconds interleaved into the audio tokens) are what the model was
|
|
trained with and keep long transcriptions from looping; leave them on."""
|
|
model = MossMusicModel.from_pretrained(name_or_path, dtype=dtype, device_map=device)
|
|
model.eval()
|
|
processor = MossMusicProcessor.from_pretrained(
|
|
name_or_path, enable_time_marker=enable_time_marker
|
|
)
|
|
return model, processor
|
|
|
|
|
|
__all__ = [
|
|
"HF_REPO",
|
|
"MelConfig",
|
|
"MossMusicConfig",
|
|
"MossMusicEncoderConfig",
|
|
"MossMusicModel",
|
|
"MossMusicProcessor",
|
|
"load_moss_music",
|
|
]
|