"""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)