| """LFM3-Audio encoder and projector, with source checkpoint key conversion.""" |
| import importlib,json,sys,types |
| from pathlib import Path |
| import torch |
| from accelerate import init_empty_weights |
| from safetensors.torch import load_file |
|
|
| def native(root): |
| root=Path(root).resolve();name='_lfm3_audio_encoder_native' |
| if name not in sys.modules: |
| p=types.ModuleType(name);p.__path__=[str(root)];sys.modules[name]=p |
| elif sys.modules[name].__path__ != [str(root)]: |
| raise RuntimeError('A different native source package is already loaded') |
| return (importlib.import_module(name+'.configuration_lfm3_audio'),importlib.import_module(name+'.modeling_lfm3_audio'),importlib.import_module(name+'.feature_extraction_lfm3_audio')) |
|
|
| def load_model(root): |
| root=Path(root);cfg,impl,_=native(root);c=cfg.Lfm3AudioConfig(**json.loads((root/'config.json').read_text())) |
| with init_empty_weights(include_buffers=False): |
| m=torch.nn.Module();m.model=torch.nn.Module();m.model.audio_encoder=impl.Lfm3AudioEncoder(c.audio_encoder_config);m.model.audio_projector=impl.Lfm3AudioProjector(c) |
| raw=load_file(str(root/'model.safetensors'),device='cpu');rules=(('encoder.downsample.','model.audio_encoder.downsample.'),('encoder.lfm.lfm.','model.audio_encoder.lfm.'),('mlp.model.0.','model.audio_projector.layer_norm.'),('mlp.model.1.','model.audio_projector.linear_1.'),('mlp.model.3.','model.audio_projector.linear_2.')) |
| state={} |
| for key,val in raw.items(): |
| candidates=[dst+key[len(src):] for src,dst in rules if key.startswith(src)] |
| if len(candidates)!=1 or candidates[0] in state:raise ValueError('Unmapped or duplicate key: '+key) |
| state[candidates[0]]=val |
| m.load_state_dict(state,strict=True,assign=True) |
| return m.eval() |
|
|
| def load_feature_extractor(root): |
| *_,features=native(root);c=json.loads((Path(root)/'preprocessor_config.json').read_text());c.pop('auto_map',None);return features.Lfm3AudioFeatureExtractor(**c) |
|
|
