Function export

kleidiai-examples/audiogen/scripts/export_sao.py:120–151

(args)

        logging.info("AutoEncoder Encoder model has been saved to %s", os.path.abspath(os.path.join(output_path, "autoencoder_encoder_model.tflite")))

def export(args) -> None:

    torch.manual_seed(0)
    device = torch.device("cpu")

    # Load the model configuration
    logging.info("Loading the AudioGen Checkpoint...")
    with open(args.model_config, encoding="utf-8") as f:
        model_config = json.load(f)

    # Load the model
    model, model_config = load_model(
        model_config = model_config,
        model_ckpt_path = args.ckpt_path,
        pretrained_name=None,
        device=device,
    )
    logging.info("Model is loaded...")

    # --------- Conditioners Model ---------
    export_conditioners(model, args.output_path)

    # --------- DiT Model ----------------
    export_dit(model, args.output_path)

    # --------- AutoEncoder Model ---------

    # Removing weight norm from the model as it is causing issues during export
    remove_weight_norm_from_model(model.pretransform)

    export_autoencoder(model, args.output_path)
    export_autoencoder_encoder(model, args.output_path)

def main():
    parser = argparse.ArgumentParser()