"""Preserve codec buffer values while giving CPU storage views their own storage."""


def instantiate_cpu_model(model_type, config):
    import amphion_utils

    original_load = amphion_utils.load_model

    def load_codec(module, *args, **kwargs):
        for name, buffer in list(module.named_buffers(remove_duplicate=False)):
            if buffer.device.type == "cpu" and (
                buffer.storage_offset() != 0 or
                buffer.untyped_storage().nbytes() != buffer.numel() * buffer.element_size()
            ):
                owner, _, key = name.rpartition(".")
                setattr(module.get_submodule(owner) if owner else module, key, buffer.clone())
        return original_load(module, *args, **kwargs)

    # The CLI and build check construct one model at a time. Always restore the
    # source package's loader, including when loading raises an exception.
    amphion_utils.load_model = load_codec
    try:
        return model_type(config)
    finally:
        amphion_utils.load_model = original_load